From 1b5daccac90b1e40520b265c451494ebbcba9b04 Mon Sep 17 00:00:00 2001 From: Hansung Kim Date: Wed, 31 Jul 2024 16:07:03 -0700 Subject: [PATCH] tensor: Generate fp16-packed matrix in script --- tests/kernel/tensor/generate_matrix.py | 71 +++++++++++++++++++------- 1 file changed, 52 insertions(+), 19 deletions(-) diff --git a/tests/kernel/tensor/generate_matrix.py b/tests/kernel/tensor/generate_matrix.py index 35ad7d73..6833d58a 100644 --- a/tests/kernel/tensor/generate_matrix.py +++ b/tests/kernel/tensor/generate_matrix.py @@ -1,7 +1,10 @@ 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.random.rand(8, 16) +A_array = np.arange(8 * 8).reshape([8, 8]) +B_array = np.arange(8 * 8).reshape([8, 8]) +# C_array = np.random.rand(16, 16) +C_array = np.zeros([8, 8]) # A_array = np.zeros((16, 8)) # B_array = np.zeros((8, 16)) # A_array[0,:] = 1.0 @@ -11,22 +14,52 @@ C_array = np.random.rand(16, 16) # for j in range(16): # C_array[i,j] = i * 16 + j -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('\n') +# 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] -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('\n') + T = array.transpose([1, 0]) + T_packed = T.reshape([cols, -1, 2]) + result = T_packed.transpose([1, 0, 2]).reshape([rows // 2, cols * 2]) + return result -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('\n') -np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array) \ No newline at end of file +if __name__ == "__main__": + 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, ') + f.write('\n') + 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, ') + f.write('\n') + 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, ') + f.write('\n') + + np.savez("abc", A_array=A_array, B_array=B_array, C_array=C_array) + + A_array.astype('float32').tofile("input.a.bin") + B_array.astype('float32').tofile("input.b.bin") + + # A_array.astype('float16').tofile("input.a.bin") + # B_array = pack_fp16_by_column(B_array) + # B_array.astype('float16').tofile("input.b.bin") + print(B_array)