Add constant prop to async regs (#1355)

* Add constant prop to async regs

* Added another test of no reset value but constant assignment

* Clarify name of updateNodeMap

* Update constant assignment of async reset to not be inferred as a latch, works with donttouch

* Revert "Update constant assignment of async reset to not be inferred as a latch, works with donttouch"

This reverts commit 952bf38127cb32f814496a2b4b3bfb173d532728.

2 files changed, 62 insertions, 7 deletions

diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index 20b24e60..66cf24b8 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -406,7 +406,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
     // (can have more than 1 of the same submodule)
     val constSubInputs = mutable.HashMap.empty[OfModule, mutable.HashMap[String, Seq[Literal]]]
     // AsyncReset registers don't have reset turned into a mux so we must be careful
-    val asyncResetRegs = mutable.HashSet.empty[String]
+    val asyncResetRegs = mutable.HashMap.empty[String, DefRegister]
 
     // Register constant propagation is intrinsically more complicated, as it is not feed-forward.
     // Therefore, we must store some memoized information about how nodes can be canditates for
@@ -468,7 +468,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       stmtx match {
         case x: DefNode if !dontTouches.contains(x.name) => propagateRef(x.name, x.value)
         case reg: DefRegister if reg.reset.tpe == AsyncResetType =>
-          asyncResetRegs += reg.name
+          asyncResetRegs(reg.name) = reg
         case Connect(_, WRef(wname, wtpe, WireKind, _), expr: Literal) if !dontTouches.contains(wname) =>
           val exprx = constPropExpression(nodeMap, instMap, constSubOutputs)(pad(expr, wtpe))
           propagateRef(wname, exprx)
@@ -478,7 +478,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
           constOutputs(pname) = paddedLit
         // Const prop registers that are driven by a mux tree containing only instances of one constant or self-assigns
         // This requires that reset has been made explicit
-        case Connect(_, lref @ WRef(lname, ltpe, RegKind, _), rhs) if !dontTouches(lname) && !asyncResetRegs(lname) =>
+        case Connect(_, lref @ WRef(lname, ltpe, RegKind, _), rhs) if !dontTouches(lname) =>
 
          /* Checks if an RHS expression e of a register assignment is convertible to a constant assignment.
           * Here, this means that e must be 1) a literal, 2) a self-connect, or 3) a mux tree of
@@ -500,9 +500,9 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
             case Mux(_, tval, fval, _) => regConstant(tval).resolve(regConstant(fval))
             case _ => RegCPEntry(NonConstant, NonConstant)
           }
-          def zero = passes.RemoveValidIf.getGroundZero(ltpe)
-          def padCPExp(e: Expression) = constPropExpression(nodeMap, instMap, constSubOutputs)(pad(e, ltpe))
-          regConstant(rhs) match {
+
+          // Updates nodeMap after analyzing the returned value from regConstant
+          def updateNodeMapIfConstant(e: Expression): Unit = regConstant(e) match {
             case RegCPEntry(BoundConstant(`lname`), litBinding) => litBinding match {
               case UnboundConstant => nodeMap(lname) = padCPExp(zero) // only self-assigns -> replace with zero
               case BoundConstant(lit) => nodeMap(lname) = padCPExp(lit) // self + lit assigns -> replace with lit
@@ -511,6 +511,16 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
             case RegCPEntry(UnboundConstant, BoundConstant(lit)) => nodeMap(lname) = padCPExp(lit) // only lit assigns
             case _ =>
           }
+          def zero = passes.RemoveValidIf.getGroundZero(ltpe)
+          def padCPExp(e: Expression) = constPropExpression(nodeMap, instMap, constSubOutputs)(pad(e, ltpe))
+
+          asyncResetRegs.get(lname) match {
+            // Normal Register
+            case None => updateNodeMapIfConstant(rhs)
+            // Async Register
+            case Some(reg: DefRegister) => updateNodeMapIfConstant(Mux(reg.reset, reg.init, rhs))
+          }
+
         // Mark instance inputs connected to a constant
         case Connect(_, lref @ WSubField(WRef(inst, _, InstanceKind, _), port, ptpe, _), lit: Literal) =>
           val paddedLit = constPropExpression(nodeMap, instMap, constSubOutputs)(pad(lit, ptpe)).asInstanceOf[Literal]
diff --git a/src/test/scala/firrtlTests/AsyncResetSpec.scala b/src/test/scala/firrtlTests/AsyncResetSpec.scala
index 98d51ac8..ebc94cc8 100644
--- a/src/test/scala/firrtlTests/AsyncResetSpec.scala
+++ b/src/test/scala/firrtlTests/AsyncResetSpec.scala
@@ -319,6 +319,52 @@ class AsyncResetSpec extends FirrtlFlatSpec {
     )
   }
 
+  "Unassigned asyncronously reset registers" should "properly constantprop" in {
+    val result = compileBody(
+      s"""
+         |input clock : Clock
+         |input reset : AsyncReset
+         |output z : UInt<1>[4]
+         |wire literal : UInt<1>[2]
+         |literal[0] <= UInt<1>("h01")
+         |literal[1] <= UInt<1>("h01")
+         |wire complex_literal : UInt<1>[4]
+         |complex_literal[0] <= literal[0]
+         |complex_literal[1] <= literal[1]
+         |complex_literal[2] <= UInt<1>("h00")
+         |complex_literal[3] <= UInt<1>("h00")
+         |reg r : UInt<1>[4], clock with : (reset => (reset, complex_literal))
+         |z <= r""".stripMargin
+    )
+    result shouldNot containLine("always @(posedge clock or posedge reset) begin")
+  }
+
+  "Constantly assigned asynchronously reset registers" should "properly constantprop" in {
+    val result = compileBody(
+      s"""
+         |input clock : Clock
+         |input reset : AsyncReset
+         |output z : UInt<1>
+         |reg r : UInt<1>, clock with : (reset => (reset, r))
+         |r <= UInt(0)
+         |z <= r""".stripMargin
+    )
+    result shouldNot containLine("always @(posedge clock or posedge reset) begin")
+  }
+
+  "Constantly assigned and initialized asynchronously reset registers" should "properly constantprop" in {
+    val result = compileBody(
+      s"""
+         |input clock : Clock
+         |input reset : AsyncReset
+         |output z : UInt<1>
+         |reg r : UInt<1>, clock with : (reset => (reset, UInt(0)))
+         |r <= UInt(0)
+         |z <= r""".stripMargin
+    )
+    result shouldNot containLine("always @(posedge clock or posedge reset) begin")
+  }
+
   "AsyncReset registers" should "emit 'else' case for reset even for trivial valued registers" in {
     val withDontTouch = s"""
       |circuit m :
@@ -341,7 +387,6 @@ class AsyncResetSpec extends FirrtlFlatSpec {
      )
 
   }
-
 }
 
 class AsyncResetExecutionTest extends ExecutionTest("AsyncResetTester", "/features")