diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 8dc0643319..6305a8268f 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -6,6 +6,45 @@ from tinygrad.helpers import prod, DEBUG from tinygrad.shape.symbolic import Variable, MulNode, NumNode, Node, SumNode, sint from tinygrad.shape.view import View +@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) + +def expr_node_mask(view:View, idx, valid=None) -> Node: + expr = [valid] if valid is not None else [] + if view.mask is not None: + acc = 1 + for ns,(x,y) in reversed(list(zip(view.shape, view.mask))): + if x != 0 or y != ns: + 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(view:View, idx=None) -> Node: + if idx is None: idx = Variable('idx', 0, prod(view.shape)-1) + ret: List[Node] = [Variable.num(view.offset) if isinstance(view.offset, int) else view.offset] if view.offset else [] + acc = 1 + for d,s in reversed(to_shape_strides(view.shape, view.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(view:View, idxs) -> Node: + assert len(idxs) == len(view.shape), f"need an idx for all dimensions {idxs} vs {view.shape}" + return Variable.sum([Variable.num(view.offset) if isinstance(view.offset, int) else view.offset] + [idx*st for idx,sh,st in zip(idxs, view.shape, view.strides) if sh != 1 and st != 0]) + @functools.lru_cache(maxsize=None) def merge_views(vm2:View, vm1:View) -> Optional[View]: if vm2.mask: return None # this isn't supported yet @@ -70,8 +109,8 @@ class ShapeTracker: def _expr_idx(self, idx, valid) -> Tuple[Node, Node]: for v in reversed(self.views[0:-1]): if valid.max == 0: return Variable.num(-1), valid - valid = v.expr_node_mask(idx, valid) - idx = v.expr_node(idx) + valid = expr_node_mask(v, idx, valid) + idx = expr_node(v, idx) return idx, valid def simplify(self): @@ -84,13 +123,13 @@ class ShapeTracker: 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))) + idx = expr_idxs(self.views[-1], tuple(idxs)) + valid = expr_node_mask(self.views[-1], 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)) + return self._expr_idx(expr_node(self.views[-1], idx), expr_node_mask(self.views[-1], idx)) def axis_is_masked(self, axis) -> bool: _, valid = self.expr_idxs() diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py index 8be7d61fb7..7dc43812ba 100644 --- a/tinygrad/shape/view.py +++ b/tinygrad/shape/view.py @@ -2,20 +2,7 @@ from __future__ import annotations import functools from typing import Tuple, List, Optional, NamedTuple from tinygrad.helpers import prod -from tinygrad.shape.symbolic import Variable, Node, NumNode, is_sym_int, sint, all_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) +from tinygrad.shape.symbolic import NumNode, is_sym_int, sint, all_int @functools.lru_cache(maxsize=None) def filter_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> Tuple[int, ...]: @@ -44,32 +31,6 @@ class View(NamedTuple): @functools.lru_cache(maxsize=None) # pylint: disable=method-cache-max-size-none def size(self): return prod([s for s,st in zip(self.shape, self.strides) if st != 0]) - 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))): - if x != 0 or y != ns: - 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(to_shape_strides(self.shape, self.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]) - # MovementOps live here now def __unsafe_resize(self, arg: Tuple[Tuple[sint, sint], ...], mask=None) -> View: