From 12752b8a447c1db6fceb43657f02db0305b27d37 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Sun, 31 May 2026 17:59:48 -0700
Subject: [PATCH] fix wmma

---
 tinygrad/renderer/cstyle.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py
index 4ca71687cb..5b07b0da56 100644
--- a/tinygrad/renderer/cstyle.py
+++ b/tinygrad/renderer/cstyle.py
@@ -9,7 +9,9 @@ from tinygrad.renderer import Renderer
 from tinygrad.codegen.late.devectorizer import no_vectorized_alu
 
 def render_index(ctx,buf,idx):
-  if buf.addrspace == AddrSpace.REG:
+  base = buf
+  while base.op is Ops.AFTER: base = base.src[0]
+  if base.addrspace == AddrSpace.REG or base.op not in {Ops.PARAM, Ops.DEFINE_LOCAL}:
     assert idx.op is Ops.CONST
     return f"{ctx[buf]}[{idx.arg}]"
   else: