From 285776404f552a790718ac09ecdb6c909988496c Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Wed, 31 Jul 2024 11:49:41 -0700 Subject: [PATCH] tensor: Fix tensor unittest kernel --- tests/kernel/tensor/main.cpp | 119 +++++++++++++++++------------------ 1 file changed, 58 insertions(+), 61 deletions(-) diff --git a/tests/kernel/tensor/main.cpp b/tests/kernel/tensor/main.cpp index d90c38be..d1deb67c 100644 --- a/tests/kernel/tensor/main.cpp +++ b/tests/kernel/tensor/main.cpp @@ -4,19 +4,22 @@ #include #include #include +#include "test_data.h" constexpr int DIM_M = 8; -inline void vx_wmma() { +// single "substep" wmma instruction +// use accum buffer 0 (f16-f23) +inline void vx_wmma_acc0() { asm volatile (".insn r %0, 0, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3)); } -inline void vx_wmma_new() { +// single "substep" wmma instruction +// use accum buffer 1 (f24-f31) +inline void vx_wmma_acc1() { asm volatile (".insn r %0, 0, 0, x1, x0, x0" :: "i"(RISCV_CUSTOM3)); } -#include "test_data.h" - inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) { const int tg = tid / 4; @@ -85,58 +88,60 @@ inline constexpr void map_c_8lanes(const int tid, int &row, int &col) { } void vx_wmma_load() { - int tid = vx_thread_id(); - int tg = tid / 4; + int tid = vx_thread_id(); + int tg = tid / 4; - int row = 0; - int col = 0; + int row = 0; + int col = 0; - map_operand_8lanes(tid, row, col); + map_operand_8lanes(tid, row, col); - // load A - // each operand element is read twice by two threadgroups (Sec. III-B); - // i.e. 8 regs * 32 lanes = 256 fp32 elements = 2 * (16 * 8) elements - asm volatile ("flw f0, %0" :: "m"(A[row][0])); - asm volatile ("flw f1, %0" :: "m"(A[row][1])); - asm volatile ("flw f2, %0" :: "m"(A[row][2])); - asm volatile ("flw f3, %0" :: "m"(A[row][3])); - asm volatile ("flw f4, %0" :: "m"(A[row][4])); - asm volatile ("flw f5, %0" :: "m"(A[row][5])); - asm volatile ("flw f6, %0" :: "m"(A[row][6])); - asm volatile ("flw f7, %0" :: "m"(A[row][7])); + // load A + // each operand element is read twice by two threadgroups (Sec. III-B); + // i.e. 8 regs * 32 lanes = 256 fp32 elements = 2 * (16 * 8) elements + asm volatile("flw f0, %0" ::"m"(A[row][0])); + asm volatile("flw f1, %0" ::"m"(A[row][1])); + asm volatile("flw f2, %0" ::"m"(A[row][2])); + asm volatile("flw f3, %0" ::"m"(A[row][3])); + asm volatile("flw f4, %0" ::"m"(A[row][4])); + asm volatile("flw f5, %0" ::"m"(A[row][5])); + asm volatile("flw f6, %0" ::"m"(A[row][6])); + asm volatile("flw f7, %0" ::"m"(A[row][7])); - // load B - asm volatile ("flw f8 , %0" :: "m"(B[0][col])); - asm volatile ("flw f9 , %0" :: "m"(B[1][col])); - asm volatile ("flw f10, %0" :: "m"(B[2][col])); - asm volatile ("flw f11, %0" :: "m"(B[3][col])); - asm volatile ("flw f12, %0" :: "m"(B[4][col])); - asm volatile ("flw f13, %0" :: "m"(B[5][col])); - asm volatile ("flw f14, %0" :: "m"(B[6][col])); - asm volatile ("flw f15, %0" :: "m"(B[7][col])); + // load B + asm volatile("flw f8 , %0" ::"m"(B[0][col])); + asm volatile("flw f9 , %0" ::"m"(B[1][col])); + asm volatile("flw f10, %0" ::"m"(B[2][col])); + asm volatile("flw f11, %0" ::"m"(B[3][col])); + asm volatile("flw f12, %0" ::"m"(B[4][col])); + asm volatile("flw f13, %0" ::"m"(B[5][col])); + asm volatile("flw f14, %0" ::"m"(B[6][col])); + asm volatile("flw f15, %0" ::"m"(B[7][col])); - map_c_8lanes(tid, row, col); + map_c_8lanes(tid, row, col); - // load C - asm volatile ("flw f16, %0" :: "m"(C[row+0][col+0])); - asm volatile ("flw f17, %0" :: "m"(C[row+0][col+1])); - asm volatile ("flw f18, %0" :: "m"(C[row+2][col+0])); - asm volatile ("flw f19, %0" :: "m"(C[row+2][col+1])); - asm volatile ("flw f20, %0" :: "m"(C[row+0][col+4])); - asm volatile ("flw f21, %0" :: "m"(C[row+0][col+5])); - asm volatile ("flw f22, %0" :: "m"(C[row+2][col+4])); - asm volatile ("flw f23, %0" :: "m"(C[row+2][col+5])); - asm volatile ("flw f24, %0" :: "m"(C[row+0][col+0])); - asm volatile ("flw f25, %0" :: "m"(C[row+0][col+1])); - asm volatile ("flw f26, %0" :: "m"(C[row+2][col+0])); - asm volatile ("flw f27, %0" :: "m"(C[row+2][col+1])); - asm volatile ("flw f28, %0" :: "m"(C[row+0][col+4])); - asm volatile ("flw f29, %0" :: "m"(C[row+0][col+5])); - asm volatile ("flw f30, %0" :: "m"(C[row+2][col+4])); - asm volatile ("flw f31, %0" :: "m"(C[row+2][col+5])); + // load C + // accum buffer 0 + asm volatile("flw f16, %0" ::"m"(C[row + 0][col + 0])); + asm volatile("flw f17, %0" ::"m"(C[row + 0][col + 1])); + asm volatile("flw f18, %0" ::"m"(C[row + 2][col + 0])); + asm volatile("flw f19, %0" ::"m"(C[row + 2][col + 1])); + asm volatile("flw f20, %0" ::"m"(C[row + 0][col + 4])); + asm volatile("flw f21, %0" ::"m"(C[row + 0][col + 5])); + asm volatile("flw f22, %0" ::"m"(C[row + 2][col + 4])); + asm volatile("flw f23, %0" ::"m"(C[row + 2][col + 5])); + // accum buffer 1 + asm volatile("flw f24, %0" ::"m"(C[row + 0][col + 0])); + asm volatile("flw f25, %0" ::"m"(C[row + 0][col + 1])); + asm volatile("flw f26, %0" ::"m"(C[row + 2][col + 0])); + asm volatile("flw f27, %0" ::"m"(C[row + 2][col + 1])); + asm volatile("flw f28, %0" ::"m"(C[row + 0][col + 4])); + asm volatile("flw f29, %0" ::"m"(C[row + 0][col + 5])); + asm volatile("flw f30, %0" ::"m"(C[row + 2][col + 4])); + asm volatile("flw f31, %0" ::"m"(C[row + 2][col + 5])); } -// float results[32*8]; +// hardcoded device address for result float *const results = reinterpret_cast(0xc0000000UL); void store_wmma_result() { @@ -150,17 +155,8 @@ void store_wmma_result() { map_c_8lanes(tid, row, col); // store C - // asm volatile ("fsw f16, %0" :: "m"(results[tid*8+0])); - // asm volatile ("fsw f17, %0" :: "m"(results[tid*8+1])); - // asm volatile ("fsw f18, %0" :: "m"(results[tid*8+2])); - // asm volatile ("fsw f19, %0" :: "m"(results[tid*8+3])); - // asm volatile ("fsw f20, %0" :: "m"(results[tid*8+4])); - // asm volatile ("fsw f21, %0" :: "m"(results[tid*8+5])); - // asm volatile ("fsw f22, %0" :: "m"(results[tid*8+6])); - // asm volatile ("fsw f23, %0" :: "m"(results[tid*8+7])); - float *const results_wid = results + (DIM_M * DIM_M * wid); - + // uncomment to have two accum buffers in rf // asm volatile("fsw f16, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 0)])); // asm volatile("fsw f17, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 1)])); // asm volatile("fsw f18, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 0)])); @@ -202,9 +198,9 @@ void wmma() { vx_wmma_load(); // #pragma GCC unroll 100 // for (int i = 0; i < 100; i++) { - // vx_wmma(); + // vx_wmma_acc0(); // } - vx_wmma_new(); + vx_wmma_acc1(); store_wmma_result(); // print_wmma_result(); @@ -214,7 +210,8 @@ void wmma() { int main() { const int num_warps = vx_num_warps(); - vx_wspawn(num_warps, wmma); + // vx_wspawn(num_warps, wmma); + vx_wspawn(1, wmma); wmma(); vx_wspawn_wait();