diff --git a/src/main/scala/radiance/tile/RadianceTile.scala b/src/main/scala/radiance/tile/RadianceTile.scala index 0b9d21c..ae7a7e7 100644 --- a/src/main/scala/radiance/tile/RadianceTile.scala +++ b/src/main/scala/radiance/tile/RadianceTile.scala @@ -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.