Skip to content

WGMMA+SP Kernel + AGU Serialization Fix #383

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
NikhilRout wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

WGMMA+SP Kernel + AGU Serialization Fix #383

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
449169e
wg+sp kernel + agu simx fix
NikhilRout Jul 24, 2026
543bbe9
agu serialization rtl fix
NikhilRout Jul 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
238 changes: 107 additions & 131 deletions hw/rtl/tcu/VX_tcu_agu.sv
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "",
Expand Down Expand Up @@ -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

// -----------------------------------------------------------------------
Expand All @@ -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).
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand Down
Loading
Loading