mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-30 20:26:07 +00:00
ShapeTracker.var_vals (#1540)
This commit is contained in:
chenyu
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
+5
-6
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user