wu-arch/kernels
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Blackwell SGEMM kernel scaffolding

Browse Source
This commit is contained in:
abnerhexu
2026-04-25 10:15:31 +08:00
parent 71f713b9fc
commit bcc566b621
2 changed files with 36 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
kernels/sgemm_tcore/sgemm_impl.hpp
View File

@@ -267,6 +267,34 @@ inline void vx_wgmma_wait() {
asm volatile (".insn r %0, 1, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3)); asm volatile (".insn r %0, 1, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3));
} }
inline void vx_tcgen05_cp(const uint32_t addr_tmem, const uint32_t addr_smem) {
asm volatile(".insn r %0, 2, 0, x0, %1, %2" ::"i"(RISCV_CUSTOM3), "r"(addr_tmem),
"r"(addr_smem));
}
inline void vx_tcgen05_cp_wait() {
asm volatile (".insn r %0, 3, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3));
}
inline void vx_bwgmma(const uint32_t addr_tmem_a, const uint32_t addr_smem_b) {
asm volatile(".insn r %0, 0, 0, x0, %1, %2" ::"i"(RISCV_CUSTOM3), "r"(addr_tmem_a),
"r"(addr_smem_b));
}
inline void vx_bwgmma_wait() {
asm volatile (".insn r %0, 1, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3));
}
inline void vx_tcgen05_ld(const uint32_t addr_tmem, const uint32_t rd_hint) {
asm volatile(".insn r %0, 4, 0, %1, %2, x0" ::"i"(RISCV_CUSTOM3), "r"(rd_hint),
"r"(addr_tmem));
}
inline void vx_tcgen05_st(const uint32_t addr_tmem, const uint32_t rd_hint) {
asm volatile(".insn r %0, 5, 0, %1, %2, x0" ::"i"(RISCV_CUSTOM3), "r"(rd_hint),
"r"(addr_tmem));
}
// Remap logical row/col coordinate of a matrix element to a memory index that // Remap logical row/col coordinate of a matrix element to a memory index that
// follows the 2-level block-row-major layout that Gemmini DMA uses // follows the 2-level block-row-major layout that Gemmini DMA uses
template <bool use_dma, uint32_t dim_col> template <bool use_dma, uint32_t dim_col>

8
kernels/sgemm_tcore_blackwell/Makefile Normal file
View File

@@ -0,0 +1,8 @@
PROJECT = sgemm_tcore_blackwell
VX_SRCS = ../sgemm_tcore/kernel.cpp
VX_INCLUDES = ../sgemm_tcore/sgemm_impl.hpp ../sgemm_tcore/common.h
OPTS ?= -n16
include ../common.mk