From 7b769e9e8490084d56361b91c9a6256a499947be Mon Sep 17 00:00:00 2001 From: Zhongdi LUO Date: Wed, 15 Jul 2026 07:26:13 +0000 Subject: [PATCH] feat: execute blackwell bwgmma on systolic array --- .../core/BlackwellSystolicArray.scala | 114 ++++++++++ .../radiance/core/TensorCoreBlackwell.scala | 194 ++++++++++-------- 2 files changed, 218 insertions(+), 90 deletions(-) create mode 100644 src/main/scala/radiance/core/BlackwellSystolicArray.scala diff --git a/src/main/scala/radiance/core/BlackwellSystolicArray.scala b/src/main/scala/radiance/core/BlackwellSystolicArray.scala new file mode 100644 index 0000000..43a6c5d --- /dev/null +++ b/src/main/scala/radiance/core/BlackwellSystolicArray.scala @@ -0,0 +1,114 @@ +// 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 systolic array for Blackwell FP8 BWGMMA. + * + * A operands move left-to-right and B operands move top-to-bottom. Each PE + * owns one FP32 accumulator and performs one FP8 E4M3 MAC per active cycle. + */ +class BlackwellSystolicArray(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 tOut = tile.FType.S + private val recWidth = 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(8.W)))) + val b = Input(Vec(kDim, Vec(arrayDim, UInt(8.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(8.W)))))) + val bPipe = RegInit(VecInit(Seq.fill(arrayDim)(VecInit(Seq.fill(arrayDim)(0.U(8.W)))))) + val accum = Reg(Vec(arrayDim, Vec(arrayDim, UInt(recWidth.W)))) + + val aFlow = Wire(Vec(arrayDim, Vec(arrayDim, UInt(8.W)))) + val bFlow = Wire(Vec(arrayDim, Vec(arrayDim, UInt(8.W)))) + + for (row <- 0 until arrayDim) { + val aActive = stepReg >= row.U && stepReg < (row + kDim).U + val aIndex = (stepReg - row.U)(kIndexWidth - 1, 0) + aFlow(row)(0) := Mux(aActive, io.a(row)(aIndex), 0.U) + for (col <- 1 until arrayDim) { + aFlow(row)(col) := aPipe(row)(col - 1) + } + } + + for (col <- 0 until arrayDim) { + val bActive = stepReg >= col.U && stepReg < (col + kDim).U + val bIndex = (stepReg - col.U)(kIndexWidth - 1, 0) + bFlow(0)(col) := Mux(bActive, io.b(bIndex)(col), 0.U) + for (row <- 1 until arrayDim) { + bFlow(row)(col) := bPipe(row - 1)(col) + } + } + + val sums = Seq.tabulate(arrayDim, arrayDim) { (row, col) => + val product = unbox( + recode(FP8E4M3MulToFloat32(aFlow(row)(col), bFlow(row)(col)), S), + S, + Some(tOut)) + 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 + 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 + } + } +} diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala index ad45d5a..84a0095 100644 --- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala +++ b/src/main/scala/radiance/core/TensorCoreBlackwell.scala @@ -60,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) @@ -129,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 } @@ -143,18 +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 issueElemReg = RegInit(0.U(log2Ceil(numLanes).W)) - val retireElemReg = 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))) @@ -166,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) @@ -211,31 +220,44 @@ 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 = false, - inputType = TensorInputType.FP8E4M3 - )) + val systolic = Module(new BlackwellSystolicArray(systolicK)) + systolic.io.start := state === State.bwArrayStart + systolic.io.c := cTile - 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) := TensorCoreBlackwellFP8Packing.selectA(operandA, k, elemM, numLanes) - dpu.io.in.bits.b(k) := TensorCoreBlackwellFP8Packing.selectB(operandB, k, elemN) + // Preserve the software-visible 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) + systolic.io.a(row)(k) := TensorCoreBlackwellFP8Packing.selectA( + aBuf(aIndex), kInSet, elemM, numLanes) + } + + 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 + systolic.io.b(k)(col) := TensorCoreBlackwellFP8Packing.selectB( + bBuf(bIndex), 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(issueElemReg) - dpu.io.stall := false.B - val dpuValid = dpu.io.out.valid when(io.writeback.fire) { wbValid := false.B @@ -248,14 +270,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 - issueElemReg := 0.U - retireElemReg := 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 } @@ -277,8 +296,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 { @@ -303,9 +321,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 @@ -323,62 +342,57 @@ class TensorCoreBlackwell( } when(state === State.bwReadCResp) { - cDataReg := io.tmemC.cRdata - issueElemReg := 0.U - retireElemReg := 0.U - state := State.bwCompute - } - - when(state === State.bwCompute) { - dpuInValid := true.B - when(issueElemReg === (numLanes - 1).U) { - state := State.bwDpuResp + 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 { - issueElemReg := issueElemReg + 1.U + cTileFragReg := cTileFragReg + 1.U + state := State.bwReadCReq } } - val inDpuResponseWindow = state === State.bwCompute || state === State.bwDpuResp + when(state === State.bwArrayStart) { + when(systolic.io.ready) { + state := State.bwArrayRun + } + } - when(dpuValid && inDpuResponseWindow) { - mmaDataReg(retireElemReg) := dpu.io.out.bits.data - when(retireElemReg === (numLanes - 1).U) { + when(state === State.bwArrayRun) { + when(systolic.io.done) { + cTileFragReg := 0.U state := State.bwWriteCReq - }.otherwise { - retireElemReg := retireElemReg + 1.U } } 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 + } } } }