diff --git a/kernels/blackwell_fp8_e4m3/Makefile b/kernels/blackwell_fp8_e4m3/test_case_01_gemm/Makefile similarity index 51% rename from kernels/blackwell_fp8_e4m3/Makefile rename to kernels/blackwell_fp8_e4m3/test_case_01_gemm/Makefile index 57cd1c04..30edaa5b 100644 --- a/kernels/blackwell_fp8_e4m3/Makefile +++ b/kernels/blackwell_fp8_e4m3/test_case_01_gemm/Makefile @@ -4,7 +4,7 @@ VX_SRCS = kernel.cpp VX_INCLUDES = fp8_common.hpp OPTS ?= -n1 -include ../common.mk +include ../../common.mk -args.bin input.a.bin input.b.bin input.c.bin: ../wu_arch_cases/zero.bin +args.bin input.a.bin input.b.bin input.c.bin: ../../wu_arch_cases/zero.bin cp $< $@ diff --git a/kernels/blackwell_fp8_e4m3/README.md b/kernels/blackwell_fp8_e4m3/test_case_01_gemm/README.md similarity index 100% rename from kernels/blackwell_fp8_e4m3/README.md rename to kernels/blackwell_fp8_e4m3/test_case_01_gemm/README.md diff --git a/kernels/blackwell_fp8_e4m3/fp8_common.hpp b/kernels/blackwell_fp8_e4m3/test_case_01_gemm/fp8_common.hpp similarity index 100% rename from kernels/blackwell_fp8_e4m3/fp8_common.hpp rename to kernels/blackwell_fp8_e4m3/test_case_01_gemm/fp8_common.hpp diff --git a/kernels/blackwell_fp8_e4m3/kernel.cpp b/kernels/blackwell_fp8_e4m3/test_case_01_gemm/kernel.cpp similarity index 100% rename from kernels/blackwell_fp8_e4m3/kernel.cpp rename to kernels/blackwell_fp8_e4m3/test_case_01_gemm/kernel.cpp diff --git a/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/Makefile b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/Makefile new file mode 100644 index 00000000..848781b6 --- /dev/null +++ b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/Makefile @@ -0,0 +1,13 @@ +PROJECT = blackwell_fp8_e4m3_softmax_vector + +VX_SRCS = kernel.cpp +VX_INCLUDES = fp8_common.hpp ../../wu_arch_cases/common_wu_blackwell_fa.h ../../wu_arch_cases/common_wu_min.h +OPTS ?= -n1 + +VORTEX_KN_PATH ?= $(realpath ../../../lib) +GEMMINI_SW_PATH ?= $(realpath ../../../lib/gemmini) + +include ../../common.mk + +args.bin input.a.bin input.b.bin input.c.bin: ../../wu_arch_cases/zero.bin + cp $< $@ diff --git a/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/README.md b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/README.md new file mode 100644 index 00000000..1166a49c --- /dev/null +++ b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/README.md @@ -0,0 +1,20 @@ +# test_case_02_scalar_tmem_softmax_vector + +FP8 E4M3 version of the scalar TMEM softmax vector handoff test. + +The test keeps the FP16-oriented `wu_arch_cases/case25_scalar_tmem_softmax_vector` +unchanged. This case uses FP8 E4M3 TMEM A and SMEM B layout throughout: + +- first BWGMMA computes uniform FP32 scores: `Q = 1.0`, `K = 1.0`, `C = 0.0`, + expected score `32.0`; +- scalar TMEM softmax consumes the FP32 C score tile and writes FP8 E4M3 + probabilities into TMEM A; +- uniform softmax probability is `1/32`, exactly representable as E4M3 `0x10`; +- second BWGMMA consumes that FP8 P tile with `V = 1.0`, expected FP32 output + `1.0`. + +Build: + +```bash +make -C /home/lzd/wu/wuarch/virgo-kernels/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector +``` diff --git a/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/fp8_common.hpp b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/fp8_common.hpp new file mode 100644 index 00000000..3a38f217 --- /dev/null +++ b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/fp8_common.hpp @@ -0,0 +1,41 @@ +#ifndef BLACKWELL_FP8_E4M3_COMMON_HPP +#define BLACKWELL_FP8_E4M3_COMMON_HPP + +#include +#include + +#define WU_FP8_E4M3_ZERO 0x00u +#define WU_FP8_E4M3_ONE_OVER_32 0x10u +#define WU_FP8_E4M3_HALF 0x30u +#define WU_FP8_E4M3_ONE 0x38u +#define WU_FP8_E4M3_TWO 0x40u + +#define WU_FP8_PACK4(a, b, c, d) \ + ((((uint32_t)(a) & 0xffu) << 0) | (((uint32_t)(b) & 0xffu) << 8) | \ + (((uint32_t)(c) & 0xffu) << 16) | (((uint32_t)(d) & 0xffu) << 24)) + +#define WU_FP8_REP2(x) x, x +#define WU_FP8_REP4(x) WU_FP8_REP2(x), WU_FP8_REP2(x) + +#define WU_FP8_E4M3_ONE_PACKED \ + WU_FP8_PACK4(WU_FP8_E4M3_ONE, WU_FP8_E4M3_ONE, WU_FP8_E4M3_ONE, \ + WU_FP8_E4M3_ONE) +#define WU_FP8_E4M3_ONE_OVER_32_PACKED \ + WU_FP8_PACK4(WU_FP8_E4M3_ONE_OVER_32, WU_FP8_E4M3_ONE_OVER_32, \ + WU_FP8_E4M3_ONE_OVER_32, WU_FP8_E4M3_ONE_OVER_32) + +static inline void wu_bwgmma_fp8(uint32_t addr_tmem_c, uint32_t addr_tmem_a, + uint32_t addr_smem_b) { + asm volatile(".insn r %0, 0, 0, %1, %2, %3" + : + : "i"(RISCV_CUSTOM3), "r"(addr_tmem_c), "r"(addr_tmem_a), + "r"(addr_smem_b) + : "memory"); +} + +static inline void wu_bwgmma_wait() { + asm volatile(".insn r %0, 1, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3) + : "memory"); +} + +#endif diff --git a/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/kernel.cpp b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/kernel.cpp new file mode 100644 index 00000000..a1bb5391 --- /dev/null +++ b/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/kernel.cpp @@ -0,0 +1,218 @@ +#define WU_CASE_WAIT_SPIN 16384u + +#include "fp8_common.hpp" +#include "../../wu_arch_cases/common_wu_blackwell_fa.h" + +#define WU_FP8_SOFTMAX_SCORE_READY_BASE 0xc500u +#define WU_FP8_SOFTMAX_P_READY_BASE 0xc600u +#define WU_FP8_SOFTMAX_DONE_BASE 0xc700u + +#define WU_FP8_SOFTMAX_FP32_ZERO 0x00000000u +#define WU_FP8_SOFTMAX_FP32_ONE 0x3f800000u +#define WU_FP8_SOFTMAX_FP32_THIRTY_TWO 0x42000000u + +extern "C" { +volatile uint32_t g_fp8_softmax_q_frag[4] __attribute__((aligned(16))) = { + WU_FP8_REP4(WU_FP8_E4M3_ONE_PACKED)}; +volatile uint32_t g_fp8_softmax_zero_frag[4] __attribute__((aligned(16))) = { + WU_FP8_SOFTMAX_FP32_ZERO, WU_FP8_SOFTMAX_FP32_ZERO, + WU_FP8_SOFTMAX_FP32_ZERO, WU_FP8_SOFTMAX_FP32_ZERO}; +volatile uint32_t g_fp8_softmax_score_ready __attribute__((aligned(16))); +volatile uint32_t g_fp8_softmax_p_ready __attribute__((aligned(16))); +volatile uint32_t g_fp8_softmax_done __attribute__((aligned(16))); +volatile uint32_t g_fp8_softmax_score_bits[NUM_THREADS] + __attribute__((aligned(16))); +volatile uint32_t g_fp8_softmax_p_bits[NUM_THREADS] + __attribute__((aligned(16))); +volatile uint32_t g_fp8_softmax_out[WU_BW_OUT_WORDS] + __attribute__((aligned(16))); +} + +#undef WU_FP8_REP2 +#undef WU_FP8_REP4 + +static inline int wu_fp8_softmax_wait_status(volatile uint32_t *status, + uint32_t expected) { + for (uint32_t spin = 0; spin < WU_CASE_WAIT_SPIN; ++spin) { + if (*status == expected) { + return 0; + } + } + return 1; +} + +extern "C" void __attribute__((naked, noinline, used)) +tensor_fp8_softmax_worker() { + asm volatile( + "csrr x5, %[csr_wid]\n\t" + "addi x1, x5, -%[num_scalar_warps]\n\t" + "slli x1, x1, 11\n\t" + "addi x2, x1, %[c_offset]\n\t" + "la x3, g_fp8_softmax_q_frag\n\t" + "li x7, 0\n\t" + "1:\n\t" + "add x4, x1, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 1b\n\t" + "la x3, g_fp8_softmax_zero_frag\n\t" + "li x7, 0\n\t" + "2:\n\t" + "add x4, x2, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 2b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "li x4, %[smem_base]\n\t" + ".insn r %[custom3], 0, 0, x2, x1, x4\n\t" + ".insn r %[custom3], 1, 0, x0, x0, x0\n\t" + "la x7, g_fp8_softmax_score_ready\n\t" + "li x6, %[score_ready]\n\t" + "sw x6, 0(x7)\n\t" + "la x7, g_fp8_softmax_p_ready\n\t" + "li x4, %[p_ready]\n\t" + "3:\n\t" + "lw x6, 0(x7)\n\t" + "bne x6, x4, 3b\n\t" + "la x3, g_fp8_softmax_zero_frag\n\t" + "li x7, 0\n\t" + "4:\n\t" + "add x4, x2, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 4b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "li x4, %[smem_base]\n\t" + ".insn r %[custom3], 0, 0, x2, x1, x4\n\t" + ".insn r %[custom3], 1, 0, x0, x0, x0\n\t" + "la x3, g_fp8_softmax_out\n\t" + "li x7, 0\n\t" + "5:\n\t" + "add x4, x2, x7\n\t" + "add x6, x3, x7\n\t" + ".insn r %[custom3], 6, 0, x0, x4, x6\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 5b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "la x7, g_fp8_softmax_done\n\t" + "li x6, %[done]\n\t" + "sw x6, 0(x7)\n\t" + ".insn r %[custom0], 0, 0, x0, x0, x0\n\t" + "6: j 6b\n\t" + : + : [csr_wid] "i"(VX_CSR_WARP_ID), [custom0] "i"(RISCV_CUSTOM0), + [custom3] "i"(RISCV_CUSTOM3), + [num_scalar_warps] "i"(NUM_SCALAR_WARPS), + [c_offset] "i"(WU_BW_TMEM_C_BYTE_OFFSET), + [tile_bytes] "i"(WU_BW_TMEM_TILE_BYTES), + [smem_base] "i"(WU_BW_DEV_SMEM_START_ADDR), + [score_ready] "i"(WU_FP8_SOFTMAX_SCORE_READY_BASE), + [p_ready] "i"(WU_FP8_SOFTMAX_P_READY_BASE), + [done] "i"(WU_FP8_SOFTMAX_DONE_BASE) + : "memory"); +} + +extern "C" int wu_main() { + if (vx_core_id() != 0 || vx_warp_id() != 0) { + return 0; + } + + const uint32_t tid = wu_tid(); + if (tid == 0) { + wu_case_reset(); + g_fp8_softmax_score_ready = 0; + g_fp8_softmax_p_ready = 0; + g_fp8_softmax_done = 0; + for (uint32_t i = 0; i < NUM_THREADS; ++i) { + g_fp8_softmax_score_bits[i] = 0; + g_fp8_softmax_p_bits[i] = 0; + } + for (uint32_t i = 0; i < WU_BW_OUT_WORDS; ++i) { + g_fp8_softmax_out[i] = 0; + } + wu_bw_fill_smem_tile( + reinterpret_cast(WU_BW_DEV_SMEM_START_ADDR), + WU_FP8_E4M3_ONE_PACKED); + vx_spawn_tensor(1u << NUM_SCALAR_WARPS, tensor_fp8_softmax_worker); + } + asm volatile("fence rw, rw" ::: "memory"); + + if (wu_fp8_softmax_wait_status(&g_fp8_softmax_score_ready, + WU_FP8_SOFTMAX_SCORE_READY_BASE) != 0) { + if (tid == 0) { + g_case_mem[1] = 0x81u; + } + } + asm volatile("fence rw, rw" ::: "memory"); + + const uint32_t c_frag = + wu_bw_tmem_c_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES; + const uint32_t observed_score = wu_bw_scalar_tmem_ld(c_frag); + g_fp8_softmax_score_bits[tid] = observed_score; + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0 && g_case_mem[1] == 0 && + observed_score != WU_FP8_SOFTMAX_FP32_THIRTY_TWO) { + g_aux[0] = observed_score; + g_case_mem[1] = 0x82u; + } + asm volatile("fence rw, rw" ::: "memory"); + + uint32_t softmax_token = 0; + const uint32_t p_frag = + wu_bw_tmem_a_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES; + if (g_case_mem[1] == 0) { + vx_tmc(wu_bw_all_lanes_mask()); + softmax_token = wu_bw_scalar_tmem_softmax(c_frag, p_frag); + vx_tmc_one(); + } + asm volatile("fence rw, rw" ::: "memory"); + + const uint32_t observed_p = wu_bw_scalar_tmem_ld(p_frag); + g_fp8_softmax_p_bits[tid] = observed_p; + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0 && g_case_mem[1] == 0) { + if (softmax_token != p_frag) { + g_aux[0] = softmax_token; + g_aux[1] = p_frag; + g_case_mem[1] = 0x85u; + } else if (observed_p != WU_FP8_E4M3_ONE_OVER_32_PACKED) { + g_aux[0] = observed_p; + g_aux[1] = WU_FP8_E4M3_ONE_OVER_32_PACKED; + g_case_mem[1] = 0x86u; + } else { + g_fp8_softmax_p_ready = WU_FP8_SOFTMAX_P_READY_BASE; + } + } + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0) { + if (g_case_mem[1] == 0 && + wu_fp8_softmax_wait_status(&g_fp8_softmax_done, + WU_FP8_SOFTMAX_DONE_BASE) != 0) { + g_case_mem[1] = 0x83u; + } + if (g_case_mem[1] == 0) { + volatile uint32_t bad_actual = 0; + const uint32_t bad = + wu_bw_verify_constant(g_fp8_softmax_out, WU_BW_OUT_WORDS, + WU_FP8_SOFTMAX_FP32_ONE, &bad_actual); + if (bad != WU_BW_OUT_WORDS) { + g_aux[0] = bad; + g_aux[1] = bad_actual; + g_case_mem[1] = 0x84u; + } + } + if (g_case_mem[1] != 0) { + wu_case_fail(g_case_mem[1]); + return 1; + } + wu_case_pass(); + } + return 0; +} diff --git a/kernels/wu_arch_cases/Makefile b/kernels/wu_arch_cases/Makefile index 82eadab6..cc763132 100644 --- a/kernels/wu_arch_cases/Makefile +++ b/kernels/wu_arch_cases/Makefile @@ -25,7 +25,8 @@ CASES := \ case21_moe_gating \ case22_gemm_silu \ case23_softmax_only \ - case24_flash_sw_pipeline + case24_flash_sw_pipeline \ + case25_scalar_tmem_softmax_vector SMOKE_CASES := \ case00_boot_scalar \ diff --git a/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/Makefile b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/Makefile new file mode 100644 index 00000000..58b67532 --- /dev/null +++ b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/Makefile @@ -0,0 +1,3 @@ +PROJECT = case25_scalar_tmem_softmax_vector + +include ../case.mk diff --git a/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/README.md b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/README.md new file mode 100644 index 00000000..82eb7e94 --- /dev/null +++ b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/README.md @@ -0,0 +1,11 @@ +# case25_scalar_tmem_softmax_vector + +Validates the scalar-subcore vector softmax TMEM path. + +The tensor warp writes one score tile into its TMEM C partition, scalar warp 0 +executes `wu_bw_scalar_tmem_softmax()` to read the C tile and write packed fp16 +probabilities into TMEM A. The softmax instruction returns the destination +fragment as a completion token only after the P tile has been written; scalar +warp 0 uses that token before publishing `p_ready`, then the tensor warp +consumes that P tile through BWGMMA. The case uses constant scores so the +expected probability is uniform `1/32` and the final PV output is fp32 `1.0`. diff --git a/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/kernel.cpp b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/kernel.cpp new file mode 100644 index 00000000..3c38a966 --- /dev/null +++ b/kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/kernel.cpp @@ -0,0 +1,198 @@ +#define WU_CASE_WAIT_SPIN 16384u + +#include "../common_wu_blackwell_fa.h" + +#define WU_CASE25_SCORE_READY_BASE 0xb500u +#define WU_CASE25_P_READY_BASE 0xb600u +#define WU_CASE25_DONE_BASE 0xb700u +#define WU_CASE25_FP32_ZERO 0x00000000u +#define WU_CASE25_FP32_THIRTY_SIX 0x42100000u + +extern "C" { +volatile uint32_t g_case25_q_row[4] __attribute__((aligned(16))) = { + WU_BW_FP16_ONE_PACKED, WU_BW_FP16_ONE_PACKED, WU_BW_FP16_ONE_PACKED, + WU_BW_FP16_ONE_PACKED}; +volatile uint32_t g_case25_zero_row[4] __attribute__((aligned(16))) = { + WU_CASE25_FP32_ZERO, WU_CASE25_FP32_ZERO, WU_CASE25_FP32_ZERO, + WU_CASE25_FP32_ZERO}; +volatile uint32_t g_case25_score_ready __attribute__((aligned(16))); +volatile uint32_t g_case25_p_ready __attribute__((aligned(16))); +volatile uint32_t g_case25_done __attribute__((aligned(16))); +volatile uint32_t g_case25_score_bits[NUM_THREADS] __attribute__((aligned(16))); +volatile uint32_t g_case25_out[WU_BW_OUT_WORDS] __attribute__((aligned(16))); +} + +static inline int wu_case25_wait_status(volatile uint32_t *status, + uint32_t expected) { + for (uint32_t spin = 0; spin < WU_CASE_WAIT_SPIN; ++spin) { + if (*status == expected) { + return 0; + } + } + return 1; +} + +extern "C" void __attribute__((naked, noinline, used)) tensor_case25_worker() { + asm volatile( + "csrr x5, %[csr_wid]\n\t" + "addi x1, x5, -%[num_scalar_warps]\n\t" + "slli x1, x1, 11\n\t" + "addi x2, x1, %[c_offset]\n\t" + "la x3, g_case25_q_row\n\t" + "li x7, 0\n\t" + "1:\n\t" + "add x4, x1, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 1b\n\t" + "la x3, g_case25_zero_row\n\t" + "li x7, 0\n\t" + "2:\n\t" + "add x4, x2, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 2b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "li x4, %[smem_base]\n\t" + ".insn r %[custom3], 0, 0, x2, x1, x4\n\t" + ".insn r %[custom3], 1, 0, x0, x0, x0\n\t" + "la x7, g_case25_score_ready\n\t" + "li x6, %[score_ready]\n\t" + "sw x6, 0(x7)\n\t" + "la x7, g_case25_p_ready\n\t" + "li x4, %[p_ready]\n\t" + "3:\n\t" + "lw x6, 0(x7)\n\t" + "bne x6, x4, 3b\n\t" + "la x3, g_case25_zero_row\n\t" + "li x7, 0\n\t" + "4:\n\t" + "add x4, x2, x7\n\t" + ".insn r %[custom3], 2, 0, x0, x4, x3\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 4b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "li x4, %[smem_base]\n\t" + ".insn r %[custom3], 0, 0, x2, x1, x4\n\t" + ".insn r %[custom3], 1, 0, x0, x0, x0\n\t" + "la x3, g_case25_out\n\t" + "li x7, 0\n\t" + "5:\n\t" + "add x4, x2, x7\n\t" + "add x6, x3, x7\n\t" + ".insn r %[custom3], 6, 0, x0, x4, x6\n\t" + "addi x7, x7, 16\n\t" + "li x4, %[tile_bytes]\n\t" + "blt x7, x4, 5b\n\t" + ".insn r %[custom3], 3, 0, x0, x0, x0\n\t" + "la x7, g_case25_done\n\t" + "li x6, %[done]\n\t" + "sw x6, 0(x7)\n\t" + ".insn r %[custom0], 0, 0, x0, x0, x0\n\t" + "6: j 6b\n\t" + : + : [csr_wid] "i"(VX_CSR_WARP_ID), [custom0] "i"(RISCV_CUSTOM0), + [custom3] "i"(RISCV_CUSTOM3), + [num_scalar_warps] "i"(NUM_SCALAR_WARPS), + [c_offset] "i"(WU_BW_TMEM_C_BYTE_OFFSET), + [tile_bytes] "i"(WU_BW_TMEM_TILE_BYTES), + [smem_base] "i"(WU_BW_DEV_SMEM_START_ADDR), + [score_ready] "i"(WU_CASE25_SCORE_READY_BASE), + [p_ready] "i"(WU_CASE25_P_READY_BASE), + [done] "i"(WU_CASE25_DONE_BASE) + : "memory"); +} + +extern "C" int wu_main() { + if (vx_core_id() != 0 || vx_warp_id() != 0) { + return 0; + } + + const uint32_t tid = wu_tid(); + if (tid == 0) { + wu_case_reset(); + g_case25_score_ready = 0; + g_case25_p_ready = 0; + g_case25_done = 0; + for (uint32_t i = 0; i < NUM_THREADS; ++i) { + g_case25_score_bits[i] = 0; + } + for (uint32_t i = 0; i < WU_BW_OUT_WORDS; ++i) { + g_case25_out[i] = 0; + } + wu_bw_fill_smem_tile( + reinterpret_cast(WU_BW_DEV_SMEM_START_ADDR), + WU_BW_FP16_ONE_PACKED); + vx_spawn_tensor(1u << NUM_SCALAR_WARPS, tensor_case25_worker); + } + asm volatile("fence rw, rw" ::: "memory"); + + if (wu_case25_wait_status(&g_case25_score_ready, + WU_CASE25_SCORE_READY_BASE) != 0) { + if (tid == 0) { + g_case_mem[1] = 0x91u; + } + } + asm volatile("fence rw, rw" ::: "memory"); + + const uint32_t c_frag = + wu_bw_tmem_c_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES; + const uint32_t observed = wu_bw_scalar_tmem_ld(c_frag); + g_case25_score_bits[tid] = observed; + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0 && g_case_mem[1] == 0 && + observed != WU_CASE25_FP32_THIRTY_SIX) { + g_aux[0] = observed; + g_case_mem[1] = 0x92u; + } + asm volatile("fence rw, rw" ::: "memory"); + + uint32_t softmax_token = 0; + const uint32_t p_frag = + wu_bw_tmem_a_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES; + if (g_case_mem[1] == 0) { + vx_tmc(wu_bw_all_lanes_mask()); + softmax_token = wu_bw_scalar_tmem_softmax(c_frag, p_frag); + vx_tmc_one(); + } + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0 && g_case_mem[1] == 0) { + if (softmax_token != p_frag) { + g_aux[0] = softmax_token; + g_aux[1] = p_frag; + g_case_mem[1] = 0x95u; + } else { + g_case25_p_ready = WU_CASE25_P_READY_BASE; + } + } + asm volatile("fence rw, rw" ::: "memory"); + + if (tid == 0) { + if (g_case_mem[1] == 0 && + wu_case25_wait_status(&g_case25_done, WU_CASE25_DONE_BASE) != 0) { + g_case_mem[1] = 0x93u; + } + if (g_case_mem[1] == 0) { + volatile uint32_t bad_actual = 0; + const uint32_t bad = + wu_bw_verify_constant(g_case25_out, WU_BW_OUT_WORDS, + WU_BW_FP32_ONE, &bad_actual); + if (bad != WU_BW_OUT_WORDS) { + g_aux[0] = bad; + g_aux[1] = bad_actual; + g_case_mem[1] = 0x94u; + } + } + if (g_case_mem[1] != 0) { + wu_case_fail(g_case_mem[1]); + return 1; + } + wu_case_pass(); + } + return 0; +} diff --git a/kernels/wu_arch_cases/common_wu_blackwell_fa.h b/kernels/wu_arch_cases/common_wu_blackwell_fa.h index 68f4c029..b2aad4d2 100644 --- a/kernels/wu_arch_cases/common_wu_blackwell_fa.h +++ b/kernels/wu_arch_cases/common_wu_blackwell_fa.h @@ -50,6 +50,17 @@ static inline void wu_bw_scalar_tmem_st(uint32_t frag_addr, uint32_t value) { : "memory"); } +static inline uint32_t wu_bw_scalar_tmem_softmax(uint32_t score_frag_addr, + uint32_t p_frag_addr) { + uint32_t token; + asm volatile(".insn r %[custom1], 3, 0x30, %[token], %[score], %[prob]" + : [token] "=r"(token) + : [custom1] "i"(RISCV_CUSTOM1), [score] "r"(score_frag_addr), + [prob] "r"(p_frag_addr) + : "memory"); + return token; +} + static inline uint32_t wu_bw_tmem_a_byte_base(uint32_t tensor_block) { return tensor_block * WU_BW_TMEM_WARP_BYTES; }