feat: pipeline scalar softmax and blackwell mma issue

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

@@ -126,7 +126,8 @@ class TensorCoreBlackwell(
   val bIndexReg = RegInit(0.U(log2Ceil(numBFragsPerGroup).W))
   val mGroupReg = RegInit(0.U(log2Ceil(numMGroups).W))
   val substepReg = RegInit(0.U(1.W))
-  val elemReg = RegInit(0.U(log2Ceil(numLanes).W))
+  val issueElemReg = RegInit(0.U(log2Ceil(numLanes).W))
+  val retireElemReg = RegInit(0.U(log2Ceil(numLanes).W))
   val waitCounter = RegInit(0.U(3.W))
 
   val aBuf = Reg(Vec(numAFragsPerSet, UInt(memWidth.W)))
@@ -205,8 +206,8 @@ class TensorCoreBlackwell(
     x((idx + 1) * 16 - 1, idx * 16)
   }
 
-  val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0)
-  val elemN = if (numLanes == 4) elemReg(1) else elemReg(2)
+  val elemM = if (numLanes == 4) issueElemReg(0, 0) else issueElemReg(1, 0)
+  val elemN = if (numLanes == 4) issueElemReg(1) else issueElemReg(2)
   dpu.io.in.valid := dpuInValid
   for (k <- 0 until 8) {
     dpu.io.in.bits.a(k) := (
@@ -223,7 +224,7 @@ class TensorCoreBlackwell(
     )
     dpu.io.in.bits.b(k) := Mux(elemN.asBool, halfWord(operandB, 8 + k), halfWord(operandB, k))
   }
-  dpu.io.in.bits.c := cWords(elemReg)
+  dpu.io.in.bits.c := cWords(issueElemReg)
   dpu.io.stall := false.B
   val dpuValid = dpu.io.out.valid
 
@@ -244,7 +245,8 @@ class TensorCoreBlackwell(
     bIndexReg := 0.U
     mGroupReg := 0.U
     substepReg := 0.U
-    elemReg := 0.U
+    issueElemReg := 0.U
+    retireElemReg := 0.U
     switch(io.initiate.bits.op) {
       is(Ops.bwgmma) { state := State.bwLoadAReq }
       is(Ops.tcgen05Cp) { state := State.cpRead }
@@ -313,24 +315,28 @@ class TensorCoreBlackwell(
 
   when(state === State.bwReadCResp) {
     cDataReg := io.tmemC.cRdata
-    elemReg  := 0.U
+    issueElemReg := 0.U
+    retireElemReg := 0.U
     state    := State.bwCompute
   }
 
   when(state === State.bwCompute) {
     dpuInValid := true.B
-    state := State.bwDpuResp
+    when(issueElemReg === (numLanes - 1).U) {
+      state := State.bwDpuResp
+    }.otherwise {
+      issueElemReg := issueElemReg + 1.U
+    }
   }
 
-  when(state === State.bwDpuResp) {
-    when(dpuValid) {
-      mmaDataReg(elemReg) := dpu.io.out.bits.data
-      when(elemReg === (numLanes - 1).U) {
-        state := State.bwWriteCReq
-      }.otherwise {
-        elemReg := elemReg + 1.U
-        state := State.bwCompute
-      }
+  when(dpuValid) {
+    assert(state === State.bwCompute || state === State.bwDpuResp,
+      "BWGMMA DPU response arrived outside the compute states")
+    mmaDataReg(retireElemReg) := dpu.io.out.bits.data
+    when(retireElemReg === (numLanes - 1).U) {
+      state := State.bwWriteCReq
+    }.otherwise {
+      retireElemReg := retireElemReg + 1.U
     }
   }
 

src/main/scala/radiance/tile/VortexCore.scala

@@ -213,6 +213,7 @@ class Vortex(tile: RadianceTile)(implicit p: Parameters)
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch_unit.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch_unit_sane.sv")
+  addResource("/vsrc/vortex/hw/rtl/core/VX_tmem_softmax_unit.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_execute.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_fetch.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_gather_unit.sv")

src/test/scala/radiance/TensorCoreBlackwellTest.scala

@@ -283,6 +283,7 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester {
 
       var pendingB = Option.empty[(BigInt, BigInt)]
       var sawWriteback = false
+      var cycles = 0
 
       for (_ <- 0 until 20000 if !sawWriteback) {
         // Drive TMEM reads/writes
@@ -306,11 +307,14 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester {
           } else None
 
           c.clock.step()
+          cycles += 1
           pendingB = nextB
         }
       }
 
       assert(sawWriteback, "BWGMMA did not complete")
+      assert(cycles < 5000,
+        s"BWGMMA took $cycles cycles; fragment elements are not issuing back-to-back")
       c.io.writeback.bits.wid.expect(1.U)
       // Verify all 32 C frags in TMEM
       for (i <- 0 until 32) {