Fixed width inference bug where later constraints on the output width were not propogating to the input widths, for primops

2 files changed, 27 insertions, 8 deletions

diff --git a/src/main/stanza/primop.stanza b/src/main/stanza/primop.stanza
index 14589135..a1e9633f 100644
--- a/src/main/stanza/primop.stanza
+++ b/src/main/stanza/primop.stanza
@@ -248,12 +248,12 @@ public defn lower-and-type-primop (e:DoPrim) -> DoPrim :
       BIT-XOR-REDUCE-OP: DoPrim(op(e),args(e),consts(e),u())
 
 public defn primop-gen-constraints (e:DoPrim,v:Vector<WGeq>) -> Type :
-   defn all-equal (ls:List<Expression>) -> Width : 
-      if length(ls) == 1 : width!(ls[0])
+   defn all-equal (ls:List<Width>) -> Width : 
+      if length(ls) == 1 : (ls[0])
       else :
-         val m = MaxWidth(map(width!,ls))
+         val m = MaxWidth(ls)
          for (l in ls) do :
-            add(v,WGeq(width!(l),m))
+            add(v,WGeq(l,m))
          m
    ;defn new-width (w:Width) -> Width:
    ;   val w* = VarWidth(gensym(`w))
@@ -275,9 +275,11 @@ public defn primop-gen-constraints (e:DoPrim,v:Vector<WGeq>) -> Type :
    val all-args-not-equal = list(MUX-UU-OP,MUX-SS-OP,CONCAT-OP)
    ;val consts-gte-args    = list(PAD-U-OP,PAD-S-OP)
 
+   
+   val w-var = VarWidth(gensym(`w))
    val w* = 
       if not contains?(all-args-not-equal,op(e)) :
-         val max-args-w = all-equal(args(e))
+         val max-args-w = all-equal(List(w-var,map(width!,args(e))))
          switch {op(e) == _} :
             ADD-UU-OP : PlusWidth(max-args-w,IntWidth(1))
             ADD-US-OP : PlusWidth(max-args-w,IntWidth(1))
@@ -347,14 +349,13 @@ public defn primop-gen-constraints (e:DoPrim,v:Vector<WGeq>) -> Type :
          switch {op(e) == _} :
             MUX-UU-OP : 
                add(v,WGeq(width!(args(e)[0]),IntWidth(1)))
-               all-equal(tail(args(e)))
+               all-equal(List(w-var,tail(map(width!,args(e)))))
             MUX-SS-OP : 
                add(v,WGeq(width!(args(e)[0]),IntWidth(1)))
-               all-equal(tail(args(e)))
+               all-equal(List(w-var,tail(map(width!,args(e)))))
             CONCAT-OP : 
                PlusWidth(width!(args(e)[0]),width!(args(e)[1]))
 
-   val w-var = VarWidth(gensym(`w))
    add(v,WGeq(w-var,w*))
    match(type(e)) :
       (t:UIntType) : UIntType(w-var)
diff --git a/test/passes/jacktest/Tlb.fir b/test/passes/jacktest/Tlb.fir
new file mode 100644
index 00000000..35442ac8
--- /dev/null
+++ b/test/passes/jacktest/Tlb.fir
@@ -0,0 +1,18 @@
+; RUN: firrtl -i %s -o %s.flo -x X -p cTwd | tee %s.out | FileCheck %s
+; CHECK: Done!
+circuit Tbl : 
+  module Tbl : 
+    output o : UInt(16)
+    input i : UInt(16)
+    input d : UInt(16)
+    input we : UInt(1)
+    
+    mem m : UInt(10)[256]
+    node T_12 = UInt(0, 1)
+    o := Pad(T_12,?)
+    when we : 
+      accessor T_13 = m[i]
+      T_13 := Pad(d,?)
+    else : 
+      accessor T_14 = m[i]
+      o := Pad(T_14,?)