diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala index f595956..ad45d5a 100644 --- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala +++ b/src/main/scala/radiance/core/TensorCoreBlackwell.scala @@ -338,9 +338,9 @@ class TensorCoreBlackwell( } } - when(dpuValid) { - assert(state === State.bwCompute || state === State.bwDpuResp, - "BWGMMA DPU response arrived outside the compute states") + val inDpuResponseWindow = state === State.bwCompute || state === State.bwDpuResp + + when(dpuValid && inDpuResponseWindow) { mmaDataReg(retireElemReg) := dpu.io.out.bits.data when(retireElemReg === (numLanes - 1).U) { state := State.bwWriteCReq diff --git a/src/main/scala/radiance/core/TensorDPU.scala b/src/main/scala/radiance/core/TensorDPU.scala index 8bed9b6..9d53439 100644 --- a/src/main/scala/radiance/core/TensorDPU.scala +++ b/src/main/scala/radiance/core/TensorDPU.scala @@ -381,7 +381,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters val inC = unbox(recode(io.in.bits.c, S), S, Some(tOut)) val productStageOut = StallingPipe(io.stall, io.in.valid, VecInit(productRecoded)) - val productStageC = StallingPipe(io.stall, io.in.valid, inC) + val productStageC = RegEnable(inC, !io.stall && io.in.valid) val log2Dim = log2Ceil(dim) require(dim == (1 << log2Dim), s"dim (${dim}) is not power of two!") @@ -389,7 +389,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters val interim = (log2Dim to 0 by -1).map { i => Wire(Valid(Vec(1 << i, Bits(recOutFLen.W)))) } - val interimC = (log2Dim to 0 by -1).map(_ => Wire(Valid(Bits(recOutFLen.W)))) + val interimC = (log2Dim to 0 by -1).map(_ => Wire(Bits(recOutFLen.W))) interim(0) := productStageOut interimC(0) := productStageC @@ -413,11 +413,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters } outputs := StallingPipe(io.stall, inputs.valid, VecInit(addOuts)) - outC := StallingPipe(io.stall, inputs.valid, inC.bits) - when(inputs.valid =/= inC.valid) { - printf("WARN: DotProductPipeFP8E4M3 input/C valid mismatch: inputs=%d c=%d\n", - inputs.valid, inC.valid) - } + outC := RegEnable(inC, !io.stall && inputs.valid) (outputs, outC) } @@ -427,7 +423,7 @@ class DotProductPipeFP8E4M3(dim: Int) extends Module with tile.HasFPUParameters val acc = Module(new hardfloat.AddRecFN(outExpWidth, outSigWidth)) acc.io.subOp := 0.U acc.io.a := addStageOut.bits(0) - acc.io.b := addStageC.bits + acc.io.b := addStageC acc.io.roundingMode := hardfloat.consts.round_near_even acc.io.detectTininess := hardfloat.consts.tininess_afterRounding diff --git a/src/test/scala/radiance/TensorCoreBlackwellTest.scala b/src/test/scala/radiance/TensorCoreBlackwellTest.scala index c872ccd..a481df5 100644 --- a/src/test/scala/radiance/TensorCoreBlackwellTest.scala +++ b/src/test/scala/radiance/TensorCoreBlackwellTest.scala @@ -324,4 +324,80 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester { } } } + + it should "run bwgmma with 4-lane 16-byte fragments" in { + val lanes = 4 + test(new TensorCoreBlackwell(numWarps, lanes, half = true, numSourceIds = 4)) + .withAnnotations(Seq(VerilatorBackendAnnotation)) { c => + idleIO(c) + + val fragBytes = 16 + val aBase = BigInt(0x100) + val bBase = BigInt(0x800) + val cBase = BigInt(0x1000) + val fp8One = BigInt(0x38) + val fp8Two = BigInt(0x40) + val fp32One = BigInt(0x3f800000) + val fp32SixtyFive = BigInt(0x42820000) + val aFrag = packWords(Seq.fill(16)(fp8One), 8) + val bFrag = packWords(Seq.fill(16)(fp8Two), 8) + val cFrag = packWords(Seq.fill(lanes)(fp32One), 32) + val expectedCFrag = packWords(Seq.fill(lanes)(fp32SixtyFive), 32) + + val tmem = makeTmem() + for (i <- 0 until 64) { + tmem(aBase / fragBytes + i) = aFrag + tmem(cBase / fragBytes + i) = cFrag + } + val bMem = mutable.Map[BigInt, BigInt]() + for (i <- 0 until 64) bMem(bBase + i * fragBytes) = bFrag + + c.io.reqB.ready.poke(true.B) + c.io.writeback.ready.poke(true.B) + c.io.initiate.valid.poke(true.B) + c.io.initiate.bits.op.poke(0.U) + c.io.initiate.bits.wid.poke(1.U) + c.io.initiate.bits.rd.poke(0.U) + c.io.initiate.bits.addressA.poke(aBase.U) + c.io.initiate.bits.addressB.poke(bBase.U) + c.io.initiate.bits.addressC.poke(cBase.U) + c.clock.step() + c.io.initiate.valid.poke(false.B) + + var pendingB = Option.empty[(BigInt, BigInt)] + var sawWriteback = false + var cycles = 0 + for (_ <- 0 until 30000 if !sawWriteback) { + stepTmem(c, tmem) + pendingB.foreach { case (src, data) => + c.io.respB.valid.poke(true.B) + c.io.respB.bits.source.poke(src.U) + c.io.respB.bits.data.poke(data.U) + } + if (pendingB.isEmpty) c.io.respB.valid.poke(false.B) + + if (c.io.writeback.valid.peek().litToBoolean) { + sawWriteback = true + } else { + val nextB = if (c.io.reqB.valid.peek().litToBoolean) { + val addr = c.io.reqB.bits.address.peek().litValue + val src = c.io.reqB.bits.source.peek().litValue + Some((src, bMem(addr))) + } else None + c.clock.step() + cycles += 1 + pendingB = nextB + } + } + + assert(sawWriteback, "4-lane BWGMMA did not complete") + assert(cycles < 10000, + s"4-lane BWGMMA took $cycles cycles") + for (i <- 0 until 64) { + val row = cBase / fragBytes + i + assert(tmem(row) == expectedCFrag, + s"4-lane C frag $i mismatch: got 0x${tmem(row).toString(16)}") + } + } + } }