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