FMA in FPU.

Change-Id: I36c27f7f80f7b11cbddb67e58471421143666e6e
diff --git a/hdl/chisel/src/common/BUILD b/hdl/chisel/src/common/BUILD
index 8397f95..2ba064c 100644
--- a/hdl/chisel/src/common/BUILD
+++ b/hdl/chisel/src/common/BUILD
@@ -24,6 +24,7 @@
         "Fifo.scala",
         "IDiv.scala",
         "Library.scala",
+        "MathUtil.scala",
         "Slice.scala",
     ],
     visibility = ["//visibility:public"],
@@ -52,6 +53,7 @@
         "Fpu.scala",
     ],
     deps = [
+        ":common",
         ":fp",
     ],
 )
diff --git a/hdl/chisel/src/common/Fp.scala b/hdl/chisel/src/common/Fp.scala
index 1cc1c07..93bce28 100644
--- a/hdl/chisel/src/common/Fp.scala
+++ b/hdl/chisel/src/common/Fp.scala
@@ -38,9 +38,10 @@
     (exponent === "b11111111".U) && (mantissa =/= 0.U)
   }
 
-  // The mantissa is generally the fractional component of significand.
+  // The mantissa is the fractional component of significand.
+  // There is a leading 1, except for subnormal numbers.
   def significand(): UInt = {
-    Cat(1.U(1.W), mantissa)
+    Cat(exponent.orR, mantissa)
   }
 }
 
diff --git a/hdl/chisel/src/common/Fpu.scala b/hdl/chisel/src/common/Fpu.scala
index 384a1f6..fc2f82e 100644
--- a/hdl/chisel/src/common/Fpu.scala
+++ b/hdl/chisel/src/common/Fpu.scala
@@ -17,69 +17,144 @@
 import chisel3._
 import chisel3.util._
 
-object FpuOp extends ChiselEnum {
-  val FpuMul = Value
-}
-
 class FpuCmd extends Bundle {
-  val op = FpuOp()
   val ina = new Fp32
   val inb = new Fp32
+  val inc = new Fp32
 }
 
 class FpuState1 extends Bundle {
-  val zero        = Bool()
-  val inf         = Bool()
-  val nan         = Bool()
-  val sign        = Bool()
+  // Multiply variables
+  val ab_inf      = Bool()
+  val ab_sign     = Bool()
   val exponent    = SInt(10.W)
   val significand = UInt(48.W)
+
+  // Addition variables
+  val c_inf         = Bool()
+  val c_significand = UInt(48.W)
+  val shift         = SInt(11.W)
+  val sub           = Bool()
+
+  val nan           = Bool()
+}
+
+class FpuState2 extends Bundle {
+  val ab_inf      = Bool()
+  val c_inf       = Bool()
+  val sign        = Bool()
+  val exponent    = SInt(10.W)
+  val significand = UInt(49.W)
+
+  val nan           = Bool()
 }
 
 object Fpu {
   def apply(cmd: FpuCmd): Fp32 = {
-    FpuStage2(FpuStage1(cmd))
+    FpuStage3(FpuStage2(FpuStage1(cmd)))
   }
 
   def FpuStage1(cmd: FpuCmd): FpuState1 = {
     val state = Wire(new FpuState1)
 
-    state.zero := cmd.ina.isZero() || cmd.inb.isZero()
-    state.inf := cmd.ina.isInf() || cmd.inb.isInf()
-    state.nan := cmd.ina.isNan() || cmd.inb.isNan()
+    val ab_zero = cmd.ina.isZero() || cmd.inb.isZero()
+    val ab_inf = cmd.ina.isInf() || cmd.inb.isInf()
+    state.ab_inf := ab_inf
+    state.c_inf := cmd.inc.isInf()
 
-    state.sign := cmd.ina.sign ^ cmd.inb.sign
-    state.exponent := (cmd.ina.exponent +& cmd.inb.exponent).zext - 127.S
+    // Compute ina * inb % normalization
+    val ab_sign = cmd.ina.sign ^ cmd.inb.sign
+    state.ab_sign := ab_sign
     state.significand := cmd.ina.significand() * cmd.inb.significand()
+    val product_exponent = (cmd.ina.exponent +& cmd.inb.exponent).zext - 127.S
+
+    // Preshift c for addition.
+    val sum_shift = cmd.inc.exponent.zext -& product_exponent
+    // Right pad c significand to match product, no propagation delay.
+    val padded_c_significand = cmd.inc.significand << 23.U
+    // Compute shift, saturate and take 6 bits to barrel shift.
+    // We saturate to 6 bits max as ceil(log2(48)) = 5.
+    val raw_right_shift = product_exponent -& cmd.inc.exponent.zext
+    val right_shift = Clamp(raw_right_shift, 0.S, 63.S).asUInt
+    state.c_significand := padded_c_significand >> right_shift(5, 0)
+    state.shift := raw_right_shift
+
+    // Mark next cycle as a subtraction if the signs of ab and c differ.
+    state.sub := (ab_sign ^ cmd.inc.sign)
+
+    // Take max exponent of (a*b) or c. The smaller of the two will be right
+    // shifted (a*b in stage 2 or c in stage 2)
+    state.exponent := Mux(
+        raw_right_shift > 0.S, product_exponent, cmd.inc.exponent.zext)
+
+    state.nan := cmd.ina.isNan() || cmd.inb.isNan() || cmd.inc.isNan() ||
+                 (ab_zero && ab_inf)
+
     state
   }
 
-  def FpuStage2(state: FpuState1): Fp32 = {
-    // Grab 24-bits of the mantissa for rounding. At least one of the MSB (for
-    // when the significand product >= 2) or 2nd MSB is guarenteed to be set.
-    // The below mux effectively picks the correct 25-bit truncated significand
-    // depending if the MSB is set, then returns the lower 24-bits of that
-    // result (the mantissa of the truncated significand).
-    val mantissa24 = Mux(
-        state.significand(47),
-        state.significand(46, 23),
-        state.significand(45, 22))
-    // TODO(derekjchow): Rounding modes
-    val mantissa = ((mantissa24 + 1.U(1.W)) >> 1)(22, 0)
+  def FpuStage2(state1: FpuState1): FpuState2 = {
+    val state2 = Wire(new FpuState2)
 
-    // If the significand product >= 2, we "shift the decimal" to the right
-    // by one bit. Add 1 to the exponent to compensate.
-    val exponent = state.exponent + state.significand(47).asUInt.zext
+    // Variables to forward to next cycle.
+    state2.ab_inf := state1.ab_inf
+    state2.c_inf := state1.c_inf
+    state2.exponent := state1.exponent
+    // Inf - Inf = NaN
+    state2.nan := state1.nan || (state1.ab_inf && state1.c_inf && state1.sub)
+
+    // Compute shift, saturate and take 6 bits to barrel shift ab_significand.
+    // Hopefully shift here matches propagation delay of potential C inversion.
+    val shift = (Clamp(-state1.shift, 0.S, 63.S).asUInt)(5, 0)
+    val ab_significand = (state1.significand >> shift).zext
+    assert(ab_significand.getWidth == 49)
+    // Zext and invert if necessary
+    val c_significand = Mux(
+      state1.sub, -(state1.c_significand.zext), state1.c_significand.zext)
+    assert(c_significand.getWidth == 49)
+
+    val significand_sum = ab_significand +& c_significand
+    assert(significand_sum.getWidth == 50)
+    val sign            = significand_sum(49)
+    val new_significand = (significand_sum.abs.asUInt)(48, 0)
+    assert(new_significand.getWidth == 49)
+    state2.sign := state1.ab_sign ^ sign
+    state2.significand := new_significand
+
+    state2
+  }
+
+  def FpuStage3(state: FpuState2): Fp32 = {
+    // Compute mantissa
+    val left_shamt =
+        PriorityEncoder(Cat(1.U(1.W), Reverse(state.significand)))(5,0)
+    val shifted_significand =
+        (state.significand << left_shamt)(state.significand.getWidth, 0)
+    // Grab 25 bit significand
+    val reduced_significand = shifted_significand(shifted_significand.getWidth - 1,
+                                                  shifted_significand.getWidth - 26)
+    // Perform rounding step, going to 26 bits
+    // TODO(derekjchow): Rounding mode
+    val rounded_significand = reduced_significand +& 1.U(1.W)
+    // Get new mantissa
+    val mantissa = Mux(rounded_significand(25),
+                       rounded_significand(24, 2),
+                       rounded_significand(23, 1))
+
+    // Compute new exponent
+    // The +2.S comes from two widening operations in previous stages
+    val exponent = state.exponent - left_shamt.zext + 2.S +
+                   rounded_significand(25).asUInt.zext
     // Check for overflow.
-    val inf = state.inf || (exponent >= (1 << 8).S)
+    val inf = state.ab_inf || state.c_inf || (exponent >= (1 << 8).S)
     // Check for very small numbers that should round to zero.
-    val zero = state.zero || (exponent < 0.S)
+    val zero = (reduced_significand === 0.U) || (exponent < 0.S)
+    val nan = state.nan
 
     MuxCase(
         Fp32(state.sign, exponent(7, 0), mantissa),
         Array(
-            state.nan -> Fp32(false.B, ((1<<8)-1).U, mantissa),
-            (state.zero && state.inf) -> Fp32.NaN(),
+            nan -> Fp32.NaN(),
             inf -> Fp32.Inf(state.sign),
             zero -> Fp32.Zero(state.sign)
         ))
diff --git a/hdl/chisel/src/common/FpuTest.scala b/hdl/chisel/src/common/FpuTest.scala
index c80cf7c..693ddd9 100644
--- a/hdl/chisel/src/common/FpuTest.scala
+++ b/hdl/chisel/src/common/FpuTest.scala
@@ -24,20 +24,26 @@
   val io = IO(new Bundle {
     val ina    = Input(UInt(32.W))
     val inb    = Input(UInt(32.W))
-    val op     = Input(FpuOp())
+    val inc    = Input(UInt(32.W))
+    val state1 = Output(new FpuState1)
+    val state2 = Output(new FpuState2)
     val out    = Output(new Fp32)
   })
 
   val fp_a = Fp32.fromWord(io.ina)
   val fp_b = Fp32.fromWord(io.inb)
+  val fp_c = Fp32.fromWord(io.inc)
 
   val cmd = Wire(new FpuCmd)
   cmd.ina := fp_a
   cmd.inb := fp_b
-  cmd.op := io.op
+  cmd.inc := fp_c
 
   val stage1 = Fpu.FpuStage1(cmd)
-  io.out := Fpu.FpuStage2(stage1)
+  val stage2 = Fpu.FpuStage2(stage1)
+  io.state1 := stage1
+  io.state2 := stage2
+  io.out := Fpu.FpuStage3(stage2)
 }
 
 class FpuSpec extends AnyFreeSpec with ChiselScalatestTester {
@@ -50,80 +56,157 @@
     int
   }
 
-  "Zero" in {
+  def GetFloat(exponent: Int, mantissa: Int): Float = {
+    val int_val = (exponent << 23) + mantissa
+    java.lang.Float.intBitsToFloat(int_val)
+  }
+
+  "Mul Zero" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(0))
       dut.io.inb.poke(Float2BigInt(42))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(0) { dut.io.out.exponent.peekInt() }
       assertResult(0) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "Identity" in {
+  "Mul Identity" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(1))
       dut.io.inb.poke(Float2BigInt(42))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(132) { dut.io.out.exponent.peekInt() }
       assertResult(2621440) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "Negative" in {
+  "Mul Negative" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(-1.0f))
       dut.io.inb.poke(Float2BigInt(42))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(1) { dut.io.out.sign.peekInt() }
       assertResult(132) { dut.io.out.exponent.peekInt() }
       assertResult(2621440) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "Half" in {
+  "Mul Half" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(0.5f))
       dut.io.inb.poke(Float2BigInt(42))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(131) { dut.io.out.exponent.peekInt() }
       assertResult(2621440) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "Overflow" in {
+  "Mul Overflow" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(2e30f))
       dut.io.inb.poke(Float2BigInt(2e30f))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(255) { dut.io.out.exponent.peekInt() }
       assertResult(0) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "Rounds to Zero" in {
+  "Mul Rounds to Zero" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(1e-30f))
       dut.io.inb.poke(Float2BigInt(1e-30f))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(0) { dut.io.out.exponent.peekInt() }
       assertResult(0) { dut.io.out.mantissa.peekInt() }
     }
   }
 
-  "NaN" in {
+  "Mul NaN" in {
     test(new FpuTester()) { dut =>
-      dut.io.op.poke(FpuOp.FpuMul)
       dut.io.ina.poke(Float2BigInt(Float.NaN))
       dut.io.inb.poke(Float2BigInt(4.0f))
+      dut.io.inc.poke(Float2BigInt(0))
       assertResult(0) { dut.io.out.sign.peekInt() }
       assertResult(255) { dut.io.out.exponent.peekInt() }
       assert(dut.io.out.mantissa.peekInt() != 0)
     }
   }
-}
\ No newline at end of file
+
+  "Fma" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(2.0f))
+      dut.io.inb.poke(Float2BigInt(1.5f))
+      dut.io.inc.poke(Float2BigInt(6.0f))
+
+      assertResult(0) { dut.io.out.sign.peekInt() }
+      assertResult(130) { dut.io.out.exponent.peekInt() }
+      assertResult(1048576) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+
+  "Fms" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(2.0f))
+      dut.io.inb.poke(Float2BigInt(1.5f))
+      dut.io.inc.poke(Float2BigInt(-6.0f))
+
+      assertResult(1) { dut.io.out.sign.peekInt() }
+      assertResult(128) { dut.io.out.exponent.peekInt() }
+      assertResult(4194304) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+
+  "Fnma" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(-2.0f))
+      dut.io.inb.poke(Float2BigInt(1.5f))
+      dut.io.inc.poke(Float2BigInt(13.5f))
+
+      assertResult(0) { dut.io.out.sign.peekInt() }
+      assertResult(130) { dut.io.out.exponent.peekInt() }
+      assertResult(2621440) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+
+  "Fnms" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(-2.0f))
+      dut.io.inb.poke(Float2BigInt(1.5f))
+      dut.io.inc.poke(Float2BigInt(-13.5f))
+
+      assertResult(1) { dut.io.out.sign.peekInt() }
+      assertResult(131) { dut.io.out.exponent.peekInt() }
+      assertResult(262144) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+
+  "Add" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(9000.0f))
+      dut.io.inb.poke(Float2BigInt(1.0f))
+      dut.io.inc.poke(Float2BigInt(1.0f))
+
+      assertResult(0) { dut.io.out.sign.peekInt() }
+      assertResult(140) { dut.io.out.exponent.peekInt() }
+      assertResult(828416) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+
+  "Sub" in {
+    test(new FpuTester()) { dut =>
+      dut.io.ina.poke(Float2BigInt(15.0f))
+      dut.io.inb.poke(Float2BigInt(1.0f))
+      dut.io.inc.poke(Float2BigInt(-100.0f))
+
+      assertResult(1) { dut.io.out.sign.peekInt() }
+      assertResult(133) { dut.io.out.exponent.peekInt() }
+      assertResult(2752512) { dut.io.out.mantissa.peekInt() }
+    }
+  }
+}
diff --git a/hdl/chisel/src/common/MathUtil.scala b/hdl/chisel/src/common/MathUtil.scala
new file mode 100644
index 0000000..155e7d4
--- /dev/null
+++ b/hdl/chisel/src/common/MathUtil.scala
@@ -0,0 +1,24 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package common
+
+import chisel3._
+import chisel3.util._
+
+object Clamp {
+  def apply(x: SInt, min: SInt, max: SInt): SInt = {
+    Mux(x > min, Mux(x < max, x, max), min)
+  }
+}