From 07d4731e9ffbf353ca85010b7f1a68c27fd9c059 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20O=C5=BC=C3=B3g?= <szymon.ozog@gmail.com>
Date: Fri, 18 Aug 2023 16:17:49 +0200
Subject: [PATCH] Track local size

---
 tinygrad/runtime/ops_triton.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tinygrad/runtime/ops_triton.py b/tinygrad/runtime/ops_triton.py
index f571dc3d81..f7bffb580d 100644
--- a/tinygrad/runtime/ops_triton.py
+++ b/tinygrad/runtime/ops_triton.py
@@ -69,6 +69,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
                 assert var.min == 0, "local loop must start at 0"
                 kk(f"{var.expr} = tl.arange({0}, {var.max+1})[{', '.join([':' if i == j else 'None' for j in range(len(args[0]))])}]")
                 acc_local_shape *= var.max+1
+                local_size.append(var.max+1)
               else:
                 kk(f"for {var.expr} in range({var.min}, {var.max+1}):")
                 depth += 1