feat: execute blackwell fp16 bwgmma on systolic array

src/main/scala/radiance/core/BlackwellFP16SystolicArray.scala

@@ -0,0 +1,129 @@
+// See LICENSE.SiFive for license details.
+// See LICENSE.Berkeley for license details.
+
+package radiance.core
+
+import chisel3._
+import chisel3.util._
+import freechips.rocketchip.tile
+
+/** An 8x8 output-stationary FP16 systolic array for Blackwell BWGMMA.
+  *
+  * Each PE contains an FP16 multiplier and an FP32 accumulator. A operands
+  * move left-to-right and B operands move top-to-bottom.
+  */
+class BlackwellFP16SystolicArray(val kDim: Int = 32)
+    extends Module with tile.HasFPUParameters {
+  private val arrayDim = 8
+  private val wavefrontCycles = kDim + 2 * (arrayDim - 1)
+
+  require(kDim > 0)
+
+  val fLen = 32
+  val minFLen = 16
+  def xLen = 32
+
+  private val tIn = tile.FType.H
+  private val tOut = tile.FType.S
+  private val recOutWidth = tOut.exp + tOut.sig + 1
+  private val stepWidth = log2Ceil(wavefrontCycles)
+  private val kIndexWidth = math.max(1, log2Ceil(kDim))
+
+  val io = IO(new Bundle {
+    val start = Input(Bool())
+    val ready = Output(Bool())
+    val busy = Output(Bool())
+    val done = Output(Bool())
+    val a = Input(Vec(arrayDim, Vec(kDim, UInt(16.W))))
+    val b = Input(Vec(kDim, Vec(arrayDim, UInt(16.W))))
+    val c = Input(Vec(arrayDim, Vec(arrayDim, UInt(32.W))))
+    val result = Output(Vec(arrayDim, Vec(arrayDim, UInt(32.W))))
+  })
+
+  val busyReg = RegInit(false.B)
+  val doneReg = RegInit(false.B)
+  val stepReg = RegInit(0.U(stepWidth.W))
+  val aPipe = RegInit(VecInit(Seq.fill(arrayDim)(VecInit(Seq.fill(arrayDim)(0.U(16.W))))))
+  val bPipe = RegInit(VecInit(Seq.fill(arrayDim)(VecInit(Seq.fill(arrayDim)(0.U(16.W))))))
+  val accum = Reg(Vec(arrayDim, Vec(arrayDim, UInt(recOutWidth.W))))
+
+  val aFlow = Wire(Vec(arrayDim, Vec(arrayDim, UInt(16.W))))
+  val bFlow = Wire(Vec(arrayDim, Vec(arrayDim, UInt(16.W))))
+
+  for (row <- 0 until arrayDim) {
+    val active = stepReg >= row.U && stepReg < (row + kDim).U
+    val kIndex = (stepReg - row.U)(kIndexWidth - 1, 0)
+    aFlow(row)(0) := Mux(active, io.a(row)(kIndex), 0.U)
+    for (col <- 1 until arrayDim) {
+      aFlow(row)(col) := aPipe(row)(col - 1)
+    }
+  }
+
+  for (col <- 0 until arrayDim) {
+    val active = stepReg >= col.U && stepReg < (col + kDim).U
+    val kIndex = (stepReg - col.U)(kIndexWidth - 1, 0)
+    bFlow(0)(col) := Mux(active, io.b(kIndex)(col), 0.U)
+    for (row <- 1 until arrayDim) {
+      bFlow(row)(col) := bPipe(row - 1)(col)
+    }
+  }
+
+  val sums = Seq.tabulate(arrayDim, arrayDim) { (row, col) =>
+    val aRec = unbox(recode(aFlow(row)(col), H), H, Some(tIn))
+    val bRec = unbox(recode(bFlow(row)(col), H), H, Some(tIn))
+    val multiplier = Module(new hardfloat.MulFullRawFN(tIn.exp, tIn.sig))
+    multiplier.io.a := hardfloat.rawFloatFromRecFN(tIn.exp, tIn.sig, aRec)
+    multiplier.io.b := hardfloat.rawFloatFromRecFN(tIn.exp, tIn.sig, bRec)
+
+    val product = Module(new hardfloat.RoundAnyRawFNToRecFN(
+      multiplier.io.rawOut.expWidth,
+      multiplier.io.rawOut.sigWidth,
+      tOut.exp,
+      tOut.sig,
+      0))
+    product.io.invalidExc := multiplier.io.invalidExc
+    product.io.infiniteExc := false.B
+    product.io.in := multiplier.io.rawOut
+    product.io.roundingMode := hardfloat.consts.round_near_even
+    product.io.detectTininess := hardfloat.consts.tininess_afterRounding
+
+    val adder = Module(new hardfloat.AddRecFN(tOut.exp, tOut.sig))
+    adder.io.subOp := 0.U
+    adder.io.a := accum(row)(col)
+    adder.io.b := product.io.out
+    adder.io.roundingMode := hardfloat.consts.round_near_even
+    adder.io.detectTininess := hardfloat.consts.tininess_afterRounding
+    adder.io.out
+  }
+
+  io.ready := !busyReg
+  io.busy := busyReg
+  io.done := doneReg
+  for (row <- 0 until arrayDim; col <- 0 until arrayDim) {
+    io.result(row)(col) := ieee(box(accum(row)(col), S))
+  }
+
+  doneReg := false.B
+  when(io.start && !busyReg) {
+    busyReg := true.B
+    stepReg := 0.U
+    for (row <- 0 until arrayDim; col <- 0 until arrayDim) {
+      aPipe(row)(col) := 0.U
+      bPipe(row)(col) := 0.U
+      accum(row)(col) := unbox(recode(io.c(row)(col), S), S, Some(tOut))
+    }
+  }.elsewhen(busyReg) {
+    aPipe := aFlow
+    bPipe := bFlow
+    for (row <- 0 until arrayDim; col <- 0 until arrayDim) {
+      accum(row)(col) := sums(row)(col)
+    }
+
+    when(stepReg === (wavefrontCycles - 1).U) {
+      busyReg := false.B
+      doneReg := true.B
+    }.otherwise {
+      stepReg := stepReg + 1.U
+    }
+  }
+}

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

@@ -6,6 +6,29 @@ package radiance.core
 import chisel3._
 import chisel3.util._
 
+object TensorCoreBlackwellFP16Packing {
+  def halfWord(x: UInt, idx: Int): UInt = {
+    x((idx + 1) * 16 - 1, idx * 16)
+  }
+
+  def selectA(operandA: UInt, k: Int, elemM: UInt, numLanes: Int): UInt = {
+    if (numLanes == 4) {
+      Mux(elemM.asBool, halfWord(operandA, 8 + k), halfWord(operandA, k))
+    } else {
+      MuxLookup(elemM, halfWord(operandA, k))(Seq(
+        0.U -> halfWord(operandA, k),
+        1.U -> halfWord(operandA, 8 + k),
+        2.U -> halfWord(operandA, 16 + k),
+        3.U -> halfWord(operandA, 24 + k)
+      ))
+    }
+  }
+
+  def selectB(operandB: UInt, k: Int, elemN: UInt): UInt = {
+    Mux(elemN.asBool, halfWord(operandB, 8 + k), halfWord(operandB, k))
+  }
+}
+
 class TensorCoreBlackwell(
     val numWarps: Int,
     val numLanes: Int,
@@ -37,6 +60,13 @@ class TensorCoreBlackwell(
   val numBFragsPerGroup = numSubsteps * numBFragsPerSubstep
   val numBFragsPerSet = numBGroups * numBFragsPerGroup
   val numCFrags = numBGroups * numMGroups * numSubsteps
+  val systolicDim = 8
+  val systolicK = 32
+  val numTilesPerDim = 16 / systolicDim
+  val numMGroupsPerTile = systolicDim / mElemsPerFrag
+  val numCFragsPerTile = systolicDim * systolicDim / numLanes
+  val totalAFrags = numSets * numAFragsPerSet
+  val totalBFrags = numSets * numBFragsPerSet
 
   object Ops {
     val bwgmma :: bwgmmaWait :: tcgen05Cp :: tcgen05CpWait :: tcgen05Ld :: tcgen05St :: tcgen05Cb :: Nil = Enum(7)
@@ -106,7 +136,7 @@ class TensorCoreBlackwell(
 
   object State extends ChiselEnum {
     val idle, bwLoadAReq, bwLoadAResp, bwLoadBReq, bwLoadBResp,
-        bwReadCReq, bwReadCResp, bwCompute, bwDpuResp, bwWriteCReq,
+        bwReadCReq, bwReadCResp, bwArrayStart, bwArrayRun, bwWriteCReq,
         bwWriteCWait, bwDone, cpRead, cpWrite, ldReq, stReq, stWrite, waitWb,
         cbRead, cbCapture, cbWrite = Value
   }
@@ -120,17 +150,16 @@ class TensorCoreBlackwell(
   val addrCReg = RegInit(0.U(addressWidth.W))
   val sourceCounter = RegInit(0.U(sourceWidth.W))
 
-  val setReg = RegInit(0.U(log2Ceil(numSets).W))
-  val aIndexReg = RegInit(0.U(log2Ceil(numAFragsPerSet).W))
-  val bGroupReg = RegInit(0.U(log2Ceil(numBGroups).W))
-  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 aIndexReg = RegInit(0.U(log2Ceil(totalAFrags).W))
+  val bIndexReg = RegInit(0.U(log2Ceil(totalBFrags).W))
+  val tileMReg = RegInit(0.U(log2Ceil(numTilesPerDim).W))
+  val tileNReg = RegInit(0.U(log2Ceil(numTilesPerDim).W))
+  val cTileFragReg = RegInit(0.U(log2Ceil(numCFragsPerTile).W))
   val waitCounter = RegInit(0.U(3.W))
 
-  val aBuf = Reg(Vec(numAFragsPerSet, UInt(memWidth.W)))
-  val bBuf = Reg(Vec(numBFragsPerGroup, UInt(memWidth.W)))
+  val aBuf = Reg(Vec(totalAFrags, UInt(memWidth.W)))
+  val bBuf = Reg(Vec(totalBFrags, UInt(memWidth.W)))
+  val cTile = Reg(Vec(systolicDim, Vec(systolicDim, UInt(laneWidth.W))))
   val cDataReg = Reg(UInt(memWidth.W))
   val mmaDataReg = Reg(Vec(numLanes, UInt(laneWidth.W)))
 
@@ -142,11 +171,15 @@ class TensorCoreBlackwell(
     base + (fragIndex << fragOffsetBits).asUInt
   }
 
-  val aFragIndex = (setReg * numAFragsPerSet.U) + aIndexReg
-  val bFragIndex =
-    (setReg * numBFragsPerSet.U) + (bGroupReg * numBFragsPerGroup.U) + bIndexReg
+  val localSubstep = cTileFragReg(0)
+  val localMGroup = (cTileFragReg >> 1)(log2Ceil(numMGroupsPerTile) - 1, 0)
+  val localBGroup = cTileFragReg >> log2Ceil(numMGroupsPerTile * numSubsteps)
+  val cMGroup = (tileMReg * numMGroupsPerTile.U) + localMGroup
+  val cBGroup = (tileNReg * 2.U) + localBGroup
   val cFragIndex =
-    (((bGroupReg * numMGroups.U) + mGroupReg) * numSubsteps.U) + substepReg
+    (((cBGroup * numMGroups.U) + cMGroup) * numSubsteps.U) + localSubstep
+  val aFragIndex = aIndexReg
+  val bFragIndex = bIndexReg
   val aReqAddress = byteAddress(addrAReg, aFragIndex)
   val bReqAddress = byteAddress(addrBReg, bFragIndex)
   val cReqAddress = byteAddress(addrCReg, cFragIndex)
@@ -187,45 +220,47 @@ class TensorCoreBlackwell(
   io.respB.ready := false.B
   io.initiate.ready := state === State.idle && !wbValid
 
-  val operandA = Cat(aBuf((mGroupReg << 1) + 1.U), aBuf(mGroupReg << 1))
-  val operandB =
-    if (numLanes == 4) {
-      Cat(bBuf((substepReg << 1) + 1.U), bBuf(substepReg << 1))
-    } else {
-      bBuf(substepReg)
-    }
-  val cWords = cDataReg.asTypeOf(Vec(numLanes, UInt(laneWidth.W)))
-  val dpuInValid = WireDefault(false.B)
-  val dpu = Module(new TensorDotProductUnit(
-    dim = 8,
-    half = true
-  ))
+  val systolic = Module(new BlackwellFP16SystolicArray(systolicK))
+  systolic.io.start := state === State.bwArrayStart
+  systolic.io.c := cTile
 
-  private def halfWord(x: UInt, idx: Int): UInt = {
-    x((idx + 1) * 16 - 1, idx * 16)
+  // Preserve the software-visible FP16 fragment packing while presenting
+  // logical 8x32 and 32x8 operands to the systolic array.
+  for (row <- 0 until systolicDim; k <- 0 until systolicK) {
+    val set = k / 8
+    val kInSet = k % 8
+    val logicalM = (tileMReg << 3) + row.U
+    val mGroup = logicalM >> log2Ceil(mElemsPerFrag)
+    val elemM = logicalM(log2Ceil(mElemsPerFrag) - 1, 0)
+    val aIndex = set.U * numAFragsPerSet.U + (mGroup << 1)
+    val operandA = Cat(aBuf(aIndex + 1.U), aBuf(aIndex))
+    systolic.io.a(row)(k) := TensorCoreBlackwellFP16Packing.selectA(
+      operandA, kInSet, elemM, numLanes)
   }
 
-  val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0)
-  val elemN = if (numLanes == 4) elemReg(1) else elemReg(2)
-  dpu.io.in.valid := dpuInValid
-  for (k <- 0 until 8) {
-    dpu.io.in.bits.a(k) := (
-      if (numLanes == 4) {
-        Mux(elemM.asBool, halfWord(operandA, 8 + k), halfWord(operandA, k))
-      } else {
-        MuxLookup(elemM, halfWord(operandA, k))(Seq(
-          0.U -> halfWord(operandA, k),
-          1.U -> halfWord(operandA, 8 + k),
-          2.U -> halfWord(operandA, 16 + k),
-          3.U -> halfWord(operandA, 24 + k)
-        ))
-      }
-    )
-    dpu.io.in.bits.b(k) := Mux(elemN.asBool, halfWord(operandB, 8 + k), halfWord(operandB, k))
+  for (k <- 0 until systolicK; col <- 0 until systolicDim) {
+    val set = k / 8
+    val kInSet = k % 8
+    val logicalN = (tileNReg << 3) + col.U
+    val bGroup = logicalN >> 2
+    val substep = logicalN(1)
+    val elemN = logicalN(0)
+    val bIndex = set.U * numBFragsPerSet.U +
+      bGroup * numBFragsPerGroup.U + substep * numBFragsPerSubstep.U
+    val operandB =
+      if (numLanes == 4) Cat(bBuf(bIndex + 1.U), bBuf(bIndex)) else bBuf(bIndex)
+    systolic.io.b(k)(col) := TensorCoreBlackwellFP16Packing.selectB(
+      operandB, kInSet, elemN)
+  }
+
+  val mmaWords = Wire(Vec(numLanes, UInt(laneWidth.W)))
+  for (lane <- 0 until numLanes) {
+    val elemM = lane % mElemsPerFrag
+    val elemN = lane / mElemsPerFrag
+    val row = localMGroup * mElemsPerFrag.U + elemM.U
+    val col = localBGroup * 4.U + localSubstep * 2.U + elemN.U
+    mmaWords(lane) := systolic.io.result(row)(col)
   }
-  dpu.io.in.bits.c := cWords(elemReg)
-  dpu.io.stall := false.B
-  val dpuValid = dpu.io.out.valid
 
   when(io.writeback.fire) {
     wbValid := false.B
@@ -238,13 +273,11 @@ class TensorCoreBlackwell(
     addrAReg := io.initiate.bits.addressA
     addrBReg := io.initiate.bits.addressB
     addrCReg := io.initiate.bits.addressC
-    setReg := 0.U
     aIndexReg := 0.U
-    bGroupReg := 0.U
     bIndexReg := 0.U
-    mGroupReg := 0.U
-    substepReg := 0.U
-    elemReg := 0.U
+    tileMReg := 0.U
+    tileNReg := 0.U
+    cTileFragReg := 0.U
     switch(io.initiate.bits.op) {
       is(Ops.bwgmma) { state := State.bwLoadAReq }
       is(Ops.tcgen05Cp) { state := State.cpRead }
@@ -266,8 +299,7 @@ class TensorCoreBlackwell(
 
   when(state === State.bwLoadAResp) {
     aBuf(aIndexReg) := io.tmemC.aRdata
-    when(aIndexReg === (numAFragsPerSet - 1).U) {
-      bGroupReg := 0.U
+    when(aIndexReg === (totalAFrags - 1).U) {
       bIndexReg := 0.U
       state := State.bwLoadBReq
     }.otherwise {
@@ -292,9 +324,10 @@ class TensorCoreBlackwell(
     io.respB.ready := true.B
     when(io.respB.fire) {
       bBuf(bIndexReg) := io.respB.bits.data
-      when(bIndexReg === (numBFragsPerGroup - 1).U) {
-        mGroupReg := 0.U
-        substepReg := 0.U
+      when(bIndexReg === (totalBFrags - 1).U) {
+        tileMReg := 0.U
+        tileNReg := 0.U
+        cTileFragReg := 0.U
         state := State.bwReadCReq
       }.otherwise {
         bIndexReg := bIndexReg + 1.U
@@ -312,58 +345,57 @@ class TensorCoreBlackwell(
   }
 
   when(state === State.bwReadCResp) {
-    cDataReg := io.tmemC.cRdata
-    elemReg  := 0.U
-    state    := State.bwCompute
+    val cWords = io.tmemC.cRdata.asTypeOf(Vec(numLanes, UInt(laneWidth.W)))
+    for (lane <- 0 until numLanes) {
+      val elemM = lane % mElemsPerFrag
+      val elemN = lane / mElemsPerFrag
+      val row = localMGroup * mElemsPerFrag.U + elemM.U
+      val col = localBGroup * 4.U + localSubstep * 2.U + elemN.U
+      cTile(row)(col) := cWords(lane)
+    }
+    when(cTileFragReg === (numCFragsPerTile - 1).U) {
+      state := State.bwArrayStart
+    }.otherwise {
+      cTileFragReg := cTileFragReg + 1.U
+      state := State.bwReadCReq
+    }
   }
 
-  when(state === State.bwCompute) {
-    dpuInValid := true.B
-    state := State.bwDpuResp
+  when(state === State.bwArrayStart) {
+    when(systolic.io.ready) {
+      state := State.bwArrayRun
+    }
   }
 
-  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(state === State.bwArrayRun) {
+    when(systolic.io.done) {
+      cTileFragReg := 0.U
+      state := State.bwWriteCReq
     }
   }
 
   when(state === State.bwWriteCReq) {
     io.tmemC.cWen   := true.B
     io.tmemC.cWaddr := tmemCBase + cFragIndex
-    io.tmemC.cWdata := mmaDataReg.asUInt
+    io.tmemC.cWdata := mmaWords.asUInt
     io.tmemC.cMask  := Fill(maskWidth, 1.U(1.W))
     when(io.tmemC.cWready) {
-      when(substepReg === 0.U) {
-        substepReg := 1.U
-        state := State.bwReadCReq
-      }.elsewhen(mGroupReg =/= (numMGroups - 1).U) {
-        substepReg := 0.U
-        mGroupReg := mGroupReg + 1.U
-        state := State.bwReadCReq
-      }.elsewhen(bGroupReg =/= (numBGroups - 1).U) {
-        substepReg := 0.U
-        mGroupReg := 0.U
-        bGroupReg := bGroupReg + 1.U
-        bIndexReg := 0.U
-        state := State.bwLoadBReq
-      }.elsewhen(setReg =/= (numSets - 1).U) {
-        substepReg := 0.U
-        mGroupReg := 0.U
-        bGroupReg := 0.U
-        bIndexReg := 0.U
-        setReg := setReg + 1.U
-        aIndexReg := 0.U
-        state := State.bwLoadAReq
+      mmaDataReg := mmaWords
+      when(cTileFragReg =/= (numCFragsPerTile - 1).U) {
+        cTileFragReg := cTileFragReg + 1.U
       }.otherwise {
-        waitCounter := 7.U
-        state := State.bwWriteCWait
+        cTileFragReg := 0.U
+        when(tileNReg =/= (numTilesPerDim - 1).U) {
+          tileNReg := tileNReg + 1.U
+          state := State.bwReadCReq
+        }.elsewhen(tileMReg =/= (numTilesPerDim - 1).U) {
+          tileMReg := tileMReg + 1.U
+          tileNReg := 0.U
+          state := State.bwReadCReq
+        }.otherwise {
+          waitCounter := 7.U
+          state := State.bwWriteCWait
+        }
       }
     }
   }

src/main/scala/radiance/tile/RadianceTile.scala

@@ -851,6 +851,9 @@ class RadianceTileModuleImp(outer: RadianceTile)
       core.io.tc_tmem_C_rready := DontCare
       core.io.tc_tmem_C_rdata := DontCare
       core.io.tc_tmem_C_wready := DontCare
+      core.io.sc_tmem_rready := DontCare
+      core.io.sc_tmem_rdata := DontCare
+      core.io.sc_tmem_wready := DontCare
     }
 
     def connectTensorBlackwell = {
@@ -885,59 +888,225 @@ class RadianceTileModuleImp(outer: RadianceTile)
         tcDData.foreach(_ := 0.U)
         tcDTag.foreach(_ := 0.U)
 
-        // TMEM matrix: one shared 2R1W SRAM.  read0 is operand A, read1 is C.
-        // Each warp owns 2KB: A tile and C tile are 1KB each.  The row count
-        // scales with the physical fragment width (16B for 4 lanes, 32B for 8).
+        // TMEM keeps the ISA-visible address space unified while storing the
+        // A and C halves in separate 1R1W arrays.  This avoids duplicating each
+        // bank for two read ports, and still allows common A-read/C-read pairs
+        // to proceed in parallel because they normally hit different arrays.
         val tmemBytesPerWarp = 2048
-        val tmemDepth = outer.numWarps * (tmemBytesPerWarp / outer.tcSmemSize)
-        val tmem = Module(new radiance.memory.TwoReadOneWriteSyncMem(
-          tmemDepth, UInt((outer.tcSmemSize * 8).W)))
+        val tmemFragsPerWarp = tmemBytesPerWarp / outer.tcSmemSize
+        val tmemFragsPerTile = tmemFragsPerWarp / 2
+        val tmemLogicalDepth = outer.numWarps * tmemFragsPerWarp
+        val tmemArrayDepth = outer.numWarps * tmemFragsPerTile
+        val tmemBanks = 4
+        val tmemBankBits = log2Ceil(tmemBanks)
+        val tmemFragAddrBits = log2Ceil(tmemFragsPerWarp)
+        val tmemTileAddrBits = log2Ceil(tmemFragsPerTile)
+        val tmemWarpAddrBits = log2Ceil(outer.numWarps)
+        val tmemPhysAddrBits = log2Ceil(tmemArrayDepth)
+        val tmemBankDepth = tmemArrayDepth / tmemBanks
+        require(isPow2(tmemBanks))
+        require(isPow2(tmemFragsPerWarp))
+        require(tmemFragsPerWarp == tmemFragsPerTile * 2)
+        require(tmemLogicalDepth <= (1 << tmemAddrBits))
+        require(tmemArrayDepth % tmemBanks == 0)
+        require(tmemPhysAddrBits > tmemBankBits)
+        val tmemA = Seq.fill(tmemBanks) {
+          Module(new radiance.memory.TwoPortSyncMem(
+            tmemBankDepth, UInt((outer.tcSmemSize * 8).W)))
+        }
+        val tmemC = Seq.fill(tmemBanks) {
+          Module(new radiance.memory.TwoPortSyncMem(
+            tmemBankDepth, UInt((outer.tcSmemSize * 8).W)))
+        }
 
-        val aReadArb = Module(new RRArbiter(UInt(tmemAddrBits.W), nTC))
-        val cReadArb = Module(new RRArbiter(UInt(tmemAddrBits.W), nTC))
+        class TmemReadReq extends Bundle {
+          val addr = UInt(tmemAddrBits.W)
+          val src = UInt(2.W)
+          val tc = UInt(log2Ceil(nTC max 2).W)
+        }
 
         class TmemWriteReq extends Bundle {
           val addr = UInt(tmemAddrBits.W)
           val data = UInt(tmemDataBits.W)
           val mask = UInt(tmemMaskBits.W)
-        }
-        val cWriteArb = Module(new RRArbiter(new TmemWriteReq, nTC))
-
-        (0 until nTC).foreach { tc =>
-          aReadArb.io.in(tc).valid := core.io.tc_tmem_A_ren(tc)
-          aReadArb.io.in(tc).bits := slice(core.io.tc_tmem_A_raddr, tmemAddrBits, tc)
-          cReadArb.io.in(tc).valid := core.io.tc_tmem_C_ren(tc)
-          cReadArb.io.in(tc).bits := slice(core.io.tc_tmem_C_raddr, tmemAddrBits, tc)
-          cWriteArb.io.in(tc).valid := core.io.tc_tmem_C_wen(tc)
-          cWriteArb.io.in(tc).bits.addr := slice(core.io.tc_tmem_C_waddr, tmemAddrBits, tc)
-          cWriteArb.io.in(tc).bits.data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc)
-          cWriteArb.io.in(tc).bits.mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc)
+          val src = UInt(1.W)
+          val tc = UInt(log2Ceil(nTC max 2).W)
         }
 
-        aReadArb.io.out.ready := true.B
-        cReadArb.io.out.ready := true.B
-        cWriteArb.io.out.ready := true.B
+        def tmemIsC(addr: UInt): Bool = addr(tmemTileAddrBits)
+        def tmemPhysAddr(addr: UInt): UInt = {
+          val tileOffset = addr(tmemTileAddrBits - 1, 0)
+          if (tmemWarpAddrBits == 0) {
+            tileOffset
+          } else {
+            Cat(addr(tmemFragAddrBits + tmemWarpAddrBits - 1, tmemFragAddrBits), tileOffset)
+          }
+        }
+        def bank(addr: UInt): UInt = addr(tmemBankBits - 1, 0)
+        def row(addr: UInt): UInt = addr(tmemPhysAddrBits - 1, tmemBankBits)
 
-        tmem.io.ren0 := aReadArb.io.out.fire
-        tmem.io.raddr0 := aReadArb.io.out.bits
-        tmem.io.ren1 := cReadArb.io.out.fire
-        tmem.io.raddr1 := cReadArb.io.out.bits
-        tmem.io.wen := cWriteArb.io.out.fire
-        tmem.io.waddr := cWriteArb.io.out.bits.addr
-        tmem.io.wdata := cWriteArb.io.out.bits.data
-        tmem.io.mask := cWriteArb.io.out.bits.mask
+        val aReady = Wire(Vec(nTC, Bool()))
+        val cReady = Wire(Vec(nTC, Bool()))
+        val wReady = Wire(Vec(nTC, Bool()))
+        val scReadReady = Wire(Bool())
+        val scWriteReady = Wire(Bool())
+        aReady.foreach(_ := false.B)
+        cReady.foreach(_ := false.B)
+        wReady.foreach(_ := false.B)
+        scReadReady := false.B
+        scWriteReady := false.B
 
-        val aReadGrant = RegNext(Mux(aReadArb.io.out.fire, UIntToOH(aReadArb.io.chosen, nTC), 0.U(nTC.W)))
-        val cReadGrant = RegNext(Mux(cReadArb.io.out.fire, UIntToOH(cReadArb.io.chosen, nTC), 0.U(nTC.W)))
-        core.io.tc_tmem_A_rready := VecInit(aReadArb.io.in.map(_.fire)).asUInt
-        core.io.tc_tmem_C_rready := VecInit(cReadArb.io.in.map(_.fire)).asUInt
-        core.io.tc_tmem_C_wready := VecInit(cWriteArb.io.in.map(_.fire)).asUInt
+        val aReadGrant = Wire(Vec(tmemBanks, new TmemReadReq))
+        val cReadGrant = Wire(Vec(tmemBanks, new TmemReadReq))
+        val aReadValid = Wire(Vec(tmemBanks, Bool()))
+        val cReadValid = Wire(Vec(tmemBanks, Bool()))
+        val aWriteGrant = Wire(Vec(tmemBanks, new TmemWriteReq))
+        val cWriteGrant = Wire(Vec(tmemBanks, new TmemWriteReq))
+        val aWriteValid = Wire(Vec(tmemBanks, Bool()))
+        val cWriteValid = Wire(Vec(tmemBanks, Bool()))
+        aReadGrant.foreach(_ := 0.U.asTypeOf(new TmemReadReq))
+        cReadGrant.foreach(_ := 0.U.asTypeOf(new TmemReadReq))
+        aReadValid.foreach(_ := false.B)
+        cReadValid.foreach(_ := false.B)
+        aWriteGrant.foreach(_ := 0.U.asTypeOf(new TmemWriteReq))
+        cWriteGrant.foreach(_ := 0.U.asTypeOf(new TmemWriteReq))
+        aWriteValid.foreach(_ := false.B)
+        cWriteValid.foreach(_ := false.B)
+
+        (0 until tmemBanks).foreach { b =>
+          val readRequests = (0 until nTC).flatMap { tc =>
+            val aAddr = slice(core.io.tc_tmem_A_raddr, tmemAddrBits, tc)
+            val cAddr = slice(core.io.tc_tmem_C_raddr, tmemAddrBits, tc)
+            Seq(
+              (core.io.tc_tmem_A_ren(tc).asBool, aAddr, 0.U(2.W), tc.U),
+              (core.io.tc_tmem_C_ren(tc).asBool, cAddr, 1.U(2.W), tc.U)
+            )
+          } ++ Seq(
+            (core.io.sc_tmem_ren.asBool, core.io.sc_tmem_raddr, 2.U(2.W), 0.U)
+          )
+
+          var aReadUsed = false.B
+          var cReadUsed = false.B
+          readRequests.foreach { case (valid, addr, src, tc) =>
+            val physAddr = tmemPhysAddr(addr)
+            val isC = tmemIsC(addr)
+            val aGrant = valid && !isC && bank(physAddr) === b.U && !aReadUsed
+            val cGrant = valid && isC && bank(physAddr) === b.U && !cReadUsed
+            when(aGrant) {
+              aReadGrant(b).addr := physAddr
+              aReadGrant(b).src := src
+              aReadGrant(b).tc := tc
+            }
+            when(cGrant) {
+              cReadGrant(b).addr := physAddr
+              cReadGrant(b).src := src
+              cReadGrant(b).tc := tc
+            }
+            aReadUsed = aReadUsed || aGrant
+            cReadUsed = cReadUsed || cGrant
+            when(aGrant || cGrant) {
+              when(src === 0.U) { aReady(tc) := true.B }
+              when(src === 1.U) { cReady(tc) := true.B }
+              when(src === 2.U) { scReadReady := true.B }
+            }
+          }
+          aReadValid(b) := aReadUsed
+          cReadValid(b) := cReadUsed
+
+          var aWriteUsed = false.B
+          var cWriteUsed = false.B
+          (0 until nTC).foreach { tc =>
+            val addr = slice(core.io.tc_tmem_C_waddr, tmemAddrBits, tc)
+            val physAddr = tmemPhysAddr(addr)
+            val isC = tmemIsC(addr)
+            val valid = core.io.tc_tmem_C_wen(tc).asBool && bank(physAddr) === b.U
+            val aGrant = valid && !isC && !aWriteUsed
+            val cGrant = valid && isC && !cWriteUsed
+            when(aGrant) {
+              aWriteValid(b) := true.B
+              aWriteGrant(b).addr := physAddr
+              aWriteGrant(b).data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc)
+              aWriteGrant(b).mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc)
+              aWriteGrant(b).src := 0.U
+              aWriteGrant(b).tc := tc.U
+              wReady(tc) := true.B
+            }
+            when(cGrant) {
+              cWriteValid(b) := true.B
+              cWriteGrant(b).addr := physAddr
+              cWriteGrant(b).data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc)
+              cWriteGrant(b).mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc)
+              cWriteGrant(b).src := 0.U
+              cWriteGrant(b).tc := tc.U
+              wReady(tc) := true.B
+            }
+            aWriteUsed = aWriteUsed || aGrant
+            cWriteUsed = cWriteUsed || cGrant
+          }
+
+          val scWPhysAddr = tmemPhysAddr(core.io.sc_tmem_waddr)
+          val scWIsC = tmemIsC(core.io.sc_tmem_waddr)
+          val scWValid = core.io.sc_tmem_wen.asBool && bank(scWPhysAddr) === b.U
+          val scWAGrant = scWValid && !scWIsC && !aWriteUsed
+          val scWCGrant = scWValid && scWIsC && !cWriteUsed
+          when(scWAGrant) {
+            aWriteValid(b) := true.B
+            aWriteGrant(b).addr := scWPhysAddr
+            aWriteGrant(b).data := core.io.sc_tmem_wdata
+            aWriteGrant(b).mask := core.io.sc_tmem_mask
+            aWriteGrant(b).src := 1.U
+            aWriteGrant(b).tc := 0.U
+            scWriteReady := true.B
+          }
+          when(scWCGrant) {
+            cWriteValid(b) := true.B
+            cWriteGrant(b).addr := scWPhysAddr
+            cWriteGrant(b).data := core.io.sc_tmem_wdata
+            cWriteGrant(b).mask := core.io.sc_tmem_mask
+            cWriteGrant(b).src := 1.U
+            cWriteGrant(b).tc := 0.U
+            scWriteReady := true.B
+          }
+
+          tmemA(b).io.ren := aReadValid(b)
+          tmemA(b).io.raddr := row(aReadGrant(b).addr)
+          tmemA(b).io.wen := aWriteValid(b)
+          tmemA(b).io.waddr := row(aWriteGrant(b).addr)
+          tmemA(b).io.wdata := aWriteGrant(b).data
+          tmemA(b).io.mask := aWriteGrant(b).mask
+          tmemC(b).io.ren := cReadValid(b)
+          tmemC(b).io.raddr := row(cReadGrant(b).addr)
+          tmemC(b).io.wen := cWriteValid(b)
+          tmemC(b).io.waddr := row(cWriteGrant(b).addr)
+          tmemC(b).io.wdata := cWriteGrant(b).data
+          tmemC(b).io.mask := cWriteGrant(b).mask
+        }
+
+        val aReadGrantReg = RegNext(aReadGrant)
+        val cReadGrantReg = RegNext(cReadGrant)
+        val aReadValidReg = RegNext(aReadValid)
+        val cReadValidReg = RegNext(cReadValid)
+        core.io.tc_tmem_A_rready := aReady.asUInt
+        core.io.tc_tmem_C_rready := cReady.asUInt
+        core.io.tc_tmem_C_wready := wReady.asUInt
+        core.io.sc_tmem_rready := scReadReady.asUInt
+        core.io.sc_tmem_wready := scWriteReady.asUInt
         core.io.tc_tmem_A_rdata := VecInit((0 until nTC).map { tc =>
-          Mux(aReadGrant(tc), tmem.io.rdata0, 0.U(tmemDataBits.W))
+          VecInit((0 until tmemBanks).map { b =>
+            Mux(aReadValidReg(b) && aReadGrantReg(b).src === 0.U && aReadGrantReg(b).tc === tc.U, tmemA(b).io.rdata,
+              Mux(cReadValidReg(b) && cReadGrantReg(b).src === 0.U && cReadGrantReg(b).tc === tc.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
+          }).reduce(_ | _)
         }).asUInt
         core.io.tc_tmem_C_rdata := VecInit((0 until nTC).map { tc =>
-          Mux(cReadGrant(tc), tmem.io.rdata1, 0.U(tmemDataBits.W))
+          VecInit((0 until tmemBanks).map { b =>
+            Mux(aReadValidReg(b) && aReadGrantReg(b).src === 1.U && aReadGrantReg(b).tc === tc.U, tmemA(b).io.rdata,
+              Mux(cReadValidReg(b) && cReadGrantReg(b).src === 1.U && cReadGrantReg(b).tc === tc.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
+          }).reduce(_ | _)
         }).asUInt
+        core.io.sc_tmem_rdata := VecInit((0 until tmemBanks).map { b =>
+          Mux(aReadValidReg(b) && aReadGrantReg(b).src === 2.U, tmemA(b).io.rdata,
+            Mux(cReadValidReg(b) && cReadGrantReg(b).src === 2.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
+        }).reduce(_ | _)
 
         // port 2: SMEM B, one TL client per tensor core. RadianceSharedMem arbitrates them.
         (0 until nTC).foreach { tc =>
@@ -1025,6 +1194,9 @@ class RadianceTileModuleImp(outer: RadianceTile)
         core.io.tc_tmem_C_rready := DontCare
         core.io.tc_tmem_C_rdata := DontCare
         core.io.tc_tmem_C_wready := DontCare
+        core.io.sc_tmem_rready := DontCare
+        core.io.sc_tmem_rdata := DontCare
+        core.io.sc_tmem_wready := DontCare
       }
     }
 
@@ -1197,18 +1369,6 @@ class VortexTLAdapter(
     val outResp = chiselTypeOf(outTL._1.d)
   })
   val (bundle, edge) = outTL
-  val sourceGen = Module(
-    new SourceGenerator(
-      newSourceWidth,
-      Some(inReqT.source),
-      ignoreInUse = false
-    )
-  )
-  sourceGen.io.gen := io.outReq.fire // use up a source ID only when request is created
-  sourceGen.io.reclaim.valid := io.outResp.fire
-  sourceGen.io.reclaim.bits := io.outResp.bits.source
-  sourceGen.io.meta := io.inReq.bits.source
-
   // io passthrough logic
   // TLBundleA <> VortexBundleA
   io.outReq.valid := io.inReq.valid
@@ -1217,29 +1377,70 @@ class VortexTLAdapter(
   io.outReq.bits.size := io.inReq.bits.size
   io.outReq.bits.source := io.inReq.bits.source
   io.outReq.bits.address := io.inReq.bits.address
-  // Get requires contiguous mask; only copy core's potentially-partial mask
-  // when writing
+  val outMaskWidth = io.outReq.bits.mask.getWidth
+  val inMaskWidth = io.inReq.bits.mask.getWidth
+  val outDataWidth = io.outReq.bits.data.getWidth
+  val inDataWidth = io.inReq.bits.data.getWidth
+  val byteOffset = io.inReq.bits.address(log2Ceil(outMaskWidth) - 1, 0)
+  val responseOffsetWidth = log2Ceil(outMaskWidth)
+  val responseSourceWidth = inReqT.source.getWidth
+  val sourceGen = Module(
+    new SourceGenerator(
+      newSourceWidth,
+      Some(UInt((responseSourceWidth + responseOffsetWidth).W)),
+      ignoreInUse = false
+    )
+  )
+  sourceGen.io.gen := io.outReq.fire // use up a source ID only when request is created
+  sourceGen.io.reclaim.valid := io.outResp.fire
+  sourceGen.io.reclaim.bits := io.outResp.bits.source
+  sourceGen.io.meta := Cat(byteOffset, io.inReq.bits.source)
+
+  val alignedMask = Wire(UInt(outMaskWidth.W))
+  val alignedData = Wire(UInt(outDataWidth.W))
+  if (outMaskWidth == inMaskWidth) {
+    alignedMask := io.inReq.bits.mask
+  } else {
+    val paddedMask = Wire(UInt(outMaskWidth.W))
+    paddedMask := io.inReq.bits.mask
+    alignedMask := (paddedMask << byteOffset)(outMaskWidth - 1, 0)
+  }
+  if (outDataWidth == inDataWidth) {
+    alignedData := io.inReq.bits.data
+  } else {
+    val paddedData = Wire(UInt(outDataWidth.W))
+    paddedData := io.inReq.bits.data
+    alignedData := (paddedData << (byteOffset << 3))(outDataWidth - 1, 0)
+  }
+
+  // PutFull requires the TL-canonical full mask for address+size; PutPartial
+  // can carry the core-provided byte mask.
   io.outReq.bits.mask := Mux(
-    edge.hasData(io.outReq.bits),
-    io.inReq.bits.mask,
-    // generate TL-correct mask
+    io.outReq.bits.opcode === TLMessages.PutPartialData,
+    alignedMask,
     edge.mask(io.inReq.bits.address, io.inReq.bits.size)
   )
-  io.outReq.bits.data := io.inReq.bits.data
+  io.outReq.bits.data := alignedData
   io.outReq.bits.corrupt := 0.U
   io.inReq.ready := io.outReq.ready
   // VortexBundleD <> TLBundleD
   io.inResp.valid := io.outResp.valid
   io.inResp.bits.opcode := io.outResp.bits.opcode
   io.inResp.bits.size := io.outResp.bits.size
-  io.inResp.bits.source := io.outResp.bits.source
-  io.inResp.bits.data := io.outResp.bits.data
+  val responseMeta = sourceGen.io.peek.asUInt
+  val responseSource = responseMeta(responseSourceWidth - 1, 0)
+  val responseByteOffset =
+    responseMeta(responseSourceWidth + responseOffsetWidth - 1, responseSourceWidth)
+  io.inResp.bits.source := responseSource
+  if (outDataWidth == inDataWidth) {
+    io.inResp.bits.data := io.outResp.bits.data
+  } else {
+    io.inResp.bits.data := (io.outResp.bits.data >> (responseByteOffset << 3))(inDataWidth - 1, 0)
+  }
   io.outResp.ready := io.inResp.ready
 
   // "man-in-the-middle"
   io.inReq.ready := io.outReq.ready && sourceGen.io.id.valid
   io.outReq.valid := io.inReq.valid && sourceGen.io.id.valid
   io.outReq.bits.source := sourceGen.io.id.bits
-  // translate upstream response back to its old sourceId
-  io.inResp.bits.source := sourceGen.io.peek
 }

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

@@ -120,6 +120,15 @@ class VortexBundle(tile: RadianceTile)(implicit p: Parameters) extends CoreBundl
   val tc_tmem_C_waddr  = Output(UInt((numTensorCores * 9).W))
   val tc_tmem_C_wdata  = Output(UInt((numTensorCores * numLanes * 32).W))
   val tc_tmem_C_mask   = Output(UInt((numTensorCores * numLanes * 4).W))
+  val sc_tmem_ren      = Output(UInt(1.W))
+  val sc_tmem_rready   = Input(UInt(1.W))
+  val sc_tmem_raddr    = Output(UInt(9.W))
+  val sc_tmem_rdata    = Input(UInt((numLanes * 32).W))
+  val sc_tmem_wen      = Output(UInt(1.W))
+  val sc_tmem_wready   = Input(UInt(1.W))
+  val sc_tmem_waddr    = Output(UInt(9.W))
+  val sc_tmem_wdata    = Output(UInt((numLanes * 32).W))
+  val sc_tmem_mask     = Output(UInt((numLanes * 4).W))
 
   // FIXME: hardcoded
   val barrierIdBits = tile.barrierMasterNode.out(0)._2.barrierIdBits
@@ -204,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")
@@ -351,6 +361,7 @@ class Vortex(tile: RadianceTile)(implicit p: Parameters)
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_div.sv")
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_dpi.sv")
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_dsp.sv")
+  addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_exp.sv")
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_fma.sv")
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_ncomp.sv")
   addResource("/vsrc/vortex/hw/rtl/fpu/VX_fpu_rounding.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) {