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

More flash in generate_matrix

Browse Source
This commit is contained in:
Hansung Kim
2024-08-19 21:16:26 -07:00
parent 351e17c849
commit 09afd43904
1 changed files with 40 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
tests/kernel/tensor/generate_matrix.py
View File

@@ -46,14 +46,17 @@ def pack_fp16_by_row(array):
if __name__ == "__main__": if __name__ == "__main__":
M, N, K = parse_mnk() M, N, K = parse_mnk()
# A_array = np.arange(M * K).reshape([M, K]) rand = True
# B_array = np.arange(K * N).reshape([K, N]) if not rand:
# C_array = np.zeros([M, N]) A_array = np.arange(M * K).reshape([M, K])
B_array = np.arange(K * N).reshape([K, N])
np.random.seed(0) C_array = np.arange(M * N).reshape([M, N])
A_array = np.random.rand(M, K) - 0.5 else:
B_array = np.random.rand(K, N) - 0.5 np.random.seed(0)
C_array = np.zeros([M, N]) 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: with open('a_matrix.h', 'w') as f:
for i in range(A_array.shape[0]): for i in range(A_array.shape[0]):
@@ -73,40 +76,58 @@ if __name__ == "__main__":
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)
C_expected = A_array @ B_array D_expected = A_array @ B_array
C_expected.astype('float32').tofile("c_expected.bin") D_expected.astype('float32').tofile("d_expected.bin")
print('C_expected:') print('D_expected:')
print(C_expected) print(D_expected)
fp16 = True fp16 = False
if fp16: if fp16:
A_packed = pack_fp16_by_row(A_array) A_packed = pack_fp16_by_row(A_array)
AT_packed = A_packed.transpose([1, 0, 2]) AT_packed = A_packed.transpose([1, 0, 2])
AT_array = AT_packed.reshape([-1, M * 2]) AT_array = AT_packed.reshape([-1, M * 2])
AT_array.astype('float16').tofile("input.a.bin") AT_array.astype('float16').tofile("input.a.col.bin")
print('AT:')
print(AT_array) print(AT_array)
B_packed = pack_fp16_by_column(B_array) B_packed = pack_fp16_by_column(B_array)
B_array = B_packed.reshape([-1, N * 2]) B_array = B_packed.reshape([-1, N * 2])
B_array.astype('float16').tofile("input.b.bin") B_array.astype('float16').tofile("input.b.row.bin")
print('B:')
print(B_array) print(B_array)
else: else:
A_array.astype('float32').tofile("input.a.row.bin")
AT_array = A_array.transpose([1, 0]) AT_array = A_array.transpose([1, 0])
# AT_array.astype('float32').tofile("input.a.bin") AT_array.astype('float32').tofile("input.a.col.bin")
A_array.astype('float32').tofile("input.a.bin")
B_array.astype('float32').tofile("input.b.bin") B_array.astype('float32').tofile("input.b.bin")
C_array.astype('float32').tofile("input.c.bin")
print('AT:')
print(AT_array) print(AT_array)
print('B:')
print(B_array) print(B_array)
# generate rowmax result in online softmax # generate rowmax result in online softmax
row_max = np.max(C_expected, axis=1) row_max = np.max(D_expected, axis=1)
row_max.astype('float32').tofile("rowmax.bin") row_max.astype('float32').tofile("rowmax.bin")
# subtrace row_max from each row by broadcasting # subtrace row_max from each row by broadcasting
# (placeholder for exp) # (placeholder for exp)
P = C_expected - row_max[:, np.newaxis] x = D_expected - row_max[:, np.newaxis]
P = (x**2) / 2.0 + x + 1.0
# for i in range(3, 4):
# P += (x**i) / np.math.factorial(i)
# P = np.exp(exp)
P.astype('float32').tofile("P_expected.bin") P.astype('float32').tofile("P_expected.bin")
# print('P error:')
# print(P / np.exp(x))
print('P_expected:') print('P_expected:')
print(P) print(P)
row_sum = np.sum(P, axis=1) row_sum = np.sum(P, axis=1)
row_sum.astype('float32').tofile("rowsum.bin") row_sum.astype('float32').tofile("rowsum.bin")
V = C_array
# O = P.transpose([1, 0]) @ V
O = P @ V
O.astype('float32').tofile("O_expected.bin")
print('O_expected:')
print(O)