Merge pull request #621 from freechipsproject/const-prop-regs

Const prop regs

3 files changed, 140 insertions, 4 deletions

diff --git a/src/main/scala/firrtl/passes/RemoveValidIf.scala b/src/main/scala/firrtl/passes/RemoveValidIf.scala
index 35323cbc..06f0874a 100644
--- a/src/main/scala/firrtl/passes/RemoveValidIf.scala
+++ b/src/main/scala/firrtl/passes/RemoveValidIf.scala
@@ -4,6 +4,8 @@ package firrtl
 package passes
 import firrtl.Mappers._
 import firrtl.ir._
+import Utils.throwInternalError
+import WrappedExpression.weq
 
 /** Remove ValidIf and replace IsInvalid with a connection to zero */
 object RemoveValidIf extends Pass {
@@ -18,15 +20,24 @@ object RemoveValidIf extends Pass {
   private val SIntZero = SIntLiteral(BigInt(0), IntWidth(1))
   private val ClockZero = DoPrim(PrimOps.AsClock, Seq(UIntZero), Seq.empty, UIntZero.tpe)
 
+  private def getGroundZero(tpe: Type): Expression = tpe match {
+    case _: UIntType => UIntZero
+    case _: SIntType => SIntZero
+    case ClockType => ClockZero
+    case other => throwInternalError
+  }
+
   // Recursive. Replaces IsInvalid with connecting zero
   private def onStmt(s: Statement): Statement = s map onStmt map onExp match {
     case invalid @ IsInvalid(info, loc) => loc.tpe match {
-      case _: UIntType => Connect(info, loc, UIntZero)
-      case _: SIntType => Connect(info, loc, SIntZero)
       case _: AnalogType => invalid // Unclear what we should do, can't remove or we emit invalid Firrtl
-      case ClockType => Connect(info, loc, ClockZero)
-      case other => throw new Exception("Unexpected type ${other.serialize} on LowFirrtl")
+      case tpe => Connect(info, loc, getGroundZero(tpe))
     }
+    // Register connected to itself (since reset has been made explicit) is a register with no reset
+    // and no connections, connect it to zero (to be constant propped later)
+    case Connect(info, lref: WRef, rref: WRef) if weq(lref, rref) =>
+      // We can't have an Analog reg so just get a zero
+      Connect(info, lref, getGroundZero(lref.tpe))
     case other => other
   }
 
diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index 31a6a660..bf8b1a55 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -9,6 +9,7 @@ import firrtl.ir._
 import firrtl.Utils._
 import firrtl.Mappers._
 import firrtl.PrimOps._
+import firrtl.WrappedExpression.weq
 
 import annotation.tailrec
 import collection.mutable
@@ -317,6 +318,18 @@ class ConstantPropagation extends Transform {
         case Connect(_, WRef(wname, wtpe, WireKind, _), expr) if !dontTouches.contains(wname) =>
           val exprx = constPropExpression(pad(expr, wtpe))
           propagateRef(wname, exprx)
+        // Const prop registers that are fed only a constant or a mux between and constant and the
+        // register itself
+        // This requires that reset has been made explicit
+        case Connect(_, rref @ WRef(rname, rtpe, RegKind, _), expr) => expr match {
+          case lit: Literal =>
+            nodeMap(rname) = constPropExpression(pad(lit, rtpe))
+          case Mux(_, tval: WRef, fval: Literal, _) if weq(rref, tval) =>
+            nodeMap(rname) = constPropExpression(pad(fval, rtpe))
+          case Mux(_, tval: Literal, fval: WRef, _) if weq(rref, fval) =>
+            nodeMap(rname) = constPropExpression(pad(tval, rtpe))
+          case _ =>
+        }
         case _ =>
       }
       stmtx
diff --git a/src/test/scala/firrtlTests/ConstantPropagationTests.scala b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
index 8f09ac9e..75c43cf2 100644
--- a/src/test/scala/firrtlTests/ConstantPropagationTests.scala
+++ b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
@@ -514,4 +514,116 @@ class ConstantPropagationIntegrationSpec extends LowTransformSpec {
           |    z <= UInt<1>(0)""".stripMargin
     execute(input, check, Seq.empty)
   }
+
+  it should "pad constant connections to wires when propagating" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    output z : UInt<16>
+          |    wire w : { a : UInt<8>, b : UInt<8> }
+          |    w.a <= UInt<2>("h3")
+          |    w.b <= UInt<2>("h3")
+          |    z <= cat(w.a, w.b)""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    output z : UInt<16>
+          |    z <= UInt<16>("h303")""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  it should "pad constant connections to registers when propagating" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    output z : UInt<16>
+          |    reg r : { a : UInt<8>, b : UInt<8> }, clock
+          |    r.a <= UInt<2>("h3")
+          |    r.b <= UInt<2>("h3")
+          |    z <= cat(r.a, r.b)""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    output z : UInt<16>
+          |    z <= UInt<16>("h303")""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  "Registers with no reset or connections" should "be replaced with constant zero" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    output z : UInt<8>
+          |    reg r : UInt<8>, clock
+          |    z <= r""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    output z : UInt<8>
+          |    z <= UInt<8>(0)""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  "Registers with ONLY constant reset" should "be replaced with that constant" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : UInt<8>
+          |    reg r : UInt<8>, clock with : (reset => (reset, UInt<4>("hb")))
+          |    z <= r""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : UInt<8>
+          |    z <= UInt<8>("hb")""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  "Registers with ONLY constant connection" should "be replaced with that constant" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : SInt<8>
+          |    reg r : SInt<8>, clock
+          |    r <= SInt<4>(-5)
+          |    z <= r""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : SInt<8>
+          |    z <= SInt<8>(-5)""".stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  "Registers with identical constant reset and connection" should "be replaced with that constant" in {
+      val input =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : UInt<8>
+          |    reg r : UInt<8>, clock with : (reset => (reset, UInt<4>("hb")))
+          |    r <= UInt<4>("hb")
+          |    z <= r""".stripMargin
+      val check =
+        """circuit Top :
+          |  module Top :
+          |    input clock : Clock
+          |    input reset : UInt<1>
+          |    output z : UInt<8>
+          |    z <= UInt<8>("hb")""".stripMargin
+    execute(input, check, Seq.empty)
+  }
 }