fix: stabilize 4-lane fp8 bwgmma

src/main/scala/radiance/core/TensorCoreBlackwell.scala

@@ -338,9 +338,9 @@ class TensorCoreBlackwell(
     }
   }
 
-  when(dpuValid) {
-    assert(state === State.bwCompute || state === State.bwDpuResp,
-      "BWGMMA DPU response arrived outside the compute states")
+  val inDpuResponseWindow = state === State.bwCompute || state === State.bwDpuResp
+
+  when(dpuValid && inDpuResponseWindow) {
     mmaDataReg(retireElemReg) := dpu.io.out.bits.data
     when(retireElemReg === (numLanes - 1).U) {
       state := State.bwWriteCReq

src/main/scala/radiance/core/TensorDPU.scala

@@ -381,7 +381,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters
   val inC = unbox(recode(io.in.bits.c, S), S, Some(tOut))
 
   val productStageOut = StallingPipe(io.stall, io.in.valid, VecInit(productRecoded))
-  val productStageC = StallingPipe(io.stall, io.in.valid, inC)
+  val productStageC = RegEnable(inC, !io.stall && io.in.valid)
 
   val log2Dim = log2Ceil(dim)
   require(dim == (1 << log2Dim), s"dim (${dim}) is not power of two!")
@@ -389,7 +389,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters
   val interim = (log2Dim to 0 by -1).map { i =>
     Wire(Valid(Vec(1 << i, Bits(recOutFLen.W))))
   }
-  val interimC = (log2Dim to 0 by -1).map(_ => Wire(Valid(Bits(recOutFLen.W))))
+  val interimC = (log2Dim to 0 by -1).map(_ => Wire(Bits(recOutFLen.W)))
   interim(0) := productStageOut
   interimC(0) := productStageC
 
@@ -413,11 +413,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters
       }
 
       outputs := StallingPipe(io.stall, inputs.valid, VecInit(addOuts))
-      outC := StallingPipe(io.stall, inputs.valid, inC.bits)
-      when(inputs.valid =/= inC.valid) {
-        printf("WARN: DotProductPipeFP8E4M3 input/C valid mismatch: inputs=%d c=%d\n",
-          inputs.valid, inC.valid)
-      }
+      outC := RegEnable(inC, !io.stall && inputs.valid)
 
       (outputs, outC)
     }
@@ -427,7 +423,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters
   val acc = Module(new hardfloat.AddRecFN(outExpWidth, outSigWidth))
   acc.io.subOp := 0.U
   acc.io.a := addStageOut.bits(0)
-  acc.io.b := addStageC.bits
+  acc.io.b := addStageC
   acc.io.roundingMode := hardfloat.consts.round_near_even
   acc.io.detectTininess := hardfloat.consts.tininess_afterRounding
 

src/test/scala/radiance/TensorCoreBlackwellTest.scala

@@ -324,4 +324,80 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester {
       }
     }
   }
+
+  it should "run bwgmma with 4-lane 16-byte fragments" in {
+    val lanes = 4
+    test(new TensorCoreBlackwell(numWarps, lanes, half = true, numSourceIds = 4))
+      .withAnnotations(Seq(VerilatorBackendAnnotation)) { c =>
+      idleIO(c)
+
+      val fragBytes = 16
+      val aBase = BigInt(0x100)
+      val bBase = BigInt(0x800)
+      val cBase = BigInt(0x1000)
+      val fp8One = BigInt(0x38)
+      val fp8Two = BigInt(0x40)
+      val fp32One = BigInt(0x3f800000)
+      val fp32SixtyFive = BigInt(0x42820000)
+      val aFrag = packWords(Seq.fill(16)(fp8One), 8)
+      val bFrag = packWords(Seq.fill(16)(fp8Two), 8)
+      val cFrag = packWords(Seq.fill(lanes)(fp32One), 32)
+      val expectedCFrag = packWords(Seq.fill(lanes)(fp32SixtyFive), 32)
+
+      val tmem = makeTmem()
+      for (i <- 0 until 64) {
+        tmem(aBase / fragBytes + i) = aFrag
+        tmem(cBase / fragBytes + i) = cFrag
+      }
+      val bMem = mutable.Map[BigInt, BigInt]()
+      for (i <- 0 until 64) bMem(bBase + i * fragBytes) = bFrag
+
+      c.io.reqB.ready.poke(true.B)
+      c.io.writeback.ready.poke(true.B)
+      c.io.initiate.valid.poke(true.B)
+      c.io.initiate.bits.op.poke(0.U)
+      c.io.initiate.bits.wid.poke(1.U)
+      c.io.initiate.bits.rd.poke(0.U)
+      c.io.initiate.bits.addressA.poke(aBase.U)
+      c.io.initiate.bits.addressB.poke(bBase.U)
+      c.io.initiate.bits.addressC.poke(cBase.U)
+      c.clock.step()
+      c.io.initiate.valid.poke(false.B)
+
+      var pendingB = Option.empty[(BigInt, BigInt)]
+      var sawWriteback = false
+      var cycles = 0
+      for (_ <- 0 until 30000 if !sawWriteback) {
+        stepTmem(c, tmem)
+        pendingB.foreach { case (src, data) =>
+          c.io.respB.valid.poke(true.B)
+          c.io.respB.bits.source.poke(src.U)
+          c.io.respB.bits.data.poke(data.U)
+        }
+        if (pendingB.isEmpty) c.io.respB.valid.poke(false.B)
+
+        if (c.io.writeback.valid.peek().litToBoolean) {
+          sawWriteback = true
+        } else {
+          val nextB = if (c.io.reqB.valid.peek().litToBoolean) {
+            val addr = c.io.reqB.bits.address.peek().litValue
+            val src = c.io.reqB.bits.source.peek().litValue
+            Some((src, bMem(addr)))
+          } else None
+          c.clock.step()
+          cycles += 1
+          pendingB = nextB
+        }
+      }
+
+      assert(sawWriteback, "4-lane BWGMMA did not complete")
+      assert(cycles < 10000,
+        s"4-lane BWGMMA took $cycles cycles")
+      for (i <- 0 until 64) {
+        val row = cBase / fragBytes + i
+        assert(tmem(row) == expectedCFrag,
+          s"4-lane C frag $i mismatch: got 0x${tmem(row).toString(16)}")
+      }
+    }
+  }
 }