feat: execute blackwell bwgmma on systolic array
This commit is contained in:
Zhongdi LUO
114
src/main/scala/radiance/core/BlackwellSystolicArray.scala
Normal file
114
src/main/scala/radiance/core/BlackwellSystolicArray.scala
Normal file
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user