mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-13 00:48:21 -05:00
75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
# include directory copied from https://github.com/HazyResearch/ThunderMittens
|
|
# https://hazyresearch.stanford.edu/blog/2024-11-28-tk-mlx
|
|
|
|
gemm = """
|
|
#include <metal_stdlib>
|
|
#include "include/tk.metal"
|
|
using namespace mittens;
|
|
|
|
#define GEMM_PARAMS_DEF(T) \
|
|
device T* D [[buffer(0)]], \
|
|
device T* A [[buffer(1)]], \
|
|
device T* B [[buffer(2)]], \
|
|
const constant int &N [[buffer(3)]], \
|
|
const constant int &K [[buffer(4)]], \
|
|
const constant int &M [[buffer(5)]], \
|
|
uint3 tg_id [[threadgroup_position_in_grid]], \
|
|
uint simd_lane_id [[thread_index_in_simdgroup]]
|
|
|
|
template<typename T, unsigned N_BLOCK, unsigned K_BLOCK, unsigned M_BLOCK>
|
|
kernel void matmul_naive(GEMM_PARAMS_DEF(T)) {
|
|
using global_layout = gl<T, 1, 1, -1, -1>;
|
|
global_layout gl_a(A, nullptr, nullptr, N, K);
|
|
global_layout gl_b(B, nullptr, nullptr, K, M);
|
|
global_layout gl_d(D, nullptr, nullptr, N, M);
|
|
rt<T, N_BLOCK * TILE_DIM, K_BLOCK * TILE_DIM> a_reg;
|
|
rt<T, K_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> b_reg;
|
|
rt<float, N_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> d_reg;
|
|
zero(d_reg);
|
|
#pragma clang loop unroll(full)
|
|
for (int k = 0; k < K / (K_BLOCK * TILE_DIM); k++) {
|
|
load(a_reg, gl_a, {0, 0, (int)tg_id.y, k}, simd_lane_id);
|
|
load(b_reg, gl_b, {0, 0, k, (int)tg_id.x}, simd_lane_id);
|
|
mma_AB(d_reg, a_reg, b_reg, d_reg);
|
|
}
|
|
store(gl_d, d_reg, {0, 0, (int)tg_id.y, (int)tg_id.x}, simd_lane_id);
|
|
}
|
|
|
|
#define instantiate_matmul_custom(type_name, T) \
|
|
template [[host_name("matmul_custom_" #type_name)]] [[kernel]] \
|
|
void matmul_naive<T, 4, 2, 4>(GEMM_PARAMS_DEF(T)); \
|
|
|
|
instantiate_matmul_custom(float32, float);
|
|
"""
|
|
|
|
from tinygrad import Device, Tensor, Context
|
|
|
|
if __name__ == "__main__":
|
|
device = Device["METAL"]
|
|
lib = device.compiler.compile(gemm)
|
|
prg = device.runtime("matmul_custom_float32", lib)
|
|
|
|
N = 4096
|
|
a = Tensor.randn(N, N)
|
|
b = Tensor.randn(N, N)
|
|
c = Tensor.empty(N, N)
|
|
Tensor.realize(a, b, c)
|
|
|
|
TILE_DIM = 8
|
|
N_BLOCK = 4
|
|
M_BLOCK = 4
|
|
|
|
gsz = (N // (M_BLOCK * TILE_DIM), N // (N_BLOCK * TILE_DIM), 1)
|
|
for _ in range(5):
|
|
et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
|
|
global_size=gsz, local_size=(32,1,1), vals=(N, N, N), wait=True)
|
|
print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
|
|
|
|
for _ in range(5):
|
|
with Context(DEBUG=2):
|
|
ref = (a@b).realize()
|
|
|
|
print((ref-c).mean().item())
|
|
|
|
|