Constant prop binary PrimOps with matching arguments

* Add SimplifyBinaryOp trait
* Add extra functionality to comparison folding
* Add tests
* Fix comments from review

2 files changed, 90 insertions, 3 deletions

diff --git a/src/main/scala/firrtl/transforms/ConstantPropagation.scala b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
index f450f6a6..201c3325 100644
--- a/src/main/scala/firrtl/transforms/ConstantPropagation.scala
+++ b/src/main/scala/firrtl/transforms/ConstantPropagation.scala
@@ -18,6 +18,11 @@ import annotation.tailrec
 import collection.mutable
 
 object ConstantPropagation {
+  private def litOfType(value: BigInt, t: Type): Literal = t match {
+    case UIntType(w) => UIntLiteral(value, w)
+    case SIntType(w) => SIntLiteral(value, w)
+  }
+
   private def asUInt(e: Expression, t: Type) = DoPrim(AsUInt, Seq(e), Seq(), t)
 
   /** Pads e to the width of t */
@@ -99,14 +104,22 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
 
   override val annotationClasses: Traversable[Class[_]] = Seq(classOf[DontTouchAnnotation])
 
-  trait FoldCommutativeOp {
+  sealed trait SimplifyBinaryOp {
+    def matchingArgsValue(e: DoPrim, arg: Expression): Expression
+    def apply(e: DoPrim): Expression = {
+      if (e.args.head == e.args(1)) matchingArgsValue(e, e.args.head) else e
+    }
+  }
+
+  sealed trait FoldCommutativeOp extends SimplifyBinaryOp {
     def fold(c1: Literal, c2: Literal): Expression
     def simplify(e: Expression, lhs: Literal, rhs: Expression): Expression
 
-    def apply(e: DoPrim): Expression = (e.args.head, e.args(1)) match {
+    override def apply(e: DoPrim): Expression = (e.args.head, e.args(1)) match {
       case (lhs: Literal, rhs: Literal) => fold(lhs, rhs)
       case (lhs: Literal, rhs) => pad(simplify(e, lhs, rhs), e.tpe)
       case (lhs, rhs: Literal) => pad(simplify(e, rhs, lhs), e.tpe)
+      case (lhs, rhs) if (lhs == rhs) => matchingArgsValue(e, lhs)
       case _ => e
     }
   }
@@ -121,6 +134,19 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case SIntLiteral(v, w) if v == BigInt(0) => rhs
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = e
+  }
+
+  object SimplifySUB extends SimplifyBinaryOp {
+    def matchingArgsValue(e: DoPrim, arg: Expression) = litOfType(0, e.tpe)
+  }
+
+  object SimplifyDIV extends SimplifyBinaryOp {
+    def matchingArgsValue(e: DoPrim, arg: Expression) = litOfType(1, e.tpe)
+  }
+
+  object SimplifyREM extends SimplifyBinaryOp {
+    def matchingArgsValue(e: DoPrim, arg: Expression) = litOfType(0, e.tpe)
   }
 
   object FoldAND extends FoldCommutativeOp {
@@ -131,6 +157,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case UIntLiteral(v, IntWidth(w)) if v == (BigInt(1) << bitWidth(rhs.tpe).toInt) - 1 => rhs
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = asUInt(arg, e.tpe)
   }
 
   object FoldOR extends FoldCommutativeOp {
@@ -141,6 +168,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case UIntLiteral(v, IntWidth(w)) if v == (BigInt(1) << bitWidth(rhs.tpe).toInt) - 1 => lhs
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = asUInt(arg, e.tpe)
   }
 
   object FoldXOR extends FoldCommutativeOp {
@@ -150,6 +178,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case SIntLiteral(v, _) if v == BigInt(0) => asUInt(rhs, e.tpe)
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = UIntLiteral(0, getWidth(arg.tpe))
   }
 
   object FoldEqual extends FoldCommutativeOp {
@@ -159,6 +188,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case UIntLiteral(v, IntWidth(w)) if v == BigInt(0) && w == BigInt(1) && bitWidth(rhs.tpe) == BigInt(1) => DoPrim(Not, Seq(rhs), Nil, e.tpe)
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = UIntLiteral(1)
   }
 
   object FoldNotEqual extends FoldCommutativeOp {
@@ -168,6 +198,7 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
       case UIntLiteral(v, IntWidth(w)) if v == BigInt(1) && w == BigInt(1) && bitWidth(rhs.tpe) == BigInt(1) => DoPrim(Not, Seq(rhs), Nil, e.tpe)
       case _ => e
     }
+    def matchingArgsValue(e: DoPrim, arg: Expression) = UIntLiteral(0)
   }
 
   private def foldConcat(e: DoPrim) = (e.args.head, e.args(1)) match {
@@ -267,7 +298,16 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
         case ex => ex
       }
     }
-    foldIfZeroedArg(foldIfOutsideRange(e))
+
+    def foldIfMatchingArgs(x: Expression) = x match {
+      case DoPrim(op, Seq(a, b), _, _) if (a == b) => op match {
+        case (Lt | Gt) => zero
+        case (Leq | Geq) => one
+        case _ => x
+      }
+      case _ => x
+    }
+    foldIfZeroedArg(foldIfOutsideRange(foldIfMatchingArgs(e)))
   }
 
   private def constPropPrim(e: DoPrim): Expression = e.op match {
@@ -277,6 +317,9 @@ class ConstantPropagation extends Transform with ResolvedAnnotationPaths {
     case Dshr => foldDynamicShiftRight(e)
     case Cat => foldConcat(e)
     case Add => FoldADD(e)
+    case Sub => SimplifySUB(e)
+    case Div => SimplifyDIV(e)
+    case Rem => SimplifyREM(e)
     case And => FoldAND(e)
     case Or => FoldOR(e)
     case Xor => FoldXOR(e)
diff --git a/src/test/scala/firrtlTests/ConstantPropagationTests.scala b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
index ef52507f..cc7a5e32 100644
--- a/src/test/scala/firrtlTests/ConstantPropagationTests.scala
+++ b/src/test/scala/firrtlTests/ConstantPropagationTests.scala
@@ -1393,6 +1393,50 @@ class ConstantPropagationIntegrationSpec extends LowTransformSpec {
       """.stripMargin
     execute(input, check, Seq.empty)
   }
+
+  private def matchingArgs(op: String, iType: String, oType: String, result: String): Unit = {
+    val input =
+      s"""circuit Top :
+         |  module Top :
+         |    input i : ${iType}
+         |    output o : ${oType}
+         |    o <= ${op}(i, i)
+      """.stripMargin
+    val check =
+      s"""circuit Top :
+         |  module Top :
+         |    input i : ${iType}
+         |    output o : ${oType}
+         |    o <= ${result}
+      """.stripMargin
+    execute(input, check, Seq.empty)
+  }
+
+  it should "optimize some binary operations when arguments match" in {
+    // Signedness matters
+    matchingArgs("sub", "UInt<8>", "UInt<8>", """ UInt<8>("h0")  """ )
+    matchingArgs("sub", "SInt<8>", "SInt<8>", """ SInt<8>("h0")  """ )
+    matchingArgs("div", "UInt<8>", "UInt<8>", """ UInt<8>("h1")  """ )
+    matchingArgs("div", "SInt<8>", "SInt<8>", """ SInt<8>("h1")  """ )
+    matchingArgs("rem", "UInt<8>", "UInt<8>", """ UInt<8>("h0")  """ )
+    matchingArgs("rem", "SInt<8>", "SInt<8>", """ SInt<8>("h0")  """ )
+    matchingArgs("and", "UInt<8>", "UInt<8>", """ i              """ )
+    matchingArgs("and", "SInt<8>", "UInt<8>", """ asUInt(i)      """ )
+    // Signedness doesn't matter
+    matchingArgs("or",  "UInt<8>", "UInt<8>", """ i """ )
+    matchingArgs("or",  "SInt<8>", "UInt<8>", """ asUInt(i) """ )
+    matchingArgs("xor", "UInt<8>", "UInt<8>", """ UInt<8>("h0")  """ )
+    matchingArgs("xor", "SInt<8>", "UInt<8>", """ UInt<8>("h0")  """ )
+    // Always true
+    matchingArgs("eq",  "UInt<8>", "UInt<1>", """ UInt<1>("h1")  """ )
+    matchingArgs("leq", "UInt<8>", "UInt<1>", """ UInt<1>("h1")  """ )
+    matchingArgs("geq", "UInt<8>", "UInt<1>", """ UInt<1>("h1")  """ )
+    // Never true
+    matchingArgs("neq", "UInt<8>", "UInt<1>", """ UInt<1>("h0")  """ )
+    matchingArgs("lt",  "UInt<8>", "UInt<1>", """ UInt<1>("h0")  """ )
+    matchingArgs("gt",  "UInt<8>", "UInt<1>", """ UInt<1>("h0")  """ )
+  }
+
 }