Merge branch 'kernels' of https://github.com/hansungk/vortex-private into kernels
This commit is contained in:
Richard Yan
@@ -14,7 +14,8 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
TOOLDIR=${TOOLDIR:=/opt}
|
||||
TOOLDIR=${TOOLDIR:=/scratch/hansung/build/vortex-toolchain-prebuilt}
|
||||
export TOOLDIR
|
||||
|
||||
export VERILATOR_ROOT=$TOOLDIR/verilator
|
||||
export PATH=$VERILATOR_ROOT/bin:$PATH
|
||||
@@ -25,6 +26,8 @@ export PATH=$SV2V_PATH/bin:$PATH
|
||||
export YOSYS_PATH=$TOOLDIR/yosys
|
||||
export PATH=$YOSYS_PATH/bin:$PATH
|
||||
|
||||
export LLVM_VORTEX=$TOOLDIR/llvm-vortex
|
||||
export POCL_CC_PATH=$TOOLDIR/pocl/compiler
|
||||
export POCL_RT_PATH=$TOOLDIR/pocl/runtime
|
||||
# LLVM_POCL seems to be only used in tests/opencl
|
||||
export LLVM_POCL=/scratch/hansung/build/llvm-vortex2
|
||||
export LLVM_VORTEX=/scratch/hansung/build/llvm-vortex2
|
||||
export POCL_CC_PATH=/scratch/hansung/build/pocl-vortex2/compiler
|
||||
export POCL_RT_PATH=/scratch/hansung/build/pocl-vortex2/runtime
|
||||
|
||||
@@ -13,6 +13,7 @@ MEMORY {
|
||||
DRAMARG (rwx): ORIGIN = 0x9fff0000, LENGTH = 8K
|
||||
DRAM1 (rwx): ORIGIN = 0xa0000000, LENGTH = 16M
|
||||
DRAM2 (rwx): ORIGIN = 0xa1000000, LENGTH = 16M
|
||||
DRAM3 (rwx): ORIGIN = 0xa2000000, LENGTH = 16M
|
||||
}
|
||||
|
||||
SECTIONS
|
||||
@@ -288,4 +289,8 @@ SECTIONS
|
||||
*(.operand.b)
|
||||
. += 32K;
|
||||
}> DRAM2
|
||||
.operand.c : {
|
||||
*(.operand.c)
|
||||
. += 32K;
|
||||
}> DRAM3
|
||||
}
|
||||
|
||||
@@ -6,3 +6,17 @@ DEPS += b_matrix.h
|
||||
DEPS += c_matrix.h
|
||||
|
||||
include ../common.mk
|
||||
|
||||
OBJCOPY ?= $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-objcopy
|
||||
OBJCOPY_FLAGS ?= "LOAD,ALLOC,DATA,CONTENTS"
|
||||
BINFILES := args.bin input.a.bin input.b.bin
|
||||
$(PROJECT).elf: $(SRCS) $(DEPS)
|
||||
$(CC) $(CFLAGS) $(SRCS) $(LDFLAGS) -o $(PROJECT).elf
|
||||
$(OBJCOPY) --set-section-flags .operand.a=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.b=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.c=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .args=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --update-section .operand.a=input.a.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.b=input.b.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.c=input.c.bin $@ || true
|
||||
$(OBJCOPY) --update-section .args=args.bin $@ || true
|
||||
|
||||
@@ -41,9 +41,13 @@ def C_index(threadgroup, thread, register):
|
||||
def check_sim_output():
|
||||
file = input("simulator output filename: ")
|
||||
|
||||
A_array = np.zeros((16, 8))
|
||||
B_array = np.zeros((8, 16))
|
||||
C_array = np.zeros((16, 16))
|
||||
M = 8
|
||||
N = 8
|
||||
K = 16
|
||||
|
||||
A_array = np.zeros((M, K))
|
||||
B_array = np.zeros((K, N))
|
||||
C_array = np.zeros((M, N))
|
||||
|
||||
with open(file) as f:
|
||||
for line in f.readlines():
|
||||
@@ -87,26 +91,23 @@ def check_sim_output():
|
||||
|
||||
if __name__ == "__main__":
|
||||
expected = np.load("abc.npz")
|
||||
# expected_A = expected['A_array']
|
||||
# expected_B = expected['B_array']
|
||||
# expected_C = expected['C_array']
|
||||
expected_A = expected['A_array'][0:8, 0:8]
|
||||
expected_B = expected['B_array'][0:8, 0:8]
|
||||
expected_C = expected['C_array'][0:8, 0:8]
|
||||
expected_A = expected['A_array']
|
||||
expected_B = expected['B_array']
|
||||
expected_C = expected['C_array']
|
||||
expected_C = expected_C + expected_A @ expected_B
|
||||
print('expected A:')
|
||||
print(expected_A)
|
||||
print('expected B:')
|
||||
print(expected_B)
|
||||
print('expected C:')
|
||||
print(expected_C[0:8, 0:8])
|
||||
print(expected_C)
|
||||
expected_C.astype('float32').tofile("c_expected.bin")
|
||||
|
||||
[got_A, got_B, got_C] = check_sim_output()
|
||||
print('got C:')
|
||||
print(C_array[0:8, 0:8])
|
||||
print(C_array)
|
||||
print('diff C:')
|
||||
print(expected_C[0:8, 0:8] - C_array[0:8, 0:8])
|
||||
print(expected_C - C_array)
|
||||
assert np.allclose(expected_A, got_A)
|
||||
assert np.allclose(expected_B, got_B)
|
||||
assert np.allclose(expected_C, got_C)
|
||||
|
||||
158
tests/kernel/tensor/flash_attn.py
Normal file
158
tests/kernel/tensor/flash_attn.py
Normal file
@@ -0,0 +1,158 @@
|
||||
import sys
|
||||
import numpy as np
|
||||
|
||||
def parse_mnk():
|
||||
if len(sys.argv) != 4:
|
||||
print(f"usage: {sys.argv[0]} dimM dimN dimK", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
m = int(sys.argv[1])
|
||||
n = int(sys.argv[2])
|
||||
k = int(sys.argv[3])
|
||||
return (m, n, k)
|
||||
|
||||
|
||||
# Reorder array in a way that groups two adjacent elements along the column to
|
||||
# be now adjacent along the row. This way, when the resulting fp16 array is
|
||||
# read in column-major order with 32-bit granularity, the fp16 elements will be
|
||||
# read in the same order as regular fp32 elements in column-major.
|
||||
#
|
||||
# For example:
|
||||
# [[1 2]
|
||||
# [3 4]
|
||||
# [5 6]
|
||||
# [7 8]]
|
||||
# becomes
|
||||
# [[1 3 2 4]
|
||||
# [5 7 6 8]]
|
||||
def pack_fp16_by_column(array):
|
||||
rows = array.shape[0]
|
||||
cols = array.shape[1]
|
||||
|
||||
T = array.transpose([1, 0])
|
||||
T_packed = T.reshape([cols, -1, 2])
|
||||
result = T_packed.transpose([1, 0, 2])
|
||||
return result
|
||||
|
||||
|
||||
# Do the same as pack_fp16_by_column, but for every two elements along the row.
|
||||
def pack_fp16_by_row(array):
|
||||
rows = array.shape[0]
|
||||
cols = array.shape[1]
|
||||
|
||||
result = array.reshape([rows, -1, 2])
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
seqlen, _, headdim = parse_mnk()
|
||||
|
||||
rand = True
|
||||
if not rand:
|
||||
A_array = np.arange(seqlen * headdim).reshape([seqlen, headdim])
|
||||
B_array = np.arange(headdim * seqlen).reshape([headdim, seqlen])
|
||||
C_array = np.arange(seqlen * seqlen).reshape([seqlen, headdim])
|
||||
else:
|
||||
np.random.seed(0)
|
||||
A_array = np.random.rand(seqlen, headdim) - 0.5
|
||||
B_array = np.random.rand(headdim, seqlen) - 0.5
|
||||
C_array = np.random.rand(seqlen, headdim) - 0.5
|
||||
# C_array = np.zeros([M, N])
|
||||
|
||||
fp16 = False
|
||||
if fp16:
|
||||
A_packed = pack_fp16_by_row(A_array)
|
||||
AT_packed = A_packed.transpose([1, 0, 2])
|
||||
AT_array = AT_packed.reshape([-1, seqlen * 2])
|
||||
AT_array.astype('float16').tofile("input.a.col.bin")
|
||||
# print('AT:')
|
||||
# print(AT_array)
|
||||
B_packed = pack_fp16_by_column(B_array)
|
||||
B_array = B_packed.reshape([-1, headdim * 2])
|
||||
B_array.astype('float16').tofile("input.b.row.bin")
|
||||
# print('B:')
|
||||
# print(B_array)
|
||||
else:
|
||||
A_array.astype('float32').tofile("input.a.row.bin")
|
||||
AT_array = A_array.transpose([1, 0])
|
||||
AT_array.astype('float32').tofile("input.a.col.bin")
|
||||
B_array.astype('float32').tofile("input.b.bin")
|
||||
C_array.astype('float32').tofile("input.c.bin")
|
||||
# print('AT:')
|
||||
# print(AT_array)
|
||||
# print('B:')
|
||||
# print(B_array)
|
||||
|
||||
assert((seqlen % 64) == 0)
|
||||
|
||||
Br = 64
|
||||
Bc = Br
|
||||
|
||||
rowmax = np.zeros([Br])
|
||||
rowsum = np.zeros([Br])
|
||||
O = np.zeros([Br, headdim])
|
||||
|
||||
def exp2(x):
|
||||
return (x**2) / 2.0 + x + 1.0
|
||||
|
||||
full_S = A_array @ B_array
|
||||
full_S_T = full_S.transpose([1, 0])
|
||||
full_S.astype('float32').tofile("full_S.bin")
|
||||
|
||||
col_to_save = 0
|
||||
|
||||
for col in range(0, seqlen, Bc):
|
||||
print(f"tile iteration {col}~{col + Bc} ======================================")
|
||||
|
||||
# FIXME: only work with the first 64 rows of Q for now
|
||||
Q_tile = A_array[0:64, :]
|
||||
K_tile = B_array[:, col:col+Bc]
|
||||
|
||||
S = Q_tile @ K_tile
|
||||
if col == col_to_save:
|
||||
print('S_expected:')
|
||||
print(S)
|
||||
S.astype('float32').tofile("S_expected.bin")
|
||||
|
||||
# generate rowmax result in online softmax
|
||||
rowmax_this = np.max(S, axis=1)
|
||||
rowmax_prev = rowmax.copy()
|
||||
rowmax = np.maximum(rowmax, rowmax_this)
|
||||
if col == col_to_save:
|
||||
rowmax.astype('float32').tofile("rowmax.bin")
|
||||
|
||||
# subtrace rowmax from each row by broadcasting
|
||||
# (placeholder for exp)
|
||||
x = S - rowmax[:, np.newaxis]
|
||||
P = exp2(x)
|
||||
# for i in range(3, 4):
|
||||
# P += (x**i) / np.math.factorial(i)
|
||||
# P = np.exp(exp)
|
||||
# print('P error:')
|
||||
# print(P / np.exp(x))
|
||||
if col == col_to_save:
|
||||
print('P_expected:')
|
||||
print(P)
|
||||
P.astype('float32').tofile("P_expected.bin")
|
||||
|
||||
rowsum_this = np.sum(P, axis=1)
|
||||
x = rowmax_prev - rowmax_this
|
||||
rowsum = exp2(x) * rowsum + rowsum_this
|
||||
if col == col_to_save:
|
||||
rowsum.astype('float32').tofile("rowsum.bin")
|
||||
|
||||
x = rowmax_prev - rowmax
|
||||
O = O / (exp2(x)[:, np.newaxis])
|
||||
if col == col_to_save:
|
||||
print('O_before_PV:')
|
||||
print(O)
|
||||
O.astype('float32').tofile("O_before_PV.bin")
|
||||
|
||||
V = C_array[col:col+Bc, :]
|
||||
if col == col_to_save:
|
||||
V.astype('float32').tofile("V_expected.bin")
|
||||
# O = P.transpose([1, 0]) @ V
|
||||
O = O + P @ V
|
||||
if col == col_to_save:
|
||||
print('O_after_PV:')
|
||||
print(O)
|
||||
O.astype('float32').tofile("O_after_PV.bin")
|
||||
@@ -1,32 +1,108 @@
|
||||
import sys
|
||||
import numpy as np
|
||||
A_array = np.random.rand(16, 8)
|
||||
B_array = np.random.rand(8, 16)
|
||||
C_array = np.random.rand(16, 16)
|
||||
# A_array = np.zeros((16, 8))
|
||||
# B_array = np.zeros((8, 16))
|
||||
# A_array[0,:] = 1.0
|
||||
# B_array[:,4] = 1.0
|
||||
# C_array = np.zeros((16, 16))
|
||||
# for i in range(16):
|
||||
# for j in range(16):
|
||||
# C_array[i,j] = i * 16 + j
|
||||
|
||||
with open('a_matrix.h', 'w') as f:
|
||||
def parse_mnk():
|
||||
if len(sys.argv) != 4:
|
||||
print(f"usage: {sys.argv[0]} dimM dimN dimK", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
m = int(sys.argv[1])
|
||||
n = int(sys.argv[2])
|
||||
k = int(sys.argv[3])
|
||||
return (m, n, k)
|
||||
|
||||
|
||||
# Reorder array in a way that groups two adjacent elements along the column to
|
||||
# be now adjacent along the row. This way, when the resulting fp16 array is
|
||||
# read in column-major order with 32-bit granularity, the fp16 elements will be
|
||||
# read in the same order as regular fp32 elements in column-major.
|
||||
#
|
||||
# For example:
|
||||
# [[1 2]
|
||||
# [3 4]
|
||||
# [5 6]
|
||||
# [7 8]]
|
||||
# becomes
|
||||
# [[1 3 2 4]
|
||||
# [5 7 6 8]]
|
||||
def pack_fp16_by_column(array):
|
||||
rows = array.shape[0]
|
||||
cols = array.shape[1]
|
||||
|
||||
T = array.transpose([1, 0])
|
||||
T_packed = T.reshape([cols, -1, 2])
|
||||
result = T_packed.transpose([1, 0, 2])
|
||||
return result
|
||||
|
||||
|
||||
# Do the same as pack_fp16_by_column, but for every two elements along the row.
|
||||
def pack_fp16_by_row(array):
|
||||
rows = array.shape[0]
|
||||
cols = array.shape[1]
|
||||
|
||||
result = array.reshape([rows, -1, 2])
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
M, N, K = parse_mnk()
|
||||
|
||||
rand = False
|
||||
if not rand:
|
||||
A_array = np.arange(M * K).reshape([M, K])
|
||||
B_array = np.arange(K * N).reshape([K, N])
|
||||
# C_array = np.arange(M * N).reshape([M, N])
|
||||
C_array = np.zeros([M, N])
|
||||
else:
|
||||
np.random.seed(0)
|
||||
A_array = np.random.rand(M, K) - 0.5
|
||||
B_array = np.random.rand(K, N) - 0.5
|
||||
C_array = np.random.rand(N, K) - 0.5
|
||||
# C_array = np.zeros([M, N])
|
||||
|
||||
with open('a_matrix.h', 'w') as f:
|
||||
for i in range(A_array.shape[0]):
|
||||
for j in range(A_array.shape[1]):
|
||||
f.write(f'{A_array[i,j]}f, ')
|
||||
f.write(f'{A_array[i,j]:f}f, ')
|
||||
f.write('\n')
|
||||
|
||||
with open('b_matrix.h', 'w') as f:
|
||||
with open('b_matrix.h', 'w') as f:
|
||||
for i in range(B_array.shape[0]):
|
||||
for j in range(B_array.shape[1]):
|
||||
f.write(f'{B_array[i,j]}f, ')
|
||||
f.write(f'{B_array[i,j]:f}f, ')
|
||||
f.write('\n')
|
||||
|
||||
with open('c_matrix.h', 'w') as f:
|
||||
with open('c_matrix.h', 'w') as f:
|
||||
for i in range(C_array.shape[0]):
|
||||
for j in range(C_array.shape[1]):
|
||||
f.write(f'{C_array[i,j]}f, ')
|
||||
f.write(f'{C_array[i,j]:f}f, ')
|
||||
f.write('\n')
|
||||
|
||||
np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array)
|
||||
np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array)
|
||||
|
||||
fp16 = False
|
||||
if fp16:
|
||||
A_packed = pack_fp16_by_row(A_array)
|
||||
AT_packed = A_packed.transpose([1, 0, 2])
|
||||
AT_array = AT_packed.reshape([-1, M * 2])
|
||||
AT_array.astype('float16').tofile("input.a.col.bin")
|
||||
print('AT:')
|
||||
print(AT_array)
|
||||
B_packed = pack_fp16_by_column(B_array)
|
||||
B_array = B_packed.reshape([-1, N * 2])
|
||||
B_array.astype('float16').tofile("input.b.row.bin")
|
||||
print('B:')
|
||||
print(B_array)
|
||||
else:
|
||||
A_array.astype('float32').tofile("input.a.row.bin")
|
||||
AT_array = A_array.transpose([1, 0])
|
||||
AT_array.astype('float32').tofile("input.a.col.bin")
|
||||
B_array.astype('float32').tofile("input.b.bin")
|
||||
C_array.astype('float32').tofile("input.c.bin")
|
||||
print('AT:')
|
||||
print(AT_array)
|
||||
print('B:')
|
||||
print(B_array)
|
||||
|
||||
D_expected = A_array @ B_array
|
||||
D_expected.astype('float32').tofile("d_expected.bin")
|
||||
print('D_expected:')
|
||||
print(D_expected)
|
||||
|
||||
|
||||
@@ -4,9 +4,15 @@
|
||||
#include <vx_intrinsics.h>
|
||||
#include <stdio.h>
|
||||
#include <vx_print.h>
|
||||
#include "test_data.h"
|
||||
|
||||
constexpr int DIM_M = 8;
|
||||
constexpr int DIM_N = 8;
|
||||
constexpr int DIM_K = 8;
|
||||
|
||||
// #include "test_data.h"
|
||||
const float *A = reinterpret_cast<const float *>(0xa0000000UL);
|
||||
const float *B = reinterpret_cast<const float *>(0xa1000000UL);
|
||||
const float *C = reinterpret_cast<const float *>(0xa2000000UL);
|
||||
|
||||
// single "substep" wmma instruction
|
||||
// use accum buffer 0 (f16-f23)
|
||||
@@ -97,48 +103,64 @@ void vx_wmma_load() {
|
||||
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]));
|
||||
// A is stored K-major in the memory,
|
||||
// loaded K-major into the RF.
|
||||
//
|
||||
// For 32 lanes config, 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[DIM_K * row + 0]));
|
||||
asm volatile("flw f1, %0" ::"m"(A[DIM_K * row + 1]));
|
||||
asm volatile("flw f2, %0" ::"m"(A[DIM_K * row + 2]));
|
||||
asm volatile("flw f3, %0" ::"m"(A[DIM_K * row + 3]));
|
||||
asm volatile("flw f4, %0" ::"m"(A[DIM_K * row + 4]));
|
||||
asm volatile("flw f5, %0" ::"m"(A[DIM_K * row + 5]));
|
||||
asm volatile("flw f6, %0" ::"m"(A[DIM_K * row + 6]));
|
||||
asm volatile("flw f7, %0" ::"m"(A[DIM_K * 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]));
|
||||
// B is stored N-major in the memory,
|
||||
// loaded K-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]));
|
||||
// B is stored K-major in the memory,
|
||||
// loaded K-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);
|
||||
|
||||
// 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]));
|
||||
asm volatile("flw f16, %0" ::"m"(C[DIM_N * (row + 0) + col + 0]));
|
||||
asm volatile("flw f17, %0" ::"m"(C[DIM_N * (row + 0) + col + 1]));
|
||||
asm volatile("flw f18, %0" ::"m"(C[DIM_N * (row + 2) + col + 0]));
|
||||
asm volatile("flw f19, %0" ::"m"(C[DIM_N * (row + 2) + col + 1]));
|
||||
asm volatile("flw f20, %0" ::"m"(C[DIM_N * (row + 0) + col + 4]));
|
||||
asm volatile("flw f21, %0" ::"m"(C[DIM_N * (row + 0) + col + 5]));
|
||||
asm volatile("flw f22, %0" ::"m"(C[DIM_N * (row + 2) + col + 4]));
|
||||
asm volatile("flw f23, %0" ::"m"(C[DIM_N * (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]));
|
||||
asm volatile("flw f24, %0" ::"m"(C[DIM_N * (row + 0) + col + 0]));
|
||||
asm volatile("flw f25, %0" ::"m"(C[DIM_N * (row + 0) + col + 1]));
|
||||
asm volatile("flw f26, %0" ::"m"(C[DIM_N * (row + 2) + col + 0]));
|
||||
asm volatile("flw f27, %0" ::"m"(C[DIM_N * (row + 2) + col + 1]));
|
||||
asm volatile("flw f28, %0" ::"m"(C[DIM_N * (row + 0) + col + 4]));
|
||||
asm volatile("flw f29, %0" ::"m"(C[DIM_N * (row + 0) + col + 5]));
|
||||
asm volatile("flw f30, %0" ::"m"(C[DIM_N * (row + 2) + col + 4]));
|
||||
asm volatile("flw f31, %0" ::"m"(C[DIM_N * (row + 2) + col + 5]));
|
||||
}
|
||||
|
||||
// hardcoded device address for result
|
||||
@@ -155,24 +177,24 @@ void store_wmma_result() {
|
||||
map_c_8lanes(tid, row, col);
|
||||
|
||||
// store C
|
||||
float *const results_wid = results + (DIM_M * DIM_M * wid);
|
||||
float *const results_wid = results + (DIM_M * DIM_N * 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)]));
|
||||
// asm volatile("fsw f19, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 1)]));
|
||||
// asm volatile("fsw f20, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 4)]));
|
||||
// asm volatile("fsw f21, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 5)]));
|
||||
// asm volatile("fsw f22, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 4)]));
|
||||
// asm volatile("fsw f23, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 5)]));
|
||||
asm volatile("fsw f24, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 0)]));
|
||||
asm volatile("fsw f25, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 1)]));
|
||||
asm volatile("fsw f26, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 0)]));
|
||||
asm volatile("fsw f27, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 1)]));
|
||||
asm volatile("fsw f28, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 4)]));
|
||||
asm volatile("fsw f29, %0" ::"m"(results_wid[DIM_M * (row + 0) + (col + 5)]));
|
||||
asm volatile("fsw f30, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 4)]));
|
||||
asm volatile("fsw f31, %0" ::"m"(results_wid[DIM_M * (row + 2) + (col + 5)]));
|
||||
// asm volatile("fsw f16, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 0)]));
|
||||
// asm volatile("fsw f17, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 1)]));
|
||||
// asm volatile("fsw f18, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 0)]));
|
||||
// asm volatile("fsw f19, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 1)]));
|
||||
// asm volatile("fsw f20, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 4)]));
|
||||
// asm volatile("fsw f21, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 5)]));
|
||||
// asm volatile("fsw f22, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 4)]));
|
||||
// asm volatile("fsw f23, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 5)]));
|
||||
asm volatile("fsw f24, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 0)]));
|
||||
asm volatile("fsw f25, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 1)]));
|
||||
asm volatile("fsw f26, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 0)]));
|
||||
asm volatile("fsw f27, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 1)]));
|
||||
asm volatile("fsw f28, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 4)]));
|
||||
asm volatile("fsw f29, %0" ::"m"(results_wid[DIM_N * (row + 0) + (col + 5)]));
|
||||
asm volatile("fsw f30, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 4)]));
|
||||
asm volatile("fsw f31, %0" ::"m"(results_wid[DIM_N * (row + 2) + (col + 5)]));
|
||||
}
|
||||
|
||||
void print_wmma_result() {
|
||||
@@ -211,9 +233,8 @@ int main() {
|
||||
const int num_warps = vx_num_warps();
|
||||
|
||||
// vx_wspawn(num_warps, wmma);
|
||||
vx_wspawn(1, wmma);
|
||||
wmma();
|
||||
vx_wspawn_wait();
|
||||
// vx_wspawn_wait();
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -49,7 +49,10 @@ VX_CP = $(LLVM_VORTEX)/bin/llvm-objcopy
|
||||
#VX_CP = $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-objcopy
|
||||
|
||||
VX_CFLAGS += -v -O3 -std=c++17
|
||||
VX_CFLAGS += -mcmodel=medany -fno-rtti -fno-exceptions -nostartfiles -fdata-sections -ffunction-sections -mllvm -inline-threshold=8192
|
||||
VX_CFLAGS += -mcmodel=medany -fno-rtti -fno-exceptions -nostartfiles -fdata-sections -ffunction-sections
|
||||
# comment out below for regression/basic, which uses GCC that doesn't
|
||||
# understand these flags
|
||||
VX_CFLAGS += -mllvm -inline-threshold=262144
|
||||
VX_CFLAGS += -I$(VORTEX_KN_PATH)/include -I$(VORTEX_KN_PATH)/../hw -I$(GEMMINI_SW_PATH)
|
||||
VX_CFLAGS += -DNDEBUG -DLLVM_VORTEX
|
||||
|
||||
@@ -104,23 +107,28 @@ kernel.bin: kernel.elf kernel.radiance.elf
|
||||
|
||||
OBJCOPY ?= $(RISCV_TOOLCHAIN_PATH)/bin/$(RISCV_PREFIX)-objcopy
|
||||
OBJCOPY_FLAGS ?= "LOAD,ALLOC,DATA,CONTENTS"
|
||||
kernel.elf: $(VX_SRCS)
|
||||
$(VX_CXX) $(VX_CFLAGS) $(VX_SRCS) $(VX_LDFLAGS) -o $@
|
||||
BINFILES := args.bin input.a.bin input.b.bin input.c.bin
|
||||
kernel.elf: $(VX_SRCS) $(VX_INCLUDES) $(BINFILES)
|
||||
$(VX_CXX) $(VX_CFLAGS) -o $@ $(VX_SRCS) $(VX_LDFLAGS)
|
||||
$(OBJCOPY) --set-section-flags .operand.a=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.b=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.c=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .args=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --update-section .operand.a=input.a.bin $@
|
||||
$(OBJCOPY) --update-section .operand.b=input.b.bin $@
|
||||
$(OBJCOPY) --update-section .args=args.bin $@
|
||||
$(OBJCOPY) --update-section .operand.a=input.a.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.b=input.b.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.c=input.c.bin $@ || true
|
||||
$(OBJCOPY) --update-section .args=args.bin $@ || true
|
||||
|
||||
kernel.radiance.elf: $(VX_SRCS)
|
||||
kernel.radiance.elf: $(VX_SRCS) $(VX_INCLUDES) $(BINFILES)
|
||||
$(VX_CXX) $(VX_CFLAGS) $(VX_SRCS) $(VX_LDFLAGS) -DRADIANCE -o $@
|
||||
$(OBJCOPY) --set-section-flags .operand.a=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.b=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .operand.c=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --set-section-flags .args=$(OBJCOPY_FLAGS) $@
|
||||
$(OBJCOPY) --update-section .operand.a=input.a.bin $@
|
||||
$(OBJCOPY) --update-section .operand.b=input.b.bin $@
|
||||
$(OBJCOPY) --update-section .args=args.bin $@
|
||||
$(OBJCOPY) --update-section .operand.a=input.a.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.b=input.b.bin $@ || true
|
||||
$(OBJCOPY) --update-section .operand.c=input.c.bin $@ || true
|
||||
$(OBJCOPY) --update-section .args=args.bin $@ || true
|
||||
|
||||
ifneq ($(CONFIG),)
|
||||
kernel$(CONFIGEXT).elf: kernel.elf
|
||||
|
||||
1
tests/regression/flash_attention/.gitignore
vendored
Normal file
1
tests/regression/flash_attention/.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
flash_attention
|
||||
12
tests/regression/flash_attention/Makefile
Normal file
12
tests/regression/flash_attention/Makefile
Normal file
@@ -0,0 +1,12 @@
|
||||
PROJECT = flash_attention
|
||||
|
||||
SRCS = main.cpp common.h
|
||||
|
||||
VX_SRCS = kernel.cpp
|
||||
VX_INCLUDES = ../sgemm_tcore/sgemm_impl.hpp
|
||||
|
||||
OPTS ?= -n16
|
||||
|
||||
VX_CFLAGS += -I../sgemm_tcore
|
||||
|
||||
include ../common.mk
|
||||
18
tests/regression/flash_attention/common.h
Normal file
18
tests/regression/flash_attention/common.h
Normal file
@@ -0,0 +1,18 @@
|
||||
#ifndef _COMMON_H_
|
||||
#define _COMMON_H_
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#define KERNEL_ARG_DEV_MEM_ADDR 0x9fff0000
|
||||
#define DEV_SMEM_START_ADDR 0xff000000
|
||||
|
||||
typedef struct {
|
||||
uint32_t dim_seqlen;
|
||||
uint32_t dim_headdim;
|
||||
uint64_t addr_q;
|
||||
uint64_t addr_k;
|
||||
uint64_t addr_v;
|
||||
uint64_t addr_o;
|
||||
} kernel_arg_t;
|
||||
|
||||
#endif
|
||||
4018
tests/regression/flash_attention/half.hpp
Normal file
4018
tests/regression/flash_attention/half.hpp
Normal file
File diff suppressed because it is too large
Load Diff
950
tests/regression/flash_attention/kernel.cpp
Normal file
950
tests/regression/flash_attention/kernel.cpp
Normal file
@@ -0,0 +1,950 @@
|
||||
#include <stdint.h>
|
||||
#include <vx_intrinsics.h>
|
||||
#include <vx_print.h>
|
||||
#include <vx_spawn.h>
|
||||
#include <float.h>
|
||||
#include "common.h"
|
||||
#include "sgemm_impl.hpp"
|
||||
#include "include/gemmini.h"
|
||||
#include "gemmini_mmio.h"
|
||||
|
||||
#define B_ROW BM
|
||||
#define B_COL BN
|
||||
// FIXME
|
||||
#define HEADDIM B_COL
|
||||
|
||||
constexpr uint32_t ROWMAX_SETS = 3;
|
||||
constexpr bool DEBUG = true;
|
||||
constexpr bool WARP_SPECIALIZED = true;
|
||||
|
||||
constexpr uint32_t DEV_FAKE_SMEM_START_ADDR = 0xf0000000;
|
||||
|
||||
// temporary safety stop for wrong configs
|
||||
static_assert(NUM_CORES == 4);
|
||||
static_assert(NUM_THREADS == 8);
|
||||
static_assert(NUM_WARPS == 8);
|
||||
|
||||
inline void thread_block_init_sharedmem(const uint32_t tid_in_threadblock,
|
||||
const uint32_t threads_per_threadblock,
|
||||
float *smem_O, float *smem_rowmax,
|
||||
float *smem_rowsum,
|
||||
float *smem_O_row_scale) {
|
||||
asm volatile("threadblock_init_sharedmem_start_%=:" ::);
|
||||
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
|
||||
static_assert((B_ROW % NUM_THREADS) == 0,
|
||||
"B_ROW must be a multiple of NUM_THREADS");
|
||||
static_assert(B_ROW < (NUM_THREADS * CORES_PER_CLUSTER *
|
||||
(NUM_WARPS / (WARP_SPECIALIZED ? 2 : 1))),
|
||||
"not enough warps to initialize rowmax/rowsum");
|
||||
|
||||
// each thread initializes one element in rowmax/rowsum
|
||||
// multiple warps participate for the whole vector
|
||||
constexpr uint32_t needed_warps = B_ROW / NUM_THREADS;
|
||||
if (warp_id < needed_warps /* more warps in HW than needed? */) {
|
||||
uint32_t offset = NUM_THREADS * warp_id + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < ROWMAX_SETS; i++) {
|
||||
smem_rowmax[offset + i * ROWMAX_SETS] = FLT_MIN;
|
||||
}
|
||||
smem_rowsum[offset] = 0.0f;
|
||||
smem_O_row_scale[offset] = 0.0f;
|
||||
}
|
||||
|
||||
// each warp clears out a row of smem_O
|
||||
// FIXME: dedup this pattern
|
||||
#pragma GCC unroll 1
|
||||
for (int row_offset = 0; row_offset < B_COL;
|
||||
row_offset += warps_in_threadblock) {
|
||||
const uint32_t row = row_offset + warp_id;
|
||||
uint32_t thread_offset = HEADDIM * row + tid_in_warp;
|
||||
constexpr uint32_t per_row_iter = HEADDIM / NUM_THREADS;
|
||||
const float one = 0.0f;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
smem_O[thread_offset] = 0.0f;
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
}
|
||||
|
||||
asm volatile("threadblock_init_sharedmem_finish_%=:" ::);
|
||||
}
|
||||
|
||||
inline void thread_block_copy_rowmax(const float *src, float *dest,
|
||||
const uint32_t tid_in_threadblock,
|
||||
const uint32_t threads_per_threadblock,
|
||||
const uint32_t threadblock_id_in_cluster) {
|
||||
asm volatile("threadblock_copy_rowmax_start_%=:" ::);
|
||||
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_per_threadblock_per_core =
|
||||
warps_in_threadblock / CORES_PER_CLUSTER;
|
||||
|
||||
// each thread copies one element in rowmax
|
||||
// multiple warps participate for the whole vector
|
||||
constexpr uint32_t num_warps = B_ROW / NUM_THREADS;
|
||||
if (warp_id < num_warps) {
|
||||
uint32_t offset = NUM_THREADS * warp_id + tid_in_warp;
|
||||
dest[offset] = src[offset];
|
||||
}
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
asm volatile("threadblock_copy_rowmax_finish_%=:" ::);
|
||||
}
|
||||
|
||||
inline void thread_block_copy_tile(const float *src, float *dest,
|
||||
const uint32_t tid_in_threadblock,
|
||||
const uint32_t threads_per_threadblock,
|
||||
const uint32_t threadblock_id_in_cluster) {
|
||||
asm volatile("threadblock_copy_tile_start_%=:" ::);
|
||||
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_per_threadblock_per_core =
|
||||
warps_in_threadblock / CORES_PER_CLUSTER;
|
||||
|
||||
// FIXME: dedup this pattern
|
||||
#pragma GCC unroll 1
|
||||
for (int row_offset = 0; row_offset < B_ROW;
|
||||
row_offset += warps_in_threadblock) {
|
||||
const uint32_t row = row_offset + warp_id;
|
||||
const uint32_t first_thread_offset = B_COL * row;
|
||||
|
||||
constexpr uint32_t per_row_iter = B_COL / NUM_THREADS;
|
||||
uint32_t thread_offset = first_thread_offset + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
dest[thread_offset] = src[thread_offset];
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
}
|
||||
|
||||
asm volatile("threadblock_copy_tile_finish_%=:" ::);
|
||||
}
|
||||
|
||||
template <int order>
|
||||
inline float exponential_taylor_term(const float x) {
|
||||
asm volatile("exponential_taylor_term_start_%=:" ::);
|
||||
|
||||
float res = 1.0f;
|
||||
|
||||
if constexpr (order == 1) {
|
||||
res = x;
|
||||
} else if constexpr (order == 2) {
|
||||
res = x * x;
|
||||
res /= 2.0f;
|
||||
} else if constexpr (order == 3) {
|
||||
res = x * x * x;
|
||||
res /= 6.0f;
|
||||
}
|
||||
|
||||
asm volatile("exponential_taylor_term_end_%=:" ::);
|
||||
return res;
|
||||
}
|
||||
|
||||
__attribute__((always_inline)) inline void thread_block_online_softmax(
|
||||
const float *smem_S, float *smem_P, const uint32_t tid_in_threadblock,
|
||||
const uint32_t threads_per_threadblock,
|
||||
const uint32_t threadblock_id_in_cluster, float *smem_scratchpad,
|
||||
float *smem_rowmax, float *smem_rowsum, float *smem_O_row_scale) {
|
||||
asm volatile("thread_block_online_softmax_start_%=:" ::);
|
||||
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_per_threadblock_per_core =
|
||||
warps_in_threadblock / CORES_PER_CLUSTER;
|
||||
|
||||
// float ft[8];
|
||||
// asm volatile("fmv.s %0, f16" : "=f"(ft[0]));
|
||||
// asm volatile("fmv.s %0, f17" : "=f"(ft[1]));
|
||||
// asm volatile("fmv.s %0, f18" : "=f"(ft[2]));
|
||||
// asm volatile("fmv.s %0, f19" : "=f"(ft[3]));
|
||||
// asm volatile("fmv.s %0, f20" : "=f"(ft[4]));
|
||||
// asm volatile("fmv.s %0, f21" : "=f"(ft[5]));
|
||||
// asm volatile("fmv.s %0, f22" : "=f"(ft[6]));
|
||||
// asm volatile("fmv.s %0, f23" : "=f"(ft[7]));
|
||||
|
||||
float *smem_rowmax_this = smem_rowmax + B_ROW;
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (int row_offset = 0; row_offset < B_ROW;
|
||||
row_offset += warps_in_threadblock) {
|
||||
const uint32_t row = row_offset + warp_id;
|
||||
const uint32_t first_thread_offset = B_COL * row;
|
||||
|
||||
// rowmax
|
||||
//
|
||||
// two-level tree reduction: reduce each row into NUM_THREADS intermediate
|
||||
// maxes, then reduce it down to one row max
|
||||
// one warp handles one row in tile
|
||||
|
||||
constexpr uint32_t per_row_iter = B_COL / NUM_THREADS;
|
||||
uint32_t thread_offset = first_thread_offset + tid_in_warp;
|
||||
// FIXME: threadblock_id needs to be in here too
|
||||
float *warp_smem = smem_scratchpad + (warp_id * NUM_THREADS);
|
||||
|
||||
// #define DUMB_ROWMAX
|
||||
#ifdef DUMB_ROWMAX
|
||||
// FIXME remove
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// no tree reduction; a single thread in a warp does serialized max across
|
||||
// the entire row
|
||||
if (tid_in_warp == 0) {
|
||||
float rowmax = smem_S[first_thread_offset];
|
||||
#pragma GCC unroll 16
|
||||
for (int i = 0; i < B_COL; i++) {
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(rowmax)
|
||||
: "f"(rowmax), "f"(smem_S[first_thread_offset + i]));
|
||||
}
|
||||
smem_rowmax_this[row] = rowmax;
|
||||
|
||||
// update previous rowmax
|
||||
// i.e. mi_new = max(mi, mij)
|
||||
float prev_rowmax = smem_rowmax[row];
|
||||
// stage prev rowmax in scratchpad for warp-wide broadcast
|
||||
warp_smem[0] = prev_rowmax;
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(rowmax)
|
||||
: "f"(rowmax), "f"(prev_rowmax));
|
||||
smem_rowmax[row] = rowmax;
|
||||
}
|
||||
|
||||
#else
|
||||
static_assert((B_COL % NUM_THREADS) == 0,
|
||||
"B_COL must be a multiple of NUM_THREADS");
|
||||
float per_thread_max = FLT_MIN;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
const float next = smem_S[thread_offset];
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(per_thread_max)
|
||||
: "f"(per_thread_max), "f"(next));
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
// stage per-thread max value in smem
|
||||
warp_smem[tid_in_warp] = per_thread_max;
|
||||
|
||||
// sync writes to warp_smem
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// #define PARALLEL_ROWMAX
|
||||
#ifndef PARALLEL_ROWMAX
|
||||
// elect 0-th thread to reduce all other thread's values in the warp
|
||||
if (tid_in_warp == 0) {
|
||||
float rowmax = per_thread_max;
|
||||
for (int i = 1; i < NUM_THREADS; i++) {
|
||||
float other = warp_smem[i];
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(rowmax)
|
||||
: "f"(rowmax), "f"(other));
|
||||
}
|
||||
smem_rowmax_this[row] = rowmax;
|
||||
|
||||
// update previous rowmax
|
||||
// i.e. mi_new = max(mi, mij)
|
||||
float prev_rowmax = smem_rowmax[row];
|
||||
// stage prev rowmax in scratchpad for warp-wide broadcast
|
||||
warp_smem[0] = prev_rowmax;
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(rowmax)
|
||||
: "f"(rowmax), "f"(prev_rowmax));
|
||||
smem_rowmax[row] = rowmax;
|
||||
}
|
||||
#else
|
||||
if (warp_id < warps_in_threadblock / NUM_THREADS) {
|
||||
const uint32_t row = row_offset + NUM_THREADS * warp_id + tid_in_warp;
|
||||
float *const thread_smem = smem_scratchpad + (tid_in_warp * NUM_THREADS);
|
||||
float rowmax = FLT_MIN;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < NUM_THREADS; i++) {
|
||||
const float f = thread_smem[i];
|
||||
asm volatile("fmax.s %0, %1, %2" : "=f"(rowmax) : "f"(rowmax), "f"(f));
|
||||
}
|
||||
smem_rowmax_this[row] = rowmax;
|
||||
|
||||
// update previous rowmax
|
||||
// i.e. mi_new = max(mi, mij)
|
||||
float prev_rowmax = smem_rowmax[row];
|
||||
// stage prev rowmax in scratchpad for warp-wide broadcast
|
||||
thread_smem[0] = prev_rowmax;
|
||||
asm volatile("fmax.s %0, %1, %2"
|
||||
: "=f"(rowmax)
|
||||
: "f"(rowmax), "f"(prev_rowmax));
|
||||
smem_rowmax[row] = rowmax;
|
||||
}
|
||||
#endif // PARALLEL_ROWMAX
|
||||
#endif // DUMB_ROWMAX
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// broadcast prev rowmax to all threads in the warp
|
||||
// NOTE: memory consistency is a little sketchy here
|
||||
const float rowmax_prev = warp_smem[0];
|
||||
const float rowmax_this = smem_rowmax_this[row];
|
||||
|
||||
// exponential
|
||||
//
|
||||
// B_ROW / (B_ROW * B_COL / (exp_elem * threads_per_threadblock))
|
||||
// const uint32_t row_stride =
|
||||
// (exp_elem_per_thread * threads_per_threadblock) / B_COL;
|
||||
|
||||
// broadcast updated rowmax to all threads in the warp
|
||||
const float rowmax_new = smem_rowmax[row];
|
||||
|
||||
asm volatile("flashattn_exp_p_start_%=:" ::);
|
||||
|
||||
thread_offset = first_thread_offset + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
float f0 = smem_S[thread_offset];
|
||||
|
||||
f0 -= rowmax_new;
|
||||
|
||||
// 2nd-order Taylor approximation
|
||||
float exp = 1.0f;
|
||||
exp += exponential_taylor_term<1>(f0);
|
||||
exp += exponential_taylor_term<2>(f0);
|
||||
|
||||
// Store S transposed to the shared memory
|
||||
|
||||
smem_P[thread_offset] = exp;
|
||||
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
|
||||
asm volatile("flashattn_exp_p_end_%=:" ::);
|
||||
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// rowsum
|
||||
//
|
||||
// two-level tree reduction, similar to rowmax
|
||||
|
||||
asm volatile("flashattn_rowsum_start_%=:" ::);
|
||||
|
||||
float per_thread_sum = 0.0f;
|
||||
|
||||
thread_offset = first_thread_offset + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
per_thread_sum += smem_P[thread_offset];
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
// stage per-thread sum value in smem
|
||||
// FIXME: threadblock_id needs to be in here too
|
||||
warp_smem = smem_scratchpad + (warp_id * NUM_THREADS);
|
||||
warp_smem[tid_in_warp] = per_thread_sum;
|
||||
|
||||
// sync writes to warp_smem
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// 0-th thread collects all other thread's values in the warp
|
||||
if (tid_in_warp == 0) {
|
||||
float rowsum = per_thread_sum;
|
||||
for (int iter = 1; iter < NUM_THREADS; iter++) {
|
||||
float other = warp_smem[iter];
|
||||
rowsum += other;
|
||||
}
|
||||
|
||||
const float mi_prev = rowmax_prev;
|
||||
const float mi_this = rowmax_this;
|
||||
|
||||
const float x = mi_prev - mi_this;
|
||||
// 2nd-order Taylor approximation
|
||||
float exp = 1.0f;
|
||||
exp += exponential_taylor_term<1>(x);
|
||||
exp += exponential_taylor_term<2>(x);
|
||||
|
||||
// update rowsum
|
||||
const float rowsum_prev = smem_rowsum[row];
|
||||
float rowsum_new = exp * rowsum_prev + rowsum;
|
||||
|
||||
smem_rowsum[row] = rowsum_new;
|
||||
}
|
||||
|
||||
asm volatile("flashattn_rowsum_end_%=:" ::);
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
// compute Oi rescale factor
|
||||
// FIXME: parallelize this across threads
|
||||
//
|
||||
asm volatile("flashattn_rescale_factor_start_%=:" ::);
|
||||
|
||||
thread_offset = first_thread_offset + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
const float mi_prev = rowmax_prev;
|
||||
const float mi_new = rowmax_new;
|
||||
|
||||
const float x = mi_prev - mi_new;
|
||||
// 2nd-order Taylor approximation
|
||||
float exp = 1.0f;
|
||||
exp += exponential_taylor_term<1>(x);
|
||||
exp += exponential_taylor_term<2>(x);
|
||||
|
||||
// @perf: div vs. expansion on e(-x)?
|
||||
smem_O_row_scale[row] = 1.0f / exp;
|
||||
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
|
||||
asm volatile("flashattn_rescale_factor_end_%=:" ::);
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
}
|
||||
|
||||
asm volatile("thread_block_online_softmax_finish_%=:" ::);
|
||||
}
|
||||
|
||||
__attribute__((always_inline)) inline void thread_block_O_rescale(
|
||||
const float *smem_O_in, float *smem_O_out, const float *smem_O_row_scale,
|
||||
const uint32_t tid_in_threadblock, const uint32_t threads_per_threadblock,
|
||||
const uint32_t threadblock_id_in_cluster) {
|
||||
asm volatile("thread_block_O_rescale_start_%=:" ::);
|
||||
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_per_threadblock_per_core =
|
||||
warps_in_threadblock / CORES_PER_CLUSTER;
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (int row_offset = 0; row_offset < B_ROW;
|
||||
row_offset += warps_in_threadblock) {
|
||||
const uint32_t row = row_offset + warp_id;
|
||||
const uint32_t first_thread_offset = B_COL * row;
|
||||
constexpr uint32_t per_row_iter = B_COL / NUM_THREADS;
|
||||
uint32_t thread_offset = first_thread_offset + tid_in_warp;
|
||||
|
||||
// Oi rescale
|
||||
//
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
const float o = smem_O_in[thread_offset];
|
||||
const float scale = smem_O_row_scale[row];
|
||||
smem_O_out[thread_offset] = (o * scale);
|
||||
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
}
|
||||
|
||||
asm volatile("thread_block_O_rescale_finish_%=:" ::);
|
||||
}
|
||||
|
||||
void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
|
||||
// @perf: All threads are running these compute whose result is mostly same
|
||||
// across the threadblock
|
||||
|
||||
#ifdef RADIANCE
|
||||
constexpr uint32_t cores_per_cluster = CORES_PER_CLUSTER;
|
||||
#else
|
||||
constexpr uint32_t cores_per_cluster = 1;
|
||||
#endif
|
||||
|
||||
// FIXME: headdim not considered
|
||||
constexpr uint32_t threads_per_threadblock_theoretical =
|
||||
(B_ROW * B_COL) / (ELEM_PER_THREAD);
|
||||
constexpr uint32_t hw_threads_per_cluster =
|
||||
CORES_PER_CLUSTER * NUM_THREADS * NUM_WARPS;
|
||||
// cap maximum threadblock size to # of HW threads in cluster, to prevent
|
||||
// multiple "wave" invocations which slows down the kernel
|
||||
constexpr uint32_t threads_per_threadblock =
|
||||
(threads_per_threadblock_theoretical > hw_threads_per_cluster)
|
||||
? hw_threads_per_cluster
|
||||
: threads_per_threadblock_theoretical;
|
||||
constexpr uint32_t threadblocks_per_cluster =
|
||||
hw_threads_per_cluster / threads_per_threadblock;
|
||||
constexpr uint32_t warps_per_threadblock_per_core =
|
||||
NUM_WARPS / threadblocks_per_cluster;
|
||||
|
||||
const uint32_t threadblock_id = task_id / threads_per_threadblock;
|
||||
const uint32_t threadblock_id_in_cluster =
|
||||
threadblock_id % threadblocks_per_cluster;
|
||||
const uint32_t tid_in_threadblock = task_id % threads_per_threadblock;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
constexpr uint32_t warps_in_threadblock =
|
||||
threads_per_threadblock / NUM_THREADS;
|
||||
|
||||
// warpgroup context
|
||||
constexpr uint32_t threads_per_warpgroup =
|
||||
threads_per_threadblock / (WARP_SPECIALIZED ? 2 : 1);
|
||||
constexpr uint32_t warpgroups_per_cluster =
|
||||
threadblocks_per_cluster * (WARP_SPECIALIZED ? 2 : 1);
|
||||
const uint32_t warps_per_warpgroup_per_core =
|
||||
NUM_WARPS / warpgroups_per_cluster;
|
||||
const uint32_t warpgroup_id = task_id / threads_per_warpgroup;
|
||||
const uint32_t warpgroup_id_in_cluster =
|
||||
warpgroup_id % warpgroups_per_cluster;
|
||||
const uint32_t tid_in_warpgroup = tid_in_threadblock % threads_per_warpgroup;
|
||||
|
||||
// FIXME do proper software pipelining
|
||||
// if (WARP_SPECIALIZED && warpgroup_id_in_cluster != 1) {
|
||||
// return;
|
||||
// }
|
||||
|
||||
const uint32_t dim_seqlen = arg->dim_seqlen;
|
||||
const uint32_t dim_headdim = arg->dim_headdim;
|
||||
|
||||
// get global memory addresses from kernel arguments
|
||||
const float *gmem_Q = reinterpret_cast<float *>(arg->addr_q);
|
||||
const float *gmem_K = reinterpret_cast<float *>(arg->addr_k);
|
||||
const float *gmem_V = reinterpret_cast<float *>(arg->addr_v);
|
||||
float *gmem_O = reinterpret_cast<float *>(arg->addr_o);
|
||||
|
||||
float *gmem_tmp_d0 = reinterpret_cast<float *>(0xd0000000UL);
|
||||
float *gmem_tmp_d1 = reinterpret_cast<float *>(0xd1000000UL);
|
||||
float *gmem_tmp_d2 = reinterpret_cast<float *>(0xd2000000UL);
|
||||
float *gmem_tmp_d3 = reinterpret_cast<float *>(0xd3000000UL);
|
||||
float *gmem_tmp_d4 = reinterpret_cast<float *>(0xd4000000UL);
|
||||
float *gmem_tmp_d5 = reinterpret_cast<float *>(0xd5000000UL);
|
||||
float *gmem_tmp_d6 = reinterpret_cast<float *>(0xd6000000UL);
|
||||
float *gmem_tmp_d7 = reinterpret_cast<float *>(0xd7000000UL);
|
||||
float *gmem_tmp_e0 = reinterpret_cast<float *>(0xe0000000UL);
|
||||
float *gmem_tmp_e1 = reinterpret_cast<float *>(0xe1000000UL);
|
||||
float *gmem_tmp_e2 = reinterpret_cast<float *>(0xe2000000UL);
|
||||
float *gmem_tmp_e3 = reinterpret_cast<float *>(0xe3000000UL);
|
||||
|
||||
// static shared memory allocation
|
||||
constexpr uint32_t smem_Q_size = B_ROW * HEADDIM;
|
||||
constexpr uint32_t smem_K_size = B_COL * HEADDIM;
|
||||
constexpr uint32_t smem_QK_size = B_ROW * B_COL;
|
||||
constexpr uint32_t smem_V_size = B_COL * HEADDIM;
|
||||
constexpr uint32_t smem_O_size = B_COL * HEADDIM;
|
||||
uint8_t *smem_per_threadblock = reinterpret_cast<uint8_t *>(
|
||||
DEV_SMEM_START_ADDR +
|
||||
sizeof(float_type) *
|
||||
(smem_QK_size + smem_V_size + smem_O_size) *
|
||||
threadblock_id_in_cluster);
|
||||
float *smem_cursor = reinterpret_cast<float *>(DEV_FAKE_SMEM_START_ADDR);
|
||||
float *smem_Q0 = smem_cursor;
|
||||
smem_cursor += smem_Q_size;
|
||||
float *smem_Q1 = smem_cursor;
|
||||
smem_cursor += smem_Q_size;
|
||||
float *smem_K0 = smem_cursor;
|
||||
smem_cursor += smem_K_size;
|
||||
float *smem_K1 = smem_cursor;
|
||||
smem_cursor += smem_K_size;
|
||||
float *smem_V0 = smem_cursor;
|
||||
smem_cursor += smem_V_size;
|
||||
float *smem_V1 = smem_cursor;
|
||||
smem_cursor += smem_V_size;
|
||||
float *smem_S0 = smem_cursor;
|
||||
smem_cursor += smem_QK_size;
|
||||
float *smem_S1 = smem_cursor;
|
||||
smem_cursor += smem_QK_size;
|
||||
float *smem_P0 = smem_S0; // in-place update
|
||||
float *smem_P1 = smem_S1; // in-place update
|
||||
float *smem_O0 = smem_cursor;
|
||||
smem_cursor += smem_O_size;
|
||||
float *smem_O1 = smem_cursor;
|
||||
smem_cursor += smem_O_size;
|
||||
|
||||
// allocate rowmax/rowsum storage at the end of the sharedmem address space
|
||||
constexpr uint32_t smem_rowmax_size = B_ROW * ROWMAX_SETS;
|
||||
constexpr uint32_t smem_rowsum_size = B_ROW;
|
||||
constexpr uint32_t smem_O_row_scale_size = B_ROW;
|
||||
smem_cursor = reinterpret_cast<float *>(SMEM_ADDR_END);
|
||||
|
||||
smem_cursor -= smem_rowmax_size;
|
||||
float *smem_rowmax_0 = smem_cursor;
|
||||
smem_cursor -= smem_rowmax_size;
|
||||
float *smem_rowmax_1 = smem_cursor;
|
||||
smem_cursor -= smem_rowsum_size;
|
||||
float *smem_rowsum_0 = smem_cursor;
|
||||
smem_cursor -= smem_rowsum_size;
|
||||
float *smem_rowsum_1 = smem_cursor;
|
||||
smem_cursor -= smem_O_row_scale_size;
|
||||
float *smem_O_row_scale_0 = smem_cursor;
|
||||
smem_cursor -= smem_O_row_scale_size;
|
||||
float *smem_O_row_scale_1 = smem_cursor;
|
||||
|
||||
// sharedmem "scratchpad" area to put temporary data, e.g. for tree reduction
|
||||
// in rowsum
|
||||
// NOTE: out-of bounds is not checked
|
||||
// TODO: reduce this from B_ROW to NUM_WARPS
|
||||
constexpr uint32_t smem_scratchpad_size =
|
||||
threads_per_warpgroup * 2 /*arbitrary slack*/;
|
||||
smem_cursor -= smem_scratchpad_size;
|
||||
float *smem_scratchpad_0 = smem_cursor;
|
||||
smem_cursor -= smem_scratchpad_size;
|
||||
float *smem_scratchpad_1 = smem_cursor;
|
||||
|
||||
// select the correct buffer by warpgroup
|
||||
float *smem_Q = (warpgroup_id % 2) ? smem_Q1 : smem_Q0;
|
||||
float *smem_K = (warpgroup_id % 2) ? smem_K1 : smem_K0;
|
||||
float *smem_V = (warpgroup_id % 2) ? smem_V1 : smem_V0;
|
||||
float *smem_S = (warpgroup_id % 2) ? smem_S1 : smem_S0;
|
||||
float *smem_O = (warpgroup_id % 2) ? smem_O1 : smem_O0;
|
||||
float *smem_P = smem_S;
|
||||
float *smem_O_row_scale =
|
||||
(warpgroup_id % 2) ? smem_O_row_scale_1 : smem_O_row_scale_0;
|
||||
float *smem_rowmax = (warpgroup_id % 2) ? smem_rowmax_1 : smem_rowmax_0;
|
||||
float *smem_rowsum = (warpgroup_id % 2) ? smem_rowsum_1 : smem_rowsum_0;
|
||||
float *smem_scratchpad =
|
||||
(warpgroup_id % 2) ? smem_scratchpad_1 : smem_scratchpad_0;
|
||||
|
||||
// initialize rowmax/rowsum values in sharedmem
|
||||
thread_block_init_sharedmem(tid_in_warpgroup, threads_per_warpgroup, smem_O,
|
||||
smem_rowmax, smem_rowsum, smem_O_row_scale);
|
||||
|
||||
constexpr uint32_t global_barrier_id = NUM_WARPS - 1; // arbitrary
|
||||
|
||||
// delay warpgroup 0 by 1 iteration to do ping-pong scheduling
|
||||
if (warpgroup_id == 1) {
|
||||
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
|
||||
}
|
||||
|
||||
// read Q and K into SMEM before the loop starts
|
||||
//
|
||||
static_assert(B_ROW == B_COL, "currently only supports square tiles");
|
||||
|
||||
// load Q; this stays in SMEM for the entire loop
|
||||
load_tile_to_smem<float, MemLayout::MN_major, MemLayout::MN_major, B_ROW,
|
||||
HEADDIM, threads_per_warpgroup>(
|
||||
dim_seqlen, warpgroup_id, 0 /* dim_k == headdim */, gmem_Q, smem_Q,
|
||||
tid_in_warpgroup);
|
||||
|
||||
// load K
|
||||
load_tile_to_smem<float, MemLayout::MN_major, MemLayout::MN_major, B_COL,
|
||||
HEADDIM, threads_per_warpgroup>(
|
||||
dim_seqlen, /*tile_k=*/0, 0 /* dim_k == headdim */, gmem_K, smem_K,
|
||||
tid_in_warpgroup);
|
||||
|
||||
// protect write to SMEM
|
||||
threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
|
||||
|
||||
asm volatile ("tile_loop_start_%=:" :: );
|
||||
|
||||
// "inner loop" along the columns of K^T
|
||||
const uint32_t k_tiles = (dim_seqlen / B_COL);
|
||||
for (uint32_t tile_k = 0; tile_k < k_tiles; tile_k++) {
|
||||
// float *smem_P_produce = (tile_k % 2) ? smem_P0 : smem_P1;
|
||||
// float *smem_P_consume = (tile_k % 2) ? smem_P1 : smem_P0;
|
||||
// float *smem_V_produce = (tile_k % 2) ? smem_V0 : smem_V1;
|
||||
// float *smem_V_consume = (tile_k % 2) ? smem_V1 : smem_V0;
|
||||
// float *smem_O_row_scale_produce =
|
||||
// (tile_k % 2) ? smem_O_row_scale_0 : smem_O_row_scale_1;
|
||||
// float *smem_O_row_scale_consume =
|
||||
// (tile_k % 2) ? smem_O_row_scale_1 : smem_O_row_scale_0;
|
||||
|
||||
constexpr bool skip_gemm_qk = false;
|
||||
if constexpr (!skip_gemm_qk) {
|
||||
// GEMM I: S = Q*K
|
||||
//
|
||||
// FIXME: deduplicate this between GEMM II
|
||||
if constexpr (!WARP_SPECIALIZED) {
|
||||
// clear out accumulators before GEMM
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
thread_block_gemm_single_tile<
|
||||
float, MemLayout::MN_major, MemLayout::MN_major, B_ROW, B_COL,
|
||||
HEADDIM, /*leading_dim_a=*/0, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/false,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_Q, smem_K, nullptr /*ignore accum*/, smem_S, tid_in_warpgroup,
|
||||
threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
} else {
|
||||
// when warp-specialized, there's only enough warps to do 64x32 tile
|
||||
// size so we need to do 2 GEMM calls
|
||||
static_assert(B_ROW / 2 == 32,
|
||||
"tile size assumption for warp-specialization not met");
|
||||
|
||||
// assumes smem_Q is K-major
|
||||
// FIXME: fix this to MN-major
|
||||
float *smem_Q_half0 = smem_Q;
|
||||
float *smem_Q_half1 = smem_Q + (B_ROW / 2); // MN-major
|
||||
// float *smem_Q_half1 = smem_Q + (B_ROW / 2) * HEADDIM; // K-major
|
||||
float *smem_S_half0 = smem_S;
|
||||
float *smem_S_half1 = smem_S + (B_ROW / 2) * B_COL;
|
||||
|
||||
// clear out accumulators before GEMM
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
// split by rows into 2 chunks
|
||||
thread_block_gemm_single_tile<
|
||||
float, MemLayout::MN_major, MemLayout::MN_major, B_ROW / 2,
|
||||
B_COL, HEADDIM, /*leading_dim_a=*/B_ROW, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/false,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_Q_half0, smem_K, nullptr /*ignore accum*/, smem_S_half0,
|
||||
tid_in_warpgroup, threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
thread_block_gemm_single_tile<
|
||||
float, MemLayout::MN_major, MemLayout::MN_major, B_ROW / 2,
|
||||
B_COL, HEADDIM, /*leading_dim_a=*/B_ROW, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/false,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_Q_half1, smem_K, nullptr /*ignore accum*/, smem_S_half1,
|
||||
tid_in_warpgroup, threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
} else {
|
||||
// load Q*K
|
||||
load_tile_to_smem<float, MemLayout::K_major, MemLayout::K_major, B_COL,
|
||||
HEADDIM, threads_per_warpgroup>(
|
||||
dim_seqlen, warpgroup_id /* parallelize across rows */, tile_k,
|
||||
gmem_Q /*contains S*/, smem_S, tid_in_warpgroup);
|
||||
}
|
||||
|
||||
// protect write to SMEM (smem_S) before softmax
|
||||
threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
|
||||
|
||||
if constexpr (DEBUG) {
|
||||
if (warpgroup_id == 0) {
|
||||
if (tile_k == 0) {
|
||||
thread_block_copy_tile(smem_S, gmem_tmp_d0,
|
||||
tid_in_warpgroup, threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
} else if (tile_k == 1) {
|
||||
thread_block_copy_tile(smem_S, gmem_tmp_d1,
|
||||
tid_in_warpgroup, threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
|
||||
threadblock_barrier(warpgroup_id_in_cluster,
|
||||
warps_per_warpgroup_per_core);
|
||||
}
|
||||
}
|
||||
|
||||
// inter-warpgroup barrier before online softmax
|
||||
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
|
||||
|
||||
// Online softmax
|
||||
//
|
||||
thread_block_online_softmax(smem_S, smem_P, tid_in_warpgroup,
|
||||
threads_per_warpgroup, warpgroup_id_in_cluster,
|
||||
smem_scratchpad, smem_rowmax, smem_rowsum,
|
||||
smem_O_row_scale);
|
||||
|
||||
// data movement for K and V
|
||||
//
|
||||
// Q stays in SMEM for the entire loop
|
||||
//
|
||||
// load K for the next iteration
|
||||
load_tile_to_smem<float, MemLayout::MN_major, MemLayout::MN_major, B_COL,
|
||||
HEADDIM, threads_per_warpgroup>(
|
||||
dim_seqlen, tile_k + 1, 0 /* dim_k == headdim */, gmem_K, smem_K,
|
||||
tid_in_warpgroup);
|
||||
|
||||
// load V for the current iteration
|
||||
// V dimension is [seqlen, headdim], stored N(headdim)-major
|
||||
load_tile_to_smem<float, MemLayout::MN_major, MemLayout::MN_major, B_COL,
|
||||
HEADDIM, threads_per_warpgroup>(
|
||||
HEADDIM, 0 /* full N-dimension */, tile_k, gmem_V, smem_V,
|
||||
tid_in_warpgroup);
|
||||
|
||||
// protect write to SMEM
|
||||
threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
|
||||
|
||||
if constexpr (DEBUG) {
|
||||
if (warpgroup_id == 0) {
|
||||
if (tile_k == 0) {
|
||||
thread_block_copy_rowmax(smem_rowmax, gmem_tmp_e0, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e2, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
} else if (tile_k == 1) {
|
||||
thread_block_copy_rowmax(smem_rowmax, gmem_tmp_e1, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
thread_block_copy_rowmax(smem_rowsum, gmem_tmp_e3, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
|
||||
threadblock_barrier(warpgroup_id_in_cluster,
|
||||
warps_per_warpgroup_per_core);
|
||||
}
|
||||
}
|
||||
|
||||
// inter-warpgroup barrier before GEMM II
|
||||
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
|
||||
|
||||
// GEMM II: O = O + P*V
|
||||
|
||||
// Oi rescale
|
||||
thread_block_O_rescale(smem_O, smem_O /*in-place*/,
|
||||
smem_O_row_scale, tid_in_warpgroup,
|
||||
threads_per_warpgroup, warpgroup_id_in_cluster);
|
||||
|
||||
// rescale-to-PV-GEMM barrier
|
||||
threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
|
||||
|
||||
if constexpr (DEBUG) {
|
||||
if (warpgroup_id == 0) {
|
||||
// O before PV
|
||||
if (tile_k == 0) {
|
||||
thread_block_copy_tile(smem_P, gmem_tmp_d2, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
thread_block_copy_tile(smem_O, gmem_tmp_d4, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
} else if (tile_k == 1) {
|
||||
thread_block_copy_tile(smem_P, gmem_tmp_d3, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
thread_block_copy_tile(smem_O, gmem_tmp_d5, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
|
||||
threadblock_barrier(warpgroup_id_in_cluster,
|
||||
warps_per_warpgroup_per_core);
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (!WARP_SPECIALIZED) {
|
||||
// clear out accumulators before GEMM
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
thread_block_gemm_single_tile<float, MemLayout::K_major,
|
||||
MemLayout::MN_major, B_ROW, HEADDIM, B_COL,
|
||||
/*leading_dim_a=*/0, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/true,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_P, smem_V, smem_O /*load accum*/, smem_O, tid_in_warpgroup,
|
||||
threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
|
||||
// FIXME: wrong but fast
|
||||
// thread_block_gemm_single_tile<float, MemLayout::MN_major,
|
||||
// MemLayout::MN_major,
|
||||
// B_ROW, HEADDIM, B_COL,
|
||||
// /*leading_dim_a=*/0, /*leading_dim_b=*/0,
|
||||
// /*load_accum=*/true,
|
||||
// /*write_to_smem=*/true>(
|
||||
// smem_P, smem_V, smem_O /*load accum*/, smem_O,
|
||||
// tid_in_warpgroup, threads_per_warpgroup,
|
||||
// warpgroups_per_cluster, warpgroup_id_in_cluster);
|
||||
} else {
|
||||
// when warp-specialized, there's only enough warps to do 64x32 tile
|
||||
// size so we need to do 2 GEMM calls
|
||||
static_assert(B_ROW / 2 == 32,
|
||||
"tile size assumption for warp-specialization not met");
|
||||
|
||||
// assumes smem_P is K-major
|
||||
float *smem_P_half0 = smem_P;
|
||||
float *smem_P_half1 = smem_P + (B_ROW / 2) * B_COL;
|
||||
float *smem_O_half0 = smem_O;
|
||||
float *smem_O_half1 = smem_O + (B_ROW / 2) * HEADDIM;
|
||||
|
||||
// clear out accumulators before GEMM
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
// split by rows into 2 chunks
|
||||
thread_block_gemm_single_tile<
|
||||
float, MemLayout::K_major, MemLayout::MN_major, B_ROW / 2, HEADDIM,
|
||||
B_COL, /*leading_dim_a=*/0, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/true,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_P_half0, smem_V, smem_O_half0 /*load accum*/, smem_O_half0,
|
||||
tid_in_warpgroup, threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
|
||||
initialize_accum_regs<0>();
|
||||
initialize_accum_regs<1>();
|
||||
|
||||
thread_block_gemm_single_tile<
|
||||
float, MemLayout::K_major, MemLayout::MN_major, B_ROW / 2, HEADDIM,
|
||||
B_COL, /*leading_dim_a=*/0, /*leading_dim_b=*/0,
|
||||
/*load_accum=*/true,
|
||||
/*write_to_smem=*/true>(
|
||||
smem_P_half1, smem_V, smem_O_half1 /*load accum*/, smem_O_half1,
|
||||
tid_in_warpgroup, threads_per_warpgroup, warpgroups_per_cluster,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
|
||||
threadblock_barrier(warpgroup_id_in_cluster, warps_per_warpgroup_per_core);
|
||||
|
||||
if constexpr (DEBUG) {
|
||||
if (warpgroup_id == 0) {
|
||||
// O after PV
|
||||
if (tile_k == 0) {
|
||||
thread_block_copy_tile(smem_O, gmem_tmp_d6, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
} else if (tile_k == 1) {
|
||||
thread_block_copy_tile(smem_O, gmem_tmp_d7, tid_in_warpgroup,
|
||||
threads_per_warpgroup,
|
||||
warpgroup_id_in_cluster);
|
||||
}
|
||||
|
||||
threadblock_barrier(warpgroup_id_in_cluster,
|
||||
warps_per_warpgroup_per_core);
|
||||
}
|
||||
}
|
||||
|
||||
tile_iter_end:
|
||||
// synchronize progress of two warpgroups
|
||||
// threadblock_barrier(threadblock_id_in_cluster,
|
||||
// warps_per_threadblock_per_core);
|
||||
// threadblock_barrier(3, // FIXME
|
||||
// NUM_WARPS);
|
||||
}
|
||||
|
||||
asm volatile ("tile_loop_finish_%=:" :: );
|
||||
|
||||
// wait for warpgroup 1 to finish, which called the global barrier before
|
||||
// entering the loop
|
||||
if (warpgroup_id == 0) {
|
||||
threadblock_barrier(global_barrier_id, warps_per_threadblock_per_core);
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
kernel_arg_t *arg = (kernel_arg_t *)KERNEL_ARG_DEV_MEM_ADDR;
|
||||
|
||||
// FIXME:: use actuall seqlen/headdim
|
||||
const uint32_t problem_size = (B_ROW * B_COL) / (ELEM_PER_THREAD);
|
||||
const uint32_t hw_threads_per_cluster =
|
||||
CORES_PER_CLUSTER * vx_num_threads() * vx_num_warps();
|
||||
// prevent launching more threads than the necessary problem size
|
||||
// TODO: this does not take into account multiple clusters
|
||||
const uint32_t grid_size = (problem_size > hw_threads_per_cluster)
|
||||
? hw_threads_per_cluster
|
||||
: problem_size;
|
||||
|
||||
#ifdef RADIANCE
|
||||
vx_spawn_tasks_cluster(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
|
||||
#else
|
||||
// NOTE: This kernel assumes contiguous thread scheduling for efficient shared
|
||||
// memory allocation, and therefore does not work with original vx_spawn_tasks
|
||||
vx_spawn_tasks_contiguous(grid_size, (vx_spawn_tasks_cb)kernel_body, arg);
|
||||
#endif
|
||||
return 0;
|
||||
}
|
||||
166
tests/regression/flash_attention/main.cpp
Normal file
166
tests/regression/flash_attention/main.cpp
Normal file
@@ -0,0 +1,166 @@
|
||||
#include <iostream>
|
||||
#include <fstream>
|
||||
#include <unistd.h>
|
||||
#include <string.h>
|
||||
#include <vortex.h>
|
||||
#include <vector>
|
||||
#include <cassert>
|
||||
#include "common.h"
|
||||
#include "half.hpp"
|
||||
|
||||
using half_float::half;
|
||||
using half_float::half_cast;
|
||||
|
||||
#define RT_CHECK(_expr) \
|
||||
do { \
|
||||
int _ret = _expr; \
|
||||
if (0 == _ret) \
|
||||
break; \
|
||||
printf("Error: '%s' returned %d!\n", #_expr, (int)_ret); \
|
||||
cleanup(); \
|
||||
exit(-1); \
|
||||
} while (false)
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
const char* kernel_file = "kernel.bin";
|
||||
uint32_t count = 0;
|
||||
|
||||
std::vector<float> ref_data;
|
||||
|
||||
vx_device_h device = nullptr;
|
||||
std::vector<uint8_t> staging_buf;
|
||||
kernel_arg_t kernel_arg = {};
|
||||
|
||||
static void show_usage() {
|
||||
std::cout << "Vortex Test." << std::endl;
|
||||
std::cout << "Usage: [-k: kernel] [-n words] [-h: help]" << std::endl;
|
||||
}
|
||||
|
||||
static void parse_args(int argc, char **argv) {
|
||||
int c;
|
||||
while ((c = getopt(argc, argv, "n:k:h?")) != -1) {
|
||||
switch (c) {
|
||||
case 'n':
|
||||
count = atoi(optarg);
|
||||
break;
|
||||
case 'k':
|
||||
kernel_file = optarg;
|
||||
break;
|
||||
case 'h':
|
||||
case '?': {
|
||||
show_usage();
|
||||
exit(0);
|
||||
} break;
|
||||
default:
|
||||
show_usage();
|
||||
exit(-1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void cleanup() {
|
||||
if (device) {
|
||||
// vx_mem_free(device, kernel_arg.addr_a);
|
||||
// vx_mem_free(device, kernel_arg.addr_b);
|
||||
// vx_mem_free(device, kernel_arg.addr_c);
|
||||
vx_dev_close(device);
|
||||
}
|
||||
}
|
||||
|
||||
int run_test(const kernel_arg_t& kernel_arg,
|
||||
uint32_t buf_size) {
|
||||
// start device
|
||||
std::cout << "start device" << std::endl;
|
||||
RT_CHECK(vx_start(device));
|
||||
|
||||
// wait for completion
|
||||
std::cout << "wait for completion" << std::endl;
|
||||
RT_CHECK(vx_ready_wait(device, VX_MAX_TIMEOUT));
|
||||
|
||||
// download destination buffer
|
||||
std::cout << "download destination buffer" << std::endl;
|
||||
RT_CHECK(vx_copy_from_dev(device, staging_buf.data(), kernel_arg.addr_o, buf_size));
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
int main(int argc, char *argv[]) {
|
||||
// parse command arguments
|
||||
parse_args(argc, argv);
|
||||
|
||||
if (count == 0) {
|
||||
count = 1;
|
||||
}
|
||||
|
||||
std::srand(50);
|
||||
|
||||
// open device connection
|
||||
std::cout << "open device connection" << std::endl;
|
||||
RT_CHECK(vx_dev_open(&device));
|
||||
|
||||
uint32_t dim_seqlen = 128;
|
||||
uint32_t dim_headdim = 64;
|
||||
|
||||
using float_type = half;
|
||||
|
||||
uint32_t dst_buf_size =
|
||||
dim_seqlen * dim_headdim * sizeof(ref_data[0]);
|
||||
|
||||
// upload program
|
||||
std::cout << "upload program" << std::endl;
|
||||
RT_CHECK(vx_upload_kernel_file(device, kernel_file));
|
||||
|
||||
// allocate device memory
|
||||
std::cout << "allocate device memory" << std::endl;
|
||||
kernel_arg.addr_q = 0xa0000000;
|
||||
kernel_arg.addr_k = 0xa1000000;
|
||||
kernel_arg.addr_v = 0xa2000000;
|
||||
kernel_arg.addr_o = 0xc0000000;
|
||||
|
||||
kernel_arg.dim_seqlen = dim_seqlen;
|
||||
kernel_arg.dim_headdim = dim_headdim;
|
||||
|
||||
std::cout << "dev_addr_q=0x" << std::hex << kernel_arg.addr_q << std::endl;
|
||||
std::cout << "dev_addr_k=0x" << std::hex << kernel_arg.addr_k << std::endl;
|
||||
std::cout << "dev_addr_v=0x" << std::hex << kernel_arg.addr_v << std::endl;
|
||||
std::cout << "dev_addr_o=0x" << std::hex << kernel_arg.addr_o << std::endl;
|
||||
|
||||
// allocate staging buffer
|
||||
{
|
||||
std::cout << "allocate staging buffer" << std::endl;
|
||||
uint32_t staging_buf_size = sizeof(kernel_arg_t);
|
||||
staging_buf.resize(staging_buf_size);
|
||||
}
|
||||
|
||||
// upload kernel argument
|
||||
{
|
||||
std::cout << "upload kernel argument" << std::endl;
|
||||
auto buf_ptr = staging_buf.data();
|
||||
memcpy(buf_ptr, &kernel_arg, sizeof(kernel_arg_t));
|
||||
RT_CHECK(vx_copy_to_dev(device, KERNEL_ARG_DEV_MEM_ADDR, staging_buf.data(), sizeof(kernel_arg_t)));
|
||||
|
||||
std::cout << "uploading argument buffer to device, device mem address="
|
||||
<< std::hex << KERNEL_ARG_DEV_MEM_ADDR << ", size=" << std::dec
|
||||
<< sizeof(kernel_arg_t) << " bytes\n";
|
||||
std::ofstream file("args.bin", std::ios::binary | std::ios::out);
|
||||
if (!file) {
|
||||
std::cerr << "error: failed to open args.bin for writing\n";
|
||||
exit(EXIT_FAILURE);
|
||||
}
|
||||
file.write(reinterpret_cast<char *>(staging_buf.data()),
|
||||
sizeof(kernel_arg_t));
|
||||
file.close();
|
||||
}
|
||||
|
||||
// run tests
|
||||
std::cout << "run tests" << std::endl;
|
||||
RT_CHECK(run_test(kernel_arg, dst_buf_size));
|
||||
std::cout << "PASSED!" << std::endl;
|
||||
|
||||
// cleanup
|
||||
std::cout << "cleanup" << std::endl;
|
||||
cleanup();
|
||||
|
||||
return 0;
|
||||
}
|
||||
@@ -3,6 +3,7 @@ PROJECT = sgemm_tcore
|
||||
SRCS = main.cpp common.h
|
||||
|
||||
VX_SRCS = kernel.cpp
|
||||
VX_INCLUDES = sgemm_impl.hpp
|
||||
|
||||
OPTS ?= -n16
|
||||
|
||||
|
||||
@@ -3,538 +3,44 @@
|
||||
#include <vx_print.h>
|
||||
#include <vx_spawn.h>
|
||||
#include "common.h"
|
||||
#include "util.hpp"
|
||||
#include "sgemm_impl.hpp"
|
||||
#include "include/gemmini.h"
|
||||
#include "gemmini_mmio.h"
|
||||
|
||||
#define GEMMINI_DMA 0
|
||||
#if SMEM_SIZE == 0x4000
|
||||
#define SMEM_ADDR_Q0 ((float * const) 0xff000000)
|
||||
#define SMEM_ADDR_Q1 ((float * const) 0xff001000)
|
||||
#define SMEM_ADDR_Q2 ((float * const) 0xff002000)
|
||||
#define SMEM_ADDR_Q3 ((float * const) 0xff003000)
|
||||
#define SPAD_ADDR_Q0 0x0
|
||||
#define SPAD_ADDR_Q1 0x80
|
||||
#define SPAD_ADDR_Q2 0x100
|
||||
#define SPAD_ADDR_Q3 0x180
|
||||
#define BOUND_INST 0x400040004ULL
|
||||
#elif SMEM_SIZE == 0x10000
|
||||
#define SMEM_ADDR_Q0 ((float * const) 0xff000000)
|
||||
#define SMEM_ADDR_Q1 ((float * const) 0xff004000)
|
||||
#define SMEM_ADDR_Q2 ((float * const) 0xff008000)
|
||||
#define SMEM_ADDR_Q3 ((float * const) 0xff00c000)
|
||||
#define SPAD_ADDR_Q0 0x0
|
||||
#define SPAD_ADDR_Q1 0x200
|
||||
#define SPAD_ADDR_Q2 0x400
|
||||
#define SPAD_ADDR_Q3 0x600
|
||||
#define BOUND_INST 0x800080008ULL
|
||||
#else
|
||||
#error Unsupported smem size
|
||||
#endif
|
||||
constexpr bool DEBUG = true;
|
||||
|
||||
// FIXME: NUM_THREADS and NUM_WARPS hardcoded
|
||||
#if ((BM * BN / ELEM_PER_THREAD) > (CORES_PER_CLUSTER * 8 * 8))
|
||||
#error "threadblock size too big for cluster"
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
inline void global_dmem_load(const uint32_t dim_n, const uint32_t dim_k,
|
||||
const uint32_t k, const T *A, const T *B,
|
||||
volatile T *local_a, volatile T *local_b,
|
||||
const uint32_t tid_in_threadblock,
|
||||
const uint32_t threadblock_id_x,
|
||||
const uint32_t threadblock_id_y) {
|
||||
// In fp16 mode, bit-pack two fp16 elements into each fp32 element, and do
|
||||
// data movement at the fp32 granularity. Assuming that the matrix is stored
|
||||
// row-major in GMEM, the packed fp16 pairs belong to the same row,
|
||||
// neighboring columns; therefore, it essentially becomes equivalent to
|
||||
// moving a fp32 matrix whose column dimensions (dim_k/BK/k) are compressed
|
||||
// by a factor of two.
|
||||
constexpr uint32_t packed_factor = (std::is_same_v<T, float16_t> ? 2 : 1);
|
||||
constexpr uint32_t BK_adjusted = BK / packed_factor;
|
||||
const uint32_t dim_k_adjusted = dim_k / packed_factor;
|
||||
constexpr uint32_t BN_adjusted = BN / packed_factor;
|
||||
const uint32_t dim_n_adjusted = dim_n / packed_factor;
|
||||
const uint32_t k_adjusted = k / packed_factor;
|
||||
|
||||
const uint32_t local_a_row = tid_in_threadblock / BK_adjusted;
|
||||
const uint32_t local_a_col = tid_in_threadblock % BK_adjusted;
|
||||
const uint32_t local_as_row = tid_in_threadblock / BM;
|
||||
const uint32_t local_as_col = tid_in_threadblock % BM;
|
||||
const uint32_t local_b_row = tid_in_threadblock / BN_adjusted;
|
||||
const uint32_t local_b_col = tid_in_threadblock % BN_adjusted;
|
||||
|
||||
// FIXME: need fix for fp16?
|
||||
constexpr uint32_t threads_in_threadblock = (BM * BN) / ELEM_PER_THREAD;
|
||||
|
||||
// Data move from GMEM to SMEM
|
||||
//
|
||||
// Make sure global offset values for A and B are contiguous between
|
||||
// neighboring threads to ensure GMEM coalescing.
|
||||
//
|
||||
// TODO: Sharedmem swizzling is important here
|
||||
|
||||
// move A
|
||||
if constexpr (!TRANSPOSE_AT_PRODUCE) {
|
||||
// No transpose at GMEM->SMEM movement
|
||||
// FIXME: !TRANSPOSE_AS code is old
|
||||
|
||||
const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
|
||||
// number of rows a full TB can read at a time
|
||||
// this is equivalent to threadblock_dim_y (assuming threadblock_dim_x ==
|
||||
// BK)
|
||||
constexpr uint32_t row_stride_a = threads_in_threadblock / BK_adjusted;
|
||||
const float *global_a = reinterpret_cast<float *>(A) +
|
||||
dim_k_adjusted * global_a_row +
|
||||
(k_adjusted + local_a_col);
|
||||
volatile float *local_a_tmp = reinterpret_cast<float *>(local_a) +
|
||||
BK_adjusted * local_a_row + local_a_col;
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t local_row_offset = 0; local_row_offset < BM;
|
||||
local_row_offset += row_stride_a) {
|
||||
*local_a_tmp = *global_a;
|
||||
|
||||
// move to the next "row-chunk", when threadblock is smaller than BM*BK
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
local_a_tmp += BK_adjusted * row_stride_a;
|
||||
}
|
||||
} else {
|
||||
if constexpr (!GMEM_COALESCED_A) {
|
||||
// !GMEM_COALESCED_A: threads do uncoalesced read from neighboring row in
|
||||
// GMEM, writes to neighboring cols in SMEM
|
||||
constexpr uint32_t row_stride_as = threads_in_threadblock / BM;
|
||||
const uint32_t global_a_row = BM * threadblock_id_y + local_as_col;
|
||||
const float *global_a =
|
||||
reinterpret_cast<float *>(A) + dim_k_adjusted * global_a_row + (k_adjusted + local_as_row);
|
||||
volatile float *local_a_tmp =
|
||||
reinterpret_cast<float *>(local_a) + BM * local_as_row + local_as_col;
|
||||
|
||||
static_assert(
|
||||
row_stride_as * 8 <= BK_adjusted,
|
||||
"manual loop unrolling condition not met; consider increasing BK");
|
||||
static_assert(
|
||||
(BK_adjusted % (row_stride_as * 8)) == 0,
|
||||
"manual loop unrolling condition not met; BK should be power-of-two");
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t local_row_offset = 0; local_row_offset < BK_adjusted;
|
||||
local_row_offset += row_stride_as * 8) {
|
||||
// @perf: bank conflicts here
|
||||
// const uint32_t global_a_offset =
|
||||
// dim_k_adjusted * (global_a_row) + (k + local_as_row + local_row_offset);
|
||||
// FIXME experimenting with global coalescing
|
||||
// const uint32_t global_a_offset =
|
||||
// dim_k_adjusted * (global_a_row + local_row_offset) + (k + local_as_col);
|
||||
// local_a[BM * (local_as_row + local_row_offset) + local_as_col] =
|
||||
// A[global_a_offset];
|
||||
|
||||
// *local_a_tmp = *global_a;
|
||||
asm volatile ("flw ft0, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft1, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft2, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft3, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft4, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft5, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft6, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
asm volatile ("flw ft7, (%0)" :: "r"(global_a));
|
||||
global_a += row_stride_as;
|
||||
|
||||
// NOTE: stride is fixed to word size , i.e. sizeof(float) = 4,
|
||||
// regardless of fp16 or fp32. Since Vortex core does not support fp16,
|
||||
// load things at word granularity and reinterpret bits inside the
|
||||
// tensor core.
|
||||
asm volatile ("fsw ft0, %0(%1)" :: "i"(BM * row_stride_as * 0 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft1, %0(%1)" :: "i"(BM * row_stride_as * 1 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft2, %0(%1)" :: "i"(BM * row_stride_as * 2 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft3, %0(%1)" :: "i"(BM * row_stride_as * 3 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft4, %0(%1)" :: "i"(BM * row_stride_as * 4 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft5, %0(%1)" :: "i"(BM * row_stride_as * 5 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft6, %0(%1)" :: "i"(BM * row_stride_as * 6 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft7, %0(%1)" :: "i"(BM * row_stride_as * 7 * sizeof(float)), "r"(local_a_tmp));
|
||||
local_a_tmp += BM * row_stride_as * 8;
|
||||
}
|
||||
} else {
|
||||
constexpr uint32_t row_stride_a = threads_in_threadblock / BK_adjusted;
|
||||
const uint32_t global_a_row = BM * threadblock_id_y + local_a_row;
|
||||
const float *global_a = reinterpret_cast<const float *>(A) +
|
||||
dim_k_adjusted * global_a_row +
|
||||
(k_adjusted + local_a_col);
|
||||
// NOTE that SMEM writes are transposed
|
||||
volatile float *local_a_tmp =
|
||||
reinterpret_cast<volatile float *>(local_a) + BM * local_a_col +
|
||||
local_a_row;
|
||||
|
||||
static_assert(
|
||||
row_stride_a * 8 <= BM,
|
||||
"manual loop unrolling condition not met; consider increasing BM");
|
||||
static_assert(
|
||||
(BM % (row_stride_a * 8)) == 0,
|
||||
"manual loop unrolling condition not met; BM should be power-of-two");
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t local_row_offset = 0; local_row_offset < BM;
|
||||
local_row_offset += row_stride_a * 8) {
|
||||
// const uint32_t global_a_offset =
|
||||
// dim_k_adjusted * (global_a_row + local_row_offset) + (k + local_a_col);
|
||||
// NOTE that SMEM writes are transposed
|
||||
// local_a[BM * (local_a_col) + local_a_row + local_row_offset] =
|
||||
// A[global_a_offset];
|
||||
|
||||
asm volatile ("flw ft0, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft1, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft2, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft3, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft4, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft5, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft6, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
asm volatile ("flw ft7, (%0)" :: "r"(global_a));
|
||||
global_a += dim_k_adjusted * row_stride_a;
|
||||
|
||||
// stride along columns
|
||||
asm volatile ("fsw ft0, %0(%1)" :: "i"(row_stride_a * 0 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft1, %0(%1)" :: "i"(row_stride_a * 1 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft2, %0(%1)" :: "i"(row_stride_a * 2 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft3, %0(%1)" :: "i"(row_stride_a * 3 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft4, %0(%1)" :: "i"(row_stride_a * 4 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft5, %0(%1)" :: "i"(row_stride_a * 5 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft6, %0(%1)" :: "i"(row_stride_a * 6 * sizeof(float)), "r"(local_a_tmp));
|
||||
asm volatile ("fsw ft7, %0(%1)" :: "i"(row_stride_a * 7 * sizeof(float)), "r"(local_a_tmp));
|
||||
local_a_tmp += row_stride_a * 8;
|
||||
}
|
||||
}
|
||||
} // end move A
|
||||
|
||||
// move B
|
||||
constexpr uint32_t row_stride_b = threads_in_threadblock / BN_adjusted;
|
||||
const uint32_t global_b_col = BN_adjusted * threadblock_id_x + local_b_col;
|
||||
// NOTE: not k_adjusted here; k is along the row dimension which is not
|
||||
// compressed for fp16
|
||||
const float *global_b = reinterpret_cast<const float *>(B) +
|
||||
dim_n_adjusted * (k + local_b_row) + global_b_col;
|
||||
volatile float *local_b_tmp = reinterpret_cast<volatile float *>(local_b) +
|
||||
BN_adjusted * local_b_row + local_b_col;
|
||||
|
||||
static_assert(
|
||||
row_stride_b * 8 <= BK_adjusted,
|
||||
"manual loop unrolling condition not met; consider increasing BK");
|
||||
static_assert(
|
||||
(BK_adjusted % (row_stride_b * 8)) == 0,
|
||||
"manual loop unrolling condition not met; BK should be power-of-two");
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t load_offset = 0; load_offset < BK;
|
||||
load_offset += row_stride_b * 8) {
|
||||
// equivalent code:
|
||||
//
|
||||
// *local_b_tmp = *global_b;
|
||||
// global_b += dim_n * row_stride_b;
|
||||
// local_b_tmp += BN * row_stride_b;
|
||||
|
||||
asm volatile ("flw ft0, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft1, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft2, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft3, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft4, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft5, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft6, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
asm volatile ("flw ft7, (%0)" :: "r"(global_b));
|
||||
global_b += dim_n_adjusted * row_stride_b;
|
||||
|
||||
asm volatile ("fsw ft0, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp));
|
||||
asm volatile ("fsw ft1, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp));
|
||||
local_b_tmp += BN_adjusted * row_stride_b * 2;
|
||||
asm volatile ("fsw ft2, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp));
|
||||
asm volatile ("fsw ft3, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp));
|
||||
local_b_tmp += BN_adjusted * row_stride_b * 2;
|
||||
asm volatile ("fsw ft4, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp));
|
||||
asm volatile ("fsw ft5, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp));
|
||||
local_b_tmp += BN_adjusted * row_stride_b * 2;
|
||||
asm volatile ("fsw ft6, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 0 * sizeof(float)), "r"(local_b_tmp));
|
||||
asm volatile ("fsw ft7, %0(%1)" :: "i"(BN_adjusted * row_stride_b * 1 * sizeof(float)), "r"(local_b_tmp));
|
||||
local_b_tmp += BN_adjusted * row_stride_b * 2;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void thread_block_gemm(kernel_arg_t *__UNIFORM__ arg,
|
||||
template <uint32_t tile_dim_row, uint32_t tile_dim_col>
|
||||
inline void thread_block_copy_tile(const float *src, float *dest,
|
||||
const uint32_t tid_in_threadblock,
|
||||
const uint32_t threads_per_threadblock,
|
||||
const uint32_t threadblock_dim_y,
|
||||
/*const uint32_t threadblock_id_x,
|
||||
const uint32_t threadblock_id_y,*/
|
||||
const uint32_t threadblocks_per_cluster,
|
||||
const uint32_t threadblock_id_in_cluster,
|
||||
uint8_t *sharedmem_per_threadblock) {
|
||||
const T *A = (const T *)arg->addr_a;
|
||||
const T *B = (const T *)arg->addr_b;
|
||||
float *C = (float *)arg->addr_c;
|
||||
const uint32_t threadblock_id_in_cluster) {
|
||||
asm volatile("threadblock_copy_tile_start_%=:" ::);
|
||||
|
||||
const uint32_t dim_m = arg->dim_m;
|
||||
const uint32_t dim_n = arg->dim_n;
|
||||
const uint32_t dim_k = arg->dim_k;
|
||||
|
||||
const uint32_t local_a_row = tid_in_threadblock / BK;
|
||||
const uint32_t local_a_col = tid_in_threadblock % BK;
|
||||
const uint32_t local_as_row = tid_in_threadblock / BM;
|
||||
const uint32_t local_as_col = tid_in_threadblock % BM;
|
||||
const uint32_t local_b_row = tid_in_threadblock / BN;
|
||||
const uint32_t local_b_col = tid_in_threadblock % BN;
|
||||
|
||||
// no double-buffering
|
||||
const uint32_t threads_per_warpgroup = threads_per_threadblock;
|
||||
const uint32_t warp_id_in_warpgroup = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warp_row = warp_id_in_warpgroup / (BN / WN);
|
||||
const uint32_t warp_col = warp_id_in_warpgroup % (BN / WN);
|
||||
const uint32_t tid_in_warp = tid_in_threadblock % NUM_THREADS;
|
||||
|
||||
volatile T *local_a = reinterpret_cast<T *>(sharedmem_per_threadblock);
|
||||
constexpr size_t local_a_elems = (BM * BK);
|
||||
volatile T *local_a_buf = local_a + local_a_elems;
|
||||
|
||||
volatile T *local_b = local_a_buf + local_a_elems;
|
||||
constexpr size_t local_b_elems = (BK * BN);
|
||||
volatile T *local_b_buf = local_a_buf + local_b_elems;
|
||||
|
||||
constexpr uint32_t skips =
|
||||
loop_matmul_skips(/*skip_lda=*/0, /*skip_ldb=*/0, /*skip_ldd=*/1,
|
||||
/*skip_ex=*/1, /*skip_stc=*/1);
|
||||
|
||||
#if (GEMMINI_DMA == 1)
|
||||
if (tid_in_threadblock == 0) {
|
||||
gemmini_extended_config_ex(WEIGHT_STATIONARY, 0, 0, 1, 0, 0);
|
||||
// gemmini_extended_config_ex(dataflow, act & 3, 0, 1, a_transpose,
|
||||
// b_transpose);
|
||||
|
||||
gemmini_extended3_config_ld(dim_k * sizeof(elem_t), MVIN_SCALE_IDENTITY,
|
||||
false, 0);
|
||||
gemmini_extended3_config_ld(dim_n * sizeof(elem_t), MVIN_SCALE_IDENTITY,
|
||||
false, 1);
|
||||
gemmini_extended_config_st(dim_n * sizeof(elem_t), 0, MVIN_SCALE_IDENTITY);
|
||||
|
||||
gemmini_fence();
|
||||
}
|
||||
#endif
|
||||
|
||||
// divide rows (M) by the number of threadblocks
|
||||
const uint32_t dim_m_range = (dim_m / threadblocks_per_cluster);
|
||||
const uint32_t dim_m_start = dim_m_range * threadblock_id_in_cluster;
|
||||
const uint32_t block_m_start = dim_m_start / BM;
|
||||
const uint32_t block_m_end = (dim_m_start + dim_m_range) / BM;
|
||||
const uint32_t warp_id = tid_in_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_in_threadblock = threads_per_threadblock / NUM_THREADS;
|
||||
const uint32_t warps_per_threadblock_per_core =
|
||||
warps_in_threadblock / CORES_PER_CLUSTER;
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t block_m = block_m_start; block_m < block_m_end; block_m++) {
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t block_n = 0; (block_n * BN) < dim_n; block_n++) {
|
||||
// clear out C
|
||||
initialize_C(0);
|
||||
initialize_C(1);
|
||||
for (int row_offset = 0; row_offset < tile_dim_row;
|
||||
row_offset += warps_in_threadblock) {
|
||||
const uint32_t row = row_offset + warp_id;
|
||||
const uint32_t first_thread_offset = tile_dim_col * row;
|
||||
|
||||
if constexpr (GEMMINI_DMA) {
|
||||
// pipeline initiation
|
||||
if (tid_in_threadblock == 0) {
|
||||
// configure dma gmem address to load from
|
||||
// FIXME: block_k is wrong
|
||||
ROCC_INSTRUCTION_RS1_RS2(
|
||||
XCUSTOM_ACC,
|
||||
(uint64_t)(A + block_m * BM * dim_k + /*block_k:*/0 * BK),
|
||||
(uint64_t)(B + /*block_k:*/0 * BK * dim_n + block_n * BN),
|
||||
k_LOOP_WS_CONFIG_ADDRS_AB)
|
||||
// GEMMINI_CISC(8) does k_LOOP_WS_CONFIG_STRIDES_AB
|
||||
GEMMINI_CISC_CMD_R((dim_n << 16) | (dim_k << 8) | 8);
|
||||
gemmini_fence();
|
||||
|
||||
GEMMINI_CISC_CMD_I(10);
|
||||
gemmini_fence();
|
||||
|
||||
#if 0
|
||||
// sp_tiled_matmul_full_spad_ws includes CONFIG_BOUNDS
|
||||
// FIXME: block_k is 0 for two times
|
||||
sp_tiled_matmul_full_spad_ws(
|
||||
#if 1
|
||||
SPAD_ADDR_Q0, SPAD_ADDR_Q1,
|
||||
#else
|
||||
(/*block_k:*/ 0 & 1) ? SPAD_ADDR_Q2 : SPAD_ADDR_Q0,
|
||||
(/*block_k:*/ 0 & 1) ? SPAD_ADDR_Q3 : SPAD_ADDR_Q1,
|
||||
#endif
|
||||
/*spad_D=*/0, /*spad_C=*/SPAD_ADDR_Q3,
|
||||
/*I=*/BM / DIM, /*J=*/BN / DIM, /*K=*/BK / DIM, /*pad_I=*/0,
|
||||
/*pad_J=*/0, /*pad_K=*/0,
|
||||
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
|
||||
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips)
|
||||
gemmini_fence();
|
||||
#endif
|
||||
constexpr uint32_t per_row_iter = tile_dim_col / NUM_THREADS;
|
||||
uint32_t thread_offset = first_thread_offset + tid_in_warp;
|
||||
#pragma GCC unroll
|
||||
for (int i = 0; i < per_row_iter; i++) {
|
||||
dest[thread_offset] = src[thread_offset];
|
||||
thread_offset += NUM_THREADS;
|
||||
}
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster, threadblock_dim_y);
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
}
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (uint32_t block_k = 0; (block_k * BK) < (dim_k); block_k++) {
|
||||
|
||||
// producer code: GMEM->SMEM memory movement
|
||||
// ---------------------------------------------------------------------
|
||||
//
|
||||
// this is either done using DMA or SIMT cores depending on GEMMINI_DMA
|
||||
|
||||
#if (GEMMINI_DMA == 1)
|
||||
if ((tid_in_threadblock == 0) && ((block_k * BK) != (dim_k - BK))) {
|
||||
// configure dma gmem address to load from
|
||||
// FIXME: block_k is wrong
|
||||
ROCC_INSTRUCTION_RS1_RS2(
|
||||
XCUSTOM_ACC,
|
||||
(uint64_t)(A + block_m * BM * dim_k + (block_k + 1/*runahead*/) * BK),
|
||||
(uint64_t)(B + (block_k + 1/*runahead*/) * BK * dim_n + block_n * BN),
|
||||
k_LOOP_WS_CONFIG_ADDRS_AB)
|
||||
// GEMMINI_CISC(8) does k_LOOP_WS_CONFIG_STRIDES_AB
|
||||
GEMMINI_CISC_CMD_R((dim_n << 16) | (dim_k << 8) | 8);
|
||||
// gemmini_fence();
|
||||
|
||||
// block_k is even: opcode 11 (write to local_a_buf)
|
||||
// block_k is odd: opcode 10 (write to local_a)
|
||||
const uint32_t opcode = 11 - (block_k & 1);
|
||||
GEMMINI_CISC_CMD_R(opcode);
|
||||
// // TODO: branch is probably slow
|
||||
// if (block_k & 1) {
|
||||
// GEMMINI_CISC_CMD_I(12);
|
||||
// } else { // block_k == 0 is here
|
||||
// GEMMINI_CISC_CMD_I(13);
|
||||
// }
|
||||
|
||||
// configure loop iteration bounds
|
||||
// FIXME: shouldn't be necessary
|
||||
// ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC, 0, BOUND_INST,
|
||||
// k_LOOP_WS_CONFIG_BOUNDS) ROCC_INSTRUCTION_RS1_RS2(XCUSTOM_ACC,
|
||||
// SPAD_ADDR_Q0, SPAD_ADDR_Q1, k_LOOP_WS_CONFIG_SPAD_AB)
|
||||
// ROCC_INSTRUCTION_RS1_RS2(
|
||||
// XCUSTOM_ACC,
|
||||
// ((uint64_t)(/*a_spad_id:*/ 0) << 18) |
|
||||
// ((uint64_t)(/*b_spad_id:*/ 0) << 16) |
|
||||
// ((uint64_t)(/*act:0*/ 0) << 8) | ((/*low_D:*/ 0) << 2) |
|
||||
// ((/*full_C:*/ 0) << 1) | (/*ex_accumulate:*/ 0),
|
||||
// ((uint64_t)(/*C_spad_addr:*/ A) << 32) | 0x200U | (skips) |
|
||||
// ((/*is_resadd*/ 0) << 2) | ((/*B_transpose:*/ 0) << 1) |
|
||||
// (/*A_transpose:*/ 1),
|
||||
// k_LOOP_WS)
|
||||
// gemmini_fence();
|
||||
|
||||
#if 0
|
||||
uint32_t spad_a_produce;
|
||||
uint32_t spad_b_produce;
|
||||
const uint32_t mask_odd = (block_k & 1) << 31 >> 31;
|
||||
const uint32_t mask_even = ((block_k & 1) ^ 1) << 31 >> 31;
|
||||
spad_a_produce =
|
||||
((mask_odd & (SPAD_ADDR_Q0)) | (mask_even & (SPAD_ADDR_Q2)));
|
||||
spad_b_produce =
|
||||
((mask_odd & (SPAD_ADDR_Q1)) | (mask_even & (SPAD_ADDR_Q3)));
|
||||
// sp_tiled_matmul_full_spad_ws includes CONFIG_BOUNDS
|
||||
// FIXME: block_k is 0 for two times
|
||||
sp_tiled_matmul_full_spad_ws(
|
||||
spad_a_produce,
|
||||
spad_b_produce,
|
||||
/*spad_D=*/0, /*spad_C=*/SPAD_ADDR_Q1,
|
||||
/*I=*/BM / DIM, /*J=*/BN / DIM, /*K=*/BK / DIM, /*pad_I=*/0,
|
||||
/*pad_J=*/0, /*pad_K=*/0,
|
||||
/*a_transpose=*/0, /*b_transpose=*/0, /*full_C=*/0, /*low_D=*/0,
|
||||
/*acc=*/0, /*act=*/NO_ACTIVATION, /*skips=*/skips)
|
||||
#endif
|
||||
}
|
||||
#else
|
||||
global_dmem_load<T>(dim_n, dim_k, block_k * BK, A, B, local_a, local_b,
|
||||
tid_in_threadblock, block_n, block_m);
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster, threadblock_dim_y);
|
||||
#endif
|
||||
|
||||
// consumer code: SMEM->RF and compute
|
||||
// ----------------------------------------------------------------------
|
||||
// @perf: this loop spills to stack a lot because of all the flws in
|
||||
const volatile T *local_a_consume;
|
||||
const volatile T *local_b_consume;
|
||||
if constexpr (GEMMINI_DMA) {
|
||||
// local_a_consume = (k_index % 2) ? local_a_buf : local_a;
|
||||
// local_b_consume = (k_index % 2) ? local_b_buf : local_b;
|
||||
// FIXME: swap multiply with bitshifts
|
||||
// const uint32_t mask_odd = (block_k & 1) << 31 >> 31;
|
||||
// const uint32_t mask_even = ((block_k & 1) ^ 1) << 31 >> 31;
|
||||
// local_a_consume = reinterpret_cast<volatile T *>(
|
||||
// (mask_odd & reinterpret_cast<uintmax_t>(local_a_buf)) |
|
||||
// (mask_even & reinterpret_cast<uintmax_t>(local_a)));
|
||||
// local_b_consume = reinterpret_cast<volatile T *>(
|
||||
// (mask_odd & reinterpret_cast<uintmax_t>(local_b_buf)) |
|
||||
// (mask_even & reinterpret_cast<uintmax_t>(local_b)));
|
||||
local_a_consume = local_a + (block_k & 1) * (local_a_elems);
|
||||
local_b_consume = local_b + (block_k & 1) * (local_b_elems);
|
||||
} else {
|
||||
// no double-buffering without DMA
|
||||
local_a_consume = local_a;
|
||||
local_b_consume = local_b;
|
||||
}
|
||||
|
||||
#pragma GCC unroll 1
|
||||
for (int i = 0; i < BK_LOOP; i++) {
|
||||
#pragma GCC unroll 4
|
||||
for (uint32_t local_k = 0; local_k < BK; local_k += TCK) {
|
||||
#pragma GCC unroll 2
|
||||
for (int wn_iter = 0; wn_iter < WNITER; wn_iter++) {
|
||||
// SMEM -> RF
|
||||
vx_wmma_load_b<T>(local_b_consume, local_k, warp_col, wn_iter,
|
||||
tid_in_warp);
|
||||
#pragma GCC unroll 2
|
||||
for (int wm_iter = 0; wm_iter < WMITER; wm_iter++) {
|
||||
// SMEM -> RF
|
||||
vx_wmma_load_a<T>(local_a_consume, local_k, warp_row, wm_iter,
|
||||
tid_in_warp);
|
||||
// perform mma
|
||||
vx_wmma(wm_iter);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (GEMMINI_DMA) {
|
||||
// Call gemmini fence at the end of the loop to overlap dma & wmma.
|
||||
// Hopefully by this time, dma would have finished so that this is a
|
||||
// no-op
|
||||
if (tid_in_threadblock == 0) {
|
||||
gemmini_fence();
|
||||
}
|
||||
}
|
||||
|
||||
threadblock_barrier(threadblock_id_in_cluster, threadblock_dim_y);
|
||||
}
|
||||
|
||||
#pragma GCC unroll 2
|
||||
for (int wm_iter = 0; wm_iter < WMITER; wm_iter++) {
|
||||
#pragma GCC unroll 2
|
||||
for (int wn_iter = 0; wn_iter < WNITER; wn_iter++) {
|
||||
write_results(tid_in_warp, warp_col, warp_row, wn_iter, wm_iter,
|
||||
dim_n, C, block_n, block_m);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
asm volatile("threadblock_copy_tile_finish_%=:" ::);
|
||||
}
|
||||
|
||||
void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
|
||||
@@ -547,18 +53,21 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
|
||||
constexpr uint32_t cores_per_cluster = 1;
|
||||
#endif
|
||||
|
||||
uint32_t threads_per_threadblock = (BM * BN) / (ELEM_PER_THREAD);
|
||||
const uint32_t hw_threads_per_cluster =
|
||||
cores_per_cluster * vx_num_threads() * vx_num_warps();
|
||||
constexpr uint32_t threads_per_threadblock_theoretical =
|
||||
(BM * BN) / (ELEM_PER_THREAD);
|
||||
constexpr uint32_t hw_threads_per_cluster =
|
||||
CORES_PER_CLUSTER * NUM_THREADS * NUM_WARPS;
|
||||
// cap maximum threadblock size to # of HW threads in cluster, to prevent
|
||||
// multiple "wave" invocations which slows down the kernel
|
||||
if (threads_per_threadblock > hw_threads_per_cluster) {
|
||||
threads_per_threadblock = hw_threads_per_cluster;
|
||||
}
|
||||
const uint32_t threadblocks_per_cluster =
|
||||
constexpr uint32_t threads_per_threadblock =
|
||||
(threads_per_threadblock_theoretical > hw_threads_per_cluster)
|
||||
? hw_threads_per_cluster
|
||||
: threads_per_threadblock_theoretical;
|
||||
constexpr uint32_t threadblocks_per_cluster =
|
||||
hw_threads_per_cluster / threads_per_threadblock;
|
||||
constexpr uint32_t warps_per_threadblock_per_core =
|
||||
NUM_WARPS / threadblocks_per_cluster;
|
||||
|
||||
const uint32_t threadblock_dim_y = vx_num_warps() / threadblocks_per_cluster;
|
||||
const int threadblock_id = task_id / threads_per_threadblock;
|
||||
const int threadblock_id_in_cluster =
|
||||
threadblock_id % threadblocks_per_cluster;
|
||||
@@ -572,20 +81,40 @@ void kernel_body(int task_id, kernel_arg_t *__UNIFORM__ arg) {
|
||||
const uint32_t problem_size = (dim_m * dim_n) / (ELEM_PER_THREAD);
|
||||
const uint32_t num_threadblocks = problem_size / threads_per_threadblock;
|
||||
|
||||
using float_type = float;
|
||||
|
||||
// "static" shared memory allocation. This would determine threadblock
|
||||
// occupancy of a single cluster
|
||||
uint8_t *sharedmem_per_threadblock = reinterpret_cast<uint8_t *>(
|
||||
DEV_SMEM_START_ADDR + sizeof(float_type) * 2 /*overkill for non-dma*/ *
|
||||
(2 * BM * BK) * threadblock_id_in_cluster);
|
||||
DEV_SMEM_START_ADDR +
|
||||
sizeof(float_type) * 2 * (2 * BM * BK) * threadblock_id_in_cluster);
|
||||
|
||||
thread_block_gemm<float_type>(
|
||||
arg, tid_in_threadblock, threads_per_threadblock, threadblock_dim_y,
|
||||
/*threadblock_id_x, threadblock_id_y,*/
|
||||
threadblocks_per_cluster,
|
||||
// threadblock_id,
|
||||
threadblock_id_in_cluster, sharedmem_per_threadblock);
|
||||
thread_block_gemm<float_type, threads_per_threadblock,
|
||||
/*write_to_gmem=*/true,
|
||||
/*smem_a_offset=*/0,
|
||||
/*smem_a_dbuf_offset=*/0,
|
||||
/*smem_b_offset=*/2 * BM * BK * sizeof(float),
|
||||
/*smem_b_dbuf_offset=*/2 * BM * BK * sizeof(float)>(
|
||||
(const float_type *)arg->addr_a, (const float_type *)arg->addr_b,
|
||||
(float *)arg->addr_c, arg->dim_m, arg->dim_n, arg->dim_k,
|
||||
tid_in_threadblock, threadblocks_per_cluster, threadblock_id_in_cluster,
|
||||
sharedmem_per_threadblock);
|
||||
|
||||
float *gmem_tmp_d0 = reinterpret_cast<float *>(0xd0000000UL);
|
||||
float *gmem_tmp_d1 = reinterpret_cast<float *>(0xd1000000UL);
|
||||
|
||||
const float *smem_A = reinterpret_cast<float *>(sharedmem_per_threadblock);
|
||||
const float *smem_B = smem_A + 2 * BM * BK;
|
||||
|
||||
if constexpr (DEBUG) {
|
||||
threadblock_barrier(threadblock_id_in_cluster,
|
||||
warps_per_threadblock_per_core);
|
||||
|
||||
thread_block_copy_tile<BM, BK>(smem_A, gmem_tmp_d0, tid_in_threadblock,
|
||||
threads_per_threadblock,
|
||||
threadblock_id_in_cluster);
|
||||
thread_block_copy_tile<BK, BN>(smem_B, gmem_tmp_d1, tid_in_threadblock,
|
||||
threads_per_threadblock,
|
||||
threadblock_id_in_cluster);
|
||||
}
|
||||
}
|
||||
|
||||
int main() {
|
||||
|
||||
@@ -173,7 +173,8 @@ int main(int argc, char *argv[]) {
|
||||
uint32_t dim_n = 64;
|
||||
uint32_t dim_k = 64;
|
||||
|
||||
using float_type = float;
|
||||
using float_type = half;
|
||||
|
||||
generate_source_matrix<float_type>(dim_m, dim_n, dim_k);
|
||||
generate_reference_matmul<float_type>(dim_m, dim_n, dim_k);
|
||||
|
||||
|
||||
1064
tests/regression/sgemm_tcore/sgemm_impl.hpp
Normal file
1064
tests/regression/sgemm_tcore/sgemm_impl.hpp
Normal file
File diff suppressed because it is too large
Load Diff
@@ -1,343 +0,0 @@
|
||||
#ifndef _UTIL_H_
|
||||
#define _UTIL_H_
|
||||
|
||||
#include <vx_intrinsics.h>
|
||||
#include <vx_spawn.h>
|
||||
#include "include/gemmini.h"
|
||||
#include "gemmini_mmio.h"
|
||||
|
||||
// Constraints on parameters:
|
||||
// * Memory:
|
||||
// (BM + BN) * BK * sizeof(T) <= sharedmem size.
|
||||
// BM * BK == BN * BK >= threadblock size >= NT * CORES_PER_CLUSTER
|
||||
// When larger, the kernel runs a sequential loop to read into sharedmem;
|
||||
// but smaller case is not handled.
|
||||
// * Compute:
|
||||
// ( M* N) / (TM*TN) == grid size >= NC*NW*NT
|
||||
// (BM*BN) / (TM*TN) == threadblock size < NT * NW * CORES_PER_CLUSTER
|
||||
// (BM*BN) / (TM*TN) == threadblock size >= NT * CORES_PER_CLUSTER
|
||||
// * Combining BM * BK >= (BM*BN) / (TM*TN) == threadblock yields
|
||||
// BM <= BK*TM*TN
|
||||
#define BM 64
|
||||
#define BN 64
|
||||
#define BK 64
|
||||
#define WM 16
|
||||
#define WN 8
|
||||
#define TCM 8
|
||||
#define TCN 8
|
||||
#define TCK 8
|
||||
#define WMITER (WM / TCM)
|
||||
#define WNITER (WN / TCN)
|
||||
#define ELEM_PER_THREAD (WMITER * WNITER * (TCM * TCN) / NUM_THREADS)
|
||||
|
||||
// number of loop around the inner 0..TCK..BK loop to simulate perfect-DRAM
|
||||
// scenario
|
||||
#define BK_LOOP 1
|
||||
// Whether to transpose smem A tile at GMEM->SMEM (produce), or SMEM->RF
|
||||
// (consume). This is because the tensor core expects the A tile to be stored
|
||||
// in column-major order in SMEM, whereas it will be ultimately stored in
|
||||
// row-major in the RF.
|
||||
//
|
||||
// For correctness, only one of either should be 1. E.g., PRODUCE 1 CONSUME 0
|
||||
// generates the NN kernel where both A and B are stored row-major in GMEM.
|
||||
// To model the case where the A matrix is already stored transposed in GMEM
|
||||
// ("TN" kernel), set both to 0.
|
||||
#define TRANSPOSE_AT_PRODUCE 1
|
||||
#define TRANSPOSE_AT_CONSUME 0
|
||||
// GMEM_COALESCED: When TRANSPOSE_AT_PRODUCE == 1 (i.e. transpose at
|
||||
// GMEM->SMEM), determines whether we do bank-conflict-free accesses for
|
||||
// 1: GMEM loads of A matrix, or
|
||||
// 0: SMEM stores of A matrix.
|
||||
//
|
||||
// Usually, GMEM_COALESCED==1 yields better performance since the memory
|
||||
// behavior of GMEM is more sensitive to bank conflicts.
|
||||
#define GMEM_COALESCED_A 1
|
||||
|
||||
// "fake" fp16 type that only has the correct data width.
|
||||
using float16_t = uint16_t;
|
||||
|
||||
inline constexpr void map_operand_32lanes(const int tid, int &row, int &col) {
|
||||
const int tg = tid / 4;
|
||||
|
||||
// A (row major)
|
||||
// Figure 7(a) in paper
|
||||
// row 0~ 3: threadgroups 0 and 2
|
||||
// row 4~ 7: threadgroups 4 and 6
|
||||
// row 8~11: threadgroups 1 and 3
|
||||
// row 12~15: threadgroups 5 and 7
|
||||
row = tid % 4;
|
||||
row += (tg * 8) % 16;
|
||||
row += (tg / 4) * 4;
|
||||
|
||||
// B (column major)
|
||||
// NOTE: Matrix B mapping in Figure 7(a) is incorrect; below is the
|
||||
// corrected mapping:
|
||||
// col 0~ 3: threadgroups 0 and 1
|
||||
// col 4~ 7: threadgroups 4 and 5
|
||||
// col 8~11: threadgroups 2 and 3
|
||||
// col 12~15: threadgroups 6 and 7
|
||||
col = tid % 4;
|
||||
col += ((tg % 4) / 2) * 8;
|
||||
col += (tg / 4) * 4;
|
||||
}
|
||||
|
||||
inline constexpr void map_operand_8lanes(const int tid, int &row, int &col) {
|
||||
const int tg = tid / 4;
|
||||
|
||||
// A (row major)
|
||||
// row 0~ 3: threadgroup 0
|
||||
// row 4~ 7: threadgroup 1
|
||||
row = tid % 4;
|
||||
row += tg * 4;
|
||||
|
||||
// B (column major)
|
||||
// col 0~ 3: threadgroup 0
|
||||
// col 4~ 7: threadgroup 1
|
||||
col = tid % 4;
|
||||
col += tg * 4;
|
||||
}
|
||||
|
||||
inline constexpr void map_operand(const int tid, int &row, int &col) {
|
||||
if constexpr (NUM_THREADS == 32) {
|
||||
map_operand_32lanes(tid, row, col);
|
||||
} else if constexpr (NUM_THREADS == 8) {
|
||||
map_operand_8lanes(tid, row, col);
|
||||
} else {
|
||||
// FIXME: not allowed
|
||||
}
|
||||
}
|
||||
|
||||
inline constexpr void map_c_32lanes(const int tid, int &row, int &col) {
|
||||
const int tg = tid / 4;
|
||||
|
||||
// C
|
||||
// Figure 7(b), left
|
||||
col = ((tg % 4) / 2) * 8;
|
||||
row = (tg * 8) % 16;
|
||||
row += (tg / 4) * 4;
|
||||
|
||||
// Figure 7(b), right
|
||||
row += (tid % 4) % 2;
|
||||
col += ((tid % 4) / 2) * 2;
|
||||
}
|
||||
|
||||
inline constexpr void map_c_8lanes(const int tid, int &row, int &col) {
|
||||
const int tg = tid / 4;
|
||||
|
||||
// C
|
||||
col = 0;
|
||||
row = tg * 4;
|
||||
|
||||
// Figure 7(b), right
|
||||
row += (tid % 4) % 2;
|
||||
col += ((tid % 4) / 2) * 2;
|
||||
}
|
||||
|
||||
inline constexpr void map_c(const int tid, int &row, int &col) {
|
||||
if constexpr (NUM_THREADS == 32) {
|
||||
map_c_32lanes(tid, row, col);
|
||||
} else if constexpr (NUM_THREADS == 8) {
|
||||
map_c_8lanes(tid, row, col);
|
||||
} else {
|
||||
// FIXME: not allowed
|
||||
}
|
||||
}
|
||||
|
||||
#define RISCV_CUSTOM3 0x7B
|
||||
|
||||
inline void vx_wmma(const int dest_reg) {
|
||||
if (dest_reg == 0) {
|
||||
asm volatile (".insn r %0, 0, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3));
|
||||
} else {
|
||||
asm volatile (".insn r %0, 0, 0, x1, x0, x0" :: "i"(RISCV_CUSTOM3));
|
||||
}
|
||||
}
|
||||
|
||||
// `local_k` is assumed to be multiple of TCK
|
||||
template <typename T>
|
||||
inline void vx_wmma_load_a(volatile const T *smem_A, const int local_k,
|
||||
const int warp_row, const int wm_iter, const int thread_in_warp) {
|
||||
const int tid = thread_in_warp;
|
||||
const int tg = tid / 4;
|
||||
|
||||
// @perf: this is duplicately computed in vx_wmma_load_a and vx_wmma_load_b
|
||||
int row = 0;
|
||||
int col = 0;
|
||||
map_operand(tid, row, col);
|
||||
|
||||
// In fp16 mode, bit-pack two fp16 elements into each fp32 element, and do
|
||||
// data movement at the fp32 granularity. Assuming that the matrix is stored
|
||||
// row-major in GMEM, the packed fp16 pairs belong to the same row,
|
||||
// neighboring columns; therefore, it essentially becomes equivalent to
|
||||
// moving a fp32 matrix whose column dimensions (dim_k/BK/k) are compressed
|
||||
// by a factor of two.
|
||||
constexpr uint32_t packed_factor = (std::is_same_v<T, float16_t> ? 2 : 1);
|
||||
constexpr uint32_t BK_adjusted = BK / packed_factor;
|
||||
|
||||
constexpr int smem_A_rows = BM;
|
||||
constexpr int smem_A_cols = BK_adjusted;
|
||||
constexpr int smem_AS_rows = BK_adjusted;
|
||||
constexpr int smem_AS_cols = BM;
|
||||
|
||||
if constexpr (TRANSPOSE_AT_CONSUME) {
|
||||
// int A_offset = (WM * warp_row + TCM * wm_iter + row) * smem_A_cols;
|
||||
|
||||
// @perf: bank conflicts
|
||||
// f8-f15 stores a single row of A
|
||||
const volatile uint8_t *smem_addr;
|
||||
smem_addr = reinterpret_cast<const volatile uint8_t *>(
|
||||
&reinterpret_cast<const volatile float *>(
|
||||
smem_A)[(WM * warp_row + TCM * wm_iter + row) * smem_A_cols +
|
||||
local_k]);
|
||||
// step to the next column
|
||||
// threads read from different rows; bank conflicts
|
||||
asm volatile("flw f0, %0(%1)" ::"i"(0 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f1, %0(%1)" ::"i"(1 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f2, %0(%1)" ::"i"(2 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f3, %0(%1)" ::"i"(3 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f4, %0(%1)" ::"i"(4 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f5, %0(%1)" ::"i"(5 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f6, %0(%1)" ::"i"(6 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f7, %0(%1)" ::"i"(7 * sizeof(float)), "r"(smem_addr));
|
||||
} else {
|
||||
// read smem A tile as-is; bank-conflict-free AS load
|
||||
// smem A tile is stored column-major
|
||||
// f8-f15 stores a single row of A
|
||||
const volatile uint8_t *smem_addr;
|
||||
smem_addr = reinterpret_cast<const volatile uint8_t *>(
|
||||
&reinterpret_cast<const volatile float *>(
|
||||
smem_A)[((local_k + 0) * smem_AS_cols) +
|
||||
(WM * warp_row + TCM * wm_iter) + row]);
|
||||
// step to the next row
|
||||
// threads read from different columns; no bank conflicts
|
||||
asm volatile("flw f0, %0(%1)" :: "i"(smem_AS_cols * 0 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f1, %0(%1)" :: "i"(smem_AS_cols * 1 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f2, %0(%1)" :: "i"(smem_AS_cols * 2 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f3, %0(%1)" :: "i"(smem_AS_cols * 3 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f4, %0(%1)" :: "i"(smem_AS_cols * 4 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f5, %0(%1)" :: "i"(smem_AS_cols * 5 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f6, %0(%1)" :: "i"(smem_AS_cols * 6 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f7, %0(%1)" :: "i"(smem_AS_cols * 7 * sizeof(float)), "r"(smem_addr));
|
||||
}
|
||||
}
|
||||
|
||||
// `local_k` is assumed to be multiple of TCK
|
||||
template <typename T>
|
||||
inline void vx_wmma_load_b(const volatile T *smem_B, const int local_k,
|
||||
const int warp_col, const int wn_iter,
|
||||
const int thread_in_warp) {
|
||||
const int tid = thread_in_warp;
|
||||
const int tg = tid / 4;
|
||||
|
||||
int row = 0;
|
||||
int col = 0;
|
||||
map_operand(tid, row, col);
|
||||
|
||||
// see comment in vx_wmma_load_a
|
||||
constexpr uint32_t packed_factor = (std::is_same_v<T, float16_t> ? 2 : 1);
|
||||
constexpr uint32_t BN_adjusted = BN / packed_factor;
|
||||
|
||||
constexpr int smem_B_rows = BK;
|
||||
constexpr int smem_B_cols = BN_adjusted;
|
||||
|
||||
// f8-f15 stores a single column of B
|
||||
const volatile uint8_t *smem_addr;
|
||||
smem_addr = reinterpret_cast<const volatile uint8_t *>(
|
||||
&reinterpret_cast<const volatile float *>(
|
||||
smem_B)[((local_k + 0) * smem_B_cols) +
|
||||
(WN * warp_col + TCN * wn_iter) + col]);
|
||||
// step to the next row
|
||||
// threads read from different columns; no bank conflicts
|
||||
asm volatile("flw f8, %0(%1)" :: "i"(smem_B_cols * 0 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f9, %0(%1)" :: "i"(smem_B_cols * 1 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f10, %0(%1)" :: "i"(smem_B_cols * 2 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f11, %0(%1)" :: "i"(smem_B_cols * 3 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f12, %0(%1)" :: "i"(smem_B_cols * 4 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f13, %0(%1)" :: "i"(smem_B_cols * 5 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f14, %0(%1)" :: "i"(smem_B_cols * 6 * sizeof(float)), "r"(smem_addr));
|
||||
asm volatile("flw f15, %0(%1)" :: "i"(smem_B_cols * 7 * sizeof(float)), "r"(smem_addr));
|
||||
}
|
||||
|
||||
inline void initialize_C(const int dest_reg) {
|
||||
// initialize C to zeros
|
||||
if (dest_reg == 0) {
|
||||
asm volatile("fmv.w.x f16, x0");
|
||||
asm volatile("fmv.w.x f17, x0");
|
||||
asm volatile("fmv.w.x f18, x0");
|
||||
asm volatile("fmv.w.x f19, x0");
|
||||
asm volatile("fmv.w.x f20, x0");
|
||||
asm volatile("fmv.w.x f21, x0");
|
||||
asm volatile("fmv.w.x f22, x0");
|
||||
asm volatile("fmv.w.x f23, x0");
|
||||
} else {
|
||||
asm volatile("fmv.w.x f24, x0");
|
||||
asm volatile("fmv.w.x f25, x0");
|
||||
asm volatile("fmv.w.x f26, x0");
|
||||
asm volatile("fmv.w.x f27, x0");
|
||||
asm volatile("fmv.w.x f28, x0");
|
||||
asm volatile("fmv.w.x f29, x0");
|
||||
asm volatile("fmv.w.x f30, x0");
|
||||
asm volatile("fmv.w.x f31, x0");
|
||||
}
|
||||
}
|
||||
|
||||
inline void write_results(const int thread_in_warp, const int warp_col,
|
||||
const int warp_row, const int wn_iter,
|
||||
const int wm_iter, const int dim_n,
|
||||
float *C, const int threadblock_id_x,
|
||||
const int threadblock_id_y) {
|
||||
int tid = thread_in_warp;
|
||||
|
||||
// these are [0, TCM/TCN)
|
||||
int tid_row = 0;
|
||||
int tid_col = 0;
|
||||
map_c(tid, tid_row, tid_col);
|
||||
|
||||
int local_row = (WM * warp_row + TCM * wm_iter) + tid_row;
|
||||
int local_col = (WN * warp_col + TCN * wn_iter) + tid_col;
|
||||
|
||||
float *global_offset_C =
|
||||
C + (BM * threadblock_id_y) * dim_n + BN * threadblock_id_x;
|
||||
|
||||
// @perf: this likely causes a lot of gmem bank conflicts
|
||||
if (wm_iter == 0) {
|
||||
volatile uint8_t *gmem_addr = reinterpret_cast<volatile uint8_t *>(
|
||||
&global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]);
|
||||
volatile uint8_t *gmem_addr_tmp = gmem_addr + (2 * dim_n) * sizeof(float);
|
||||
asm volatile ("fsw f16, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f17, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f18, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f19, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f20, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f21, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f22, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f23, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
// asm volatile ("fsw f16, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]));
|
||||
// asm volatile ("fsw f17, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 1)]));
|
||||
// asm volatile ("fsw f18, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 0)]));
|
||||
// asm volatile ("fsw f19, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 1)]));
|
||||
// asm volatile ("fsw f20, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 4)]));
|
||||
// asm volatile ("fsw f21, %0" :: "m"(global_offset_C[dim_n * (local_row + 0) + (local_col + 5)]));
|
||||
// asm volatile ("fsw f22, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 4)]));
|
||||
// asm volatile ("fsw f23, %0" :: "m"(global_offset_C[dim_n * (local_row + 2) + (local_col + 5)]));
|
||||
} else {
|
||||
volatile uint8_t *gmem_addr = reinterpret_cast<volatile uint8_t *>(
|
||||
&global_offset_C[dim_n * (local_row + 0) + (local_col + 0)]);
|
||||
volatile uint8_t *gmem_addr_tmp = gmem_addr + (2 * dim_n) * sizeof(float);
|
||||
asm volatile ("fsw f24, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f25, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f26, %0(%1)" :: "i"(0 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f27, %0(%1)" :: "i"(1 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f28, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f29, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr));
|
||||
asm volatile ("fsw f30, %0(%1)" :: "i"(4 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
asm volatile ("fsw f31, %0(%1)" :: "i"(5 * sizeof(float)), "r"(gmem_addr_tmp));
|
||||
}
|
||||
}
|
||||
|
||||
inline void threadblock_barrier(const uint32_t barrier_id, const uint32_t count) {
|
||||
vx_fence();
|
||||
vx_barrier(barrier_id, count);
|
||||
}
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user