kitten matmul is running

extra/thunder/cuda/include/types/global/gl.cuh

@@ -2,7 +2,7 @@
  * @file
  * @brief Templated layouts for global memory.
  */
- 
+
 #pragma once
 
 #include "../../common/common.cuh"
@@ -54,7 +54,7 @@ template<typename _T, int _axis=-9999, bool _swizzle_flag=true> struct descripto
     using T = detail::tma::descriptor_copy_helper_t<_T>;
     static_assert(ducks::st::all<T> || ducks::sv::all<T> || ducks::tma::descriptor::all<T>, "Must be a shared TK type to generate a TMA descriptor.");
     static constexpr int axis = (
-        ducks::tma::descriptor::all<_T> ? detail::tma::descriptor_copy_helper_v<_T> : // if a copy, inherit the axis from the original descriptor. 
+        ducks::tma::descriptor::all<_T> ? detail::tma::descriptor_copy_helper_v<_T> : // if a copy, inherit the axis from the original descriptor.
         (_axis != -9999) ? _axis : detail::tma::descriptor_copy_helper_v<_T>); // if a default value was provided, use it.
     static_assert((kittens::ducks::st::all<T> && axis >= 0 && axis <= 2) || (kittens::ducks::sv::all<T> && axis == -1), "Internal template error detected.");
     static constexpr bool swizzle_flag = ducks::tma::descriptor::all<_T> ? detail::tma::descriptor_copy_helper_swizzle_flag<_T> : _swizzle_flag;
@@ -65,8 +65,8 @@ template<typename _T, int _axis=-9999, bool _swizzle_flag=true> struct descripto
 namespace detail {
 template<typename... Args>
 struct descriptor_dict {
-    __host__ descriptor_dict() {}
-    template<typename T> __host__ descriptor_dict(T _, int b, int d, int r, int c) {}
+    __host__ __device__ descriptor_dict() {}
+    template<typename T> __host__ __device__ descriptor_dict(T _, int b, int d, int r, int c) {}
     __host__ __device__ descriptor_dict(const descriptor_dict &other) {}
 #ifdef KITTENS_HOPPER
     template<typename T, int U> __device__ const CUtensorMap* get() const {
@@ -85,8 +85,8 @@ struct descriptor_dict<_T, Args...> {
     using DESC = kittens::tma::descriptor<_T>; // copy or initialize with a default value
     CUtensorMap tma_desc;
     descriptor_dict<Args...> other_descs;
-    __host__ descriptor_dict() {}
-    __host__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
+    __host__ __device__ descriptor_dict() {}
+    __host__ __device__ descriptor_dict(typename DESC::T::dtype *data, int b, int d, int r, int c): other_descs(data, b, d, r, c) {
         kittens::detail::tma::create_tensor_map<typename DESC::T, DESC::axis, DESC::swizzle_flag>(&tma_desc, data, b, d, r, c);
     }
     __host__ __device__ inline descriptor_dict(const descriptor_dict &other) :
@@ -135,7 +135,7 @@ struct gl {
 
     detail::descriptor_dict<TMA_Types...> tma_descs;
 
-    __host__ inline gl(T *_data,
+    __host__ __device__ inline gl(T *_data,
                         ducks::gl::make_arg_t<b> _batch,
                         ducks::gl::make_arg_t<d> _depth,
                         ducks::gl::make_arg_t<r> _rows,
@@ -160,7 +160,7 @@ struct gl {
         else if constexpr (axis==2) { return size_t(rows()); }
         else if constexpr (axis==3) { return size_t(cols()); }
     }
-    template<int axis> __device__ inline size_t stride() const { 
+    template<int axis> __device__ inline size_t stride() const {
         static_assert(axis==0 || axis==1 || axis==2 || axis==3, "Axis must be 0, 1, 2, or 3.");
         if      constexpr (axis==0) { return depth()*rows()*cols(); }
         else if constexpr (axis==1) { return rows()*cols(); }
@@ -198,7 +198,7 @@ template<int N> auto make_unsafe_gl_arg(int param) { // typename std::conditiona
     if constexpr (N > 0) { return nullptr; }
     else                 { return param;   }
 }
-template<ducks::gl::all GL, bool safe=true> __host__ inline GL make_gl(uint64_t data, int b, int d, int r, int c) {
+template<ducks::gl::all GL, bool safe=true> __host__ __device__ inline GL make_gl(uint64_t data, int b, int d, int r, int c) {
     if constexpr (safe) {
         if(GL::__b__ > 0 && b != GL::__b__) {
             throw std::runtime_error("Batch dimension mismatch. Expected: " + std::to_string(GL::__b__) + ", Got: " + std::to_string(b));

extra/thunder/cuda/include/types/global/tma.cuh

@@ -45,7 +45,7 @@ __host__ static inline std::string format_tma_error(
     oss << "\n  cols: " << cols;
     if (!extra_info.empty())
         oss << "\n  " << extra_info;
-    
+
     oss << "\ncuTensorMapEncodeTiled arguments:";
     oss << "\n  tma_map: " << reinterpret_cast<uintptr_t>(tma_map);
     oss << "\n  tma_format: " << tma_format;
@@ -74,27 +74,27 @@ __host__ static inline std::string format_tma_error(
     for (size_t i = 0; i < gmem_shape_size; ++i)
         oss << gmem_shape[i] << (i < gmem_shape_size - 1 ? ", " : "");
     oss << "]";
-    
+
     oss << "\n  gmem_stride: " << reinterpret_cast<uintptr_t>(gmem_stride) << " [";
     for (size_t i = 0; i < gmem_stride_size; ++i)
         oss << gmem_stride[i] << (i < gmem_stride_size - 1 ? ", " : "");
     oss << "]";
-    
+
     oss << "\n  smem_shape: " << reinterpret_cast<uintptr_t>(smem_shape) << " [";
     for (size_t i = 0; i < smem_shape_size; ++i)
         oss << smem_shape[i] << (i < smem_shape_size - 1 ? ", " : "");
     oss << "]";
-    
+
     oss << "\n  smem_stride: " << reinterpret_cast<uintptr_t>(smem_stride) << " [";
     for (size_t i = 0; i < smem_stride_size; ++i)
         oss << smem_stride[i] << (i < smem_stride_size - 1 ? ", " : "");
     oss << "]";
-    
+
     oss << "\n  tma_interleave: " << tma_interleave;
     oss << "\n  tma_swizzle: " << tma_swizzle;
     oss << "\n  tma_l2Promotion: " << tma_l2Promotion;
     oss << "\n  tma_oobFill: " << tma_oobFill;
-    
+
     return oss.str();
 }
 
@@ -117,7 +117,7 @@ template<ducks::st::all ST, int axis, bool enable_swizzle = true>
 __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typename ST::dtype *src, int batch, int depth, int rows, int cols) {
     using dtype = typename ST::dtype;
     static_assert(axis==0 || axis==1 || axis==2, "axis must be 0, 1, or 2");
-    
+
     constexpr uint32_t  tma_dim = enable_swizzle ? 5 : 4;
     void *global_addr = (void*)(src);
 
@@ -138,7 +138,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
     constexpr CUtensorMapSwizzle      tma_swizzle     = enable_swizzle ? (
         ST::swizzle_bytes == 32  ? CU_TENSOR_MAP_SWIZZLE_32B  :
         ST::swizzle_bytes == 64  ? CU_TENSOR_MAP_SWIZZLE_64B  :
-        ST::swizzle_bytes == 128 ? CU_TENSOR_MAP_SWIZZLE_128B : 
+        ST::swizzle_bytes == 128 ? CU_TENSOR_MAP_SWIZZLE_128B :
         CU_TENSOR_MAP_SWIZZLE_NONE
     ) : CU_TENSOR_MAP_SWIZZLE_NONE;
 
@@ -148,7 +148,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
     uint32_t smem_shape [5] = {0, 0, 0, 0, 0};
     uint32_t smem_stride[5] = {1, 1, 1, 1, 1};
 
-    constexpr uint64_t shared_tile_height = ST::rows; 
+    constexpr uint64_t shared_tile_height = ST::rows;
     constexpr uint64_t shared_tile_width  = ST::cols;
 
     constexpr int swizzle_elements = ST::swizzle_bytes / sizeof(dtype);
@@ -160,7 +160,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
             gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; // round up, note this can potentially screw up out of bounds access handling :/
             gmem_shape[3] = (uint64_t)depth;
             gmem_shape[4] = (uint64_t)batch;
-    
+
             gmem_stride[0] = (uint64_t)cols * sizeof(dtype);
             gmem_stride[1] = ST::swizzle_bytes;
             gmem_stride[2] = (uint64_t)rows * cols * sizeof(dtype);
@@ -172,12 +172,12 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
             gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; // round up, note this can potentially screw up out of bounds access handling :/
             gmem_shape[3] = (uint64_t)rows;
             gmem_shape[4] = (uint64_t)batch;
-    
+
             gmem_stride[0] = (uint64_t)rows * cols * sizeof(dtype);
             gmem_stride[1] = ST::swizzle_bytes;
             gmem_stride[2] = (uint64_t)cols * sizeof(dtype);
             gmem_stride[3] = (uint64_t)depth * rows * cols * sizeof(dtype);
-    
+
         }
         else {
             gmem_shape[0] = swizzle_elements;
@@ -185,7 +185,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
             gmem_shape[2] = (uint64_t)(cols+swizzle_elements-1) / swizzle_elements; // round up, note this can potentially screw up out of bounds access handling :/
             gmem_shape[3] = (uint64_t)rows;
             gmem_shape[4] = (uint64_t)depth;
-    
+
             gmem_stride[0] = (uint64_t)depth * rows * cols * sizeof(dtype);
             gmem_stride[1] = ST::swizzle_bytes;
             gmem_stride[2] = (uint64_t)cols * sizeof(dtype);
@@ -212,7 +212,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
         smem_shape[3] = 1;
     }
 
-    // ensure that the global address is always 16-byte aligned 
+    // ensure that the global address is always 16-byte aligned
     assert((reinterpret_cast<uint64_t>(global_addr) & 0b1111) == 0);
 
     assert(gmem_stride[0] % 16 == 0); // gmem_stride[0] elements must be a multiple of 16B
@@ -239,7 +239,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
     }
 
     const uint64_t *gmem_shape_ptr = &gmem_shape[0];
-    const uint64_t *gmem_stride_ptr = &gmem_stride[0]; 
+    const uint64_t *gmem_stride_ptr = &gmem_stride[0];
     const uint32_t *smem_shape_ptr = &smem_shape[0];
     const uint32_t *smem_stride_ptr = &smem_stride[0];
 
@@ -249,7 +249,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
         tma_dim,
         global_addr,
         gmem_shape_ptr,
-        gmem_stride_ptr, 
+        gmem_stride_ptr,
         smem_shape_ptr,
         smem_stride_ptr,
         tma_interleave,
@@ -331,7 +331,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
     // There is technically a way around ^ that involves instantiating two separate TMA descriptors, one of size 256
     // and the other of size %256, but this is a fairly mild restriction and the other approach is a real PITA and incurs other costs.
     static_assert(disable_swizzle, "for vector TMA, swizzle should be disabled");
-    
+
     constexpr uint32_t  tma_dim     = 4;
     void               *global_addr = (void*)(src);
 
@@ -359,13 +359,13 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
     uint32_t smem_shape [4] = {(uint32_t)dim1, 1, 1, 1};
     uint32_t smem_stride[4] = {1, 1, 1, 1};
 
-    // ensure that the global address is always 16-byte aligned 
+    // ensure that the global address is always 16-byte aligned
     assert((reinterpret_cast<uint64_t>(global_addr) & 0b1111) == 0);
 
     assert(smem_shape[0] <= 256); // smem_shape[0] elements must be <= 256.
 
     const uint64_t *gmem_shape_ptr = &gmem_shape[0];
-    const uint64_t *gmem_stride_ptr = &gmem_stride[0]; 
+    const uint64_t *gmem_stride_ptr = &gmem_stride[0];
     const uint32_t *smem_shape_ptr = &smem_shape[0];
     const uint32_t *smem_stride_ptr = &smem_stride[0];
 
@@ -375,7 +375,7 @@ __host__ static inline void create_tensor_map(CUtensorMap *tma_map, const typena
         tma_dim,
         global_addr,
         gmem_shape_ptr,
-        gmem_stride_ptr, 
+        gmem_stride_ptr,
         smem_shape_ptr,
         smem_stride_ptr,
         tma_interleave,

extra/thunder/cuda/simple.cu

@@ -0,0 +1,45 @@
+// https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/matmul/educational/level_04.cu
+#include "kittens.cuh"
+using namespace kittens;
+
+constexpr int g_N = 1024;
+constexpr int BLOCK_SIZE = 32;
+#define NUM_WORKERS (1)
+#define NUM_THREADS (NUM_WORKERS*kittens::WARP_THREADS)
+
+using sub_tile = st_bf<BLOCK_SIZE,BLOCK_SIZE>;
+using tile_gl =  gl<bf16,  1, 1, g_N, g_N, sub_tile>;
+
+__global__ void kernel(bf16 *c_ptr, bf16 *a_ptr, bf16 *b_ptr) {
+  tile_gl g_C{c_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_A{a_ptr, nullptr, nullptr, nullptr, nullptr};
+  tile_gl g_B{b_ptr, nullptr, nullptr, nullptr, nullptr};
+
+  extern __shared__ alignment_dummy __shm[];
+  shared_allocator al((int*)&__shm[0]);
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &As = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+  st_bf<BLOCK_SIZE,BLOCK_SIZE> &Bs = al.allocate<st_bf<BLOCK_SIZE,BLOCK_SIZE>>();
+
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> A_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE> B_reg;
+  rt_bf<BLOCK_SIZE,BLOCK_SIZE, ducks::rt_layout::col> B_reg_col;
+  rt_fl<BLOCK_SIZE,BLOCK_SIZE> C_accum;
+
+  int col = blockIdx.x;
+  int row = blockIdx.y;
+
+  warp::zero(C_accum);
+  int num_tiles = (g_N + BLOCK_SIZE - 1) / BLOCK_SIZE;
+  for (int tile = 0; tile < num_tiles; ++tile) {
+    warp::load(As, g_A, {0, 0, row, tile});
+    warp::load(Bs, g_B, {0, 0, tile, col});
+    __syncthreads();
+    warp::load(A_reg, As);
+    warp::load(B_reg, Bs);
+    warp::swap_layout(B_reg_col, B_reg);
+    __syncthreads();
+    warp::mma_AB(C_accum, A_reg, B_reg_col, C_accum);
+    __syncthreads();
+  }
+  warp::store(g_C, C_accum, {0, 0, row, col});
+}

extra/thunder/cuda/simple.py

@@ -0,0 +1,39 @@
+import os, pathlib
+os.environ["CUDA_NVCC"] = '1'
+from tinygrad import Device, Tensor
+from tinygrad.helpers import Context, getenv
+from tinygrad.runtime.support.compiler_cuda import pretty_ptx
+
+if __name__ == "__main__":
+  code = pathlib.Path("simple.cu").read_text()
+  device = Device["CUDA"]
+  lib = device.compiler.compile(code)
+  kernel_name = lib.decode().split(".globl\t")[1].split("\n")[0]
+  print("kernel name", kernel_name)
+  #print(pretty_ptx(lib.decode()))
+
+  prg = device.runtime(kernel_name, lib)
+  prg.smem = 10000
+
+  N = 1024
+  a = Tensor.randn(N, N, device='CUDA')
+  b = Tensor.randn(N, N, device='CUDA')
+  c = Tensor.empty(N, N, device='CUDA')
+  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), 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(), (ref-c).max().item())
+

tinygrad/runtime/ops_cuda.py

@@ -116,10 +116,11 @@ class CUDADevice(Compiled):
     self.pending_copyin: list[tuple[int, int, BufferSpec|None]] = []
     CUDADevice.devices.append(self)
 
+    kitten_args = ["-I./include", "-std=c++20", "--expt-relaxed-constexpr", "-DKITTENS_HOPPER"]
     from tinygrad.runtime.graph.cuda import CUDAGraph
     compilers:list[CompilerPairT] = [(functools.partial(CUDARenderer, self.arch), functools.partial(CUDACompiler, self.arch)),
                                      (functools.partial(PTXRenderer, self.arch), functools.partial(PTXCompiler, self.arch)),
-                                     (functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch))]
+                                     (functools.partial(CUDARenderer, self.arch), functools.partial(NVCCCompiler, self.arch, kitten_args))]
     super().__init__(device, CUDAAllocator(self), compilers, functools.partial(CUDAProgram, self), None if MOCKGPU else CUDAGraph)
 
   def synchronize(self):