From a978855cefe94c332f38dcf61426a6410426ea27 Mon Sep 17 00:00:00 2001 From: Zhongdi LUO Date: Thu, 16 Jul 2026 00:19:31 +0000 Subject: [PATCH] feat: execute blackwell fp16 bwgmma on systolic array --- .../core/BlackwellFP16SystolicArray.scala | 129 ++++++++++ .../radiance/core/TensorCoreBlackwell.scala | 232 ++++++++++-------- 2 files changed, 258 insertions(+), 103 deletions(-) create mode 100644 src/main/scala/radiance/core/BlackwellFP16SystolicArray.scala diff --git a/src/main/scala/radiance/core/BlackwellFP16SystolicArray.scala b/src/main/scala/radiance/core/BlackwellFP16SystolicArray.scala new file mode 100644 index 0000000..51074a1 --- /dev/null +++ b/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 + } + } +} diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala index 0224d92..d1c58de 100644 --- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala +++ b/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,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))) @@ -143,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) @@ -188,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) 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) := ( - 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(issueElemReg) - dpu.io.stall := false.B - val dpuValid = dpu.io.out.valid when(io.writeback.fire) { wbValid := false.B @@ -239,14 +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 - 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 } @@ -268,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 { @@ -294,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 @@ -314,62 +345,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 } } - 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) { + when(state === State.bwArrayStart) { + when(systolic.io.ready) { + state := State.bwArrayRun + } + } + + 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 + } } } }