tensor: Change B in-memory layout to column-major

This commit is contained in:
Hansung Kim
2024-08-12 15:20:55 -07:00
parent 07dd9e35a0
commit 95e3e96c6c
6 changed files with 42 additions and 24 deletions

View File

@@ -122,14 +122,24 @@ void vx_wmma_load() {
// load B
// B is stored row-major in the memory,
// loaded column-major into the RF.
asm volatile("flw f8 , %0" ::"m"(B[DIM_N * 0 + col]));
asm volatile("flw f9 , %0" ::"m"(B[DIM_N * 1 + col]));
asm volatile("flw f10, %0" ::"m"(B[DIM_N * 2 + col]));
asm volatile("flw f11, %0" ::"m"(B[DIM_N * 3 + col]));
asm volatile("flw f12, %0" ::"m"(B[DIM_N * 4 + col]));
asm volatile("flw f13, %0" ::"m"(B[DIM_N * 5 + col]));
asm volatile("flw f14, %0" ::"m"(B[DIM_N * 6 + col]));
asm volatile("flw f15, %0" ::"m"(B[DIM_N * 7 + col]));
// asm volatile("flw f8 , %0" ::"m"(B[DIM_N * 0 + col]));
// asm volatile("flw f9 , %0" ::"m"(B[DIM_N * 1 + col]));
// asm volatile("flw f10, %0" ::"m"(B[DIM_N * 2 + col]));
// asm volatile("flw f11, %0" ::"m"(B[DIM_N * 3 + col]));
// asm volatile("flw f12, %0" ::"m"(B[DIM_N * 4 + col]));
// asm volatile("flw f13, %0" ::"m"(B[DIM_N * 5 + col]));
// asm volatile("flw f14, %0" ::"m"(B[DIM_N * 6 + col]));
// asm volatile("flw f15, %0" ::"m"(B[DIM_N * 7 + col]));
// B is stored column-major in the memory,
// loaded column-major into the RF.
asm volatile("flw f8 , %0" ::"m"(B[DIM_K * row + 0]));
asm volatile("flw f9 , %0" ::"m"(B[DIM_K * row + 1]));
asm volatile("flw f10, %0" ::"m"(B[DIM_K * row + 2]));
asm volatile("flw f11, %0" ::"m"(B[DIM_K * row + 3]));
asm volatile("flw f12, %0" ::"m"(B[DIM_K * row + 4]));
asm volatile("flw f13, %0" ::"m"(B[DIM_K * row + 5]));
asm volatile("flw f14, %0" ::"m"(B[DIM_K * row + 6]));
asm volatile("flw f15, %0" ::"m"(B[DIM_K * row + 7]));
map_c_8lanes(tid, row, col);