feat: pipeline fp8 scalar softmax and blackwell mma issue

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

@@ -149,7 +149,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)))
@@ -225,14 +226,14 @@ class TensorCoreBlackwell(
     inputType = TensorInputType.FP8E4M3
   ))
 
-  val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0)
+  val elemM = if (numLanes == 4) issueElemReg(0, 0) else issueElemReg(1, 0)
-  val elemN = if (numLanes == 4) elemReg(1) else elemReg(2)
+  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) := TensorCoreBlackwellFP8Packing.selectA(operandA, k, elemM, numLanes)
     dpu.io.in.bits.b(k) := TensorCoreBlackwellFP8Packing.selectB(operandB, k, elemN)
   }
-  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
 
@@ -253,7 +254,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 }
@@ -322,24 +324,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
+    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
+    assert(state === State.bwCompute || state === State.bwDpuResp,
-      when(elemReg === (numLanes - 1).U) {
+      "BWGMMA DPU response arrived outside the compute states")
+    mmaDataReg(retireElemReg) := dpu.io.out.bits.data
+    when(retireElemReg === (numLanes - 1).U) {
       state := State.bwWriteCReq
     }.otherwise {
-        elemReg := elemReg + 1.U
+      retireElemReg := retireElemReg + 1.U
-        state := State.bwCompute
-      }
     }
   }
 

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) {