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

tensor: Fix tensor unittest kernel

Browse Source
This commit is contained in:
Hansung Kim
2024-07-31 11:49:41 -07:00
parent 29f7290948
commit 285776404f
1 changed files with 58 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
tests/kernel/tensor/main.cpp
View File

@@ -4,19 +4,22 @@
#include <vx_intrinsics.h> #include <vx_intrinsics.h>
#include <stdio.h> #include <stdio.h>
#include <vx_print.h> #include <vx_print.h>
#include "test_data.h"
constexpr int DIM_M = 8; 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)); 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)); 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) { inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) {
const int tg = tid / 4; const int tg = tid / 4;
@@ -118,6 +121,7 @@ void vx_wmma_load() {
map_c_8lanes(tid, row, col); map_c_8lanes(tid, row, col);
// load C // load C
// accum buffer 0
asm volatile("flw f16, %0" ::"m"(C[row + 0][col + 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 f17, %0" ::"m"(C[row + 0][col + 1]));
asm volatile("flw f18, %0" ::"m"(C[row + 2][col + 0])); asm volatile("flw f18, %0" ::"m"(C[row + 2][col + 0]));
@@ -126,6 +130,7 @@ void vx_wmma_load() {
asm volatile("flw f21, %0" ::"m"(C[row + 0][col + 5])); 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 f22, %0" ::"m"(C[row + 2][col + 4]));
asm volatile("flw f23, %0" ::"m"(C[row + 2][col + 5])); 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 f24, %0" ::"m"(C[row + 0][col + 0]));
asm volatile("flw f25, %0" ::"m"(C[row + 0][col + 1])); 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 f26, %0" ::"m"(C[row + 2][col + 0]));
@@ -136,7 +141,7 @@ void vx_wmma_load() {
asm volatile("flw f31, %0" ::"m"(C[row + 2][col + 5])); 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<float *>(0xc0000000UL); float *const results = reinterpret_cast<float *>(0xc0000000UL);
void store_wmma_result() { void store_wmma_result() {
@@ -150,17 +155,8 @@ void store_wmma_result() {
map_c_8lanes(tid, row, col); map_c_8lanes(tid, row, col);
// store C // 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); 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 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 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)])); // asm volatile("fsw f18, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 0)]));
@@ -202,9 +198,9 @@ void wmma() {
vx_wmma_load(); vx_wmma_load();
// #pragma GCC unroll 100 // #pragma GCC unroll 100
// for (int i = 0; i < 100; i++) { // for (int i = 0; i < 100; i++) {
// vx_wmma(); // vx_wmma_acc0();
// } // }
vx_wmma_new(); vx_wmma_acc1();
store_wmma_result(); store_wmma_result();
// print_wmma_result(); // print_wmma_result();
@@ -214,7 +210,8 @@ void wmma() {
int main() { int main() {
const int num_warps = vx_num_warps(); const int num_warps = vx_num_warps();
vx_wspawn(num_warps, wmma); // vx_wspawn(num_warps, wmma);
vx_wspawn(1, wmma);
wmma(); wmma();
vx_wspawn_wait(); vx_wspawn_wait();