mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-10 06:23:43 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
218 lines
11 KiB
Python
218 lines
11 KiB
Python
import numpy as np, os
|
|
from tinygrad.helpers import getenv, flat_mv
|
|
from tinygrad import dtypes
|
|
|
|
# for copied uops
|
|
from tinygrad import dtypes
|
|
from tinygrad.dtype import DTYPES_DICT
|
|
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# problem variations
|
|
DTYPE_IN = DTYPES_DICT[getenv("DTYPE_IN", "half")]
|
|
DTYPE_OUT = DTYPES_DICT[getenv("DTYPE_OUT", "half")]
|
|
DTYPE_ACC = DTYPES_DICT[getenv("DTYPE_ACC", "float")]
|
|
N = getenv("N", 4096)
|
|
M = getenv("M", N)
|
|
K = getenv("K", N)
|
|
CNT = getenv("CNT", 10)
|
|
ATOL = getenv("ATOL", 5e-3 if DTYPE_IN == dtypes.float else 1e-2)
|
|
RTOL = getenv("RTOL", 1e-4 if DTYPE_IN == dtypes.float else 1e-3)
|
|
FLOPS = M * N * K * 2
|
|
BW = 2 * ((M*K) + (K*N) + (M*N))
|
|
|
|
# algorithm variations
|
|
INPUT = getenv("INPUT", "RAND")
|
|
GEMM_VARIATION = getenv("GEMM_VARIATION", "nv_hcopt")
|
|
|
|
def randoms():
|
|
if INPUT == "RAND":
|
|
na = np.random.default_rng().normal(scale=1.0, size=(M,K)).astype(dtype=np.float32)
|
|
nb = np.random.default_rng().normal(scale=1.0, size=(K,N)).astype(dtype=np.float32)
|
|
elif INPUT == "IDENTITY" and M==N==K:
|
|
na = np.identity(K, dtype=np.float32)
|
|
nb = np.identity(K, dtype=np.float32)
|
|
elif INPUT == "OUTPUTONES" and M==K:
|
|
na = np.identity(K, dtype=np.float32)
|
|
nb = np.ones((K,N), dtype=np.float32)
|
|
else:
|
|
na = np.ones((M,K), dtype=np.float32)
|
|
nb = np.ones((K,N), dtype=np.float32)
|
|
nc = np.zeros(M*N, np.float32)
|
|
if DTYPE_IN != dtypes.float:
|
|
na = na.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
|
|
nb = nb.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
|
|
if DTYPE_OUT != dtypes.float:
|
|
nc = nc.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
|
|
return na, nb, nc
|
|
|
|
if __name__ == "__main__":
|
|
print(f"gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
|
|
prog, global_size, local_size = None, None, None
|
|
|
|
if getenv("CUDA") == 1:
|
|
from tinygrad.runtime.ops_cuda import CUDAAllocator, CUDADevice, CUDAProgram, CUDACompiler
|
|
device = CUDADevice("cuda:0")
|
|
compiler = CUDACompiler(device.arch)
|
|
cudaalloc = CUDAAllocator(device)
|
|
|
|
a = cudaalloc.alloc(M*K*DTYPE_IN.itemsize)
|
|
b = cudaalloc.alloc(K*N*DTYPE_IN.itemsize)
|
|
c = cudaalloc.alloc(M*N*DTYPE_OUT.itemsize)
|
|
|
|
if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
|
|
print("Using CUDA and triton-generated kernel")
|
|
# See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py`
|
|
# this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS. theo max is 165TFLOPS.
|
|
|
|
# WMMA element size is (M, N, K) = (16, 8, 16)
|
|
# warpgroup size in WMMA tiles is (B_M, B_N, B_K) = (2, 8, 4) so 64 HMMA calls per threadgroup reduce iteration
|
|
# thread block size is (T_M, T_N, T_K) = (2, 2, 1), i.e. macro blocks in M and N, so 256 HMMA calls per kernel reduce iteration
|
|
# kernel reduce iteration size in elements = (64, 128, 64)
|
|
# single iteration SMEM_A = (64 * 64) * (2 bytes / half) = 8192 bytes, SMEM_B = (128 * 64) * (2 bytes / half) = 16384 bytes
|
|
# double-buffer smem = (8192 + 16384) * 2 = 49152 bytes
|
|
# reduce for_loop size = [1, 1, (4096 // 16 // 4)==64]
|
|
# NOTE: T_K > 0 would be group_for_reduce
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.max.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//64, N//128, 1],
|
|
'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "2_stage_swizzled_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
|
|
print("Using CUDA, 2-stage reduce pipeline, swizzled SMEM inputs")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.2_stage_swizzled_smem_input.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//64, N//128, 1],
|
|
'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "swizzled_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
|
|
print("Using CUDA, swizzled SMEM inputs")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.swizzled_smem_input.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//64, N//128, 1],
|
|
'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "flat_smem_input" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
|
|
print("Using CUDA, flat SMEM inputs")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp32.flat_smem_input.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//64, N//128, 1],
|
|
'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float:
|
|
print("Using CUDA and generated hcopt")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [32, 64, 1],
|
|
'local_size': [16, 2, 4], # 16,2 are warp, 4 workgroups upcasted to axis=1
|
|
'wait': True,
|
|
}
|
|
elif GEMM_VARIATION == "2_stage" and (M%64)== 0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
|
|
print("Using CUDA and un-optimized 2-stage, swizzled SMEM inputs and direct acc to output kernel")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.2_stage.cu')).read()))
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//64, N//128, 1],
|
|
'local_size': [128, 1, 1], # 4 warpgroups == (T_M:=2) * (T_N:=2)
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "3_stage" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
|
|
print("Using CUDA and 3-stage (interleave global copies and ldmatrix)")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.3_stage.cu')).read()), 73728)
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//256, N//128, 1],
|
|
'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "3_stage_swizzled" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
|
|
print("Using CUDA and 3-stage (interleave global copies and ldmatrix) and swizzled SMEM inputs")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.3_stage_swizzled.cu')).read()), 73728)
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//256, N//128, 1],
|
|
'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "max" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
|
|
print("Using CUDA and 3-stage (interleave global copies and ldmatrix), swizzled SMEM inputs and epilogue")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.max.cu')).read()), 73728)
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//256, N//128, 1],
|
|
'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
elif GEMM_VARIATION == "no_xor" and (M%256)== 0 and (N%128)==0 and (K%32)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.half:
|
|
print("Using CUDA and 3-stage (interleave global copies and ldmatrix), swizzled SMEM inputs and epilogue")
|
|
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp16_fp16.no_xor.cu')).read()), 73728)
|
|
args = (c, a, b)
|
|
kwargs = {
|
|
'global_size': [M//256, N//128, 1],
|
|
'local_size': [32, 4, 2], # 8 warpgroups, WG_M=4 and WG_N=2
|
|
'wait': True,
|
|
'vals': (N, K),
|
|
}
|
|
else:
|
|
raise RuntimeError(f"invalid gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
|
|
|
|
tms = []
|
|
na, nb, nc = randoms()
|
|
cudaalloc._copyin(a, memoryview(bytearray(na)))
|
|
cudaalloc._copyin(b, memoryview(bytearray(nb)))
|
|
for i in range(CNT):
|
|
tms.append(prog(*args, **kwargs))
|
|
cudaalloc._copyout(flat_mv(nc.data), c)
|
|
comp = na.astype(np.float32) @ nb.astype(np.float32)
|
|
result = nc.reshape(M, N).astype(np.float32)
|
|
|
|
print(f"{N*N:10d} {min(tms)*1e6:9.2f} us, would be {FLOPS*1e-9/min(tms):9.2f} GFLOPS matmul, {BW*1e-9/min(tms):.2f} GB/s")
|
|
try:
|
|
np.testing.assert_allclose(result, comp, atol=ATOL, rtol=RTOL)
|
|
except AssertionError as e:
|
|
if getenv("DEBUG_VALUES") > 0:
|
|
indices = np.where(~np.isclose(result, comp, rtol=RTOL, atol=ATOL))
|
|
non_matching_elements_result = result[indices]
|
|
non_matching_elements_comp = comp[indices]
|
|
print("valid :", np.where(np.isclose(result, comp, rtol=RTOL, atol=ATOL)))
|
|
print("invalid :", indices)
|
|
print("result :", non_matching_elements_result)
|
|
print("ground truth:", non_matching_elements_comp)
|
|
print("result sum :", np.sum(result))
|
|
print("ground sum :", np.sum(comp))
|
|
raise e
|
|
|
|
if getenv("DEBUG_VALUES") > 0:
|
|
print(comp)
|
|
print("ground sum :", np.sum(comp))
|
|
print(result)
|
|
print("result sum :", np.sum(result))
|
|
|
|
elif getenv("AMD") == 1:
|
|
# note: https://hipfft.readthedocs.io/en/rocm-6.1.2/how-to/fine-tuning-llms/optimizing-triton-kernel.html
|
|
|
|
# also this is different than the rocblas/tensile approach to GEMM
|
|
# see: https://github.com/ROCm/Tensile/blob/develop/Tensile/KernelWriterAssembly.py
|
|
raise RuntimeError("invalid max_matmul device")
|
|
|
|
else:
|
|
raise RuntimeError("invalid max_matmul device")
|
|
|