OpenBLAS/install/generate.py

import numpy as np
import torch
# 设置矩阵尺寸
M, K, N = 31, 31, 31  # 可修改为更大规模

# 生成随机输入矩阵，类型为float16
A = np.random.randint(0, 11, size=(M, K)).astype(np.float16)
B = np.random.randint(0, 11, size=(K, N)).astype(np.float16)
A_torch = torch.tensor(A, dtype=torch.float16, device='cuda')
B_torch = torch.tensor(B, dtype=torch.float16, device='cuda')
C_torch = torch.matmul(A_torch, B_torch)
C_ref = C_torch.cpu().numpy().astype(np.float32)

def format_array_c(name, array, c_type="hfloat16"):
    flat = array.flatten()
    elements = ", ".join(f"{x:.5f}" for x in flat)
    return f"{c_type} {name}[{len(flat)}] = {{ {elements} }};\n"

def format_array_c_float(name, array):
    flat = array.flatten()
    elements = ", ".join(f"{x:.5f}" for x in flat)
    return f"float {name}[{len(flat)}] = {{ {elements} }};\n"

# 写入C文件
with open("generated_test.c", "w") as f:
    f.write('#include <stdio.h>\n')
    f.write('#include <stdlib.h>\n')
    f.write('#include <string.h>\n')
    f.write('#include <cblas.h>\n\n')

    f.write(f"const int M = {M}, K = {K}, N = {N};\n")
    f.write("const float alpha = 1.0f, beta = 0.0f;\n\n")

    f.write(format_array_c("A", A))
    f.write(format_array_c("B", B))
    f.write(f"float C[{M*N}] = {{ 0 }};\n\n")

    f.write("int main() {\n")
    f.write("    cblas_shgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,\n")
    f.write("                 M, N, K,\n")
    f.write("                 alpha,\n")
    f.write("                 A, K,\n")
    f.write("                 B, N,\n")
    f.write("                 beta,\n")
    f.write("                 C, N);\n\n")

    f.write('    printf("Result C = A * B:\\n");\n')
    f.write("    for (int i = 0; i < M * N; i++) {\n")
    f.write("        printf(\"%.5f \", C[i]);\n")
    f.write("        if ((i + 1) % N == 0) printf(\"\\n\");\n")
    f.write("    }\n")
    f.write("    return 0;\n")
    f.write("}\n\n")

    f.write("// Reference result computed in Python:\n")
    c_ref_flat = ", ".join(f"{x:.5f}" for x in C_ref.flatten())
    f.write(f"// C_ref = {{ {c_ref_flat} }}\n")