IQ.Pilot Release Commit @ b6534c0

artifacts/package_sources/tinygrad/extra/thunder/metal/gemm.py

@@ -0,0 +1,74 @@
+# 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())
+
+