From bb1459a2090ad6ec6e9457ab0597a497b951912b Mon Sep 17 00:00:00 2001 From: Zhongdi LUO Date: Mon, 13 Jul 2026 06:29:04 +0000 Subject: [PATCH] feat: add 4-lane pre-WU Blackwell configuration --- src/main/resources/vsrc/vortex | 2 +- .../radiance/core/TensorCoreBlackwell.scala | 54 ++++++++++++------- .../scala/radiance/tile/RadianceTile.scala | 29 +++++----- src/main/scala/radiance/tile/VortexCore.scala | 9 ++-- 4 files changed, 57 insertions(+), 37 deletions(-) diff --git a/src/main/resources/vsrc/vortex b/src/main/resources/vsrc/vortex index 323ed7d..9560f9c 160000 --- a/src/main/resources/vsrc/vortex +++ b/src/main/resources/vsrc/vortex @@ -1 +1 @@ -Subproject commit 323ed7d7e9c3fd403a5ceb5bba7e4371aa37f6fc +Subproject commit 9560f9cab6294c55a4127404645715dec40d3ff5 diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala index 2068867..e33c37a 100644 --- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala +++ b/src/main/scala/radiance/core/TensorCoreBlackwell.scala @@ -14,7 +14,8 @@ class TensorCoreBlackwell( val numFPRegs: Int = 32 ) extends Module { require(half, "Blackwell MMA currently supports FP16 inputs only") - require(numLanes == 8, "Blackwell MMA currently assumes 8 lanes") + require(numLanes == 4 || numLanes == 8, + s"Blackwell MMA currently supports 4 or 8 lanes, got ${numLanes}") val numWarpBits = log2Ceil(numWarps) val sourceWidth = log2Ceil(numSourceIds) @@ -26,11 +27,16 @@ class TensorCoreBlackwell( val fragOffsetBits = log2Ceil(memWidth / 8) val numSets = 4 - val numAFragsPerSet = 8 val numBGroups = 4 - val numBFragsPerGroup = 2 - val numMGroups = 4 - val numCFrags = 32 + val numSubsteps = 2 + val mElemsPerFrag = if (numLanes == 4) 2 else 4 + val numMGroups = 16 / mElemsPerFrag + val numAFragsPerMGroup = 2 + val numAFragsPerSet = numMGroups * numAFragsPerMGroup + val numBFragsPerSubstep = if (numLanes == 4) 2 else 1 + val numBFragsPerGroup = numSubsteps * numBFragsPerSubstep + val numBFragsPerSet = numBGroups * numBFragsPerGroup + val numCFrags = numBGroups * numMGroups * numSubsteps object Ops { val bwgmma :: bwgmmaWait :: tcgen05Cp :: tcgen05CpWait :: tcgen05Ld :: tcgen05St :: tcgen05Cb :: Nil = Enum(7) @@ -128,10 +134,11 @@ class TensorCoreBlackwell( base + (fragIndex << fragOffsetBits).asUInt } - val aFragIndex = (setReg << 3) + aIndexReg - val bFragIndex = (setReg << 3) + (bGroupReg << 1) + bIndexReg - val stepIndex = Cat(bGroupReg, mGroupReg) - val cFragIndex = (stepIndex << 1) + substepReg + val aFragIndex = (setReg * numAFragsPerSet.U) + aIndexReg + val bFragIndex = + (setReg * numBFragsPerSet.U) + (bGroupReg * numBFragsPerGroup.U) + bIndexReg + val cFragIndex = + (((bGroupReg * numMGroups.U) + mGroupReg) * numSubsteps.U) + substepReg val aReqAddress = byteAddress(addrAReg, aFragIndex) val bReqAddress = byteAddress(addrBReg, bFragIndex) val cReqAddress = byteAddress(addrCReg, cFragIndex) @@ -171,7 +178,12 @@ class TensorCoreBlackwell( io.initiate.ready := state === State.idle && !wbValid val operandA = Cat(aBuf((mGroupReg << 1) + 1.U), aBuf(mGroupReg << 1)) - val operandB = bBuf(substepReg) + 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( @@ -183,16 +195,22 @@ class TensorCoreBlackwell( x((idx + 1) * 16 - 1, idx * 16) } - val elemM = elemReg(1, 0) - val elemN = elemReg(2) + val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0) + val elemN = if (numLanes == 4) elemReg(1) else elemReg(2) dpu.io.in.valid := dpuInValid for (k <- 0 until 8) { - dpu.io.in.bits.a(k) := 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.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)) } dpu.io.in.bits.c := cWords(elemReg) diff --git a/src/main/scala/radiance/tile/RadianceTile.scala b/src/main/scala/radiance/tile/RadianceTile.scala index d74e91b..2c02ce2 100644 --- a/src/main/scala/radiance/tile/RadianceTile.scala +++ b/src/main/scala/radiance/tile/RadianceTile.scala @@ -285,7 +285,7 @@ class RadianceTile private ( ) } - val tcSmemSize = 32 + val tcSmemSize = numLsuLanes * 4 val tensorUsesAsyncMem = radianceParams.core.tensorCoreDecoupled || radianceParams.core.tensorCoreBlackwell val tcSmemNodeCount = if (radianceParams.core.tensorCoreDecoupled) 2 else if (radianceParams.core.tensorCoreBlackwell) 1 else 0 val tcSmemNodes = Seq.tabulate(tcSmemNodeCount) { i => @@ -828,9 +828,10 @@ class RadianceTileModuleImp(outer: RadianceTile) if (outer.radianceParams.core.tensorCoreBlackwell) { require(outer.tcSmemNodes.nonEmpty) - // TMEM C matrix: direct SRAM (no TileLink), connected via VortexCore IO - // Each warp needs 2 tiles (A + C), each tile = 32 frags × 32B = 1KB - val tmemDepth = outer.numWarps * outer.tcSmemSize * 2 // numWarps × 64 rows + // TMEM matrix: direct SRAM (no TileLink), connected via VortexCore IO. + // Each warp owns 2KB; row count scales with the lane-dependent fragment width. + val tmemBytesPerWarp = 2048 + val tmemDepth = outer.numWarps * (tmemBytesPerWarp / outer.tcSmemSize) val tmem = Module(new radiance.memory.TwoReadOneWriteSyncMem( tmemDepth, UInt((outer.tcSmemSize * 8).W))) tmem.io.ren0 := core.io.tc_tmem_C_ren @@ -848,23 +849,23 @@ class RadianceTileModuleImp(outer: RadianceTile) val addr = core.io.tc_a_bits_address(95, 64) val tag = core.io.tc_a_bits_tag(8 + outer.tensorTagWidth - 1, 8) val write = core.io.tc_a_bits_write(2) - val mask = core.io.tc_a_bits_mask(95, 64) - val data = core.io.tc_a_bits_data(767, 512) + val mask = core.io.tc_a_bits_mask(3 * outer.tcSmemSize - 1, 2 * outer.tcSmemSize) + val data = core.io.tc_a_bits_data(3 * outer.tcSmemSize * 8 - 1, 2 * outer.tcSmemSize * 8) val aValid = core.io.tc_a_valid(2) val dReady = core.io.tc_d_ready(2) } val client = outer.tcSmemNodes.head.out.head val adapter = Module(new VortexTLAdapter( outer.smemSourceWidth, - new VortexBundleA(tagWidth = outer.tensorTagWidth, dataWidth = 32 * 8), - new VortexBundleD(tagWidth = outer.tensorTagWidth, dataWidth = 32 * 8), + new VortexBundleA(tagWidth = outer.tensorTagWidth, dataWidth = outer.tcSmemSize * 8), + new VortexBundleD(tagWidth = outer.tensorTagWidth, dataWidth = outer.tcSmemSize * 8), client )) adapter.io.inReq.bits <> DontCare adapter.io.inReq.valid := smemBBundle.aValid adapter.io.inReq.bits.address := smemBBundle.addr adapter.io.inReq.bits.source := smemBBundle.tag - adapter.io.inReq.bits.size := 5.U + adapter.io.inReq.bits.size := log2Ceil(outer.tcSmemSize).U adapter.io.inReq.bits.opcode := Mux(smemBBundle.write.asBool, TLMessages.PutFullData, TLMessages.Get) adapter.io.inReq.bits.mask := smemBBundle.mask adapter.io.inReq.bits.data := smemBBundle.data @@ -876,18 +877,18 @@ class RadianceTileModuleImp(outer: RadianceTile) val gmemClient = outer.tcGmemNode.get.out.head val gmemAdapter = Module(new VortexTLAdapter( outer.dmemSourceWidth, - new VortexBundleA(tagWidth = outer.tensorTagWidth, dataWidth = 32 * 8), - new VortexBundleD(tagWidth = outer.tensorTagWidth, dataWidth = 32 * 8), + new VortexBundleA(tagWidth = outer.tensorTagWidth, dataWidth = outer.tcSmemSize * 8), + new VortexBundleD(tagWidth = outer.tensorTagWidth, dataWidth = outer.tcSmemSize * 8), gmemClient )) gmemAdapter.io.inReq.bits <> DontCare gmemAdapter.io.inReq.valid := core.io.tc_a_valid(0) gmemAdapter.io.inReq.bits.address := core.io.tc_a_bits_address(31, 0) gmemAdapter.io.inReq.bits.source := core.io.tc_a_bits_tag(outer.tensorTagWidth - 1, 0) - gmemAdapter.io.inReq.bits.size := 5.U + gmemAdapter.io.inReq.bits.size := log2Ceil(outer.tcSmemSize).U gmemAdapter.io.inReq.bits.opcode := Mux(core.io.tc_a_bits_write(0).asBool, TLMessages.PutFullData, TLMessages.Get) - gmemAdapter.io.inReq.bits.mask := core.io.tc_a_bits_mask(31, 0) - gmemAdapter.io.inReq.bits.data := core.io.tc_a_bits_data(255, 0) + gmemAdapter.io.inReq.bits.mask := core.io.tc_a_bits_mask(outer.tcSmemSize - 1, 0) + gmemAdapter.io.inReq.bits.data := core.io.tc_a_bits_data(outer.tcSmemSize * 8 - 1, 0) gmemAdapter.io.inResp.ready := core.io.tc_d_ready(0) gmemClient._1.a <> gmemAdapter.io.outReq gmemAdapter.io.outResp <> gmemClient._1.d diff --git a/src/main/scala/radiance/tile/VortexCore.scala b/src/main/scala/radiance/tile/VortexCore.scala index 91b65b3..cb7bc78 100644 --- a/src/main/scala/radiance/tile/VortexCore.scala +++ b/src/main/scala/radiance/tile/VortexCore.scala @@ -95,11 +95,11 @@ class VortexBundle(tile: RadianceTile)(implicit p: Parameters) extends CoreBundl val tc_a_bits_write = Output(UInt(tcPortCount.W)) val tc_a_bits_address = Output(UInt((tcPortCount * 32).W)) val tc_a_bits_tag = Output(UInt((tcPortCount * 4).W)) - val tc_a_bits_mask = Output(UInt((tcPortCount * 32).W)) - val tc_a_bits_data = Output(UInt((tcPortCount * 32 * 8).W)) + val tc_a_bits_mask = Output(UInt((tcPortCount * tile.numLsuLanes * 4).W)) + val tc_a_bits_data = Output(UInt((tcPortCount * tile.numLsuLanes * 32).W)) val tc_a_ready = Input(UInt(tcPortCount.W)) val tc_d_valid = Input(UInt(tcPortCount.W)) - val tc_d_bits_data = Input(UInt((tcPortCount * 32 * 8).W)) + val tc_d_bits_data = Input(UInt((tcPortCount * tile.numLsuLanes * 32).W)) val tc_d_bits_tag = Input(UInt((tcPortCount * 4).W)) val tc_d_ready = Output(UInt(tcPortCount.W)) @@ -147,7 +147,8 @@ class Vortex(tile: RadianceTile)(implicit p: Parameters) "CORE_ID" -> tile.radianceParams.coreId, "TENSOR_FP16" -> (if (tile.radianceParams.core.tensorCoreFP16) 1 else 0), "STARTUP_ADDR" -> tile.radianceParams.core.startupAddress, - "NUM_THREADS" -> tile.numLsuLanes + "NUM_THREADS" -> tile.numLsuLanes, + "TC_DATA_WIDTH" -> (tile.numLsuLanes * 32) ) ) with HasBlackBoxResource with HasBlackBoxPath {