diff --git a/hw/rtl/tcu/VX_tcu_agu.sv b/hw/rtl/tcu/VX_tcu_agu.sv index 77df74f53e..087b825c37 100644 --- a/hw/rtl/tcu/VX_tcu_agu.sv +++ b/hw/rtl/tcu/VX_tcu_agu.sv @@ -15,24 +15,8 @@ `ifdef TCU_META_ENABLE -// Warp-level AGU for TCU_LD instructions. -// -// One TCU_LD = one memory fetch that lands in one VX_tcu_meta slot -// (selected by op_args.tcu.fmt_d). Software emits N TCU_LD ops to -// fill N namespace-local slots. -// -// State machine: IDLE → ISSUE → WAIT_RSP → COMMIT. -// IDLE : monitor per-block execute_if for INST_TCU_LD. -// ISSUE : drive client_if.req with rs1 (base address). Picks the -// lowest-indexed block with a pending TCU_LD (single AGU, -// shared across blocks). -// WAIT_RSP : wait for the matching client_if.rsp. -// COMMIT : drive meta_wr_* (broadcast to all blocks' tcu_meta -// instances) and assert the originating block's result_if. -// op_args.tcu.fmt_d → meta slot index; warp id from the -// execute header. -// -// Single-request AGU; multi-request stride patterns are not implemented. +// Warp-level AGU for TCU_LD instructions. Each TCU block can hold one +// transaction, while the LSU request and metadata write ports are shared. module VX_tcu_agu import VX_gpu_pkg::*, VX_tcu_pkg::*; #( parameter `STRING INSTANCE_ID = "", @@ -69,123 +53,115 @@ module VX_tcu_agu import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `UNUSED_SPARAM (INSTANCE_ID) localparam BLOCK_IDX_BITS = (BLOCK_SIZE > 1) ? $clog2(BLOCK_SIZE) : 1; - // ----------------------------------------------------------------------- - // State - // ----------------------------------------------------------------------- - typedef enum logic [1:0] { - S_IDLE = 2'd0, - S_ISSUE = 2'd1, - S_WAIT = 2'd2, - S_COMMIT= 2'd3 - } state_e; + logic [BLOCK_SIZE-1:0] busy_r; + logic [BLOCK_SIZE-1:0] issued_r; + logic [BLOCK_SIZE-1:0] complete_r; + tcu_header_t header_r [BLOCK_SIZE]; + logic [4:0] slot_r [BLOCK_SIZE]; + logic [`VX_CFG_XLEN-1:0] addr_r [BLOCK_SIZE]; + logic [NUM_LANES-1:0][(LSU_WORD_SIZE*8)-1:0] + rsp_data_r [BLOCK_SIZE]; + logic [NUM_LANES-1:0] rsp_lane_done_r [BLOCK_SIZE]; - state_e state_r; - logic [BLOCK_IDX_BITS-1:0] owner_block_r; // which block's TCU_LD is in flight - // Latch only the fields we need (avoids unused-bit warnings on the - // full tcu_execute_t struct). - tcu_header_t owner_header_r; - logic [4:0] owner_slot_r; // op_args.tcu.fmt_d - logic [4:0] owner_fmt_r; // op_args.tcu.fmt_s - logic [`VX_CFG_XLEN-1:0] owner_addr_r; // rs1_data[0] - logic [NUM_LANES-1:0][(LSU_WORD_SIZE*8)-1:0] rsp_data_r; - // Track which lanes have received valid response data. The - // lsu_scheduler may stream responses across multiple cycles when - // the mem_scheduler coalesces same-line accesses; we transition to - // S_COMMIT only when all requested lanes have responded. - logic [NUM_LANES-1:0] rsp_lane_done_r; + wire [BLOCK_SIZE-1:0] req_pending = busy_r & ~issued_r; + wire [BLOCK_IDX_BITS-1:0] req_block; + wire [BLOCK_SIZE-1:0] req_grant; + wire req_valid; - // ----------------------------------------------------------------------- - // Lowest-block priority pick (single AGU shared across blocks). - // ----------------------------------------------------------------------- - wire [BLOCK_SIZE-1:0] block_valid_v = per_block_ld_valid; + VX_generic_arbiter #( + .NUM_REQS (BLOCK_SIZE), + .TYPE ("R") + ) req_arb ( + .clk (clk), + .reset (reset), + .requests (req_pending), + .grant_index (req_block), + .grant_onehot (req_grant), + .grant_valid (req_valid), + .grant_ready (client_if.req_ready) + ); - logic [BLOCK_IDX_BITS-1:0] first_valid_idx; - logic any_valid; - always_comb begin - first_valid_idx = '0; - any_valid = 1'b0; - for (int bi = BLOCK_SIZE-1; bi >= 0; bi--) begin - if (block_valid_v[bi]) begin - first_valid_idx = BLOCK_IDX_BITS'(bi); - any_valid = 1'b1; - end - end - end + wire [BLOCK_SIZE-1:0] commit_pending = complete_r; + wire [BLOCK_IDX_BITS-1:0] commit_block; + wire [BLOCK_SIZE-1:0] commit_grant; + wire commit_valid; + wire commit_ready = result_ready[commit_block]; + + VX_generic_arbiter #( + .NUM_REQS (BLOCK_SIZE), + .TYPE ("R") + ) commit_arb ( + .clk (clk), + .reset (reset), + .requests (commit_pending), + .grant_index (commit_block), + .grant_onehot (commit_grant), + .grant_valid (commit_valid), + .grant_ready (commit_ready) + ); + + wire [BLOCK_IDX_BITS-1:0] rsp_block = + client_if.rsp_data.tag[BLOCK_IDX_BITS-1:0]; + wire rsp_fire = client_if.rsp_valid && client_if.rsp_ready; + wire req_fire = req_valid && client_if.req_ready; + wire commit_fire = commit_valid && commit_ready; + `UNUSED_VAR (req_grant) + `UNUSED_VAR (commit_grant) - // ----------------------------------------------------------------------- - // FSM - // ----------------------------------------------------------------------- always @(posedge clk) begin if (reset) begin - state_r <= S_IDLE; - owner_block_r <= '0; - owner_header_r <= '0; - owner_slot_r <= '0; - owner_fmt_r <= '0; - owner_addr_r <= '0; - rsp_data_r <= '0; - rsp_lane_done_r <= '0; + busy_r <= '0; + issued_r <= '0; + complete_r <= '0; + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + header_r[bi] <= '0; + slot_r[bi] <= '0; + addr_r[bi] <= '0; + rsp_data_r[bi] <= '0; + rsp_lane_done_r[bi] <= '0; + end end else begin - case (state_r) - S_IDLE: begin - if (any_valid) begin - owner_block_r <= first_valid_idx; - owner_header_r <= per_block_ld_data[first_valid_idx].header; - owner_slot_r <= per_block_ld_data[first_valid_idx].op_args.tcu.fmt_d; - owner_fmt_r <= per_block_ld_data[first_valid_idx].op_args.tcu.fmt_s; - owner_addr_r <= per_block_ld_data[first_valid_idx].rs1_data[0]; - rsp_lane_done_r <= '0; - state_r <= S_ISSUE; - end + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + if (per_block_ld_valid[bi] && per_block_ld_ready[bi]) begin + busy_r[bi] <= 1'b1; + issued_r[bi] <= 1'b0; + complete_r[bi] <= 1'b0; + header_r[bi] <= per_block_ld_data[bi].header; + slot_r[bi] <= per_block_ld_data[bi].op_args.tcu.fmt_d; + addr_r[bi] <= per_block_ld_data[bi].rs1_data[0]; + rsp_data_r[bi] <= '0; + rsp_lane_done_r[bi] <= '0; end - S_ISSUE: begin - if (client_if.req_ready) begin - state_r <= S_WAIT; - end - end - S_WAIT: begin - if (client_if.rsp_valid) begin - // Merge this packet's lane data into rsp_data_r. - for (int t = 0; t < NUM_LANES; ++t) begin - if (client_if.rsp_data.mask[t]) begin - rsp_data_r[t] <= client_if.rsp_data.data[t]; - rsp_lane_done_r[t] <= 1'b1; - end - end - // Only commit when EOP arrives AND every requested - // lane has its data. With multi-cycle responses - // (lsu_scheduler may coalesce same-line lanes and - // stream the demux), keep waiting until all lanes - // accounted for. - if (client_if.rsp_data.eop - && ((rsp_lane_done_r | client_if.rsp_data.mask) == {NUM_LANES{1'b1}})) begin - state_r <= S_COMMIT; - end + end + if (req_fire) begin + issued_r[req_block] <= 1'b1; + end + if (rsp_fire) begin + for (int t = 0; t < NUM_LANES; ++t) begin + if (client_if.rsp_data.mask[t]) begin + rsp_data_r[rsp_block][t] <= client_if.rsp_data.data[t]; + rsp_lane_done_r[rsp_block][t] <= 1'b1; end end - S_COMMIT: begin - if (result_ready[owner_block_r]) begin - state_r <= S_IDLE; - end + if (client_if.rsp_data.eop + && ((rsp_lane_done_r[rsp_block] | client_if.rsp_data.mask) + == {NUM_LANES{1'b1}})) begin + complete_r[rsp_block] <= 1'b1; end - default: state_r <= S_IDLE; - endcase + end + if (commit_fire) begin + busy_r[commit_block] <= 1'b0; + issued_r[commit_block] <= 1'b0; + complete_r[commit_block] <= 1'b0; + rsp_lane_done_r[commit_block] <= '0; + end end end - // Unused sub-fields of rsp_data (we consume valid/eop/data only). - `UNUSED_VAR (client_if.rsp_data.tag) `UNUSED_VAR (client_if.rsp_data.sop) - // ----------------------------------------------------------------------- - // execute_if.ready: only assert for the picked block in S_IDLE when we - // latch its data (transition to S_ISSUE). - // ----------------------------------------------------------------------- for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_ld_ready - assign per_block_ld_ready[bi] = - (state_r == S_IDLE) - && per_block_ld_valid[bi] - && (first_valid_idx == BLOCK_IDX_BITS'(bi)); + assign per_block_ld_ready[bi] = ~busy_r[bi]; end // ----------------------------------------------------------------------- @@ -198,14 +174,13 @@ module VX_tcu_agu import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // VX_tcu_meta at (bank=T%PWD, col_in_group=T/PWD), so lane T must // load `meta_base[(T % PWD) * stride + (T / PWD)]`. // ----------------------------------------------------------------------- - wire [`VX_CFG_XLEN-1:0] base_addr = owner_addr_r; + wire [`VX_CFG_XLEN-1:0] base_addr = addr_r[req_block]; `UNUSED_VAR (base_addr[`CLOG2(LSU_WORD_SIZE)-1:0]) wire [LSU_ADDR_WIDTH-1:0] base_word_addr = LSU_ADDR_WIDTH'(base_addr[`VX_CFG_XLEN-1:`CLOG2(LSU_WORD_SIZE)]); localparam PWD = TCU_META_PER_WARP_DEPTH; - `UNUSED_VAR (owner_fmt_r) // Per-thread metadata layout: lane T loads h_meta[slot*NT + T]. The // host packs h_meta so the 32-bit word at offset (slot*NT + T) holds // the metadata destined for SRAM cell (bank = T%PWD, col = T/PWD). @@ -273,30 +248,31 @@ module VX_tcu_agu import VX_gpu_pkg::*, VX_tcu_pkg::*; #( req_w.data[i] = '0; req_w.attr[i] = a; end - // Tag pattern: tag width must match LSU_CLIENT_TAG_WIDTH but the - // AGU doesn't read tags back as anything meaningful — just sized. req_w.tag = '0; + req_w.tag[BLOCK_IDX_BITS-1:0] = req_block; end `UNUSED_PARAM (PWD) - assign client_if.req_valid = (state_r == S_ISSUE); + `STATIC_ASSERT (LSU_CLIENT_TAG_WIDTH >= BLOCK_IDX_BITS, ("LSU client tag cannot encode TCU block")) + assign client_if.req_valid = req_valid; assign client_if.req_data = req_w; - assign client_if.rsp_ready = (state_r == S_WAIT); + assign client_if.rsp_ready = busy_r[rsp_block] && issued_r[rsp_block] + && !complete_r[rsp_block]; // ----------------------------------------------------------------------- // VX_tcu_meta write (broadcast to all blocks' meta instances). // ----------------------------------------------------------------------- - assign meta_wr_en = (state_r == S_COMMIT); - assign meta_wr_wid = owner_header_r.wid; - assign meta_wr_idx = owner_slot_r; + assign meta_wr_en = commit_valid; + assign meta_wr_wid = header_r[commit_block].wid; + assign meta_wr_idx = slot_r[commit_block]; // Repack response data into TCU_BLOCK_CAP × XLEN. The downstream // meta SRAM consumes only the low 32 bits, so for XLEN=64 we mux the // half of the LSU response that holds this lane's logical 32-bit word. for (genvar i = 0; i < TCU_BLOCK_CAP; ++i) begin : g_meta_data if (i < NUM_LANES) begin : g_in if (LG_W32_RATIO == 0) begin : g_w32_passthru - assign meta_wr_data[i] = `VX_CFG_XLEN'(rsp_data_r[i]); + assign meta_wr_data[i] = `VX_CFG_XLEN'(rsp_data_r[commit_block][i]); end else begin : g_w32_pick - wire [LSU_W32_RATIO-1:0][31:0] rsp_w32_v = rsp_data_r[i]; + wire [LSU_W32_RATIO-1:0][31:0] rsp_w32_v = rsp_data_r[commit_block][i]; assign meta_wr_data[i] = `VX_CFG_XLEN'(rsp_w32_v[half_sel_w[i]]); end end else begin : g_pad @@ -314,12 +290,12 @@ module VX_tcu_agu import VX_gpu_pkg::*, VX_tcu_pkg::*; #( tcu_result_t commit_data_w; always_comb begin commit_data_w = '0; - commit_data_w.header = owner_header_r; + commit_data_w.header = header_r[commit_block]; commit_data_w.data = '0; end for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_result - assign result_valid[bi] = (state_r == S_COMMIT) && (owner_block_r == BLOCK_IDX_BITS'(bi)); + assign result_valid[bi] = commit_valid && (commit_block == BLOCK_IDX_BITS'(bi)); assign result_data[bi] = commit_data_w; end diff --git a/sim/simx/tcu/tcu_unit.cpp b/sim/simx/tcu/tcu_unit.cpp index 12e6a9312c..bd055b145f 100644 --- a/sim/simx/tcu/tcu_unit.cpp +++ b/sim/simx/tcu/tcu_unit.cpp @@ -390,16 +390,14 @@ class TcuUnit::Impl { cta_owner_b_ = -1; cur_block_ = 0; #ifdef TCU_META_ENABLE - agu_ = agu_state_t{}; + agu_.fill(agu_state_t{}); + agu_issue_rr_ = 0; #endif } #ifdef TCU_META_ENABLE - // Metadata AGU: one transaction at a time. NT metadata words at base+T*4 - // are fetched through the LSU client port in ceil(NT/LSU_LANES) beats and - // deposited into the metadata SRAM on completion. The TCU_LD trace retires - // only after the deposit, so consuming MMAs behind it in the block queue - // observe complete metadata. + // Each TCU block tracks one metadata transaction. Requests share the LSU + // client port, while tags allow their memory latency to overlap. static constexpr uint32_t kAguBeats = (VX_CFG_NUM_THREADS + VX_CFG_NUM_LSU_LANES - 1) / VX_CFG_NUM_LSU_LANES; @@ -414,84 +412,92 @@ class TcuUnit::Impl { std::array addrs{}; std::array words{}; }; - agu_state_t agu_; + std::array agu_; + uint32_t agu_issue_rr_ = 0; - void agu_start(uint32_t wid, uint32_t selector, uint64_t base_addr, - instr_trace_t* trace) { + void agu_start(uint32_t block, uint32_t wid, uint32_t selector, + uint64_t base_addr, instr_trace_t* trace) { assert((base_addr & 3) == 0 && "metadata base must be word-aligned"); - agu_ = agu_state_t{}; - agu_.busy = true; - agu_.trace = trace; - agu_.wid = wid; - agu_.selector = selector; + auto& agu = agu_.at(block); + agu = agu_state_t{}; + agu.busy = true; + agu.trace = trace; + agu.wid = wid; + agu.selector = selector; for (uint32_t t = 0; t < VX_CFG_NUM_THREADS; ++t) { - agu_.addrs.at(t) = base_addr + uint64_t(t) * 4; + agu.addrs.at(t) = base_addr + uint64_t(t) * 4; } } void agu_step() { - // Issue one request beat per cycle while the port accepts. - if (agu_.busy && agu_.next_beat < kAguBeats && !simobject_->agu_req_out.full()) { - LsuReq req(VX_CFG_NUM_LSU_LANES); - req.op = MemOp::LD; - req.tag = agu_.next_beat; // client tag = beat index, restored on the response - req.wid = agu_.wid; - req.cid = agu_.trace->cid; - req.uuid = agu_.trace->uuid; - for (uint32_t i = 0; i < VX_CFG_NUM_LSU_LANES; ++i) { - uint32_t idx = agu_.next_beat * VX_CFG_NUM_LSU_LANES + i; - if (idx >= VX_CFG_NUM_THREADS) { - break; + if (!simobject_->agu_req_out.full()) { + for (uint32_t offset = 0; offset < VX_CFG_NUM_TCU_BLOCKS; ++offset) { + uint32_t block = (agu_issue_rr_ + offset) % VX_CFG_NUM_TCU_BLOCKS; + auto& agu = agu_.at(block); + if (!agu.busy || agu.next_beat >= kAguBeats) { + continue; } - req.mask.set(i); - req.addrs.at(i) = agu_.addrs.at(idx); - req.tids.at(i) = idx % VX_CFG_NUM_THREADS; + LsuReq req(VX_CFG_NUM_LSU_LANES); + req.op = MemOp::LD; + req.tag = block * kAguBeats + agu.next_beat; + req.wid = agu.wid; + req.cid = agu.trace->cid; + req.uuid = agu.trace->uuid; + for (uint32_t i = 0; i < VX_CFG_NUM_LSU_LANES; ++i) { + uint32_t idx = agu.next_beat * VX_CFG_NUM_LSU_LANES + i; + if (idx >= VX_CFG_NUM_THREADS) { + break; + } + req.mask.set(i); + req.addrs.at(i) = agu.addrs.at(idx); + req.tids.at(i) = idx % VX_CFG_NUM_THREADS; + } + simobject_->agu_req_out.send(req); + ++agu.next_beat; + agu_issue_rr_ = (block + 1) % VX_CFG_NUM_TCU_BLOCKS; + break; } - simobject_->agu_req_out.send(req); - ++agu_.next_beat; } - // Accumulate one response fragment per cycle. This drain is - // unconditional (not gated on busy): the LSU response demux stalls its - // whole block-0 channel while agu_rsp_in backs up, so fragments must - // always be consumed. + if (!simobject_->agu_rsp_in.empty()) { auto& rsp = simobject_->agu_rsp_in.peek(); - uint32_t beat = (uint32_t)rsp.tag; + uint32_t block = (uint32_t)rsp.tag / kAguBeats; + uint32_t beat = (uint32_t)rsp.tag % kAguBeats; + auto& agu = agu_.at(block); for (uint32_t lane = 0; lane < rsp.mask.size(); ++lane) { if (!rsp.mask.test(lane)) { continue; } uint32_t idx = beat * VX_CFG_NUM_LSU_LANES + lane; assert(idx < VX_CFG_NUM_THREADS && rsp.data.at(lane)); - uint32_t off = agu_.addrs.at(idx) & (VX_CFG_MEM_BLOCK_SIZE - 1); + uint32_t off = agu.addrs.at(idx) & (VX_CFG_MEM_BLOCK_SIZE - 1); uint32_t word = 0; std::memcpy(&word, rsp.data.at(lane)->data() + off, 4); - agu_.words.at(idx) = word; - ++agu_.rsp_count; + agu.words.at(idx) = word; + ++agu.rsp_count; } simobject_->agu_rsp_in.pop(); - // All words arrived — deposit into the metadata SRAM. - if (agu_.busy && !agu_.complete && agu_.rsp_count == VX_CFG_NUM_THREADS) { - this->agu_deposit(); - agu_.complete = true; + if (agu.busy && !agu.complete && agu.rsp_count == VX_CFG_NUM_THREADS) { + this->agu_deposit(agu); + agu.complete = true; } } } // Deposit the fetched words into the selected metadata SRAM. - void agu_deposit() { - uint32_t wid = agu_.wid; + void agu_deposit(const agu_state_t& agu) { + uint32_t wid = agu.wid; #ifdef VX_CFG_TCU_MX_ENABLE - if (agu_.selector & 0x10) { - auto& dst = (agu_.selector & 1) ? mx_meta_b_.at(wid) : mx_meta_a_.at(wid); + if (agu.selector & 0x10) { + auto& dst = (agu.selector & 1) ? mx_meta_b_.at(wid) : mx_meta_a_.at(wid); for (uint32_t lane = 0; lane < VX_CFG_NUM_THREADS; ++lane) { - dst.at(lane) = agu_.words.at(lane); + dst.at(lane) = agu.words.at(lane); } return; } #endif #ifdef VX_CFG_TCU_SPARSE_ENABLE - uint32_t slot_idx = agu_.selector & 0xf; + uint32_t slot_idx = agu.selector & 0xf; // Map lane T to its (bank, col) metadata cell using the host pack layout: // flat_store = slot*cols_per_load + T/BPS // col = flat_store / stores_per_col @@ -511,13 +517,13 @@ class TcuUnit::Impl { if (bank >= PWD || col >= kMaxMetaCols) { continue; } - uint32_t word = agu_.words.at(T); + uint32_t word = agu.words.at(T); sparse_meta_.at(wid).at(bank * kMaxMetaCols + col) = word; // Trace: META_SRAM write (CSV: wid,bank,col,addr,value). if (const char* p = std::getenv("VORTEX_TCU_TRACE")) { if (p[0] == '1') { fprintf(stderr, "META_TRC,%u,%u,%u,0x%lx,0x%08x\n", - wid, bank, col, (unsigned long)agu_.addrs.at(T), word); + wid, bank, col, (unsigned long)agu.addrs.at(T), word); } } } @@ -707,14 +713,12 @@ class TcuUnit::Impl { #endif #ifdef TCU_META_ENABLE case TcuType::TCU_LD: { - // One AGU transaction at a time — if it is busy with another - // trace, hold this block until it frees. - if (agu_.busy) { + if (agu_.at(b).busy) { continue; } // rs1 is a full-width address (.u64); use u64 to avoid truncation on XLEN=64. uint64_t base_addr = rs1_data.empty() ? 0 : rs1_data.at(0).u64; - this->agu_start(wid, tpuArgs.fmt_d, base_addr, trace); + this->agu_start(b, wid, tpuArgs.fmt_d, base_addr, trace); } break; #endif default: @@ -743,11 +747,11 @@ class TcuUnit::Impl { } #ifdef TCU_META_ENABLE // TCU_LD retires only after its words are deposited into the - // metadata SRAM. Single-transaction AGU: a block only reaches this - // point for the trace that owns the AGU. + // metadata SRAM. if (tcu_type == TcuType::TCU_LD) { - assert(agu_.busy && agu_.trace == trace); - if (!agu_.complete) { + auto& agu = agu_.at(b); + assert(agu.busy && agu.trace == trace); + if (!agu.complete) { continue; // metadata still in flight — hold the block } } @@ -769,7 +773,7 @@ class TcuUnit::Impl { #ifdef TCU_META_ENABLE // TCU_LD retired: free the AGU for the next metadata load. if (tcu_type == TcuType::TCU_LD) { - agu_ = agu_state_t{}; + agu_.at(b) = agu_state_t{}; } #endif DT(3, simobject_->name() << " execute: op=" << tcu_type << ", " << *trace); diff --git a/tests/regression/sgemm_tcu_wg_sp/kernel.cpp b/tests/regression/sgemm_tcu_wg_sp/kernel.cpp index d2f40ded89..8c60045c2e 100644 --- a/tests/regression/sgemm_tcu_wg_sp/kernel.cpp +++ b/tests/regression/sgemm_tcu_wg_sp/kernel.cpp @@ -7,8 +7,7 @@ namespace vt = vortex::tensor; using ctx = vt::wgmma_context; -// Per-warp smem layout: [A_compressed][metadata] (bank-aligned section) -// Then shared B after all warp sections +// Per-warp smem layout: [A_compressed][reserved metadata row], then shared B. static constexpr uint32_t smem_a_elems = ctx::xtileM * (ctx::tileK / 2); static constexpr uint32_t smem_a_bytes = smem_a_elems * sizeof(ctx::input_t); static constexpr uint32_t smem_b_elems = ctx::tileK * ctx::xtileN; @@ -19,18 +18,20 @@ static constexpr uint32_t smem_b_elems = ctx::tileK * ctx::xtileN; // land on a fresh LMEM bank-row — otherwise abuf's single-row fetch on XLEN=64 NT=16 // NRC=32 sparse pulls the previous warp's meta into the next warp's A stripe. static constexpr uint32_t smem_bank_bytes = VX_CFG_NUM_THREADS * (VX_CFG_XLEN / 8); -static constexpr uint32_t per_warp_section = ((smem_a_bytes + ctx::wg_meta_total_bytes + smem_bank_bytes - 1) / smem_bank_bytes) * smem_bank_bytes; +static constexpr uint32_t per_warp_section = + ((smem_a_bytes + ctx::wg_meta_total_bytes + smem_bank_bytes - 1) + / smem_bank_bytes) + * smem_bank_bytes; __kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { auto pA = reinterpret_cast(arg->A_addr); // compressed A auto pB = reinterpret_cast(arg->B_addr); // dense B - auto pC = reinterpret_cast(arg->C_addr); - auto pMetaSp = reinterpret_cast(arg->meta_sp_addr); uint32_t N = arg->N; uint32_t K = arg->K; uint32_t tid = threadIdx.x; + uint32_t lane = tid % VX_CFG_NUM_THREADS; uint32_t num_threads = blockDim.x; uint32_t warp_rank = tid / VX_CFG_NUM_THREADS; uint32_t num_warps = num_threads / VX_CFG_NUM_THREADS; @@ -41,54 +42,51 @@ __kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { uint32_t tile_col = blockIdx.x * ctx::xtileN; auto smem_base = reinterpret_cast(__local_mem()); - uint32_t smem_b_off = ((num_warps * per_warp_section + smem_bank_bytes - 1) / smem_bank_bytes) * smem_bank_bytes; + uint32_t smem_b_off = + ((num_warps * per_warp_section + smem_bank_bytes - 1) + / smem_bank_bytes) * smem_bank_bytes; auto B_smem = reinterpret_cast(smem_base + smem_b_off); // Initialize accumulator to zero ctx::fragment_acc fragC; ctx::fill_fragment(fragC, 0); - // Strides in global memory - uint32_t a_sp_stride = K / 2; // compressed A: K/2 elements per row + uint32_t a_sp_stride = K / 2; uint32_t num_k_tiles = K / ctx::tileK; uint32_t meta_words_per_tile = ctx::wg_meta_total_bytes / 4; for (uint32_t k = 0; k < K; k += ctx::tileK) { uint32_t k_tile = k / ctx::tileK; - // Cooperative load: compressed A (block-major) and metadata (flat) for all warps - for (uint32_t w = 0; w < num_warps; ++w) { - auto A_smem_w = reinterpret_cast(smem_base + w * per_warp_section); - for (uint32_t i = tid; i < smem_a_elems; i += num_threads) { - uint32_t r = i / (ctx::tileK / 2); - uint32_t c = i % (ctx::tileK / 2); - A_smem_w[ctx::a_sp_blockmajor_idx(r, c)] = - pA[(tile_row + w * ctx::xtileM + r) * a_sp_stride + (k / 2 + c)]; - } - - uint32_t tile_row_idx_w = blockIdx.y * num_warps + w; - auto pMeta_tile_w = pMetaSp + (tile_row_idx_w * num_k_tiles + k_tile) * meta_words_per_tile; - auto meta_smem_w = reinterpret_cast(smem_base + w * per_warp_section + smem_a_bytes); - for (uint32_t i = tid; i < meta_words_per_tile; i += num_threads) { - meta_smem_w[i] = pMeta_tile_w[i]; - } + // Each warp loads its compressed A tile into its private TBUF section. + auto A_smem_w = reinterpret_cast( + smem_base + warp_rank * per_warp_section); + for (uint32_t i = lane; i < smem_a_elems; i += VX_CFG_NUM_THREADS) { + uint32_t r = i / (ctx::tileK / 2); + uint32_t c = i % (ctx::tileK / 2); + A_smem_w[i] = + pA[(tile_row + warp_rank * ctx::xtileM + r) * a_sp_stride + + (k / 2 + c)]; } - // Cooperative load: sparse B written in FEDP candidate-pair (flat) - // layout — the bbuf stores each bank-row verbatim and reads it back as - // straight wiring, so no transpose is needed in hardware. + // Cooperative load: K-major B, matching the TBUF's strided descriptor. for (uint32_t i = tid; i < smem_b_elems; i += num_threads) { uint32_t r = i / ctx::xtileN; uint32_t c = i % ctx::xtileN; - B_smem[ctx::b_sp_flat_idx(r, c)] = pB[(k + r) * N + (tile_col + c)]; + B_smem[c * ctx::tileK + r] = + pB[(k + r) * N + (tile_col + c)]; } __syncthreads(); - // Each warp's A section in smem; metadata immediately follows A. + // Each warp consumes its A section from smem and metadata from global memory. auto A_warp = reinterpret_cast(smem_base + warp_rank * per_warp_section); - auto meta_sp = smem_base + warp_rank * per_warp_section + smem_a_bytes; - auto desc_b = vt::vx_make_smem_desc(B_smem, 0); // flat layout: stride field unused + uint32_t tile_row_idx_w = blockIdx.y * num_warps + warp_rank; + auto pMetaSp = reinterpret_cast(arg->meta_sp_addr); + auto meta_sp = pMetaSp + + (tile_row_idx_w * num_k_tiles + k_tile) * meta_words_per_tile; + auto desc_b = vt::vx_make_smem_desc( + B_smem, ctx::tileK * sizeof(ctx::input_t)); // TCU_LD loads sparse metadata before the WGMMA dispatch (both RS and SS modes). ctx::fragment_a fragA; @@ -96,13 +94,12 @@ __kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { #if defined(WGMMA_RS) && (WGMMA_NRC <= 16) // RS: A from registers, B from smem (NRC <= 16 only). - // ldm=0 selects block-major (matches the cooperative-load via - // a_sp_blockmajor_idx above). - ctx::load_matrix_sync(fragA, A_warp, 0); + ctx::load_matrix_sync(fragA, A_warp, ctx::tileK / 2); ctx::wgmma_sync(fragC, fragA, desc_b, fragC); #else // SS: both A and B from smem descriptors - auto desc_a = vt::vx_make_smem_desc(A_warp, 0); // stride field unused under block-major + auto desc_a = vt::vx_make_smem_desc( + A_warp, (ctx::tileK / 2) * sizeof(ctx::input_t)); ctx::wgmma_sync(fragC, desc_a, desc_b, fragC); #endif @@ -110,6 +107,7 @@ __kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { } // Store C tile to global memory + auto pC = reinterpret_cast(arg->C_addr); auto pTileC = pC + (tile_row + warp_rank * ctx::xtileM) * N + tile_col; ctx::store_matrix_sync(pTileC, fragC, N); }