Fix width of signed addition when input to mux (#2440)

Fix bugs related to arithmetic ops inlined into a mux leg.  Add formal
equivalence checks to lock in this behavior.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>

3 files changed, 137 insertions, 1 deletions

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 3d5d4988..d50137ce 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -90,7 +90,10 @@ jobs:
         # By having this "if" here, this job returns success so that all_tests_passed will succeed too
         if: github.event_name == 'pull_request' &&
             ! contains(github.event.pull_request.labels.*.name, 'Skip Formal CI')
-        run: ./.run_formal_checks.sh ${{ matrix.design }}
+        run: |
+          echo ${{ github.event_name }}
+          echo ${{ github.event.pull_request.labels }}
+          ./.run_formal_checks.sh ${{ matrix.design }}
 
   # Sentinel job to simplify how we specify which checks need to pass in branch
   # protection and in Mergify
diff --git a/src/main/scala/firrtl/passes/SplitExpressions.scala b/src/main/scala/firrtl/passes/SplitExpressions.scala
index 26088e9c..c0d6e3cb 100644
--- a/src/main/scala/firrtl/passes/SplitExpressions.scala
+++ b/src/main/scala/firrtl/passes/SplitExpressions.scala
@@ -27,6 +27,14 @@ object SplitExpressions extends Pass {
     case _ => false
   }
 
+  private def isSignedArithmetic(e: Expression): Boolean = e match {
+    case DoPrim(PrimOps.Add, _, _, _: SIntType) => true
+    case DoPrim(PrimOps.Sub, _, _, _: SIntType) => true
+    case DoPrim(PrimOps.Mul, _, _, _: SIntType) => true
+    case DoPrim(PrimOps.Div, _, _, _: SIntType) => true
+    case _ => false
+  }
+
   private def onModule(m: Module): Module = {
     val namespace = Namespace(m)
     def onStmt(s: Statement): Statement = {
@@ -53,6 +61,10 @@ object SplitExpressions extends Pass {
       def onExp(e: Expression): Expression =
         e.map(onExp) match {
           case ex: DoPrim => ex.map(split)
+          // Arguably we should be splitting all Mux expressions but this has a negative impact on
+          // Verilog, instead this is a focused fix for
+          // https://github.com/chipsalliance/firrtl/issues/2439
+          case ex: Mux if isSignedArithmetic(ex.tval) || isSignedArithmetic(ex.fval) => ex.map(split)
           case ex => ex
         }
 
diff --git a/src/test/scala/firrtlTests/VerilogEquivalenceSpec.scala b/src/test/scala/firrtlTests/VerilogEquivalenceSpec.scala
index 747f6689..3ec19e77 100644
--- a/src/test/scala/firrtlTests/VerilogEquivalenceSpec.scala
+++ b/src/test/scala/firrtlTests/VerilogEquivalenceSpec.scala
@@ -120,4 +120,125 @@ class VerilogEquivalenceSpec extends FirrtlFlatSpec {
     firrtlEquivalenceWithVerilog(input1, expected)
     firrtlEquivalenceWithVerilog(input2, expected)
   }
+
+  case class SignedMuxTest(
+    operatorFIRRTL:  String,
+    operatorVerilog: String,
+    inputWidth:      Int,
+    outputWidth:     Int,
+    secondArgSigned: Boolean) {
+    private val moduleName = s"Signed${operatorFIRRTL.capitalize}Mux"
+    private val headerFIRRTL =
+      s"""|circuit $moduleName :
+          |  module $moduleName :
+          |    input sel : UInt<1>
+          |    input is0 : SInt<$inputWidth>
+          |    input is1 : ${if (secondArgSigned) "S" else "U"}Int<$inputWidth>
+          |    output os : SInt<$outputWidth>""".stripMargin
+    private val expressionFIRRTL = s"$operatorFIRRTL(is0, is1)"
+    private val headerVerilog =
+      s"""|module Reference(
+          |  input sel,
+          |  input signed [$inputWidth-1:0] is0,
+          |  input ${if (secondArgSigned) "signed" else ""}[$inputWidth-1:0] is1,
+          |  output signed [$outputWidth-1:0] os
+          |);""".stripMargin
+    private val expressionVerilog = s"is0 $operatorVerilog is1"
+
+    def firrtlWhen: String =
+      s"""|$headerFIRRTL
+          |    os <= SInt(0)
+          |    when sel :
+          |      os <= $expressionFIRRTL
+          |""".stripMargin
+
+    def firrtlTrue: String =
+      s"""|$headerFIRRTL
+          |    os <= mux(sel, $expressionFIRRTL, SInt(0))
+          |""".stripMargin
+
+    def firrtlFalse: String =
+      s"""|$headerFIRRTL
+          |    os <= mux(sel, SInt(0), $expressionFIRRTL)
+          |""".stripMargin
+
+    def verilogTrue: String =
+      s"""|$headerVerilog
+          |  assign os = sel ? $expressionVerilog : $outputWidth'sh0;
+          |endmodule
+          |""".stripMargin
+    def verilogFalse: String =
+      s"""|$headerVerilog
+          |  assign os = sel ? $outputWidth'sh0 : $expressionVerilog;
+          |endmodule
+          |""".stripMargin
+  }
+
+  Seq(
+    SignedMuxTest("add", "+", 2, 3, true),
+    SignedMuxTest("sub", "-", 2, 3, true),
+    SignedMuxTest("mul", "*", 2, 4, true),
+    SignedMuxTest("div", "/", 2, 3, true),
+    SignedMuxTest("rem", "%", 2, 2, true),
+    SignedMuxTest("dshl", "<<", 2, 3, false)
+  ).foreach {
+    case t =>
+      s"signed ${t.operatorFIRRTL} followed by a mux" should "be correct" in {
+        info(s"'when' where '${t.operatorFIRRTL}' used in the 'when' body okay")
+        firrtlEquivalenceWithVerilog(t.firrtlWhen, t.verilogTrue)
+        info(s"'mux' where '${t.operatorFIRRTL}' used in the true leg okay")
+        firrtlEquivalenceWithVerilog(t.firrtlTrue, t.verilogTrue)
+        info(s"'mux' where '${t.operatorFIRRTL}' used in the false leg okay")
+        firrtlEquivalenceWithVerilog(t.firrtlFalse, t.verilogFalse)
+      }
+  }
+
+  "signed shl followed by a mux" should "be correct" in {
+    val headerFIRRTL =
+      """|circuit SignedShlMux :
+         |  module SignedShlMux :
+         |    input sel : UInt<1>
+         |    input is0 : SInt<2>
+         |    output os : SInt<5>""".stripMargin
+    val headerVerilog =
+      """|module Reference(
+         |  input sel,
+         |  input signed [1:0] is0,
+         |  output signed [4:0] os
+         |);""".stripMargin
+
+    val firrtlWhen =
+      s"""|$headerFIRRTL
+          |    os <= SInt(0)
+          |    when sel :
+          |      os <= shl(is0, 2)
+          |""".stripMargin
+
+    val firrtlTrue =
+      s"""|$headerFIRRTL
+          |    os <= mux(sel, shl(is0, 2), SInt(0))
+          |""".stripMargin
+
+    val firrtlFalse =
+      s"""|$headerFIRRTL
+          |    os <= mux(sel, SInt(0), shl(is0, 2))
+          |""".stripMargin
+
+    val verilogTrue =
+      s"""|$headerVerilog
+          |  assign os = sel ? (is0 << 2'h2) : 5'sh0;
+          |endmodule
+          |""".stripMargin
+
+    val verilogFalse =
+      s"""|$headerVerilog
+          |  assign os = sel ? 5'sh0 : (is0 << 2'h2);
+          |endmodule
+          |""".stripMargin
+
+    firrtlEquivalenceWithVerilog(firrtlWhen, verilogTrue)
+    firrtlEquivalenceWithVerilog(firrtlTrue, verilogTrue)
+    firrtlEquivalenceWithVerilog(firrtlFalse, verilogFalse)
+  }
+
 }