losing lines (#678)

* losing lines * FLIP -> STRIDE * shapetracker refactor
2023-03-10 21:57:05 -08:00
parent d7cb8e3e56
commit 0b03216cc3
13 changed files with 103 additions and 109 deletions
@@ -158,12 +158,12 @@ You no longer need to write mlops for a new accelerator
 The autodiff stuff is all in mlops now so you can focus on the raw operations

 ```
-Buffer                                                     # class of memory on this device
-unary_op  (NOOP, NEG, NOT, EXP, LOG)                       # A -> A
-reduce_op (SUM, MAX)                                       # A -> B (smaller size, B has 1 in shape)
-binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX)            # A + A -> A (all the same size)
-movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, FLIP)  # A -> B (different size)
-fused_op [[optional]] (MULACC)                             # A * A -> B
+Buffer                                                       # class of memory on this device
+unary_op  (NOOP, NEG, NOT, EXP, LOG)                         # A -> A
+reduce_op (SUM, MAX)                                         # A -> B (smaller size, B has 1 in shape)
+binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX)              # A + A -> A (all the same size)
+movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE)  # A -> B (different size)
+fused_op [[optional]] (MULACC)                               # A * A -> B
 ```

 ## ImageNet inference
@@ -269,6 +269,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1)), lambda x: x.flip(axis=(0,1)))
    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)), lambda x: x.flip(axis=(0,1,3)))
    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(3,)))
+    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)).flip((0,)), lambda x: x.flip(axis=(0,1,3)).flip(0))
+    helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(-1,)))

  def test_unsqueeze(self):
    helper_test_op([(4,3,6,6)], lambda x: torch.unsqueeze(x, 0), lambda x: x.unsqueeze(dim=0))
@@ -23,31 +23,31 @@ class CheckingShapeTracker:
  def simplify(self): self.st.simplify()

  def reshape(self, new_shape):
-    self.st.reshape(new_shape)
+    self.st._reshape(new_shape)
    self.t = self.t.reshape(new_shape)

  def permute(self, axis):
-    self.st.permute(axis)
+    self.st._permute(axis)
    self.t = np.transpose(self.t, axis)

  def expand(self, new_shape):
-    self.st.expand(new_shape)
+    self.st._expand(new_shape)
    self.t = np.broadcast_to(self.t, new_shape)

  def flip(self, axis):
-    self.st.flip(axis)
+    self.st._stride(tuple(-1 if i in axis else 1 for i in range(len(self.shape))))
    self.t = np.flip(self.t, axis)

  def shrink(self, arg):
-    self.st.shrink(arg)
+    self.st._shrink(arg)
    self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])]

  def pad(self, arg):
-    self.st.pad(arg)
+    self.st._pad(arg)
    self.t = np.pad(self.t, arg, constant_values=-1)

  def stride(self, arg):
-    self.st.stride(arg)
+    self.st._stride(arg)
    self.t = self.t[tuple([slice(None, None, x) for x in arg])]

  def __getitem__(self, val):
@@ -148,7 +148,7 @@ class TestSimplifyingShapeTracker(unittest.TestCase):

 class TestComplexShapeTracker(unittest.TestCase):
  def test_add_1s(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
    self.st.permute((1,0))
    self.st.reshape((1,4,1,4,1))
    assert not self.st.contiguous
@@ -156,20 +156,20 @@ class TestComplexShapeTracker(unittest.TestCase):
    assert self.st.contiguous

  def test_permute_1s_simple(self):
-    self.st = ShapeTracker((1, 16, 9,9))
+    self.st = CheckingShapeTracker((1, 16, 9,9))
    self.st.permute((1,0,2,3))
    assert self.st.contiguous
-    self.st = ShapeTracker((2, 16, 9,9))
+    self.st = CheckingShapeTracker((2, 16, 9,9))
    self.st.permute((1,0,2,3))
    assert not self.st.contiguous

  def test_remove_1s_simple(self):
-    self.st = ShapeTracker((1, 16, 1, 1))
+    self.st = CheckingShapeTracker((1, 16, 1, 1))
    self.st.reshape((16,))
    assert self.st.contiguous

  def test_remove_1s(self):
-    self.st = ShapeTracker((1, 4, 1, 4, 1))
+    self.st = CheckingShapeTracker((1, 4, 1, 4, 1))
    self.st.permute((0,3,2,1,4))
    self.st.reshape((4,4))
    assert not self.st.contiguous
@@ -177,46 +177,46 @@ class TestComplexShapeTracker(unittest.TestCase):
    assert self.st.contiguous

  def test_permute_reshape(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
    self.st.permute((1,0))
    self.st.reshape((2, 2, 2, 2))
    # TODO: should also be tested by test_super_complex
    assert len(self.st.views) == 1

  def test_factorize_split(self):
-    self.st = ShapeTracker((4, 4))
+    self.st = CheckingShapeTracker((4, 4))
    self.st.permute((1,0))
    self.st.reshape((2, 2, 2, 2))
    self.st.permute((2,3,0,1))
    assert self.st.contiguous

  def test_factorize_combine(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
    self.st.permute((2, 0, 1))
    self.st.reshape((4, 16))
    self.st.permute((1, 0))
    assert self.st.contiguous

  def test_factorize_combine_add_ones(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
    self.st.permute((2, 0, 1))
    self.st.reshape((4, 16, 1, 1))
    self.st.permute((1, 0, 2, 3))
    assert self.st.contiguous

  def test_fancy_factorize(self):
-    self.st = ShapeTracker((32, 3, 3, 1))
+    self.st = CheckingShapeTracker((32, 3, 3, 1))
    self.st.reshape((8, 4, 3, 3))
    assert len(self.st.views) == 1

  def test_super_complex_2_fail(self):
-    self.st = ShapeTracker((4, 4, 4))
+    self.st = CheckingShapeTracker((4, 4, 4))
    self.st.permute((2, 0, 1))
    self.st.reshape((16, 4))
    assert len(self.st.views) != 1

  def test_work(self):
-    self.st = ShapeTracker((64, 1024, 4))
+    self.st = CheckingShapeTracker((64, 1024, 4))
    self.st.reshape((1, 64, 128, 32))
    self.st.permute((0, 3, 1, 2))
    self.st.reshape((1, 32, 1, 64, 128))
@@ -224,7 +224,7 @@ class TestComplexShapeTracker(unittest.TestCase):
    assert self.st.contiguous

  def test_work2(self):
-    self.st = ShapeTracker((64, 1024, 4))
+    self.st = CheckingShapeTracker((64, 1024, 4))
    self.st.reshape((1, 64, 128, 32))
    self.st.permute((0, 3, 1, 2))
    self.st.reshape((1, 1, 32, 64, 128))
@@ -166,7 +166,7 @@ class ASTKernel:
        if mergeable: rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i])
        else: rets[j].append((shapes[j][i], strides[j][i]))

-    for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x))
+    for i,x in enumerate(rets): self.sts[i]._reshape(tuple(y[0] for y in x))
    self.first_reduce = get_first_reduce([x.shape for x in self.sts])

  # this should be aware of the three parts to the shape
@@ -175,8 +175,8 @@ class ASTKernel:
  #  * the size outputted by each kernel
  def reshape_and_permute(self, new_shape_fxn, axis):
    for st in self.sts:
-      if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape)))
-      if axis is not None: st.permute(tuple(axis))
+      if new_shape_fxn is not None: st._reshape(tuple(new_shape_fxn(st.shape)))
+      if axis is not None: st._permute(tuple(axis))

  # axis : the axis to pull from
  # amount : the amount to take
@@ -188,7 +188,7 @@ class LazyBuffer:
      if op in [MovementOps.RESHAPE, MovementOps.EXPAND, MovementOps.SHRINK]: return self.op.src[0].movement_op(op, arg)
      if op == MovementOps.PERMUTE: return self.op.src[0].movement_op(op, tuple(self.op.arg[i] for i in arg))
      if op == MovementOps.PAD: return self.op.src[0].movement_op(op, tuple((b1+b2, e1+e2) for (b1,e1),(b2,e2) in zip(self.op.arg, arg)))
-      if op == MovementOps.FLIP: return self.op.src[0].movement_op(op, tuple(i for i in arg+self.op.arg if not (i in arg and i in self.op.arg)))
+      if op == MovementOps.STRIDE: return self.op.src[0].movement_op(op, tuple(i*j for i,j in zip(arg, self.op.arg)))

    # push permutes before reduce ops
    if op == MovementOps.PERMUTE and PUSH_PERMUTES and self.realized is None and self.optype == ReduceOps:
@@ -155,8 +155,8 @@ class Shrink(Function):

 class Flip(Function):
  def forward(self, x, axis):
-    self.axis = axis
-    return x.movement_op(MovementOps.FLIP, axis)
+    self.arg = tuple(-1 if i in axis else 1 for i in range(len(x.shape)))
+    return x.movement_op(MovementOps.STRIDE, self.arg)

  def backward(self, grad_output):
-    return grad_output.movement_op(MovementOps.FLIP, self.axis)
+    return grad_output.movement_op(MovementOps.STRIDE, self.arg)
@@ -4,14 +4,13 @@ import numpy as np
 from enum import Enum, auto
 from typing import Union, Type, NamedTuple, Tuple, Any, List, ClassVar, Optional, Callable, Dict, TypeVar, Set
 from tinygrad.helpers import prod, DEBUG, getenv, DType, dtypes
-from tinygrad.shape import ShapeTracker
+from tinygrad.shape import ShapeTracker, MovementOps

 # these are the llops your accelerator must implement, along with toCpu
 # the Enum class doesn't work with mypy, this is static. sorry it's ugly
 class UnaryOps(Enum): NOOP = auto(); NEG = auto(); EXP = auto(); LOG = auto(); NOT = auto() # noqa: E702
 class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto() # noqa: E702
 class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702
-class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); FLIP = auto(); PAD = auto(); SHRINK = auto() # noqa: E702
 class FusedOps(Enum): MULACC = auto() # noqa: E702
 class LoadOps(Enum): FROMCPU = auto(); CONTIGUOUS = auto(); TOCPU = auto(); CUSTOM = auto() # noqa: E702

@@ -55,6 +54,11 @@ class RawBufferCopyIn(RawBuffer):
    ret.copyin(x)
    return ret

+class RawBufferMapped(RawBufferCopyIn):
+  def _buffer(self) -> memoryview: raise NotImplementedError("must be implemented")
+  def toCPU(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=self.dtype.np)
+  def copyin(self, x:np.ndarray) -> None: np.copyto(self.toCPU(), x.reshape(-1))
+
 class RawBufferCopyInOut(RawBufferCopyIn):
  def copyout(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented")

@@ -130,7 +134,9 @@ class ASTRunner:
    if DEBUG >= 1:
      print(f"*** {GlobalCounters.kernel_count:4d} {self.name:20s} arg {len(rawbufs):3d} sz {str(self.global_size):18s} {str(self.local_size):12s} OPs {self.op_estimate/1e6:7.1f}M/{GlobalCounters.global_ops/1e9:7.2f}G  mem {GlobalCounters.mem_used/1e9:5.2f} GB " +
            (str() if et is None else f"tm {et*1e6:9.2f}us/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({self.op_estimate/(et*1e9):8.2f} GFLOPS, {self.mem_estimate/(et*1e9):6.2f} GB/s)"))
-    GlobalCounters.log_kernel(self.op_estimate, self.mem_estimate)
+    GlobalCounters.kernel_count += 1
+    GlobalCounters.global_ops += self.op_estimate
+    GlobalCounters.global_mem += self.mem_estimate
    if getenv("EARLY_STOPPING") and GlobalCounters.kernel_count == getenv("EARLY_STOPPING"): exit(0)
    return et

@@ -215,9 +221,4 @@ class GlobalCounters:
  mem_used : ClassVar[int] = 0   # NOTE: this is not reset
  cache : ClassVar[Optional[List[Tuple[Callable, Any]]]] = None
  @staticmethod
-  def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count, GlobalCounters.cache = 0,0,0.0,0,None
-  @staticmethod
-  def log_kernel(op_estimate:int, mem_estimate:int):
-    GlobalCounters.kernel_count += 1
-    GlobalCounters.global_ops += op_estimate
-    GlobalCounters.global_mem += mem_estimate
+  def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count, GlobalCounters.cache = 0,0,0.0,0,None
@@ -1,18 +1,15 @@
 import os, time, ctypes, hashlib, subprocess, platform
-import numpy as np
 from collections import defaultdict
 from typing import Final, Dict
 from tinygrad.helpers import dtypes, DType
-from tinygrad.ops import CompiledBuffer, RawBufferCopyIn
+from tinygrad.ops import CompiledBuffer, RawBufferMapped
 from tinygrad.codegen.gpu import GPUCodegen, GPULanguage

-class RawMallocBuffer(RawBufferCopyIn):
+class RawMallocBuffer(RawBufferMapped):
  def __init__(self, size, dtype : DType):
    super().__init__(size, dtype)
    self._buf = ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16}[dtype] * size)()
-  def _buffer(self): return self._buf
-  def copyin(self, x:np.ndarray): ctypes.memmove(self._buf, x.ctypes.data, x.size*np.dtype(x.dtype).itemsize)
-  def toCPU(self): return np.frombuffer(self._buf, dtype=self.dtype.np)
+  def _buffer(self): return memoryview(self._buf)

 class ClangProgram:
  kernel_cnt : Final[Dict[str, int]] = defaultdict(int)
@@ -29,9 +29,9 @@ def einsum_mulacc(einsum, get_strides, expand):
 numpy_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{
  UnaryOps.NOOP: np.ascontiguousarray, UnaryOps.EXP: np.exp, UnaryOps.LOG: np.log,
  BinaryOps.MAX: np.maximum, BinaryOps.CMPEQ: lambda x,y: (x==y).astype(np.float32),
-  MovementOps.FLIP: np.flip, MovementOps.PERMUTE: lambda x, order: x.transpose(order),
-  MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to,
-  FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, np.broadcast_to)
+  MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to,
+  MovementOps.STRIDE: lambda x, arg: x.__getitem__(tuple(slice(None, None, i) for i in arg)),
+  FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, np.broadcast_to),
 }}

 class CPUBuffer(InterpretedBuffer):
@@ -1,11 +1,10 @@
 # pip3 install pyobjc-framework-Metal pyobjc-framework-Cocoa pyobjc-framework-libdispatch
 import os, subprocess, pathlib, functools
 import Metal, Cocoa, libdispatch # type: ignore
-import numpy as np
 from typing import List, Any
 from tinygrad.codegen.gpu import GPUCodegen, GPULanguage
 from tinygrad.helpers import prod, getenv, DEBUG, DType
-from tinygrad.ops import CompiledBuffer, RawBufferCopyIn
+from tinygrad.ops import CompiledBuffer, RawBufferMapped

 METAL_XCODE = getenv("METAL_XCODE")

@@ -19,20 +18,17 @@ class _METAL:
    return METAL.device.newCommandQueue()
 METAL = _METAL()

-class RawMetalBuffer(RawBufferCopyIn):
+class RawMetalBuffer(RawBufferMapped):
  def __init__(self, size:int, dtype:DType):
    super().__init__(size, dtype)
    self._cl = METAL.device.newBufferWithLength_options_(size*dtype.itemsize, Metal.MTLResourceStorageModeShared)
  def __del__(self):
    self._cl.release()
    super().__del__()
-  def _buffer(self): return self._cl.contents().as_buffer(self._cl.length())
-  def _as_np(self): return np.frombuffer(self._buffer(), dtype=self.dtype.np)
-  def copyin(self, x:np.ndarray): np.copyto(self._as_np(), x.reshape(-1).data)
-  def toCPU(self) -> np.ndarray:
+  def _buffer(self):
    for cbuf in METAL.mtl_buffers_in_flight: cbuf.waitUntilCompleted()
    METAL.mtl_buffers_in_flight = []
-    return self._as_np()  # no copy!
+    return self._cl.contents().as_buffer(self._cl.length())

 def unwrap(x):
  ret, err = x
@@ -8,7 +8,8 @@ torch_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{
  UnaryOps.NOOP: lambda x: x.contiguous(), UnaryOps.EXP: lambda x: x.exp(), UnaryOps.LOG: lambda x: x.log(),
  BinaryOps.MAX: torch.maximum, BinaryOps.CMPEQ: lambda x,y: (x==y).float(),
  MovementOps.PAD: lambda x, padding: torch.nn.functional.pad(x, [item for sublist in padding[::-1] for item in sublist]),
-  FusedOps.MULACC: einsum_mulacc(lambda s,a,b: torch.einsum(s, a.float(), b.float()).type(a.dtype), lambda x: x.stride(), lambda x,s: x.expand(s))
+  FusedOps.MULACC: einsum_mulacc(lambda s,a,b: torch.einsum(s, a.float(), b.float()).type(a.dtype), lambda x: x.stride(), lambda x,s: x.expand(s)),
+  MovementOps.STRIDE: lambda x, arg: x.__getitem__(tuple(slice(None, None, abs(i)) for i in arg)).flip([i for i,a in enumerate(arg) if a < 0])
 }}

 device = torch.device("cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu"))
@@ -1,10 +1,14 @@
 # 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 Tuple, Union, List, Optional, cast
+from enum import Enum, auto
+from typing import Tuple, Union, List, Optional, cast, Dict, Callable
 from tinygrad.helpers import prod, DEBUG
 from tinygrad.shape.symbolic import Variable, MulNode, NumNode, Node

+# these ops live here
+class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto() # noqa: E702
+
@functools.lru_cache(maxsize=None)
 def to_shape_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> List[Tuple[int, int]]:
  assert len(shape) == len(strides)
@@ -90,6 +94,8 @@ def merge_views(vm2:View, vm1:View) -> Optional[View]:
  return View(vm1.shape, tuple(new_strides), new_offset.b) if len(new_strides) == len(vm1.strides) else None

 class ShapeTracker:
+  dispatch : Dict[MovementOps, Callable]
+
  def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], views:Optional[List[ViewTypes]]=None):
    self.views : List[ViewTypes] = views if views is not None else (shape.views[:] if isinstance(shape, ShapeTracker) else [view_from_shape(shape)])
  def __repr__(self): return f"ShapeTracker(shape={self.shape}, views={self.views})"
@@ -133,13 +139,35 @@ class ShapeTracker:
  def expr_node(self, idx='idx', offset=0):
    return self._expr_idx(self.views[-1].expr_node(Variable(idx, 0, prod(self.shape)-1), offset))

-  def movement_op(self, op, arg:Union[Tuple[int, ...], Tuple[Tuple[int, int], ...]]) -> ShapeTracker:
-    return getattr(self, str(op).split(".")[1].lower())(arg)
  def needs_valid(self) -> bool:
    return any(isinstance(v, ZeroView) for v in self.views)

-  def reshape(self, new_shape : Tuple[int, ...]) -> ShapeTracker:
-    assert isinstance(new_shape, tuple)
+  # *** under this line are the movement ops ***
+
+  def __unsafe_resize(self, arg : Tuple[Tuple[int, int], ...]):
+    offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)])
+    self.views[-1] = View(tuple(y-x for x,y in arg), self.strides, self.offset+offset)
+
+  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 all(b==0 and e==0 for b,e in arg): return self   # ZeroView is expensive if we don't need it
+    zvarg = tuple((-b,s+e) for s,(b,e) in zip(self.shape, arg))
+    zeroview = ZeroView(self.shape, zvarg)
+    self.__unsafe_resize(zvarg)
+    # if we add a ZeroView, we add another (stock) view also for modding
+    self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))]
+
+  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)
+
+  def _expand(self, new_shape : Tuple[int, ...]):
+    assert all(isinstance(x, int) for x in new_shape), f"non ints for expand in {new_shape}"
+    assert all(x == y or x == 1 for x,y in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}"
+    strides : Tuple[int, ...] = tuple(s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape)))
+    self.views[-1] = View(new_shape, strides, self.offset)
+
+  def _reshape(self, new_shape : Tuple[int, ...]):
    if self.shape == new_shape: return self
    assert all(isinstance(x, int) and x != 0 for x in new_shape), f"shape must be ints and can't contain 0 {new_shape}"
    assert prod(self.shape) == prod(new_shape), f"can't reshape {self.shape} -> {new_shape}"
@@ -150,61 +178,30 @@ class ShapeTracker:
      # NOTE: the last view in self.views is never a ZeroView
      if (merged_view := merge_views(cast(View, self.views[-1]), view)) is not None: self.views[-1] = merged_view
      else: self.views.append(view)
-    return self

-  def permute(self, axis : Tuple[int, ...]) -> ShapeTracker:
-    assert isinstance(axis, tuple)
+  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.shape[a] for a in axis), tuple(self.strides[a] for a in axis), self.offset)
-    return self

-  # TODO: this is a special case of slice with strides, remove it
-  # though it's nice that it can't change size
-  def flip(self, axis : Tuple[int, ...]) -> ShapeTracker:
-    return self.stride(tuple(-1 if i in axis else 1 for i in range(len((self.shape)))))
-
-  # *** under this line are not invertible ***
-
-  def _resize(self, arg : Tuple[Tuple[int, int], ...]):
-    offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)])
-    self.views[-1] = View(tuple(y-x for x,y in arg), self.strides, self.offset+offset)
-
-  def pad(self, arg : Tuple[Tuple[int, int], ...]) -> ShapeTracker:
-    assert isinstance(arg, tuple)
-    assert all((b>=0 and e>=0) for b,e in arg) and len(arg) == len(self.shape)
-    if all(b==0 and e==0 for b,e in arg): return self   # ZeroView is expensive if we don't need it
-    zvarg = tuple((-b,s+e) for s,(b,e) in zip(self.shape, arg))
-    zeroview = ZeroView(self.shape, zvarg)
-    self._resize(zvarg)
-    # if we add a ZeroView, we add another (stock) view also for modding
-    self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))]
-    return self
-
-  def shrink(self, arg : Tuple[Tuple[int, int], ...]) -> ShapeTracker:
-    assert isinstance(arg, tuple)
-    assert all((b>=0 and e<=s) for s,(b,e) in zip(self.shape,arg)) and len(arg) == len(self.shape)
-    self._resize(arg)
-    return self
-
-  def expand(self, new_shape : Tuple[int, ...]) -> ShapeTracker:
-    assert isinstance(new_shape, tuple)
-    assert all(isinstance(x, int) for x in new_shape), f"non ints for expand in {new_shape}"
-    assert all(x == y or x == 1 for x,y in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}"
-    strides : Tuple[int, ...] = tuple(s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape)))
-    self.views[-1] = View(new_shape, strides, self.offset)
-    return self
-
-  # TODO: combine with flip? this is more generic than we need
-  def stride(self, mul : Tuple[int, ...]) -> ShapeTracker:
-    assert isinstance(mul, tuple)
+  # 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) for x in mul)
    strides = tuple(z*m for z,m in zip(self.strides, mul))
    new_shape = tuple((s+(abs(m)-1))//abs(m) for s,m in zip(self.shape, mul))
    offset = sum([(s-1)*z for s,z,m in zip(self.shape, self.strides, mul) if m < 0])
    self.views[-1] = View(new_shape, strides, self.offset + offset)
+
+  # *** entry point for external ***
+
+  def movement_op(self, op, arg:Union[Tuple[int, ...], Tuple[Tuple[int, int], ...]]) -> ShapeTracker:
+    assert isinstance(arg, tuple) and (len(arg) == len(self.shape) or op == MovementOps.RESHAPE), f"arg {arg} for {op} doesn't match dim of shape {self.shape}"
+    ShapeTracker.dispatch[op](self, arg)
    return self

+# populate dispatch
+ShapeTracker.dispatch = {op:getattr(ShapeTracker, "_"+str(op).split(".")[1].lower()) for op in MovementOps}
+
 # 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, ...]):
  # Pre-allocate all groups.
@@ -191,7 +191,7 @@ class Tensor:
    return mlops.Reshape.apply(self, shape=tuple(-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape))
  def expand(self, shape, *args) -> Tensor: return mlops.Expand.apply(self, shape=tuple(x if x != -1 else s for s,x in zip(self.shape, argfix(shape, *args))))
  def permute(self, order, *args) -> Tensor: return mlops.Permute.apply(self, order=argfix(order, *args))
-  def flip(self, axis, *args) -> Tensor: return mlops.Flip.apply(self, axis=argfix(axis, *args))
+  def flip(self, axis, *args) -> Tensor: return mlops.Flip.apply(self, axis=[x if x >= 0 else x+len(self.shape) for x in argfix(axis, *args)])
  def pad(self, arg:Tuple[Tuple[int, int], ...]) -> Tensor: return mlops.Pad.apply(self, arg=arg) if any(x != (0,0) for x in arg) else self
  def shrink(self, arg:Tuple[Tuple[int, int], ...]) -> Tensor: return mlops.Shrink.apply(self, arg=arg) if any(x != (0,s) for x,s in zip(arg, self.shape)) else self