wu-arch/radiance
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add fp16 scalar tmem softmax and split tmem

Browse Source
This commit is contained in:
Zhongdi LUO
2026-07-10 13:03:12 +00:00
parent 007350fd5a
commit 3e8976490d
2 changed files with 143 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/main/resources/vsrc/vortex

Submodule src/main/resources/vsrc/vortex updated: 9251ba0a24...2bfc6c4bde

225
src/main/scala/radiance/tile/RadianceTile.scala
View File

@@ -888,19 +888,34 @@ class RadianceTileModuleImp(outer: RadianceTile)
tcDData.foreach(_ := 0.U) tcDData.foreach(_ := 0.U)
tcDTag.foreach(_ := 0.U) tcDTag.foreach(_ := 0.U)
// TMEM matrix: four banked 2R1W SRAMs. Tensor A/C reads and scalar // TMEM keeps the ISA-visible address space unified while storing the
// reads can proceed together when bank placement avoids conflicts. // A and C halves in separate 1R1W arrays. This avoids duplicating each
// Each warp owns 2KB: A tile and C tile are 1KB each. The row count // bank for two read ports, and still allows common A-read/C-read pairs
// scales with the physical fragment width (16B for 4 lanes, 32B for 8). // to proceed in parallel because they normally hit different arrays.
val tmemBytesPerWarp = 2048 val tmemBytesPerWarp = 2048
val tmemDepth = outer.numWarps * (tmemBytesPerWarp / outer.tcSmemSize) val tmemFragsPerWarp = tmemBytesPerWarp / outer.tcSmemSize
val tmemFragsPerTile = tmemFragsPerWarp / 2
val tmemLogicalDepth = outer.numWarps * tmemFragsPerWarp
val tmemArrayDepth = outer.numWarps * tmemFragsPerTile
val tmemBanks = 4 val tmemBanks = 4
val tmemBankBits = log2Ceil(tmemBanks) val tmemBankBits = log2Ceil(tmemBanks)
val tmemBankDepth = tmemDepth / tmemBanks val tmemFragAddrBits = log2Ceil(tmemFragsPerWarp)
val tmemTileAddrBits = log2Ceil(tmemFragsPerTile)
val tmemWarpAddrBits = log2Ceil(outer.numWarps)
val tmemPhysAddrBits = log2Ceil(tmemArrayDepth)
val tmemBankDepth = tmemArrayDepth / tmemBanks
require(isPow2(tmemBanks)) require(isPow2(tmemBanks))
require(tmemDepth % tmemBanks == 0) require(isPow2(tmemFragsPerWarp))
val tmem = Seq.fill(tmemBanks) { require(tmemFragsPerWarp == tmemFragsPerTile * 2)
Module(new radiance.memory.TwoReadOneWriteSyncMem( require(tmemLogicalDepth <= (1 << tmemAddrBits))
require(tmemArrayDepth % tmemBanks == 0)
require(tmemPhysAddrBits > tmemBankBits)
val tmemA = Seq.fill(tmemBanks) {
Module(new radiance.memory.TwoPortSyncMem(
tmemBankDepth, UInt((outer.tcSmemSize * 8).W)))
}
val tmemC = Seq.fill(tmemBanks) {
Module(new radiance.memory.TwoPortSyncMem(
tmemBankDepth, UInt((outer.tcSmemSize * 8).W))) tmemBankDepth, UInt((outer.tcSmemSize * 8).W)))
} }
@@ -918,8 +933,17 @@ class RadianceTileModuleImp(outer: RadianceTile)
val tc = UInt(log2Ceil(nTC max 2).W) val tc = UInt(log2Ceil(nTC max 2).W)
} }
def tmemIsC(addr: UInt): Bool = addr(tmemTileAddrBits)
def tmemPhysAddr(addr: UInt): UInt = {
val tileOffset = addr(tmemTileAddrBits - 1, 0)
if (tmemWarpAddrBits == 0) {
tileOffset
} else {
Cat(addr(tmemFragAddrBits + tmemWarpAddrBits - 1, tmemFragAddrBits), tileOffset)
}
}
def bank(addr: UInt): UInt = addr(tmemBankBits - 1, 0) def bank(addr: UInt): UInt = addr(tmemBankBits - 1, 0)
def row(addr: UInt): UInt = addr(tmemAddrBits - 1, tmemBankBits) def row(addr: UInt): UInt = addr(tmemPhysAddrBits - 1, tmemBankBits)
val aReady = Wire(Vec(nTC, Bool())) val aReady = Wire(Vec(nTC, Bool()))
val cReady = Wire(Vec(nTC, Bool())) val cReady = Wire(Vec(nTC, Bool()))
@@ -932,101 +956,136 @@ class RadianceTileModuleImp(outer: RadianceTile)
scReadReady := false.B scReadReady := false.B
scWriteReady := false.B scWriteReady := false.B
val read0Grant = Wire(Vec(tmemBanks, new TmemReadReq)) val aReadGrant = Wire(Vec(tmemBanks, new TmemReadReq))
val read1Grant = Wire(Vec(tmemBanks, new TmemReadReq)) val cReadGrant = Wire(Vec(tmemBanks, new TmemReadReq))
val read0Valid = Wire(Vec(tmemBanks, Bool())) val aReadValid = Wire(Vec(tmemBanks, Bool()))
val read1Valid = Wire(Vec(tmemBanks, Bool())) val cReadValid = Wire(Vec(tmemBanks, Bool()))
val writeGrant = Wire(Vec(tmemBanks, new TmemWriteReq)) val aWriteGrant = Wire(Vec(tmemBanks, new TmemWriteReq))
val writeValid = Wire(Vec(tmemBanks, Bool())) val cWriteGrant = Wire(Vec(tmemBanks, new TmemWriteReq))
read0Grant.foreach(_ := 0.U.asTypeOf(new TmemReadReq)) val aWriteValid = Wire(Vec(tmemBanks, Bool()))
read1Grant.foreach(_ := 0.U.asTypeOf(new TmemReadReq)) val cWriteValid = Wire(Vec(tmemBanks, Bool()))
read0Valid.foreach(_ := false.B) aReadGrant.foreach(_ := 0.U.asTypeOf(new TmemReadReq))
read1Valid.foreach(_ := false.B) cReadGrant.foreach(_ := 0.U.asTypeOf(new TmemReadReq))
writeGrant.foreach(_ := 0.U.asTypeOf(new TmemWriteReq)) aReadValid.foreach(_ := false.B)
writeValid.foreach(_ := false.B) cReadValid.foreach(_ := false.B)
aWriteGrant.foreach(_ := 0.U.asTypeOf(new TmemWriteReq))
cWriteGrant.foreach(_ := 0.U.asTypeOf(new TmemWriteReq))
aWriteValid.foreach(_ := false.B)
cWriteValid.foreach(_ := false.B)
(0 until tmemBanks).foreach { b => (0 until tmemBanks).foreach { b =>
val requests = (0 until nTC).flatMap { tc => val readRequests = (0 until nTC).flatMap { tc =>
val aAddr = slice(core.io.tc_tmem_A_raddr, tmemAddrBits, tc) val aAddr = slice(core.io.tc_tmem_A_raddr, tmemAddrBits, tc)
val cAddr = slice(core.io.tc_tmem_C_raddr, tmemAddrBits, tc) val cAddr = slice(core.io.tc_tmem_C_raddr, tmemAddrBits, tc)
Seq( Seq(
(core.io.tc_tmem_A_ren(tc).asBool && bank(aAddr) === b.U, aAddr, 0.U(2.W), tc.U), (core.io.tc_tmem_A_ren(tc).asBool, aAddr, 0.U(2.W), tc.U),
(core.io.tc_tmem_C_ren(tc).asBool && bank(cAddr) === b.U, cAddr, 1.U(2.W), tc.U) (core.io.tc_tmem_C_ren(tc).asBool, cAddr, 1.U(2.W), tc.U)
) )
} ++ Seq( } ++ Seq(
(core.io.sc_tmem_ren.asBool && bank(core.io.sc_tmem_raddr) === b.U, (core.io.sc_tmem_ren.asBool, core.io.sc_tmem_raddr, 2.U(2.W), 0.U)
core.io.sc_tmem_raddr, 2.U(2.W), 0.U)
) )
var used0 = false.B var aReadUsed = false.B
var used1 = false.B var cReadUsed = false.B
requests.foreach { case (valid, addr, src, tc) => readRequests.foreach { case (valid, addr, src, tc) =>
val grant0 = valid && !used0 val physAddr = tmemPhysAddr(addr)
val grant1 = valid && used0 && !used1 val isC = tmemIsC(addr)
when(grant0) { val aGrant = valid && !isC && bank(physAddr) === b.U && !aReadUsed
read0Grant(b).addr := addr val cGrant = valid && isC && bank(physAddr) === b.U && !cReadUsed
read0Grant(b).src := src when(aGrant) {
read0Grant(b).tc := tc aReadGrant(b).addr := physAddr
aReadGrant(b).src := src
aReadGrant(b).tc := tc
} }
when(grant1) { when(cGrant) {
read1Grant(b).addr := addr cReadGrant(b).addr := physAddr
read1Grant(b).src := src cReadGrant(b).src := src
read1Grant(b).tc := tc cReadGrant(b).tc := tc
} }
used0 = used0 || grant0 aReadUsed = aReadUsed || aGrant
used1 = used1 || grant1 cReadUsed = cReadUsed || cGrant
when(grant0 || grant1) { when(aGrant || cGrant) {
when(src === 0.U) { aReady(tc) := true.B } when(src === 0.U) { aReady(tc) := true.B }
when(src === 1.U) { cReady(tc) := true.B } when(src === 1.U) { cReady(tc) := true.B }
when(src === 2.U) { scReadReady := true.B } when(src === 2.U) { scReadReady := true.B }
} }
} }
read0Valid(b) := used0 aReadValid(b) := aReadUsed
read1Valid(b) := used1 cReadValid(b) := cReadUsed
var writeUsed = false.B var aWriteUsed = false.B
var cWriteUsed = false.B
(0 until nTC).foreach { tc => (0 until nTC).foreach { tc =>
val addr = slice(core.io.tc_tmem_C_waddr, tmemAddrBits, tc) val addr = slice(core.io.tc_tmem_C_waddr, tmemAddrBits, tc)
val valid = core.io.tc_tmem_C_wen(tc).asBool && bank(addr) === b.U val physAddr = tmemPhysAddr(addr)
val grant = valid && !writeUsed val isC = tmemIsC(addr)
when(grant) { val valid = core.io.tc_tmem_C_wen(tc).asBool && bank(physAddr) === b.U
writeValid(b) := true.B val aGrant = valid && !isC && !aWriteUsed
writeGrant(b).addr := addr val cGrant = valid && isC && !cWriteUsed
writeGrant(b).data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc) when(aGrant) {
writeGrant(b).mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc) aWriteValid(b) := true.B
writeGrant(b).src := 0.U aWriteGrant(b).addr := physAddr
writeGrant(b).tc := tc.U aWriteGrant(b).data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc)
aWriteGrant(b).mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc)
aWriteGrant(b).src := 0.U
aWriteGrant(b).tc := tc.U
wReady(tc) := true.B wReady(tc) := true.B
} }
writeUsed = writeUsed || grant when(cGrant) {
cWriteValid(b) := true.B
cWriteGrant(b).addr := physAddr
cWriteGrant(b).data := slice(core.io.tc_tmem_C_wdata, tmemDataBits, tc)
cWriteGrant(b).mask := slice(core.io.tc_tmem_C_mask, tmemMaskBits, tc)
cWriteGrant(b).src := 0.U
cWriteGrant(b).tc := tc.U
wReady(tc) := true.B
}
aWriteUsed = aWriteUsed || aGrant
cWriteUsed = cWriteUsed || cGrant
} }
val scWValid = core.io.sc_tmem_wen.asBool && bank(core.io.sc_tmem_waddr) === b.U val scWPhysAddr = tmemPhysAddr(core.io.sc_tmem_waddr)
val scWGrant = scWValid && !writeUsed val scWIsC = tmemIsC(core.io.sc_tmem_waddr)
when(scWGrant) { val scWValid = core.io.sc_tmem_wen.asBool && bank(scWPhysAddr) === b.U
writeValid(b) := true.B val scWAGrant = scWValid && !scWIsC && !aWriteUsed
writeGrant(b).addr := core.io.sc_tmem_waddr val scWCGrant = scWValid && scWIsC && !cWriteUsed
writeGrant(b).data := core.io.sc_tmem_wdata when(scWAGrant) {
writeGrant(b).mask := core.io.sc_tmem_mask aWriteValid(b) := true.B
writeGrant(b).src := 1.U aWriteGrant(b).addr := scWPhysAddr
writeGrant(b).tc := 0.U aWriteGrant(b).data := core.io.sc_tmem_wdata
aWriteGrant(b).mask := core.io.sc_tmem_mask
aWriteGrant(b).src := 1.U
aWriteGrant(b).tc := 0.U
scWriteReady := true.B
}
when(scWCGrant) {
cWriteValid(b) := true.B
cWriteGrant(b).addr := scWPhysAddr
cWriteGrant(b).data := core.io.sc_tmem_wdata
cWriteGrant(b).mask := core.io.sc_tmem_mask
cWriteGrant(b).src := 1.U
cWriteGrant(b).tc := 0.U
scWriteReady := true.B scWriteReady := true.B
} }
tmem(b).io.ren0 := read0Valid(b) tmemA(b).io.ren := aReadValid(b)
tmem(b).io.raddr0 := row(read0Grant(b).addr) tmemA(b).io.raddr := row(aReadGrant(b).addr)
tmem(b).io.ren1 := read1Valid(b) tmemA(b).io.wen := aWriteValid(b)
tmem(b).io.raddr1 := row(read1Grant(b).addr) tmemA(b).io.waddr := row(aWriteGrant(b).addr)
tmem(b).io.wen := writeValid(b) tmemA(b).io.wdata := aWriteGrant(b).data
tmem(b).io.waddr := row(writeGrant(b).addr) tmemA(b).io.mask := aWriteGrant(b).mask
tmem(b).io.wdata := writeGrant(b).data tmemC(b).io.ren := cReadValid(b)
tmem(b).io.mask := writeGrant(b).mask tmemC(b).io.raddr := row(cReadGrant(b).addr)
tmemC(b).io.wen := cWriteValid(b)
tmemC(b).io.waddr := row(cWriteGrant(b).addr)
tmemC(b).io.wdata := cWriteGrant(b).data
tmemC(b).io.mask := cWriteGrant(b).mask
} }
val read0GrantReg = RegNext(read0Grant) val aReadGrantReg = RegNext(aReadGrant)
val read1GrantReg = RegNext(read1Grant) val cReadGrantReg = RegNext(cReadGrant)
val read0ValidReg = RegNext(read0Valid) val aReadValidReg = RegNext(aReadValid)
val read1ValidReg = RegNext(read1Valid) val cReadValidReg = RegNext(cReadValid)
core.io.tc_tmem_A_rready := aReady.asUInt core.io.tc_tmem_A_rready := aReady.asUInt
core.io.tc_tmem_C_rready := cReady.asUInt core.io.tc_tmem_C_rready := cReady.asUInt
core.io.tc_tmem_C_wready := wReady.asUInt core.io.tc_tmem_C_wready := wReady.asUInt
@@ -1034,19 +1093,19 @@ class RadianceTileModuleImp(outer: RadianceTile)
core.io.sc_tmem_wready := scWriteReady.asUInt core.io.sc_tmem_wready := scWriteReady.asUInt
core.io.tc_tmem_A_rdata := VecInit((0 until nTC).map { tc => core.io.tc_tmem_A_rdata := VecInit((0 until nTC).map { tc =>
VecInit((0 until tmemBanks).map { b => VecInit((0 until tmemBanks).map { b =>
Mux(read0ValidReg(b) && read0GrantReg(b).src === 0.U && read0GrantReg(b).tc === tc.U, tmem(b).io.rdata0, Mux(aReadValidReg(b) && aReadGrantReg(b).src === 0.U && aReadGrantReg(b).tc === tc.U, tmemA(b).io.rdata,
Mux(read1ValidReg(b) && read1GrantReg(b).src === 0.U && read1GrantReg(b).tc === tc.U, tmem(b).io.rdata1, 0.U(tmemDataBits.W))) Mux(cReadValidReg(b) && cReadGrantReg(b).src === 0.U && cReadGrantReg(b).tc === tc.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
}).reduce(_ | _) }).reduce(_ | _)
}).asUInt }).asUInt
core.io.tc_tmem_C_rdata := VecInit((0 until nTC).map { tc => core.io.tc_tmem_C_rdata := VecInit((0 until nTC).map { tc =>
VecInit((0 until tmemBanks).map { b => VecInit((0 until tmemBanks).map { b =>
Mux(read0ValidReg(b) && read0GrantReg(b).src === 1.U && read0GrantReg(b).tc === tc.U, tmem(b).io.rdata0, Mux(aReadValidReg(b) && aReadGrantReg(b).src === 1.U && aReadGrantReg(b).tc === tc.U, tmemA(b).io.rdata,
Mux(read1ValidReg(b) && read1GrantReg(b).src === 1.U && read1GrantReg(b).tc === tc.U, tmem(b).io.rdata1, 0.U(tmemDataBits.W))) Mux(cReadValidReg(b) && cReadGrantReg(b).src === 1.U && cReadGrantReg(b).tc === tc.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
}).reduce(_ | _) }).reduce(_ | _)
}).asUInt }).asUInt
core.io.sc_tmem_rdata := VecInit((0 until tmemBanks).map { b => core.io.sc_tmem_rdata := VecInit((0 until tmemBanks).map { b =>
Mux(read0ValidReg(b) && read0GrantReg(b).src === 2.U, tmem(b).io.rdata0, Mux(aReadValidReg(b) && aReadGrantReg(b).src === 2.U, tmemA(b).io.rdata,
Mux(read1ValidReg(b) && read1GrantReg(b).src === 2.U, tmem(b).io.rdata1, 0.U(tmemDataBits.W))) Mux(cReadValidReg(b) && cReadGrantReg(b).src === 2.U, tmemC(b).io.rdata, 0.U(tmemDataBits.W)))
}).reduce(_ | _) }).reduce(_ | _)
// port 2: SMEM B, one TL client per tensor core. RadianceSharedMem arbitrates them. // port 2: SMEM B, one TL client per tensor core. RadianceSharedMem arbitrates them.