Merge branch 'master' into modern_scan

test/unit/test_modern_scan.py

@@ -0,0 +1,36 @@
+import unittest
+from tinygrad import Tensor
+from tinygrad.uop.ops import UOp, AddrSpace
+
+class TestModernScan(unittest.TestCase):
+  def test_copy_local(self):
+    N = 256
+    state = Tensor.empty(N)
+    tmp = UOp.placeholder((N,), state.dtype, slot=-1, addrspace=AddrSpace.LOCAL)
+    tmp = tmp.after(tmp.store(state.uop))
+    state.assign(tmp)
+    state.realize()
+
+  """
+  def test_scan_gemv(self):
+    N = 256
+    gemvs = Tensor.empty(3, N, N)
+    state = Tensor.empty(N)
+    Tensor.realize(gemvs, state)
+
+    #tmp = UOp.placeholder((N,), state.dtype, slot=-1, addrspace=AddrSpace.REG)
+    tmp = Tensor.empty(N, dtype=state.dtype).uop
+    tmp = tmp.after(tmp.store(state.uop))
+    #rng = UOp.range(3, -1)
+    #tmp = tmp.after(tmp.store(state.uop, rng))
+    #tmp = tmp.after(tmp.store(tmp @ gemvs.uop[rng]).end(rng))
+    state.assign(tmp)
+
+    state.realize()
+  """
+
+if __name__ == '__main__':
+  unittest.main()
+
+
+

tinygrad/codegen/__init__.py

@@ -175,9 +175,13 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # EXPAND on scalar -> STACK
   (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
    lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
+  # TODO: make this all generic
   # INDEX on INDEX is INDEX
   (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
-   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
+   lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None),
+  # INDEX on shaped INDEX
+  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx1_arg"))),), allow_any_len=True, name="idx2"),
+   lambda buf,idx1_arg,idx2: buf.index(idx1_arg.index(*idx2.src[1:])) if len(idx1_arg.shape) == len(idx2.src[1:]) else None),
 ])
 
 def fix_group_for_reduce(x:UOp):

tinygrad/schedule/rangeify.py

@@ -434,6 +434,7 @@ class LocalAddBufferContext:
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
+  if buf.addrspace != AddrSpace.GLOBAL: return None
   param = UOp(Ops.PARAM, src=(UOp.const(dtypes.int, prod(buf.max_shape)),),
               arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
@@ -522,7 +523,7 @@ def split_store(x:UOp) -> UOp|None:
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 
   kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
-  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
+  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND and x.device is not None]):
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
   return kernel