diff --git a/VX_config.toml b/VX_config.toml index 64c8f9d29..af18fa3db 100644 --- a/VX_config.toml +++ b/VX_config.toml @@ -254,6 +254,8 @@ VX_CFG_TCU_NVFP4_ENABLE = false VX_CFG_TCU_SPARSE_ENABLE = false VX_CFG_TCU_DSM_ENABLE = false VX_CFG_TCU_WGMMA_ENABLE = false +VX_CFG_TCU_FEDP2K = false +VX_CFG_TCU_FEDP_DELAY = 4 VX_CFG_TCU_SPARSE_ENABLED = "expr: 1 if $VX_CFG_TCU_SPARSE_ENABLE else 0" VX_CFG_TCU_META_ENABLE = "expr: $VX_CFG_TCU_MX_ENABLE or $VX_CFG_TCU_SPARSE_ENABLE" VX_CFG_TCU_WGMMA_ENABLED = "expr: 1 if $VX_CFG_TCU_WGMMA_ENABLE else 0" diff --git a/ci/testcases/tensor_mx.yaml b/ci/testcases/tensor_mx.yaml index 17825013d..ddf731fae 100644 --- a/ci/testcases/tensor_mx.yaml +++ b/ci/testcases/tensor_mx.yaml @@ -96,28 +96,28 @@ tests: drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k32" - id: sgemm_tcu_mx-simx-nvfp4-nt8 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-simx-nvfp4-nt16 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-simx-nvfp4-nt32 via: blackbox drivers: - simx app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 + configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS args: "-k128" - id: sgemm_tcu_mx-simx-mxint8-nt4 via: blackbox @@ -140,7 +140,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32 args: "-k32" - known_issue: "mxint8 NT16 simx numerical mismatch (NT8/NT32 pass) — pending MX fix" - id: sgemm_tcu_mx-simx-mxint8-nt32 via: blackbox drivers: @@ -155,7 +154,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt8 via: blackbox drivers: @@ -163,7 +161,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt16 via: blackbox drivers: @@ -171,7 +168,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k32" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxfp8-nt32 via: blackbox drivers: @@ -179,7 +175,6 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxfp8 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI args: "-k64" - known_issue: "mxfp8 rtlsim DPI-path numerical mismatch (simx passes) — pending MX fix" - id: sgemm_tcu_mx-rtlsim-mxbf8-nt4 via: blackbox drivers: @@ -241,28 +236,28 @@ tests: drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k32" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt8 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=8 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt16 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=16 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k64" - id: sgemm_tcu_mx-rtlsim-nvfp4-nt32 via: blackbox drivers: - rtlsim app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI + configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS args: "-k128" - id: sgemm_tcu_mx-rtlsim-mxint8-nt4 via: blackbox @@ -292,17 +287,3 @@ tests: app: sgemm_tcu_mx configs: -DVX_CFG_NUM_THREADS=32 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=mxint8 -DOTYPE=int32 -DVX_CFG_TCU_TYPE_DPI args: "-k64" -- id: sgemm_tcu_mx-simx-nvfp4-nt4-tls - via: blackbox - drivers: - - simx - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DTCU_MX_TLS - args: "-k32" -- id: sgemm_tcu_mx-rtlsim-nvfp4-nt4-tls - via: blackbox - drivers: - - rtlsim - app: sgemm_tcu_mx - configs: -DVX_CFG_NUM_THREADS=4 -DVX_CFG_EXT_TCU_ENABLE -DVX_CFG_TCU_MX_ENABLE -DVX_CFG_TCU_FP4_ENABLE -DVX_CFG_TCU_NVFP4_ENABLE -DITYPE=nvfp4 -DOTYPE=fp32 -DVX_CFG_TCU_TYPE_DPI -DTCU_MX_TLS - args: "-k32" diff --git a/ci/testcases/tensor_sp_mx.yaml b/ci/testcases/tensor_sp_mx.yaml index 572096f46..478133363 100644 --- a/ci/testcases/tensor_sp_mx.yaml +++ b/ci/testcases/tensor_sp_mx.yaml @@ -5,10 +5,6 @@ defaults: - 32 - 64 tier: fast - # Seed coverage only — the sparse-MX path has no validated reference yet, so - # every case currently mismatches. Tracked xfail until the sparse-MX result is - # validated; drop this once the cases pass. - known_issue: "sparse-MX (sgemm_tcu_sp_mx) lacks a validated reference; results mismatch" tests: - id: sgemm_tcu_sp_mx-simx-mxfp8-nt4 via: blackbox diff --git a/hw/rtl/core/VX_scoreboard.sv b/hw/rtl/core/VX_scoreboard.sv index d9433482a..d0b39a7ff 100644 --- a/hw/rtl/core/VX_scoreboard.sv +++ b/hw/rtl/core/VX_scoreboard.sv @@ -241,14 +241,14 @@ module VX_scoreboard import VX_gpu_pkg::*; #( wire [PER_ISSUE_WARPS-1:0] arb_ready_in; // FU lock: a sequence must not interleave with another warp at the same FU. - // fu_locked ('1 = open, one-hot = locked) gates arb_valid_in so only the lock - // holder is requested while it holds the lock. - reg [PER_ISSUE_WARPS-1:0] fu_locked; + // Each FU has an independent mask: '1 = open, one-hot = locked. + reg [NUM_EX_UNITS-1:0][PER_ISSUE_WARPS-1:0] fu_locked; for (genvar w = 0; w < PER_ISSUE_WARPS; ++w) begin : g_arb_data_in // operands_ready carries data-hazard + FU-congestion; fu_locked adds the // FU-lock gate so only the lock holder is requested during a sequence. - assign arb_valid_in[w] = staging_if[w].valid && operands_ready[w] && fu_locked[w]; + wire [EX_BITS-1:0] staging_ex = staging_if[w].data.ex_type; + assign arb_valid_in[w] = staging_if[w].valid && operands_ready[w] && fu_locked[staging_ex][w]; assign arb_data_in[w] = { staging_if[w].data.uuid, @@ -281,7 +281,7 @@ module VX_scoreboard import VX_gpu_pkg::*; #( VX_generic_arbiter #( .NUM_REQS (PER_ISSUE_WARPS), .TYPE ((PER_ISSUE_WARPS > 8) ? "M" : "R"), - .STICKY (1) // Greedy + .STICKY (0) ) out_arb ( .clk (clk), .reset (reset), @@ -335,9 +335,9 @@ module VX_scoreboard import VX_gpu_pkg::*; #( fu_locked <= '1; end else if (issue_fire) begin if (issue_fu_lock && ~issue_fu_unlock) begin - fu_locked <= arb_onehot; + fu_locked[issue_ex] <= arb_onehot; end else if (issue_fu_unlock) begin - fu_locked <= '1; + fu_locked[issue_ex] <= '1; end end end diff --git a/hw/rtl/tcu/VX_tcu_abuf.sv b/hw/rtl/tcu/VX_tcu_abuf.sv index fd8cd63c3..9a896b7a7 100644 --- a/hw/rtl/tcu/VX_tcu_abuf.sv +++ b/hw/rtl/tcu/VX_tcu_abuf.sv @@ -62,11 +62,13 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Execute-side observation (req_valid is already gated to WGMMA) input wire [NW_WIDTH-1:0] req_wid, input wire req_valid, + input wire req_setup, input wire [3:0] req_step_m, input wire [3:0] req_step_n, input wire [3:0] req_step_k, input wire [`VX_CFG_XLEN-1:0] req_desc_a, input wire req_a_is_smem, + input wire req_is_sparse, input wire [UUID_WIDTH-1:0] req_uuid, // LMEM bank-parallel read port @@ -74,7 +76,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Outputs output wire abuf_ready, - output wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data + output wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data ); `UNUSED_SPARAM (INSTANCE_ID) `UNUSED_VAR (req_wid) @@ -85,10 +87,14 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_SEL_BITS = $clog2(NUM_BANKS); localparam WORD_SIZE_LOG2 = $clog2(`VX_CFG_XLEN / 8); - localparam A_BLOCK_WORDS = TCU_TC_M * TCU_TC_K; + localparam A_BLOCK_WORDS = TCU_TC_M * TCU_WG_FEDP_K; + localparam A_BLOCK_WORDS_SP = TCU_TC_M * TCU_TC_K; localparam A_BLOCK_BANK_ROWS = (A_BLOCK_WORDS + NUM_BANKS - 1) / NUM_BANKS; + localparam A_BLOCK_BANK_ROWS_SP = (A_BLOCK_WORDS_SP + NUM_BANKS - 1) / NUM_BANKS; localparam BLOCK_WORDS_PADDED = A_BLOCK_BANK_ROWS * NUM_BANKS; + localparam BLOCK_WORDS_PADDED_SP = A_BLOCK_BANK_ROWS_SP * NUM_BANKS; localparam A_STRIPE_BANK_ROWS = TCU_WG_M_STEPS * A_BLOCK_BANK_ROWS; + localparam A_STRIPE_BANK_ROWS_SP = TCU_WG_M_STEPS * A_BLOCK_BANK_ROWS_SP; localparam A_STRIPE_WORDS = A_STRIPE_BANK_ROWS * NUM_BANKS; // XLEN ratio: each physical LMEM bank-row carries XLEN_RATIO logical // 32-bit bank-rows side-by-side. Smem layout is XLEN-independent, so the @@ -97,6 +103,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_ROW_WORDS = NUM_BANKS * XLEN_RATIO; // 32-bit words per LMEM bank-row localparam BANK_ROW_WORDS_LOG2 = $clog2(BANK_ROW_WORDS); localparam A_STRIPE_LMEM_ROWS = (A_STRIPE_WORDS + BANK_ROW_WORDS - 1) / BANK_ROW_WORDS; + localparam A_STRIPE_LMEM_ROWS_SP = (A_STRIPE_BANK_ROWS_SP * NUM_BANKS + BANK_ROW_WORDS - 1) / BANK_ROW_WORDS; // Row-major path issues one LMEM read per logical row (M_STEPS*TC_M rows). localparam A_TOTAL_ROWS = TCU_WG_M_STEPS * TCU_TC_M; localparam FETCH_CTR_W_BM = `CLOG2(A_STRIPE_LMEM_ROWS + 1); @@ -107,10 +114,8 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // per-row 32-bit-word stride. Cap at 14 bits (= 64KB smem / 4B). localparam LDM_W = 14; - // Canonical-config invariant: one A-block fits one (32-bit-equivalent) - // bank-row (TC_M*TC_K == NUM_BANKS). Non-canonical configs requiring - // A_SUB_BLOCKS packing in the output mux are not supported. - `STATIC_ASSERT (A_BLOCK_BANK_ROWS == 1, ("VX_tcu_abuf assumes one A-block per bank-row")) + // Canonical configs fit an A block in one bank-row normally and two + // bank-rows with dense FEDP2K. // ----------------------------------------------------------------------- // Resident slot state @@ -125,29 +130,28 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // alongside fetch_base_r and slot_step_k_r. logic slot_row_major_r; logic [LDM_W-1:0] slot_ldm_words_r; + logic slot_is_sparse_r; wire [`UP(K_STEPS_W)-1:0] req_step_k_trunc = `UP(K_STEPS_W)'(req_step_k); if (4 > K_STEPS_W) begin : g_step_k_upper_unused `UNUSED_VAR (req_step_k[3:`UP(K_STEPS_W)]) end - // The uop expander only reads rs1 (desc_a) on uop 0 of a WGMMA expansion. - // On non-first uops, req_desc_a is garbage and cannot participate in - // the residency check. - wire is_first_uop = (req_step_m == '0) && (req_step_n == '0) && (req_step_k == '0); + // The WGMMA wrapper supplies the setup-latched descriptor on compute uops. + wire is_first_compute_uop = (req_step_m == '0) && (req_step_n == '0) && (req_step_k == '0); `UNUSED_VAR (req_step_n) // only used in is_first_uop computation - // Force a refetch on every WGMMA's first uop. A cooperative-load + // Force a refetch on every WGMMA's first compute uop. A cooperative-load // pattern (K-tile loop rewrites A_warp_smem in place, issues a fresh // WGMMA with an unchanged descriptor) would incorrectly hit the cached // stripe and serve stale A data if residency were checked by desc_a alone. - // refetched_for_first_uop_r gates stripe_resident: cleared until the - // current WGMMA's first_uop fetch completes, then set; cleared again on - // the first non-first uop so the next WGMMA always triggers a fresh fetch. - reg refetched_for_first_uop_r; + // refetched_for_first_compute_r gates stripe_resident: cleared until the + // current WGMMA's first compute fetch completes, then set; cleared again + // on later compute uops so the next WGMMA always triggers a fresh fetch. + reg refetched_for_first_compute_r; wire stripe_resident = slot_valid_r && (slot_step_k_r == req_step_k_trunc) - && (!is_first_uop || refetched_for_first_uop_r); + && (!is_first_compute_uop || refetched_for_first_compute_r); // RS mode (a_from_smem=0): A from registers, abuf bypassed → always ready. wire need_smem = req_valid && req_a_is_smem; @@ -180,14 +184,12 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam STRIPE_STRIDE_BANK_ROWS = A_STRIPE_LMEM_ROWS; - // Use latched desc_a base on non-first uops (req_desc_a is garbage there - // because the uop expander gates the rs1 register read on uop 0 only). - // Without this, k-stripe-transition refills mid-WGMMA would compute the - // wrong fetch_base. + // Compute uops use the setup-latched descriptor. Without this, + // k-stripe-transition refills mid-WGMMA would compute the wrong fetch_base. wire [BANK_ADDR_WIDTH-1:0] effective_desc_a_row_base = - is_first_uop ? desc_a_row_base : slot_desc_a_row_base_r; + is_first_compute_uop ? desc_a_row_base : slot_desc_a_row_base_r; wire [LDM_W-1:0] effective_ldm_words = - is_first_uop ? desc_a_ldm_words : slot_ldm_words_r; + is_first_compute_uop ? desc_a_ldm_words : slot_ldm_words_r; wire effective_row_major = (effective_ldm_words != '0); // Block-major stripe base (unchanged): one fetch covers M_STEPS A-blocks. @@ -220,7 +222,8 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // row-major: A_TOTAL_ROWS per-row reads (one per logical row) wire [FETCH_CTR_W-1:0] target_fetches = slot_row_major_r ? FETCH_CTR_W'(A_TOTAL_ROWS) - : FETCH_CTR_W'(A_STRIPE_LMEM_ROWS); + : (slot_is_sparse_r ? FETCH_CTR_W'(A_STRIPE_LMEM_ROWS_SP) + : FETCH_CTR_W'(A_STRIPE_LMEM_ROWS)); wire all_requested = (req_ctr_r >= target_fetches); wire can_issue = in_fetch && !all_requested @@ -231,9 +234,10 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-row 32-bit-word offset (relative to fetch_base_r * BANK_ROW_WORDS). // Width = LDM (row stride) + counter + a bit for step_k*TC_K headroom. localparam ROW_OFF_W = LDM_W + FETCH_CTR_W + 4; + wire [3:0] row_k_words = slot_is_sparse_r ? 4'(TCU_TC_K) : 4'(TCU_WG_FEDP_K); wire [ROW_OFF_W-1:0] row_word_off_req = ROW_OFF_W'(req_ctr_r) * ROW_OFF_W'(slot_ldm_words_r) - + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(TCU_TC_K); + + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(row_k_words); wire [BANK_ADDR_WIDTH-1:0] row_lmem_addr = fetch_base_r + BANK_ADDR_WIDTH'(row_word_off_req >> BANK_ROW_WORDS_LOG2); @@ -264,63 +268,71 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( fetch_base_r <= '0; slot_row_major_r <= 1'b0; slot_ldm_words_r <= '0; - refetched_for_first_uop_r <= 1'b0; + slot_is_sparse_r <= 1'b0; + refetched_for_first_compute_r <= 1'b0; end else begin - // Inflight tracker (single outstanding request at a time) - if (tcu_lmem_if.rsp_valid) - req_inflight_r <= 1'b0; - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_inflight_r <= 1'b1; - - // Latch desc_a's row base + ldm_words on first uop of every WGMMA - // so non-first uops (k-stripe-transition refills) can reuse them - // without needing the gated req_desc_a bus. - if (req_valid && is_first_uop) begin - slot_desc_a_row_base_r <= desc_a_row_base; - slot_ldm_words_r <= desc_a_ldm_words; - slot_row_major_r <= (desc_a_ldm_words != '0); - end + if (req_setup) begin + fsm_state_r <= S_IDLE; + req_ctr_r <= '0; + rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + slot_valid_r <= 1'b0; + slot_fetching_r <= 1'b0; + refetched_for_first_compute_r <= 1'b0; + end else begin + // Inflight tracker (single outstanding request at a time) + if (tcu_lmem_if.rsp_valid) + req_inflight_r <= 1'b0; + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_inflight_r <= 1'b1; + + // Latch descriptor fields on the first compute uop. + if (req_valid && is_first_compute_uop) begin + slot_desc_a_row_base_r <= desc_a_row_base; + slot_ldm_words_r <= desc_a_ldm_words; + slot_row_major_r <= (desc_a_ldm_words != '0); + end - // refetched_for_first_uop_r: set when fetch completes for - // this WGMMA's first uop; cleared on the next non-first uop. - if (last_rsp && is_first_uop) - refetched_for_first_uop_r <= 1'b1; - else if (fire && !is_first_uop) - refetched_for_first_uop_r <= 1'b0; - - case (fsm_state_r) - S_IDLE: begin - if (alloc_en) begin - fsm_state_r <= S_FETCH; - slot_fetching_r <= 1'b1; - slot_valid_r <= 1'b0; - slot_step_k_r <= req_step_k_trunc; - // Row-major base is the A_warp start bank-row (no - // step_k offset; per-row arithmetic derives the - // exact LMEM addr). Block-major base is the stripe - // origin (step_k already factored in). - fetch_base_r <= effective_row_major - ? effective_desc_a_row_base - : stripe_base; - req_ctr_r <= '0; - rsp_ctr_r <= '0; - req_inflight_r <= 1'b0; + if (last_rsp && is_first_compute_uop) + refetched_for_first_compute_r <= 1'b1; + else if (fire && !is_first_compute_uop) + refetched_for_first_compute_r <= 1'b0; + + case (fsm_state_r) + S_IDLE: begin + if (alloc_en) begin + fsm_state_r <= S_FETCH; + slot_fetching_r <= 1'b1; + slot_valid_r <= 1'b0; + slot_step_k_r <= req_step_k_trunc; + slot_is_sparse_r <= req_is_sparse; + // Row-major base is the A_warp start bank-row (no + // step_k offset; per-row arithmetic derives the + // exact LMEM addr). Block-major base is the stripe + // origin (step_k already factored in). + fetch_base_r <= effective_row_major + ? effective_desc_a_row_base + : stripe_base; + req_ctr_r <= '0; + rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + end end - end - S_FETCH: begin - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_ctr_r <= req_ctr_r + FETCH_CTR_W'(1); - if (last_rsp) begin - fsm_state_r <= S_IDLE; - slot_fetching_r <= 1'b0; - slot_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - end else if (tcu_lmem_if.rsp_valid) begin - rsp_ctr_r <= rsp_ctr_r + FETCH_CTR_W'(1); + S_FETCH: begin + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_ctr_r <= req_ctr_r + FETCH_CTR_W'(1); + if (last_rsp) begin + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + slot_valid_r <= 1'b1; + req_inflight_r <= 1'b0; + end else if (tcu_lmem_if.rsp_valid) begin + rsp_ctr_r <= rsp_ctr_r + FETCH_CTR_W'(1); + end end - end - default: fsm_state_r <= S_IDLE; - endcase + default: fsm_state_r <= S_IDLE; + endcase + end end end @@ -337,7 +349,7 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // within the LMEM response. wire [ROW_OFF_W-1:0] row_word_off_rsp = ROW_OFF_W'(rsp_ctr_r) * ROW_OFF_W'(slot_ldm_words_r) - + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(TCU_TC_K); + + ROW_OFF_W'(slot_step_k_r) * ROW_OFF_W'(row_k_words); wire [BANK_ROW_WORDS_LOG2:0] row_lane_rsp = (BANK_ROW_WORDS_LOG2+1)'( row_word_off_rsp & ROW_OFF_W'(BANK_ROW_WORDS - 1)); @@ -346,13 +358,13 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( storage_wren = '0; if (in_fetch && tcu_lmem_if.rsp_valid) begin if (slot_row_major_r) begin - // Row-major: write TC_K words of row rsp_ctr_r into - // storage[rsp_ctr_r * TC_K .. + TC_K). Source words start + // Row-major: write k_words of row rsp_ctr_r into + // storage[rsp_ctr_r * k_words .. + k_words). Source words start // at row_lane_rsp inside the LMEM response. - for (int k = 0; k < TCU_TC_K; ++k) begin - automatic int dst = int'(rsp_ctr_r) * TCU_TC_K + k; + for (int k = 0; k < TCU_WG_FEDP_K; ++k) begin + automatic int dst = int'(rsp_ctr_r) * int'(row_k_words) + k; automatic int src = int'(row_lane_rsp) + k; - if (dst < A_STRIPE_WORDS && src < BANK_ROW_WORDS) begin + if (k < int'(row_k_words) && dst < A_STRIPE_WORDS && src < BANK_ROW_WORDS) begin storage_wren[dst] = 1'b1; storage_wdata[dst * 32 +: 32] = tcu_lmem_if.rsp_data.data[src * 32 +: 32]; @@ -395,17 +407,27 @@ module VX_tcu_abuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); // ----------------------------------------------------------------------- - // Output: select A-block based on step_m, pass through as 32-bit words. - // Storage holds blocks at [m_blk * BLOCK_WORDS_PADDED .. +BLOCK_WORDS_PADDED). + // Output: select A-block based on step_m. Sparse data is stored compactly, + // then expanded into FEDP_K row stride with the unavailable tail lanes zero. // ----------------------------------------------------------------------- - logic [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] rs1_mux; + logic [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] rs1_mux; always_comb begin rs1_mux = '0; - for (int lane = 0; lane < TCU_BLOCK_CAP; ++lane) begin - if (lane < int'(A_BLOCK_WORDS)) begin - int src_idx; - src_idx = int'(req_step_m) * BLOCK_WORDS_PADDED + lane; + for (int lane = 0; lane < TCU_WG_A_DATA_SIZE; ++lane) begin + int src_idx; + if (req_is_sparse) begin + automatic int row = lane / int'(TCU_WG_FEDP_K); + automatic int k = lane % int'(TCU_WG_FEDP_K); + if (row < int'(TCU_TC_M) && k < int'(TCU_TC_K)) begin + src_idx = int'(req_step_m) * int'(BLOCK_WORDS_PADDED_SP) + + row * int'(TCU_TC_K) + + k; + if (src_idx < int'(A_STRIPE_WORDS)) + rs1_mux[lane] = `VX_CFG_XLEN'(storage_rdata[src_idx]); + end + end else if (lane < int'(A_BLOCK_WORDS)) begin + src_idx = int'(req_step_m) * int'(BLOCK_WORDS_PADDED) + lane; if (src_idx < int'(A_STRIPE_WORDS)) rs1_mux[lane] = `VX_CFG_XLEN'(storage_rdata[src_idx]); end diff --git a/hw/rtl/tcu/VX_tcu_bbuf.sv b/hw/rtl/tcu/VX_tcu_bbuf.sv index 1755dc68e..0983ac556 100644 --- a/hw/rtl/tcu/VX_tcu_bbuf.sv +++ b/hw/rtl/tcu/VX_tcu_bbuf.sv @@ -53,6 +53,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // TB-level uop observation (req_valid is already gated to WGMMA at wrapper) input wire req_valid, + input wire req_setup, input wire req_is_first_uop, input wire req_is_sparse, input wire [3:0] req_step_m, @@ -77,7 +78,8 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam BANK_SEL_BITS = $clog2(NUM_BANKS); localparam WORD_SIZE_LOG2 = $clog2(`VX_CFG_XLEN / 8); - localparam B_BLOCK_WORDS = TCU_TC_K * TCU_TC_N; + localparam B_BLOCK_WORDS = TCU_WG_FEDP_K * TCU_TC_N; + localparam B_BLOCK_WORDS_SP = TCU_TC_K * TCU_TC_N; localparam B_BUF_WORDS = NUM_BANKS; // storage per slot // (= 1 logical 32-bit bank-row) localparam LG_B_SUB_BLOCKS = $clog2(TCU_WG_B_SUB_BLOCKS); @@ -91,6 +93,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam DENSE_BLOCKS_PER_ROW = TCU_WG_B_SUB_BLOCKS * XLEN_RATIO; localparam LG_DENSE_BLOCKS_PER_ROW = (DENSE_BLOCKS_PER_ROW > 1) ? $clog2(DENSE_BLOCKS_PER_ROW) : 0; + localparam SPARSE_ONE_FETCH = (TCU_WG_FEDP_K == (2 * TCU_TC_K)); // Canonical-config invariant: 1 logical (32-bit-equivalent) bank-row // holds B_SUB_BLOCKS blocks (the smem layout is XLEN-independent). @@ -101,7 +104,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // WGMMA SS-descriptor's canonical layout) fetch path. Engaged // when desc_b's stride field (bits [31:16]) is non-zero. Performs // TCU_TC_N per-N-row LMEM reads per (step_k, step_n) WGMMA uop, writing - // TCU_TC_K 32-bit words per row into storage at the b_off offset + // TCU_WG_FEDP_K 32-bit words per row into storage at the b_off offset // tcu_core's `b_off + j*tcK + k` indexing will read. Mirrors // VX_tcu_abuf.sv's row-major fetch for A. localparam LDM_W = 14; @@ -136,7 +139,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // blocks; XLEN_RATIO of them fit in one physical bank-row. // So advance the LMEM addr every // (B_SUB_BLOCKS * XLEN_RATIO) blocks. - // Sparse: a single sparse step_n needs TC_K*TC_N*2 = 2*B_BLOCK_WORDS + // Sparse: a single sparse step_n needs TC_K*TC_N*2 = 2*B_BLOCK_WORDS_SP // 32-bit words, split across TWO dense K-blocks at the SAME // n_blk. Those two blocks live in different physical bank-rows // separated by sp_k_stride (in bank-rows). The two bbuf slots @@ -154,15 +157,15 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( : 5'({1'b0, req_step_n} >> TOTAL_SHIFT); // Sparse K-block stride in physical LMEM bank-rows. - // = n_steps * B_BLOCK_WORDS / (NUM_BANKS * XLEN_RATIO) + // = n_steps * B_BLOCK_WORDS_SP / (NUM_BANKS * XLEN_RATIO) // n_steps = NRC / 2 for the canonical (tcM*tcN == BLOCK_CAP) configs: // NRC=8 → 4, NRC=16 → 8, NRC=32 → 16. - // Old hardcoded table only matched NT=8; NT=16 has B_BLOCK_WORDS == NUM_BANKS + // Old hardcoded table only matched NT=8; NT=16 has B_BLOCK_WORDS_SP == NUM_BANKS // and so needs strides 2× larger. NT=32 happens to match NT=8. localparam SP_STRIDE_DEN = NUM_BANKS * XLEN_RATIO; - localparam SP_STRIDE_NR8 = (4 * B_BLOCK_WORDS) / SP_STRIDE_DEN; - localparam SP_STRIDE_NR16 = (8 * B_BLOCK_WORDS) / SP_STRIDE_DEN; - localparam SP_STRIDE_NR32 = (16 * B_BLOCK_WORDS) / SP_STRIDE_DEN; + localparam SP_STRIDE_NR8 = (4 * B_BLOCK_WORDS_SP) / SP_STRIDE_DEN; + localparam SP_STRIDE_NR16 = (8 * B_BLOCK_WORDS_SP) / SP_STRIDE_DEN; + localparam SP_STRIDE_NR32 = (16 * B_BLOCK_WORDS_SP) / SP_STRIDE_DEN; logic [5:0] sp_k_stride; always_comb begin case (req_cd_nregs) @@ -181,8 +184,7 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( : SUB_HALF_W'(({27'b0, block_index} >> LG_B_SUB_BLOCKS) & ((1 << LG_XLEN_RATIO) - 1)); // Sparse within-physical-bank-row selector. Picks which of the - // DENSE_BLOCKS_PER_ROW dense blocks to extract (B_BLOCK_WORDS 32-bit - // words at offset sparse_pos * B_BLOCK_WORDS in the LMEM response). + // DENSE_BLOCKS_PER_ROW dense blocks to extract. localparam SPARSE_POS_W = (LG_DENSE_BLOCKS_PER_ROW == 0) ? 1 : LG_DENSE_BLOCKS_PER_ROW; wire [SPARSE_POS_W-1:0] sparse_pos_w = (LG_DENSE_BLOCKS_PER_ROW == 0) @@ -227,9 +229,8 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Mode the slot pair was filled under (sparse vs dense). On mode // transition for the same warpgroup we must refill. logic slot_is_sparse_r; - // K-major mode + per-WGMMA latched fields. Latched at alloc_en so - // non-first-uop refills can re-derive addresses (rs2 bus is invalid - // on non-first uops). + // K-major mode + per-WGMMA latched fields. Descriptor fields are latched + // on setup; per-compute step fields are latched when a fetch is allocated. logic slot_row_major_r; logic [LDM_W-1:0] slot_ldm_words_r; logic [3:0] slot_step_k_r; @@ -239,19 +240,14 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( logic [KM_CTR_W-1:0] km_req_ctr_r; logic [KM_CTR_W-1:0] km_rsp_ctr_r; - // req_desc_b is only valid on the first uop of a WGMMA expansion; latch - // desc_b on first uop and use the latched base for subsequent uops. - // is_first_uop is provided by op_args.tcu (set alongside fu_lock in - // VX_tcu_uops), not re-derived here. + // The WGMMA wrapper supplies the setup-latched descriptor on compute uops. `UNUSED_VAR (req_step_m) - wire is_first_uop = req_is_first_uop; wire [BANK_ADDR_WIDTH-1:0] effective_desc_b_row_base = - is_first_uop ? desc_b_row_base : slot_desc_b_row_base_r; + req_is_first_uop ? desc_b_row_base : slot_desc_b_row_base_r; // K-major slot fields (slot_row_major_r / slot_ldm_words_r / // slot_step_k_r / slot_step_n_r) are latched at alloc_en — see the // always_ff below. The K-major addressing arithmetic reads them - // directly. desc_b_ldm_words on the bus is used only for the - // first-uop residency / mode-select comparison. + // directly. desc_b_ldm_words on the bus is used only for setup. // Per-mode fetch addresses. wire [BANK_ADDR_WIDTH-1:0] fetch_addr_dense = @@ -276,10 +272,9 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( && (slot_b_addr_r == fetch_addr_b_sparse) && (slot_b_sparse_pos_r == sparse_pos_w); - // K-major residency (dense + sparse): same (step_k, step_n) already - // in the slot pair. For sparse, BOTH slots must be filled (slot_a holds - // K-pair 0, slot_b holds K-pair 1; both written from the same S_FETCH_A - // multi-fetch — see the storage_write block). + // K-major residency (dense + sparse): same (step_k, step_n) already in + // the slot pair. For sparse, both slots must be filled because sparse B + // is split across the two dense candidate streams. wire bank_row_resident_kmajor = slot_a_valid_r && slot_row_major_r && (!slot_is_sparse_r || slot_b_valid_r) @@ -287,10 +282,10 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( && (slot_step_n_r == req_step_n); // Block-major residency is the existing dense/sparse branch; K-major - // overrides it when the LIVE descriptor (uop 0) or the latched slot mode - // (non-first uops) selects K-major. + // overrides it when the first compute descriptor or latched slot mode + // selects K-major. wire req_wants_kmajor = - is_first_uop ? (desc_b_ldm_words != '0) : slot_row_major_r; + req_is_first_uop ? (desc_b_ldm_words != '0) : slot_row_major_r; wire bank_row_resident = req_wants_kmajor ? bank_row_resident_kmajor @@ -304,23 +299,24 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // K-major address generation // ----------------------------------------------------------------------- // For row r (= km_req_ctr_r) of the current (step_k, step_n) block: - // word_off = (step_n × TCU_TC_N + r) × ldm_words + step_k × TCU_TC_K + // word_off = (step_n × TCU_TC_N + r) × ldm_words + step_k × k_words // bank_row = base + (word_off >> log2(BANK_ROW_WORDS)) // lane = word_off & (BANK_ROW_WORDS - 1) - // The fetched bank-row's `lane..lane + tcK` words are the K-pair operands - // for FEDP slot (j = r, k_pair = 0..tcK-1). + // k_words is fedpK for dense and tcK for sparse + + wire [3:0] km_k_words = slot_is_sparse_r ? 4'(TCU_TC_K) : 4'(TCU_WG_FEDP_K); wire [KM_OFF_W-1:0] km_word_off_req = KM_OFF_W'(slot_step_n_r) * KM_OFF_W'(TCU_TC_N) * KM_OFF_W'(slot_ldm_words_r) + KM_OFF_W'(km_req_ctr_r) * KM_OFF_W'(slot_ldm_words_r) - + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(TCU_TC_K); + + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(km_k_words); wire [BANK_ADDR_WIDTH-1:0] km_lmem_addr = slot_desc_b_row_base_r + BANK_ADDR_WIDTH'(km_word_off_req >> BANK_ROW_WORDS_LOG2); wire [KM_OFF_W-1:0] km_word_off_rsp = KM_OFF_W'(slot_step_n_r) * KM_OFF_W'(TCU_TC_N) * KM_OFF_W'(slot_ldm_words_r) + KM_OFF_W'(km_rsp_ctr_r) * KM_OFF_W'(slot_ldm_words_r) - + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(TCU_TC_K); + + KM_OFF_W'(slot_step_k_r) * KM_OFF_W'(km_k_words); wire [BANK_ROW_WORDS_LOG2:0] km_lane_rsp = (BANK_ROW_WORDS_LOG2+1)'( km_word_off_rsp & KM_OFF_W'(BANK_ROW_WORDS - 1)); @@ -405,85 +401,85 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( km_req_ctr_r <= '0; km_rsp_ctr_r <= '0; end else begin - if (tcu_lmem_if.rsp_valid) - req_inflight_r <= 1'b0; - if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - req_inflight_r <= 1'b1; - - // Latch desc_b base + ldm_words on first uop of every WGMMA so - // non-first uops can re-derive fetch_addr without needing the - // gated req_desc_b bus. - if (req_valid && is_first_uop) begin - slot_desc_b_row_base_r <= desc_b_row_base; - slot_ldm_words_r <= desc_b_ldm_words; - slot_row_major_r <= (desc_b_ldm_words != '0); - end + if (req_setup) begin + slot_a_valid_r <= 1'b0; + slot_b_valid_r <= 1'b0; + slot_fetching_r <= 1'b0; + fsm_state_r <= S_IDLE; + req_inflight_r <= 1'b0; + km_req_ctr_r <= '0; + km_rsp_ctr_r <= '0; + end else begin + if (tcu_lmem_if.rsp_valid) + req_inflight_r <= 1'b0; + if (tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + req_inflight_r <= 1'b1; + + // Latch descriptor fields on the first compute uop. + if (req_valid && req_is_first_uop) begin + slot_desc_b_row_base_r <= desc_b_row_base; + slot_ldm_words_r <= desc_b_ldm_words; + slot_row_major_r <= (desc_b_ldm_words != '0); + end - // K-major req/rsp counters advance independently of FSM state - // (single-outstanding still enforced via req_inflight_r). - if (slot_row_major_r && tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) - km_req_ctr_r <= km_req_ctr_r + KM_CTR_W'(1); - if (slot_row_major_r && tcu_lmem_if.rsp_valid && !km_final_rsp) - km_rsp_ctr_r <= km_rsp_ctr_r + KM_CTR_W'(1); - - case (fsm_state_r) - S_IDLE: begin - if (alloc_en) begin - fsm_state_r <= S_FETCH_A; - slot_fetching_r <= 1'b1; - slot_a_valid_r <= 1'b0; - slot_b_valid_r <= 1'b0; - slot_a_addr_r <= fetch_addr_a; - slot_a_sub_half_r <= dense_sub_half; - slot_a_sparse_pos_r <= sparse_pos_w; - slot_b_addr_r <= fetch_addr_b_sparse; - slot_b_sparse_pos_r <= sparse_pos_w; - slot_is_sparse_r <= req_is_sparse; - // K-major slot state — latched here so non-first - // uops (a different (step_k, step_n) in the same - // WGMMA) can compute their own addressing. - slot_step_k_r <= req_step_k; - slot_step_n_r <= req_step_n; - // slot_row_major_r is latched separately from - // desc_b_ldm_words above (covers refills mid-WGMMA - // where is_first_uop is false but the WGMMA's mode - // is what was set at uop 0). - km_req_ctr_r <= '0; - km_rsp_ctr_r <= '0; - req_inflight_r <= 1'b0; + // K-major req/rsp counters advance independently of FSM state + // (single-outstanding still enforced via req_inflight_r). + if (slot_row_major_r && tcu_lmem_if.req_valid && tcu_lmem_if.req_ready) + km_req_ctr_r <= km_req_ctr_r + KM_CTR_W'(1); + if (slot_row_major_r && tcu_lmem_if.rsp_valid && !km_final_rsp) + km_rsp_ctr_r <= km_rsp_ctr_r + KM_CTR_W'(1); + + case (fsm_state_r) + S_IDLE: begin + if (alloc_en) begin + fsm_state_r <= S_FETCH_A; + slot_fetching_r <= 1'b1; + slot_a_valid_r <= 1'b0; + slot_b_valid_r <= 1'b0; + slot_a_addr_r <= fetch_addr_a; + slot_a_sub_half_r <= dense_sub_half; + slot_a_sparse_pos_r <= sparse_pos_w; + slot_b_addr_r <= fetch_addr_b_sparse; + slot_b_sparse_pos_r <= sparse_pos_w; + slot_is_sparse_r <= req_is_sparse; + // K-major per-compute slot state for this + // (step_k, step_n) block. + slot_step_k_r <= req_step_k; + slot_step_n_r <= req_step_n; + // slot_row_major_r is latched from the setup + // descriptor and reused for all compute refills. + km_req_ctr_r <= '0; + km_rsp_ctr_r <= '0; + req_inflight_r <= 1'b0; + end end - end - S_FETCH_A: begin - if (last_rsp) begin - slot_a_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - // K-major sparse fills BOTH slots from the same - // S_FETCH_A multi-fetch (each response carries one - // N-row's worth of K-pair 0 + K-pair 1 words; the - // storage_write block routes the second half to - // slot_b). Skip S_FETCH_B. - if (slot_is_sparse_r && slot_row_major_r) begin - slot_b_valid_r <= 1'b1; - fsm_state_r <= S_IDLE; - slot_fetching_r <= 1'b0; - end else if (slot_is_sparse_r) begin - fsm_state_r <= S_FETCH_B; - end else begin + S_FETCH_A: begin + if (last_rsp) begin + slot_a_valid_r <= 1'b1; + req_inflight_r <= 1'b0; + if (slot_is_sparse_r && (slot_row_major_r || SPARSE_ONE_FETCH)) begin + slot_b_valid_r <= 1'b1; + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + end else if (slot_is_sparse_r) begin + fsm_state_r <= S_FETCH_B; + end else begin + fsm_state_r <= S_IDLE; + slot_fetching_r <= 1'b0; + end + end + end + S_FETCH_B: begin + if (last_rsp) begin fsm_state_r <= S_IDLE; slot_fetching_r <= 1'b0; + slot_b_valid_r <= 1'b1; + req_inflight_r <= 1'b0; end end - end - S_FETCH_B: begin - if (last_rsp) begin - fsm_state_r <= S_IDLE; - slot_fetching_r <= 1'b0; - slot_b_valid_r <= 1'b1; - req_inflight_r <= 1'b0; - end - end - default: fsm_state_r <= S_IDLE; - endcase + default: fsm_state_r <= S_IDLE; + endcase + end end end @@ -497,14 +493,8 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // XLEN>32, the lower NUM_BANKS otherwise. tcu_core's // b_off then picks within those words at execute time. // - // Sparse write: copies exactly one dense block (B_BLOCK_WORDS words) - // from the physical response at offset - // sparse_pos × B_BLOCK_WORDS. The other words in the - // physical bank-row are dropped — a different step_n - // needs them and will refill on the next request. The - // upper (NUM_BANKS - B_BLOCK_WORDS) storage entries are - // unused in sparse mode but the dense width is retained - // so the LUTRAM can serve both modes. + // Sparse write: extracts the compressed candidate lanes from one or two + // dense B blocks and stores them in the compact rs2 order. // ----------------------------------------------------------------------- logic [B_BUF_WORDS*32-1:0] storage_a_wdata, storage_b_wdata; @@ -523,23 +513,27 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ? (OFF_W'(slot_a_sparse_pos_r) * OFF_W'(B_BLOCK_WORDS)) : (OFF_W'(slot_a_sub_half_r) * OFF_W'(NUM_BANKS)); wire [OFF_W-1:0] b_off_words = slot_is_sparse_r - ? (OFF_W'(slot_b_sparse_pos_r) * OFF_W'(B_BLOCK_WORDS)) + ? (OFF_W'(slot_b_sparse_pos_r) * OFF_W'(B_BLOCK_WORDS_SP)) : '0; - wire [OFF_W-1:0] write_count = slot_is_sparse_r ? OFF_W'(B_BLOCK_WORDS) + wire [OFF_W-1:0] write_count = slot_is_sparse_r ? OFF_W'(B_BLOCK_WORDS_SP) : OFF_W'(NUM_BANKS); // Sparse storage permutation: B_smem block is J-major (n_in outer, k_word // inner), but the FEDP indexes rs2[k*tcN*2 + j*2 + cand] (K-major). // Permute on write so storage[b] holds the word the FEDP expects at slot b: // storage[k_pair*tcN*2 + j*2 + cand] = block[j*tcK + k_pair*2 + cand] - // i.e. src(b) = (b/2/tcN)*2 + (b%2) + ((b/2)%tcN)*tcK. - logic [OFF_W-1:0] sparse_src [B_BUF_WORDS]; + // i.e. src(b) = (b/2/tcN)*2 + (b%2) + ((b/2)%tcN)*k_words. + logic [OFF_W-1:0] sparse_src_a [B_BUF_WORDS]; + logic [OFF_W-1:0] sparse_src_b [B_BUF_WORDS]; always_comb begin for (int b = 0; b < B_BUF_WORDS; ++b) begin automatic int unsigned cand_b = b & 1; automatic int unsigned j_b = (b >> 1) % TCU_TC_N; automatic int unsigned k_pair_b = (b >> 1) / TCU_TC_N; - sparse_src[b] = OFF_W'(j_b * TCU_TC_K + k_pair_b * 2 + cand_b); + sparse_src_a[b] = OFF_W'(j_b * TCU_WG_FEDP_K + k_pair_b * 2 + cand_b); + sparse_src_b[b] = SPARSE_ONE_FETCH + ? OFF_W'(j_b * TCU_WG_FEDP_K + TCU_TC_K + k_pair_b * 2 + cand_b) + : OFF_W'(j_b * TCU_TC_K + k_pair_b * 2 + cand_b); end end @@ -555,21 +549,21 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // two sparse candidates adjacent — word (z*2 + cand) holds the // FEDP's bword{cand} for K-step z. // The FEDP reads rs2[k_idx*(TC_N*2) + j*2 + cand] with k_idx=z; - // slot_a||slot_b splits at B_BLOCK_WORDS (= TC_K*TC_N). + // slot_a||slot_b splits at B_BLOCK_WORDS_SP (= TC_K*TC_N). // Drive each (z,cand) word to its exact flat target and route by slot. for (int z = 0; z < TCU_TC_K; ++z) begin for (int c = 0; c < 2; ++c) begin automatic int src = int'(km_lane_rsp) + z * 2 + c; automatic int tgt = z * (TCU_TC_N * 2) + int'(km_rsp_ctr_r) * 2 + c; if (src < (NUM_BANKS * XLEN_RATIO) && in_fetch_a) begin - if (tgt < int'(B_BLOCK_WORDS)) begin + if (tgt < int'(B_BLOCK_WORDS_SP)) begin if (tgt < B_BUF_WORDS) begin storage_a_wren[tgt] = 1'b1; storage_a_wdata[tgt * 32 +: 32] = tcu_lmem_if.rsp_data.data[src * 32 +: 32]; end end else begin - automatic int tgt_b = tgt - int'(B_BLOCK_WORDS); + automatic int tgt_b = tgt - int'(B_BLOCK_WORDS_SP); if (tgt_b < B_BUF_WORDS) begin storage_b_wren[tgt_b] = 1'b1; storage_b_wdata[tgt_b * 32 +: 32] = @@ -580,12 +574,12 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( end end end else if (slot_row_major_r) begin - // K-major dense: write tcK words for this row (km_rsp_ctr_r) - // into storage[km_b_off + km_rsp_ctr_r * tcK .. + tcK), + // K-major dense: write fedpK words for this row (km_rsp_ctr_r) + // into storage[km_b_off + km_rsp_ctr_r * fedpK .. + fedpK), // sourced from the response at lane km_lane_rsp. - for (int k = 0; k < TCU_TC_K; ++k) begin + for (int k = 0; k < TCU_WG_FEDP_K; ++k) begin automatic int dst = int'(km_b_off) - + int'(km_rsp_ctr_r) * TCU_TC_K + + int'(km_rsp_ctr_r) * TCU_WG_FEDP_K + k; automatic int src = int'(km_lane_rsp) + k; if (dst < B_BUF_WORDS && src < (NUM_BANKS * XLEN_RATIO) && in_fetch_a) begin @@ -597,18 +591,23 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( end else begin for (int b = 0; b < B_BUF_WORDS; ++b) begin if (b < int'(write_count)) begin - automatic logic [OFF_W-1:0] src_off = slot_is_sparse_r - ? sparse_src[b] - : OFF_W'(b); if (in_fetch_a) begin + automatic logic [OFF_W-1:0] src_off = slot_is_sparse_r + ? sparse_src_a[b] + : OFF_W'(b); storage_a_wren[b] = 1'b1; storage_a_wdata[b * 32 +: 32] = tcu_lmem_if.rsp_data.data[(int'(a_off_words) + int'(src_off)) * 32 +: 32]; + if (slot_is_sparse_r && SPARSE_ONE_FETCH) begin + storage_b_wren[b] = 1'b1; + storage_b_wdata[b * 32 +: 32] = + tcu_lmem_if.rsp_data.data[(int'(a_off_words) + int'(sparse_src_b[b])) * 32 +: 32]; + end end - if (in_fetch_b) begin + if (in_fetch_b && !SPARSE_ONE_FETCH) begin storage_b_wren[b] = 1'b1; storage_b_wdata[b * 32 +: 32] = - tcu_lmem_if.rsp_data.data[(int'(b_off_words) + int'(src_off)) * 32 +: 32]; + tcu_lmem_if.rsp_data.data[(int'(b_off_words) + int'(sparse_src_b[b])) * 32 +: 32]; end end end @@ -648,12 +647,11 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .clk (clk), .reset (reset), .read (1'b1), - // In K-major sparse mode both slots are written from the same - // in_fetch_a response (each response carries one N-row's K-pair - // 0 + K-pair 1 words). Block-major sparse keeps the legacy - // S_FETCH_A → S_FETCH_B sequence. + // Sparse FEDP2K block-major and K-major sparse write both slots + // from FETCH_A; legacy block-major sparse keeps S_FETCH_B. .write ((in_fetch_b && tcu_lmem_if.rsp_valid) - || (in_fetch_a && tcu_lmem_if.rsp_valid && slot_row_major_r && slot_is_sparse_r)), + || (in_fetch_a && tcu_lmem_if.rsp_valid && slot_is_sparse_r + && (slot_row_major_r || SPARSE_ONE_FETCH))), .wren (storage_b_wren), .waddr (1'b0), .wdata (storage_b_wdata), @@ -665,8 +663,8 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Output mux. // Dense: rs2[0..NUM_BANKS-1] = storage_A (legacy). tcu_core's b_off // picks within at execute time. - // Sparse: rs2[0..B_BLOCK_WORDS-1] = storage_A[0..B-1] - // rs2[B_BLOCK_WORDS..2*B_BLOCK_WORDS-1] = storage_B[0..B-1] + // Sparse: rs2[0..B_BLOCK_WORDS_SP-1] = storage_A[0..B-1] + // rs2[B_BLOCK_WORDS_SP..2*B_BLOCK_WORDS_SP-1] = storage_B[0..B-1] // Each slot already holds exactly one dense block at the // sparse_pos selected at fetch time, so the mux is a static // concat. This matches tcu_core's sparse indexing @@ -678,11 +676,11 @@ module VX_tcu_bbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( rs2_mux = '0; for (int lane = 0; lane < TCU_WG_RS2_WIDTH; ++lane) begin if (slot_is_sparse_r) begin - if (lane < int'(B_BLOCK_WORDS)) begin + if (lane < int'(B_BLOCK_WORDS_SP)) begin rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[lane]); - end else if (lane < int'(2 * B_BLOCK_WORDS)) begin + end else if (lane < int'(2 * B_BLOCK_WORDS_SP)) begin rs2_mux[lane] = `VX_CFG_XLEN'( - storage_b_rdata[lane - int'(B_BLOCK_WORDS)]); + storage_b_rdata[lane - int'(B_BLOCK_WORDS_SP)]); end end else if (lane < int'(B_BUF_WORDS)) begin rs2_mux[lane] = `VX_CFG_XLEN'(storage_a_rdata[lane]); diff --git a/hw/rtl/tcu/VX_tcu_core.sv b/hw/rtl/tcu/VX_tcu_core.sv index 87da97894..c3c35b368 100644 --- a/hw/rtl/tcu/VX_tcu_core.sv +++ b/hw/rtl/tcu/VX_tcu_core.sv @@ -22,7 +22,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( input wire reset, `ifdef VX_CFG_TCU_WGMMA_ENABLE - input wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + input wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, input wire [TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data, input wire tbuf_ready, `endif @@ -43,23 +43,25 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); `UNUSED_SPARAM (INSTANCE_ID); + localparam FEDP_K = TCU_WG_FEDP_K; + `ifdef VX_CFG_TCU_TYPE_DSP localparam FCVT_LATENCY = 1; localparam FMUL_LATENCY = 8; localparam FADD_LATENCY = 11; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K + 1) * FADD_LATENCY; - localparam FEDP_LATENCY = FCVT_LATENCY + FMUL_LATENCY + FACC_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * FADD_LATENCY; + localparam FEDP_LATENCY = FCVT_LATENCY + FMUL_LATENCY + FACC_LATENCY + FADD_LATENCY; `elsif VX_CFG_TCU_TYPE_BHF localparam FMUL_LATENCY = 2; localparam FADD_LATENCY = 2; localparam FRND_LATENCY = 1; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K + 1) * (FADD_LATENCY + FRND_LATENCY); - localparam FEDP_LATENCY = (FMUL_LATENCY + FRND_LATENCY) + 1 + FACC_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * (FADD_LATENCY + FRND_LATENCY); + localparam FEDP_LATENCY = (FMUL_LATENCY + FRND_LATENCY) + 1 + FACC_LATENCY + (FADD_LATENCY + FRND_LATENCY); `elsif VX_CFG_TCU_TYPE_FPNEW localparam FMUL_LATENCY = 6; localparam FMUX_LATENCY = 1; localparam FADD_LATENCY = 7; - localparam FACC_LATENCY = $clog2(2 * TCU_TC_K) * FADD_LATENCY; + localparam FACC_LATENCY = $clog2(2 * FEDP_K) * FADD_LATENCY; localparam FEDP_LATENCY = FMUL_LATENCY + FMUX_LATENCY + FACC_LATENCY + FADD_LATENCY; `elsif VX_CFG_TCU_TYPE_DPI localparam FMUL_LATENCY = 2; @@ -79,6 +81,10 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam LG_A_BS = $clog2(TCU_A_BLOCK_SIZE); localparam LG_B_BS = $clog2(TCU_B_BLOCK_SIZE); localparam OFF_W = $clog2(TCU_BLOCK_CAP); +`ifdef VX_CFG_TCU_WGMMA_ENABLE + localparam LG_WG_B_BS = $clog2(TCU_WG_B_BLOCK_SIZE); + localparam WG_B_OFF_W = $clog2(TCU_WG_RS2_WIDTH); +`endif `ifdef VX_CFG_TCU_SPARSE_ENABLE localparam LG_B_BS_SP = $clog2(TCU_B_BLOCK_SIZE_SP); @@ -90,11 +96,6 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif `ifdef VX_CFG_TCU_MX_ENABLE - wire is_wmma = (execute_if.data.op_type == INST_TCU_WMMA) - `ifdef VX_CFG_TCU_SPARSE_ENABLE - || (execute_if.data.op_type == INST_TCU_WMMA_SP) - `endif - ; `ifdef VX_CFG_TCU_SPARSE_ENABLE wire mx_is_sparse = is_sparse; `else @@ -103,6 +104,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( localparam FEDP_SF = TCU_MX_MAX_SF; `else localparam FEDP_SF = 1; + `UNUSED_PARAM (FEDP_SF) `endif // ----------------------------------------------------------------------- @@ -112,7 +114,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // common interface. Downstream code uses only these wires and never // references tbuf_* or is_wgmma directly. - wire [TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] rs1_data; + wire [TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] rs1_data; `ifdef VX_CFG_TCU_WGMMA_ENABLE wire [TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] rs2_data; `else @@ -132,9 +134,11 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // RF-side rs2_data is NUM_THREADS lanes wide; the WGMMA bbuf can be // wider (TCU_WG_RS2_WIDTH lanes). Pad/truncate to the wgmma width on // the false branch so both arms match TCU_WG_RS2_WIDTH * XLEN bits. + localparam WG_RS1_BITS = TCU_WG_A_DATA_SIZE * `VX_CFG_XLEN; localparam WG_RS2_BITS = TCU_WG_RS2_WIDTH * `VX_CFG_XLEN; + wire [WG_RS1_BITS-1:0] rs1_data_rf = WG_RS1_BITS'(execute_if.data.rs1_data); wire [WG_RS2_BITS-1:0] rs2_data_rf = WG_RS2_BITS'(execute_if.data.rs2_data); - assign rs1_data = (is_wgmma && wg_a_smem) ? tbuf_rs1_data : execute_if.data.rs1_data; + assign rs1_data = (is_wgmma && wg_a_smem) ? tbuf_rs1_data : rs1_data_rf; assign rs2_data = is_wgmma ? tbuf_rs2_data : rs2_data_rf; `ifdef VX_CFG_TCU_SPARSE_ENABLE @@ -160,6 +164,13 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [4:0] fmt_d = execute_if.data.op_args.tcu.fmt_d; wire execute_fire = execute_if.valid && execute_if.ready; +`ifdef VX_CFG_TCU_WGMMA_ENABLE + wire is_wgmma_setup = is_wgmma && !execute_if.data.header.wb; +`else + wire is_wgmma_setup = 1'b0; +`endif + wire setup_enqueue = execute_fire && is_wgmma_setup; + wire fedp_enqueue = execute_fire && !is_wgmma_setup; // ----------------------------------------------------------------------- // Sparse metadata: VX_tcu_sp_meta (for WMMA_SP) + optional tile-buffer mux @@ -198,9 +209,29 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire mdata_queue_full; - wire result_fire = result_if.valid && result_if.ready; wire fedp_enable, fedp_done; + reg setup_valid_r; + tcu_header_t setup_header_r; + tcu_header_t mdata_queue_out; + + wire setup_result_fire = setup_valid_r && result_if.ready; + wire fedp_result_fire = fedp_done && result_if.ready && !setup_valid_r; + + always @(posedge clk) begin + if (reset) begin + setup_valid_r <= 1'b0; + end else begin + if (setup_result_fire) begin + setup_valid_r <= 1'b0; + end + if (setup_enqueue) begin + setup_valid_r <= 1'b1; + setup_header_r <= execute_if.data.header; + end + end + end + reg [PIPE_LATENCY-1:0] fedp_delay_pipe; always @(posedge clk) begin if (reset) begin @@ -209,18 +240,20 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( if (fedp_enable) begin fedp_delay_pipe <= fedp_delay_pipe >> 1; end - if (execute_fire) begin + if (fedp_enqueue) begin fedp_delay_pipe[PIPE_LATENCY-1] <= 1; end end end assign fedp_done = fedp_delay_pipe[0]; - assign result_if.valid = fedp_done; - assign fedp_enable = ~fedp_done || result_if.ready; - assign execute_if.ready = ~mdata_queue_full && fedp_enable && exe_ready_extra; + assign result_if.valid = setup_valid_r || fedp_done; + assign fedp_enable = ~fedp_done || fedp_result_fire; + assign execute_if.ready = is_wgmma_setup + ? ((~setup_valid_r || result_if.ready) && exe_ready_extra) + : (~mdata_queue_full && fedp_enable && exe_ready_extra); - wire mdata_push = execute_fire; + wire mdata_push = fedp_enqueue; VX_fifo_queue #( .DATAW ($bits(tcu_header_t)), @@ -230,9 +263,9 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .clk (clk), .reset (reset), .push (mdata_push), - .pop (result_fire), + .pop (fedp_result_fire), .data_in(mdata_queue_in), - .data_out(result_if.data.header), + .data_out(mdata_queue_out), `UNUSED_PIN(empty), `UNUSED_PIN(alm_empty), .full (mdata_queue_full), @@ -240,17 +273,24 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `UNUSED_PIN(size) ); + assign result_if.data.header = setup_valid_r ? setup_header_r : mdata_queue_out; + // ----------------------------------------------------------------------- // Operand offset computation // ----------------------------------------------------------------------- wire [OFF_W-1:0] a_off = (OFF_W'(step_m) & OFF_W'(TCU_A_SUB_BLOCKS-1)) << LG_A_BS; + wire [OFF_W-1:0] b_off_wm; `ifdef VX_CFG_TCU_SPARSE_ENABLE - wire [OFF_W-1:0] b_off = is_sparse + assign b_off_wm = is_sparse ? (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS_SP-1)) << LG_B_BS_SP : (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; `else - wire [OFF_W-1:0] b_off = (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; + assign b_off_wm = (OFF_W'(step_n) & OFF_W'(TCU_B_SUB_BLOCKS-1)) << LG_B_BS; +`endif +`ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [WG_B_OFF_W-1:0] b_off_wg = + (WG_B_OFF_W'(step_n) & WG_B_OFF_W'(TCU_WG_B_SUB_BLOCKS-1)) << LG_WG_B_BS; `endif // ----------------------------------------------------------------------- @@ -298,13 +338,17 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .meta_b (mx_meta_b) ); - localparam MX_IDX_W = $clog2(TCU_TILE_M > TCU_TILE_N ? TCU_TILE_M : TCU_TILE_N); - localparam MX_K_IDX_W = `LOG2UP(TCU_TILE_K * TCU_MAX_ELT_RATIO); - localparam MX_SCALE_IDX_W = $clog2(TCU_BLOCK_CAP * 4); + localparam MX_MAX_MN = TCU_TILE_M > TCU_TILE_N ? TCU_TILE_M : TCU_TILE_N; + localparam MX_IDX_W = $clog2(MX_MAX_MN); + localparam MX_TILE_K_MAX = `MAX(TCU_TILE_K, TCU_WG_K_STEPS * TCU_WG_FEDP_K); + localparam MX_K_IDX_W = `LOG2UP(MX_TILE_K_MAX * TCU_MAX_ELT_RATIO); + localparam MX_SCALE_BLOCKS_MAX = mx_scale_blocks_k_words(TCU_NVFP4_ID, MX_TILE_K_MAX); + localparam MX_SCALE_IDX_W = $clog2(MX_MAX_MN * MX_SCALE_BLOCKS_MAX); function automatic [7:0] mx_scale_at( input logic [TCU_BLOCK_CAP-1:0][31:0] meta, input logic [4:0] fmt, + input logic [MX_SCALE_IDX_W-1:0] scale_blocks_k, input logic [MX_IDX_W-1:0] mn_idx, input logic [MX_K_IDX_W-1:0] k_base_idx ); @@ -314,8 +358,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( logic [1:0] byte_idx; begin scale_k = MX_SCALE_IDX_W'(k_base_idx / mx_scale_block_size(fmt)); - scale_idx = MX_SCALE_IDX_W'(mn_idx) * MX_SCALE_IDX_W'(mx_scale_blocks_k(fmt)) - + MX_SCALE_IDX_W'(scale_k); + scale_idx = MX_SCALE_IDX_W'(mn_idx) * scale_blocks_k + scale_k; word_idx = `LOG2UP(TCU_BLOCK_CAP)'(scale_idx >> 2); byte_idx = scale_idx[1:0]; mx_scale_at = meta[word_idx][byte_idx * 8 +: 8]; @@ -325,24 +368,40 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [TCU_TC_M-1:0][FEDP_SF-1:0][7:0] mx_sf_a; wire [TCU_TC_N-1:0][FEDP_SF-1:0][7:0] mx_sf_b; wire [3:0] mx_elems_per_word = 4'(32 / tcu_fmt_width(fmt_s)); - wire [MX_K_IDX_W:0] mx_fedp_elems = (MX_K_IDX_W+1)'( - (MX_K_IDX_W+1)'(TCU_TC_K) * (MX_K_IDX_W+1)'(mx_elems_per_word) + wire [MX_SCALE_IDX_W-1:0] mx_scale_blocks_k_eff = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_WG_K_STEPS * TCU_WG_FEDP_K)) + : MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_TILE_K)); + `else + MX_SCALE_IDX_W'(mx_scale_blocks_k_words(fmt_s, TCU_TILE_K)); + `endif + wire [MX_K_IDX_W:0] mx_uop_k_words = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? (MX_K_IDX_W+1)'(TCU_WG_FEDP_K) : (MX_K_IDX_W+1)'(TCU_TC_K); + `else + (MX_K_IDX_W+1)'(TCU_TC_K); + `endif + wire [MX_K_IDX_W:0] mx_uop_k_elems = (MX_K_IDX_W+1)'( + mx_uop_k_words * (MX_K_IDX_W+1)'(mx_elems_per_word) * (MX_K_IDX_W+1)'(mx_is_sparse ? 2 : 1)); - wire [MX_K_IDX_W-1:0] mx_k_base_idx = MX_K_IDX_W'(step_k * mx_fedp_elems); + wire [MX_K_IDX_W:0] mx_fedp_k_elems = (MX_K_IDX_W+1)'( + (MX_K_IDX_W+1)'(FEDP_K) * (MX_K_IDX_W+1)'(mx_elems_per_word) + * (MX_K_IDX_W+1)'(mx_is_sparse ? 2 : 1)); + wire [MX_K_IDX_W-1:0] mx_k_base_idx = MX_K_IDX_W'(step_k * mx_uop_k_elems); for (genvar i = 0; i < TCU_TC_M; ++i) begin : g_mx_sf_a_i wire [MX_IDX_W-1:0] mx_a_idx = MX_IDX_W'(step_m) * MX_IDX_W'(TCU_TC_M) + MX_IDX_W'(i); for (genvar s = 0; s < FEDP_SF; ++s) begin : g_s - wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_elems) / FEDP_SF); - assign mx_sf_a[i][s] = is_wmma ? mx_scale_at(mx_meta_a, fmt_s, mx_a_idx, mx_k_idx) : '0; + wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_k_elems) / FEDP_SF); + assign mx_sf_a[i][s] = mx_scale_at(mx_meta_a, fmt_s, mx_scale_blocks_k_eff, mx_a_idx, mx_k_idx); end end for (genvar j = 0; j < TCU_TC_N; ++j) begin : g_mx_sf_b_j wire [MX_IDX_W-1:0] mx_b_idx = MX_IDX_W'(step_n) * MX_IDX_W'(TCU_TC_N) + MX_IDX_W'(j); for (genvar s = 0; s < FEDP_SF; ++s) begin : g_s - wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_elems) / FEDP_SF); - assign mx_sf_b[j][s] = is_wmma ? mx_scale_at(mx_meta_b, fmt_s, mx_b_idx, mx_k_idx) : '0; + wire [MX_K_IDX_W-1:0] mx_k_idx = mx_k_base_idx + MX_K_IDX_W'((s * mx_fedp_k_elems) / FEDP_SF); + assign mx_sf_b[j][s] = mx_scale_at(mx_meta_b, fmt_s, mx_scale_blocks_k_eff, mx_b_idx, mx_k_idx); end end `endif @@ -356,18 +415,80 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( for (genvar i = 0; i < TCU_TC_M; ++i) begin : g_i for (genvar j = 0; j < TCU_TC_N; ++j) begin : g_j `ifdef VX_CFG_TCU_SPARSE_ENABLE - wire [TCU_TC_K-1:0][31:0] a_row, b_col, b_col_dense, b_col_sparse, b_col_1, b_col_2; + wire [FEDP_K-1:0][31:0] a_row, b_col, b_col_dense; + wire [TCU_TC_K-1:0][31:0] b_col_sparse, b_col_1, b_col_2; `else - wire [TCU_TC_K-1:0][31:0] a_row, b_col; + wire [FEDP_K-1:0][31:0] a_row, b_col; `endif `ifdef VX_CFG_TCU_MX_ENABLE wire [FEDP_SF-1:0][7:0] sf_a = mx_sf_a[i]; wire [FEDP_SF-1:0][7:0] sf_b = mx_sf_b[j]; `endif - for (genvar k_idx = 0; k_idx < TCU_TC_K; ++k_idx) begin : g_slice_assign - assign a_row[k_idx] = 32'(rs1_data[a_off + i * TCU_TC_K + k_idx]); - `ifdef VX_CFG_TCU_SPARSE_ENABLE - assign b_col_dense[k_idx] = 32'(rs2_data[b_off + j * TCU_TC_K + k_idx]); + for (genvar k_idx = 0; k_idx < FEDP_K; ++k_idx) begin : g_slice_assign + `ifdef VX_CFG_TCU_WGMMA_ENABLE + localparam int WG_B_IDX = j * TCU_WG_FEDP_K + k_idx; + `endif + if (k_idx < TCU_TC_K) begin : g_lo + `ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [31:0] a_wgmma_smem = 32'(rs1_data[i * TCU_WG_FEDP_K + k_idx]); + wire [31:0] a_wgmma_reg = 32'(execute_if.data.rs1_data[i * TCU_TC_K + k_idx]); + assign a_row[k_idx] = is_wgmma + ? (wg_a_smem ? a_wgmma_smem : a_wgmma_reg) + : 32'(execute_if.data.rs1_data[a_off + i * TCU_TC_K + k_idx]); + `else + assign a_row[k_idx] = 32'(rs1_data[a_off + i * TCU_TC_K + k_idx]); + `endif + `ifdef VX_CFG_TCU_SPARSE_ENABLE + assign b_col_dense[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'(rs2_data[b_off_wm + j * TCU_TC_K + k_idx]); + `else + assign b_col[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'(rs2_data[b_off_wm + j * TCU_TC_K + k_idx]); + `endif + end else begin : g_hi + `ifdef VX_CFG_TCU_WGMMA_ENABLE + wire [31:0] a_wgmma_smem = 32'(rs1_data[i * TCU_WG_FEDP_K + k_idx]); + wire [31:0] a_wgmma_reg = + `ifdef VX_CFG_TCU_FEDP2K + `ifdef VX_CFG_TCU_SPARSE_ENABLE + is_sparse ? 32'b0 : + `endif + 32'(execute_if.data.rs2_data[i * TCU_TC_K + (k_idx - TCU_TC_K)]); + `else + 32'b0; + `endif + assign a_row[k_idx] = (is_wgmma + `ifdef VX_CFG_TCU_SPARSE_ENABLE + && !is_sparse + `endif + ) ? (wg_a_smem ? a_wgmma_smem : a_wgmma_reg) : 32'b0; + `else + assign a_row[k_idx] = 32'b0; + `endif + `ifdef VX_CFG_TCU_SPARSE_ENABLE + assign b_col_dense[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + (is_wgmma && !is_sparse) ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'b0; + `else + assign b_col[k_idx] = + `ifdef VX_CFG_TCU_WGMMA_ENABLE + is_wgmma ? 32'(rs2_data[int'(b_off_wg) + WG_B_IDX]) : + `endif + 32'b0; + `endif + end + end + + `ifdef VX_CFG_TCU_SPARSE_ENABLE + for (genvar k_idx = 0; k_idx < TCU_TC_K; ++k_idx) begin : g_sparse_slice_assign localparam J_SP = SYM_SPARSE ? (j % (TCU_TC_N / 2)) : j; // rs2_data sparse-pair layout differs by op: // WGMMA_SP: source is tbuf (shared mem), K-major → @@ -376,20 +497,18 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // idx = J_SP*(TC_K*2) + k_idx*2 + lane // The two layouts are incompatible; separate formulas are required. `ifdef VX_CFG_TCU_WGMMA_ENABLE - wire [31:0] b_col_1_wg = 32'(rs2_data[b_off + k_idx * TCU_TC_N * 2 + J_SP * 2]); - wire [31:0] b_col_2_wg = 32'(rs2_data[b_off + k_idx * TCU_TC_N * 2 + J_SP * 2 + 1]); - wire [31:0] b_col_1_wm = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2]); - wire [31:0] b_col_2_wm = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); + wire [31:0] b_col_1_wg = 32'(rs2_data[k_idx * TCU_TC_N * 2 + J_SP * 2]); + wire [31:0] b_col_2_wg = 32'(rs2_data[k_idx * TCU_TC_N * 2 + J_SP * 2 + 1]); + wire [31:0] b_col_1_wm = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2]); + wire [31:0] b_col_2_wm = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); assign b_col_1[k_idx] = is_wgmma ? b_col_1_wg : b_col_1_wm; assign b_col_2[k_idx] = is_wgmma ? b_col_2_wg : b_col_2_wm; `else - assign b_col_1[k_idx] = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2]); - assign b_col_2[k_idx] = 32'(rs2_data[b_off + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); - `endif - `else - assign b_col[k_idx] = 32'(rs2_data[b_off + j * TCU_TC_K + k_idx]); + assign b_col_1[k_idx] = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2]); + assign b_col_2[k_idx] = 32'(rs2_data[b_off_wm + J_SP * TCU_TC_K * 2 + k_idx * 2 + 1]); `endif end + `endif wire [31:0] c_val = 32'(execute_if.data.rs3_data[i * TCU_TC_N + j]); @@ -404,7 +523,13 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .vld_mask (vld_meta_block), .b_col_out (b_col_sparse) ); - assign b_col = is_sparse ? b_col_sparse : b_col_dense; + for (genvar k_idx = 0; k_idx < FEDP_K; ++k_idx) begin : g_sparse_b_select + if (k_idx < TCU_TC_K) begin : g_lo + assign b_col[k_idx] = is_sparse ? b_col_sparse[k_idx] : b_col_dense[k_idx]; + end else begin : g_hi + assign b_col[k_idx] = is_sparse ? 32'b0 : b_col_dense[k_idx]; + end + end `ifdef VX_TCU_LD_TRACE // GATHER trace: GATHER,wid,step_m,step_n,i,k,bword0,bword1,lo,hi,gathered @@ -450,7 +575,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif wire [4:0] fmt_s_r, fmt_d_r; - wire [TCU_TC_K-1:0][31:0] a_row_r, b_col_r; + wire [FEDP_K-1:0][31:0] a_row_r, b_col_r; `ifdef VX_CFG_TCU_MX_ENABLE wire [FEDP_SF-1:0][7:0] sf_a_r, sf_b_r; `endif @@ -458,7 +583,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `ifdef VX_CFG_TCU_MX_ENABLE VX_pipe_register #( - .DATAW (32 + 5 + 5 + TCU_TC_K * 32 + TCU_TC_K * 32 + 2 * FEDP_SF * 8) + .DATAW (32 + 5 + 5 + FEDP_K * 32 + FEDP_K * 32 + 2 * FEDP_SF * 8) ) pipe_fedp ( .clk (clk), .reset (reset), @@ -468,7 +593,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); `else VX_pipe_register #( - .DATAW (32 + 5 + 5 + TCU_TC_K * 32 + TCU_TC_K * 32) + .DATAW (32 + 5 + 5 + FEDP_K * 32 + FEDP_K * 32) ) pipe_fedp ( .clk (clk), .reset (reset), @@ -482,7 +607,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_dpi #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K), + .N (FEDP_K), .SF (FEDP_SF) ) fedp ( .clk (clk), @@ -503,7 +628,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_bhf #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), @@ -519,7 +644,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_fpnew #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), @@ -535,7 +660,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_tfr #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K), + .N (FEDP_K), .SF (FEDP_SF) ) fedp ( .clk (clk), @@ -557,7 +682,7 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_tcu_fedp_dsp #( .INSTANCE_ID (INSTANCE_ID), .LATENCY (FEDP_LATENCY), - .N (TCU_TC_K) + .N (FEDP_K) ) fedp ( .clk (clk), .reset (reset), @@ -582,9 +707,9 @@ module VX_tcu_core import VX_gpu_pkg::*, VX_tcu_pkg::*; #( always @(posedge clk) begin if (execute_if.valid && execute_if.ready) begin `TRACE(3, ("%t: %s FEDP-enq: wid=%0d, cta_id=%0d, i=%0d, j=%0d, m=%0d, n=%0d, a_row=", $time, INSTANCE_ID, execute_if.data.header.wid, execute_if.data.header.cta_id, i, j, step_m, step_n)) - `TRACE_ARRAY1D(2, "0x%0h", a_row, TCU_TC_K) + `TRACE_ARRAY1D(2, "0x%0h", a_row, FEDP_K) `TRACE(3, (", b_col=")); - `TRACE_ARRAY1D(2, "0x%0h", b_col, TCU_TC_K) + `TRACE_ARRAY1D(2, "0x%0h", b_col, FEDP_K) `TRACE(3, (", c_val=0x%0h (#%0d)\n", c_val, execute_if.data.header.uuid)); end if (result_if.valid && result_if.ready) begin diff --git a/hw/rtl/tcu/VX_tcu_lockstep.sv b/hw/rtl/tcu/VX_tcu_lockstep.sv index 42cf9d9ed..ff3d43876 100644 --- a/hw/rtl/tcu/VX_tcu_lockstep.sv +++ b/hw/rtl/tcu/VX_tcu_lockstep.sv @@ -23,12 +23,10 @@ // drain before a new CTA enters. Same-CTA across blocks (the production // case for a warpgroup at the same uop) remains free. // -// Single-owner design: by construction only one CTA can occupy the bbuf -// at a time, so a single tcu_owner_r / tcu_owned_r pair is sufficient. -// Multiple blocks firing leader uops on the same cycle for the same CTA -// (the production warpgroup case) agree on the cta, so the priority- -// encoded pick agrees with all of them. No same-cycle multi-CTA-firing -// race exists in current configs. +// Per-block ownership matches the SimX model: a block records its CTA when +// its first compute uop enters WGMMA, and releases on its last compute uop. +// A block is deferred only while another block is mid-WGMMA for a different +// CTA. // // Contract: the consumer must AND `cta_conflict` into the request // validity presented to bbuf and into any downstream ready that gates @@ -56,29 +54,24 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // ----------------------------------------------------------------------- // State // ----------------------------------------------------------------------- - reg tcu_owned_r; - reg [NCTA_WIDTH-1:0] tcu_owner_r; - // Per-block "WGMMA expansion in progress" — set on first sub-uop fire, - // cleared on last sub-uop fire. Persists across LMEM-stall gaps so the - // gate keeps a different CTA's WGMMA from sneaking in mid-expansion - // and corrupting the shared bbuf descriptor latch. + // Per-block "WGMMA expansion in progress" — first compute uop joins the + // CTA-owned region; the last compute uop releases it. This persists + // across LMEM-stall gaps so a different CTA cannot enter mid-expansion + // and corrupt the shared bbuf state. reg [BLOCK_SIZE-1:0] in_expansion_r; + reg [BLOCK_SIZE-1:0][NCTA_WIDTH-1:0] cta_owner_r; // ----------------------------------------------------------------------- // Owner update. // - // Leader-fire detection: any block firing its first sub-uop this cycle. + // Leader-fire detection: any block joining a WGMMA expansion this cycle. // The owner cta is picked from the lowest-indexed leader-firing block. // Under same-CTA warpgroup fire (production path), every leader-firing // block carries the same cta_id, so the priority pick agrees with all. // ----------------------------------------------------------------------- wire [BLOCK_SIZE-1:0] leader_fire_b = exec_fire_b & is_first_uop_b & is_wgmma_b; - wire any_leader_fire = |leader_fire_b; - wire any_in_expansion_next; - // in_expansion_next reflects this-cycle updates (set on leader fire, - // clear on last-uop fire). Released the cycle after the last block - // exits expansion. + // clear on last-uop fire). wire [BLOCK_SIZE-1:0] in_expansion_next; for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_inexp_next wire is_wgmma_fire = exec_fire_b[bi] && is_wgmma_b[bi]; @@ -86,33 +79,16 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( (in_expansion_r[bi] || (is_wgmma_fire && is_first_uop_b[bi])) && !(is_wgmma_fire && is_last_uop_b[bi]); end - assign any_in_expansion_next = |in_expansion_next; - - // first_fire_cta: priority-encoded pick of the lowest leader-firing block. - logic [NCTA_WIDTH-1:0] first_fire_cta; - always_comb begin - first_fire_cta = '0; - for (int b = BLOCK_SIZE-1; b >= 0; b--) begin - if (leader_fire_b[b]) first_fire_cta = new_cta_b[b]; - end - end always @(posedge clk) begin if (reset) begin - tcu_owned_r <= 1'b0; - tcu_owner_r <= '0; in_expansion_r <= '0; + cta_owner_r <= '0; end else begin in_expansion_r <= in_expansion_next; - if (!tcu_owned_r) begin - if (any_leader_fire) begin - tcu_owned_r <= 1'b1; - tcu_owner_r <= first_fire_cta; - end - end else begin - // Release once no block is still mid-expansion. - if (!any_in_expansion_next) begin - tcu_owned_r <= 1'b0; + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + if (leader_fire_b[bi]) begin + cta_owner_r[bi] <= new_cta_b[bi]; end end end @@ -122,9 +98,16 @@ module VX_tcu_lockstep import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-block conflict — flat combinational, no propagation chain. // ----------------------------------------------------------------------- for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_conflict - assign cta_conflict[bi] = is_wgmma_b[bi] - && tcu_owned_r - && (tcu_owner_r != new_cta_b[bi]); + logic cta_conflict_r; + always_comb begin + cta_conflict_r = 1'b0; + for (int k = 0; k < BLOCK_SIZE; ++k) begin + if (k != bi && in_expansion_r[k] && (cta_owner_r[k] != new_cta_b[bi])) begin + cta_conflict_r = 1'b1; + end + end + end + assign cta_conflict[bi] = is_wgmma_b[bi] && cta_conflict_r; end // ----------------------------------------------------------------------- diff --git a/hw/rtl/tcu/VX_tcu_pkg.sv b/hw/rtl/tcu/VX_tcu_pkg.sv index 1f3203439..532112a78 100644 --- a/hw/rtl/tcu/VX_tcu_pkg.sv +++ b/hw/rtl/tcu/VX_tcu_pkg.sv @@ -90,23 +90,29 @@ package VX_tcu_pkg; // WGMMA per-warp tile dimensions (NRA=4 fixed, NRC=NR variable). // Derived from block geometry: xtileM = 2*tcM, xtileK = 2*tcK. - // m_steps = k_steps = 2 always. + // Dense FEDP width doubles under VX_CFG_TCU_FEDP2K localparam TCU_WG_TILE_M = 2 * TCU_TC_M; localparam TCU_WG_TILE_K = 2 * TCU_TC_K; +`ifdef VX_CFG_TCU_FEDP2K + localparam TCU_WG_FEDP_K = 2 * TCU_TC_K; +`else + localparam TCU_WG_FEDP_K = TCU_TC_K; +`endif localparam TCU_WG_TILE_N = (TCU_WG_NR * TCU_NT) / TCU_WG_TILE_M; // WG step counts: block geometry (TC_M/TC_N/TC_K) unchanged, tile is larger localparam TCU_WG_M_STEPS = TCU_WG_TILE_M / TCU_TC_M; localparam TCU_WG_N_STEPS = TCU_WG_TILE_N / TCU_TC_N; - localparam TCU_WG_K_STEPS = TCU_WG_TILE_K / TCU_TC_K; + localparam TCU_WG_K_STEPS = TCU_WG_TILE_K / TCU_WG_FEDP_K; localparam TCU_WG_UOPS = TCU_WG_M_STEPS * TCU_WG_N_STEPS * TCU_WG_K_STEPS; // WG A/B micro-tiling (block geometry is shared with non-WG) localparam TCU_WG_A_BLOCK_SIZE = TCU_TC_M * TCU_TC_K; + localparam TCU_WG_A_DATA_SIZE = TCU_TC_M * TCU_WG_FEDP_K; localparam TCU_WG_A_SUB_BLOCKS = TCU_BLOCK_CAP / TCU_WG_A_BLOCK_SIZE; - localparam TCU_WG_B_BLOCK_SIZE = TCU_TC_K * TCU_TC_N; + localparam TCU_WG_B_BLOCK_SIZE = TCU_WG_FEDP_K * TCU_TC_N; localparam TCU_WG_B_SUB_BLOCKS = TCU_BLOCK_CAP / TCU_WG_B_BLOCK_SIZE; // Symmetric sparse flag (NT=4, NT=16: block_em == block_en) @@ -121,9 +127,11 @@ package VX_tcu_pkg; localparam TCU_B_SUB_BLOCKS_SP = TCU_BLOCK_CAP / TCU_B_BLOCK_SIZE_SP; // WGMMA_SP always needs the full candidate lane set, regardless of SYM_SPARSE. localparam TCU_WG_B_BLOCK_SIZE_SP = TCU_TC_K * TCU_TC_N * 2; - // Width of the tbuf_rs2_data port: wider only when SPARSE is enabled (WGMMA_SP path). - // Without SPARSE, only TCU_BLOCK_CAP lanes are ever consumed, so keep the port narrow. - localparam TCU_WG_RS2_WIDTH = `VX_CFG_TCU_SPARSE_ENABLED ? TCU_WG_B_BLOCK_SIZE_SP : TCU_BLOCK_CAP; + localparam TCU_WG_RS2_WIDTH_DENSE = + (TCU_WG_B_BLOCK_SIZE > TCU_BLOCK_CAP) ? TCU_WG_B_BLOCK_SIZE : TCU_BLOCK_CAP; + localparam TCU_WG_RS2_WIDTH = `VX_CFG_TCU_SPARSE_ENABLED + ? ((TCU_WG_B_BLOCK_SIZE_SP > TCU_WG_RS2_WIDTH_DENSE) ? TCU_WG_B_BLOCK_SIZE_SP : TCU_WG_RS2_WIDTH_DENSE) + : TCU_WG_RS2_WIDTH_DENSE; localparam TCU_MIN_FMT_WIDTH = 4; //int4 localparam TCU_MAX_ELT_RATIO = 32 / TCU_MIN_FMT_WIDTH; @@ -171,7 +179,7 @@ package VX_tcu_pkg; input int unsigned block_elems ); automatic int unsigned sparse_ratio = `VX_CFG_TCU_SPARSE_ENABLED ? 2 : 1; - automatic int unsigned fedp_elems = TCU_TC_K * (32 / data_bits) * sparse_ratio; + automatic int unsigned fedp_elems = TCU_WG_FEDP_K * (32 / data_bits) * sparse_ratio; return (fedp_elems + block_elems - 1) / block_elems; endfunction @@ -303,13 +311,20 @@ package VX_tcu_pkg; endcase endfunction - function automatic int unsigned mx_scale_blocks_k(input logic [TCU_FMT_WIDTH-1:0] fmt); + function automatic int unsigned mx_scale_blocks_k_words( + input logic [TCU_FMT_WIDTH-1:0] fmt, + input int unsigned tile_k_words + ); automatic int unsigned data_bits = tcu_fmt_width(fmt); automatic int unsigned block_elems = mx_scale_block_size(fmt); - automatic int unsigned tile_elems = (data_bits != 0) ? TCU_TILE_K * (32 / data_bits) : 0; + automatic int unsigned tile_elems = (data_bits != 0) ? tile_k_words * (32 / data_bits) : 0; return (tile_elems + block_elems - 1) / block_elems; endfunction + function automatic int unsigned mx_scale_blocks_k(input logic [TCU_FMT_WIDTH-1:0] fmt); + return mx_scale_blocks_k_words(fmt, TCU_TILE_K); + endfunction + function automatic logic [4:0] meta_num_cols(input logic [TCU_FMT_WIDTH-1:0] fmt); automatic int hw = tcu_fmt_width(fmt) / 2; return 5'((TCU_BLOCK_CAP + hw - 1) / hw); @@ -424,6 +439,7 @@ package VX_tcu_pkg; // step_m/step_n/step_k at every consumer. logic is_first_uop; logic is_last_uop; + logic setup_fire; `ifdef VX_CFG_TCU_SPARSE_ENABLE logic is_sparse; `endif diff --git a/hw/rtl/tcu/VX_tcu_tbuf.sv b/hw/rtl/tcu/VX_tcu_tbuf.sv index ec01e2d9a..724866b04 100644 --- a/hw/rtl/tcu/VX_tcu_tbuf.sv +++ b/hw/rtl/tcu/VX_tcu_tbuf.sv @@ -60,7 +60,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_mem_bus_if.master tcu_lmem_if, // Per-block operand outputs (rs2 is broadcast — bbuf is shared) - output wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + output wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, output wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0]tbuf_rs2_data, output wire [BLOCK_SIZE-1:0] tbuf_ready ); @@ -84,7 +84,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( // Per-block abufs // ----------------------------------------------------------------------- - wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data_w; + wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] abuf_rs1_data_w; wire [BLOCK_SIZE-1:0] abuf_ready_w; `ifdef PERF_ENABLE @@ -106,11 +106,17 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif .req_wid (req[b].wid), .req_valid (req[b].valid), + .req_setup (req[b].setup_fire), .req_step_m (req[b].step_m), .req_step_n (req[b].step_n), .req_step_k (req[b].step_k), .req_desc_a (req[b].desc_a), .req_a_is_smem (req[b].a_is_smem), + `ifdef VX_CFG_TCU_SPARSE_ENABLE + .req_is_sparse (req[b].is_sparse), + `else + .req_is_sparse (1'b0), + `endif .req_uuid (req[b].uuid), .tcu_lmem_if (lmem_masters[b]), .abuf_ready (abuf_ready_w[b]), @@ -136,6 +142,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( `endif wire bbuf_req_valid; + wire bbuf_req_setup; wire bbuf_req_is_first_uop; wire bbuf_req_is_sparse; wire [3:0] bbuf_req_step_m; @@ -148,6 +155,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( if (BLOCK_SIZE == 1) begin : g_bbuf_inputs_n1 assign bbuf_req_uuid = req[0].uuid; assign bbuf_req_valid = req[0].valid; + assign bbuf_req_setup = req[0].setup_fire; assign bbuf_req_is_first_uop = req[0].is_first_uop; `ifdef VX_CFG_TCU_SPARSE_ENABLE assign bbuf_req_is_sparse = req[0].is_sparse; @@ -161,8 +169,10 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign bbuf_req_desc_b = req[0].desc_b; end else begin : g_bbuf_inputs_pe wire [BLOCK_SIZE-1:0] req_valid_v; + wire [BLOCK_SIZE-1:0] req_setup_v; for (genvar b = 0; b < BLOCK_SIZE; ++b) begin : g_req_valid_v assign req_valid_v[b] = req[b].valid; + assign req_setup_v[b] = req[b].setup_fire; end wire [BLOCK_SIZE-1:0] bbuf_sel_oh; @@ -213,6 +223,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign bbuf_req_uuid = sel_uuid; assign bbuf_req_valid = bbuf_sel_valid; + assign bbuf_req_setup = |req_setup_v; assign bbuf_req_is_first_uop = sel_is_first_uop; assign bbuf_req_is_sparse = sel_is_sparse; assign bbuf_req_step_m = sel_step_m; @@ -235,6 +246,7 @@ module VX_tcu_tbuf import VX_gpu_pkg::*, VX_tcu_pkg::*; #( .lmem_reads (bbuf_lmem_reads_w), `endif .req_valid (bbuf_req_valid), + .req_setup (bbuf_req_setup), .req_is_first_uop (bbuf_req_is_first_uop), .req_is_sparse (bbuf_req_is_sparse), .req_step_m (bbuf_req_step_m), diff --git a/hw/rtl/tcu/VX_tcu_unit.sv b/hw/rtl/tcu/VX_tcu_unit.sv index 26148747c..281099873 100644 --- a/hw/rtl/tcu/VX_tcu_unit.sv +++ b/hw/rtl/tcu/VX_tcu_unit.sv @@ -151,7 +151,7 @@ module VX_tcu_unit import VX_gpu_pkg::*, VX_tcu_pkg::*; #( assign exec_data_w[bi] = core_execute_if[bi].data; end - wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data; + wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data; wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data; wire [BLOCK_SIZE-1:0] tbuf_ready_eff; diff --git a/hw/rtl/tcu/VX_tcu_uops.sv b/hw/rtl/tcu/VX_tcu_uops.sv index 16599d3aa..292d75e26 100644 --- a/hw/rtl/tcu/VX_tcu_uops.sv +++ b/hw/rtl/tcu/VX_tcu_uops.sv @@ -49,11 +49,15 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam IDX_BITS = LG_N + LG_M + LG_K; `ifdef VX_CFG_TCU_WGMMA_ENABLE `ifdef VX_CFG_TCU_SPARSE_ENABLE - localparam MAX_WG_UOPS_SP = TCU_WG_UOPS / 2; - localparam CTR_W_MAX_WMMA = MAX_UOPS > TCU_WG_UOPS ? MAX_UOPS : TCU_WG_UOPS; - localparam CTR_W_BASE = $clog2(CTR_W_MAX_WMMA > MAX_WG_UOPS_SP ? CTR_W_MAX_WMMA : MAX_WG_UOPS_SP); + localparam TCU_WG_K_STEPS_SP = (TCU_WG_K_STEPS > 1) ? (TCU_WG_K_STEPS / 2) : 1; + localparam MAX_WG_UOPS_SP = TCU_WG_M_STEPS * TCU_WG_N_STEPS * TCU_WG_K_STEPS_SP; + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS + 1; + localparam MAX_WG_UOPS_SPARSE = MAX_WG_UOPS_SP + 1; + localparam CTR_W_MAX_WMMA = MAX_UOPS > MAX_WG_UOPS_DENSE ? MAX_UOPS : MAX_WG_UOPS_DENSE; + localparam CTR_W_BASE = $clog2(CTR_W_MAX_WMMA > MAX_WG_UOPS_SPARSE ? CTR_W_MAX_WMMA : MAX_WG_UOPS_SPARSE); `else - localparam CTR_W_BASE = $clog2(MAX_UOPS > TCU_WG_UOPS ? MAX_UOPS : TCU_WG_UOPS); + localparam MAX_WG_UOPS_DENSE = TCU_WG_UOPS + 1; + localparam CTR_W_BASE = $clog2(MAX_UOPS > MAX_WG_UOPS_DENSE ? MAX_UOPS : MAX_WG_UOPS_DENSE); `endif `else localparam CTR_W_BASE = $clog2(MAX_UOPS); @@ -78,11 +82,14 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( || (ibuf_in.op_type == INST_TCU_WGMMA_SP) `endif ; +`ifdef VX_CFG_TCU_SPARSE_ENABLE + wire wg_is_sparse = (ibuf_in.op_type == INST_TCU_WGMMA_SP); +`endif wire wg_a_from_smem = ibuf_in.op_args.tcu.a_from_smem; // Variable NRC based on cd_nregs: 0→8, 1→16, 2→32 // Loop order: m (inner) → n → k (outer) [K-outer] - // m_steps=2 and k_steps=2 are fixed; n varies (middle). + // m_steps=2; k_steps is 2 normally and 1 with FEDP2K. // K-outer lets independent (m,n) tiles overlap FEDP latency. localparam LG_M_WG = $clog2(TCU_WG_M_STEPS); // 1 localparam LG_K_WG = $clog2(TCU_WG_K_STEPS); // 1 @@ -93,19 +100,36 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam WG_MN_NR16 = 16; // nrc=16: n_steps=8, mn=16 localparam WG_MN_NR32 = 32; // nrc=32: n_steps=16, mn=32 - // Pre-computed total uop counts for each cd_nregs value (dense) + // Pre-computed compute uop counts for each cd_nregs value (dense) localparam WG_UOPS_NR8 = WG_MN_NR8 * TCU_WG_K_STEPS; localparam WG_UOPS_NR16 = WG_MN_NR16 * TCU_WG_K_STEPS; localparam WG_UOPS_NR32 = WG_MN_NR32 * TCU_WG_K_STEPS; +`ifdef VX_CFG_TCU_SPARSE_ENABLE + localparam WG_UOPS_SP_NR8 = WG_MN_NR8 * TCU_WG_K_STEPS_SP; + localparam WG_UOPS_SP_NR16 = WG_MN_NR16 * TCU_WG_K_STEPS_SP; + localparam WG_UOPS_SP_NR32 = WG_MN_NR32 * TCU_WG_K_STEPS_SP; +`endif // Mux uop count based on cd_nregs: 0→8, 1→16, 2→32 reg [UOP_CTR_W-1:0] wg_uop_cnt; + reg [UOP_CTR_W-1:0] wg_compute_uop_cnt; always_comb begin case (ibuf_in.op_args.tcu.cd_nregs) - 2'd0: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR8); - 2'd1: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR16); - default: wg_uop_cnt = UOP_CTR_W'(WG_UOPS_NR32); + 2'd0: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR8); + 2'd1: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR16); + default: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_NR32); endcase + `ifdef VX_CFG_TCU_SPARSE_ENABLE + if (wg_is_sparse) begin + case (ibuf_in.op_args.tcu.cd_nregs) + 2'd0: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR8); + 2'd1: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR16); + default: wg_compute_uop_cnt = UOP_CTR_W'(WG_UOPS_SP_NR32); + endcase + wg_uop_cnt = wg_compute_uop_cnt + UOP_CTR_W'(1); + end else + `endif + wg_uop_cnt = wg_compute_uop_cnt + UOP_CTR_W'(1); end // K-outer index extraction: @@ -113,13 +137,11 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( // Sparse: ctr = n * m_steps + m (k always 0) // Since n_steps varies by cd_nregs, k-index bit position shifts. // m is always bit 0 (m_steps=2). n and k extracted via mux. - wire [`UP(CTR_W)-1:0] wg_idx_ctr = ctr; + wire is_wg_setup_uop = (ctr == '0); + wire [`UP(CTR_W)-1:0] wg_idx_ctr = is_wg_setup_uop ? '0 : (ctr - `UP(CTR_W)'(1)); wire [`UP(LG_M_WG)-1:0] wg_m_index = wg_idx_ctr[0 +: `UP(LG_M_WG)]; reg [`UP(LG_K_WG)-1:0] wg_k_index; reg [`UP(LG_N_WG_MAX)-1:0] wg_n_index; -`ifdef VX_CFG_TCU_SPARSE_ENABLE - wire wg_is_sparse = (ibuf_in.op_type == INST_TCU_WGMMA_SP); -`endif always_comb begin `ifdef VX_CFG_TCU_SPARSE_ENABLE if (wg_is_sparse) begin @@ -154,10 +176,9 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( localparam [4:0] wg_ra_base = TCU_WG_RA; // Register offsets for from-reg mode - // Dense A: rs1_off = m * k_steps + k (NRA=4 regs at ra_base: f24..f27) + // Dense A: rs1_off = m * 2 + k (NRA=4 regs at ra_base: f24..f27) // Sparse A: rs1_off = m (NRA compressed to 2 regs f24,f25; f26,f27 hold metadata) - localparam LG_WG_A_SB = $clog2(`UP(TCU_WG_A_SUB_BLOCKS)); - wire [`UP(CTR_W)-1:0] wg_rs1_reg_off_dense = ((`UP(CTR_W)'(wg_m_index) >> LG_WG_A_SB) << `UP(LG_K_WG)) | `UP(CTR_W)'(wg_k_index); + wire [`UP(CTR_W)-1:0] wg_rs1_reg_off_dense = (`UP(CTR_W)'(wg_m_index) << 1) | `UP(CTR_W)'(wg_k_index); `ifdef VX_CFG_TCU_SPARSE_ENABLE wire [`UP(CTR_W)-1:0] wg_rs1_reg_off = (wg_is_sparse && !wg_a_from_smem) ? `UP(CTR_W)'(wg_m_index) @@ -179,15 +200,11 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( ; `endif - // Sparse WGMMA (RS) expands to wg_uop_cnt/2 MMA uops; WMMA_SP expands + // Sparse WGMMA (RS) expands to max(1, k_steps/2) MMA uops plus setup; WMMA_SP expands // to TCU_UOPS (sym) or TCU_UOPS/2 (asym). assign uop_count = `ifdef VX_CFG_TCU_WGMMA_ENABLE - is_wgmma ? ( - `ifdef VX_CFG_TCU_SPARSE_ENABLE - wg_is_sparse ? (wg_uop_cnt >> 1) : - `endif - wg_uop_cnt) : + is_wgmma ? wg_uop_cnt : `endif `ifdef VX_CFG_TCU_SPARSE_ENABLE is_sparse ? (SYM_SPARSE ? UOP_CTR_W'(TCU_UOPS) : UOP_CTR_W'(TCU_UOPS / 2)) : @@ -312,25 +329,43 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( `endif `ifdef VX_CFG_TCU_WGMMA_ENABLE if (is_wgmma) begin - ibuf_r.op_args.tcu.step_m = 4'(wg_m_index); - ibuf_r.op_args.tcu.step_n = 4'(wg_n_index); - ibuf_r.op_args.tcu.step_k = 4'(wg_k_index); - ibuf_r.wb = 1'b1; + ibuf_r.op_args.tcu.step_m = is_wg_setup_uop ? 4'd0 : 4'(wg_m_index); + ibuf_r.op_args.tcu.step_n = is_wg_setup_uop ? 4'd0 : 4'(wg_n_index); + ibuf_r.op_args.tcu.step_k = is_wg_setup_uop ? 4'd0 : 4'(wg_k_index); + ibuf_r.wb = !is_wg_setup_uop; ibuf_r.rd = make_reg_num(REG_TYPE_F, TCU_WG_RC + wg_rs3_off); ibuf_r.rs3 = make_reg_num(REG_TYPE_F, TCU_WG_RC + wg_rs3_off); - ibuf_r.used_rs[2] = 1'b1; + ibuf_r.used_rs[2] = !is_wg_setup_uop; // Smem descriptors are invariant across the whole WGMMA expansion, - // so only fetch them on the first MMA uop. - if (wg_a_from_smem) begin + // so only fetch them on the setup uop. + if (is_wg_setup_uop) begin + ibuf_r.rs1 = make_reg_num(REG_TYPE_I, 5'd10); + ibuf_r.used_rs[0] = 1'b1; + end else if (wg_a_from_smem) begin ibuf_r.rs1 = make_reg_num(REG_TYPE_I, 5'd10); - ibuf_r.used_rs[0] = (wg_idx_ctr == '0); + ibuf_r.used_rs[0] = 1'b0; end else begin ibuf_r.rs1 = make_reg_num(REG_TYPE_F, wg_ra_base + 5'(wg_rs1_reg_off)); ibuf_r.used_rs[0] = 1'b1; end - // B source: always smem descriptor (x11), fetched on first MMA uop + // Setup reads desc_b from x11. Dense FEDP2K RS compute uses rs2 as the upper A half ibuf_r.rs2 = make_reg_num(REG_TYPE_I, 5'd11); - ibuf_r.used_rs[1] = (wg_idx_ctr == '0); + if (is_wg_setup_uop) begin + ibuf_r.used_rs[1] = 1'b1; + end else if (!wg_a_from_smem + `ifdef VX_CFG_TCU_FEDP2K + `ifdef VX_CFG_TCU_SPARSE_ENABLE + && !wg_is_sparse + `endif + `else + && 1'b0 + `endif + ) begin + ibuf_r.rs2 = make_reg_num(REG_TYPE_F, wg_ra_base + 5'(wg_rs1_reg_off) + 5'd1); + ibuf_r.used_rs[1] = 1'b1; + end else begin + ibuf_r.used_rs[1] = 1'b0; + end end else `endif begin @@ -370,12 +405,12 @@ module VX_tcu_uops import VX_tcu_pkg::*, VX_gpu_pkg::*; ( end `ifdef VX_CFG_TCU_WGMMA_ENABLE if (is_wgmma) begin - ibuf_r.fu_lock = (uop_idx == '0); + ibuf_r.fu_lock = is_wg_setup_uop; ibuf_r.fu_unlock = (uop_idx == (uop_count - UOP_CTR_W'(1))); // Expose first/last-uop predicates via op_args.tcu so downstream // consumers (bbuf, lockstep) read a single source of truth. - ibuf_r.op_args.tcu.is_first_uop = (uop_idx == '0); - ibuf_r.op_args.tcu.is_last_uop = (uop_idx == (uop_count - UOP_CTR_W'(1))); + ibuf_r.op_args.tcu.is_first_uop = !is_wg_setup_uop && (wg_idx_ctr == '0); + ibuf_r.op_args.tcu.is_last_uop = !is_wg_setup_uop && (uop_idx == (uop_count - UOP_CTR_W'(1))); end else `endif begin diff --git a/hw/rtl/tcu/VX_tcu_wgmma.sv b/hw/rtl/tcu/VX_tcu_wgmma.sv index 8b9d5fc9d..b8888f25f 100644 --- a/hw/rtl/tcu/VX_tcu_wgmma.sv +++ b/hw/rtl/tcu/VX_tcu_wgmma.sv @@ -46,7 +46,7 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( VX_mem_bus_if.master tcu_lmem_if, // Outputs to tcu_core (consumed by the wrapper). - output wire [BLOCK_SIZE-1:0][TCU_BLOCK_CAP-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, + output wire [BLOCK_SIZE-1:0][TCU_WG_A_DATA_SIZE-1:0][`VX_CFG_XLEN-1:0] tbuf_rs1_data, output wire [BLOCK_SIZE-1:0][TCU_WG_RS2_WIDTH-1:0][`VX_CFG_XLEN-1:0] tbuf_rs2_data, output wire [BLOCK_SIZE-1:0] tbuf_ready_eff ); @@ -62,19 +62,39 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( wire [BLOCK_SIZE-1:0] exec_fire_b_w; wire [BLOCK_SIZE-1:0] is_first_uop_b_w; wire [BLOCK_SIZE-1:0] is_last_uop_b_w; + wire [BLOCK_SIZE-1:0] is_setup_uop_b_w; + + logic [BLOCK_SIZE-1:0][`VX_CFG_XLEN-1:0] desc_a_r; + logic [BLOCK_SIZE-1:0][`VX_CFG_XLEN-1:0] desc_b_r; + for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_lockstep_inputs wire is_wgmma_op = (exec_data[bi].op_type == INST_TCU_WGMMA) `ifdef VX_CFG_TCU_SPARSE_ENABLE || (exec_data[bi].op_type == INST_TCU_WGMMA_SP) `endif ; - assign is_wgmma_b_w[bi] = exec_valid[bi] && is_wgmma_op; + assign is_setup_uop_b_w[bi] = is_wgmma_op && !exec_data[bi].header.wb; + assign is_wgmma_b_w[bi] = exec_valid[bi] && is_wgmma_op && !is_setup_uop_b_w[bi]; assign new_cta_b_w[bi] = exec_data[bi].header.cta_id; assign exec_fire_b_w[bi] = exec_valid[bi] && exec_ready[bi]; assign is_first_uop_b_w[bi] = exec_data[bi].op_args.tcu.is_first_uop; assign is_last_uop_b_w[bi] = exec_data[bi].op_args.tcu.is_last_uop; end + always_ff @(posedge clk) begin + if (reset) begin + desc_a_r <= '0; + desc_b_r <= '0; + end else begin + for (int bi = 0; bi < BLOCK_SIZE; ++bi) begin + if (exec_valid[bi] && exec_ready[bi] && is_setup_uop_b_w[bi]) begin + desc_a_r[bi] <= exec_data[bi].rs1_data[0]; + desc_b_r[bi] <= exec_data[bi].rs2_data[0]; + end + end + end + end + VX_tcu_lockstep #( .INSTANCE_ID (`SFORMATF(("%s-lockstep", INSTANCE_ID))), .BLOCK_SIZE (BLOCK_SIZE) @@ -102,18 +122,20 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( || (exec_data[bi].op_type == INST_TCU_WGMMA_SP) `endif ; - assign req[bi].valid = exec_valid[bi] && is_wgmma_b && !cta_conflict[bi]; + wire is_setup_uop = is_wgmma_b && !exec_data[bi].header.wb; + assign req[bi].valid = exec_valid[bi] && is_wgmma_b && !is_setup_uop && !cta_conflict[bi]; assign req[bi].uuid = exec_data[bi].header.uuid; assign req[bi].wid = exec_data[bi].header.wid; assign req[bi].step_m = exec_data[bi].op_args.tcu.step_m; assign req[bi].step_k = exec_data[bi].op_args.tcu.step_k; assign req[bi].step_n = exec_data[bi].op_args.tcu.step_n; assign req[bi].cd_nregs = exec_data[bi].op_args.tcu.cd_nregs; - assign req[bi].desc_a = exec_data[bi].rs1_data[0]; - assign req[bi].desc_b = exec_data[bi].rs2_data[0]; + assign req[bi].desc_a = desc_a_r[bi]; + assign req[bi].desc_b = desc_b_r[bi]; assign req[bi].a_is_smem = exec_data[bi].op_args.tcu.a_from_smem; assign req[bi].is_first_uop = exec_data[bi].op_args.tcu.is_first_uop; assign req[bi].is_last_uop = exec_data[bi].op_args.tcu.is_last_uop; + assign req[bi].setup_fire = exec_valid[bi] && exec_ready[bi] && is_setup_uop; `ifdef VX_CFG_TCU_SPARSE_ENABLE assign req[bi].is_sparse = (exec_data[bi].op_type == INST_TCU_WGMMA_SP); `endif @@ -148,7 +170,7 @@ module VX_tcu_wgmma import VX_gpu_pkg::*, VX_tcu_pkg::*; #( ); for (genvar bi = 0; bi < BLOCK_SIZE; ++bi) begin : g_tbuf_eff - assign tbuf_ready_eff[bi] = tbuf_ready[bi] && !cta_conflict[bi]; + assign tbuf_ready_eff[bi] = is_setup_uop_b_w[bi] ? 1'b1 : (tbuf_ready[bi] && !cta_conflict[bi]); end // ----------------------------------------------------------------------- diff --git a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv index c79303549..943c06638 100644 --- a/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv +++ b/hw/rtl/tcu/tfr/VX_tcu_tfr_mul_f8.sv @@ -189,20 +189,31 @@ module VX_tcu_tfr_mul_f8 import VX_tcu_pkg::*; wire [EXP_W-1:0] final_exp; `ifdef VX_CFG_TCU_MX_ENABLE + wire [EXP_W-1:0] bias_cpa_sum, bias_cpa_carry; wire [3*EXP_W-1:0] sf_comp = fmt_f[3] ? {EXP_W'(sf_a), EXP_W'(sf_b), -EXP_W'(254)} : (3*EXP_W)'(0); VX_csa_tree #( - .N(5), + .N(4), .W(EXP_W), .S(EXP_W) ) exp_sf_csa ( - .operands ({EXP_W'(max_pre_sum), EXP_W'(bias_sel), sf_comp}), - .sum (max_pre_sum_cpa), - .carry (bias_sel_cpa) + .operands ({EXP_W'(bias_sel), sf_comp}), + .sum (bias_cpa_sum), + .carry (bias_cpa_carry) + ); + VX_ks_adder #( + .N(EXP_W), + .BYPASS(`FORCE_BUILTIN_ADDER(EXP_W)) + ) exp_bias_add ( + .dataa(bias_cpa_sum), + .datab(bias_cpa_carry), + .cin(1'b0), + .sum(bias_sel_cpa), + `UNUSED_PIN(cout) ); `else - assign max_pre_sum_cpa = EXP_W'(max_pre_sum); assign bias_sel_cpa = EXP_W'(bias_sel); `endif + assign max_pre_sum_cpa = EXP_W'(max_pre_sum); VX_ks_adder #( .N(EXP_W), .BYPASS(`FORCE_BUILTIN_ADDER(EXP_W)) diff --git a/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv new file mode 100644 index 000000000..b463e5db1 --- /dev/null +++ b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_desc_uop.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,32,8.0,100,8.0,100,0,0,77,2,6,1,0,0,2,84,8,3,0,5,6336,0,0,4.02,9.39,59904,45568,168032,234768,0.716,PASS +baseline,4,RS,32,8,4,fp16,fp32,43,8.0,100,32.0,100,0,0,66,1,5,4,1,0,2,83,8,3,1,5,816,0,0,4.68,22.53,31232,25088,21568,36126,0.597,PASS +baseline,4,SS,8,8,4,fp16,fp32,32,8.0,100,8.0,100,0,0,73,2,6,0,0,0,3,84,7,3,0,5,6336,0,0,4.02,9.88,43520,45568,160864,224501,0.717,PASS +baseline,4,SS,32,8,4,fp16,fp32,44,8.0,100,32.0,100,0,0,69,1,5,4,1,0,4,83,7,3,1,5,816,0,0,4.71,27.09,23040,25088,20672,35326,0.585,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,31,8.0,100,8.0,100,0,0,71,3,5,2,1,0,1,79,12,4,1,4,6336,0,0,4.1,9.69,59904,45568,110944,155274,0.715,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,47,8.0,100,32.0,100,0,0,56,3,4,6,1,0,1,78,12,5,1,4,816,0,0,5.02,22.5,31232,25088,14208,25786,0.551,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,30,8.0,100,8.0,100,0,0,71,3,5,1,0,0,3,79,11,5,1,5,6336,0,0,4.03,9.86,43520,45568,100128,137848,0.726,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,48,7.9,99,32.0,100,0,0,53,1,3,6,1,0,4,78,11,5,1,4,816,0,0,5.1,27.13,23040,25088,13104,24392,0.537,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,31,8.0,100,8.0,100,0,0,73,3,5,1,1,0,2,79,12,4,1,4,6336,0,0,4.1,9.61,59904,45568,110944,155220,0.715,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,47,8.0,100,32.0,100,0,0,56,3,4,6,1,0,1,78,12,5,1,4,816,0,0,5.02,22.5,31232,25088,14208,25786,0.551,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,30,8.0,100,8.0,100,0,0,68,3,5,1,1,0,4,79,11,5,1,5,6336,0,0,4.03,10.05,43520,45568,100128,137867,0.726,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,48,8.0,99,32.0,100,0,0,61,1,3,6,1,0,4,78,11,5,1,4,816,0,0,5.1,27.26,23040,25088,13104,24291,0.539,PASS diff --git a/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv new file mode 100644 index 000000000..cd1cd4e52 --- /dev/null +++ b/perf/results/dse/wgmma/desc/sgemm_tcu_wg_fedp2k_tcu_meta_ld.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,32,8.0,100,7.9,99,0,0,69,1,5,1,0,0,2,82,8,5,0,5,7936,512,6,4.02,9.32,59904,46080,172960,243040,0.712,PASS +baseline,4,RS,32,8,4,fp16,fp32,43,7.9,99,31.5,99,0,0,62,1,5,4,1,0,2,81,8,5,1,5,1024,64,6,4.67,21.85,31232,25152,22240,37409,0.595,PASS +baseline,4,SS,8,8,4,fp16,fp32,33,8.0,100,7.8,98,0,0,70,1,6,0,0,0,3,82,7,5,0,6,7936,512,6,4.02,9.94,43520,46592,167392,239417,0.699,PASS +baseline,4,SS,32,8,4,fp16,fp32,43,8.0,99,31.3,98,0,0,68,1,5,4,1,0,4,81,8,5,1,5,1024,64,6,4.69,26.6,23040,25216,21552,36545,0.59,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,34,8.0,100,7.8,98,0,0,66,3,5,1,1,0,1,77,12,7,0,4,7936,512,6,4.1,9.6,59904,46080,115872,169531,0.683,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,47,8.0,100,31.3,98,0,0,54,2,3,5,1,0,1,76,12,7,1,4,1024,64,6,4.98,22.05,31232,25152,14880,27394,0.543,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,30,7.9,99,7.7,97,0,0,61,3,4,1,1,0,2,76,11,7,1,5,7936,512,6,4.03,9.48,43520,46592,106656,147193,0.725,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,48,7.9,99,31.0,97,0,0,50,2,3,5,1,0,4,75,12,8,1,5,1024,64,6,5.04,26.94,23040,25216,13984,26141,0.535,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,33,8.0,100,7.8,98,0,0,66,3,5,1,1,0,2,77,12,7,0,4,7936,512,6,4.11,9.7,59904,46080,115872,167617,0.691,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,47,8.0,100,31.3,98,0,0,54,2,3,5,1,0,2,76,12,7,1,4,1024,64,6,4.98,22.02,31232,25152,14880,27285,0.545,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,30,7.9,99,7.7,97,0,0,60,3,4,0,1,0,3,76,11,7,1,5,7936,512,6,4.03,9.52,43520,46592,106656,147023,0.725,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,48,7.9,99,31.0,97,0,0,51,2,3,5,1,0,5,75,12,8,1,5,1024,64,6,5.04,26.75,23040,25216,13984,26256,0.533,PASS diff --git a/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv new file mode 100644 index 000000000..9e740871f --- /dev/null +++ b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp16.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp16,fp32,60,8.0,99,8.0,100,0,0,57,0,1,9,1,0,1,71,9,16,0,4,13056,0,0,4.23,41.85,61440,36864,138592,335658,0.413,PASS +baseline,4,SS,8,8,4,fp16,fp32,61,8.0,100,8.0,100,0,0,57,0,0,9,1,0,1,70,9,17,0,4,13056,0,0,4.23,46.96,53248,36864,132384,331183,0.400,PASS +baseline,4,RS,32,8,4,fp16,fp32,90,7.9,99,32.0,100,0,0,66,0,0,17,3,0,0,69,10,16,1,4,1664,0,0,12.75,173.62,35328,23552,17760,169411,0.105,PASS +baseline,4,SS,32,8,4,fp16,fp32,90,7.9,99,32.0,100,0,0,66,0,0,18,3,0,0,68,10,17,1,4,1664,0,0,13.08,208.59,31232,23552,16992,168734,0.101,PASS +fedp2k,4,RS,8,8,4,fp16,fp32,65,8.0,100,8.0,100,0,0,59,1,0,8,2,0,0,69,10,17,0,4,13056,0,0,4.29,45.74,66048,37376,131104,360829,0.363,PASS +fedp2k,4,SS,8,8,4,fp16,fp32,68,8.0,100,8.0,100,0,0,64,0,0,9,4,0,1,68,10,18,0,4,13056,0,0,4.29,54.27,57856,37376,124960,382144,0.327,PASS +fedp2k,4,RS,32,8,4,fp16,fp32,90,7.9,98,32.0,100,0,0,64,0,0,16,4,0,0,68,11,17,1,4,1664,0,0,12.97,171.89,35328,23552,17376,169277,0.103,PASS +fedp2k,4,SS,32,8,4,fp16,fp32,91,7.8,98,32.0,100,0,0,67,0,0,18,2,0,0,67,10,18,1,4,1664,0,0,13.30,208.78,31232,23552,16608,170203,0.098,PASS +fedp2k,8,RS,8,8,4,fp16,fp32,64,7.9,99,8.0,100,0,0,60,1,0,9,2,0,1,69,10,17,0,4,13056,0,0,4.28,45.83,66048,37376,131104,356623,0.368,PASS +fedp2k,8,SS,8,8,4,fp16,fp32,68,8.0,100,8.0,100,0,0,64,0,0,9,4,0,1,68,10,18,0,4,13056,0,0,4.29,54.27,57856,37376,124960,382144,0.327,PASS +fedp2k,8,RS,32,8,4,fp16,fp32,90,7.9,98,32.0,100,0,0,64,0,0,16,4,0,0,68,11,17,1,4,1664,0,0,12.97,172.00,35328,23552,17376,169287,0.103,PASS +fedp2k,8,SS,32,8,4,fp16,fp32,91,7.9,98,32.0,100,0,0,67,0,0,17,2,0,0,67,10,18,1,4,1664,0,0,13.30,209.51,31232,23552,16608,170430,0.097,PASS diff --git a/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv new file mode 100644 index 000000000..72313478f --- /dev/null +++ b/perf/results/dse/wgmma/sp_rs_ss/sgemm_tcu_wg_sp_fedp2k_fp8.csv @@ -0,0 +1,13 @@ +model,fedp_latency,mode,threads,warps,issue_width,itype,otype,idle_pct,occupancy,occupancy_pct,simt_util,simt_util_pct,stall_fetch_pct,stall_ibuf_pct,stall_scrb_pct,stall_opds_pct,stall_alu_pct,stall_lsu_pct,stall_sfu_pct,stall_fpu_pct,stall_tcu_pct,mix_alu_pct,mix_lsu_pct,mix_sfu_pct,mix_fpu_pct,mix_tcu_pct,branches,divergent_branches,divergence_pct,ifetch_latency,load_latency,loads,stores,instrs,cycles,ipc,status +baseline,4,RS,8,8,4,fp8,fp32,59,7.9,99,8.0,100,0,0,56,1,0,9,1,0,0,71,11,15,1,3,8192,0,0,4.35,36.23,45056,34816,90336,214859,0.420,PASS +baseline,4,SS,8,8,4,fp8,fp32,59,7.9,99,8.0,100,0,0,56,1,0,9,1,0,1,71,11,15,1,3,8192,0,0,4.36,39.24,40960,34816,88032,211130,0.417,PASS +baseline,4,RS,32,8,4,fp8,fp32,88,7.7,96,32.0,100,0,0,60,0,0,16,1,0,0,69,12,15,1,3,1056,0,0,14.43,138.98,24064,21504,11840,94625,0.125,PASS +baseline,4,SS,32,8,4,fp8,fp32,88,7.7,97,32.0,100,0,0,58,0,0,16,1,0,0,69,12,15,1,3,1056,0,0,14.61,138.39,22016,21504,11552,95324,0.121,PASS +fedp2k,4,RS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,53,1,0,9,2,0,0,70,11,15,1,3,8192,0,0,4.36,36.50,45056,34816,87264,211279,0.413,PASS +fedp2k,4,SS,8,8,4,fp8,fp32,60,8.0,99,8.0,100,0,0,55,1,0,9,1,0,1,70,11,16,1,3,8192,0,0,4.40,39.41,40960,34816,84960,209156,0.406,PASS +fedp2k,4,RS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,61,0,0,17,1,0,0,69,12,15,1,3,1056,0,0,14.43,128.49,24064,21504,11664,94822,0.123,PASS +fedp2k,4,SS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,63,0,0,17,0,0,0,69,12,15,1,3,1056,0,0,14.74,145.60,22016,21504,11376,93752,0.121,PASS +fedp2k,8,RS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,52,1,0,9,2,0,1,70,11,15,1,3,8192,0,0,4.36,35.99,45056,34816,87264,212092,0.411,PASS +fedp2k,8,SS,8,8,4,fp8,fp32,60,7.9,99,8.0,100,0,0,54,1,0,9,1,0,1,70,11,16,1,3,8192,0,0,4.40,38.96,40960,34816,84960,206133,0.412,PASS +fedp2k,8,RS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,61,0,0,17,1,0,0,69,12,15,1,3,1056,0,0,14.43,128.49,24064,21504,11664,94866,0.123,PASS +fedp2k,8,SS,32,8,4,fp8,fp32,88,7.9,99,32.0,100,0,0,63,0,0,17,0,0,0,69,12,15,1,3,1056,0,0,14.74,147.17,22016,21504,11376,93702,0.121,PASS diff --git a/sim/rtlsim/Makefile b/sim/rtlsim/Makefile index 53487a8e3..eb61c0c0b 100644 --- a/sim/rtlsim/Makefile +++ b/sim/rtlsim/Makefile @@ -75,7 +75,16 @@ ifneq (,$(filter -DVX_CFG_EXT_TCU_ENABLE, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/dsp else ifneq (,$(filter -DVX_CFG_TCU_TYPE_BHF, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/bhf - RTL_INCLUDE += -J$(THIRD_PARTY_DIR)/hardfloat/source/RISCV + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/fNToRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/addRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/bsg_hardfloat_pkg.sv + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/HardFloat_primitives.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/HardFloat_rawFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/isSigNaNRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/mulRecFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/recFNToFN.v + RTL_PRE_SRCS += $(THIRD_PARTY_DIR)/hardfloat/source/recFNToRecFN.v + RTL_INCLUDE += -I$(THIRD_PARTY_DIR)/hardfloat/source/RISCV RTL_INCLUDE += -I$(THIRD_PARTY_DIR)/hardfloat/source else ifneq (,$(filter -DVX_CFG_TCU_TYPE_FPNEW, $(XCONFIGS))) RTL_INCLUDE += -I$(RTL_DIR)/tcu/fpnew @@ -137,6 +146,7 @@ VL_FLAGS += $(CONFIGS) VL_FLAGS += $(TCU_META_DCFG) VL_FLAGS += $(RTL_INCLUDE) VL_FLAGS += $(RTL_PKGS) +VL_FLAGS += $(RTL_PRE_SRCS) VL_FLAGS += --cc $(TOP) --top-module $(TOP) # Extract RTL directories from include directories diff --git a/sim/simx/decode.cpp b/sim/simx/decode.cpp index 7d2e7af56..70a825496 100644 --- a/sim/simx/decode.cpp +++ b/sim/simx/decode.cpp @@ -889,7 +889,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t fmt_d = rd, fmt_s = rs1; bool is_sparse = (rs2 & 1) != 0; instr->set_op_type(is_sparse ? TcuType::WMMA_SP : TcuType::WMMA); - instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, 0, 0, 0, 0, 0, 0}); instr->set_macro_op(); instr->set_wstall(true); } break; @@ -900,7 +900,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t cd_nregs = (rs2 >> 1) & 0x3; bool is_a_smem = (rs2 >> 3) & 1; instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); - instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, fmt_s, fmt_d, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, fmt_s, fmt_d, 0, 0, 0, 0, 0, 0}); instr->set_macro_op(); instr->set_wstall(true); } break; @@ -910,7 +910,7 @@ Instr::Ptr Decoder::decode(uint32_t code, uint64_t uuid) { uint32_t fmt_s = rs2; uint32_t slot = rd; instr->set_op_type(TcuType::TCU_LD); - instr->set_args(IntrTcuArgs{0, 0, fmt_s, slot, 0, 0, 0, 0, 0}); + instr->set_args(IntrTcuArgs{0, 0, fmt_s, slot, 0, 0, 0, 0, 0, 0}); // rs1 holds the warp-broadcast base address (real I-reg read). instr->set_src_reg(0, rs1, RegType::Integer); } break; diff --git a/sim/simx/tcu/tcu_unit.cpp b/sim/simx/tcu/tcu_unit.cpp index aa551dfd4..41811ef1f 100644 --- a/sim/simx/tcu/tcu_unit.cpp +++ b/sim/simx/tcu/tcu_unit.cpp @@ -32,6 +32,7 @@ using namespace vortex; namespace vt = vortex::tensor; using cfg = vt::wmma_config_t; using wg_cfg = vt::wgmma_config_t; +static constexpr uint32_t kFedpWords = wg_cfg::fedpK; inline uint64_t nan_box(uint32_t value) { return value | 0xffffffff00000000; @@ -148,11 +149,15 @@ template struct FEDP { using itype = typename It::dtype; static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { constexpr uint32_t i_ratio = sizeof(uint32_t) / sizeof(itype); static_assert(i_ratio * sizeof(itype) == sizeof(uint32_t), "FEDP: tcK * i_ratio must be <= 32"); if constexpr (std::is_same_v) { uint32_t acc = 0; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = reinterpret_cast(&a_row[z].u32); auto b = reinterpret_cast(&b_col[z].u32); uint32_t prod = 0; @@ -164,7 +169,7 @@ struct FEDP { return rv_fadd_s(c_val, acc, 0, nullptr); } else { uint32_t acc = c_val; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = reinterpret_cast(&a_row[z].u32); auto b = reinterpret_cast(&b_col[z].u32); for (uint32_t i = 0; i < i_ratio; ++i) { @@ -179,8 +184,12 @@ struct FEDP { template <> struct FEDP{ static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { auto acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = a_row[z].u32; auto b = b_col[z].u32; for (uint32_t i = 0; i < 8; ++i) { // 8 * 4 bits = 32 bits @@ -202,8 +211,12 @@ struct FEDP{ template <> struct FEDP{ static uint32_t eval(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val) { + return eval_n(a_row, b_col, c_val, cfg::tcK); + } + + static uint32_t eval_n(const reg_data_t *a_row, const reg_data_t *b_col, uint32_t c_val, uint32_t k_words) { auto acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < k_words; ++z) { auto a = a_row[z].u32; auto b = b_col[z].u32; for (uint32_t i = 0; i < 8; ++i) { // 8 * 4 bits = 32 bits @@ -216,91 +229,65 @@ struct FEDP{ } }; -using PFN_FEDP = uint32_t (*)(const reg_data_t*, const reg_data_t*, uint32_t); +using PFN_FEDP_N = uint32_t (*)(const reg_data_t*, const reg_data_t*, uint32_t, uint32_t); -static PFN_FEDP select_FEDP(uint32_t IT, uint32_t OT) { +static PFN_FEDP_N select_FEDP_N(uint32_t IT, uint32_t OT) { switch (OT) { case vt::fp32::id: switch (IT) { case vt::fp16::id: - return FEDP::eval; + return FEDP::eval_n; case vt::bf16::id: - return FEDP::eval; + return FEDP::eval_n; case vt::fp8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::bf8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::tf32::id: - return FEDP::eval; + return FEDP::eval_n; default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); + break; } break; case vt::fp16::id: - switch (IT) { - case vt::fp16::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::fp16::id) + return FEDP::eval_n; break; case vt::bf16::id: - switch (IT) { - case vt::bf16::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::bf16::id) + return FEDP::eval_n; break; case vt::fp8::id: - switch (IT) { - case vt::fp8::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::fp8::id) + return FEDP::eval_n; break; case vt::bf8::id: - switch (IT) { - case vt::bf8::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::bf8::id) + return FEDP::eval_n; break; case vt::tf32::id: - switch (IT) { - case vt::tf32::id: - return FEDP::eval; - default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); - } + if (IT == vt::tf32::id) + return FEDP::eval_n; break; case vt::int32::id: switch (IT) { case vt::int8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::uint8::id: - return FEDP::eval; + return FEDP::eval_n; case vt::int4::id: - return FEDP::eval; + return FEDP::eval_n; case vt::uint4::id: - return FEDP::eval; + return FEDP::eval_n; default: - std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; - std::abort(); + break; } break; default: - std::cout << "Error: unsupported output type: " << OT << "!" << std::endl; - std::abort(); + break; } + std::cout << "Error: unsupported mma format: " << IT << " -> " << OT << "!" << std::endl; + std::abort(); } // Format-agnostic sparse gather: for each bword, iterate over its elem_count packed @@ -355,6 +342,7 @@ class TcuUnit::Impl { exec_done_.fill(false); wgmma_planned_warps_.fill(0); in_wgmma_.fill(false); + wgmma_desc_.fill({0, 0}); } ~Impl() {} @@ -375,6 +363,8 @@ class TcuUnit::Impl { exec_done_.fill(false); wgmma_planned_warps_.fill(0); in_wgmma_.fill(false); + lmem_desc_.clear(); + wgmma_desc_.fill({0, 0}); cta_owner_a_.fill(-1); cta_owner_b_ = -1; cur_block_ = 0; @@ -391,28 +381,21 @@ class TcuUnit::Impl { if (input.empty()) continue; auto trace = input.peek(); if (!tcu_is_wgmma(std::get(trace->op_type))) continue; - wgmma_active |= (1u << b); uint32_t wid = trace->wid; uint64_t wid_bit = (uint64_t(1) << wid); - if (wgmma_planned_warps_.at(b) & wid_bit) continue; + int32_t new_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; auto& instr = *trace->instr_ptr; auto tpuArgs = std::get(instr.get_args()); - if (!(tpuArgs.step_m == 0 && tpuArgs.step_n == 0 && tpuArgs.step_k == 0)) { - // Non-first uop arrived without a prior plan: first uop already drained. - // Mark planned and continue (descriptors persist in lmem_desc_[wid]). - wgmma_planned_warps_.at(b) |= wid_bit; + + if (tpuArgs.is_setup_uop) continue; - } - auto& rs1_data = trace->src_data[0]; - auto& rs2_data = trace->src_data[1]; - uint32_t a_desc = rs1_data.empty() ? 0 : rs1_data.at(0).u32; - uint32_t b_desc = rs2_data.empty() ? 0 : rs2_data.at(0).u32; + + wgmma_active |= (1u << b); // CTA-overlap fence — defer this block's WGMMA if any other block // is mid-flight with a different CTA. The shared B buffer assumes // single-CTA occupancy across all blocks. - int32_t new_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; bool block_other_cta_inflight = false; for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { if (k == b) continue; @@ -426,6 +409,16 @@ class TcuUnit::Impl { continue; } + if (wgmma_planned_warps_.at(b) & wid_bit) continue; + if (!(tpuArgs.step_m == 0 && tpuArgs.step_n == 0 && tpuArgs.step_k == 0)) { + // Non-first uop arrived without a prior plan: first uop already drained. + // Mark planned and continue (descriptors persist in lmem_desc_[wid]). + wgmma_planned_warps_.at(b) |= wid_bit; + continue; + } + uint32_t a_desc = wgmma_desc_[wid][0]; + uint32_t b_desc = wgmma_desc_[wid][1]; + // Drop the shared B buffer only when no other block is mid-WGMMA — // otherwise we'd evict their resident bytes mid-flight. bool any_in_wgmma = false; @@ -476,19 +469,31 @@ class TcuUnit::Impl { continue; auto trace = input.peek(); auto tcu_type = std::get(trace->op_type); + auto tpuArgs = std::get(trace->instr_ptr->get_args()); - #ifdef VX_CFG_TCU_WGMMA_ENABLE + #ifdef VX_CFG_TCU_WGMMA_ENABLE // CTA-overlap fence deferred this block — skip until pass 1 plans it. - if (tcu_is_wgmma(tcu_type) && + if (tcu_is_wgmma(tcu_type) && !tpuArgs.is_setup_uop && !(wgmma_planned_warps_.at(b) & (uint64_t(1) << trace->wid))) continue; + if (tcu_is_wgmma(tcu_type) && !tpuArgs.is_setup_uop) { + int32_t this_cta = (int32_t)core_->scheduler().warp(trace->wid).cta_csrs.cta_id; + bool block_other_cta_inflight = false; + for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { + if (k == b) continue; + if (in_wgmma_.at(k) && cta_owner_a_.at(k) != this_cta) { + block_other_cta_inflight = true; + break; + } + } + if (block_other_cta_inflight) + continue; + } #endif // Execute once per trace; results persist across backpressure retries // via exec_done_[b]. if (!exec_done_.at(b)) { - auto& instr = *trace->instr_ptr; - auto tpuArgs = std::get(instr.get_args()); uint32_t wid = trace->wid; uint32_t num_threads = VX_CFG_NUM_THREADS; auto& rs1_data = trace->src_data[0]; @@ -513,7 +518,7 @@ class TcuUnit::Impl { cur_block_ = b; // CTA lockstep invariant: no block may execute a WGMMA uop for a // different cta_id while another block is mid-WGMMA. - { + if (!tpuArgs.is_setup_uop) { int32_t this_cta = (int32_t)core_->scheduler().warp(wid).cta_csrs.cta_id; for (uint32_t k = 0; k < VX_CFG_NUM_TCU_BLOCKS; ++k) { if (k == b) continue; @@ -528,9 +533,9 @@ class TcuUnit::Impl { } this->wgmma(wid, tpuArgs.fmt_s, tpuArgs.fmt_d, tpuArgs.step_m, tpuArgs.step_n, tpuArgs.step_k, - a_desc, b_desc, rs1_data, rs3_data, rd_data, + a_desc, b_desc, rs1_data, rs2_data, rs3_data, rd_data, tcu_is_sparse(tcu_type), - tpuArgs.cd_nregs, tpuArgs.is_a_smem); + tpuArgs.cd_nregs, tpuArgs.is_a_smem, tpuArgs.is_setup_uop); } break; #endif #ifdef VX_CFG_TCU_META_ENABLE @@ -552,7 +557,7 @@ class TcuUnit::Impl { case TcuType::WMMA_SP: case TcuType::WGMMA: case TcuType::WGMMA_SP: - delay = 4; + delay = tpuArgs.is_setup_uop ? 0 : VX_CFG_TCU_FEDP_DELAY; break; #ifdef VX_CFG_TCU_META_ENABLE case TcuType::TCU_LD: @@ -591,18 +596,16 @@ class TcuUnit::Impl { uint32_t fmt_s = args.fmt_s; bool is_a_smem = args.is_a_smem; uint32_t e_bits = elem_bits(fmt_s); - if (e_bits < 8) return; - uint32_t e_bytes = e_bits / 8; // NRC: cd_nregs 0/1/2 → 8/16/32; xtileN = NRC * NT / xtileM. uint32_t nrc = (args.cd_nregs == 0) ? 8 : (args.cd_nregs == 1) ? 16 : 32; uint32_t xtile_n = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; lmem_desc_t sd_a{}, sd_b{}; if (is_a_smem) { - sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) / e_bytes, false}; + sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) * 8 / e_bits, false}; lmem_desc_[wid][0] = sd_a; } - sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) / e_bytes, false}; + sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) * 8 / e_bits, false}; lmem_desc_[wid][1] = sd_b; // tileK = xtileK × ratio (ratio = 32/e_bits); sparse compresses K on A only. @@ -611,9 +614,11 @@ class TcuUnit::Impl { uint32_t a_k = is_sparse ? (tile_k / 2) : tile_k; // ldm==0 → block-major layout; ldm!=0 → row-major (stride in elements). - uint32_t k_blk_dim = cfg::tcK * ratio; - uint32_t a_blk_elems = cfg::tcM * k_blk_dim; - uint32_t b_blk_elems = k_blk_dim * cfg::tcN; + uint32_t fedp_words = kFedpWords; + uint32_t b_k_blk_dim = fedp_words * ratio; + uint32_t a_k_blk_dim = is_sparse ? (cfg::tcK * ratio) : b_k_blk_dim; + uint32_t a_blk_elems = cfg::tcM * a_k_blk_dim; + uint32_t b_blk_elems = b_k_blk_dim * cfg::tcN; uint32_t n_steps = xtile_n / cfg::tcN; auto& tbuf = simobject_->tbuf(); @@ -629,14 +634,14 @@ class TcuUnit::Impl { if (a_block_major) { uint32_t m_blk = r / cfg::tcM; uint32_t i_in = r % cfg::tcM; - uint32_t k_blk = c / k_blk_dim; - uint32_t k_in = c % k_blk_dim; + uint32_t k_blk = c / a_k_blk_dim; + uint32_t k_in = c % a_k_blk_dim; elem_off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems - + i_in * k_blk_dim + k_in; + + i_in * a_k_blk_dim + k_in; } else { elem_off = uint64_t(r) * sd_a.ldm + c; } - uint64_t addr = sd_a.base + elem_off * e_bytes; + uint64_t addr = sd_a.base + elem_off * e_bits / 8; a_lines.push_back(addr & ~uint64_t(VX_CFG_MEM_BLOCK_SIZE - 1)); } } @@ -654,17 +659,17 @@ class TcuUnit::Impl { for (uint32_t c = 0; c < xtile_n; ++c) { uint64_t elem_off; if (b_block_major) { - uint32_t k_blk = r / k_blk_dim; - uint32_t r_in = r % k_blk_dim; + uint32_t k_blk = r / b_k_blk_dim; + uint32_t r_in = r % b_k_blk_dim; uint32_t n_blk = c / cfg::tcN; uint32_t n_in = c % cfg::tcN; // Within-block layout: N outer, K inner. elem_off = (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * k_blk_dim + r_in; + + n_in * b_k_blk_dim + r_in; } else { elem_off = uint64_t(c) * sd_b.ldm + r; } - uint64_t addr = sd_b.base + elem_off * e_bytes; + uint64_t addr = sd_b.base + elem_off * e_bits / 8; b_lines.push_back(addr & ~uint64_t(VX_CFG_MEM_BLOCK_SIZE - 1)); } } @@ -763,16 +768,15 @@ class TcuUnit::Impl { ? (step_n % cfg::b_sub_blocks_sp) * cfg::b_block_size_sp : (step_n % cfg::b_sub_blocks) * cfg::b_block_size; - // Prepare A tile [tcM][tcK] - reg_data_t a_tile[cfg::tcM * cfg::tcK]; + // WMMA occupies the low half of a widened FEDP and zeros the unused half. + reg_data_t a_tile[cfg::tcM * kFedpWords] = {}; for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t z = 0; z < cfg::tcK; ++z) { - a_tile[i * cfg::tcK + z] = rs1_data.at(a_off + i * cfg::tcK + z); + a_tile[i * kFedpWords + z] = rs1_data.at(a_off + i * cfg::tcK + z); } } - // Prepare B tile [tcM][tcN][tcK] - reg_data_t b_tile[cfg::tcM * cfg::tcN * cfg::tcK]; + reg_data_t b_tile[cfg::tcM * cfg::tcN * kFedpWords] = {}; if (is_sparse) { constexpr uint32_t kCompression = 2; uint32_t ebits = elem_bits(fmt_s); @@ -792,7 +796,7 @@ class TcuUnit::Impl { lo |= meta_bit(row_base + meta_bits * z + b) << b; hi |= meta_bit(row_base + meta_bits * (cfg::tcK + z) + b) << b; } - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = + b_tile[(i * cfg::tcN + j) * kFedpWords + z].u32 = gather_sparse(rs2_data.at(b_idx).u32, rs2_data.at(b_idx + 1).u32, lo, hi, ebits); } } @@ -801,14 +805,16 @@ class TcuUnit::Impl { for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t j = 0; j < cfg::tcN; ++j) { for (uint32_t z = 0; z < cfg::tcK; ++z) { - b_tile[(i * cfg::tcN + j) * cfg::tcK + z] = rs2_data.at(b_off + j * cfg::tcK + z); + b_tile[(i * cfg::tcN + j) * kFedpWords + z] = rs2_data.at(b_off + j * cfg::tcK + z); } } } } fedp_tile(wid, step_m, step_n, step_k, fmt_s, fmt_d, - a_tile, b_tile, rs3_data, rd_data, is_sparse, true); + a_tile, b_tile, rs3_data, rd_data, + is_sparse, kFedpWords, cfg::tcK, + cfg::k_steps * cfg::tcK); } void wgmma(uint32_t wid, @@ -820,12 +826,13 @@ class TcuUnit::Impl { uint32_t a_desc, uint32_t b_desc, const std::vector& rs1_data, + const std::vector& rs2_data, const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, uint32_t cd_nregs, - uint32_t is_a_smem) { - __unused(cd_nregs); + uint32_t is_a_smem, + uint32_t is_setup_uop) { if (is_sparse && !vt::sparse_format_supported(fmt_s)) { std::cout << "Error: WGMMA_SP unsupported input format: " << vt::fmt_string(fmt_s) << " (id=" << fmt_s << ")" << std::endl; @@ -833,46 +840,51 @@ class TcuUnit::Impl { } uint32_t ratio = elem_ratio(fmt_s); - uint32_t k_words = cfg::tcK; - uint32_t e_bytes = elem_bits(fmt_s) / 8; + uint32_t k_words = is_sparse ? cfg::tcK : kFedpWords; + uint32_t e_bits = elem_bits(fmt_s); // Decode smem descriptors (B always from smem, A optionally). - lmem_desc_t sd_a, sd_b; - if (step_k == 0 && step_m == 0 && step_n == 0) { + if (is_setup_uop) { + wgmma_desc_[wid][0] = a_desc; + wgmma_desc_[wid][1] = b_desc; if (is_a_smem) { - sd_a = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) / e_bytes, false}; - lmem_desc_[wid][0] = sd_a; + lmem_desc_[wid][0] = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (a_desc & 0xFFFF), (a_desc >> 16) * 8 / e_bits, false}; } - sd_b = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) / e_bytes, false}; - lmem_desc_[wid][1] = sd_b; - } else { - sd_a = lmem_desc_[wid][0]; - sd_b = lmem_desc_[wid][1]; + lmem_desc_[wid][1] = {uint64_t(VX_MEM_LMEM_BASE_ADDR) + (b_desc & 0xFFFF), (b_desc >> 16) * 8 / e_bits, false}; + uint32_t nrc = (cd_nregs == 0) ? 8 : (cd_nregs == 1) ? 16 : 32; + cur_xtile_n_ = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; + return; } + + lmem_desc_t sd_a = lmem_desc_[wid][0]; + lmem_desc_t sd_b = lmem_desc_[wid][1]; // load_lmem_word distinguishes A from B by descriptor base. cur_a_desc_base_ = is_a_smem ? sd_a.base : ~uint64_t(0); + cur_is_sparse_ = is_sparse; // NRC: cd_nregs 0/1/2 → 8/16/32; xtileN = NRC * NT / xtileM. { uint32_t nrc = (cd_nregs == 0) ? 8 : (cd_nregs == 1) ? 16 : 32; cur_xtile_n_ = (nrc * VX_CFG_NUM_THREADS) / wg_cfg::xtileM; } - // Prepare A tile [tcM][tcK] - reg_data_t a_tile[cfg::tcM * cfg::tcK]; + reg_data_t a_tile[cfg::tcM * kFedpWords] = {}; for (uint32_t i = 0; i < cfg::tcM; ++i) { uint32_t a_row_idx = step_m * cfg::tcM + i; for (uint32_t z = 0; z < k_words; ++z) { if (is_a_smem) { uint32_t k_elem = (step_k * k_words + z) * ratio; - a_tile[i * cfg::tcK + z].u32 = load_lmem_word(sd_a, a_row_idx, k_elem, fmt_s, false); + a_tile[i * k_words + z].u32 = load_lmem_word(sd_a, a_row_idx, k_elem, fmt_s, false); } else { - a_tile[i * cfg::tcK + z] = rs1_data.at(i * cfg::tcK + z); + if (z < cfg::tcK) { + a_tile[i * k_words + z] = rs1_data.at(i * cfg::tcK + z); + } else { + a_tile[i * k_words + z] = rs2_data.at(i * cfg::tcK + (z - cfg::tcK)); + } } } } - // Prepare B tile [tcM][tcN][tcK] - reg_data_t b_tile[cfg::tcM * cfg::tcN * cfg::tcK]; + reg_data_t b_tile[cfg::tcM * cfg::tcN * kFedpWords] = {}; if (is_sparse) { uint32_t ebits = elem_bits(fmt_s); uint32_t rtl_i_ratio = 32 / ebits; @@ -907,12 +919,12 @@ class TcuUnit::Impl { uint32_t bword0 = load_lmem_word(sd_b, k_elem_b0, b_col_idx, fmt_s, true); uint32_t bword1 = load_lmem_word(sd_b, k_elem_b0 + ratio, b_col_idx, fmt_s, true); uint32_t gathered = gather_sparse(bword0, bword1, lo, hi, ebits); - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = gathered; + b_tile[(i * cfg::tcN + j) * k_words + z].u32 = gathered; // Trace: B-gather (CSV: wid,step_m,step_n,i,lane,bword0,bword1,lo,hi,gathered). if (const char* p = std::getenv("VORTEX_TCU_TRACE")) { if (p[0] == '1') { fprintf(stderr, "GATHER,%u,%u,%u,%u,%u,0x%08x,0x%08x,%u,%u,0x%08x\n", - wid, step_m, step_n, i, j*cfg::tcK+z, + wid, step_m, step_n, i, j*k_words+z, bword0, bword1, lo, hi, gathered); } } @@ -925,7 +937,7 @@ class TcuUnit::Impl { uint32_t b_col_idx = step_n * cfg::tcN + j; for (uint32_t z = 0; z < k_words; ++z) { uint32_t k_elem = (step_k * k_words + z) * ratio; - b_tile[(i * cfg::tcN + j) * cfg::tcK + z].u32 = + b_tile[(i * cfg::tcN + j) * k_words + z].u32 = load_lmem_word(sd_b, k_elem, b_col_idx, fmt_s, true); } } @@ -933,7 +945,9 @@ class TcuUnit::Impl { } fedp_tile(wid, step_m, step_n, step_k, fmt_s, fmt_d, - a_tile, b_tile, rs3_data, rd_data, is_sparse, false); + a_tile, b_tile, rs3_data, rd_data, + is_sparse, k_words, k_words, + wg_cfg::k_steps * wg_cfg::fedpK); __unused(b_desc); } @@ -975,10 +989,10 @@ class TcuUnit::Impl { return 32 / elem_bits(fmt_s); } - uint32_t mx_tile_scale_blocks(uint32_t fmt_s) const { - uint32_t logical_tile_k = cfg::k_steps * cfg::tcK * elem_ratio(fmt_s); + uint32_t mx_tile_scale_blocks(uint32_t fmt_s, uint32_t tile_k_words) const { + uint32_t logical_tile_k = tile_k_words * elem_ratio(fmt_s); uint32_t block_size = vt::mx_scale_block_size(fmt_s); - return std::max(1u, logical_tile_k / block_size); + return std::max(1u, (logical_tile_k + block_size - 1) / block_size); } static uint8_t meta_byte(const std::vector& words, uint32_t index) { @@ -997,13 +1011,15 @@ class TcuUnit::Impl { uint32_t eval_mx_fedp(uint32_t wid, uint32_t fmt_s, uint32_t fmt_d, uint32_t step_m, uint32_t step_n, uint32_t step_k, uint32_t i, uint32_t j, const reg_data_t* a_row, - const reg_data_t* b_col, uint32_t c_val, bool is_sparse) const { + const reg_data_t* b_col, uint32_t c_val, bool is_sparse, + uint32_t fedp_words, uint32_t tile_k_words) const { #ifndef VX_CFG_TCU_MX_ENABLE - __unused(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, is_sparse); + __unused(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, + is_sparse, fedp_words, tile_k_words); return c_val; #else uint32_t ratio = elem_ratio(fmt_s); - uint32_t scale_blocks_k = mx_tile_scale_blocks(fmt_s); + uint32_t scale_blocks_k = mx_tile_scale_blocks(fmt_s, tile_k_words); uint32_t block_size = vt::mx_scale_block_size(fmt_s); auto scale_a = [&](uint32_t elem_k) { uint32_t row = step_m * cfg::tcM + i; @@ -1016,10 +1032,10 @@ class TcuUnit::Impl { if (fmt_d == vt::fp32::id) { uint32_t acc = c_val; - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < fedp_words; ++z) { for (uint32_t e = 0; e < ratio; ++e) { uint32_t sparse_ratio = is_sparse ? 2 : 1; - uint32_t elem_k = ((step_k * cfg::tcK + z) * ratio + e) * sparse_ratio; + uint32_t elem_k = ((step_k * fedp_words + z) * ratio + e) * sparse_ratio; uint32_t xa, xb; if (fmt_s == vt::mxfp8::id) { xa = rv_mxfp8tof_s((a_row[z].u32 >> (8 * e)) & 0xff, scale_a(elem_k), 0, nullptr); @@ -1044,9 +1060,9 @@ class TcuUnit::Impl { if (fmt_s == vt::mxint8::id && fmt_d == vt::int32::id) { int32_t acc = bit_cast(c_val); - for (uint32_t z = 0; z < cfg::tcK; ++z) { + for (uint32_t z = 0; z < fedp_words; ++z) { for (uint32_t e = 0; e < ratio; ++e) { - uint32_t elem_k = ((step_k * cfg::tcK + z) * ratio + e) * (is_sparse ? 2 : 1); + uint32_t elem_k = ((step_k * fedp_words + z) * ratio + e) * (is_sparse ? 2 : 1); auto a = static_cast((a_row[z].u32 >> (8 * e)) & 0xff); auto b = static_cast((b_col[z].u32 >> (8 * e)) & 0xff); int32_t shift = int32_t(scale_a(elem_k)) + int32_t(scale_b(elem_k)) - 266; @@ -1069,19 +1085,21 @@ class TcuUnit::Impl { template uint32_t gather_word(ReadLine read_line, const lmem_desc_t& desc, uint32_t row, uint32_t col, - uint32_t fmt_s, bool pack_along_row) const { + uint32_t fmt_s, bool pack_along_row, + bool sparse_a_layout = false) const { uint32_t e_bits = elem_bits(fmt_s); uint32_t ratio = (e_bits >= 32) ? 1 : (32 / e_bits); - uint32_t e_bytes = (e_bits >= 8) ? (e_bits / 8) : 1; uint32_t result = 0; for (uint32_t r = 0; r < ratio; ++r) { uint32_t cur_row = pack_along_row ? (row + r) : row; uint32_t cur_col = pack_along_row ? col : (col + r); - uint64_t byte_addr; + uint64_t elem_off; if (desc.ldm == 0) { // Block-major SMEM. K dimension is along col for A (pack_along_row // false) and along row for B (pack_along_row true). - uint32_t k_blk_dim = cfg::tcK * ratio; + uint32_t k_blk_dim = (sparse_a_layout && !pack_along_row) + ? (cfg::tcK * ratio) + : (kFedpWords * ratio); if (pack_along_row) { // B: r is K coord, c is N coord. Within-block layout: N outer, K inner. uint32_t k_blk = cur_row / k_blk_dim; @@ -1090,9 +1108,8 @@ class TcuUnit::Impl { uint32_t n_in = cur_col % cfg::tcN; uint32_t b_blk_elems = k_blk_dim * cfg::tcN; uint32_t n_steps = cur_xtile_n_ / cfg::tcN; - uint64_t off = (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * k_blk_dim + r_in; - byte_addr = desc.base + off * e_bytes; + elem_off = (k_blk * n_steps + n_blk) * b_blk_elems + + n_in * k_blk_dim + r_in; } else { // A: r is M coord, c is K coord. uint32_t m_blk = cur_row / cfg::tcM; @@ -1100,20 +1117,20 @@ class TcuUnit::Impl { uint32_t k_blk = cur_col / k_blk_dim; uint32_t k_in = cur_col % k_blk_dim; uint32_t a_blk_elems = cfg::tcM * k_blk_dim; - uint64_t off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems - + i_in * k_blk_dim + k_in; - byte_addr = desc.base + off * e_bytes; + elem_off = (k_blk * wg_cfg::m_steps + m_blk) * a_blk_elems + + i_in * k_blk_dim + k_in; } } else if (desc.col_major) { - byte_addr = desc.base + (uint64_t(cur_col) * desc.ldm + cur_row) * e_bytes; + elem_off = uint64_t(cur_col) * desc.ldm + cur_row; } else if (pack_along_row) { // B: K-major (N-outer K-inner). cur_row=K, cur_col=N; // ldm = stride in elements between N rows. - byte_addr = desc.base + (uint64_t(cur_col) * desc.ldm + cur_row) * e_bytes; + elem_off = uint64_t(cur_col) * desc.ldm + cur_row; } else { // A: row-major (M-outer K-inner). cur_row=M, cur_col=K. - byte_addr = desc.base + (uint64_t(cur_row) * desc.ldm + cur_col) * e_bytes; + elem_off = uint64_t(cur_row) * desc.ldm + cur_col; } + uint64_t byte_addr = desc.base + elem_off * e_bits / 8; auto line = read_line(byte_addr); if (!line) { std::cout << "Error: TCU buffer miss at 0x" << std::hex << byte_addr @@ -1132,9 +1149,10 @@ class TcuUnit::Impl { result |= (val & 0xFFFF) << (r * 16); } else if (e_bits == 8) { result |= uint32_t((*line)[off]) << (r * 8); + } else if (e_bits == 4) { + uint32_t val = (uint32_t((*line)[off]) >> (4 * (elem_off & 1))) & 0xf; + result |= val << (r * 4); } else { - // 4-bit (int4/uint4) not supported. - std::cout << "Error: TCU 4-bit operand gather not supported" << std::endl; std::abort(); } } @@ -1148,7 +1166,7 @@ class TcuUnit::Impl { if (desc.base == cur_a_desc_base_) { uint32_t b = cur_block_; return gather_word([&](uint64_t addr) { return tbuf->read_a(b, addr); }, - desc, row, col, fmt_s, pack_along_row); + desc, row, col, fmt_s, pack_along_row, cur_is_sparse_); } return gather_word([&](uint64_t addr) { return tbuf->read_b(addr); }, desc, row, col, fmt_s, pack_along_row); @@ -1159,8 +1177,8 @@ class TcuUnit::Impl { static constexpr uint32_t kMaxMetaCols = VX_CFG_NUM_THREADS / 2; // FEDP tile computation for both WMMA and WGMMA. - // a_tile: [tcM][tcK] flat array of pre-loaded A operand words. - // b_tile: [tcM][tcN][tcK] flat array of pre-loaded B operand words + // a_tile: [tcM][k_words] flat array of pre-loaded A operand words. + // b_tile: [tcM][tcN][k_words] flat array of pre-loaded B operand words // (for dense, each i-slice is identical; for sparse, already gathered). void fedp_tile(uint32_t wid, uint32_t step_m, uint32_t step_n, uint32_t step_k, @@ -1170,23 +1188,23 @@ class TcuUnit::Impl { const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, - bool allow_mx) { - if (!allow_mx && vt::mx_scale_format(fmt_s)) { - std::cout << "Error: MX formats are supported only by WMMA." << std::endl; - std::abort(); - } - PFN_FEDP fedp = vt::mx_scale_format(fmt_s) ? nullptr : select_FEDP(fmt_s, fmt_d); + uint32_t k_words = cfg::tcK, + uint32_t fedp_words = cfg::tcK, + uint32_t tile_k_words = cfg::k_steps * cfg::tcK) { + PFN_FEDP_N fedp = vt::mx_scale_format(fmt_s) ? nullptr : select_FEDP_N(fmt_s, fmt_d); for (uint32_t i = 0; i < cfg::tcM; ++i) { for (uint32_t j = 0; j < cfg::tcN; ++j) { auto t = i * cfg::tcN + j; auto c_val = rs3_data.at(t).u32; - auto a_row = &a_tile[i * cfg::tcK]; - auto b_col = &b_tile[(i * cfg::tcN + j) * cfg::tcK]; + auto a_row = &a_tile[i * k_words]; + auto b_col = &b_tile[(i * cfg::tcN + j) * k_words]; - uint32_t d_val = (allow_mx && vt::mx_scale_format(fmt_s)) - ? eval_mx_fedp(wid, fmt_s, fmt_d, step_m, step_n, step_k, i, j, a_row, b_col, c_val, is_sparse) - : fedp(a_row, b_col, c_val); + uint32_t d_val = vt::mx_scale_format(fmt_s) + ? eval_mx_fedp(wid, fmt_s, fmt_d, step_m, step_n, step_k, + i, j, a_row, b_col, c_val, is_sparse, + fedp_words, tile_k_words) + : fedp(a_row, b_col, c_val, k_words); rd_data.at(t).u64 = nan_box(d_val); DTH(3, simobject_->name() << " FEDP" @@ -1207,6 +1225,7 @@ class TcuUnit::Impl { std::vector> mx_meta_b_; #endif std::unordered_map lmem_desc_; + std::array, VX_CFG_NUM_WARPS> wgmma_desc_; mutable PerfStats perf_stats_; // Per-block guard: execute already happened for this trace; reset on pop(). std::array exec_done_; @@ -1218,6 +1237,7 @@ class TcuUnit::Impl { uint32_t cur_block_ = 0; // A-descriptor base for the current wgmma(); distinguishes A from B in load_lmem_word. uint64_t cur_a_desc_base_ = ~uint64_t(0); + bool cur_is_sparse_ = false; // xtileN for the active WGMMA (derived from NRC). uint32_t cur_xtile_n_ = 8; // CTA owner per block's A buffer and the shared B buffer (-1 = unowned). @@ -1277,9 +1297,9 @@ uint32_t TcuUopGen::uop_count(const Instr& instr) { if (tcu_is_wgmma(tcu_type)) { bool is_sparse = tcu_is_sparse(tcu_type); uint32_t nrc = (args.cd_nregs == 0) ? 8 : (args.cd_nregs == 1) ? 16 : 32; - uint32_t k_count = is_sparse ? (wg_cfg::k_steps / 2) : wg_cfg::k_steps; + uint32_t k_count = is_sparse ? std::max(1u, wg_cfg::k_steps / 2) : wg_cfg::k_steps; uint32_t mma_uops = k_count * nrc; - return mma_uops; + return mma_uops + 1; } #endif @@ -1337,7 +1357,7 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { uint32_t actual_n = lg_k ? (n_sp >> lg_k) : n_sp; uint32_t reg_rs3 = rc_base + (mma_idx >> 1); uop_instr->set_op_type(TcuType::WMMA_SP); - uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m_sp, actual_n, 0, 0, 0}); + uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m_sp, actual_n, 0, 0, 0, 0}); uop_instr->set_dest_reg(reg_rs3, RegType::Float); uop_instr->set_src_reg(0, ra_base + m_sp, RegType::Float); uop_instr->set_src_reg(1, rb_base + n_sp, RegType::Float); @@ -1393,7 +1413,7 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { reg_rs3 = rc_base + c_off; } uop_instr->set_op_type(is_sparse ? TcuType::WMMA_SP : TcuType::WMMA); - uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m, n, k, 0, 0}); + uop_instr->set_args(IntrTcuArgs{0, 0, fmt_s, fmt_d, m, n, k, 0, 0, 0}); uop_instr->set_dest_reg(reg_rs3, RegType::Float); uop_instr->set_src_reg(0, reg_rs1, RegType::Float); uop_instr->set_src_reg(1, reg_rs2, RegType::Float); @@ -1405,52 +1425,62 @@ Instr::Ptr TcuUopGen::get(const Instr& macro_instr, uint32_t uop_index) { else if (tcu_is_wgmma(tcu_type)) { constexpr uint32_t m_steps = wg_cfg::m_steps; constexpr uint32_t k_steps = wg_cfg::k_steps; + constexpr uint32_t a_reg_k_steps = 2; uint32_t fmt_s = args.fmt_s; uint32_t fmt_d = args.fmt_d; bool is_sparse = tcu_is_sparse(tcu_type); bool is_a_smem = args.is_a_smem; uint32_t cd_nregs = args.cd_nregs; - uint32_t k_count = is_sparse ? (k_steps / 2) : k_steps; + uint32_t k_count = is_sparse ? std::max(1u, k_steps / 2) : k_steps; constexpr uint32_t a0 = 10, a1 = 11; { + uop_instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); + + if (uop_index == 0) { + uop_instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, + fmt_s, fmt_d, 0, 0, 0, 1, 0, 1}); + uop_instr->set_src_reg(0, a0, RegType::Integer); + uop_instr->set_src_reg(1, a1, RegType::Integer); + uop_instr->set_fu_lock(true); + uop_instr->set_fu_unlock(false); + return uop_instr; + } + // MMA phase - uint32_t mma_idx = uop_index; + uint32_t mma_idx = uop_index - 1; uint32_t ra_base = is_a_smem ? 10 : 24; // Loop order: m (inner) -> n (middle) -> k (outer). K-outer maximizes // per-block A-buffer reuse: each A_w[m,k] is consumed across the entire // (n,m) inner sweep, and each shared B[k,n] is consumed for m_steps // consecutive uops. - uint32_t mn = total / k_count; + uint32_t compute_total = total - 1; + uint32_t mn = compute_total / k_count; uint32_t k = mma_idx / mn; uint32_t rem = mma_idx % mn; uint32_t n = rem / m_steps; uint32_t m = rem % m_steps; uint32_t r = n * m_steps + m; - uop_instr->set_op_type(is_sparse ? TcuType::WGMMA_SP : TcuType::WGMMA); - bool first = (uop_index == 0); + bool first = (mma_idx == 0); bool last = (uop_index == (total - 1)); uop_instr->set_args(IntrTcuArgs{is_a_smem ? 1u : 0u, cd_nregs, - fmt_s, fmt_d, m, n, k, first ? 1u : 0u, last ? 1u : 0u}); + fmt_s, fmt_d, m, n, k, first ? 1u : 0u, last ? 1u : 0u, 0}); uop_instr->set_dest_reg(r, RegType::Float); - if (mma_idx == 0) { - if (is_a_smem) { - uop_instr->set_src_reg(0, a0, RegType::Integer); - } else { - uint32_t rs1_off = m * k_count + k; - uop_instr->set_src_reg(0, ra_base + rs1_off, RegType::Float); - } - uop_instr->set_src_reg(1, a1, RegType::Integer); - } else if (!is_a_smem) { - uint32_t rs1_off = m * k_count + k; + if (!is_a_smem) { + uint32_t rs1_off = is_sparse ? m : (m * a_reg_k_steps + k); uop_instr->set_src_reg(0, ra_base + rs1_off, RegType::Float); + if constexpr (k_steps == 1) { + if (!is_sparse) { + uop_instr->set_src_reg(1, ra_base + rs1_off + 1, RegType::Float); + } + } } uop_instr->set_src_reg(2, r, RegType::Float); } // fu_lock on first uop, fu_unlock on last uop - uop_instr->set_fu_lock(uop_index == 0); + uop_instr->set_fu_lock(false); uop_instr->set_fu_unlock(uop_index == (total - 1)); } #endif @@ -1507,20 +1537,22 @@ void TcuUnit::wmma(uint32_t wid, } void TcuUnit::wgmma(uint32_t wid, - uint32_t fmt_s, - uint32_t fmt_d, - uint32_t step_m, - uint32_t step_n, - uint32_t step_k, - uint32_t a_desc, - uint32_t b_desc, - const std::vector& rs1_data, - const std::vector& rs3_data, - std::vector& rd_data, - bool is_sparse, - uint32_t cd_nregs, - uint32_t is_a_smem) { + uint32_t fmt_s, + uint32_t fmt_d, + uint32_t step_m, + uint32_t step_n, + uint32_t step_k, + uint32_t a_desc, + uint32_t b_desc, + const std::vector& rs1_data, + const std::vector& rs2_data, + const std::vector& rs3_data, + std::vector& rd_data, + bool is_sparse, + uint32_t cd_nregs, + uint32_t is_a_smem, + uint32_t is_setup_uop) { impl_->wgmma(wid, fmt_s, fmt_d, step_m, step_n, step_k, a_desc, b_desc, - rs1_data, rs3_data, rd_data, - is_sparse, cd_nregs, is_a_smem); + rs1_data, rs2_data, rs3_data, rd_data, + is_sparse, cd_nregs, is_a_smem, is_setup_uop); } diff --git a/sim/simx/tcu/tcu_unit.h b/sim/simx/tcu/tcu_unit.h index 8b8dacad7..2728e977b 100644 --- a/sim/simx/tcu/tcu_unit.h +++ b/sim/simx/tcu/tcu_unit.h @@ -90,11 +90,13 @@ class TcuUnit : public FuncUnit { uint32_t a_desc, uint32_t b_desc, const std::vector& rs1_data, + const std::vector& rs2_data, const std::vector& rs3_data, std::vector& rd_data, bool is_sparse, uint32_t cd_nregs, - uint32_t is_a_smem); + uint32_t is_a_smem, + uint32_t is_setup_uop); // Tile-buffer subsystem (owns abuf×Q + bbuf + LMEM arb). // Exposed so that `Core` can bind its single LMEM port pair. diff --git a/sim/simx/types.h b/sim/simx/types.h index 37904edfd..516d91278 100644 --- a/sim/simx/types.h +++ b/sim/simx/types.h @@ -711,6 +711,7 @@ struct IntrTcuArgs { uint32_t step_k : 4; uint32_t is_first_uop : 1; // set per-uop by tcu_uops expansion (C4) uint32_t is_last_uop : 1; + uint32_t is_setup_uop : 1; // WGMMA descriptor setup uop, no FEDP compute }; // Helper: is_sparse derived from op_type (no per-uop bit). diff --git a/sw/common/tensor_cfg.h b/sw/common/tensor_cfg.h index 18e5116fa..d856ccc7b 100644 --- a/sw/common/tensor_cfg.h +++ b/sw/common/tensor_cfg.h @@ -345,7 +345,7 @@ struct wmma_config_t { // All geometry derived from NT and NRC alone (NRA=4 fixed): // tcM = 2^ceil(log2(NT)/2), tcN = tcK = 2^floor(log2(NT)/2) // xtileM = 2*tcM, xtileN = NRC*NT/xtileM, xtileK = 2*tcK -// m_steps = k_steps = 2 (always) +// m_steps = 2; k_steps is 2 normally, or 1 with a doubled-K WGMMA FEDP. template struct wgmma_config_t { private: @@ -360,12 +360,18 @@ struct wgmma_config_t { static constexpr uint32_t tcM = 1u << ((lg_NT + 1) / 2); static constexpr uint32_t tcN = 1u << (lg_NT / 2); static constexpr uint32_t tcK = tcN; + static constexpr uint32_t fedpK = +#ifdef VX_CFG_TCU_FEDP2K + 2 * tcK; +#else + tcK; +#endif static constexpr uint32_t xtileM = 2 * tcM; static constexpr uint32_t xtileN = (NRC_ * NT) / xtileM; static constexpr uint32_t xtileK = 2 * tcK; static constexpr uint32_t tileK = xtileK * i_ratio; static constexpr uint32_t m_steps = 2; - static constexpr uint32_t k_steps = 2; + static constexpr uint32_t k_steps = xtileK / fedpK; static constexpr uint32_t NRC = NRC_; }; diff --git a/sw/kernel/include/vx_tensor.h b/sw/kernel/include/vx_tensor.h index f067ea041..61e977289 100644 --- a/sw/kernel/include/vx_tensor.h +++ b/sw/kernel/include/vx_tensor.h @@ -715,6 +715,9 @@ struct wgmma_context { // Fragments using fragment_acc = typename ctx_c::fragment_acc; using fragment_a = typename ctx_a::fragment_a; + using fragment_b = typename ctx_c::fragment_b; + + static constexpr bool input_is_mx = mx_scale_format(It::id); // Block (micro-tile) geometry — derived from NT alone static constexpr uint32_t i_ratio = XB / sizeof(typename It::dtype); @@ -722,19 +725,27 @@ struct wgmma_context { static constexpr uint32_t tcN = 1u << (lg_NT / 2); static constexpr uint32_t tcK = tcN; - // Per-warp tile geometry — m_steps = k_steps = 2 always + // Per-warp tile geometry. WGMMA-2K keeps the logical tile K fixed and + // consumes both K halves inside one widened FEDP uop. + static constexpr uint32_t fedpK = +#ifdef VX_CFG_TCU_FEDP2K + 2 * tcK; +#else + tcK; +#endif static constexpr uint32_t m_steps = 2; - static constexpr uint32_t k_steps = 2; + static constexpr uint32_t k_steps = (2 * tcK) / fedpK; + static constexpr uint32_t a_reg_k_steps = 2; static constexpr uint32_t xtileM = m_steps * tcM; static constexpr uint32_t xtileN = (NRC_ * NT) / xtileM; - static constexpr uint32_t tileK = k_steps * tcK * i_ratio; + static constexpr uint32_t tileK = k_steps * fedpK * i_ratio; static constexpr uint32_t n_steps = xtileN / tcN; // Block-major SMEM constants. // BLOCK = micro-tile (tcM × tcK or tcK × tcN), measured in input_t element units. - static constexpr uint32_t a_blk_elems = tcM * tcK * i_ratio; // elements per A block + static constexpr uint32_t a_blk_elems = tcM * fedpK * i_ratio; // elements per A block static constexpr uint32_t a_warp_elems = xtileM * tileK; // elements per warp's A slice - static constexpr uint32_t b_blk_elems = tcK * i_ratio * tcN; // elements per B block + static constexpr uint32_t b_blk_elems = fedpK * i_ratio * tcN; // elements per B block // Cooperative-load index into A_smem for an (r, c) target in the // row-major-equivalent A view (r ∈ [0, cta_M), c ∈ [0, tileK)). @@ -744,11 +755,11 @@ struct wgmma_context { uint32_t r_in_warp = r % xtileM; uint32_t m_blk = r_in_warp / tcM; uint32_t i_in = r_in_warp % tcM; - uint32_t k_blk = c / (tcK * i_ratio); - uint32_t k_in = c % (tcK * i_ratio); + uint32_t k_blk = c / (fedpK * i_ratio); + uint32_t k_in = c % (fedpK * i_ratio); return warp_idx * a_warp_elems + (k_blk * m_steps + m_blk) * a_blk_elems - + i_in * (tcK * i_ratio) + k_in; + + i_in * (fedpK * i_ratio) + k_in; } // Cooperative-load index into B_smem for an (r, c) target in the @@ -756,12 +767,12 @@ struct wgmma_context { // Within-block layout: N outer, K inner — each 32-bit word packs i_ratio // K-elements at one (j, k_word) cell, matching tcu_core's b_off + j*TC_K + k. static __attribute__((always_inline)) uint32_t b_blockmajor_idx(uint32_t r, uint32_t c) { - uint32_t k_blk = r / (tcK * i_ratio); - uint32_t r_in = r % (tcK * i_ratio); + uint32_t k_blk = r / (fedpK * i_ratio); + uint32_t r_in = r % (fedpK * i_ratio); uint32_t n_blk = c / tcN; uint32_t n_in = c % tcN; return (k_blk * n_steps + n_blk) * b_blk_elems - + n_in * (tcK * i_ratio) + r_in; + + n_in * (fedpK * i_ratio) + r_in; } // Cooperative-load index into per-warp sparse A_smem_w for an (r, c) target @@ -769,11 +780,12 @@ struct wgmma_context { // Sparse A is K/2 compressed; same per-block shape as dense A but only // K_STEPS/2 half-k blocks. Caller passes a per-warp pointer. static __attribute__((always_inline)) uint32_t a_sp_blockmajor_idx(uint32_t r, uint32_t c) { + constexpr uint32_t sparse_blk_elems = tcM * tcK * i_ratio; uint32_t half_k_blk = c / (tcK * i_ratio); uint32_t k_in_elem = c % (tcK * i_ratio); uint32_t m_blk = r / tcM; uint32_t i_in = r % tcM; - return (half_k_blk * m_steps + m_blk) * a_blk_elems + return (half_k_blk * m_steps + m_blk) * sparse_blk_elems + i_in * (tcK * i_ratio) + k_in_elem; } @@ -784,7 +796,8 @@ struct wgmma_context { // so wg_meta_total_bytes = NT*4. Unused SRAM cells are packed with zero // to keep the load shape uniform. static constexpr uint32_t sp_rtl_i_ratio = 32 / It::bits; - static constexpr uint32_t wg_meta_banks = m_steps * (k_steps / 2); + static constexpr uint32_t sparse_k_steps = (k_steps > 1) ? (k_steps / 2) : 1; + static constexpr uint32_t wg_meta_banks = m_steps * sparse_k_steps; static constexpr uint32_t wg_meta_row_bits = tcK * 2 * sp_rtl_i_ratio; static constexpr uint32_t wg_meta_stride_words = (tcM * wg_meta_row_bits + 31) / 32; static constexpr uint32_t wg_meta_stride_bytes = wg_meta_stride_words * 4; @@ -816,6 +829,7 @@ struct wgmma_context { static_assert(tcM * tcK == NT, "wgmma block-major load assumes canonical config (TC_M*TC_K == NT)"); constexpr uint32_t k_row_elems = tcK * i_ratio; + constexpr uint32_t sparse_blk_elems = tcM * k_row_elems; uint32_t lane = vx_thread_id(); uint32_t i_in = lane / tcK; uint32_t k_in_elem = (lane % tcK) * i_ratio; @@ -824,7 +838,8 @@ struct wgmma_context { uint32_t elem_off; if (ldm == 0) { // Block-major SMEM: blocks contiguous, k-block outer. - elem_off = (k_blk * m_steps + m_blk) * a_blk_elems + constexpr uint32_t block_elems = is_sparse ? sparse_blk_elems : a_blk_elems; + elem_off = (k_blk * m_steps + m_blk) * block_elems + i_in * k_row_elems + k_in_elem; } else { @@ -837,7 +852,7 @@ struct wgmma_context { }; if constexpr (is_sparse) { - constexpr uint32_t sp_k_steps_local = k_steps / 2; + constexpr uint32_t sp_k_steps_local = sparse_k_steps; constexpr uint32_t a_regs = m_steps * sp_k_steps_local; detail::unroll_for([&](auto r) { uint32_t m_blk = r / sp_k_steps_local; @@ -846,9 +861,24 @@ struct wgmma_context { }); } else { detail::unroll_for([&](auto r) { - uint32_t m_blk = r / k_steps; - uint32_t k_blk = r % k_steps; - dst.data[r] = load_reg(m_blk, k_blk); + uint32_t m_blk = r / a_reg_k_steps; + uint32_t k_blk = r % a_reg_k_steps; + if constexpr (k_steps == 1) { + uint32_t elem_off; + if (ldm == 0) { + elem_off = m_blk * a_blk_elems + + i_in * (fedpK * i_ratio) + + k_blk * k_row_elems + + k_in_elem; + } else { + elem_off = (m_blk * tcM + i_in) * uint32_t(ldm) + + (k_blk * k_row_elems + k_in_elem); + } + dst.data[r] = *reinterpret_cast( + reinterpret_cast(src) + elem_off); + } else { + dst.data[r] = load_reg(m_blk, k_blk); + } }); } } @@ -876,6 +906,11 @@ struct wgmma_context { ); } + template + static __attribute__((always_inline)) void load_mx_metadata(Frag& frag, const void* meta_mx_ptr) { + ctx_c::load_mx_metadata(frag, meta_mx_ptr); + } + // Store accumulator with n-major register layout: r = n * m_steps + m template static __attribute__((always_inline)) void store_matrix_sync(void *dst, const Frag &src, size_t ldm) { @@ -1040,6 +1075,7 @@ struct wgmma_context { // --- RS path: A from registers, B from smem (NRC <= 16, NRA = 4) --- else if constexpr (!a_is_smem && b_is_smem) { static_assert(OpA::NR == 4, "WGMMA RS requires NRA=4"); + register uint32_t ra __asm__("a0") = 0; register uint32_t rb __asm__("a1") = op_b.value; if constexpr (NRC_ == 16) { @@ -1072,7 +1108,7 @@ struct wgmma_context { "+f"(fd12), "+f"(fd13), "+f"(fd14), "+f"(fd15) : [insn]"i"(RISCV_CUSTOM0), [fmd]"i"(Ot::id), [fms]"i"(It::id), [flags]"i"(flags), "f"(fa0), "f"(fa1), "f"(fa2), "f"(fa3), - "r"(rb) + "r"(ra), "r"(rb) ); frag_d.data = { @@ -1099,7 +1135,7 @@ struct wgmma_context { "+f"(fd4), "+f"(fd5), "+f"(fd6), "+f"(fd7) : [insn]"i"(RISCV_CUSTOM0), [fmd]"i"(Ot::id), [fms]"i"(It::id), [flags]"i"(flags), "f"(fa0), "f"(fa1), "f"(fa2), "f"(fa3), - "r"(rb) + "r"(ra), "r"(rb) ); frag_d.data = {fd0, fd1, fd2, fd3, fd4, fd5, fd6, fd7}; diff --git a/tests/regression/Makefile b/tests/regression/Makefile index d11fca954..be60b6363 100644 --- a/tests/regression/Makefile +++ b/tests/regression/Makefile @@ -8,7 +8,7 @@ TESTS := \ vecadd vecadd_v1 sgemm sgemmx sgemm_v1 conv3 relu sgemv \ sgemm2 sgemm2_v1 madmax stencil3d raycast bfs jacobi pathfinder \ softmax \ - sgemm_tcu sgemm_tcu_mx sgemm_tcu_sp_mx sgemm_tcu_wg sgemm_tcu_sp sgemm_tcu_wg_sp \ + sgemm_tcu sgemm_tcu_mx sgemm_tcu_sp_mx sgemm_tcu_wg sgemm_tcu_wg_mx sgemm_tcu_sp sgemm_tcu_wg_sp \ sgemm2_dxa sgemm2_tcu sgemm_tcu_wg_dxa sgemm_tcu_wg_sp_dxa \ sgemm2_dxa_mcast sgemm_tcu_wg_dxa_mcast \ dxa_copy dxa_copy_mcast dxa_kmajor_check \ diff --git a/tests/regression/sgemm_tcu_mx/main.cpp b/tests/regression/sgemm_tcu_mx/main.cpp index 8887c6014..fd7edc87a 100644 --- a/tests/regression/sgemm_tcu_mx/main.cpp +++ b/tests/regression/sgemm_tcu_mx/main.cpp @@ -59,10 +59,8 @@ static void pack_mx_a_metadata(std::vector &packed, uint32_t logical_tileK = cfg::tileK * kElemsPerByte; uint32_t num_k_tiles = K_logical / logical_tileK; uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; - uint32_t tile_scale_blocks_k = logical_tileK / vt::ITYPE::ele_block; - if (tile_scale_blocks_k == 0) { - tile_scale_blocks_k = 1; - } + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); packed.assign(num_tile_rows * num_k_tiles * VX_CFG_NUM_THREADS, 0); for (uint32_t tr = 0; tr < num_tile_rows; ++tr) { @@ -91,10 +89,8 @@ static void pack_mx_b_metadata(std::vector &packed, uint32_t num_tile_cols = N / cfg::tileN; uint32_t logical_tileK = cfg::tileK * kElemsPerByte; uint32_t num_k_tiles = K_logical / logical_tileK; - uint32_t tile_scale_blocks_k = logical_tileK / vt::ITYPE::ele_block; - if (tile_scale_blocks_k == 0) { - tile_scale_blocks_k = 1; - } + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); packed.assign(num_tile_cols * num_k_tiles * VX_CFG_NUM_THREADS, 0); for (uint32_t tc = 0; tc < num_tile_cols; ++tc) { diff --git a/tests/regression/sgemm_tcu_wg/Makefile b/tests/regression/sgemm_tcu_wg/Makefile index 66934cf47..d67265b4c 100644 --- a/tests/regression/sgemm_tcu_wg/Makefile +++ b/tests/regression/sgemm_tcu_wg/Makefile @@ -13,11 +13,9 @@ VX_SRCS := $(SRC_DIR)/kernel.cpp CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) -CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DWGMMA_RS +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS) -# RTLsim is roughly two orders of magnitude slower than SimX; keep the -# default matrix size small enough for a 60s sweep budget. Override -# -m/-n/-k at the command line for full-size correctness sweeps. OPTS := $(OPTS) $(if $(findstring -m,$(OPTS)),,-m 16) OPTS := $(OPTS) $(if $(findstring -n,$(OPTS)),,-n 16) OPTS := $(OPTS) $(if $(findstring -k,$(OPTS)),,-k 16) diff --git a/tests/regression/sgemm_tcu_wg_mx/Makefile b/tests/regression/sgemm_tcu_wg_mx/Makefile new file mode 100644 index 000000000..7b985e493 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/Makefile @@ -0,0 +1,23 @@ +ROOT_DIR := $(realpath ../../..) +include $(ROOT_DIR)/config.mk + +CONFIGS := $(if $(findstring -DVX_CFG_EXT_TCU_ENABLE,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_EXT_TCU_ENABLE) +CONFIGS := $(if $(findstring -DVX_CFG_TCU_MX_ENABLE,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_TCU_MX_ENABLE) +CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) +CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS) + +PROJECT := sgemm_tcu_wg_mx +SRC_DIR := $(VORTEX_HOME)/tests/regression/$(PROJECT) +SRCS := $(SRC_DIR)/main.cpp $(SW_COMMON_DIR)/rvfloats.cpp $(SW_COMMON_DIR)/softfloat_ext.cpp +VX_SRCS := $(SRC_DIR)/kernel.cpp + +OPTS := $(OPTS) $(if $(findstring -m,$(OPTS)),,-m 64) +OPTS := $(OPTS) $(if $(findstring -n,$(OPTS)),,-n 64) +OPTS := $(OPTS) $(if $(findstring -k,$(OPTS)),,-k 64) +CXXFLAGS += -I$(THIRD_PARTY_DIR)/softfloat/source/include +LDFLAGS += $(THIRD_PARTY_DIR)/softfloat/build/Linux-x86_64-GCC/softfloat.a +KERNEL_LIB := vortex2 + +include ../common.mk diff --git a/tests/regression/sgemm_tcu_wg_mx/common.h b/tests/regression/sgemm_tcu_wg_mx/common.h new file mode 100644 index 000000000..7120f73aa --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/common.h @@ -0,0 +1,25 @@ +#ifndef _COMMON_H_ +#define _COMMON_H_ + +#include + +#ifndef WGMMA_NRC +#define WGMMA_NRC 8 +#endif +#ifndef ITYPE +#define ITYPE mxfp8 +#endif +#ifndef OTYPE +#define OTYPE fp32 +#endif + +typedef struct { + uint32_t M, N, K; + uint64_t A_addr, B_addr, C_addr; + uint64_t MX_A_addr, MX_B_addr; +#ifdef VX_CFG_TCU_MX_TLS + uint64_t A_tensor_scale_addr, B_tensor_scale_addr; +#endif +} kernel_arg_t; + +#endif diff --git a/tests/regression/sgemm_tcu_wg_mx/kernel.cpp b/tests/regression/sgemm_tcu_wg_mx/kernel.cpp new file mode 100644 index 000000000..fe4b67ad7 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/kernel.cpp @@ -0,0 +1,83 @@ +#include "common.h" +#include +#include +#include + +namespace vt = vortex::tensor; +using ctx = vt::wgmma_context; + +__kernel void kernel_main(kernel_arg_t* __UNIFORM__ arg) { + static_assert(sizeof(ctx::input_t) == 1, "MX WGMMA inputs use byte storage"); + auto pA = reinterpret_cast(arg->A_addr); + auto pB = reinterpret_cast(arg->B_addr); + auto pC = reinterpret_cast(arg->C_addr); + auto pMxA = reinterpret_cast(arg->MX_A_addr); + auto pMxB = reinterpret_cast(arg->MX_B_addr); +#ifdef VX_CFG_TCU_MX_TLS + auto pATensorScale = reinterpret_cast(arg->A_tensor_scale_addr); + auto pBTensorScale = reinterpret_cast(arg->B_tensor_scale_addr); +#endif + + uint32_t N = arg->N; + uint32_t K = arg->K; + uint32_t tid = threadIdx.x; + 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; + uint32_t cta_m = num_warps * ctx::xtileM; + uint32_t tile_row = blockIdx.y * cta_m; + uint32_t tile_col = blockIdx.x * ctx::xtileN; + uint32_t num_k_tiles = K / ctx::tileK; + + auto A_smem = reinterpret_cast(__local_mem()); + auto B_smem = A_smem + cta_m * ctx::tileK; + ctx::fragment_a fragA; + ctx::fragment_b fragB; + ctx::fragment_acc fragC; + ctx::fill_fragment(fragC, 0); + + for (uint32_t k = 0; k < K; k += ctx::tileK) { + uint32_t k_tile = k / ctx::tileK; + uint32_t a_bytes = cta_m * ctx::tileK; + for (uint32_t idx = tid; idx < a_bytes; idx += num_threads) { + uint32_t row = idx / ctx::tileK; + uint32_t col = idx % ctx::tileK; + A_smem[idx] = pA[(tile_row + row) * K + k + col]; + } + uint32_t b_bytes = ctx::xtileN * ctx::tileK; + for (uint32_t idx = tid; idx < b_bytes; idx += num_threads) { + uint32_t col = idx / ctx::tileK; + uint32_t k_byte = idx % ctx::tileK; + B_smem[idx] = pB[(tile_col + col) * K + k + k_byte]; + } + __syncthreads(); + + uint32_t tile_m = blockIdx.y * num_warps + warp_rank; + auto pTileMxA = pMxA + (tile_m * num_k_tiles + k_tile) * VX_CFG_NUM_THREADS; + auto pTileMxB = pMxB + (blockIdx.x * num_k_tiles + k_tile) * VX_CFG_NUM_THREADS; + ctx::load_mx_metadata(fragA, pTileMxA); + ctx::load_mx_metadata(fragB, pTileMxB); + + auto A_warp = A_smem + warp_rank * ctx::xtileM * ctx::tileK; + auto desc_b = vt::vx_make_smem_desc(B_smem, ctx::tileK); +#if defined(WGMMA_RS) && (WGMMA_NRC <= 16) + ctx::load_matrix_sync(fragA, A_warp, ctx::tileK); + ctx::wgmma_sync(fragC, fragA, desc_b, fragC); +#else + auto desc_a = vt::vx_make_smem_desc(A_warp, ctx::tileK); + ctx::wgmma_sync(fragC, desc_a, desc_b, fragC); +#endif + __syncthreads(); + } + +#ifdef VX_CFG_TCU_MX_TLS + if constexpr (std::is_same::value) { + float tensor_scale = (*pATensorScale) * (*pBTensorScale); + for (uint32_t r = 0; r < ctx::fragment_acc::NR; ++r) + fragC.data[r] *= tensor_scale; + } +#endif + auto out = pC + (tile_row + warp_rank * ctx::xtileM) * N + tile_col; + ctx::store_matrix_sync(out, fragC, N); +} diff --git a/tests/regression/sgemm_tcu_wg_mx/main.cpp b/tests/regression/sgemm_tcu_wg_mx/main.cpp new file mode 100644 index 000000000..5b3e80164 --- /dev/null +++ b/tests/regression/sgemm_tcu_wg_mx/main.cpp @@ -0,0 +1,465 @@ +#include "common.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define MAX_ERRORS 100 + +#define RT_CHECK(_expr) \ + do { \ + int _ret = _expr; \ + if (0 == _ret) \ + break; \ + printf("Error: '%s' returned %d!\n", #_expr, (int)_ret); \ + cleanup(); \ + exit(-1); \ + } while (false) + +using namespace vortex; +namespace vt = tensor; + +using cfg = vt::wgmma_config_t; +static constexpr uint32_t kTileM = cfg::xtileM; +static constexpr uint32_t kTileN = cfg::xtileN; +using itype_t = typename vt::ITYPE::dtype; +using otype_t = typename vt::OTYPE::dtype; + +static constexpr uint32_t kElemsPerByte = (vt::ITYPE::bits < 8) ? (8 / vt::ITYPE::bits) : 1; +static constexpr uint32_t kScalePack = 4; + +static uint8_t read_nibble(const uint8_t *ptr, uint32_t offset) { + uint8_t value = ptr[offset / 2]; + return (offset & 1) ? (value >> 4) : (value & 0x0f); +} + +static uint32_t pack4(const std::vector &scales, uint32_t start) { + uint32_t word = 0; + for (uint32_t i = 0; i < kScalePack; ++i) { + if ((start + i) < scales.size()) { + word |= static_cast(scales[start + i]) << (8 * i); + } + } + return word; +} + +static void pack_mx_a_metadata(std::vector &packed, + const std::vector &scales, + uint32_t M, + uint32_t K_logical) { + uint32_t num_tile_rows = M / kTileM; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t num_k_tiles = K_logical / logical_tileK; + uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); + + packed.assign(num_tile_rows * num_k_tiles * VX_CFG_NUM_THREADS, 0); + for (uint32_t tr = 0; tr < num_tile_rows; ++tr) { + for (uint32_t kt = 0; kt < num_k_tiles; ++kt) { + std::vector tile_scales; + tile_scales.reserve(kTileM * tile_scale_blocks_k); + uint32_t first_block = (kt * logical_tileK) / vt::ITYPE::ele_block; + for (uint32_t m = 0; m < kTileM; ++m) { + uint32_t row = tr * kTileM + m; + for (uint32_t kb = 0; kb < tile_scale_blocks_k; ++kb) { + tile_scales.push_back(scales[row * scale_blocks_k + first_block + kb]); + } + } + uint32_t base = (tr * num_k_tiles + kt) * VX_CFG_NUM_THREADS; + for (uint32_t w = 0; w < VX_CFG_NUM_THREADS && (w * kScalePack) < tile_scales.size(); ++w) { + packed[base + w] = pack4(tile_scales, w * kScalePack); + } + } + } +} + +static void pack_mx_b_metadata(std::vector &packed, + const std::vector &scales, + uint32_t N, + uint32_t K_logical) { + uint32_t num_tile_cols = N / kTileN; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t num_k_tiles = K_logical / logical_tileK; + uint32_t tile_scale_blocks_k = + std::max(1u, (logical_tileK + vt::ITYPE::ele_block - 1) / vt::ITYPE::ele_block); + + packed.assign(num_tile_cols * num_k_tiles * VX_CFG_NUM_THREADS, 0); + for (uint32_t tc = 0; tc < num_tile_cols; ++tc) { + for (uint32_t kt = 0; kt < num_k_tiles; ++kt) { + std::vector tile_scales; + tile_scales.reserve(kTileN * tile_scale_blocks_k); + uint32_t first_block = (kt * logical_tileK) / vt::ITYPE::ele_block; + for (uint32_t n = 0; n < kTileN; ++n) { + uint32_t col = tc * kTileN + n; + for (uint32_t kb = 0; kb < tile_scale_blocks_k; ++kb) { + tile_scales.push_back(scales[(first_block + kb) * N + col]); + } + } + uint32_t base = (tc * num_k_tiles + kt) * VX_CFG_NUM_THREADS; + for (uint32_t w = 0; w < VX_CFG_NUM_THREADS && (w * kScalePack) < tile_scales.size(); ++w) { + packed[base + w] = pack4(tile_scales, w * kScalePack); + } + } + } +} + +static float dequantize_mx_value(const itype_t *data, + const std::vector &scales, + uint32_t offset, + uint32_t scale_index, + float tensor_scale) { + uint8_t sf = scales[scale_index]; + switch (vt::ITYPE::id) { + case vt::mxfp8::id: + return bit_cast(rv_mxfp8tof_s(data[offset], sf, 0, nullptr)); + case vt::mxbf8::id: + return bit_cast(rv_mxbf8tof_s(data[offset], sf, 0, nullptr)); + case vt::mxint8::id: + { + float scale = std::ldexp(1.0f, static_cast(sf) - 127); + return (static_cast(data[offset]) / 64.0f) * scale; + } + case vt::mxfp4::id: + { + uint8_t q = read_nibble(reinterpret_cast(data), offset); + return bit_cast(rv_mxfp4tof_s(q, sf, 0, nullptr)); + } + case vt::nvfp4::id: + { + uint8_t q = read_nibble(reinterpret_cast(data), offset); + return bit_cast(rv_nvfp4tof_s(q, sf, 0, nullptr)) * tensor_scale; + } + default: + std::abort(); + } +} + +static int32_t trunc_shift(int32_t value, int32_t shift) { + if (shift >= 0) { + return value << shift; + } + uint32_t abs_shift = static_cast(-shift); + uint32_t mag = value < 0 ? static_cast(-value) : static_cast(value); + int32_t scaled = static_cast(mag >> abs_shift); + return value < 0 ? -scaled : scaled; +} + +static void matmul_cpu(otype_t *C, + const itype_t *A, + const itype_t *B, + const std::vector &scale_a, + const std::vector &scale_b, + float tensor_scale_a, + float tensor_scale_b, + uint32_t M, + uint32_t N, + uint32_t K_logical) { + uint32_t scale_blocks_k = K_logical / vt::ITYPE::ele_block; + for (uint32_t m = 0; m < M; ++m) { + for (uint32_t n = 0; n < N; ++n) { + if constexpr (std::is_same::value) { + int32_t sum = 0; + for (uint32_t k = 0; k < K_logical; ++k) { + uint32_t scale_k = k / vt::ITYPE::ele_block; + uint8_t sf_a = scale_a[m * scale_blocks_k + scale_k]; + uint8_t sf_b = scale_b[scale_k * N + n]; + int32_t shift = static_cast(sf_a) - 133 + static_cast(sf_b) - 133; + int32_t product = static_cast(A[m * K_logical + k]) * static_cast(B[n * K_logical + k]); + sum += trunc_shift(product, shift); + } + C[m * N + n] = sum; + } else { + float sum = 0.0f; + for (uint32_t k = 0; k < K_logical; ++k) { + uint32_t scale_k = k / vt::ITYPE::ele_block; + auto a = dequantize_mx_value(A, scale_a, m * K_logical + k, + m * scale_blocks_k + scale_k, tensor_scale_a); + auto b = dequantize_mx_value(B, scale_b, n * K_logical + k, + scale_k * N + n, tensor_scale_b); + sum += a * b; + } + C[m * N + n] = sum; + } + } + } +} + +template +static bool quantize_inputs(typename FormatT::dtype *A, + typename FormatT::dtype *B, + std::vector &scale_a, + std::vector &scale_b, + float &A_tensor_scale, + float &B_tensor_scale, + const float *A_dense, + const float *B_dense, + uint32_t M, + uint32_t N, + uint32_t K_logical) { + if constexpr (std::is_same::value) { + return vt::quantize_mx_a_rowmajor( + reinterpret_cast(A), scale_a, A_tensor_scale, + A_dense, M, K_logical) + && vt::quantize_mx_b_colmajor( + reinterpret_cast(B), scale_b, B_tensor_scale, + B_dense, K_logical, N); + } else { + return vt::quantize_mx_a_rowmajor( + A, scale_a, A_dense, M, K_logical) + && vt::quantize_mx_b_colmajor( + B, scale_b, B_dense, K_logical, N); + } +} + +const char *kernel_file = "kernel.vxbin"; + +uint32_t xm = 32; +uint32_t xn = 32; +uint32_t xk = 32; + +vx_device_h device = nullptr; +vx_queue_h queue = nullptr; +vx_module_h module_ = nullptr; +vx_kernel_h kernel = nullptr; +vx_buffer_h A_buffer = nullptr; +vx_buffer_h B_buffer = nullptr; +vx_buffer_h C_buffer = nullptr; +vx_buffer_h MX_A_buffer = nullptr; +vx_buffer_h MX_B_buffer = nullptr; +#ifdef VX_CFG_TCU_MX_TLS +vx_buffer_h A_tensor_scale_buffer = nullptr; +vx_buffer_h B_tensor_scale_buffer = nullptr; +#endif +kernel_arg_t kernel_arg = {}; + +static void show_usage() { + std::cout << "Vortex SGEMM TCU MX Test." << std::endl; + std::cout << "Usage: [-m M] [-n N] [-k K] [-h]" << std::endl; +} + +static void parse_args(int argc, char **argv) { + int c; + while ((c = getopt(argc, argv, "m:n:k:h")) != -1) { + switch (c) { + case 'm': xm = atoi(optarg); break; + case 'n': xn = atoi(optarg); break; + case 'k': xk = atoi(optarg); break; + case 'h': + show_usage(); + exit(0); + default: + show_usage(); + exit(-1); + } + } +} + +void cleanup() { + if (device) { + vx_mem_free(A_buffer); + vx_mem_free(B_buffer); + vx_mem_free(C_buffer); + vx_mem_free(MX_A_buffer); + vx_mem_free(MX_B_buffer); +#ifdef VX_CFG_TCU_MX_TLS + vx_mem_free(A_tensor_scale_buffer); + vx_mem_free(B_tensor_scale_buffer); +#endif + if (kernel) vx_kernel_release(kernel); + if (module_) vx_module_release(module_); + if (queue) vx_queue_release(queue); + vx_dump_perf(device, stdout); + vx_dev_close(device); + } +} + +int main(int argc, char *argv[]) { + parse_args(argc, argv); + std::srand(50); + + RT_CHECK(vx_dev_open(&device)); + vx_queue_info_t qi = {sizeof(qi), nullptr, VX_QUEUE_PRIORITY_NORMAL, 0}; + RT_CHECK(vx_queue_create(device, &qi, &queue)); + + uint64_t isa_flags; + RT_CHECK(vx_dev_caps(device, VX_CAPS_ISA_FLAGS, &isa_flags)); + if ((isa_flags & VX_ISA_EXT_TCU) == 0) { + std::cout << "TCU extension not supported!" << std::endl; + cleanup(); + return -1; + } + + uint64_t NT; + RT_CHECK(vx_dev_caps(device, VX_CAPS_NUM_THREADS, &NT)); + if (NT != VX_CFG_NUM_THREADS) { + std::cout << "Error: device warp size (" << NT << ") must match VX_CFG_NUM_THREADS=" << VX_CFG_NUM_THREADS << "!" << std::endl; + return -1; + } + + uint32_t M = xm; + uint32_t N = xn; + uint32_t K_logical = xk; + uint32_t logical_tileK = cfg::tileK * kElemsPerByte; + uint32_t K_storage = K_logical / kElemsPerByte; + + uint32_t warps_per_cta = VX_CFG_ISSUE_WIDTH; + uint32_t cta_m = warps_per_cta * kTileM; + if ((M % cta_m) != 0 || (N % kTileN) != 0 || (K_logical % logical_tileK) != 0) { + std::cout << "Error: M/N/K must be multiples of CTA M=" << cta_m + << " N=" << kTileN << " K=" << logical_tileK << std::endl; + return -1; + } + if ((K_logical % vt::ITYPE::ele_block) != 0) { + std::cout << "Error: K must be a multiple of MX block size=" << vt::ITYPE::ele_block << std::endl; + return -1; + } + + size_t sizeA = M * K_storage; + size_t sizeB = K_storage * N; + size_t sizeC = M * N; + uint32_t grid_dim[2] = {N / kTileN, M / cta_m}; + uint32_t block_dim[2] = {warps_per_cta * (uint32_t)NT, 1}; + + std::cout << "input data type: " << vt::ITYPE::name << " (id=" << vt::ITYPE::id << ")" << std::endl; + std::cout << "output data type: " << vt::OTYPE::name << " (id=" << vt::OTYPE::id << ")" << std::endl; + std::cout << "MMA Tile Dimension: M=" << kTileM << ", N=" << kTileN + << ", K(logical)=" << logical_tileK << ", K(storage)=" << cfg::tileK << std::endl; + + kernel_arg.M = M; + kernel_arg.N = N; + kernel_arg.K = K_storage; +#ifdef VX_CFG_TCU_MX_TLS + kernel_arg.A_tensor_scale_addr = 0; + kernel_arg.B_tensor_scale_addr = 0; +#endif + float A_tensor_scale = 1.0f; + float B_tensor_scale = 1.0f; + + RT_CHECK(vx_mem_alloc(device, sizeA * sizeof(itype_t), VX_MEM_READ, &A_buffer)); + RT_CHECK(vx_mem_address(A_buffer, &kernel_arg.A_addr)); + RT_CHECK(vx_mem_alloc(device, sizeB * sizeof(itype_t), VX_MEM_READ, &B_buffer)); + RT_CHECK(vx_mem_address(B_buffer, &kernel_arg.B_addr)); + RT_CHECK(vx_mem_alloc(device, sizeC * sizeof(otype_t), VX_MEM_WRITE, &C_buffer)); + RT_CHECK(vx_mem_address(C_buffer, &kernel_arg.C_addr)); + + std::vector h_A_dense(M * K_logical); + std::vector h_B_dense(K_logical * N); + for (auto &v : h_A_dense) { + v = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + } + for (auto &v : h_B_dense) { + v = (static_cast(std::rand()) / RAND_MAX) * 2.0f - 1.0f; + } + + std::vector h_A(sizeA); + std::vector h_B(sizeB); + std::vector scale_a; + std::vector scale_b; + + bool ok = quantize_inputs( + h_A.data(), h_B.data(), scale_a, scale_b, A_tensor_scale, B_tensor_scale, + h_A_dense.data(), h_B_dense.data(), M, N, K_logical); + if (!ok) { + std::cout << "Error: MX quantization failed!" << std::endl; + return -1; + } + + std::vector h_mx_a; + std::vector h_mx_b; + pack_mx_a_metadata(h_mx_a, scale_a, M, K_logical); + pack_mx_b_metadata(h_mx_b, scale_b, N, K_logical); + + RT_CHECK(vx_mem_alloc(device, h_mx_a.size() * sizeof(uint32_t), VX_MEM_READ, &MX_A_buffer)); + RT_CHECK(vx_mem_address(MX_A_buffer, &kernel_arg.MX_A_addr)); + RT_CHECK(vx_mem_alloc(device, h_mx_b.size() * sizeof(uint32_t), VX_MEM_READ, &MX_B_buffer)); + RT_CHECK(vx_mem_address(MX_B_buffer, &kernel_arg.MX_B_addr)); +#ifdef VX_CFG_TCU_MX_TLS + if constexpr (std::is_same::value) { + RT_CHECK(vx_mem_alloc(device, sizeof(float), VX_MEM_READ, &A_tensor_scale_buffer)); + RT_CHECK(vx_mem_address(A_tensor_scale_buffer, &kernel_arg.A_tensor_scale_addr)); + RT_CHECK(vx_mem_alloc(device, sizeof(float), VX_MEM_READ, &B_tensor_scale_buffer)); + RT_CHECK(vx_mem_address(B_tensor_scale_buffer, &kernel_arg.B_tensor_scale_addr)); + } +#endif + + RT_CHECK(vx_copy_to_dev(A_buffer, h_A.data(), 0, sizeA * sizeof(itype_t))); + RT_CHECK(vx_copy_to_dev(B_buffer, h_B.data(), 0, sizeB * sizeof(itype_t))); + RT_CHECK(vx_copy_to_dev(MX_A_buffer, h_mx_a.data(), 0, h_mx_a.size() * sizeof(uint32_t))); + RT_CHECK(vx_copy_to_dev(MX_B_buffer, h_mx_b.data(), 0, h_mx_b.size() * sizeof(uint32_t))); +#ifdef VX_CFG_TCU_MX_TLS + if constexpr (std::is_same::value) { + RT_CHECK(vx_copy_to_dev(A_tensor_scale_buffer, &A_tensor_scale, 0, sizeof(float))); + RT_CHECK(vx_copy_to_dev(B_tensor_scale_buffer, &B_tensor_scale, 0, sizeof(float))); + } +#endif + + RT_CHECK(vx_module_load_file(device, kernel_file, &module_)); + RT_CHECK(vx_module_get_kernel(module_, "main", &kernel)); + + auto time_start = std::chrono::high_resolution_clock::now(); + vx_launch_info_t li = {}; + li.struct_size = sizeof(li); + li.kernel = kernel; + li.args_host = &kernel_arg; + li.args_size = sizeof(kernel_arg); + li.ndim = 2; + li.grid_dim[0] = grid_dim[0]; + li.grid_dim[1] = grid_dim[1]; + li.block_dim[0] = block_dim[0]; + li.block_dim[1] = block_dim[1]; + li.lmem_size = (cta_m + kTileN) * cfg::tileK * sizeof(itype_t); + + vx_event_h launch_ev = nullptr; + RT_CHECK(vx_enqueue_launch(queue, &li, 0, nullptr, &launch_ev)); + RT_CHECK(vx_event_wait_value(launch_ev, 1, VX_TIMEOUT_INFINITE)); + vx_event_release(launch_ev); + auto time_end = std::chrono::high_resolution_clock::now(); + double elapsed = std::chrono::duration_cast(time_end - time_start).count(); + printf("Elapsed time: %lg ms\n", elapsed); + + std::vector h_C(sizeC); + RT_CHECK(vx_copy_from_dev(h_C.data(), C_buffer, 0, sizeC * sizeof(otype_t))); + + std::vector h_ref(sizeC); + matmul_cpu(h_ref.data(), h_A.data(), h_B.data(), scale_a, scale_b, + A_tensor_scale, B_tensor_scale, M, N, K_logical); + + int errors = 0; + float rel_tol = std::is_same::value ? 0.0f : + (std::is_same::value || std::is_same::value) ? 0.25f : 0.05f; + for (uint32_t i = 0; i < h_ref.size(); ++i) { + float actual = static_cast(h_C[i]); + float expected = static_cast(h_ref[i]); + float diff = std::abs(actual - expected); + float tol = std::is_same::value ? 0.0f : + std::max(1.0e-3f, std::abs(expected) * rel_tol); + if (!std::isfinite(actual) || !std::isfinite(expected) || !std::isfinite(diff) || diff > tol) { + if (errors < MAX_ERRORS) { + printf("*** error: [%d] expected=%f, actual=%f, diff=%f, tol=%f\n", + i, expected, actual, diff, tol); + } + ++errors; + } + } + + cleanup(); + if (errors != 0) { + std::cout << "Found " << errors << " / " << sizeC << " errors!" << std::endl; + std::cout << "FAILED!" << std::endl; + return errors; + } + + std::cout << "PASSED!" << std::endl; + return 0; +} diff --git a/tests/regression/sgemm_tcu_wg_sp/Makefile b/tests/regression/sgemm_tcu_wg_sp/Makefile index fa25578b9..54bb27c1a 100644 --- a/tests/regression/sgemm_tcu_wg_sp/Makefile +++ b/tests/regression/sgemm_tcu_wg_sp/Makefile @@ -13,7 +13,10 @@ VX_SRCS := $(SRC_DIR)/kernel.cpp CONFIGS := $(if $(findstring -DVX_CFG_NUM_WARPS=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_NUM_WARPS=4) CONFIGS := $(if $(findstring -DVX_CFG_ISSUE_WIDTH=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DVX_CFG_ISSUE_WIDTH=4) -CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE -DITYPE=fp16 -DOTYPE=fp32 -DWGMMA_RS +CONFIGS += -DVX_CFG_TCU_WGMMA_ENABLE -DVX_CFG_TCU_SPARSE_ENABLE +CONFIGS := $(if $(findstring -DITYPE=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DITYPE=fp16) +CONFIGS := $(if $(findstring -DOTYPE=,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DOTYPE=fp32) +CONFIGS := $(if $(findstring -DWGMMA_SS,$(CONFIGS)),$(CONFIGS),$(if $(findstring -DWGMMA_RS,$(CONFIGS)),$(CONFIGS),$(CONFIGS) -DWGMMA_RS)) CXXFLAGS += -I$(THIRD_PARTY_DIR)/softfloat/source/include diff --git a/tests/regression/sgemm_tcu_wg_sp/main.cpp b/tests/regression/sgemm_tcu_wg_sp/main.cpp index 372f5338a..30b2a1324 100644 --- a/tests/regression/sgemm_tcu_wg_sp/main.cpp +++ b/tests/regression/sgemm_tcu_wg_sp/main.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -36,7 +37,7 @@ static constexpr uint32_t kTcK = wg_cfg_t::tcK; static constexpr uint32_t kTcM = wg_cfg_t::tcM; static constexpr uint32_t kMSteps = wg_cfg_t::m_steps; static constexpr uint32_t kKSteps = wg_cfg_t::k_steps; -static constexpr uint32_t kHalfKSteps = kKSteps / 2; +static constexpr uint32_t kHalfKSteps = (kKSteps > 1) ? (kKSteps / 2) : 1; static constexpr uint32_t kMetaRowBits = kTcK * 2 * kRtlIRatio; static constexpr uint32_t kMetaStrWords = (kTcM * kMetaRowBits + 31) / 32; static constexpr uint32_t kWgMetaBanks = kMSteps * kHalfKSteps; @@ -60,6 +61,67 @@ static constexpr uint32_t kDensePerSpStep = kTcK * kRtlIRatio * 2; using itype_t = vt::ITYPE::dtype; using otype_t = vt::OTYPE::dtype; +template +struct type_ops { + using dtype = typename T::dtype; + + static dtype generate() { + return dtype(float(rand()) / RAND_MAX); + } + + static float to_float(dtype value) { + return static_cast(value); + } +}; + +template <> +struct type_ops { + static uint16_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoh_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint16_t value) { + return bit_cast(rv_htof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint16_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftob_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint16_t value) { + return bit_cast(rv_btof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint8_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoe4m3_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint8_t value) { + return bit_cast(rv_e4m3tof_s(value, 0, nullptr)); + } +}; + +template <> +struct type_ops { + static uint8_t generate() { + auto value = float(rand()) / RAND_MAX; + return rv_ftoe5m2_s(bit_cast(value), 0, nullptr); + } + + static float to_float(uint8_t value) { + return bit_cast(rv_e5m2tof_s(value, 0, nullptr)); + } +}; + // CPU reference matmul using pruned (zero-padded) A static void matmul_cpu(otype_t *C, const itype_t *A_pruned, const itype_t *B, uint32_t M, uint32_t N, uint32_t K) { @@ -69,8 +131,8 @@ static void matmul_cpu(otype_t *C, const itype_t *A_pruned, const itype_t *B, for (uint32_t k = 0; k < K; ++k) { auto a = A_pruned[m * K + k]; auto b = B[k * N + n]; - auto fa = bit_cast(rv_htof_s(a, 0, nullptr)); - auto fb = bit_cast(rv_htof_s(b, 0, nullptr)); + auto fa = type_ops::to_float(a); + auto fb = type_ops::to_float(b); sum += fa * fb; } C[m * N + n] = sum; @@ -289,12 +351,10 @@ int main(int argc, char *argv[]) { std::vector h_A_full(sizeA_full); std::vector h_B(sizeB); for (uint32_t i = 0; i < sizeA_full; ++i) { - auto fv = float(rand()) / RAND_MAX; - h_A_full[i] = rv_ftoh_s(bit_cast(fv), 0, nullptr); + h_A_full[i] = type_ops::generate(); } for (uint32_t i = 0; i < sizeB; ++i) { - auto fv = float(rand()) / RAND_MAX; - h_B[i] = rv_ftoh_s(bit_cast(fv), 0, nullptr); + h_B[i] = type_ops::generate(); } // prune A in-place (2:4) then compress @@ -380,12 +440,20 @@ int main(int argc, char *argv[]) { int errors = 0; for (uint32_t i = 0; i < sizeC; ++i) { - union fi_t { float f; int32_t i; }; - fi_t fa, fb; - fa.f = h_C[i]; - fb.f = h_ref[i]; - auto d = std::abs(fa.i - fb.i); - if (d > FLOAT_ULP) { + bool match; + if constexpr (std::is_same::value || + std::is_same::value) { + match = (h_C[i] == 0.0f && h_ref[i] == 0.0f) + || (std::abs((h_C[i] - h_ref[i]) / h_ref[i]) < 0.01f); + } else { + union fi_t { float f; int32_t i; }; + fi_t fa, fb; + fa.f = h_C[i]; + fb.f = h_ref[i]; + auto d = std::abs(fa.i - fb.i); + match = d <= FLOAT_ULP; + } + if (!match) { if (errors < MAX_ERRORS) { printf("*** error: [%u] expected=%f, actual=%f\n", i, h_ref[i], h_C[i]); }