From a89142e46f579727329b33271b72db68cb4848b1 Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 14 Aug 2023 18:53:37 -0700 Subject: [PATCH] ShapeTracker.var_vals (#1540) --- test/test_lazybuffer.py | 24 ------- test/test_symbolic_shapetracker.py | 106 ++++++++++++++++++++--------- test/unit/test_symbolic.py | 10 +-- tinygrad/lazy.py | 11 ++- tinygrad/shape/shapetracker.py | 28 ++++---- tinygrad/shape/symbolic.py | 2 - 6 files changed, 96 insertions(+), 85 deletions(-) diff --git a/test/test_lazybuffer.py b/test/test_lazybuffer.py index cc71a0d0f4..19afc25de5 100644 --- a/test/test_lazybuffer.py +++ b/test/test_lazybuffer.py @@ -68,29 +68,5 @@ class TestLazyBuffer(unittest.TestCase): assert GlobalCounters.cache[2][0].name.startswith("E_") GlobalCounters.cache = None -class TestVariableBuffer(unittest.TestCase): - def test_get_variable_buffers_no_variable(self): - t = Tensor.rand(2, 3) - assert t.lazydata.get_variable_buffers() == {} - - def test_get_variable_buffers_one_variable(self): - v = Variable("v", 1, 10) - t = Tensor.rand(2, 3).reshape(v, 3) - buffers = t.lazydata.get_variable_buffers() - assert len(buffers) == 1 and buffers[v].realize().realized.toCPU() == 2 - v = Variable("v", 1, 10) - t = Tensor.rand(2, 3).reshape(2, v) - buffers = t.lazydata.get_variable_buffers() - assert len(buffers) == 1 and buffers[v].realize().realized.toCPU() == 3 - - def test_get_variable_buffers_cat(self): - v1 = Variable("v1", 1, 10) - v2 = Variable("v2", 1, 10) - t1 = Tensor.rand(2, 3).reshape(v1, 3) - t2 = Tensor.rand(6, 3).reshape(v2, 3) - t = t1.cat(t2) - buffers = t.lazydata.get_variable_buffers() - assert len(buffers) == 2 and buffers[v1].realize().realized.toCPU() == 2 and buffers[v2].realize().realized.toCPU() == 6 - if __name__ == "__main__": unittest.main() diff --git a/test/test_symbolic_shapetracker.py b/test/test_symbolic_shapetracker.py index 7a0ee876bb..8a8857b97b 100644 --- a/test/test_symbolic_shapetracker.py +++ b/test/test_symbolic_shapetracker.py @@ -40,65 +40,91 @@ class TestSymbolic(unittest.TestCase): class TestSymbolicReshape(unittest.TestCase): def test_reshape_into_symbols_simple(self): - for i in range(1, 5): - vi = Variable("i", 1, 10) - assert Tensor.rand(i, 4).reshape(vi, 4).shape == (vi, 4) - assert vi.val == i - vi = Variable("i", 1, 10) - assert Tensor.rand(i, 6).reshape(vi, 2, 3).shape == (vi, 2, 3) - assert vi.val == i + vi = Variable("i", 1, 5) + for i in range(1, 6): + t = Tensor.rand(i, 4).reshape(vi, 4) + assert t.shape == (vi, 4) + assert t.lazydata.st.var_vals[vi] == i + t = Tensor.rand(i, 6).reshape(vi, 2, 3) + assert t.shape == (vi, 2, 3) + assert t.lazydata.st.var_vals[vi] == i def test_reshape_symbols_reshape_ints(self): - for i in range(1, 5): - vi = Variable("i", 1, 10) - assert Tensor.rand(i, 4).reshape(vi, 4).reshape(i, 4).shape == (i, 4) - assert Tensor.rand(i, 4).reshape(vi, 4).reshape(i*4,).shape == (i*4,) - assert Tensor.rand(i, 6).reshape(vi, 6).reshape(i*2, 3).shape == (i*2, 3) - with self.assertRaises(AssertionError): - Tensor.rand(i, 6).reshape(vi, 6).reshape(1, 77).shape + vi = Variable("i", 1, 5) + for i in range(1, 6): + t = Tensor.rand(i, 4).reshape(vi, 4) + assert t.shape == (vi, 4) + assert t.lazydata.st.var_vals == {vi: i} + t = t.reshape(i, 4) + assert t.shape == (i, 4) + assert t.lazydata.st.var_vals == {} def test_reshape_reuse_var_same_value_ok(self): - for i in range(1, 5): - vi = Variable("i", 1, 10) + vi = Variable("i", 1, 5) + for i in range(1, 6): a = Tensor.rand(i, 4).reshape(vi, 4) b = Tensor.rand(i, 3).reshape(vi, 3) - assert vi.val == i + assert a.lazydata.st.var_vals[vi] == i + assert b.lazydata.st.var_vals[vi] == i - def test_reshape_reuse_var_different_value_fail(self): - for i in range(1, 5): - vi = Variable("i", 1, 10) + def test_reshape_reuse_var_different_value_ok(self): + vi = Variable("i", 1, 10) + for i in range(1, 6): a = Tensor.rand(i, 4).reshape(vi, 2) - with self.assertRaises(AssertionError): - b = Tensor.rand(i, 3).reshape(vi, 3) + b = Tensor.rand(i, 3).reshape(vi, 3) + # a and b have different values of vi + assert a.lazydata.st.var_vals[vi] == 2 * i + assert b.lazydata.st.var_vals[vi] == i def test_reshape_into_symbols_bad_shape(self): vi = Variable("i", 1, 10) vj = Variable("j", 1, 10) with self.assertRaises(AssertionError): - t = Tensor.rand(3, 4).reshape(vi, vj) + t = Tensor.rand(3, 4).reshape(vi, vj) # reshape into two variables with self.assertRaises(AssertionError): - t = Tensor.rand(4, 4).reshape(vi, vi) + t = Tensor.rand(4, 4).reshape(vi, vi) # reshape into same variable in 2 dimensions with self.assertRaises(AssertionError): - t = Tensor.rand(4, 6).reshape(vi, 6).reshape(vi, 4) + t = Tensor.rand(4, 6).reshape(vi, 6).reshape(vi, 4) # conflicted implied variable values + with self.assertRaises(AssertionError): + t = Tensor.rand(4, 6).reshape(vi, 6).reshape(1, 77) # reshape to a different size new shape through symbolic shape with self.assertRaises(AssertionError): t = Tensor.rand(100, 4).reshape(Variable("too_small", 1, 10), 4) with self.assertRaises(AssertionError): t = Tensor.rand(3, 4).reshape(Variable("too_big", 100, 200), 4) + with self.assertRaises(AssertionError): + t = Tensor.rand(3, 4).reshape(3, (vi+1)) # reshape into non-Variable Node + + def test_two_symbol_reshape(self): + vi = Variable("i", 1, 5) + vj = Variable("j", 1, 5) + for i in range(1, 6): + for j in range(1, 6): + t1 = Tensor.rand(i, 5).reshape(vi, 5) + t2 = Tensor.rand(5, j).reshape(5, vj) + t = t1@t2 + assert t.shape == (vi, vj) + t = t.reshape(1, vi*vj) + assert t.shape == (1, vi*vj) + t = t.reshape(vj, vi) + assert t.shape == (vj, vi) class TestSymbolicExpand(unittest.TestCase): def test_expand_into_symbols(self): - vi = Variable("i", 1, 10) + vi = Variable("i", 1, 5) + vj = Variable("j", 1, 5) a = Tensor([[1], [2], [3]]).expand((3, vi)) assert a.shape == (3, vi) - vj = Variable("j", 1, 10) + assert a.lazydata.st.var_vals == {} a = a.reshape(3, vi, 1).expand((3, vi, vj)) assert a.shape == (3, vi, vj) + assert a.lazydata.st.var_vals == {} def test_plus_expands_constant(self): - vi = Variable("i", 1, 10) - a = Tensor.rand(3, 4).reshape(3, vi) - a = a + 1 - assert a.shape == (3, vi) + vi = Variable("i", 1, 5) + for i in range(1, 6): + a = Tensor.rand(3, i).reshape(3, vi) + a = a + 1 + assert a.shape == (3, vi) class TestSymbolicShapeExpr(unittest.TestCase): def test_symbolic_expr_idxs(self): @@ -114,5 +140,23 @@ class TestSymbolicShapeExpr(unittest.TestCase): idx, valid = st.expr_idxs(idx) assert idx.render() == "(((1+i)*1)+(lidx1*((i*4)+4))+gidx0)" +class TestShapeTrackerVarVals(unittest.TestCase): + def test_reshape_reshape_updates_var_vals(self): + vi = Variable("i", 1, 5) + vj = Variable("j", 1, 5) + t = Tensor.rand(3, 4).reshape(3, vi).reshape(4, vj) + assert t.lazydata.st.var_vals == {vi: 4, vj: 3} + + def test_lazy_check_var_vals(self): + vi = Variable("i", 1, 5) + a = Tensor.rand(3, 4).reshape(3, vi) + b = Tensor.rand(5, 6).reshape(vi, 6) + assert a.lazydata.st.var_vals == {vi: 4} + assert b.lazydata.st.var_vals == {vi: 5} + c = a@b + # shapetracker works with symbolic shape and doesn't check / propagate the underlying variable values + assert c.shape == (3, 6) + assert c.lazydata.st.var_vals == {} + if __name__ == '__main__': unittest.main() \ No newline at end of file diff --git a/test/unit/test_symbolic.py b/test/unit/test_symbolic.py index 22c64d2434..d2ec6810c5 100644 --- a/test/unit/test_symbolic.py +++ b/test/unit/test_symbolic.py @@ -1,6 +1,6 @@ #!/usr/bin/env python import unittest -from tinygrad.shape.symbolic import MulNode, SumNode, Variable, NumNode, LtNode, sym_vars, sym_render +from tinygrad.shape.symbolic import MulNode, SumNode, Variable, NumNode, LtNode, sym_render class TestSymbolic(unittest.TestCase): def helper_test_variable(self, v, n, m, s): @@ -261,14 +261,6 @@ class TestSymbolicVars(unittest.TestCase): assert (a % 3 + b // 5).vars() == [a, b] assert (a + b + c - a).vars() == [b, c] - def test_sym_vars(self): - a = Variable("a", 0, 10) - b = Variable("b", 0, 10) - assert sym_vars(1) == [] - assert sym_vars(a) == [a] - assert sym_vars(a+b) == [a, b] - assert sym_vars(a*3) == [a] - class TestSymbolicMinMax(unittest.TestCase): def test_min_max_known(self): a = Variable("a", 1, 8) diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 1fdb63b4a0..99960ad2d1 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -9,7 +9,7 @@ from tinygrad.helpers import GRAPH, DEBUG, prod, getenv, DType, dtypes, flatten, from tinygrad.runtime.ops_cpu import RawNumpyBuffer from tinygrad.runtime.ops_disk import RawDiskBuffer from tinygrad.shape.shapetracker import MovementOps, ShapeTracker, View, get_contraction -from tinygrad.shape.symbolic import Variable, sym_vars +from tinygrad.shape.symbolic import Node from tinygrad.ops import Compiled, Interpreted, UnaryOps, BinaryOps, TernaryOps, ReduceOps, LoadOps, OpType, LazyOp from tinygrad.runtime.lib import RawBufferMapped, RawConst, RawBuffer @@ -214,7 +214,7 @@ class LazyBuffer: if not self.realized and self.op.op == LoadOps.CONTIGUOUS: return self # two CONTIGUOUS in a row is one return create_lazybuffer(self.device, ShapeTracker(self.shape), LoadOps, LazyOp(LoadOps.CONTIGUOUS, (self,), None), self.dtype) - def shuffle_and_prune_movement_ops(self, st: ShapeTracker, op: MovementOps, arg: Union[Tuple[int, ...], Tuple[Tuple[int, int], ...]]) -> LazyBuffer: + def shuffle_and_prune_movement_ops(self, st: ShapeTracker, op: MovementOps, arg: Union[Tuple[Union[Node,int], ...], Tuple[Tuple[int, int], ...]]) -> LazyBuffer: if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and not self.realized and (op in {MovementOps.SHRINK, MovementOps.STRIDE, MovementOps.PERMUTE} or (op == MovementOps.RESHAPE and self.op.op in UnaryOps)) and len(self.children) == 0: return self.op.replace_with_movement_ops([(op, arg)]) ret = create_lazybuffer(self.device, st, MovementOps, LazyOp(op, (self,), arg), self.dtype) @@ -231,13 +231,13 @@ class LazyBuffer: return create_lazybuffer(self.device, ShapeTracker(new_shape), ReduceOps, LazyOp(op, srcs, new_shape), self.dtype) def reduce_op(self:LazyBuffer, op:ReduceOps, new_shape:Tuple[int, ...]) -> LazyBuffer: - if prod(self.shape) // prod(new_shape) < 32768: return self._reduce_op(op, new_shape) # The amount of work should be big enough to take the benefit of "2 kernels" approach. + if any(not isinstance(s, int) for s in self.shape) or prod(self.shape) // prod(new_shape) < 32768: return self._reduce_op(op, new_shape) # The amount of work should be big enough to take the benefit of "2 kernels" approach. heuristic, divisor, dim_to_split = max(((divisor := math.gcd(256, old))/(stride or math.inf), divisor, i) for i, (old, new, stride) in enumerate(zip(self.shape, new_shape, self.st.real_strides())) if old != new) # type: ignore if divisor < 16 or heuristic < 0.125: return self._reduce_op(op, new_shape) # Choose largest divisor (>=16) to split on, penalize large strides. def splitted_shape(dim_aft_div): return self.shape[:dim_to_split] + (self.shape[dim_to_split]//divisor,) + dim_aft_div + self.shape[dim_to_split+1:] return self.reshape(splitted_shape((divisor,)))._reduce_op(op, splitted_shape((1,))).reshape(splitted_shape(()))._reduce_op(op, new_shape) - def reshape(self:LazyBuffer, arg:Tuple[int, ...]) -> LazyBuffer: + def reshape(self:LazyBuffer, arg:Tuple[Union[Node, int], ...]) -> LazyBuffer: if self.shape == arg: return self if not self.realized and self.op.op == MovementOps.RESHAPE: self.op.src[0].children.discard(self) # NOTE: this is only required in reshape and when pushing permutes, why?? @@ -249,7 +249,7 @@ class LazyBuffer: if not self.realized and self.op.op == MovementOps.PAD: return self.op.src[0].pad(tuple([(b1+b2, e1+e2) for (b1,e1),(b2,e2) in zip(self.op.arg, arg)])) return self.shuffle_and_prune_movement_ops(ShapeTracker(self.st).pad(arg), MovementOps.PAD, arg) - def expand(self: LazyBuffer, arg:Tuple[int, ...]) -> LazyBuffer: + def expand(self: LazyBuffer, arg:Tuple[Union[Node,int], ...]) -> LazyBuffer: if self.shape == arg: return self if not self.realized and self.op.op == MovementOps.EXPAND: return self.op.src[0].expand(arg) @@ -293,7 +293,6 @@ class LazyBuffer: def buffers(self) -> Tuple[LazyBuffer, ...]: return (self,) def map_buffers(self, real_srcs: Dict[Any, Any]): return real_srcs.get(self, self) def get_lazyops(self) -> List[Any]: return [] - def get_variable_buffers(self) -> Dict[Variable, LazyBuffer]: return {v:LazyBuffer.loadop(LoadOps.FROM, (1,), dtypes.int32, self.device, src=LazyBuffer.fromCPU(np.array([v.val], dtype=np.int32))) for s in self.shape for v in sym_vars(s)} def replace_with_movement_ops(self: LazyBuffer, ops:List[Tuple[MovementOps, Any]]) -> LazyBuffer: y = self for op, arg in ops: y = MOVEMENT_OPS_DISPATCHER[op](y, arg) diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 7428f453aa..beba5c4f17 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -3,7 +3,7 @@ from __future__ import annotations from enum import Enum, auto import functools from typing import Dict, Tuple, Union, List, Optional, Callable, cast, NamedTuple -from tinygrad.helpers import prod, DEBUG +from tinygrad.helpers import prod, DEBUG, partition from tinygrad.shape.symbolic import Variable, MulNode, NumNode, Node, SumNode, is_sym_int # these ops live here @@ -129,17 +129,18 @@ def get_unsafe_resize_offset(strides, arg): return sum([s * x[0] for s, x in zip(strides,arg)]) class ShapeTracker: - __slots__ = "views" - def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], views:Optional[List[View]]=None): + __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)]) - def __repr__(self): return f"ShapeTracker(shape={self.views[-1].shape}, views={self.views})" + 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]) @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 + def shape(self) -> Tuple[int, ...]: return self.views[-1].shape # NOTE: real type is Tuple[Union[Node, int], ...] but mypy complains about prod(shape) @property def key(self) -> Tuple[View, ...]: return tuple(self.views) @@ -231,15 +232,16 @@ class ShapeTracker: return self def reshape(self, new_shape: Tuple[Union[Node,int], ...]): - # reshape into symbolic shape, update the variable value - if all(isinstance(s, int) for s in self.shape) and len(new_vars:=list(s for s in new_shape if isinstance(s, Variable))) > 0: - assert len(new_vars) == 1, "only one variable is supported in a shape" - new_var, new_val = new_vars[0], prod(self.shape) // prod(s for s in new_shape if isinstance(s, int)) - if new_var.val is None: + 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}]" - new_var.val = new_val - else: assert new_var.val == new_val, f"value conflicts, was {new_var.val}, set to {new_val}" - + 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 = {} if self.views[-1].shape == new_shape: return self 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 diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index b5fb04a2d3..5d5e866015 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -9,7 +9,6 @@ from typing import List, Dict, Callable, Tuple, Type, Union, Optional, Any # symbolic matches the Python behavior, but the code output is agnostic, and will never have negative numbers in div or mod def is_sym_int(x: Any) -> bool: return isinstance(x, (int, Node)) -def sym_vars(x: Union[Node, int]) -> List[Variable]: return [] if isinstance(x, int) else x.vars() class Node: b: Union[Node, int] @@ -141,7 +140,6 @@ class Variable(Node): def __init__(self, expr:Optional[str], nmin:int, nmax:int): self.expr, self.min, self.max = expr, nmin, nmax - self.val: Optional[int] = None def vars(self): return [self] class NumNode(Node):