36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import numpy as np
|
|
|
|
# Function to generate random fp16 values
|
|
def generate_fp16_matrix(size):
|
|
return np.random.rand(size, size).astype(np.float16)
|
|
|
|
# Function to save the matrix to a binary file
|
|
def save_matrix_to_bin(file_name, matrix):
|
|
matrix.tofile(file_name)
|
|
|
|
# Function to perform matrix multiplication and truncate to specified size
|
|
def truncated_matrix_multiplication(matrix_a, matrix_b, size):
|
|
truncated_a = matrix_a.flatten()[:size * size].reshape(size, size)
|
|
truncated_b = matrix_b.flatten()[:size * size].reshape(size, size)
|
|
result = np.matmul(truncated_a, truncated_b)
|
|
return result.astype(np.float16)
|
|
|
|
# Generate the 512x512 matrices
|
|
size = 512
|
|
matrix_a = generate_fp16_matrix(size)
|
|
matrix_b = generate_fp16_matrix(size)
|
|
|
|
# Save the operand matrices to binary files
|
|
# save_matrix_to_bin("input.a.bin", matrix_a)
|
|
# save_matrix_to_bin("input.b.bin", matrix_b)
|
|
|
|
# Generate and save the reference matrices for 128x128, 256x256, and 512x512 sizes
|
|
sizes = [128, 256, 512]
|
|
for s in sizes:
|
|
ref_matrix = truncated_matrix_multiplication(matrix_a, matrix_b, s)
|
|
print(ref_matrix)
|
|
# save_matrix_to_bin(f"ref{s}.bin", ref_matrix)
|
|
|
|
print("All files generated successfully.")
|
|
|