From 6131c0aad3f3d9477e5b237c0667da4279b1f59b Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Thu, 14 Aug 2025 13:17:49 -0700
Subject: [PATCH] work

---
 tinygrad/schedule/kernelize.py | 10 +++++-----
 tinygrad/schedule/rangeify.py  | 12 +++++-------
 tinygrad/uop/ops.py            |  3 ++-
 3 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/tinygrad/schedule/kernelize.py b/tinygrad/schedule/kernelize.py
index 6d6e815aa6..295a98deec 100644
--- a/tinygrad/schedule/kernelize.py
+++ b/tinygrad/schedule/kernelize.py
@@ -331,12 +331,11 @@ new_fixups = mops_merge+PatternMatcher([
 
 def split_load(ctx:list[UOp], s:UOp):
   if len(s.src) == 1 or s.src[0].src[0].op is not Ops.DEFINE_GLOBAL: return None
-  ctx.extend(s.src[1:])
   return s.replace(src=s.src[0:1])
 
 def debuf(ctx:list[int], b:UOp):
-  ret = UOp(Ops.DEFINE_GLOBAL, b.dtype.ptr(b.arg), arg=ctx[0])
-  ctx[0] += 1
+  ret = UOp(Ops.DEFINE_GLOBAL, b.dtype.ptr(b.arg), arg=len(ctx))
+  ctx.append(b)
   return ret
 
 to_define_global = PatternMatcher([
@@ -349,10 +348,10 @@ def split_store(x:UOp):
   name = "k_"+'_'.join([str(s) for s in shape])
 
   b = x.src[0].src[0]
-  ctx = [0]
+  ctx = []
   ret = graph_rewrite(x, to_define_global, ctx=ctx, name="* kernel split")
   ret = ret.sink(arg=KernelInfo(name=name))
-  kernel = UOp(Ops.KERNEL, src=(b,), arg=Kernel(ret, ()))
+  kernel = UOp(Ops.KERNEL, src=(b,)+tuple(ctx), arg=Kernel(ret, ()))
   return b.assign(kernel)
 
 split_kernels = PatternMatcher([
@@ -383,6 +382,7 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
     tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, ctx=AddBufferContext(), bottom_up=True, input_map=tensor_map, name="* buffer")
     tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, bottom_up=True, input_map=tensor_map, name="* split kernels")
 
+    if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph")
     #rsink = graph_rewrite(rsink, sym, name="* symbolic")
 
     #from tinygrad.codegen.devectorizer import pm_reduce, ReduceContext
diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py
index 659530f587..1862e9039a 100644
--- a/tinygrad/schedule/rangeify.py
+++ b/tinygrad/schedule/rangeify.py
@@ -1,6 +1,6 @@
 from typing import Any
 from dataclasses import dataclass, field
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, PtrDType
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady
 from tinygrad.helpers import argsort, prod, all_same
 
@@ -228,13 +228,11 @@ def add_store(ctx:AddBufferContext, x:UOp):
     buf = UOp(Ops.DEFINE_LOCAL, dtype=x.dtype.ptr(size=prod(shape), addrspace=AddrSpace.LOCAL), arg=ctx.dg)
   ctx.map[buf] = (buf.op, ctx.dg)
   ctx.dg += 1
-  return buf.reshape(shape).index(*rngs).store(x.src[0], *rngs)
+  return buf.reshape(shape).index(*rngs, dtype=x.dtype.ptr(size=prod(shape))).store(x.src[0], *rngs)
 
 def add_load(ctx:AddBufferContext, x:UOp, b:UOp, idx:UOp):
-  if b not in ctx.map:
-    ctx.map[b] = (Ops.DEFINE_GLOBAL, ctx.dg)
-    ctx.dg += 1
-  return UOp(ctx.map[b][0], dtype=x.dtype.ptr(size=b.arg), arg=ctx.map[b][1]).index(idx).load()
+  if isinstance(x.dtype, PtrDType): return None
+  return x.replace(dtype=x.dtype.ptr(b.size)).load()
 
 def add_load_on_store(ctx:AddBufferContext, x:UOp, st:UOp):
   rngs = x.src[1:]
@@ -244,7 +242,7 @@ def add_load_on_store(ctx:AddBufferContext, x:UOp, st:UOp):
 pm_add_buffers = pm_mops+PatternMatcher([
   (UPat(Ops.CONTIGUOUS, name="x"), add_store),
   (UPat(Ops.ENDRANGE, name="x"), lambda x: x.src[0]),
-  #(UPat(Ops.INDEX, src=(UPat(Ops.BUFFER, name="b"), UPat(name="idx")), name="x"), add_load),
+  (UPat(Ops.INDEX, src=(UPat(Ops.BUFFER, name="b"), UPat(name="idx")), name="x"), add_load),
   (UPat(Ops.INDEX, src=(UPat(Ops.STORE, name="st"),), allow_any_len=True, name="x"), add_load_on_store),
   (UPat(Ops.BIND, name="b"), lambda b: b.src[0]),
   # HACK: ignore copy
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 1224827565..f4e07e49f4 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -234,7 +234,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return ret
   def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
-  def index(self, *srcs:UOp|None): return UOp(Ops.INDEX, self.dtype, (self,)+tuple([x for x in srcs if x is not None]))
+  def index(self, *srcs:UOp|None, **kwargs):
+    return UOp(Ops.INDEX, kwargs.pop("dtype", self.dtype), (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def __getitem__(self, idx): return self.index(idx)
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source