K2: 2D grid + scf.for runtime K-loop + shared memory tiling

Author: ganeshnj

CMakeLists.txt

@@ -11,7 +11,8 @@ find_package(LLVM REQUIRED CONFIG)
 # LLVM 18 from apt.llvm.org exports zstd::libzstd_shared but libzstd-dev may
 # not be installed.  Manually define the cmake target if the .so is present.
 if(NOT TARGET zstd::libzstd_shared)
-  find_library(ZSTD_LIB NAMES zstd)
+  find_library(ZSTD_LIB NAMES zstd zstd.so.1 libzstd.so.1.4.8
+               HINTS /usr/lib/aarch64-linux-gnu)
   if(ZSTD_LIB)
     add_library(zstd::libzstd_shared SHARED IMPORTED)
     set_target_properties(zstd::libzstd_shared PROPERTIES

README.md

@@ -152,9 +152,9 @@ Mirrors [Modular's Blackwell series](https://www.modular.com/blog/matrix-multipl
 - [x] Bug fix: `reduce_sum` partial-warp shuffle (`blockSize < 32`) — passes correct `width` to `gpu::ShuffleOp` instead of hardcoded 32
 - [x] Bug fix: `emit_mul` scalar promotion — `_promote_scalar` in `jit.py` prevents `TypeError` when a constexpr int is passed as an IR operand
 - [x] Benchmark notebook — `examples/gemm_benchmark.ipynb` with cuBLAS reference numbers and gap analysis
-- [ ] K0: Naive GEMM — add `//` (FloorDiv) and `%` (Mod) to JIT `BinOp` handler
-- [ ] K1: Row GEMM — rename/reframe `tiled_matmul_kernel` in notebook
-- [ ] K2: Shmem GEMM — `program_id(1)` (2D grid) + `scf.for` runtime K-loop
+- [x] K0: Naive GEMM — `//` (FloorDiv) and `%` (Mod) in JIT `BinOp` handler
+- [x] K1: Row GEMM — one block per output row, constexpr K-loop
+- [x] K2: Shmem GEMM — `program_id(1)` (2D grid) + `ForRangeOp` → `scf.for` runtime K-loop + shared memory tiling
 - [ ] K3: Swizzled GEMM — 128-byte XOR swizzle to eliminate shmem bank conflicts
 - [ ] K4: Vectorized GEMM — `LDG.128` vectorized loads
 - [ ] K5: Pipelined GEMM — `cp.async` to overlap load and compute

bindings/python_bindings.cpp

@@ -155,6 +155,30 @@ PYBIND11_MODULE(_tiny_ton_core, m) {
              return PyValue{self.emitSharedLoad(idx.val, bufferSize, et)};
            },
            py::arg("idx"), py::arg("buffer_size"), py::arg("dtype") = "f32")
+      .def("begin_for_range",
+           [](tinyton::IRBuilder &self, PyValue start, PyValue stop,
+              PyValue step, const std::vector<PyValue> &initArgs) {
+             std::vector<mlir::Value> inits;
+             for (auto &pv : initArgs)
+               inits.push_back(pv.val);
+             auto results =
+                 self.beginForRange(start.val, stop.val, step.val, inits);
+             std::vector<PyValue> out;
+             for (auto v : results)
+               out.push_back(PyValue{v});
+             return out;
+           })
+      .def("end_for_range",
+           [](tinyton::IRBuilder &self, const std::vector<PyValue> &yieldVals) {
+             std::vector<mlir::Value> vals;
+             for (auto &pv : yieldVals)
+               vals.push_back(pv.val);
+             auto results = self.endForRange(vals);
+             std::vector<PyValue> out;
+             for (auto v : results)
+               out.push_back(PyValue{v});
+             return out;
+           })
       .def("emit_ret", [](tinyton::IRBuilder &self) { self.emitRet(); })
       .def("dump_mlir",
            [](tinyton::IRBuilder &self) {
@@ -201,7 +225,10 @@ PYBIND11_MODULE(_tiny_ton_core, m) {
       .def("set_args", &tinyton::SimulatedGPU::setArgs)
       .def("write_memory", &tinyton::SimulatedGPU::writeMemory)
       .def("read_memory", &tinyton::SimulatedGPU::readMemory)
-      .def("run", &tinyton::SimulatedGPU::run);
+      .def("run",
+           [](tinyton::SimulatedGPU &self, int gridX, int gridY,
+              int threadsPerBlock) { self.run(gridX, gridY, threadsPerBlock); },
+           py::arg("grid_x"), py::arg("grid_y"), py::arg("threads_per_block"));
 
 #ifdef TTN_ENABLE_CUDA
   m.def("has_cuda", [] { return tinyton::CUDARuntime::isAvailable(); });
@@ -231,13 +258,13 @@ PYBIND11_MODULE(_tiny_ton_core, m) {
            })
       .def("launch",
            [](tinyton::CUDARuntime &self, const std::string &ptx,
-              const std::string &kernelName, int gridX, int blockX,
+              const std::string &kernelName, int gridX, int gridY, int blockX,
               const std::vector<uintptr_t> &args) {
              std::vector<void *> voidArgs;
              voidArgs.reserve(args.size());
              for (auto a : args)
                voidArgs.push_back(reinterpret_cast<void *>(a));
-             self.launch(ptx, kernelName, gridX, blockX, voidArgs);
+             self.launch(ptx, kernelName, gridX, gridY, blockX, voidArgs);
            });
 #else
   m.def("has_cuda", [] { return false; });

docs/18-shmem-gemm.md

@@ -0,0 +1,130 @@
+# 18 — K2 Shmem GEMM: 2D Grid + Runtime scf.for + Shared Memory Tiling
+
+## Overview
+
+K2 is the third GEMM kernel in tiny-ton's progression toward cuBLAS-level performance.
+It introduces three new compiler features that are independently useful:
+
+| Feature | What it unlocks |
+|---|---|
+| **2D grid launch** — `tt.program_id(1)` | Natural 2D indexing for matrix tiles |
+| **Runtime `scf.for` loop** — `ForRangeOp` | K can be a runtime argument; no IR explosion |
+| **Shared memory tiling** — `tt.shared_store / tt.shared_load` | Block-local caching of A and B tiles |
+
+## Algorithm
+
+```
+Grid  : (M // TM, N // TN)   one block per output tile
+Block : TK threads
+
+For each block (bm, bn):
+  acc = 0
+  for k0 in range(0, K, TK):         ← runtime loop, K is not constexpr
+    shmem_A[0..TK-1]  = A[bm, k0..k0+TK-1]   ← global → shared
+    shmem_B[TK..2TK-1] = B[k0..k0+TK-1, bn]  ← global → shared
+    sync()
+    acc += reduce_sum(shmem_A * shmem_B)
+    sync()
+  C[bm, bn] = acc
+```
+
+With TM = TN = 1 (current), each block computes **one scalar** of C.
+Increasing TM / TN (K3+) harvests the shared memory reuse benefit.
+
+## Architecture: Python DSL → PTX
+
+```
+@tt.jit kernel
+    │
+    ▼  jit.py: visit_For detects runtime bounds → ForRangeOp
+TinyTon IR
+    tinyton.for_range %start, %K, %step iter_args(%acc = %init) {
+      ^bb(%iv: i32, %acc_: f32): ... tinyton.yield %acc2
+    }
+    │
+    ▼  TinyTonToGPU.cpp: lowerOneOp (recursive lambda)
+GPU dialect + SCF
+    scf.for %iv = %start to %K step %step iter_args(%acc = %init) -> f32 {
+      ... scf.yield %acc2
+    }
+    │
+    ▼  Pipeline.cpp: createConvertSCFToCFPass()
+CF dialect
+    cf.cond_br %done, ^exit, ^body
+    ^body: ... cf.br ^header
+    │
+    ▼  CombinedGPULoweringPass → NVPTX backend
+PTX
+    $loop_header:
+      setp.ge.s32 %p0, %k0, %K
+      @%p0 bra $loop_exit
+      ...
+      add.s32 %k0, %k0, 16
+      bra $loop_header
+```
+
+## Compiler changes
+
+### Workstream A — 2D Grid
+
+| File | Change |
+|---|---|
+| `CUDARuntime.h/.cpp` | `launch(gridX, gridY, blockX, ...)` — `gridY` was hardcoded 1 |
+| `Simulator.h/.cpp` | `run(gridX, gridY, threadsPerBlock)` — PID axis=1 → `blockId / gridX` |
+| `python_bindings.cpp` | Updated `launch` and `run` bindings |
+| `jit.py` | `_launch_cuda` / `_launch_simulator` read `grid[1]` |
+
+### Workstream B — Runtime scf.for
+
+| File | Change |
+|---|---|
+| `TinyTonOps.td` | `ForRangeOp` (body region + iter_args) + `YieldOp` |
+| `Builder.h/.cpp` | `beginForRange` / `endForRange` |
+| `TinyTonToGPU.cpp` | Recursive `lowerOneOp` lambda converts `ForRangeOp` → `scf.for` |
+| `CMakeLists.txt` (Conversion) | `MLIRSCFDialect` |
+| `Pipeline.cpp` | `createConvertSCFToCFPass()` before `CombinedGPULoweringPass` |
+| `CMakeLists.txt` (Compiler) | `MLIRSCFToControlFlow` |
+| `python_bindings.cpp` | `begin_for_range` / `end_for_range` |
+| `jit.py` | `visit_For` falls through to `ForRangeOp` when bounds are runtime |
+
+## Kernel source
+
+```python
+@tt.jit
+def shmem_gemm(A_ptr, B_ptr, C_ptr, M, N, K,
+               TM: tt.constexpr, TN: tt.constexpr, TK: tt.constexpr):
+    bm  = tt.program_id(0)
+    bn  = tt.program_id(1)
+    tid = tt.arange(0, TK)
+    acc = 0.0
+    for k0 in range(0, K, TK):
+        a_val = tt.load(A_ptr + bm * TM * K + k0 + tid)
+        tt.shared_store(tid, a_val, buffer_size=2*TK)
+        b_val = tt.load(B_ptr + (k0 + tid) * N + bn * TN)
+        tt.shared_store(TK + tid, b_val, buffer_size=2*TK)
+        tt.sync()
+        a_sh = tt.shared_load(tid,      buffer_size=2*TK)
+        b_sh = tt.shared_load(TK + tid, buffer_size=2*TK)
+        acc  = acc + tt.reduce_sum(a_sh * b_sh)
+        tt.sync()
+    tt.store(C_ptr + bm * TM * N + bn * TN, acc)
+```
+
+## Key result: compile time vs K1
+
+| Kernel | K=128 compile | K=256 compile |
+|---|---|---|
+| K1 Row GEMM (constexpr unroll) | ~200 ms | **hangs (>60 s)** |
+| K2 Shmem GEMM (runtime scf.for) | ~200 ms | ~200 ms |
+
+K2's IR size is **constant** regardless of K — `scf.for` emits a single loop
+instead of N×(K/TK) unrolled copies of the body.
+
+## Next: K3 Swizzled GEMM
+
+XOR-swizzle the shared memory addresses to eliminate bank conflicts:
+```python
+swizzled_idx = idx ^ ((idx >> 3) & 0x7)
+tt.shared_store(swizzled_idx, val, buffer_size=2*TK)
+```
+This requires an address-swizzle helper in the JIT and no new compiler passes.