scan op work

* write scan in uops

* ops range

* no need for variable

* meh, later

* shorter

test/test_outerworld.py

@@ -50,13 +50,67 @@ class TestOuterRange(unittest.TestCase):
     # 3 matmuls with outer world range
     i = UOp.range(3, -100, AxisType.OUTER)
     vec_i = Tensor(vec.uop.after(i))
-    vi = UOp.variable("i", i.vmin, i.vmax).bind(i)
-    out = Tensor(vec.uop.after(vec_i.uop.store((vec_i.contiguous() @ mats[vi]).uop).end(i)))
+    comp = vec_i.contiguous() @ mats[i]
+    store = vec_i.uop.store(comp.uop).end(i)
+    out = Tensor(vec.uop.after(store))
     out.realize()
 
     # TODO: testing allclose
     assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
 
+class TestOuterScan(unittest.TestCase):
+  def _test_scan(self):
+    vec = Tensor.randn(1, 10).realize()
+    mats = Tensor.randn(3, 10, 10).realize()
+
+    # 3 matmuls in "scan"
+    vec1 = vec @ mats[0]
+    vec2 = vec1 @ mats[1]
+    vec3 = vec2 @ mats[2]
+    ref = Tensor.stack(vec1, vec2, vec3)
+    ref.realize()
+    return vec, mats, ref
+
+  def test_uop_fold_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with FOLD
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out.uop))
+    comp = phi @ mats[i]
+    store = out.uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref[2], out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+  def test_uop_scan_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with SCAN
+    i = UOp.range(3, -100, AxisType.OUTER)
+    out = Tensor.empty(3, 1, 10)
+    phi = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop))
+    comp = phi @ mats[i]
+    store = out[i].uop.store(comp.uop).end(i)
+    out = Tensor(out.uop.after(store))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+
+  def test_fold_matmul(self):
+    vec, mats, ref = self._test_scan()
+
+    # 3 matmuls with SCAN
+    i = UOp.range(3, -100, AxisType.OUTER)
+    phi = vec._apply_uop(UOp.phi)
+    comp = phi @ mats[i]
+    scan = comp._apply_uop(UOp.fold, phi, extra_args=(i,))
+    scan.realize()
+
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
     t = Tensor.arange(100).reshape(10,10).realize()

tinygrad/engine/schedule.py

@@ -23,10 +23,12 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
   in_degree: dict[UOp, int] = {}
   var_vals: dict[str, int] = {}
   for u in sched_sink.toposort():
-    if u.op is not Ops.AFTER: continue  # anything that's not an ASSIGN doesn't write a kernel, so we can skip
+    if u.op is Ops.RANGE:
+      in_degree.setdefault(u, 0)
+      continue
+    if u.op is not Ops.AFTER or u.src[1].op is Ops.RANGE: continue
     k = u.src[1]
     in_degree.setdefault(k, 0)
-    if k.op is Ops.RANGE: continue
     for s in k.src[0].src if k.op is Ops.END else k.src:
       if s.op is Ops.AFTER:
         children[s.src[1]].append(k)

tinygrad/runtime/ops_amd.py

@@ -12,7 +12,7 @@ from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
 from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
+from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler, AMDLLVMCompiler
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_ip_offsets, import_pmc
@@ -909,7 +909,8 @@ class AMDDevice(HCQCompiled):
     self.sdma_queue = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20))
 
     compilers:list[CompilerPairT] = [(functools.partial(AMDRenderer, self.arch), functools.partial(HIPCompiler, self.arch)),
-                                     (functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch))]
+                                     (functools.partial(AMDLLVMRenderer, self.arch), functools.partial(AMDLLVMCompiler, self.arch)),
+                                     (functools.partial(AMDRenderer, self.arch), functools.partial(HIPCCCompiler, self.arch))]
 
     super().__init__(device, AMDAllocator(self), compilers, functools.partial(AMDProgram, self), AMDSignal,
                      functools.partial(AMDComputeAQLQueue if self.is_aql else AMDComputeQueue, self),

tinygrad/runtime/support/compiler_amd.py

@@ -1,4 +1,4 @@
-import ctypes
+import ctypes, hashlib, tempfile, subprocess, pathlib
 from tinygrad.helpers import system
 from tinygrad.runtime.autogen import comgr
 try:
@@ -90,6 +90,24 @@ class HIPCompiler(Compiler):
     except RuntimeError as e: raise CompileError(e) from e
   def disassemble(self, lib:bytes): amdgpu_disassemble(lib)
 
+class HIPCCCompiler(Compiler):
+  def __init__(self, arch:str, extra_options:list[str]=[]):
+    self.arch, self.extra_options = arch, extra_options
+    super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+  def compile(self, src:str) -> bytes:
+    with tempfile.NamedTemporaryFile(suffix=".cpp") as srcf, tempfile.NamedTemporaryFile(suffix=".bc") as bcf:
+      with tempfile.NamedTemporaryFile(suffix=".hsaco") as libf:
+        srcf.write(src.encode())
+        srcf.flush()
+
+        subprocess.run(["hipcc", "-c", "-emit-llvm", "--cuda-device-only", "-O3", "-mcumode",
+                        f"--offload-arch={self.arch}", "-I/opt/rocm/include/hip", "-o", bcf.name, srcf.name] + self.extra_options, check=True)
+        subprocess.run(["hipcc", "-target", "amdgcn-amd-amdhsa", f"-mcpu={self.arch}",
+                        "-O3", "-mllvm", "-amdgpu-internalize-symbols", "-c", "-o", libf.name, bcf.name], check=True)
+
+        return pathlib.Path(libf.name).read_bytes()
+  def disassemble(self, lib:bytes): amdgpu_disassemble(lib)
+
 class AMDLLVMCompiler(LLVMCompiler):
   jit = False
   target_arch = "AMDGPU"

tinygrad/schedule/rangeify.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 import itertools
 from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, ssimplify, KernelInfo
-from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags
+from tinygrad.uop.ops import track_rewrites, graph_rewrite, identity_element, sint, AxisType, BottomUpGate, Kernel, _remove_all_tags, range_str
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY
 from tinygrad.helpers import PCONTIG, partition, get_single_element, unwrap, disable_gc
@@ -397,6 +397,9 @@ def handle_after(ctx:LocalAddBufferContext, after:UOp):
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.tag != (): return None
+  if r.arg[-1] == AxisType.OUTER:
+    # for outer range, we replace with a bound variable
+    return UOp.variable("range_"+range_str(r), r.vmin, r.vmax).bind(r.replace(tag=None))
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=None)
   ctx.range += 1
   return ret
@@ -571,12 +574,12 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))
   if assign_rep: tsink = graph_rewrite(tsink, _substitute, ctx=assign_rep, bottom_up=True, name="fix_assign")
 
-  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-
   # TODO: we can probably get this earlier
   sink_tags = [s.tag for s in tsink.src]
   tsink = graph_rewrite(tsink, _remove_all_tags, name="remove all tags")
 
+  if getenv("VIZ"): graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+
   becomes_map: dict[UOp, UOp] = {}
   for tag, s in zip(sink_tags, tsink.src):
     assert tag is not None

tinygrad/uop/__init__.py

@@ -92,6 +92,8 @@ class Ops(FastEnum):
   # reduce
   REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto()
 
+  PHI = auto(); SCAN = auto(); FOLD = auto()
+
   # errors/placeholders
   REWRITE_ERROR = auto(); SENTINEL = auto()
 

tinygrad/uop/ops.py

@@ -25,7 +25,7 @@ axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL:
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
                AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5, AxisType.OUTER: -2}
 
-range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1}
+range_start = {Ops.BUFFERIZE: 1, Ops.REDUCE: 1, Ops.STORE: 2, Ops.WMMA: 3, Ops.END: 1, Ops.FOLD: 2}
 
 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> ConstType: return dtypes.as_const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dtypes.min(dt)}[op], dt)
@@ -219,9 +219,14 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return (self.ptrdtype.size,)
 
       # passthrough ops
-      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END:
+      case Ops.REDUCE | Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.END | Ops.PHI | Ops.FOLD:
         return self.src[0]._shape
 
+      # scan adds dims to the front
+      case Ops.SCAN:
+        if self.src[0]._shape is None: return None
+        return tuple(x.vmax+1 for x in self.src[2:]) + self.src[0]._shape
+
       # ops with custom handling
       case Ops.KERNEL: return self.arg.ast._shape
 
@@ -443,6 +448,10 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid)
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
 
+  def fold(self, *src:UOp, **kwargs): return UOp(Ops.FOLD, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
+  def scan(self, *src:UOp, **kwargs): return UOp(Ops.SCAN, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
+  def phi(self, *src:UOp, **kwargs): return UOp(Ops.PHI, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
+
   def is_contiguous(self):
     # TODO: this is is_realized
     if self.op is Ops.RESHAPE: return self.src[0].is_contiguous()
@@ -1265,15 +1274,12 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat(Ops.SPECIAL, src=(UPat.var("var").cast(dtypes.index),), name="u"), lambda u,var: u.replace(dtype=dtypes.int, src=(var,)).cast(dtypes.index)),
   (UPat(Ops.DEFINE_VAR, dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=dtypes.int).cast(dtypes.index)),
   (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.index), UPat.cvar("val").cast(dtypes.index))), lambda var,val: var.bind(val).cast(dtypes.index)),
-  (UPat(Ops.CAST, src=(UPat(name="x").cast(dtypes.index),), name="c"), lambda x,c: x.cast(c.dtype)),
   # lower Invalid
   (UPat.var("buf").index(UPat.var("cond").where(UPat.var("idx"), UPat(Ops.CONST, arg=Invalid))), lambda buf,idx,cond: buf.index(idx, cond, ptr=True)),
   # remove hanging casts
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx, ptr=True)),
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("valid"))),
    lambda buf,idx,valid: buf.index(idx, valid, ptr=True)),
-  (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
-    lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
   (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"),
    lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.index else s for s in n.src))),
 ])

tinygrad/uop/spec.py

@@ -42,7 +42,7 @@ shared_spec = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:]) or None),
 
   # RANGE/SPECIAL define loops, END closes them
-  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE)), dtype=dtypes.void), lambda: True),
+  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE))), lambda: True),
 ])
 
 # ***** UOp spec in the Tensor graph *****
@@ -171,7 +171,7 @@ kernel_spec = PatternMatcher([
   (UPat(Ops.UNROLL, name="x"), lambda x: x.src[0].dtype.count == prod(y[1] for y in x.arg)),
 
   # END can end multiple axes here
-  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True, dtype=dtypes.void), lambda: True),
+  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),
 
   # bufferize can be on anything
   (UPat(Ops.BUFFERIZE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: True),