diff --git a/src/main/resources/vsrc/vortex b/src/main/resources/vsrc/vortex index 2bfc6c4..eb1b2ea 160000 --- a/src/main/resources/vsrc/vortex +++ b/src/main/resources/vsrc/vortex @@ -1 +1 @@ -Subproject commit 2bfc6c4bdef3a8b33bb8b53fae5b82e44df480f0 +Subproject commit eb1b2ea330ccb24d8e5fadcc517dce9c267c73a2 diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala index 30f2f48..0224d92 100644 --- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala +++ b/src/main/scala/radiance/core/TensorCoreBlackwell.scala @@ -126,7 +126,8 @@ class TensorCoreBlackwell( 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 issueElemReg = RegInit(0.U(log2Ceil(numLanes).W)) + val retireElemReg = RegInit(0.U(log2Ceil(numLanes).W)) val waitCounter = RegInit(0.U(3.W)) val aBuf = Reg(Vec(numAFragsPerSet, UInt(memWidth.W))) @@ -205,8 +206,8 @@ class TensorCoreBlackwell( x((idx + 1) * 16 - 1, idx * 16) } - val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0) - val elemN = if (numLanes == 4) elemReg(1) else elemReg(2) + 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) := ( @@ -223,7 +224,7 @@ class TensorCoreBlackwell( ) dpu.io.in.bits.b(k) := Mux(elemN.asBool, halfWord(operandB, 8 + k), halfWord(operandB, k)) } - dpu.io.in.bits.c := cWords(elemReg) + dpu.io.in.bits.c := cWords(issueElemReg) dpu.io.stall := false.B val dpuValid = dpu.io.out.valid @@ -244,7 +245,8 @@ class TensorCoreBlackwell( bIndexReg := 0.U mGroupReg := 0.U substepReg := 0.U - elemReg := 0.U + issueElemReg := 0.U + retireElemReg := 0.U switch(io.initiate.bits.op) { is(Ops.bwgmma) { state := State.bwLoadAReq } is(Ops.tcgen05Cp) { state := State.cpRead } @@ -313,24 +315,28 @@ class TensorCoreBlackwell( when(state === State.bwReadCResp) { cDataReg := io.tmemC.cRdata - elemReg := 0.U + issueElemReg := 0.U + retireElemReg := 0.U state := State.bwCompute } when(state === State.bwCompute) { dpuInValid := true.B - state := State.bwDpuResp + when(issueElemReg === (numLanes - 1).U) { + state := State.bwDpuResp + }.otherwise { + issueElemReg := issueElemReg + 1.U + } } - 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(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) { + state := State.bwWriteCReq + }.otherwise { + retireElemReg := retireElemReg + 1.U } } diff --git a/src/main/scala/radiance/tile/VortexCore.scala b/src/main/scala/radiance/tile/VortexCore.scala index 1b06a88..dfef1bb 100644 --- a/src/main/scala/radiance/tile/VortexCore.scala +++ b/src/main/scala/radiance/tile/VortexCore.scala @@ -213,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") diff --git a/src/test/scala/radiance/TensorCoreBlackwellTest.scala b/src/test/scala/radiance/TensorCoreBlackwellTest.scala index feb8008..8291a96 100644 --- a/src/test/scala/radiance/TensorCoreBlackwellTest.scala +++ b/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) {