tinygrad 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

tinygrad/runtime/ops_python.py

@@ -0,0 +1,204 @@
+# a python uops emulator
+# works to test the tensor cores, and all the uops in general
+# this is the (living) definition of uops
+from typing import Tuple, List, Optional, Any, Dict
+import pickle, base64, itertools, time, struct
+from tinygrad.dtype import DType, dtypes, ImageDType
+from tinygrad.helpers import all_same, getenv, flatten
+from tinygrad.device import Compiled, Compiler, Allocator
+from tinygrad.codegen.uops import UOpGraph, UOps
+from tinygrad.ops import BinaryOps, TernaryOps, exec_alu
+from tinygrad.renderer import Renderer
+from tinygrad.renderer.cstyle import CUDARenderer, MetalRenderer, HIPRenderer
+
+def _load(m, i):
+  if i < 0 or i >= len(m): raise IndexError(f"load out of bounds, size is {len(m)} and access is {i}")
+  return m[i]
+
+def load(inp, j=0):
+  if len(inp) == 4: return [_load(m, x+j) if gate else default for m,x,gate,default in zip(*inp)]
+  else: return [_load(m, x+j) for m,x in zip(inp[0], inp[1])]
+
+def _store(m, i, v):
+  if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
+  m[i] = v
+
+class PythonProgram:
+  def __init__(self, name:str, lib:bytes):
+    self.uops: List[Tuple[UOps, Optional[DType], List[int], Any]] = pickle.loads(lib)
+  def __call__(self, *bufs, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
+    st = time.perf_counter()
+    warp = list(itertools.product(*[range(x) for x in local_size[::-1]]))
+    warp_size = len(warp)
+    for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
+      ul: Dict[int, Any] = {}
+      dl: Dict[int, DType] = {}
+      pbufs: List[memoryview] = list(bufs)
+      pvals: List[int] = list(vals)
+      i = 0
+      loop_ends: Dict[int, int] = {}
+      while i < len(self.uops):
+        uop, dtype, idp, arg = self.uops[i]
+        void_ops = {UOps.STORE, UOps.ENDRANGE, UOps.BARRIER, UOps.IF, UOps.ENDIF}
+        if uop is UOps.DEFINE_ACC: idp.clear()
+        inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops]
+        dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
+        if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp)
+        if uop is UOps.STORE:
+          if len(inp) == 3: inp.append([True] * len(inp[0]))  # set the gate to True
+          if isinstance(dtp[0], ImageDType):
+            # image store
+            assert dtp[2].count == 4
+            for j,val in enumerate(inp[2]):
+              for m,ox,oy,v,g in zip(inp[0], inp[1][0], inp[1][1], val, inp[3]):
+                assert ox >= 0 and ox < dtp[0].shape[1] and oy >= 0 and oy < dtp[0].shape[0]
+                if g: _store(m, ox*4 + oy*dtp[0].shape[1]*4 + j, v)
+          elif dtp[2].count > 1:
+            for j,val in enumerate(inp[2]):
+              for m,o,v,g in zip(inp[0], inp[1], val, inp[3]):
+                if g: _store(m, o+j, v)
+          else:
+            for m,o,v,g in zip(*inp):
+              if g: _store(m, o, v)
+          i += 1
+          continue
+        elif uop is UOps.ENDRANGE:
+          loop_ends[idp[0]] = i
+          i = idp[0]
+          continue
+        elif uop in (UOps.BARRIER, UOps.IF, UOps.ENDIF):
+          # in the python emulator, the warp is always in sync
+          i += 1
+          continue
+        assert dtype is not None, f"{uop} is missing a dtype"
+        dl[i] = dtype
+        if uop is UOps.DEFINE_GLOBAL:
+          assert dtype.fmt is not None
+          ul[i] = [pbufs.pop(0).cast(dtype.fmt)] * warp_size
+        elif uop is UOps.DEFINE_LOCAL:
+          assert dtype.fmt is not None
+          lbuf = memoryview(bytearray(arg[1]*dtype.itemsize))
+          ul[i] = [lbuf.cast(dtype.fmt)] * warp_size
+        elif uop is UOps.DEFINE_VAR:
+          ul[i] = [pvals.pop(0)] * warp_size
+        elif uop is UOps.SPECIAL:
+          if arg[1][0] == 'g':
+            ul[i] = [idxs[2-arg[0]]] * warp_size
+          elif arg[1][0] == 'l':
+            ul[i] = [x[2-arg[0]] for x in warp]
+        elif uop is UOps.CONST:
+          ul[i] = [[arg] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [arg] * warp_size
+        elif uop is UOps.DEFINE_ACC:
+          ul[i] = [[arg[0]] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [arg[0]] * warp_size
+        elif uop is UOps.RANGE:
+          if i not in ul: ul[i] = [inp[0][0]] * warp_size
+          else:
+            for j in range(len(ul[i])):
+              ul[i][j] += 1
+            if ul[i][0] == inp[1][0]:
+              del ul[i]
+              i = loop_ends[i] + 1
+              continue
+        elif uop in {UOps.CAST, UOps.BITCAST}:
+          if dtype.count > 1: ul[i] = inp
+          else:
+            assert dtp[0].fmt and dtype.fmt
+            pack_format, unpack_format = str(warp_size) + dtp[0].fmt, str(warp_size) + dtype.fmt
+            if uop is UOps.BITCAST: ul[i] = list(struct.unpack(unpack_format, struct.pack(pack_format, *inp[0])))
+            else:
+              casted = [dtypes.as_const(x, dtype) for x in inp[0]]
+              overflow_adjust = 2**(dtype.itemsize*8 - 1) if (dtypes.is_int(dtype) and not dtypes.is_unsigned(dtype)) else 0
+              overflow_fixed = [((x + overflow_adjust) % 2**(dtype.itemsize*8) - overflow_adjust) if dtypes.is_int(dtype) else x for x in casted]
+              ul[i] = list(struct.unpack(unpack_format, struct.pack(unpack_format, *overflow_fixed)))
+        elif uop is UOps.LOAD:
+          if isinstance(dtp[0], ImageDType):
+            assert dtype.count == 4
+            ul[i] = []
+            for j in range(dtype.count):
+              ret = []
+              for m,ox,oy in zip(inp[0], inp[1][0], inp[1][1]):
+                if ox < 0 or ox >= dtp[0].shape[1] or oy < 0 or oy >= dtp[0].shape[0]: ret.append(0)
+                else: ret.append(_load(m, ox*4 + oy*dtp[0].shape[1]*4 + j))
+              ul[i].append(ret)
+          elif dtype.count > 1:
+            ul[i] = [load([inp[i][j] if dtp[i].count > 1 else inp[i] for i in range(len(inp))], j) for j in range(dtype.count)]
+          else:
+            ul[i] = load(inp)
+        elif uop is UOps.PHI:
+          for j in range(len(inp[0])): inp[0][j] = inp[1][j]
+          ul[i] = inp[0]
+        elif uop is UOps.GEP:
+          ul[i] = inp[0][arg]
+        elif uop is UOps.WMMA:
+          # here are the models for the WMMA instruction on the different hardware
+          def wmma_helper(WARP_THREADS, K, NUM_A, NUM_B, NUM_C, a_elem, b_elem, c_map):
+            assert len(inp[0]) == NUM_A, f"A must have {NUM_A} elements per thread"
+            assert len(inp[1]) == NUM_B, f"B must have {NUM_B} elements per thread"
+            assert len(inp[2]) == NUM_C, f"C must have {NUM_C} elements per thread"
+            assert len(flatten(inp[0])) == NUM_A * warp_size, f"WMMA must have {NUM_A * warp_size} total elements for A in WMMA"
+            assert len(flatten(inp[1])) == NUM_B * warp_size, f"WMMA must have {NUM_B * warp_size} total elements for B in WMMA"
+            assert len(flatten(inp[2])) == NUM_C * warp_size, f"WMMA must have {NUM_C * warp_size} total elements for C in WMMA"
+            assert warp_size > 0 and warp_size % WARP_THREADS == 0, f"must have multiples of {WARP_THREADS} warp threads"
+            out = [inp[2][elem_idx][:] for elem_idx in range(NUM_C)]
+            for goff in range(0, warp_size, WARP_THREADS):
+              for lane_id in range(WARP_THREADS):
+                for elem_idx in range(NUM_C): # calculate new muls and add to acc
+                  (c_i, c_j) = c_map(lane_id, elem_idx)
+                  out[elem_idx][goff+lane_id] += sum(a_elem(inp[0], _k, c_j, goff) * b_elem(inp[1], c_i, _k, goff) for _k in range(K))
+            return out
+
+          # TODO: refactor these to a shared TensorCoreLayout in kernel.py
+          if arg[5] == "METAL":
+            # A (2 elements on 32 threads): row major
+            def a_b_elem(x, i, j, goff): return x[(i%2)][goff+(i//2)%2+(j%4)*2+(i//4)*8+(j//4)*16]
+            # (i, j), C, D (2 elements on 32 threads): row major same as A/B
+            def c_map(lane, elem): return (elem + ((lane%2)*2) + ((lane//8)%2)*4, ((lane//2)%4) + (lane//16)*4)
+            ul[i] = wmma_helper(32, 8, 2, 2, 2, a_b_elem, a_b_elem, c_map)
+          elif arg[5] == "HSA":
+            # A (16 elements on 32 threads): col major, lane 16-32 == lane 0-15
+            def a_elem(x, i, j, goff):
+              assert x[i][goff+j] == x[i][goff+j+16], "warp elements not duplicated properly across lanes"
+              return x[i][goff+j]
+            # B (16 elements on 32 threads): row major, lane 16-32 == lane 0-15
+            def b_elem(x, i, j, goff): return a_elem(x, j, i, goff)
+            def c_map(lane, elem): return (lane%16, lane//16+elem*2) # (i, j), C, D (8 elements on 32 threads): row major
+            ul[i] = wmma_helper(32, 16, 16, 16, 8, a_elem, b_elem, c_map)
+          elif arg[5] == "CUDA":
+            # A (8 elements on 32 threads)
+            def a_elem(x, i, j, goff): return x[(i%2)+(j//8)*2+(i//8)*4][goff+((i//2)%4)+(j%8)*4]
+            # B (4 elements on 32 threads)
+            def b_elem(x, i, j, goff): return x[(j%2)+(j//8)*2][goff+(j//2)%4+(i)*4]
+            # (i, j), C, D (4 elements on 32 threads)
+            def c_map(lane, elem): return ((elem%2)+(lane%4)*2, (lane//4)+(elem//2)*8)
+            ul[i] = wmma_helper(32, 16, 8, 4, 4, a_elem, b_elem, c_map)
+          else: raise NotImplementedError(f"unimplemented tensor core {arg}")
+        elif uop is UOps.ALU:
+          assert all_same([len(x) for x in inp]), f"{[len(x) for x in inp]} doesn't match on {arg}"
+          assert all_same([dtype] + dtp) or arg in {BinaryOps.CMPEQ, BinaryOps.CMPLT, TernaryOps.WHERE}, f"dtype mismatch on {arg}"
+          ul[i] = [exec_alu(arg, dtype, p) for p in zip(*inp)]
+        assert i in ul, (uop, dtype, idp, arg)
+        i += 1
+    return time.perf_counter() - st
+
+class PythonRenderer(Renderer):
+  device = "PYTHON"
+  def __init__(self):
+    if getenv("EMULATE_METAL"): self.device, self.tensor_cores = "METAL", MetalRenderer.tensor_cores
+    if getenv("EMULATE_HSA"): self.device, self.tensor_cores = "HSA", HIPRenderer.tensor_cores
+    if getenv("EMULATE_CUDA"): self.device, self.tensor_cores = "CUDA", CUDARenderer.tensor_cores
+
+  def render(self, name:str, uops:UOpGraph) -> str:
+    lops = [(u.uop, u.dtype, [uops.uops.index(v) for v in u.vin], u.arg) for u in uops]
+    return base64.b64encode(pickle.dumps(lops)).decode()
+
+class PythonCompiler(Compiler):
+  def compile(self, src:str) -> bytes: return base64.b64decode(src)
+
+class PythonAllocator(Allocator):
+  def _alloc(self, size, options): return memoryview(bytearray(size))
+  def copyin(self, dest, src:memoryview): dest[:] = src
+  def copyout(self, dest:memoryview, src): dest[:] = src
+
+class PythonDevice(Compiled):
+  def __init__(self, device:str):
+    super().__init__(device, PythonAllocator(), PythonRenderer(), PythonCompiler(), PythonProgram)

tinygrad/shape/shapetracker.py

@@ -1,286 +1,118 @@
 # ShapeTracker allows movement operations to a buffer that don't require a copy to be made.
 from __future__ import annotations
-import functools
-from typing import Dict, Tuple, Union, List, Optional, cast, NamedTuple
-from tinygrad.helpers import prod, DEBUG, partition
-from tinygrad.shape.symbolic import Variable, MulNode, NumNode, Node, SumNode, is_sym_int
-
-@functools.lru_cache(maxsize=None)
-def to_shape_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> Tuple[Tuple[int, int], ...]:
-  assert len(shape) == len(strides)
-  ret = [(shape[0], strides[0])] if shape else []
-  for i in range(1, len(shape)):
-    if ret[-1][1] == shape[i]*strides[i] or ret[-1][0] == 1:
-      ret[-1] = (ret[-1][0] * shape[i], strides[i])
-    elif shape[i] == 1:
-      continue
-    else:
-      ret.append((shape[i], strides[i]))
-  return tuple(ret)
-
-@functools.lru_cache(maxsize=None)
-def is_contiguous(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> bool: return all(s1 == s2 or s == 1 for s,s1,s2 in zip(shape, strides, strides_for_shape(shape)))
-
-@functools.lru_cache(maxsize=None)
-def filter_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> Tuple[int, ...]:
-  return tuple(stride if shp != 1 else 0 for stride, shp in zip(strides, shape))
-
-class ViewInternal(NamedTuple):
-  shape:Tuple[int, ...]
-  strides:Tuple[int, ...]
-  offset:int
-  mask:Optional[Tuple[Tuple[int, int]]]
-  contiguous:bool
-  shape_strides:Tuple[Tuple[int, int], ...]
-
-@functools.lru_cache(maxsize=None)
-class View(ViewInternal):
-  def __new__(cls, shape, strides=None, offset=0, mask=None):
-    strides_from_shape = strides_for_shape(shape)
-    strides = strides_from_shape if not strides else filter_strides(shape, strides)
-    contiguous = offset == 0 and is_contiguous(shape, strides) and mask is None
-    return super().__new__(cls, shape, strides, offset, mask, contiguous, to_shape_strides(shape, strides))
-  def __init__(self, shape, strides=None, offset=0, mask=None, contiguous=False, shape_strides=()): super().__init__()
-
-  def expr_node_mask(self, idx, valid=None) -> Node:
-    expr = [valid] if valid is not None else []
-    if self.mask is not None:
-      acc = 1
-      for ns,(x,y) in reversed(list(zip(self.shape, self.mask))):
-        base = ((idx//acc) % ns)
-        expr += [base >= x, base < y]
-        acc *= ns
-    return Variable.ands(expr)
-
-  # generate an expression if you have a single idx variable
-  def expr_node(self, idx=None) -> Node:
-    if idx is None: idx = Variable('idx', 0, prod(self.shape)-1)
-    ret: List[Node] = [Variable.num(self.offset) if isinstance(self.offset, int) else self.offset] if self.offset else []
-    acc = 1
-    for d,s in reversed(self.shape_strides):
-      ret.append(((idx//acc)%d)*s)
-      acc *= d
-    return Variable.sum(ret)
-
-  # generate an expression if you have a variable or expression for each index
-  def expr_idxs(self, idxs) -> Node:
-    assert len(idxs) == len(self.shape), f"need an idx for all dimensions {idxs} vs {self.shape}"
-    return Variable.sum([Variable.num(self.offset) if isinstance(self.offset, int) else self.offset] + [idx*st for idx,sh,st in zip(idxs, self.shape, self.strides) if sh != 1 and st != 0])
-
-@functools.lru_cache(maxsize=None)
-def idxs_to_idx(shape:Tuple[int, ...], idxs) -> Node:
-  assert len(idxs) == len(shape), "need an idx for all dimensions"
-  acc = 1
-  ret = []
-  for tidx,d in reversed(list(zip(idxs, shape))):
-    ret.append(tidx * acc)
-    acc *= d
-  return Variable.sum(ret)
-
-@functools.lru_cache(maxsize=None)
-def strides_for_shape(shape:Tuple[int, ...]) -> Tuple[int, ...]:
-  strides = [1] if shape else []
-  for d in shape[::-1][:-1]: strides = [d*strides[0]] + strides
-  return filter_strides(shape, tuple(strides))
-
-@functools.lru_cache(maxsize=None)
-def merge_views(vm2:View, vm1:View) -> Optional[View]:
-  if vm2.mask: return None  # this isn't supported yet
-  mst = ShapeTracker(vm1.shape, [vm2, vm1])
-  strides = mst.real_strides()
-  if None in strides: return None
-  return View(vm1.shape, strides, mst.real_offset(), vm1.mask)
-
-@functools.lru_cache(maxsize=None)
-def _reshape(view: View, new_shape:Tuple[int, ...]) -> Tuple[View, bool]:
-  shape, mask, strides, offset = view.shape, view.mask, view.strides, view.offset
-  # check if this is adding or removing 1s (only)
-  # NOTE: this is optional, but removes most calls to (expensive!) merge_views (with mask, not optional)
-  if [x for x in shape if x != 1] == [x for x in new_shape if x != 1]:
-    new_strides: List[int] = [y for x,y in zip(shape, strides) if x != 1]
-    new_strides_tuple: Tuple[int, ...] = tuple([0 if x == 1 else new_strides.pop(0) for x in new_shape])
-    new_mask_tuple = None
-    if mask:
-      for x,y in zip(shape, mask):
-        if x == 1 and y != (0, 1):
-          new_mask_tuple = ((0,0),) * len(new_shape)
-          break
-      else:
-        new_mask: List[Tuple[int, int]] = [y for x,y in zip(shape, mask) if x != 1]
-        new_mask_tuple = tuple([(0,1) if x == 1 else new_mask.pop(0) for x in new_shape])
-    return View(new_shape, new_strides_tuple, offset, new_mask_tuple), False
+from dataclasses import dataclass
+from typing import Tuple, List, Optional, Dict, Set, Iterable, cast
+from tinygrad.helpers import merge_dicts, getenv
+from tinygrad.shape.symbolic import Variable, MulNode, Node, SumNode, NumNode, create_lt_node, create_ge_node, sint
+from tinygrad.shape.view import View, strides_for_shape
+
+def _expr_view(view:View, idxs:List[Node], valid:Optional[Node]=None) -> Tuple[Node, Node]:
+  assert len(idxs) == len(view.shape), f"need an idx for all dimensions {idxs} vs {view.shape}"
+  iexpr: List[Node] = [NumNode(view.offset) if isinstance(view.offset, int) else view.offset]
+  vexpr: List[Node] = [valid] if valid is not None else []
+  for idx,sh,st,m in zip(idxs, view.shape, view.strides, view.mask if view.mask is not None else [None]*len(view.shape)):
+    if sh != 1 and st != 0: iexpr.append(idx*st)
+    if m is not None: vexpr += [create_ge_node(idx, m[0]), create_lt_node(idx, m[1])]  # idx >= m[0], idx < m[1]
+  return Node.sum(iexpr), Node.ands(vexpr)
+
+@dataclass(frozen=True)
+class ShapeTracker:
+  views: Tuple[View, ...]
 
-  new_view = View(new_shape)
-  if view.contiguous: return new_view, False # NOTE: if it's contiguous it can't have an offset
-  if (merged_view := merge_views(view, new_view)) is not None: return merged_view, False
-  if DEBUG >= 4: print(f"WARNING: creating new view with reshape {view} -> {new_shape}")
-  return new_view, True
+  def __add__(self, st:ShapeTracker) -> ShapeTracker:
+    ret = self
+    for v in st.views: ret = ShapeTracker(ret.views + (v,)).simplify() # one view at a time = better simplification
+    return ret
 
-@functools.lru_cache(maxsize=None)
-def get_pad_args(shape:Tuple[int,...], arg:Tuple[Tuple[int, int], ...]):
-  return tuple([(-b,s+e) for s,(b,e) in zip(shape, arg)]), tuple([(b,s+b) for s,(b,_) in zip(shape, arg)])
+  def invert(self, out_shape:Tuple[sint, ...]) -> Optional[ShapeTracker]:
+    ret = tuple(v.invert(s) for v,s in zip(self.views[::-1], [x.shape for x in self.views[::-1][1:]]+[out_shape]))
+    return ShapeTracker(cast(Tuple[View, ...], ret)).reshape(out_shape) if all(x is not None for x in ret) else None
 
-@functools.lru_cache(maxsize=None)
-def get_unsafe_resize_offset(strides, arg):
-  return sum([s * x[0] for s, x in zip(strides,arg)])
-
-class ShapeTracker:
-  __slots__ = "views", "var_vals"
-  def __init__(self, shape:Union[ShapeTracker, Tuple[Union[Node,int], ...]], views:Optional[List[View]]=None):
-    self.views: List[View] = views if views is not None else ([*cast(ShapeTracker, shape).views] if shape.__class__ is ShapeTracker else [View(shape)])
-    self.var_vals: Dict[Variable, int] = shape.var_vals if isinstance(shape, ShapeTracker) else {}
-  def __repr__(self): return f"ShapeTracker(shape={self.views[-1].shape}, views={self.views}, var_vals={self.var_vals})"
-  def copy(self) -> ShapeTracker: return ShapeTracker(self.views[-1].shape, [*self.views])
+  @staticmethod
+  def from_shape(shape:Tuple[sint, ...]): return ShapeTracker((View.create(shape),))
 
   @property
   def contiguous(self) -> bool: return len(self.views) == 1 and self.views[0].contiguous
 
   @property
-  def shape(self) -> Tuple[int, ...]: return self.views[-1].shape # NOTE: real type is Tuple[Union[Node, int], ...] but mypy complains about prod(shape)
+  def consecutive(self) -> bool: return len(self.views) == 1 and (v:=self.views[0]).mask is None and v.strides == strides_for_shape(v.shape)
+
+  @property
+  def shape(self) -> Tuple[sint, ...]: return self.views[-1].shape
 
   @property
-  def key(self) -> Tuple[View, ...]: return tuple(self.views)
+  def size(self) -> int: return self.views[-1].size()
+
+  def real_size(self) -> int:
+    if 0 in self.shape: return 0
+    idx, valid = self.expr_idxs()
+    if not valid: return 0
+    # TODO: it's possible that the real_size is smaller condition on valid being true
+    ret = idx.max
+    if not isinstance(ret, int): ret = ret.max  # might be represent by symbolic shape, take one more max for int max
+    assert isinstance(ret, int), f"ret must be integer, {ret=} isn't"
+    return ret+1
 
-  # this is the real size (ish)
-  def size(self): return prod([s for s,st in zip(self.views[-1].shape, self.views[-1].strides) if st != 0])
+  def vars(self) -> Set[Variable]: return set.union(*[v.vars() for v in self.views], set())
 
-  # these are multiview strides, value is None if it's not a simple strided dimension
-  # TODO: this can be shared code between simplify and merge_views
-  def real_offset(self) -> int:
-    real_offset, mask = self.expr_node(Variable('zero', 0, 0))
-    assert real_offset.__class__ is NumNode, f"how is the offset not a number? {real_offset} {mask}"
-    return real_offset.b
+  @property
+  def var_vals(self) -> Dict[Variable, int]: return merge_dicts([dict([v.unbind()]) for v in self.vars()])
+
+  def unbind(self) -> Tuple[ShapeTracker, Dict[Variable, int]]:
+    unbound_views, var_vals = zip(*[v.unbind() for v in self.views])
+    return ShapeTracker(tuple(unbound_views)), merge_dicts(var_vals)
 
   # NOTE: if a stride is not always valid, it will be None
-  def real_strides(self, ignore_valid=False) -> Tuple[Optional[Union[Node, int]], ...]:
+  def real_strides(self, ignore_valid=False) -> Tuple[Optional[sint], ...]:
     if len(self.views) == 1 and self.views[-1].mask is None: return self.views[-1].strides
-    idxs = [Variable(f"idx{i}", 0, s-1) for i,s in enumerate(self.shape)]
+    idxs: List[Node] = [Variable(f"idx{i}", 0, s-1) for i,s in enumerate(self.shape)]
     idx, valid = self.expr_idxs(idxs)
-    ret: List[Optional[Union[Node, int]]] = [None] * len(self.views[-1].shape)
+    ret: List[Optional[sint]] = [None] * len(self.views[-1].shape)
+    bad_idx_vars: Set[Variable] = set()
     for this_dim in (idx.nodes if isinstance(idx, SumNode) else [idx]):
-      if isinstance(this_dim, MulNode) and isinstance(this_dim.a, Variable) and this_dim.a in idxs:
-        ret[idxs.index(this_dim.a)] = this_dim.b
-      elif isinstance(this_dim, Variable):
-        ret[idxs.index(this_dim)] = 1
+      idx_maybe, stride_maybe = (this_dim.a, this_dim.b) if isinstance(this_dim, MulNode) else (this_dim, 1)
+      try: ret[idxs.index(idx_maybe)] = cast(sint, stride_maybe)
+      except ValueError: bad_idx_vars = bad_idx_vars.union(idx_maybe.vars())
     idx_vars, valid_vars = idx.vars(), valid.vars()
     for i,tidx in enumerate(idxs):
-      if tidx in valid_vars and not ignore_valid: ret[i] = None
+      if tidx in bad_idx_vars or (tidx in valid_vars and not ignore_valid): ret[i] = None
       elif tidx not in idx_vars: ret[i] = 0
     return tuple(ret)
+
   def unit_stride_axes(self, ignore_valid=False) -> List[int]: return [i for i,st in enumerate(self.real_strides(ignore_valid)) if st == 1]
 
-  def _expr_idx(self, idx, valid):
-    for v in reversed(self.views[0:-1]):
-      valid = v.expr_node_mask(idx, valid)
-      idx = v.expr_node(idx)
+  def expr_idxs(self, idxs:Optional[Iterable[Node]]=None) -> Tuple[Node, Node]:
+    idxs = [Variable(f"idx{i}", 0, s-1) for i,s in enumerate(self.shape)] if idxs is None else list(idxs)
+    idx, valid = _expr_view(self.views[-1], idxs)
+    for view in reversed(self.views[0:-1]):
+      if valid.max == 0: return NumNode(-1), valid
+      view = view.minify()
+      acc, idxs = 1, []
+      for d in reversed(view.shape):
+        idxs.append((idx//acc)%d)
+        acc *= d
+      idx, valid = _expr_view(view, idxs[::-1], valid)
+    assert not isinstance(idx.min, int) or idx.min >= -2**31, f"idx.min too small. {idx=}, {idx.min=}"
+    assert not isinstance(idx.max, int) or idx.max < 2**31, f"idx.max too big. {idx=}, {idx.max=}"
     return idx, valid
 
-  def simplify(self):
-    if len(self.views) >= 2:
-      new_view = merge_views(self.views[-2], self.views[-1])
-      if new_view:
-        if DEBUG >= 4: print(f"st simplify : {self.views[-2]} + {self.views[-1]} = {new_view}")
-        self.views = self.views[:-2] + [new_view]
-        self.simplify()
-
-  def expr_idxs(self, idxs=None):
-    if idxs is None: idxs = [Variable(f"idx{i}", 0, s-1) for i,s in enumerate(self.shape)]
-    idx = self.views[-1].expr_idxs(tuple(idxs))
-    valid = self.views[-1].expr_node_mask(idxs_to_idx(self.views[-1].shape, tuple(idxs)))
-    return self._expr_idx(idx, valid)
-
-  def expr_node(self, idx='idx'):
-    if idx.__class__ is str: idx = Variable(idx, 0, prod(self.shape)-1)
-    return self._expr_idx(self.views[-1].expr_node(idx), self.views[-1].expr_node_mask(idx))
-
-  def needs_valid(self) -> bool:
-    return any(v.mask is not None for v in self.views)
-
-  # *** under this line are the movement ops ***
-
-  def __unsafe_resize(self, arg: Tuple[Tuple[int, int], ...], mask=None):
-    offset = get_unsafe_resize_offset(self.views[-1].strides, arg)
-    if self.views[-1].mask:
-      # move the old mask
-      nmask = tuple([(max(mx-ax, 0), min(my-ax, ay-ax)) for (mx,my),(ax,ay) in zip(self.views[-1].mask, arg)])
-      # merge the masks if we have two
-      mask = tuple([(max(mx1, mx2), min(my1, my2)) for (mx1, my1), (mx2, my2) in zip(nmask, mask)]) if mask is not None else nmask
-    self.views[-1] = View(tuple([y-x for x,y in arg]), self.views[-1].strides, self.views[-1].offset+offset, mask)
+  def axis_is_masked(self, axis:int) -> bool:
+    _, valid = self.expr_idxs()
+    return f'idx{axis}' in [v.expr for v in valid.vars()]
 
-  def pad(self, arg: Tuple[Tuple[int, int], ...]):
-    assert all((b>=0 and e>=0) for b,e in arg) and len(arg) == len(self.shape)
-    if any(b or e for b, e in arg):
-      zvarg, mask = get_pad_args(self.shape, arg)
-      self.__unsafe_resize(zvarg, mask=mask)
+  def simplify(self) -> ShapeTracker:
+    if len(self.views) >= 2 and (new_view := self.views[-2] + self.views[-1]) is not None:
+      return ShapeTracker(self.views[:-2] + (new_view,)).simplify()
     return self
 
-  def shrink(self, arg: Tuple[Tuple[int, int], ...]):
-    assert all((b>=0 and e<=s) for s,(b,e) in zip(self.shape,arg)) and len(arg) == len(self.shape)
-    self.__unsafe_resize(arg)
-    return self
-
-  def expand(self, new_shape: Tuple[Union[Node,int], ...]) -> ShapeTracker:
-    assert len(new_shape) == len(self.views[-1].shape)
-    assert all(is_sym_int(x) and (s == x or (s == 1 and st == 0)) for s,x,st in zip(self.shape, new_shape, self.views[-1].strides)), f"can't expand {self.shape} into {new_shape}"
-    # NOTE: can the mask ever be (0,0)?
-    mask = tuple([(((0,0) if m != (0,1) else (0,ns)) if s != ns else m) for m,s,ns in zip(self.views[-1].mask, self.shape, new_shape)]) if self.views[-1].mask else None
-    self.views[-1] = View(new_shape, self.views[-1].strides, self.views[-1].offset, mask)
-    return self
-
-  def reshape(self, new_shape: Tuple[Union[Node,int], ...]):
-    if self.views[-1].shape == new_shape: return self
-    new_ints, new_nodes = partition(new_shape, lambda s: isinstance(s, int))
-    if new_nodes and all(isinstance(s, int) for s in self.shape):
-      # reshape from all int shape into shape with a variable, update the variable value
-      assert len(new_nodes) == 1 and isinstance(new_nodes[0], Variable), "only support adding one Variable to the int shape"
-      new_var, new_val = new_nodes[0], prod(self.shape) // prod(new_ints)
-      if new_var not in self.var_vals:
-        assert new_var.min <= new_val <= new_var.max, f"variable value {new_val} out of range [{new_var.min}, {new_var.max}]"
-        self.var_vals[new_var] = new_val
-      else: assert self.var_vals[new_var] == new_val, f"value conflicts, was {self.var_vals[new_var]}, set to {new_val}"
-    elif not new_nodes: self.var_vals = {}
-    assert all(is_sym_int(x) and x > 0 for x in new_shape), f"shape must be symbolic ints and can't contain 0 or negative numbers {new_shape}"
-    # only check size for int shapes. we don't check symbolic here as long as the reshape itself can be done
-    if all(isinstance(s, int) for s in self.shape) and all(isinstance(s, int) for s in new_shape):
-      assert prod(self.shape) == prod(new_shape), f"can't reshape {self.shape} -> {new_shape}" # type: ignore  # mypy cannot resolve, all ints here
-    new_view, extra = _reshape(self.views[-1], new_shape)
-    if extra: self.views.append(new_view)
-    else: self.views[-1] = new_view
-    return self
-
-  def permute(self, axis: Tuple[int, ...]):
-    assert all(isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis), f"invalid permute {axis} for {self.shape}"
-    assert len(set(axis)) == len(axis) and len(axis) == len(self.shape), f"can't permute {self.shape} with {axis}"
-    self.views[-1] = View(tuple([self.views[-1].shape[a] for a in axis]), tuple([self.views[-1].strides[a] for a in axis]), self.views[-1].offset, tuple([self.views[-1].mask[a] for a in axis]) if self.views[-1].mask is not None else None)
-    return self
+  # *** under this line are the movement ops ***
 
-  # except for the negative case, you can build this from the others. invertible in the negative case
-  def stride(self, mul: Tuple[int, ...]):
-    assert all(isinstance(x, int) and x != 0 for x in mul), f"invalid stride {mul} for {self.shape}"
-    strides = tuple([z*m for z,m in zip(self.views[-1].strides, mul)])
-    new_shape = tuple([(s+(abs(m)-1))//abs(m) for s,m in zip(self.views[-1].shape, mul)])
-    offset = sum([(s-1)*z for s,z,m in zip(self.views[-1].shape, self.views[-1].strides, mul) if m < 0])
-    mask = tuple([(((mx if m > 0 else s-my)+(abs(m)-1))//abs(m), ((my if m > 0 else s-mx)+(abs(m)-1))//abs(m)) for (mx,my),s,m in zip(self.views[-1].mask, self.views[-1].shape, mul)]) if self.views[-1].mask is not None else None
-    self.views[-1] = View(new_shape, strides, self.views[-1].offset + offset, mask)
-    return self
+  def pad(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].pad(arg), ))
+  def shrink(self, arg: Tuple[Tuple[sint, sint], ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].shrink(arg), ))
+  def expand(self, new_shape: Tuple[sint, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].expand(new_shape), ))
+  def permute(self, axis: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].permute(axis), ))
+  def stride(self, mul: Tuple[int, ...]) -> ShapeTracker: return ShapeTracker(self.views[0:-1] + (self.views[-1].stride(mul), ))
 
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:Tuple[int, ...], new_shape:Tuple[int, ...]) -> Optional[List[List[int]]]:
-  # Pre-allocate all groups.
-  axis_groups: List[List[int]] = [[] for _ in range(len(new_shape))]
-  # Index for new_shape and axis_groups.
-  i: int = 0
-  old_shape_i: int = 0
-  while old_shape_i < len(old_shape):
-    # 1s exist in new_shape only will lead to empty axes group creations.
-    if new_shape[i] == 1 and old_shape[old_shape_i] != 1:
-      if i < len(new_shape) - 1: i += 1
-    else:
-      axis_groups[i].append(old_shape_i)
-      axis_group_size = prod([old_shape[x] for x in axis_groups[i]])
-      # Move to next axes group if total size of all dimensions match.
-      if axis_group_size == new_shape[i]:
-        if i < len(new_shape) - 1: i += 1
-      elif axis_group_size > new_shape[i]: return None
-      old_shape_i += 1
-  return axis_groups
+  def reshape(self, new_shape: Tuple[sint, ...]) -> ShapeTracker:
+    if getenv("MERGE_VIEW", 1) and (new_view := self.views[-1].reshape(new_shape)) is not None: return ShapeTracker(self.views[0:-1] + (new_view,))
+    return ShapeTracker(self.views + (View.create(new_shape), ))