feat: add fp16 scalar tmem softmax and split tmem

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
+        // 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 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)
+        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)))
         }
 
@@ -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 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 =>
-          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_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(
-            (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 === 1.U) { cReady(tc) := 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 =>
             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
             }
-            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
           }
 
-          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_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(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(_ | _)
         }).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(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(_ | _)
         }).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(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(_ | _)
 
         // port 2: SMEM B, one TL client per tensor core. RadianceSharedMem arbitrates them.