_apply_movement_op cache

* fix: revive torch backend

* as_strided view vs copy

* Revert "as_strided view vs copy"

This reverts commit 82a61223f2.

* add extra tests (move inplace, add fusion tests)

* better fusion with inplace_op

* no optimizer hooks (break mnist training fusion)

* split off fusion tests in separate file, assert on resnet fusion

fix: remove comments

* cleanup, reduce diff

* reduce diff

* better fusion and identity checks

---------

Co-authored-by: George Hotz <[email protected]>

.github/workflows/test.yml

@@ -86,65 +86,67 @@ jobs:
         clang -O2 recognize.c -lm -o recognize
         cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock
 
-  # TODO: fix the torch backend and reenable
-  # torchbackend:
-  #   name: Torch Backend Tests
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       pydeps: "pillow torchvision expecttest"
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Lint with ruff
-  #     run: |
-  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
-  #       python3 -m ruff check extra/torch_backend/backend.py
-  #   - name: Test one op
-  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-  #   - name: Test ResNet-18
-  #     run: DEBUG=2 python3 extra/torch_backend/example.py
-  #   - name: My (custom) tests
-  #     run: python3 extra/torch_backend/test.py
-  #   - name: Test one op in torch tests
-  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-  #   - name: Test Ops with TINY_BACKEND
-  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-  #   - name: Test in-place operations on views
-  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-  #   - name: Test multi-gpu
-  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  torchbackend:
+    name: Torch Backend Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
+      run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: My (custom) tests
+      run: python3 extra/torch_backend/test.py
+    - name: Test one op in torch tests
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+    - name: Test kernel fusion
+      run: python3 extra/torch_backend/test_kernel_fusion.py
 
-  # torchbackendmore:
-  #   name: Torch Backend Tests More
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
-  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-  #   - name: Test some torch tests (expect failure)
-  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+
+  torchbackendmore:
+    name: Torch Backend Tests More
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+    - name: Test some torch tests (expect failure)
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
   bepython:
     name: Python Backend
@@ -306,6 +308,7 @@ jobs:
       with:
         key: spec-unit
         deps: testing_unit
+        python-version: '3.14'
     - name: Test SPEC=2
       run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 

extra/nv_gpu_driver/nv_ioctl.py

@@ -64,14 +64,17 @@ nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(
           x.startswith("NV") and x[6:].startswith("_CTRL_") and isinstance(getattr(nv_gpu, x), int)}
 
 def get_classes():
-  hdrpy = (pathlib.Path(__file__).parent.parent.parent / "tinygrad/runtime/autogen/nv_570.py").read_text()
-  clss = re.search(r'NV01_ROOT.*?NV_SEMAPHORE_SURFACE = \(0x000000da\) # macro', hdrpy, re.DOTALL).group()
-  pattern = r'([0-9a-zA-Z_]*) = +\((0x[0-9a-fA-F]+)\)'
-  matches = re.findall(pattern, clss, re.MULTILINE)
-  return {int(num, base=16):name for name, num in matches}
+  res = {}
+  known_classes = {"NV01_DEVICE_0", "NV01_ROOT", "NV1_MEMORY_SYSTEM", "NV01_MEMORY_VIRTUAL", "NV1_MEMORY_USER", "NV50_MEMORY_VIRTUAL", "NV_FERMI_VASPACE_A",
+                   "NV20_SUBDEVICE_0"}
+  for nm,val in nv_gpu.__dict__.items():
+    if not isinstance(val, int): continue
+    if 0x3000 < val < 0xffff: res[val] = nm
+    if nm in known_classes: res[val] = nm
+  return res
 nvclasses = get_classes()
 nvuvms = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("UVM_") and nv_gpu.__dict__.get(x+"_PARAMS")}
-nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
+nvqcmds = {int(getattr(nv_gpu, x)):x for x in dir(nv_gpu) if x[:7] in {"NVC9B0_", "NVC6C0_", "NVC56F_", "NVC6B5_"} and isinstance(getattr(nv_gpu, x), int)}
 
 global_ioctl_id = 0
 gpus_user_modes = []

extra/sqtt/roc.py

@@ -1,4 +1,4 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools, threading
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
@@ -94,14 +94,14 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
   ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
 
   @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
+  def copy_cb(buf, buf_size, _):
     if (prof_info:=ROCParseCtx.next_sqtt()) is None: return 0
     buf[0] = ctypes.cast(prof_info, ctypes.POINTER(ctypes.c_ubyte))
     buf_size[0] = len(prof_info)
     return len(prof_info)
 
   @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
+  def trace_cb(record_type, events_ptr, n, _):
     match record_type:
       case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
         for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
@@ -112,7 +112,7 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, _):
     instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
@@ -126,9 +126,11 @@ def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
 
     return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
 
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  def worker():
+    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+    except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install") from e
+  (t:=threading.Thread(target=worker, daemon=True)).start()
+  t.join()
   return ROCParseCtx
 
 if __name__ == "__main__":

extra/thunder/tiny/tk/group.py

@@ -56,7 +56,7 @@ class Group:
     self.ker.push_store(dst_store, dst)
     return dst.after(dst_store).reshape(dst.shape)
 
-  def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
+  def mma_AB(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
     c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
     assert self.warps == 1
 
@@ -77,9 +77,9 @@ class Group:
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape) if after else c_store
+    return c.after(c_store).reshape(c.shape)
 
-  def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT, after=True):
+  def mma_ABt(self, c:UOp|RT, a:UOp|RT, b:UOp|RT):
     c, a, b = cast(UOp, c), cast(UOp, a), cast(UOp, b)
     assert self.warps == 1
 
@@ -100,7 +100,7 @@ class Group:
           c_store = UOp.group(*c_i).end(height, width, inner)
 
     self.ker.push_store(c_store, c)
-    return c.after(c_store).reshape(c.shape) if after else c_store
+    return c.after(c_store).reshape(c.shape)
 
   map_rid = 400
   def map(self, a:ALL_TILES, op:Callable[[UOp], UOp]|Callable[[UOp, tuple], UOp]):

extra/thunder/tiny/tk/kernel.py

@@ -80,7 +80,11 @@ class Kernel(AbstractContextManager):
     rngs = []
     while self.range_stack: rngs.append(self.range_stack.pop(0)._rng)
 
-    return self.store_stack.pop()[0]._uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
+    last_store = self.store_stack.pop()[0]
+    if hasattr(last_store, '_uop'): uop = last_store._uop
+    else: uop = last_store
+
+    return uop.end(*rngs).sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
   def endrange(self):
     last_store = self.store_stack.pop()

extra/thunder/tiny/tk/tiles.py

@@ -11,9 +11,9 @@ def unwrap(x):
   if isinstance(x, dict): return {k: unwrap(v) for k,v in x.items()}
   return x
 
-def wrap(x, ker, cls):
-  if isinstance(x, UOp): return cls(x, ker)
-  if isinstance(x, (list, tuple)): return type(x)(wrap(y, ker, cls) for y in x)
+def wrap(x, s):
+  if isinstance(x, UOp): return s.ruop(x)
+  if isinstance(x, (list, tuple)): return type(x)(wrap(y, s) for y in x)
   return x
 
 def autowrap(source_cls, blacklist=None):
@@ -31,10 +31,10 @@ def autowrap(source_cls, blacklist=None):
       if callable(val):
         @functools.wraps(val)
         def proxy(*args, **kwargs):
-          return wrap(val(*unwrap(args), **unwrap(kwargs)), self.ker, cls)
+          return wrap(val(*unwrap(args), **unwrap(kwargs)), self)
         return proxy
       if name in UOp.__slots__: return val
-      return wrap(val, self.ker, cls)
+      return wrap(val, self)
     cls.__getattr__ = __getattr__
 
     for name in dir(source_cls):
@@ -46,9 +46,9 @@ def autowrap(source_cls, blacklist=None):
       else:
         original = getattr(source_cls, name)
         if callable(original):
-          def make_proxy(op_name, func):
+          def make_proxy(_, func):
             def proxy(self, *args, **kwargs):
-              return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self.ker, cls)
+              return wrap(func(self._uop, *unwrap(args), **unwrap(kwargs)), self)
             return proxy
           setattr(cls, name, make_proxy(name, original))
 
@@ -69,7 +69,7 @@ class TileMathMixin(MathMixin):
         if isinstance(self, RT) and isinstance(src[0], RV): uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[idx[0], 0, (idx[2]%4)//2])))
         else: uop = self.ker.warp.map(self._uop, lambda x, idx: UOp.alu(x, op, inner_op(src[0]._uop[*idx])))
     else: raise NotImplementedError
-    return type(self)(uop, self.ker)
+    return self.ruop(uop)
   def const_like(self, b): return b
 
   # override ops that do compute on the src uop
@@ -83,6 +83,9 @@ class GL:
   def __init__(self, uop, ker):
     self._uop, self.ker = uop, ker
 
+  def ruop(self, uop):
+    return GL(uop, self.ker)
+
   @classmethod
   def create(cls, shape, dtype, ker):
     uop = ker.alloc(shape, dtype, AddrSpace.GLOBAL)
@@ -93,6 +96,9 @@ class ST:
   def __init__(self, uop, ker):
     self._uop, self.ker = uop, ker
 
+  def ruop(self, uop):
+    return ST(uop, self.ker)
+
   @classmethod
   def create(cls, shape, dtype, ker):
     uop = ker.alloc(shape, dtype, AddrSpace.LOCAL)
@@ -107,6 +113,9 @@ class RT(TileMathMixin):
   def __init__(self, uop, ker):
     self._uop, self.ker = uop, ker
 
+  def ruop(self, uop):
+    return RT(uop, self.ker)
+
   @classmethod
   def create(cls, shape, dtype, ker):
     assert len(shape) == 2
@@ -121,8 +130,11 @@ class RT(TileMathMixin):
 
 @autowrap(UOp)
 class RV(TileMathMixin):
-  def __init__(self, uop, ker):
-    self._uop, self.ker = uop, ker
+  def __init__(self, uop, layout, ker):
+    self._uop, self.layout, self.ker = uop, layout, ker
+
+  def ruop(self, uop):
+    return RV(uop, self.layout, self.ker)
 
   @classmethod
   def create(cls, length, dtype, layout, ker):
@@ -138,6 +150,6 @@ class RV(TileMathMixin):
       case _: raise NotImplementedError(f"rv layout {layout} not implemented")
 
     uop = ker.alloc((outer_dim, inner_dim, 2), dtype, AddrSpace.REG)
-    return RV(uop, ker)
+    return RV(uop, layout, ker)
 
 ALL_TILES = UOp | GL | ST | RT | RV

extra/torch_backend/backend.py

@@ -4,10 +4,10 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops
-from tinygrad.helpers import getenv, prod
+from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
-import torch, pathlib, math, operator, functools, inspect
+import torch, pathlib, math, operator, functools, weakref
 torch.autograd.grad_mode.set_multithreading_enabled(False)
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
 
@@ -18,7 +18,17 @@ def _to_torch_device(device: str): return torch.device("tiny", int(device.partit
 
 import torch.utils.cpp_extension
 mod = torch.utils.cpp_extension.load(name="custom_device_extension", sources=[str(pathlib.Path(__file__).parent / "wrapped_tensor.cpp")])
-def wrap(x:Tensor) -> torch.Tensor: return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
+def calculate_storage_offset(x: Tensor) -> int:
+  offset = 0
+  for u in x.uop.toposort():
+    if u.op == Ops.SHRINK:
+      u_strides = strides_for_shape(u.src[0].shape)
+      for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+  return offset
+def wrap(x: Tensor) -> torch.Tensor:
+  x._strides = strides_for_shape(x.shape) # always recalculate
+  if (not hasattr(x, '_storage_offset')) or (not x.uop.is_realized): x._storage_offset = calculate_storage_offset(x)
+  return mod.wrap(x, _to_torch_dtype(x.dtype), _to_torch_device(x.device).index)
 def unwrap(x:torch.Tensor) -> Tensor:
   assert isinstance(x, torch.Tensor), f"x isn't {type(x)}"
   return mod.unwrap(x)
@@ -35,17 +45,20 @@ torch.utils.generate_methods_for_privateuse1_backend()
 aten = torch.ops.aten
 
 # track view relationships for in place operations
-def is_view(tensor: Tensor): return hasattr(tensor, "_view_base")
 def canonical_base(view: Tensor): return getattr(view, "_view_base", view)
 def derived_views(base: Tensor): return [t for tref in getattr(base, "_views", set()) if (t:=tref()) is not None]
+def unwrap_args(args, kwargs):
+  return [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args], {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
 def wrap_view_op(fn):
-  def _wrap(*args,**kwargs):
-    args = [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args]
-    kwargs = {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
-    ret = fn(*args,**kwargs)
-    ret._view_base = base = canonical_base(args[0])
-    if not hasattr(base, "_views"): base._views = set()
+  @functools.wraps(fn)
+  def _wrap(*args, **kwargs):
+    args, kwargs = unwrap_args(args, kwargs)
+    ret = fn(*args, **kwargs)
+    base = canonical_base(args[0])
+    ret._view_base = base
+    base._views = getattr(base, "_views", set())
     base._views.add(weakref.ref(ret))
+    ret._view_ops = _get_view_ops(args[0]) + [(fn, args[1:], kwargs)]
     return wrap(ret)
   return _wrap
 
@@ -60,46 +73,79 @@ view_ops = {
   "aten.unsqueeze": Tensor.unsqueeze,
   "aten.detach": Tensor.detach,
   "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
-}
+  "aten.permute": Tensor.permute,
+  "aten.alias": lambda self: self,
+  }
 
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
 
-# in place operations with views
-def realize_with_views(self: Tensor, views: Tensor):
-  if not self.uop.st.contiguous: self.replace(self.contiguous())
-  self.replace(self.clone().realize())
-  for v in views:
-    if v.uop.base.op is Ops.BUFFER_VIEW: continue # skip subbuffer, we just use the real buffer view
-    ret = self
-    st = ShapeTracker(self.uop.st.views + v.uop.st.views) # TODO: is this right?
-    for mo in cached_to_movement_ops(self.shape, st): ret = apply_mop(ret, mo)
-    v.replace(ret)
-def maybe_realize_storage(self: Tensor) -> bool:
-  if realize:=is_view(self): realize_with_views((base:=canonical_base(self)), derived_views(base))
-  return realize
-def inplace_fn(outvars: str|list[str]):
-  if type(outvars) is str: outvars = [outvars]
-  def decorator(fn):
-    sig = inspect.signature(fn)
-    def wrapper(*args, **kwargs):
-      bound = sig.bind(*args, **kwargs)
-      outs = [kwargs.get(v, bound.arguments.get(v)) for v in outvars]
-      outs = [unwrap(o) if isinstance(o, torch.Tensor) else o for o in outs]
-      realize = any(maybe_realize_storage(o) for o in outs)
-      ret = fn(*args, **kwargs)
-      if realize: Tensor.realize(*(o for o in outs))
-      return ret
-    return wrapper
-  return decorator
+def _get_view_ops(view): return getattr(view, "_view_ops", [])
+
+def _apply_view_ops(target, ops):
+  for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
+  return target
+
+# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
+def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
+  if not (req := argfix(*args)): return None
+  new_shape, infer_idx = [], -1
+  for i, s in enumerate(req):
+    if s is None: s = shape[i] if i < len(shape) else None
+    if not isinstance(s, int): return None
+    if s == -1:
+      if infer_idx != -1: return None
+      infer_idx = len(new_shape)
+    new_shape.append(s)
+  total = prod(shape)
+  if infer_idx != -1:
+    known = prod(x for x in new_shape if x != -1)
+    if known == 0:
+      if total != 0: return None
+      new_shape[infer_idx] = 0
+    else: new_shape[infer_idx] = total // known
+  return tuple(new_shape) if prod(new_shape) == total else None
+
+# TODO: can we get rid of this? only for test_flatten_reshape_add
+def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
+  if not (ops := _get_view_ops(view)): return False
+  shapes = [base.shape]
+  for fn, args, _ in ops:
+    if fn is Tensor.reshape:
+      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
+      shapes.append(next_shape)
+  if shapes[-1] != view.shape: return False
+  for s in reversed(shapes[:-1]): val = val.reshape(s)
+  base.assign(val)
+  return True
+
+def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
+  val = value if value.dtype == base.dtype else value.cast(base.dtype)
+  if view.shape == base.shape: return base.assign(val)
+  if _try_simple_reshape_view_write(base, view, val): return
+  idx_base = Tensor.arange(base.numel(), device=base.device, dtype=dtypes.int32).reshape(base.shape)
+  idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
+  flat_base = base.reshape(base.numel()).contiguous()
+  flat_base[idx_view] = val.reshape(-1)
+  base.assign(flat_base.reshape(base.shape))
+
+def _apply_inplace(target: Tensor, value: Tensor) -> None:
+  val = value if value.dtype == target.dtype else value.cast(target.dtype)
+  base = canonical_base(target)
+  views = derived_views(base)
+  if not views: return target.assign(val)
+  view_ops_map = {v: _get_view_ops(v) for v in views}
+  if target is base or target.uop is base.uop: base.assign(val)
+  else: _view_write(base, target, val)
+  for v in views: v.replace(_apply_view_ops(base, view_ops_map[v]))
 
 # *** bad functions on CPU ***
 
 @torch.library.impl("aten::_index_put_impl_", "privateuseone")
-@inplace_fn("self")
 def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
   # TODO: move to tinygrad
   ret = aten._index_put_impl_(self.cpu(), [x.cpu() if isinstance(x, torch.Tensor) else None for x in indices], values.cpu(), accumulate, unsafe).to(self.device)
-  return wrap(unwrap(self).assign(unwrap(ret)))
+  unwrap(self).assign(unwrap(ret))
+  return self
 
 @torch.library.impl("aten::index_put", "privateuseone")
 def index_put(self, indices, values, accumulate=False):
@@ -150,43 +196,23 @@ for i in [
 def index_tensor(x, y):
   return wrap(unwrap(x)[[unwrap(_y.to(x.device)) if _y is not None else slice(None) for _y in y]])
 
-@torch.library.impl("aten::zero_", "privateuseone")
-@inplace_fn("x")
-def zero_(x):
-  if TORCH_DEBUG: print(f"zero_ {x.shape}")
-  tt = unwrap(x)
-  tt.assign(tt.zeros_like())
-
-@torch.library.impl("aten::fill_.Scalar", "privateuseone")
-@inplace_fn("x")
-def fill_scalar(x, y):
-  if TORCH_DEBUG: print(f"fill_.Scalar {x.shape} {y}")
-  tt = unwrap(x)
-  tt.assign(tt.full_like(y))
 
 @torch.library.impl("aten::_local_scalar_dense", "privateuseone")
 def _local_scalar_dense(tensor): return unwrap(tensor).item()
 
-@functools.cache
-def cached_to_movement_ops(shape, st) -> list:
-  mops = to_movement_ops(st)
-  if mops[0] == (MovementOps.RESHAPE, shape): mops = mops[1:]
-  return mops
-
-from tinygrad.shape.shapetracker import ShapeTracker, View
-from extra.to_movement_ops import to_movement_ops, apply_mop, MovementOps
-
 @wrap_view_op
-def _as_strided(tensor:Tensor, size, stride, storage_offset=None):
-  # multiple as_strided do not compound
-  base = canonical_base(tensor)
-  # TODO: this is heavyweight
-  st = ShapeTracker(base.uop.st.views + (View.create(tuple(size), tuple(stride), storage_offset),))
-  ret = base
-  if TORCH_DEBUG >= 1: print("**** as_strided", tensor.shape, size, stride, st)
-  if prod(size) == 1: return ret.flatten()[storage_offset].reshape(size)
-  for mo in cached_to_movement_ops(tuple(base.shape), st): ret = apply_mop(ret, mo)
-  return ret
+def _as_strided(tensor:Tensor, size, stride, storage_offset=0):
+  base = getattr(tensor, "_as_strided_base", canonical_base(tensor)).flatten()
+  if prod(size) == 1: return base[storage_offset].reshape(size)
+  indices = Tensor.zeros(size, dtype=dtypes.int32, device=base.device) + storage_offset
+  for dim, (sz, st) in enumerate(zip(size, stride)):
+    if st != 0:
+      dim_range = Tensor.arange(sz, device=base.device, dtype=dtypes.int32) * st
+      shape_for_broadcast = [1] * dim + [sz] + [1] * (len(size) - dim - 1)
+      indices = indices + dim_range.reshape(shape_for_broadcast)
+  result = base[indices.flatten()].reshape(size)
+  result._as_strided_base = base
+  return result
 
 @torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
@@ -245,15 +271,14 @@ def convolution_overrideable(input, weight, bias, stride, padding, dilation, tra
   if TORCH_DEBUG >= 1:
     print(f"convolution {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
   input, weight, bias = unwrap(input), unwrap(weight), unwrap(bias) if bias is not None else None
-  # TODO: fix test_biased_conv2d fails without realize()
-  if not transposed: return wrap(input.conv2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding).realize())
-  return wrap(input.conv_transpose2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding).realize())
+  if not transposed: return wrap(input.conv2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding))
+  return wrap(input.conv_transpose2d(weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding, output_padding=output_padding))
 
 @torch.library.impl("aten::convolution_backward_overrideable", "privateuseone")
 def convolution_backward_overrideable(grad_out, input, weight, stride, padding, dilation, transposed, output_padding, groups, output_mask):
   if TORCH_DEBUG >= 1:
     print(f"convolution_backward {input.shape=} {weight.shape=} {stride=} {padding=} {dilation=} {transposed=} {output_padding=} {groups=}")
-  grad_out, input, weight, bias = unwrap(grad_out), unwrap(input), unwrap(weight), Tensor.zeros(weight.shape[0], device=_from_torch_device(weight.device))
+  grad_out, input, weight, bias = unwrap(grad_out).detach(), unwrap(input).detach(), unwrap(weight).detach(), Tensor.zeros(weight.shape[0], device=_from_torch_device(weight.device))
   if not transposed: out = Tensor.conv2d(input, weight, bias, groups=groups, stride=stride, dilation=dilation, padding=padding)
   else:
     bias = Tensor.zeros(weight.shape[1] * groups)
@@ -315,55 +340,57 @@ for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
 
 @torch.library.impl("aten::scatter_add.out", "privateuseone")
-@inplace_fn("out")
 def scatter_add(self, dim, index, src, out):
-  self, index, src, out = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
-  if self.shape == (): return wrap(out.assign(src))
-  return wrap(out.assign(Tensor.scatter_reduce(self, dim, index, src, reduce='sum')))
+  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
+  if self.shape == (): _apply_inplace(out_unwrapped, src)
+  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
+  return out
 
-@torch.library.impl("aten::_copy_from", "privateuseone")
-def _copy_from(src: torch.Tensor, dest, non_blocking=False):
-  realize = dest.is_tiny and maybe_realize_storage(unwrap(dest))
-  cast_dtype = _from_torch_dtype(dest.dtype)
+def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
-    to_device = _from_torch_device(dest.device)
-    src,dest = unwrap(src),unwrap(dest)
-    # TODO we need to properly match dest shape and strides, not blindly assign
-    if dest.uop.st.contiguous or dest.uop.is_realized: src = src.contiguous() # this only solves some cases
-    dest.assign(src.cast(cast_dtype).to(to_device))
-    if realize: Tensor.realize(dest)
+    src_t, dest_t = unwrap(src), unwrap(dest)
+    if dest_t.uop.is_contiguous() or dest_t.uop.is_realized: src_t = src_t.contiguous()
+    _apply_inplace(dest_t, src_t.cast(cast_dtype).to(to_device))
   elif src.is_tiny and dest.is_cpu:
-    # TODO: is there a better way?
     dest.resize_(src.numel()).resize_(src.shape)
     dest.copy_(torch.from_numpy(unwrap(src).cast(cast_dtype).numpy()))
   elif src.is_cpu and dest.is_tiny:
-    to_device = _from_torch_device(dest.device)
-    # TODO we need to properly match dest shape and strides, not blindly assign
     unwrap(dest).assign(Tensor(src.numpy()).cast(cast_dtype).to(to_device))
-    if realize: Tensor.realize(unwrap(dest))
   else:
     raise NotImplementedError(f"can't copy from {src.device} -> {dest.device}")
 
+@torch.library.impl("aten::_copy_from", "privateuseone")
+def _copy_from(src: torch.Tensor, dest, non_blocking=False):
+  cast_dtype = _from_torch_dtype(dest.dtype)
+  to_device = _from_torch_device(dest.device)
+  _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking)
+  return dest
+
+@torch.library.impl("aten::copy_", "privateuseone")
+def copy_(self, src, non_blocking=False):
+  cast_dtype = _from_torch_dtype(self.dtype)
+  to_device = _from_torch_device(self.device)
+  _copy_between_devices(src, self, cast_dtype, to_device, non_blocking)
+  return self
+
 @torch.library.impl("aten::cat.out", "privateuseone")
-@inplace_fn("out")
 def cat_out(tensors, dim=0, out=None):
-  unwrap(out).assign(Tensor.cat(*[unwrap(x) for x in tensors], dim=dim))
+  _apply_inplace(unwrap(out), Tensor.cat(*[unwrap(x) for x in tensors], dim=dim))
+  return out
 
 @torch.library.impl("aten::topk.values", "privateuseone")
-@inplace_fn(["values", "indices"])
 def topk_values(input, k, dim=None, largest=True, sorted=True, values=None, indices=None):
   out_values, out_indices = unwrap(input).topk(k, dim if dim is not None else -1, largest, sorted)
-  unwrap(values).assign(out_values)
-  unwrap(indices).assign(out_indices.cast(dtypes.int64))
-  return wrap(out_values), wrap(out_indices)
+  _apply_inplace(unwrap(values), out_values)
+  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
+  return values, indices
 
 @torch.library.impl("aten::sort.values_stable", "privateuseone")
-@inplace_fn(["values", "indices"])
 def sort_values(input, dim=-1, descending=False, stable=True, values=None, indices=None):
   out_values, out_indices = unwrap(input).sort(dim, descending)
-  unwrap(values).assign(out_values)
-  unwrap(indices).assign(out_indices.cast(dtypes.int64))
-  return wrap(out_values), wrap(out_indices)
+  _apply_inplace(unwrap(values), out_values)
+  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
+  return values, indices
 
 @torch.library.impl("aten::_linalg_svd", "privateuseone")
 def _linalg_svd(self, full_matrices=False):
@@ -373,7 +400,6 @@ def _linalg_svd(self, full_matrices=False):
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
-  aten.native_batch_norm, aten.native_batch_norm_backward,
   aten.native_layer_norm_backward,
   aten.linalg_cross,
   aten.addmm,
@@ -510,7 +536,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
 
 # we add the "out" here
 def wrap_out(f):
-  @inplace_fn("out")
   def _wrap_out(*args, **kwargs):
     out = kwargs.pop('out')
     assigned = f(*args, **kwargs)
@@ -518,22 +543,33 @@ def wrap_out(f):
     assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
     assert out.device == assigned.device, f"device mismatch: {assigned.device} -> {out.device}"
     assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    if out.uop.is_realized: assigned = assigned.contiguous() # TODO: how does this map to torch's semantics
     return out.assign(assigned)
   return _wrap_out
 
+def _inplace_op(t, new_value):
+  if not hasattr(t, "_view_base") and not getattr(canonical_base(t), "_views", set()): t.replace(new_value)
+  else: _apply_inplace(t, new_value)
+  return t
+
 tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
   "aten.floor_divide": lambda x,y: x//y,
-  "aten.floor_divide_.Tensor": inplace_fn("x")(lambda x,y: x.assign(x//y)),
+  "aten.floor_divide_.Tensor": lambda x,y: x//y,
   # TODO: use tinygrad methods, but they require x to be unsigned
   "aten.__lshift__.Scalar": lambda x,y: x*(2**y),
-  "aten.__ilshift__.Scalar": inplace_fn("x")(lambda x,y: x.assign(x*(2**y))),
+  "aten.__ilshift__.Scalar": lambda x,y: x*(2**y),
   "aten.__rshift__.Scalar": lambda x,y: x//(2**y),
-  "aten.__irshift__.Scalar": inplace_fn("x")(lambda x,y: x.assign(x//(2**y))),
+  "aten.__irshift__.Scalar": lambda x,y: x//(2**y),
+  # inplace ops using replace for fusion
+  "aten.zero_": lambda x: x.zeros_like(),
+  "aten.fill_.Scalar": lambda x, y: x.full_like(y),
+  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.mul_.Tensor": lambda self, other: self * other,
+  "aten.mul_.Scalar": lambda self, other: self * other,
   # relu doesn't have an out form?
   "aten.relu": Tensor.relu,
-  "aten.relu_": inplace_fn("x")(lambda x: x.assign(x.relu())),
+  "aten.relu_": lambda x: x.relu(),
   "aten.mean": Tensor.mean,
   "aten.mean.dim": Tensor.mean,
   "aten.min": Tensor.min,
@@ -554,19 +590,17 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
   "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
   "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": inplace_fn("self")(lambda self:
-    self.assign(Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype))),
-  "aten.random_.from": inplace_fn("self")(lambda self, from_, to:
-    self.assign(Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype))),
-  "aten.uniform_": inplace_fn("self")(lambda self, low=0, high=1: self.assign(Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype))),
-  "aten.normal_": inplace_fn("self")(lambda self, mean=0, std=1: self.assign(Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype))),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
+  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
+  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
+  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
   # these don't work in out form, they have size 0
   "aten.abs": Tensor.abs,
   "aten.logical_not": Tensor.logical_not,
-  "aten.logical_or_": inplace_fn("x")(lambda x, y: x.assign(x | y)),
+  "aten.logical_or_": lambda x, y: x | y,
   "aten.multinomial": Tensor.multinomial,
-  "aten.masked_fill_.Scalar": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
-  "aten.masked_fill_.Tensor": inplace_fn("self")(lambda self, mask, value: self.assign(self.masked_fill(mask, value))),
+  "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
+  "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
   "aten.masked_fill.Scalar": Tensor.masked_fill,
   "aten.masked_fill.Tensor": Tensor.masked_fill,
   "aten.masked_select": Tensor.masked_select,
@@ -580,7 +614,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.asinh": Tensor.asinh,
   "aten.mul": Tensor.mul,
   "aten.atanh": Tensor.atanh,
-  "aten.fill_.Tensor": Tensor.full, # TODO: looks wrong
+  "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
   "aten.squeeze_.dim": lambda self, dim: self.replace(self.squeeze(dim), allow_shape_mismatch=True), # TODO: inplace view op, here?
@@ -601,20 +635,51 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.unfold": Tensor.unfold,
 }}
 
+# operations that need inplace treatment (use _inplace_op instead of wrap_fxn) AKA return original tensor
+inplace_ops = {
+  "aten.zero_",
+  "aten.fill_.Scalar",
+  "aten.fill_.Tensor",
+  "aten.add_.Tensor",
+  "aten.add_.Scalar",
+  "aten.mul_.Tensor",
+  "aten.mul_.Scalar",
+  "aten.floor_divide_.Tensor",
+  "aten.__ilshift__.Scalar",
+  "aten.__irshift__.Scalar",
+  "aten.relu_",
+  "aten.random_",
+  "aten.random_.from",
+  "aten.uniform_",
+  "aten.normal_",
+  "aten.logical_or_",
+  "aten.masked_fill_.Scalar",
+  "aten.masked_fill_.Tensor",
+}
+
 def wrap_fxn(k,f):
   def nf(*args, **kwargs):
     if TORCH_DEBUG:
       print(k, len(args), [x.shape if isinstance(x, torch.Tensor) else x for x in args],
                           {k:v.shape if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()})
-    args = [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args]
-    kwargs = {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
+    args, kwargs = unwrap_args(args, kwargs)
     out = f(*args, **kwargs)
     if isinstance(out, Tensor): return wrap(out)
     elif isinstance(out, tuple): return tuple(wrap(x) for x in out)
     else: raise RuntimeError(f"unknown output type {type(out)}")
   return nf
 
-for k,v in tiny_backend.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_fxn(k,v))
+def wrap_inplace(k,f):
+  def nf(*args, **kwargs):
+    orig = args[0]
+    args, kwargs = unwrap_args(args, kwargs)
+    _inplace_op(args[0], f(*args, **kwargs))
+    return orig
+  return nf
+
+for k,v in tiny_backend.items():
+  wrapper = wrap_inplace if k in inplace_ops else wrap_fxn
+  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
 
 @torch.library.impl("aten::equal", "privateuseone")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -628,42 +693,72 @@ if TORCH_DEBUG:
       return func(*args, **(kwargs or {}))
   (_dispatch_log:=DispatchLog()).__enter__() # NOTE: must be kept alive
 
-# NOTE: patch torch optimizer step to avoid continously growing the computation graph
-import weakref
-_torch_modules_with_buffers: weakref.WeakSet[torch.nn.Module] = weakref.WeakSet()
-def register_torch_buffer(mod, _name, _buffer): _torch_modules_with_buffers.add(mod)
-def get_real_tinygrad_buffers():
-  res = set()
-  for mod in _torch_modules_with_buffers:
-    for _,b in mod.named_buffers(recurse=False):
-      if b is not None and b.is_tiny:
-        res.add(unwrap(b))
-  return res
-torch.nn.modules.module.register_module_buffer_registration_hook(register_torch_buffer)
+# this implementation is needed to allow the batchnorm kernels to fuse in e.g. mnist training
+# aten::native_batch_norm does more than Tensor.batchnorm
+@torch.library.impl("aten::native_batch_norm", "privateuseone")
+def native_batch_norm(input, weight, bias, running_mean, running_var, training, momentum, eps):
+  input_t, weight_t, bias_t = unwrap(input), unwrap(weight) if weight is not None else None, unwrap(bias) if bias is not None else None
+  running_mean_t, running_var_t = unwrap(running_mean) if running_mean is not None else None, unwrap(running_var) if running_var is not None else None
+  if training:
+    batch_var, batch_mean = input_t.var_mean(axis=tuple(x for x in range(input_t.ndim) if x != 1), correction=0)
+    batch_invstd = batch_var.add(eps).rsqrt()
+    out = input_t.batchnorm(weight_t, bias_t, batch_mean, batch_invstd)
+    if running_mean_t is not None and running_var_t is not None:
+      numel_ratio = input_t.numel() / (input_t.numel() - input_t.shape[1])
+      running_mean_t.assign((1 - momentum) * running_mean_t + momentum * batch_mean.detach())
+      running_var_t.assign((1 - momentum) * running_var_t + momentum * numel_ratio * batch_var.detach())
+    return wrap(out), wrap(batch_mean), wrap(batch_invstd)
+  else:
+    out = input_t.batchnorm(weight_t, bias_t, running_mean_t, running_var_t.add(eps).rsqrt())
+    return wrap(out), wrap(running_mean_t), wrap(running_var_t.add(eps).rsqrt())
 
-from torch.nn.modules import Module
-def param_hook(_grad):
-  if _grad is not None and _grad.is_tiny: Tensor.realize(unwrap(_grad))
-def module_hook(module:Module, _name, _submodule):
-  for param in _submodule.parameters(recurse=False):
-    if param.requires_grad: param.register_hook(param_hook)
-torch.nn.modules.module.register_module_module_registration_hook(module_hook)
+@torch.library.impl("aten::native_batch_norm_backward", "privateuseone")
+def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask):
+  grad_out_t, input_t = unwrap(grad_out), unwrap(input)
+  weight_t = unwrap(weight) if weight is not None else None
+  save_mean_t = unwrap(save_mean)
+  save_invstd_t = unwrap(save_invstd)
+  out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t)
+  targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m]
+  if targets:
+    grads = out.gradient(*targets, gradient=grad_out_t)
+    grad_input = grads.pop(0) if output_mask[0] else None
+    grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None
+  else:
+    grad_input, grad_weight = None, None
+  grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None
+  return (wrap(grad_input) if grad_input is not None else None,
+          wrap(grad_weight) if grad_weight is not None else None,
+          wrap(grad_bias) if grad_bias is not None else None)
 
-def realize_optimizer_step(optimizer: torch.optim.Optimizer, *args, **kwargs):
-  tinygrad_tensors = []
-  for param_group in optimizer.param_groups:
-    for param in param_group["params"]:
-      if param is None: continue
-      tinygrad_tensors.append(param.data)
-  for state_dict in optimizer.state.values():
-    for _, value in state_dict.items():
-      if torch.is_tensor(value): tinygrad_tensors.append(value)
-  real_tinygrad_tensors = [unwrap(x) for x in tinygrad_tensors if x.is_tiny]
-  real_tinygrad_tensors += get_real_tinygrad_buffers()
-  if len(real_tinygrad_tensors): Tensor.realize(*real_tinygrad_tensors)
+# _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
+# we need torch.autograd.Function with explicit AutogradPrivateUse1 registration
+class _PadCircular(torch.autograd.Function):
+  @staticmethod
+  def forward(ctx, input, padding):
+    ctx.save_for_backward(input)
+    ctx.padding = padding
+    return pad_forward(input, padding, mode="circular")
+  @staticmethod
+  def backward(ctx, grad_output):
+    input, = ctx.saved_tensors
+    return pad_backward(grad_output, input, ctx.padding, mode="circular"), None
 
-_optimizer_init = torch.optim.Optimizer.__init__
-def _optimizer_patched_init(self, *args, **kwargs):
-  _optimizer_init(self, *args, **kwargs)
-  self.register_step_post_hook(realize_optimizer_step)
-torch.optim.Optimizer.__init__ = _optimizer_patched_init
+@torch.library.impl("aten::_pad_circular", "privateuseone")
+def _pad_circular(self, padding): return _PadCircular.apply(self, padding)
+
+@torch.library.impl("aten::_pad_circular", "AutogradPrivateUse1")
+def _pad_circular_autograd(self, padding): return _PadCircular.apply(self, padding)
+
+# only needed for test_diag_backward_gradient_values
+# was going through torch before, but now we are using tinygrad directly and tracking views
+# Tensor.diagonal does not support all cases tests in the tests
+@torch.library.impl("aten::diagonal", "privateuseone")
+@wrap_view_op
+def diagonal(self, offset=0, dim1=0, dim2=1):
+  if offset != 0: raise NotImplementedError(f"diagonal with {offset=} not implemented")
+  dim1, dim2 = dim1 % self.ndim, dim2 % self.ndim
+  if dim1 != self.ndim - 2 or dim2 != self.ndim - 1: raise NotImplementedError(f"diagonal with {dim1=}, {dim2=} not implemented, only last two dims supported")
+  batch_shape, m, n = self.shape[:-2], self.shape[-2], self.shape[-1]
+  diag_len = min(m, n)
+  return self.reshape(*batch_shape, m*n).pad(tuple((0,0) for _ in batch_shape) + ((0, diag_len),)).reshape(*batch_shape, diag_len, n+1)[..., :, 0]

extra/torch_backend/example.py

@@ -1,12 +1,13 @@
 from PIL import Image
-from tinygrad.helpers import getenv
-import torch, torchvision, pathlib
+from tinygrad.helpers import getenv, GlobalCounters
+import torch, torchvision, pathlib, warnings
 import torchvision.transforms as transforms
 import extra.torch_backend.backend
 device = "tiny"
 torch.set_default_device(device)
 
 if __name__ == "__main__":
+  GlobalCounters.reset()
   img = Image.open(pathlib.Path(__file__).parent.parent.parent / "test/models/efficientnet/Chicken.jpg").convert('RGB')
   transform = transforms.Compose([
     transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(),
@@ -19,3 +20,10 @@ if __name__ == "__main__":
   out = model(img).detach().cpu().numpy()
   print("output:", out.shape, out.argmax())
   assert out.argmax() == 7  # cock
+
+  kernel_count = GlobalCounters.kernel_count
+  assert kernel_count > 0, "No kernels, test failed"
+  expected_kernels = 228
+  expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
+  if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
+  assert kernel_count <= expected_kernels, f"{expectation}"

extra/torch_backend/test.py

@@ -2,7 +2,7 @@
 import unittest
 import torch
 import numpy as np
-from tinygrad.helpers import getenv, Context, GlobalCounters
+from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
   import extra.torch_backend.backend2
   device = "cpu"
@@ -25,7 +25,7 @@ class TestTorchBackend(unittest.TestCase):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
 
-  def test_numpy_ones(self):
+  def test_numpy_ones_int32(self):
     a = torch.ones(4, dtype=torch.int32, device=device)
     assert a.dtype == torch.int32
     np.testing.assert_equal(a.cpu().numpy(), [1,1,1,1])
@@ -219,7 +219,6 @@ class TestTorchBackend(unittest.TestCase):
     a = torch.ones(4, device=device)
     print(str(a))
 
-  @unittest.skip("failed")
   def test_floor_div(self):
     a = torch.tensor([10., 7., 5.], device=device)
     b = torch.tensor([3., 2., 2.], device=device)
@@ -248,5 +247,672 @@ class TestTorchBackend(unittest.TestCase):
   def test_diagonal_rectangular(self): self._test_diagonal(4, 5, 6)
   def test_diagonal_4d(self): self._test_diagonal(2, 3, 4, 5)
 
+  def test_pad_circular_simple(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2)
+    padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
+    expected = np.array([[[[3.,2.,3.,2.], [1.,0.,1.,0.], [3.,2.,3.,2.], [1.,0.,1.,0.]]]], dtype=np.float32)
+    np.testing.assert_allclose(padded.cpu().numpy(), expected)
+
+  def test_pad_circular_backward(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1,1,2,2).requires_grad_(True)
+    padded = torch.nn.functional.pad(a, (1,1,1,1), mode="circular")
+    loss = padded.sum()
+    loss.backward()
+    expected_grad = np.array([[[[4., 4.], [4., 4.]]]], dtype=np.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad)
+
+
+  def test_matmul_backward(self):
+    x = torch.randn(3, 4, device=device, dtype=torch.float32, requires_grad=True)
+    y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
+    z = (x @ y).sum()
+    z.backward()
+    assert x.grad is not None
+    assert y.grad is not None
+    assert x.grad.shape == x.shape
+    assert y.grad.shape == y.shape
+
+  def test_matmul_broadcast_backward(self):
+    x = torch.randn(2, 3, 4, device=device, dtype=torch.float32, requires_grad=True)
+    y = torch.randn(4, 5, device=device, dtype=torch.float32, requires_grad=True)
+    z = (x @ y).sum()
+    z.backward()
+    assert x.grad is not None
+    assert y.grad is not None
+    assert x.grad.shape == x.shape
+    assert y.grad.shape == y.shape
+
+  def test_diag_vector_to_matrix(self):
+    vec = torch.tensor([1., 2., 3., 4., 5.], dtype=torch.float32, device=device)
+    mat = torch.diag(vec)
+    expected = np.diag([1., 2., 3., 4., 5.])
+    np.testing.assert_allclose(mat.cpu().numpy(), expected, rtol=1e-5)
+    assert mat.shape == (5, 5)
+
+  def test_diagonal_matrix_to_vector(self):
+    mat = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device)
+    vec = torch.linalg.diagonal(mat)
+    expected = np.array([1., 5., 9.])
+    np.testing.assert_allclose(vec.cpu().numpy(), expected, rtol=1e-5)
+    assert vec.shape == (3,)
+
+  def test_permute_2(self):
+    a = torch.randn(2, 3, 4, dtype=torch.float32, device=device)
+    b = a.permute(2, 0, 1)
+    assert b.shape == (4, 2, 3)
+    np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))
+
+  def test_batchnorm_unsqueeze(self):
+    bn = torch.nn.BatchNorm2d(4).to(device)
+    x = torch.randn(8, 4, 3, 3, device=device)
+    out = bn(x)
+    self.assertEqual(out.shape, x.shape)
+
+  def test_slice_inplace_zero(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b.zero_()
+    expected = np.array([[1., 1., 1.],
+                         [1., 0., 0.],
+                         [1., 0., 0.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_slice_inplace_fill(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b.fill_(5.0)
+    expected = np.array([[1., 1., 1.],
+                         [1., 5., 5.],
+                         [1., 5., 5.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_fill_tensor_value(self):
+    a = torch.zeros((2, 2), dtype=torch.float32, device=device)
+    value = torch.tensor(3, dtype=torch.int64, device=device)
+    a.fill_(value)
+    expected = np.full((2, 2), 3, dtype=np.float32)
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_slice_inplace_mul(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:]
+    b *= 2
+    expected = np.array([[1., 1., 1.],
+                         [1., 2., 2.],
+                         [1., 2., 2.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_permute_slice_zero(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:].permute(1, 0)
+    b.zero_()
+    expected = np.array([[1., 1., 1.],
+                         [1., 0., 0.],
+                         [1., 0., 0.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_permute_slice_mul(self):
+    a = torch.ones((3, 3), device=device)
+    b = a[1:, 1:].permute(1, 0)
+    b *= 2
+    expected = np.array([[1., 1., 1.],
+                         [1., 2., 2.],
+                         [1., 2., 2.]])
+    np.testing.assert_equal(a.cpu().numpy(), expected)
+
+  def test_simple_slice_setitem(self):
+    a = torch.tensor([10, 20, 30], device=device)
+    a[1] = 99
+    np.testing.assert_equal(a.cpu().numpy(), [10, 99, 30])
+
+  def test_2d_slice_setitem(self):
+    a = torch.zeros((3, 3), device=device)
+    a[1, 2] = 99
+    self.assertEqual(a[1, 2].item(), 99)
+    self.assertEqual(a.sum().item(), 99)
+
+  def test_view_copy(self):
+    a = torch.tensor([10, 20, 30], device=device)
+    view = a[1]
+    view.copy_(torch.tensor(88, device=device))
+    np.testing.assert_equal(a.cpu().numpy(), [10, 88, 30])
+
+  def test_diag_2d_input(self):
+    a = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device)
+    d = torch.diag(a)
+    np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
+
+  def test_diag_1d_input(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    d = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(d.cpu().numpy(), expected)
+
+  def test_permute_view_tracking(self):
+    a = torch.ones((2, 3, 4), device=device)
+    b = a.permute(2, 0, 1)
+    self.assertEqual(b.shape, (4, 2, 3))
+
+  def test_detach_view_creation(self):
+    a = torch.tensor([1.0, 2.0, 3.0], device=device)
+    b = a.detach()
+    np.testing.assert_equal(b.cpu().numpy(), [1.0, 2.0, 3.0])
+
+  def test_view_zero_inplace(self):
+    a = torch.ones((4, 4), device=device)
+    view = a[1:3, 1:3]
+    view.zero_()
+    self.assertEqual(view.sum().item(), 0)
+
+  def test_view_fill_inplace(self):
+    a = torch.zeros((4, 4), device=device)
+    view = a[1:3, 1:3]
+    view.fill_(5)
+    self.assertEqual(view.sum().item(), 20)
+
+  def test_permute_contiguous(self):
+    a = torch.tensor([[1, 2], [3, 4]], device=device)
+    b = a.permute(1, 0)
+    c = b.contiguous()
+    expected = [[1, 3], [2, 4]]
+    np.testing.assert_equal(c.cpu().numpy(), expected)
+
+  def test_diag_2d_extract_diagonal(self):
+    a = torch.tensor([[1, 2], [3, 4]], device=device)
+    result = torch.diag(a)
+    np.testing.assert_equal(result.cpu().numpy(), [1, 4])
+
+  def test_slice_inplace_multiply_offset_preservation(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    a[1:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [1, 4, 6])
+
+  def test_slice_inplace_mul_pattern(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[:2] *= 3
+    a[2:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [3, 6, 6, 8])
+
+  def test_chained_slice_column(self):
+    a = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4)
+    torch_res = a[:, 1:2][:, 0:1].cpu().numpy()
+    cpu_res = torch.arange(16, dtype=torch.float32).reshape(4, 4)[:, 1:2][:, 0:1].numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_slice_with_step(self):
+    a = torch.arange(20, dtype=torch.float32, device=device)
+    torch_res = a[::2][1:4].cpu().numpy()
+    cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_slice_negative_dim(self):
+    a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
+    torch_chunks = a.chunk(3, -1)
+    cpu_chunks = torch.arange(13, dtype=torch.int32).repeat(8, 1).chunk(3, -1)
+    assert len(torch_chunks) == len(cpu_chunks)
+    for i in range(len(torch_chunks)):
+      np.testing.assert_equal(torch_chunks[i].cpu().numpy(), cpu_chunks[i].numpy())
+
+  def test_dot_vector_matrix(self):
+    a = torch.arange(65, dtype=torch.float32, device=device)
+    b = torch.arange(65*45, dtype=torch.float32, device=device).reshape(65, 45)
+    torch_res = a.matmul(b).reshape(-1).cpu().numpy()
+    cpu_res = torch.arange(65, dtype=torch.float32).matmul(torch.arange(65*45, dtype=torch.float32).reshape(65, 45)).numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_alias_passthrough(self):
+    a = torch.randn(3, 3, device=device)
+    alias_view = torch.ops.aten.alias(a)
+    alias_view += 1
+    np.testing.assert_equal(a.cpu().numpy(), alias_view.cpu().numpy())
+
+  def test_split_simple_vector(self):
+    a = torch.arange(10, dtype=torch.float32, device=device)
+    torch_chunks = a.split([1,4,5])
+    cpu_chunks = torch.arange(10, dtype=torch.float32).split([1,4,5])
+    for tc, cc in zip(torch_chunks, cpu_chunks):
+      np.testing.assert_equal(tc.cpu().numpy(), cc.cpu().numpy())
+
+  def test_split_matches_torch(self):
+    a = torch.arange(10, dtype=torch.float32, device=device)
+    torch_chunks = a.split([1,4,5])
+    tiny_chunks = [chunk.cpu().numpy() for chunk in torch_chunks]
+    cpu_chunks = [torch.arange(10, dtype=torch.float32).split([1,4,5])[i].numpy() for i in range(3)]
+    for tr, cr in zip(tiny_chunks, cpu_chunks): np.testing.assert_equal(tr, cr)
+
+  def test_sum_matches_torch(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2,3)
+    torch_res = a.sum().cpu().numpy()
+    cpu_res = torch.arange(6, dtype=torch.float32).reshape(2,3).sum().numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_view_matches_torch(self):
+    a = torch.arange(6, dtype=torch.float32, device=device)
+    torch_res = a.view(2, 3).cpu().numpy()
+    cpu_res = torch.arange(6, dtype=torch.float32).view(2, 3).numpy()
+    np.testing.assert_equal(torch_res, cpu_res)
+
+  def test_view_zero_with_indices(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[1:3].zero_()
+    np.testing.assert_equal(a.cpu().numpy(), [1, 0, 0, 4])
+
+  def test_view_fill_with_indices(self):
+    a = torch.tensor([1, 2, 3, 4], device=device)
+    a[::2].fill_(9)
+    np.testing.assert_equal(a.cpu().numpy(), [9, 2, 9, 4])
+
+  def test_nested_slice_inplace_ops(self):
+    a = torch.tensor([1, 2, 3, 4, 5, 6], device=device)
+    a[:3] += 10
+    a[3:] *= 2
+    np.testing.assert_equal(a.cpu().numpy(), [11, 12, 13, 8, 10, 12])
+
+  def test_diag_1d(self):
+    a = torch.tensor([1, 2, 3], device=device)
+    result = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(result.cpu().numpy(), expected)
+
+  def test_diag_backward(self):
+    a = torch.randn(5, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_diagonal(self):
+    a = torch.tensor([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    expected = torch.tensor([1., 5., 9.], dtype=torch.float32)
+    self.assertEqual(b.shape, (3,))
+    np.testing.assert_allclose(b.detach().cpu().numpy(), expected.numpy(), rtol=1e-5)
+
+  def test_diagonal_backward(self):
+    a = torch.randn(5, 5, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_expand_backward(self):
+    a = torch.randn(4, 3, 1, 6, dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(4, 3, 2, 6)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_einsum_backward(self):
+    a = torch.randn(10, 10, dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.einsum('ij->ji', a)
+    b.sum().backward()
+    assert a.grad is not None
+
+  def test_diag_backward_gradient_values(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.ones(3, dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_backward_gradient_values_2d_to_1d(self):
+    a = torch.tensor([[1.0, 2.0, 3.0],
+                      [4.0, 5.0, 6.0],
+                      [7.0, 8.0, 9.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diagonal(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[1.0, 0.0, 0.0],
+                                   [0.0, 1.0, 0.0],
+                                   [0.0, 0.0, 1.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_backward_gradient_values(self):
+    a = torch.tensor([[1.0], [2.0], [3.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(3, 4)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[4.0], [4.0], [4.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_backward_with_leading_dims(self):
+    a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(3, 1, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[3.0, 3.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_2d_to_1d_backward(self):
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_expand_complex_backward(self):
+    a = torch.tensor([[[1.0, 2.0]]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.expand(2, 3, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[[6.0, 6.0]]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_diag_backward_with_scaling(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    loss = (b * torch.tensor([[2.0, 0.0, 0.0],
+                               [0.0, 3.0, 0.0],
+                               [0.0, 0.0, 4.0]], device=device)).sum()
+    loss.backward()
+    expected_grad = torch.tensor([2.0, 3.0, 4.0], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_repeat_basic(self):
+    a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
+    b = a.repeat(2, 1)
+    expected = torch.tensor([[1, 2, 3], [1, 2, 3]], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_repeat_multidim(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
+    b = a.repeat(2, 3)
+    expected = torch.arange(6, dtype=torch.float32).reshape(2, 3).repeat(2, 3)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_repeat_backward(self):
+    a = torch.tensor([[1.0, 2.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = a.repeat(3, 2)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([[6.0, 6.0]], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_cumsum_1d(self):
+    a = torch.tensor([1, 2, 3, 4], dtype=torch.float32, device=device)
+    b = torch.cumsum(a, dim=0)
+    expected = torch.tensor([1, 3, 6, 10], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_cumsum_2d(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    b = torch.cumsum(a, dim=0)
+    expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=0)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+    c = torch.cumsum(a, dim=1)
+    expected = torch.arange(12, dtype=torch.float32).reshape(3, 4).cumsum(dim=1)
+    np.testing.assert_equal(c.cpu().numpy(), expected.numpy())
+
+  def test_cumsum_backward(self):
+    a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.tensor([4.0, 3.0, 2.0, 1.0], dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_constant_pad_nd_1d(self):
+    a = torch.tensor([1, 2, 3], dtype=torch.float32, device=device)
+    b = torch.nn.functional.pad(a, (1, 2), mode='constant', value=0)
+    expected = torch.tensor([0, 1, 2, 3, 0, 0], dtype=torch.float32)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_constant_pad_nd_2d(self):
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(2, 3)
+    b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
+    expected = torch.nn.functional.pad(torch.arange(6, dtype=torch.float32).reshape(2, 3), (1, 1, 1, 1), mode='constant', value=0)
+    np.testing.assert_equal(b.cpu().numpy(), expected.numpy())
+
+  def test_constant_pad_nd_2d_backward(self):
+    a = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.nn.functional.pad(a, (1, 1, 1, 1), mode='constant', value=0)
+    loss = b.sum()
+    loss.backward()
+    expected_grad = torch.ones((2, 2), dtype=torch.float32)
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected_grad.numpy(), rtol=1e-5)
+
+  def test_negative_strides_cumsum_backward(self):
+    a = torch.randn(5, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    b.sum().backward()
+    grad = a.grad.cpu().numpy()
+    self.assertEqual(len(grad), 5)
+
+  def test_cumsum_fix_gradient_values(self):
+    a = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.cumsum(a, dim=0)
+    loss = b.sum()
+    loss.backward()
+    expected = np.array([4.0, 3.0, 2.0, 1.0])
+    np.testing.assert_allclose(a.grad.cpu().numpy(), expected, rtol=1e-5)
+
+  def test_diag_1d_to_2d(self):
+    a = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32, device=device, requires_grad=True)
+    b = torch.diag(a)
+    expected = [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    np.testing.assert_equal(b.detach().cpu().numpy(), expected)
+
+  def test_diag_2d_to_1d(self):
+    c = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=torch.float32, device=device)
+    d = torch.diag(c)
+    np.testing.assert_equal(d.cpu().numpy(), [1, 5, 9])
+
+  def test_biased_conv2d(self):
+    # Test case for two sequential conv2d with same weights/bias and ReLU in between, this is as special case from test_ops.py
+    torch.manual_seed(0)
+    C = 8
+    x_cpu = torch.randn(1, C, 5, 5, requires_grad=True)
+    w_cpu = torch.randn(C, C, 1, 1, requires_grad=True)
+    b_cpu = torch.randn(C, requires_grad=True)
+    x_tiny = x_cpu.detach().to(device).requires_grad_(True)
+    w_tiny = w_cpu.detach().to(device).requires_grad_(True)
+    b_tiny = b_cpu.detach().to(device).requires_grad_(True)
+    out_cpu = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_cpu, w_cpu, b_cpu).relu(), w_cpu, b_cpu)
+    out_tiny = torch.nn.functional.conv2d(torch.nn.functional.conv2d(x_tiny, w_tiny, b_tiny).relu(), w_tiny, b_tiny)
+    grad_out = torch.randn_like(out_cpu)
+    out_cpu.backward(grad_out)
+    out_tiny.backward(grad_out.to(device))
+    np.testing.assert_allclose(x_tiny.grad.cpu().numpy(), x_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+    np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+    np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
+
+
+from tinygrad import Tensor
+class TestBackendHelpers(unittest.TestCase):
+
+  def test_calculate_storage_offset_no_shrink(self):
+    t = Tensor.ones(3, 4)
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 0
+
+  def test_calculate_storage_offset_with_shrink(self):
+    t = Tensor.ones(10, 10)[2:5, 3:7]
+    # strides for (10, 10) are [10, 1]
+    # offset = 2*10 + 3*1 = 23
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 23
+
+  def test_calculate_storage_offset_multiple_shrinks(self):
+    t = Tensor.ones(5, 6, 7)[1:3, 2:4, 3:5]
+    # strides for (5, 6, 7) are [42, 7, 1]
+    # offset = 1*42 + 2*7 + 3*1 = 42 + 14 + 3 = 59
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == 59
+
+  def test_calculate_storage_offset_with_reshape(self):
+    t = Tensor.ones(10, 10)
+    orig_offset = extra.torch_backend.backend.calculate_storage_offset(t)
+    assert orig_offset == 0
+    t = t.reshape(100)
+    assert extra.torch_backend.backend.calculate_storage_offset(t) == orig_offset
+
+  def test_slice_values_match_torch(self):
+    torch_cpu = torch.arange(100, dtype=torch.float32).reshape(10, 10)
+    torch_tiny = torch_cpu.to(device)
+    sliced_cpu = torch_cpu[2:5, 3:7]
+    sliced_tiny = torch_tiny[2:5, 3:7]
+    np.testing.assert_equal(sliced_tiny.cpu().numpy(), sliced_cpu.numpy())
+
+  def test_slice_values_match_torch_3d(self):
+    torch_cpu_3d = torch.arange(210, dtype=torch.float32).reshape(5, 6, 7)
+    torch_tiny_3d = torch_cpu_3d.to(device)
+    sliced_cpu_3d = torch_cpu_3d[1:3, 2:4, 3:5]
+    sliced_tiny_3d = torch_tiny_3d[1:3, 2:4, 3:5]
+    np.testing.assert_equal(sliced_tiny_3d.cpu().numpy(), sliced_cpu_3d.numpy())
+
+  def test_topk_out(self):
+    a = torch.tensor([1, 3, 2, 4], device=device)
+    values = torch.empty(2, device=device)
+    indices = torch.empty(2, dtype=torch.int64, device=device)
+    ret_values, ret_indices = torch.topk(a, k=2, out=(values, indices))
+    np.testing.assert_equal(values.cpu().numpy(), [4, 3])
+    np.testing.assert_equal(indices.cpu().numpy(), [3, 1])
+    assert ret_values is values
+    assert ret_indices is indices
+
+  def test_sort_out(self):
+    a = torch.tensor([3, 1, 4, 2], device=device)
+    values = torch.empty(4, device=device)
+    indices = torch.empty(4, dtype=torch.int64, device=device)
+    ret_values, ret_indices = torch.sort(a, out=(values, indices))
+    np.testing.assert_equal(values.cpu().numpy(), [1, 2, 3, 4])
+    np.testing.assert_equal(indices.cpu().numpy(), [1, 3, 0, 2])
+    assert ret_values is values
+    assert ret_indices is indices
+
+  def test_cat_out(self):
+    a = torch.tensor([1, 2], device=device)
+    b = torch.tensor([3, 4], device=device)
+    out = torch.empty(4, device=device)
+    ret = torch.cat([a, b], out=out)
+    np.testing.assert_equal(out.cpu().numpy(), [1, 2, 3, 4])
+    assert ret is out
+
+  def test_scatter_add_out(self):
+    src = torch.tensor([[1, 2, 3], [4, 5, 6]], device=device, dtype=torch.float32)
+    index = torch.tensor([[0, 1, 2], [0, 1, 2]], device=device)
+    input = torch.zeros(3, 3, device=device, dtype=torch.float32)
+    out = torch.zeros(3, 3, device=device, dtype=torch.float32)
+    ret = torch.scatter_add(input, 0, index, src, out=out)
+    expected = torch.tensor([[5, 0, 0], [0, 7, 0], [0, 0, 9]], dtype=torch.float32)
+    np.testing.assert_allclose(out.cpu().numpy(), expected.cpu().numpy())
+    assert ret is out
+
+  def test_floor_divide_inplace_identity(self):
+    x = torch.tensor([10, 20, 30, 40], dtype=torch.int32, device=device)
+    y = torch.tensor([2, 4, 5, 8], dtype=torch.int32, device=device)
+    ret = x.floor_divide_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [5, 5, 6, 5])
+
+  def test_lshift_inplace_identity(self):
+    x = torch.tensor([1, 2, 3, 4], dtype=torch.int32, device=device)
+    ret = x.__ilshift__(2)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
+
+  def test_rshift_inplace_identity(self):
+    x = torch.tensor([16, 32, 48, 64], dtype=torch.int32, device=device)
+    ret = x.__irshift__(2)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [4, 8, 12, 16])
+
+  def test_relu_inplace_identity(self):
+    x = torch.tensor([-1.0, 2.0, -3.0, 4.0], device=device)
+    ret = x.relu_()
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
+
+  def test_random_inplace_identity(self):
+    x = torch.zeros(10, dtype=torch.int32, device=device)
+    ret = x.random_()
+    assert ret is x
+    assert x.shape == (10,)
+
+  def test_random_from_inplace_identity(self):
+    x = torch.zeros(10, dtype=torch.int32, device=device)
+    ret = x.random_(5, 10)
+    assert ret is x
+    # values should be in range [5, 10)
+    assert torch.all(x >= 5).item() and torch.all(x < 10).item()
+
+  def test_uniform_inplace_identity(self):
+    x = torch.zeros(10, device=device)
+    ret = x.uniform_(0.0, 1.0)
+    assert ret is x
+    # values should be in range [0, 1)
+    assert torch.all(x >= 0.0).item() and torch.all(x < 1.0).item()
+
+  def test_normal_inplace_identity(self):
+    x = torch.zeros(100, device=device)
+    ret = x.normal_(0.0, 1.0)
+    assert ret is x
+    # just check that values changed from zeros
+    assert not torch.all(x == 0.0).item()
+
+  def test_logical_or_inplace_identity(self):
+    x = torch.tensor([True, False, True, False], device=device)
+    y = torch.tensor([False, False, True, True], device=device)
+    ret = x.logical_or_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [True, False, True, True])
+
+  def test_masked_fill_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    mask = torch.tensor([True, False, True, False], device=device)
+    ret = x.masked_fill_(mask, 0.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 2.0, 0.0, 4.0])
+
+  def test_masked_fill_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    mask = torch.tensor([True, False, True, False], device=device)
+    value = torch.tensor(99.0, device=device)
+    ret = x.masked_fill_(mask, value)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [99.0, 2.0, 99.0, 4.0])
+
+  def test_zero_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.zero_()
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [0.0, 0.0, 0.0, 0.0])
+
+  def test_fill_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.fill_(5.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [5.0, 5.0, 5.0, 5.0])
+
+  def test_fill_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    value = torch.tensor(7.0, device=device)
+    ret = x.fill_(value)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [7.0, 7.0, 7.0, 7.0])
+
+  def test_add_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    y = torch.tensor([10.0, 20.0, 30.0, 40.0], device=device)
+    ret = x.add_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [11.0, 22.0, 33.0, 44.0])
+
+  def test_add_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.add_(10.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [11.0, 12.0, 13.0, 14.0])
+
+  def test_mul_tensor_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    y = torch.tensor([2.0, 3.0, 4.0, 5.0], device=device)
+    ret = x.mul_(y)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [2.0, 6.0, 12.0, 20.0])
+
+  def test_mul_scalar_inplace_identity(self):
+    x = torch.tensor([1.0, 2.0, 3.0, 4.0], device=device)
+    ret = x.mul_(2.0)
+    assert ret is x
+    np.testing.assert_equal(x.cpu().numpy(), [2.0, 4.0, 6.0, 8.0])
+
 if __name__ == "__main__":
   unittest.main()

extra/torch_backend/test_kernel_fusion.py

@@ -0,0 +1,144 @@
+# simple tests
+import unittest
+import torch
+import warnings
+from tinygrad.helpers import getenv, GlobalCounters
+if getenv("TINY_BACKEND2"):
+  import extra.torch_backend.backend2
+  device = "cpu"
+else:
+  import extra.torch_backend.backend
+  device = "tiny"
+
+
+class TestKernelFusionRegression(unittest.TestCase):
+  def _realize(self, t): _ = t.detach().cpu().numpy()
+
+  def _check_kernel_count(self, fn, expected_kernels):
+    torch.manual_seed(42)
+    GlobalCounters.reset()
+    fn().detach().cpu().numpy()
+    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
+    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
+    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
+
+  def test_elementwise_fusion(self):
+    def fn():
+      x = torch.randn(128, 128, device=device)
+      return (x + 1.0) * 2.0 - 0.5
+    self._check_kernel_count(fn, 6)
+
+  def test_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 3, 32, 32, device=device)
+      conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
+      with torch.no_grad():
+        return torch.nn.functional.relu(conv(x))
+    self._check_kernel_count(fn, 8)
+
+  def test_batchnorm_fusion(self):
+    def fn():
+      x = torch.randn(2, 3, 16, 16, device=device)
+      conv = torch.nn.Conv2d(3, 8, 3, padding=1).to(device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.eval()
+      with torch.no_grad():
+        return torch.nn.functional.relu(bn(conv(x)))
+    self._check_kernel_count(fn, 16)
+
+  def test_reduce_fusion(self):
+    def fn():
+      x = torch.randn(64, 64, device=device)
+      return (x * 2.0).sum()
+    self._check_kernel_count(fn, 7)
+
+  def test_matmul_elementwise_fusion(self):
+    def fn():
+      x = torch.randn(32, 32, device=device)
+      w = torch.randn(32, 32, device=device)
+      return torch.nn.functional.relu(x @ w + 1.0)
+    self._check_kernel_count(fn, 6)
+
+  def test_pooling_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      return torch.nn.functional.max_pool2d(x * 2.0, 2)
+    self._check_kernel_count(fn, 5)
+
+  def test_residual_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      identity = torch.randn(1, 8, 16, 16, device=device)
+      out = x + identity
+      return torch.nn.functional.relu(out)
+    self._check_kernel_count(fn, 6)
+
+  def test_inplace_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 16, 32, 32, device=device)
+      y = torch.randn(1, 16, 32, 32, device=device)
+      x += y
+      return torch.nn.functional.relu(x)
+    self._check_kernel_count(fn, 6)
+
+  def test_conv_bn_add_relu_fusion(self):
+    def fn():
+      x = torch.randn(1, 8, 16, 16, device=device)
+      identity = torch.randn(1, 8, 16, 16, device=device)
+      conv = torch.nn.Conv2d(8, 8, 3, padding=1, bias=False).to(device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.eval()
+      with torch.no_grad():
+        out = bn(conv(x))
+        out += identity
+        return torch.nn.functional.relu(out)
+    self._check_kernel_count(fn, 16)
+
+  def test_multiple_inplace_ops_fusion(self):
+    def fn():
+      x = torch.randn(64, 64, device=device)
+      x += 1.0
+      x *= 2.0
+      return torch.nn.functional.relu(x)
+    self._check_kernel_count(fn, 4)
+
+  def test_view_inplace_no_fusion_break(self):
+    def fn():
+      x = torch.randn(4, 64, device=device)
+      view = x[1:3]
+      view += 1.0
+      return x.sum()
+    self._check_kernel_count(fn, 8)
+
+  def test_batchnorm_running_stats_update(self):
+    def fn():
+      x = torch.randn(2, 8, 8, 8, device=device)
+      bn = torch.nn.BatchNorm2d(8).to(device)
+      bn.train()
+      with torch.no_grad():
+        return bn(x)
+    self._check_kernel_count(fn, 10)
+
+  # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
+  def test_mnist_training_fusion(self):
+    def fn():
+      model = torch.nn.Sequential(
+        torch.nn.Conv2d(1, 8, 3, padding=1),
+        torch.nn.ReLU(),
+        torch.nn.MaxPool2d(2),
+        torch.nn.Flatten(),
+        torch.nn.Linear(8*14*14, 10)
+      ).to(device)
+      optimizer = torch.optim.Adam(model.parameters(), 1e-3)
+      x = torch.randn(32, 1, 28, 28, device=device)
+      labels = torch.randint(0, 10, (32,), device=device)
+      out = model(x)
+      loss = torch.nn.functional.cross_entropy(out, labels)
+      optimizer.zero_grad()
+      loss.backward()
+      optimizer.step()
+      return loss
+    self._check_kernel_count(fn, 33)
+
+if __name__ == "__main__":
+  unittest.main()

extra/torch_backend/wrapped_tensor.cpp

@@ -113,16 +113,9 @@ int register_hook() {
 int temp_register_hook = register_hook();
 
 at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceIndex device_index) {
-  // TODO: we have to get the dtype and the shape from the tinygrad Tensor
   std::vector<int64_t> sizes = py_obj.attr("shape").cast<std::vector<int64_t>>();
-
-  py::list views = py_obj.attr("uop").attr("st").attr("views");
-  std::vector<int64_t> strides = views[views.size() - 1].attr("strides").cast<std::vector<int64_t>>();
-  int64_t storage_offset = 0;
-  for (auto& v: views) {
-    storage_offset += v.attr("offset").cast<int64_t>(); // TODO: is this correct?
-  }
-
+  std::vector<int64_t> strides = py_obj.attr("_strides").cast<std::vector<int64_t>>();
+  int64_t storage_offset = py_obj.attr("_storage_offset").cast<int64_t>();
   return at::detail::make_tensor<at::TinyOpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>>(
     at::DispatchKeySet(at::DispatchKey::PrivateUse1),
     c10::scalarTypeToTypeMeta(dtype),

test/external/external_uop_gc.py

@@ -1,7 +1,8 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.engine.realize import method_cache, get_program
-from tinygrad.schedule.indexing import apply_movement_op
+from tinygrad.schedule.indexing import _apply_movement_op
+from tinygrad.uop.divandmod import fold_divmod_general
 from test.test_tiny import TestTiny
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
@@ -68,7 +69,8 @@ if __name__ == "__main__":
 
     # these caches will keep uops alive
     method_cache.clear()
-    apply_movement_op.cache_clear()
+    _apply_movement_op.cache_clear()
+    fold_divmod_general.cache_clear()
     Tensor._device_seeds.clear()
     Tensor._device_rng_counters.clear()
 

test/external/process_replay/process_replay.py

@@ -36,7 +36,9 @@ def trunc_log(x):
   logging.info("\n".join(lines))
 
 # user config
-SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+# NOTE: process replay is slow so it's now disabled by default. add [pr] to enable it
+#SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
+SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
 if REF == "master": SKIP_PROCESS_REPLAY = True
 class ProcessReplayWarning(Warning): pass
 

test/test_outerworld.py

@@ -1,5 +1,6 @@
 import unittest
-from tinygrad import Tensor, UOp
+import numpy as np
+from tinygrad import Tensor, UOp, nn
 from tinygrad.uop.ops import AxisType, Ops
 
 class TestOuterworldReduce(unittest.TestCase):
@@ -77,7 +78,8 @@ class TestOuterScan(unittest.TestCase):
     # 3 matmuls with SCAN
     i = UOp.range(3, -100, AxisType.OUTER)
     out = Tensor.empty(3, 1, 10)
-    comp = Tensor(i.eq(0).where(vec.uop, out[(i-1).maximum(0)].uop)) @ mats[i]
+    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()
@@ -144,5 +146,85 @@ class TestOuterworld(unittest.TestCase):
     out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
     self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
+class TestVmap(unittest.TestCase):
+  def test_vmap_inner(self, axis_type=AxisType.LOOP, fuse=False, grad=False):
+    x = Tensor.ones(1, 10).contiguous().requires_grad_()
+    mats = Tensor.ones(3, 10, 10).contiguous().requires_grad_()
+
+    ref = x @ mats
+    if fuse: ref = ref * 2
+
+    # vmap across axis 0
+    a = UOp.range(3, -1, axis_type)
+    out = x @ mats[a]
+    out = out.reshape(1, 10).pad(((a,(3-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    if fuse: out = out * 2
+    if grad:
+      out.mean().backward()
+      np.testing.assert_allclose(mats.grad.numpy(), (2./30) if fuse else (1./30))
+    out.realize()
+
+    # TODO: testing allclose
+    assert Tensor.allclose(ref, out, atol=1e-6), f"{ref.numpy()=}, {out.numpy()=}"
+  def test_vmap_inner_fuse(self): self.test_vmap_inner(fuse=True)
+  def test_vmap_outer(self): self.test_vmap_inner(AxisType.OUTER)
+  def test_vmap_outer_fuse(self): self.test_vmap_inner(AxisType.OUTER, fuse=True)
+
+  def test_vmap_inner_grad(self): self.test_vmap_inner(grad=True)
+  def test_vmap_inner_fuse_grad(self): self.test_vmap_inner(fuse=True, grad=True)
+  def test_vmap_outer_grad(self): self.test_vmap_inner(AxisType.OUTER, grad=True)
+
+  def test_vmap_convs(self):
+    layers = [
+      nn.Conv2d(1, 8, 3), Tensor.relu,
+      nn.Conv2d(8, 8, 3), Tensor.relu]
+    img = Tensor.randn(4, 1, 16, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None, None, None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  def test_vmap_gemm(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    img = Tensor.randn(4, 16).realize(*nn.state.get_parameters(layers))
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.realize()
+    np.testing.assert_allclose(out.numpy(), img.sequential(layers).numpy(), atol=1e-6)
+
+  @unittest.skip("this is broken, we need to lower the outer reduce in the outer graph")
+  def test_vmap_gemm_grad(self):
+    layers = [
+      nn.Linear(16, 16, bias=False), Tensor.relu,
+      nn.Linear(16, 16, bias=False), Tensor.relu]
+    layer_tensors = nn.state.get_parameters(layers)
+    img = Tensor.randn(4, 16).realize(*layer_tensors)
+    for l in layer_tensors: l.requires_grad_()
+    a = UOp.range(4, -1, AxisType.OUTER)
+    out = img[a:a+1].sequential(layers)
+    out = out.pad(((a,(4-a)-1), None))
+    out = Tensor(out.uop.reduce(a, arg=Ops.ADD))
+    out.mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    out_grads = [x.numpy() for x in grads]
+
+    # compute reference grads
+    for l in layer_tensors: l.grad = None
+    img.sequential(layers).mean().backward()
+    grads = [l.grad for l in layer_tensors]
+    out.realize(*grads)
+    ref_grads = [x.numpy() for x in grads]
+
+    # compare
+    for o,r in zip(out_grads, ref_grads): np.testing.assert_allclose(o, r, atol=1e-6)
+
 if __name__ == '__main__':
   unittest.main()

test/test_uops.py

@@ -517,7 +517,7 @@ class TestUOpStr(unittest.TestCase):
 
 class TestUPatHelpers(unittest.TestCase):
   def test_location(self):
-    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "math.py")
+    self.assertEqual(sym.patterns[-1][0].location[0].replace("\\", "/").split("/")[-1], "symbolic.py")
     self.assertEqual(shared_spec.patterns[0][0].location[0].replace("\\", "/").split("/")[-1], "spec.py")
     test_upat = UPat(Ops.CONST, dtypes.bool)
     self.assertEqual(test_upat.location[0].split("/")[-1], __file__.replace("\\", "/").split("/")[-1])

test/unit/test_autogen.py

@@ -5,16 +5,16 @@ from tinygrad.runtime.support.c import Struct
 class TestAutogen(unittest.TestCase):
   def test_packed_struct_sizeof(self):
     layout = [('a', ctypes.c_char), ('b', ctypes.c_int, 5), ('c', ctypes.c_char)]
-    class X(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Y(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Z(Struct): _packed_, _fields_ = True, layout
-    self.assertNotEqual(ctypes.sizeof(X), 4) # ctypes bug! gcc-13.3.0 says this should have size 4
+    class Z(Struct): pass
+    Z._packed_, Z._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Y), 6)
     self.assertEqual(ctypes.sizeof(Z), 3)
     layout = [('a', ctypes.c_int, 31), ('b', ctypes.c_int, 31), ('c', ctypes.c_int, 1), ('d', ctypes.c_int, 1)]
     class Foo(ctypes.Structure): _fields_, _layout_ = layout, 'gcc-sysv'
     class Bar(ctypes.Structure): _fields_, _pack_, _layout_ = layout, 1, 'ms'
-    class Baz(Struct): _fields_, _packed_ = layout, True
+    class Baz(Struct): pass
+    Baz._packed_, Baz._fields_ = True, layout
     self.assertEqual(ctypes.sizeof(Foo), 12)
     self.assertEqual(ctypes.sizeof(Bar), 12)
     self.assertEqual(ctypes.sizeof(Baz), 8)

test/unit/test_microbenchmarks.py

@@ -1,4 +1,5 @@
 import unittest, time
+from tinygrad.helpers import Profiling
 from tinygrad.uop.ops import UOp
 from tinygrad.dtype import dtypes
 
@@ -38,6 +39,14 @@ class TestMicrobenchmarks(unittest.TestCase):
     a = UOp.const(dtypes.int, 2)
     for _ in range(N): (a+a).simplify()
 
+class TestMicroprofile(unittest.TestCase):
+  def test_uop_simplify_complex(self):
+    x = UOp.variable("x", 0, 10)
+    y = UOp.variable("y", 0, 10)
+    expr = (x*2)+5+(x*4)+(y*2)+y
+    with Profiling():
+      for _ in range(1000): expr.simplify()
+
 if __name__ == '__main__':
   unittest.main()
 

test/unit/test_symbolic_failures.py

@@ -159,3 +159,38 @@ class TestFuzzFailure(unittest.TestCase):
     num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
     self.assertEqual(num, rn)
+
+  def test_fuzz_failure11(self):
+    v1=Variable("v1", 0, 16)
+    v2=Variable("v2", 0, 128)
+    v3=Variable("v3", 0, 5)
+    expr = UOp(Ops.MOD, dtypes.index, arg=None, src=(
+      UOp(Ops.ADD, dtypes.index, arg=None, src=(
+        UOp(Ops.MOD, dtypes.index, arg=None, src=(
+          UOp(Ops.ADD, dtypes.index, arg=None, src=(
+            UOp(Ops.MAX, dtypes.index, arg=None, src=(
+              UOp(Ops.MUL, dtypes.index, arg=None, src=(
+                x5:=UOp(Ops.DEFINE_VAR, dtypes.index, arg=('v2', 0, 128), src=()),
+                UOp(Ops.CONST, dtypes.index, arg=0, src=()),)),
+              UOp(Ops.CONST, dtypes.index, arg=8, src=()),)),
+            UOp(Ops.MUL, dtypes.index, arg=None, src=(
+              x5,
+              UOp(Ops.CONST, dtypes.index, arg=-2, src=()),)),)),
+          x10:=UOp(Ops.CONST, dtypes.index, arg=5, src=()),)),
+        UOp(Ops.ADD, dtypes.index, arg=None, src=(
+          UOp(Ops.ADD, dtypes.index, arg=None, src=(
+            UOp(Ops.IDIV, dtypes.index, arg=None, src=(
+              x14:=UOp(Ops.DEFINE_VAR, dtypes.index, arg=('v1', 0, 16), src=()),
+              UOp(Ops.CONST, dtypes.index, arg=6, src=()),)),
+            UOp(Ops.CONST, dtypes.index, arg=4, src=()),)),
+          UOp(Ops.ADD, dtypes.index, arg=None, src=(
+            x14,
+            UOp(Ops.CONST, dtypes.index, arg=1, src=()),)),)),)),
+      x10,))
+    v1_val, v2_val, v3_val = UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 7),UOp.const(dtypes.int, 0)
+    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
+    self.assertEqual(num, rn)
+
+if __name__ == '__main__':
+  unittest.main()

test/unit/test_tqdm.py

@@ -128,6 +128,7 @@ class TestProgressBar(unittest.TestCase):
         self._compare_bars(tinytqdm_output, tqdm_output)
         if n > 5: break
 
+  @unittest.skip("this is flaky")
   @patch('sys.stderr', new_callable=StringIO)
   @patch('shutil.get_terminal_size')
   def test_set_description(self, mock_terminal_size, mock_stderr):

tinygrad/codegen/opt/postrange.py

@@ -18,6 +18,7 @@ class Scheduler:
     self.ast, self.ren = ast, ren
     self.dont_use_locals = self.ast.arg.dont_use_locals if self.ast.arg is not None else False
     self.applied_opts = list(self.ast.arg.applied_opts) if self.ast.arg is not None else []
+    self.opt_range = itertools.count(start=max([x.arg[0] for x in self.rngs], default=0)+1)
 
   @property
   def rngs(self):
@@ -29,8 +30,6 @@ class Scheduler:
   def full_shape(self): return [ssimplify(x.src[0]) for x in self.rngs]
   @property
   def axis_types(self): return [x.arg[-1] for x in self.rngs]
-  @property
-  def maxarg(self): return max([x.arg[0] for x in self.rngs], default=0)
 
   # strings like ['g0', 'g1', 'l0', 'l1', 'l2', 'l3', 'l4', 'l5', 'R0', 'r0', 'r1', 'r2', 'u0', 'u1', 'u2']
   def shape_str(self) -> list[str]:
@@ -95,7 +94,7 @@ class Scheduler:
   def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng=None):
     if (old_sz:=rng.src[0].divides(amount)) is None:
       raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
-    new_rng = UOp.range(amount, self.maxarg+1, new_type) if input_new_rng is None else input_new_rng
+    new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
     replaced_rng = rng.replace(src=(UOp.const(dtypes.int, old_sz),))
     sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
     self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
@@ -231,9 +230,9 @@ class Scheduler:
       for tc in tensor_cores:
         if tc.dtype_in == in0.dtype.scalar() and tc.dtype_in == in1.dtype.scalar() and tc.dtype_out == reduceop.dtype.scalar():
           # tensor cores have three ranges. X, Y, and REDUCE
-          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: -x.arg[0])
-          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: -x.arg[0])
-          red_ranges = sorted(reduceop.src[1:], key=lambda x: -x.arg[0])
+          in0_ranges = sorted([u for u in in0.ranges if u not in in1.ranges], key=lambda x: x.arg[0], reverse=True)
+          in1_ranges = sorted([u for u in in1.ranges if u not in in0.ranges], key=lambda x: x.arg[0], reverse=True)
+          red_ranges = sorted(reduceop.src[1:], key=lambda x: x.arg[0], reverse=True)
           if DEBUG >= 3:
             print(f"TC({axis}): {[(x.arg[0],x.vmax+1) for x in in0_ranges]}",
                               f"{[(x.arg[0],x.vmax+1) for x in in1_ranges]} {[(x.arg[0],x.vmax+1) for x in red_ranges]}")

tinygrad/engine/schedule.py

@@ -90,7 +90,7 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[
       if rk.op is Ops.END: schedule.append(rk)
     else:
       raise RuntimeError(f"can't schedule {k.op}")
-    for x in children[k]:
+    for x in children[rk]:
       in_degree[x] -= 1
       if in_degree[x] == 0: queues[_heuristic(x)].append(x)
 

tinygrad/gradient.py

@@ -3,14 +3,15 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata
 from tinygrad.helpers import argsort
 
-def reduce_gradient(ctx:UOp, ret:UOp):
+def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   def broadcast_to_input(x): return x.reshape(x.shape+(1,)*(len(ret.src[0].shape)-len(x.shape))).expand(ret.src[0].shape)
-  if ret.arg[0] == Ops.ADD: return (broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MAX:
+  if op == Ops.ADD: return (broadcast_to_input(ctx),)
+  if op == Ops.MAX:
+    assert ret.op is Ops.REDUCE_AXIS, "only works on REDUCE_AXIS"
     mask = ret.src[0].eq(broadcast_to_input(ret)).cast(ctx.dtype)
     count = mask.r(Ops.ADD, ret.arg[1])
     return ((mask/broadcast_to_input(count)) * broadcast_to_input(ctx),)
-  if ret.arg[0] == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
+  if op == Ops.MUL: return (broadcast_to_input(ctx * ret) / ret.src[0],)
 
 # ctx is grad_output
 pm_gradient = PatternMatcher([
@@ -28,7 +29,8 @@ pm_gradient = PatternMatcher([
     ((x>y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)), (x<y).where(ctx, (x.eq(y)).where(ctx * 0.5, 0)))),
   (UPat(Ops.MUL, name="ret"), lambda ctx, ret: (ret.src[1]*ctx, ret.src[0]*ctx)),
   (UPat(Ops.WHERE, name="ret"), lambda ctx, ret: (None, ret.src[0].where(ctx, ctx.const_like(0)), ret.src[0].where(ctx.const_like(0), ctx))),
-  (UPat(Ops.REDUCE_AXIS, name="ret"), reduce_gradient),
+  (UPat(Ops.REDUCE_AXIS, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg[0])),
+  (UPat(Ops.REDUCE, name="ret"), lambda ctx, ret: reduce_gradient(ctx, ret, ret.arg) + (None,)*(len(ret.src)-1)),
   (UPat(Ops.CONTIGUOUS), lambda ctx: (ctx,)),
   (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)),
   (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)),
@@ -68,4 +70,8 @@ def compute_gradient(root:UOp, root_grad:UOp, targets:set[UOp]) -> dict[UOp, UOp
         # we add the backward metadata to everything new in the graph
         for bw_uop in v.toposort(lambda x: x not in (t0, *t0.src, grads[t0])):
           all_metadata[bw_uop] = all_metadata.get(bw_uop, ())+backward_metadata
+  # end any ranges on grads with a reduce sum
+  for k,v in grads.items():
+    if len(v.ranges):
+      grads[k] = v.reduce(*v.ranges, arg=Ops.ADD)
   return grads

tinygrad/runtime/autogen/am/smu_v14_0_2.py

@@ -476,7 +476,7 @@ PP_GRTAVFS_FW_SEP_FUSE_FREQUENCY_TO_COUNT_SCALER_4 = PP_GRTAVFS_FW_SEP_FUSE_e.de
 PP_GRTAVFS_FW_SEP_FUSE_COUNT = PP_GRTAVFS_FW_SEP_FUSE_e.define('PP_GRTAVFS_FW_SEP_FUSE_COUNT', 19)
 
 class SviTelemetryScale_t(Struct): pass
-int8_t = ctypes.c_char
+int8_t = ctypes.c_byte
 SviTelemetryScale_t._fields_ = [
   ('Offset', int8_t),
   ('Padding', uint8_t),

tinygrad/runtime/autogen/mesa.py

@@ -89,7 +89,7 @@ NIR_CMAT_C_SIGNED = nir_cmat_signed.define('NIR_CMAT_C_SIGNED', 4)
 NIR_CMAT_RESULT_SIGNED = nir_cmat_signed.define('NIR_CMAT_RESULT_SIGNED', 8)
 
 class nir_const_value(ctypes.Union): pass
-int8_t = ctypes.c_char
+int8_t = ctypes.c_byte
 uint8_t = ctypes.c_ubyte
 int16_t = ctypes.c_int16
 uint16_t = ctypes.c_uint16
@@ -3723,7 +3723,7 @@ struct__IO_FILE._fields_ = [
   ('_flags2', ctypes.c_int32),
   ('_old_offset', ctypes.c_int64),
   ('_cur_column', ctypes.c_uint16),
-  ('_vtable_offset', ctypes.c_char),
+  ('_vtable_offset', ctypes.c_byte),
   ('_shortbuf', (ctypes.c_char * 1)),
   ('_lock', ctypes.POINTER(_IO_lock_t)),
   ('_offset', ctypes.c_int64),

tinygrad/runtime/ops_amd.py

@@ -8,6 +8,7 @@ from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filte
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, DMAFdRef, BufferSpec, CompilerPairT
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, suppress_finalizing, lo32, hi32, colored, prod, ContextVar
+from tinygrad.helpers import VIZ
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, pci, sqtt
@@ -19,7 +20,7 @@ from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_so
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta, PCIDevice, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
 
-SQTT, SQTT_ITRACE_SE_MASK, PMC = ContextVar("SQTT", 0), ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("PMC", 0)
+SQTT, SQTT_ITRACE_SE_MASK, PMC = ContextVar("SQTT", VIZ.value>=2), ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("PMC", 0)
 EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=

tinygrad/runtime/support/autogen.py

@@ -103,7 +103,7 @@ def gen(dll, files, args=[], prolog=[], rules=[], epilog=[], recsym=False, use_e
     suggested_name = anon_names.get(f"{loc_file(loc(decl:=clang.clang_getTypeDeclaration(t)))}:{loc_line(loc(decl))}", suggested_name)
     nonlocal lines, types, anoncnt, objc
     tmap = {clang.CXType_Void:"None", clang.CXType_Char_U:"ctypes.c_ubyte", clang.CXType_UChar:"ctypes.c_ubyte", clang.CXType_Char_S:"ctypes.c_char",
-            clang.CXType_SChar:"ctypes.c_char",
+            clang.CXType_SChar:"ctypes.c_byte",
             **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{k.lower()}" for k in ["Bool", "WChar", "Float", "Double", "LongDouble"]},
             **{getattr(clang, f'CXType_{k}'):f"ctypes.c_{'u' if 'U' in k else ''}int{sz}" for sz,k in
                [(16, "UShort"), (16, "Short"), (32, "UInt"), (32, "Int"), (64, "ULong"), (64, "Long"), (64, "ULongLong"), (64, "LongLong")]}}

tinygrad/runtime/support/compiler_amd.py

@@ -103,7 +103,7 @@ class HIPCCCompiler(Compiler):
         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)
+                        "-O3", "-mllvm", "-amdgpu-internalize-symbols", "-c", "-o", libf.name, bcf.name] + self.extra_options, check=True)
 
         return pathlib.Path(libf.name).read_bytes()
   def disassemble(self, lib:bytes): amdgpu_disassemble(lib)

tinygrad/schedule/indexing.py

@@ -26,6 +26,8 @@ pm_generate_realize_map = PatternMatcher([
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize COPY/BUFFER_VIEW/CONTIGUOUS/STORE
   (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
+  # always realize REDUCE on outer ranges
+  (UPat(Ops.REDUCE, name="r"), lambda ctx,r: realize(ctx, r) if any(tr.arg[-1] == AxisType.OUTER for tr in r.src[1:]) else None),
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)
@@ -115,7 +117,7 @@ pm_apply_rangeify = PatternMatcher([
 
 # this is the definition of the movement ops
 @functools.cache
-def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
+def _apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
   match op:
     case Ops.SHRINK:  rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg))
     case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg))
@@ -143,6 +145,13 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
     case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 
+def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
+  # for PAD and RESHAPE, we replace the ranges with PLACEHOLDERS
+  if op not in (Ops.PAD, Ops.RESHAPE): return _apply_movement_op(op, in_shape, arg, rngs)
+  sink = UOp.sink(*rngs)
+  real_ranges = {r:UOp.range(r.src[0], i, AxisType.PLACEHOLDER) for i,r in enumerate(sink.ranges)}
+  return UOp.sink(*_apply_movement_op(op, in_shape, arg, sink.substitute(real_ranges).src)).substitute({v:k for k,v in real_ranges.items()}).src
+
 @profile_matches
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
   if debug: print("**************************")

tinygrad/schedule/rangeify.py

@@ -325,6 +325,18 @@ def bufferize_to_store(ctx:itertools.count|None, x:UOp, idx:UOp, allow_locals=Tr
     for m in mops[::-1]: ret = ret._mop(*m)
     return ret
 
+  # lower outerworld reduce here
+  if x.src[0].op is Ops.REDUCE and len(x.src[0].src) == 2 and x.src[0].src[1].arg[-1] == AxisType.OUTER:
+    assert sdtype.addrspace == AddrSpace.GLOBAL
+    outer_range = x.src[0].src[1]
+    buf = UOp.new_buffer(x.arg.device, size, x.dtype)
+    # NOTE: this has the same number as the outer range, we need string ranges!
+    zero_range = outer_range.replace(src=(UOp.const(dtypes.index, size),), arg=outer_range.arg[:-1]+(AxisType.LOOP,))
+    buf = buf.after(buf.index(zero_range).store(0).end(zero_range))
+    bufi = buf.index(idx, dtype=sdtype)
+    do_store = bufi.store(bufi.load() + x.src[0].src[0], tag=x.tag).end(*rngs).end(outer_range)
+    return buf.after(do_store)
+
   # NOTE: the DEFINE_LOCAL needs to be disambiguated here
   if sdtype.addrspace == AddrSpace.GLOBAL:
     buf = UOp.new_buffer(x.arg.device, size, x.dtype)
@@ -472,6 +484,7 @@ pm_add_range_tags = PatternMatcher([
 ])
 
 def split_store(ctx:list[UOp], x:UOp) -> UOp|None:
+  # if we have any outer ranges open here, we don't split
   if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
 
   # ends of outer range don't go in kernels
@@ -543,7 +556,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   # convert movement ops to ranges
   tsink, rctx = run_rangeify(tsink, DEBUG_RANGEIFY)
 
-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")  # this does const folding
+  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse")
   tsink = graph_rewrite(tsink, pm_remove_bufferize, bottom_up=True, name="remove bufferize with cost function")
   tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding, name="symbolic+reduce_collapse pt 2")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

tinygrad/uop/divandmod.py

@@ -0,0 +1,112 @@
+import functools
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
+
+# NOTE: this cache is only on index UOps and matches the cache in the old ShapeTracker in spirit
+@functools.cache
+def fold_divmod_general(d: UOp, correct_divmod_folding: bool) -> UOp|None:
+  x, y = d.src
+
+  # cancel_divmod: simple cancel div/mod case when the range of the numerator lies within a single denominator interval
+  x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
+  assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
+  if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
+  if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
+    return x - q*y if d.op is Ops.MOD else d.const_like(q)
+
+  # split uops for the rest of the processing
+  x_peeled, const = x.pop_const()
+  uops_no_const = list(x_peeled.split_uop(Ops.ADD))
+
+  # ** Constant Denominator Rules **
+  # these rules strictly require y to be a scalar constant > 0
+  if y.op is Ops.CONST and (c := y.arg) > 0:
+    # remove_nested_mod: remove nested mod in case the inner mod is a multiple of the outer mod, example: (a%4 + b)%2 -> (a+b)%2
+    if d.op is Ops.MOD and x.vmin >= 0:
+      new_xs, changed = [], False
+      for u in uops_no_const:
+        if u.op is Ops.MOD and u.src[1].divides(c) is not None:
+          u = u.src[0]
+          changed = True
+        new_xs.append(u)
+      if changed and (new_x:=(UOp.sum(*new_xs) + const)).vmin >= 0: return new_x % y
+
+    # Shared decomposition for folding rules
+    decomp = [(u.divides(f:=u.const_factor()),f) for u in uops_no_const]
+    terms, factors = zip(*decomp)
+
+    # fold_binary_numerator: fold if expression has one non-constant term that takes on two values
+    if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
+      y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
+      y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
+      return (y2-y1)*(v-v.vmin) + y1
+
+    # fold_divmod_congruence: fold if a is congruent to an expression whose range is between 0 and c
+    if not (x.vmin<0 and correct_divmod_folding):
+      rems = [min((r:=f%c), r-c, key=abs) for f in factors]
+      if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c==rem.vmax//c:
+        if d.op is Ops.MOD: return rem - rem.vmin//c*c
+        return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
+
+    # gcd_with_remainder: factor out common gcd from numerator
+    # Note: this rule uses uops_no_const to exclude the additive constant from the GCD calculation
+    if x.vmin >= 0:
+      gcd = UOp.gcd(*uops_no_const, y).simplify()
+      if gcd.op is Ops.CONST and gcd.arg > 1:
+        new_x = unwrap(x_peeled.divide_exact(gcd)).simplify() + (const%c)//gcd.arg
+        if new_x.vmin >= 0:
+          ret = new_x.alu(d.op, x.ufix(c//gcd.arg))
+          return ret*gcd + const%gcd.arg if d.op is Ops.MOD else ret+const//c
+
+    # nest_div_by_smallest_factor: try and nest the div and see if it allows the numerator to be simplified
+    if d.op is Ops.IDIV and x.vmin >= 0:
+      div = min([c] + [abs(f) for u, f in zip(uops_no_const, factors) if u.op not in (Ops.CONST, Ops.VCONST) and abs(f) > 1 and (c%f)==0])
+      # NOTE: this is recursive!
+      if div < c and (newxs := fold_divmod_general(x//div, correct_divmod_folding)) is not None and newxs.vmin >= 0:
+        return newxs // (c // div)
+
+  # ** Variable Denominator / Fallback Rules **
+  # These rules apply to variables OR constants that failed the checks above.
+  # Reconstruct all uops including const for these checks.
+  all_uops = uops_no_const + ([x.const_like(const)] if const != 0 else [])
+
+  # divide_by_gcd: x//y -> (x//gcd)//(y//gcd)
+  gcd = UOp.gcd(*all_uops, y).simplify()
+  if not (gcd.op is Ops.CONST and gcd.arg==1):
+    ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd)))
+    return ret*gcd if d.op is Ops.MOD else ret
+
+  # factor_remainder: (d*x+y)//d -> x+y//d
+  if y.vmin<0 or x.vmin<0: return None
+  quo, rem = [], []
+  for u in all_uops:
+    if (q:=u.divide_exact(y)) is not None: quo.append(q)
+    elif d.op is Ops.MOD and y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c:
+      rem.append(u.divides(c)*(c%y.arg))
+      quo.append(u.const_like(0))
+    else: rem.append(u)
+
+  if not quo: return None
+  new_x = sum(rem)+x.const_like(0)
+  if new_x.vmin<0: return None
+  return new_x%y if d.op is Ops.MOD else new_x//y+sum(quo)
+
+div_and_mod_symbolic = PatternMatcher([
+  # ** 1. Fast Inline Rules **
+  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
+    if c.vmin>0 and d.vmin>0 and ((x.vmin>=0 and a.vmin>=0) or (x.vmax<=0 and a.vmax<=0)) else None),  # (x//c+a)//d -> (x+a*c)//(c*d)
+  (UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
+  (UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax <= 0 else None),
+  ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
+    lambda x,c,n,d: ((x+c.arg%d.arg)//d + c.arg//d.arg) if c.arg%d.arg!=c.arg and x.vmin>=0 and n.vmin>=0 and d.arg>0 else None),
+  ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
+    lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
+
+  # ** 2. Slow Rules **
+  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d"), lambda d: fold_divmod_general(d, bool(CORRECT_DIVMOD_FOLDING))),
+
+  # NOTE: these have to go at the bottom or TestSymbolicOps.test_var loops
+  (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
+  (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax < 0 else None),
+])

tinygrad/uop/ops.py

@@ -14,12 +14,12 @@ if TYPE_CHECKING:
 class AxisType(Enum):
   def __repr__(self): return str(self)
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
-  THREAD = auto(); OUTER = auto() # noqa: E702
+  THREAD = auto(); OUTER = auto(); PLACEHOLDER = auto() # noqa: E702
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
-                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.OUTER: "O"}
+                AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r", AxisType.OUTER: "O", AxisType.PLACEHOLDER: "P"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
                AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta",
-               AxisType.OUTER: "green"}
+               AxisType.OUTER: "green", AxisType.PLACEHOLDER: "white"}
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
 axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
@@ -89,8 +89,8 @@ class UOpMetaClass(type):
     if SPEC > 1:
       from tinygrad.uop.spec import full_spec, test_pyrender
       if SPEC > 2: test_pyrender(created)
-      with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created)
-      if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}")
+      with Context(IGNORE_OOB=1): fret = cast(bool|None, full_spec.rewrite(created))
+      if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}")
     return created
 
 # some uops map to other stuff
@@ -583,7 +583,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
     return UOp(Ops.BUFFER, dtype, (UOp.unique(num), UOp(Ops.DEVICE, arg=device)), size)
   @property
-  def device(self) -> str|tuple[str, ...]: return cast(str|tuple[str, ...], unwrap(self._device))
+  def device(self) -> str|tuple[str, ...]: return unwrap(self._device)
   @recursive_property
   def _device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.DEVICE: return self.arg
@@ -615,6 +615,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
 
   def buf_target(self) -> UOp:
     # the buffer that's being loaded from or store to
+    # NOTE: this is the good one to keep
     match self.op:
       case Ops.DEFINE_GLOBAL | Ops.DEFINE_LOCAL | Ops.DEFINE_REG: return self
       case Ops.AFTER | Ops.INDEX | Ops.STORE | Ops.LOAD: return self.src[0].buf_target()
@@ -866,8 +867,8 @@ def print_uops(uops:list[UOp]):
 
 def get_location() -> tuple[str, int]:
   frm = sys._getframe(1)
-  # skip over ops.py/mathtraits.py (unless there's nothing but ops.py/mathtraits.py)
-  while pathlib.Path(frm.f_code.co_filename).name in ("ops.py", "mathtraits.py") and frm.f_back is not None and \
+  # skip over ops.py and anything in mixin
+  while ((codepath:=pathlib.Path(frm.f_code.co_filename)).name == "ops.py" or codepath.parent.name == "mixin") and frm.f_back is not None and \
       not frm.f_back.f_code.co_filename.startswith("<frozen"):
     frm = frm.f_back
   return frm.f_code.co_filename, frm.f_lineno
@@ -1077,20 +1078,22 @@ def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=Fal
 
 active_rewrites:list[TrackedGraphRewrite] = []
 def profile_matches(fxn:Callable):
-  def wrap(*args, **kwargs):
-    name = str(kwargs.get("name", None) or fxn.__name__)
-    assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}"
-    if tracking:=(TRACK_MATCH_STATS >= 2):
+  def wrap_profile_matches(*args, **kwargs):
+    if TRACK_MATCH_STATS >= 2:
+      name = str(kwargs.get("name", None) or fxn.__name__)
+      assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}"
       loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
       depth = len(active_rewrites)
       if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}"))
       tracked_ctxs[-1].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False)))
       active_rewrites.append(ctx)
-    with cpu_profile(name, "TINY", display=tracking):
-      ret = fxn(*args, **kwargs)
-    if tracking: active_rewrites.pop()
-    return ret
-  return wrap
+      with cpu_profile(name, "TINY"):
+        ret = fxn(*args, **kwargs)
+      active_rewrites.pop()
+      return ret
+    # without tracking, we just call the function
+    return fxn(*args, **kwargs)
+  return wrap_profile_matches
 
 class TrackedPatternMatcher(PatternMatcher):
   def rewrite(self, uop:UOp, ctx=None) -> UOp|None:
@@ -1164,12 +1167,12 @@ class RewriteContext:
 
   def cached_pm_rewrite(self, x:UOp):
     if (ret:=self.pm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
-    ret = self.pm_cache[x] = cast(PatternMatcher, self.pm).rewrite(x, self.ctx)
+    ret = self.pm_cache[x] = unwrap(self.pm).rewrite(x, self.ctx)
     return ret
 
   def cached_bpm_rewrite(self, x:UOp):
     if (ret:=self.bpm_cache.get(x,SENTINEL)) is not SENTINEL: return ret
-    ret = self.bpm_cache[x] = cast(PatternMatcher, self.bpm).rewrite(x, self.ctx)
+    ret = self.bpm_cache[x] = unwrap(self.bpm).rewrite(x, self.ctx)
     return ret
 
   def unified_rewrite(self, root:UOp) -> UOp:
@@ -1350,7 +1353,7 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.REDUCE_AXIS, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}.r({r.arg[0]}, {r.arg[1]})"),
   # NOTE: range has srcs sometimes after control flow
   (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
-    "UOp.range("+', '.join([str(c.arg)] + [str(y) for y in x.arg])+
+    "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+
       (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.index else '')+")"),
   # TODO: index shouldn't mismatch dtype
   (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:

tinygrad/uop/symbolic.py

@@ -2,9 +2,10 @@
 import math, operator, struct, functools
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import ConstType, dtypes, PtrDType, AddrSpace, can_safe_cast, Invalid
-from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, cdiv, cmod, CORRECT_DIVMOD_FOLDING, unwrap
+from tinygrad.dtype import ConstType, dtypes, PtrDType, can_safe_cast, Invalid
+from tinygrad.helpers import partition, all_same, prod, flatten, get_single_element, unwrap
 from tinygrad.uop.decompositions import xpow
+from tinygrad.uop.divandmod import div_and_mod_symbolic
 
 # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********
 
@@ -24,19 +25,16 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 
+# this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
 propagate_invalid = PatternMatcher([
-  # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
   # propagate invalid, push it past children
-  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype) if cast.dtype is not dtypes.index else None),
+  (invalid_gate.cast(name="cast"), lambda i,x,cond,cast: x.cast(cast.dtype)),
   *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: cond.where(x.alu(alu.op,y), i))
     for op in GroupOp.Binary-GroupOp.Comparison),
+  # TODO: when can this happen? and is it always safe to just drop invalid?
   *((invalid_gate.alu(op, UPat.var("y")).named("alu"), lambda cond,x,y,alu,i: x.alu(alu.op,y)) for op in GroupOp.Comparison),
-  # invalid + y -> y same for other ops
+  # invalid + y -> invalid same for other ops
   *((invalid_pat.alu(op, UPat(dtype=dtypes.index)).named("alu"), lambda alu,i: i) for op in GroupOp.Binary-GroupOp.Comparison),
-  # i < y -> a_bool_value_that_will_never_be_used: we choose a random bool const
-  *((invalid_pat.alu(op, UPat(dtype=dtypes.index)), lambda i: UOp.const(dtypes.bool, True)) for op in GroupOp.Comparison),
-  # a.where(b.where(c, d), d) -> (a & b).where(c, d)
-  (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
 ])
 
 symbolic_simple = propagate_invalid + PatternMatcher([
@@ -105,22 +103,17 @@ symbolic_simple = propagate_invalid + PatternMatcher([
   # positive const ** x
   (UPat.cvar("c", vec=False).alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None),
   # rules for threefry
-  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)&0xFFFFFFFF), # TODO: why is the and needed?
+  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),
   (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
   (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  # hacks for threefry long removal when padded (TODO: genericize)
-  (UPat.var('x', dtypes.uint32).cast(dtypes.uint64) * UPat.var('y').where(UPat.const(dtypes.uint64, 1<<32), UPat.const(dtypes.uint64, 0)),
-   lambda x,y: y.where(x, 0).cast(dtypes.uint64) * (1<<32)),
-  ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
-   lambda x,y: y.where(x.cast(dtypes.uint32), 0)),
-  # new decomp rules for threefry
   (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
   (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x),
-  (UPat.var('b').where(UPat.var('x', dtypes.uint32).cast(dtypes.uint64), UPat.const(dtypes.uint64, 0)).cast(dtypes.uint32), lambda b,x: b.where(x,0)),
   # ** simple where folding **
   # a conditional with the same results either way is a noop, also fold const conditionals
   (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
   (UPat.cvar("gate", vec=False).where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1),
+  # a.where(b.where(c, d), d) -> (a & b).where(c, d)
+  (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
 ])
 
 # ******** phase 2 builds on phase 1, it includes the old "symbolic", rules that match deeper ********
@@ -144,101 +137,6 @@ def canonicalize_simplex(X:UOp) -> UOp|None:
     ret.append(u)
   return UOp.sum(*ret) if changed else None
 
-def cancel_divmod(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # simple cancel div/mod case when the range of the numerator lies within a single denominator interval
-  x_min, x_max, y_min, y_max = x.vmin, x.vmax, y.vmin, y.vmax
-  assert isinstance(x_min, int) and isinstance(x_max, int) and isinstance(y_min, int) and isinstance(y_max, int)
-  if y_min==y_max==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.IDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
-  if y_min*y_max > 0 and (q:=cdiv(x_min,y_min)) == cdiv(x_min,y_max) == cdiv(x_max,y_min) == cdiv(x_max,y_max):
-    return x - q*y if d.op is Ops.MOD else d.const_like(q)
-  return None
-
-def remove_nested_mod(m: UOp, x: UOp, y: UOp) -> UOp|None:
-  # remove nested mod in case the inner mod is a multiple of the outer mod
-  # example: (a%4 + b)%2 -> (a+b)%2
-  if ((c := y.arg) < 0) or x.vmin<0: return None
-  new_xs = []
-  something_changed = False
-  for u in x.split_uop(Ops.ADD):
-    if u.op is Ops.MOD:
-      if u.src[1].divides(c) is not None:
-        something_changed = True
-        u = u.src[0]
-    new_xs.append(u)
-  new_x: UOp = UOp.sum(*new_xs)
-  if something_changed and new_x.vmin>=0: return new_x % y
-  return None
-
-def fold_binary_numerator(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # we can fold if the expression has only one non-constant term and this term can only take on two values
-  if ((c := y.arg) < 0): return None
-  x,const = x.pop_const()
-  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)])
-  if len(terms)==1 and (v:=terms[0]).vmax-v.vmin == 1:
-    y1 = cmod(factors[0]*v.vmin+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmin+const, c)
-    y2 = cmod(factors[0]*v.vmax+const, c) if d.op is Ops.MOD else cdiv(factors[0]*v.vmax+const, c)
-    return (y2-y1)*(v-v.vmin) + y1
-  return None
-
-def fold_divmod_congruence(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # within a mod we can freely subtract multiples of c, we use this to see if a is congruent to an expression whose vmin/vmax are between 0 and c
-  if (x.vmin<0 and CORRECT_DIVMOD_FOLDING) or ((c := y.arg) < 0): return None
-  x,const = x.pop_const()
-  terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in x.split_uop(Ops.ADD)])
-  # a//c = (a-a%c)/c, if we can fold a%c, we can fold a//c
-  rems = [min((r:=f%c), r-c, key=abs) for f in factors]
-  if (rem:=sum(r*v for r,v in zip(rems,terms))+const%c).vmin//c!=rem.vmax//c: return None
-  if d.op is Ops.MOD: return rem - rem.vmin//c*c
-  return sum((f-r)//c * v for f,r,v in zip(factors,rems,terms)) + (const-const%c+rem.vmin//c*c)//c
-
-def divide_by_gcd(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # x//y -> (x//gcd)//(y//gcd) or x%y -> gcd*(x//gcd)%(y//gcd)
-  gcd = UOp.gcd(*x.split_uop(Ops.ADD), y).simplify()
-  if gcd.op is Ops.CONST and gcd.arg==1: return None
-  ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd)))
-  return ret*gcd if d.op is Ops.MOD else ret
-
-def gcd_with_remainder(d: UOp, x: UOp, y: UOp):
-  # (gcd*x+r)//(gcd*d) -> (x+(r%d)//gcd)//d + r//(gcd*d)
-  # (gcd*x+r)%(gcd*d) -> gcd*(x+(r%d)//gcd)%d + r%gcd
-  # These only work for floordiv (and the corresponding remainder)! Thats why we check the sign of x,y and new_x
-  if ((c := y.arg) < 0) or x.vmin<0: return None
-  x_no_const, const = x.pop_const()
-  gcd = UOp.gcd(*x_no_const.split_uop(Ops.ADD), y).simplify()
-  assert gcd.op is Ops.CONST
-  if gcd.arg==1: return None
-  new_x = unwrap(x_no_const.divide_exact(gcd)).simplify() + (const%c)//gcd
-  if new_x.vmin<0: return None
-  ret = new_x.alu(d.op, x.ufix(c//gcd.arg))
-  return ret*gcd + const%gcd.arg if d.op is Ops.MOD else ret+const//c
-
-def factor_remainder(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # (d*x+y)//d -> x+y//d  or  (d*x+y)%d
-  # for mod we go further and take the remainder of all factors to reduce their size
-  # These only work for floordiv (and the corresponding remainder)! Thats why we check the sign of x,y and new_x
-  if y.vmin<0 or x.vmin<0: return None
-  quo, rem = [], []
-  for u in x.split_uop(Ops.ADD):
-    if (q:=u.divide_exact(y)) is not None: quo.append(q)
-    # if this is mod and y is a const, we can make the remainder factor sm
-    elif d.op is Ops.MOD and y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c:
-      rem.append(u.divides(c)*(c%y.arg))
-      quo.append(u.const_like(0))  # we append this so we can check if something changed
-    else: rem.append(u)
-  new_x = sum(rem)+x.const_like(0)
-  if len(quo)==0 or new_x.vmin<0: return None
-  return new_x%y if d.op is Ops.MOD else new_x//y+sum(quo)
-
-def nest_div_by_smallest_factor(d: UOp, x: UOp, y: UOp) -> UOp|None:
-  # we try and nest the div and see if it allows the numerator to be simplified
-  if ((c := y.arg) < 0): return None
-  factors = [u.const_factor() for u in x.split_uop(Ops.ADD) if u.op not in (Ops.CONST, Ops.VCONST)]
-  div = min([y.arg]+[abs(f) for f in factors if abs(f) > 1 and (c%f)==0])
-  newxs = fold_divmod_congruence(newx:=(x//div), x, y.const_like(div))
-  if newxs is None: newxs = factor_remainder(newx, x, y.const_like(div))
-  if div==y.arg or newxs is None or x.vmin<0 or newx.vmin<0: return None
-  return newxs//(c//div)
-
 def gep_through_wmma(gep:UOp, wmma:UOp):
   out_sz = prod(x[1] for x in wmma.arg[6][-1])
   wmma_idxs = gep.arg[::out_sz]
@@ -341,31 +239,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   # canonicalize a simplex with positive coefficients > 0
   # not x < 1 -> X > 0
   ((UPat.var("x", dtypes.index)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None),
-  # ** div **
-  # div folding
-  ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d)
-    if c.vmin>0 and d.vmin>0 and ((x.vmin>=0 and a.vmin>=0) or (x.vmax<=0 and a.vmax<=0)) else None),  # (x//c+a)//d -> (x+a*c)//(c*d)
   # a range mod its own upper bound is just the range
   (UPat(Ops.RANGE, src=UPat.var("end"), name="r")%UPat.var("end"), lambda r,end: r),
   (UPat(Ops.RANGE, src=UPat.var("end"), name="r")//UPat.var("end"), lambda r,end: r.const_like(0)),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), cancel_divmod),
-  (UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -(x//(-d)) if d.vmax < 0 else None),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_binary_numerator),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), fold_divmod_congruence),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), divide_by_gcd),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), gcd_with_remainder),
-  (UPat(Ops.MOD, dtypes.index, name="m", src=(UPat.var("x"), UPat.cvar("y", vec=False))), remove_nested_mod),
-  (UPat((Ops.IDIV), dtypes.index, name="d", src=(UPat.var("x"), UPat.cvar("y", vec=False))), nest_div_by_smallest_factor),
-  (UPat((Ops.IDIV, Ops.MOD), dtypes.index, name="d", src=(UPat.var("x"), UPat.var("y"))), factor_remainder),
-  (UPat.var("x", dtypes.index) // UPat.var("d"), lambda x,d: -((-x)//d) if x.vmax<=0 else None),
-  ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
-    lambda x,c,n,d: ((x+c.arg%d.arg)//d + c.arg//d.arg) if c.arg%d.arg!=c.arg and x.vmin>=0 and n.vmin>=0 and d.arg>0 else None),
-  ((UPat.var("x", dtypes.index)+UPat.cvar("c", vec=False)).named("n")//UPat.cvar("d", vec=False),
-    lambda x,c,n,d: (-(-(c.arg%d.arg + x - (d.arg-1))//d) + c.arg//d.arg) if x.vmax<=0 and n.vmin>=0 and d.arg>0 else None),
-  # ** mod **
-  # mod folding
-  (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: -((-x)%d) if x.vmax <= 0 else None),
-  (UPat.var("x", dtypes.index) % UPat.var("d"), lambda x,d: (x%(-d)) if d.vmax <  0 else None),
   # cast/long folding
   # if the intermediate cast doesnt narrow we can do it in one cast
   (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_safe_cast(x.dtype, a.dtype) else None),
@@ -382,7 +258,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
   # VECTORIZE/CONST
   (UPat(Ops.VECTORIZE, src=UPat(Ops.CONST), name="vec"), lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src))),
-])+gep_pushing
+])+div_and_mod_symbolic+gep_pushing
 
 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
 
@@ -407,14 +283,10 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
     expr, is_upper, c = res
     bounds[expr][int(is_upper)] = c
 
-  # don't simplify any other gates, can lead to OOB, we substitute them back later
-  uop = uop.substitute((load_subs:={u: UOp(Ops.NOOP, dtype=u.dtype, arg=u) for u in uop.toposort() if u.op is Ops.INDEX}))
-
   # simplify uop given that valid is True
   all_candidates = []
   for i,(expr,v) in enumerate(bounds.items()):
     v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
-    expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
     # try checking the whole clause
     all_candidates.append((expr, UOp.variable(f"fake{i}", v0, v1, expr.dtype)))
 
@@ -438,8 +310,6 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
   # try all the valids together (but only the whole expressions)
   if (s_uop:=uop.substitute(sub_dict:=dict(all_candidates))) is not uop:
     uop = s_uop.simplify().substitute({newX:X for X,newX in sub_dict.items()}).simplify(full_symbolic=False)
-  # put the loads back in
-  uop = uop.substitute({v:k for k,v in load_subs.items()})
   return uop
 
 def _valid_priority(v: UOp, valids:list[UOp]):
@@ -456,7 +326,7 @@ def simplify_valid(valid:UOp) -> UOp|None:
     if ret[-1] is not stmt: something_changed = True
   return UOp.prod(*ret) if something_changed else None
 
-# ******** phase 3 is the complete symbolic, and deals with very complex things like loop rewriting and threefry transform ********
+# ******** phase 3 is the complete symbolic ********
 
 def reduce_mul_chain(r:UOp):
   if r.arg not in {Ops.ADD, Ops.MAX}: return None
@@ -485,6 +355,8 @@ def where_on_load(c1, buf, x):
   # aditionally we can drop the clause on the where if it already exists in the load
   remaining_clause = UOp.const(dtypes.bool, True).prod(*[c for c in c1.split_uop(Ops.AND) if c not in removed])
   return remaining_clause.where(buf.index(x.get_idx().valid(functools.reduce(operator.and_, moved_clauses, c2))), 0)
+
+# where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
 pm_move_where_on_load = PatternMatcher([
   (UPat.var("c1").where(UPat.var("buf").index(UPat.var("x")), 0), where_on_load),
   (UPat.var("c1").where(0, UPat.var("buf").index(UPat.var("x"))), lambda c1,buf,x: where_on_load(c1.logical_not(),buf,x)),
@@ -500,21 +372,9 @@ pm_simplify_valid = PatternMatcher([
 # this is symbolic 2.0
 REMOVE_FROM_SINK_LIKE = {Ops.UNROLL, Ops.NOOP, Ops.VECTORIZE, Ops.SINK}
 sym = symbolic+pm_simplify_valid+PatternMatcher([
-  # LOAD/STORE -> NOOP
-  (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]),
-  (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c),
-  # VECTORIZE/GEP
-  (UPat(Ops.VECTORIZE, src=UPat(Ops.GEP, src=(UPat.var("x"),)), name="vec"), lambda vec,x: x.gep(tuple(y.arg[0] for y in vec.src))),
   # reorder ALU/VECTORIZE
   (UPat(GroupOp.ALU, src=(UPat(Ops.VECTORIZE, src=UPat(name='x')), UPat(Ops.VECTORIZE, src=UPat(name='y'))), name='alu'),
    lambda x,y,alu: UOp(Ops.VECTORIZE, alu.dtype, (UOp(alu.op, alu.dtype.scalar(), (x,y)),)*alu.dtype.count)),
-  # VECTORIZE of a single element is just that element
-  (UPat(Ops.VECTORIZE, src=(UPat(name='x'),)), lambda x: x),
-  # VECTORIZE void is GROUP
-  (UPat(Ops.VECTORIZE, dtype=dtypes.void, name='x'), lambda x: UOp.group(*x.src)),
-  # tensor core with a 0 input is acc
-  (UPat(Ops.WMMA, src=(UPat.const(None, 0.0), UPat.var(), UPat.var("acc"))), lambda acc: acc),
-  (UPat(Ops.WMMA, src=(UPat.var(), UPat.const(None, 0.0), UPat.var("acc"))), lambda acc: acc),
   # ** self folding **
   # x!=0 -> (bool)x
   (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
@@ -531,7 +391,6 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([
   # fold gated LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat.const(dtypes.index, Invalid)).or_casted(),), allow_any_len=True, name="x"),
     lambda x: UOp(Ops.NOOP) if x.op is Ops.STORE else x.const_like(0)), # invalid store does nothing. invalid load produces 0
-  # # Where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer
   ((UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()),  # 1/(x^c) -> (1/x)^c
   ((UPat.var("x") * UPat.var("x") * UPat.var("x")).reciprocal(), lambda x: x.reciprocal()*x.reciprocal()*x.reciprocal()),
   ((UPat.var("x") * UPat.cvar("c")).reciprocal(), lambda x,c: x.reciprocal()*c.reciprocal()), # 1/(x*c) -> (1/c)*(1/x)

tinygrad/viz/index.html

@@ -271,7 +271,7 @@
   }
   #device-list > div {
     min-height: 32px;
-    width: 132px;
+    width: 134px;
     overflow-x: auto;
     overflow-y: hidden;
     white-space: nowrap;

tinygrad/viz/js/index.js

@@ -36,7 +36,7 @@ const updateProgress = ({ start, err }) => {
   d3.select("#custom").html("");
   if (err) {
     displaySelection("#custom");
-    d3.select("#custom").append(() => d3.create("div").classed("raw-text", true).call(s => s.append(() => codeBlock(err, "txt"))).node());
+    d3.select("#custom").append("div").classed("raw-text", true).call(s => s.append(() => codeBlock(err, "txt"))).node();
   }
 }
 
@@ -198,6 +198,8 @@ function focusShape(shape) {
   return metadata.replaceChildren(shapeMetadata.get(focusedShape) ?? "");
 }
 
+const EventTypes = { EXEC:0, BUF:1 };
+
 async function renderProfiler(path, unit) {
   displaySelection("#profiler");
   metadata.replaceChildren(shapeMetadata.get(focusedShape) ?? "");
@@ -238,12 +240,11 @@ async function renderProfiler(path, unit) {
     const { y:baseY, height:baseHeight } = rect(div.node());
     const offsetY = baseY-canvasTop+padding/2;
     const shapes = [], visible = [];
-    const EventTypes = {TIMELINE:0, MEMORY:1};
     const eventType = u8(), eventsLen = u32();
-    if (eventType === EventTypes.TIMELINE) {
+    if (eventType === EventTypes.EXEC) {
       const levelHeight = baseHeight-padding;
       const levels = [];
-      data.tracks.set(k, { shapes, visible, offsetY, pcolor:"#9ea2ad" });
+      data.tracks.set(k, { shapes, eventType, visible, offsetY, pcolor:"#9ea2ad" });
       let colorKey, ref;
       for (let j=0; j<eventsLen; j++) {
         const e = {name:strings[u32()], ref:optional(u32()), key:optional(u32()), st:u32(), dur:f32(), info:strings[u32()] || null};
@@ -366,7 +367,8 @@ async function renderProfiler(path, unit) {
         sum.x.push(allX[i], allX[i+1]);
         const y = maxY.get(allX[i]); sum.y1.push(y, y); sum.y0.push(base0, base0);
       }
-      data.tracks.set(k, { shapes:[sum], visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height, views:[[sum], shapes], valueMap });
+      data.tracks.set(k, { shapes:[sum], eventType, visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height,
+                           views:[[sum], shapes], valueMap });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
@@ -396,11 +398,11 @@ async function renderProfiler(path, unit) {
     xscale.domain(visibleX);
     // draw shapes
     const paths = [];
-    for (const [_, { offsetY, shapes, visible, valueMap, pcolor }] of data.tracks) {
+    for (const [_, { shapes, eventType, visible, offsetY, valueMap, pcolor }] of data.tracks) {
       visible.length = 0;
       for (const e of shapes) {
         const p = new Path2D();
-        if (e.width == null) { // generic polygon
+        if (eventType === EventTypes.BUF) { // generic polygon
           if (e.x[0]>et || e.x.at(-1)<st) continue;
           const x = e.x.map(xscale);
           p.moveTo(x[0], offsetY+e.y0[0]);