diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 61642bc795..1f2128737c 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -404,7 +404,7 @@ jobs: testcommalatest: name: comma Benchmark (0.11.0) runs-on: [self-hosted, Linux, comma] - timeout-minutes: 10 + timeout-minutes: 12 defaults: run: shell: bash -e -o pipefail {0} @@ -435,7 +435,7 @@ jobs: testcommaold: name: comma Benchmark (0.10.1) runs-on: [self-hosted, Linux, comma] - timeout-minutes: 10 + timeout-minutes: 12 defaults: run: shell: bash -e -o pipefail {0} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index dbf55f991c..cca9359d25 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -291,7 +291,7 @@ jobs: llvm: 'true' - name: Test openpilot model kernel count and gate usage run: | - ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=55 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 + ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1391 ALLOWED_GATED_READ_IMAGE=58 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916 # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630 - name: Test openpilot CL compile fp32 (test correctness) run: | diff --git a/spec/tinyspec.pdf b/spec/tinyspec.pdf index 852aec4f4f..f8b3b7cb9d 100644 Binary files a/spec/tinyspec.pdf and b/spec/tinyspec.pdf differ diff --git a/spec/tinyspec.tex b/spec/tinyspec.tex index a9dda1c071..1ca4d333ca 100644 --- a/spec/tinyspec.tex +++ b/spec/tinyspec.tex @@ -73,7 +73,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op \op{Permute} & $(T,)$ & axis order $\pi$ & Reorder axes. $\pi = (1,0)$ is transpose. \\ \op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\ \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\ -\op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\ +\op{Expand} & $(T, \mathbf{n})$ & --- & Prepend axes $\mathbf{n}$ on the left. Output shape is $\mathbf{n} + T.\mathrm{shape}$. \\ \op{Pad} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Place $T$ at offset $o_k$ in an invalid-filled output of shape $s'_k$. \\ \op{Shrink} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Keep $s'_k$ elements starting at offset $o_k$ per axis. Inverse of \op{Pad}. \\ \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\ @@ -284,7 +284,7 @@ Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$. \medskip \textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim). -\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. +\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$. \op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}. %% ============================================================ diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 40c6fd3c6b..2ab0b5b66a 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -252,8 +252,8 @@ class TestViz(unittest.TestCase): self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)]) def test_const_reshape_expand_folded(self): - # CONST->RESHAPE->EXPAND should be folded into the ALU node, not shown as separate RESHAPE/EXPAND nodes - c = UOp.const(dtypes.float, 1.0, shape=(3,4)) # creates CONST->RESHAPE->EXPAND chain + # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes + c = UOp.const(dtypes.float, 1.0, shape=(3,4)) # creates CONST->EXPAND chain a = UOp.variable("a", 0.0, 10.0, dtypes.float) alu = a + c with save_viz() as viz: @@ -262,7 +262,6 @@ class TestViz(unittest.TestCase): excluded_nodes = {v["label"].split("\n")[0] for v in graph.values() if v["exclude"]} self.assertIn("CONST", excluded_nodes) self.assertIn("STACK", excluded_nodes) - self.assertIn("RESHAPE", excluded_nodes) self.assertIn("EXPAND", excluded_nodes) self.assertIn("CONST1 1", graph[id(alu)]["label"]) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 86a7e764cf..013e734aa3 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -172,9 +172,9 @@ devectorizer2 = pm_mops+PatternMatcher([ (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]), # reshape of a single element shaped value to scalar is an index (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None), - # RESHAPE+EXPAND -> STACK - (UPat(Ops.EXPAND, src=(UPat(Ops.RESHAPE, src=(UPat.var("x"), UPat())), UPat()), name="out"), - lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if out.shape == (out.max_numel(),) else None), + # EXPAND on scalar -> STACK + (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"), + lambda x,out: UOp.vectorize(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None), # INDEX on INDEX is INDEX (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"), lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])), diff --git a/tinygrad/mixin/gradient.py b/tinygrad/mixin/gradient.py index aece6850d8..bbafd38b67 100644 --- a/tinygrad/mixin/gradient.py +++ b/tinygrad/mixin/gradient.py @@ -75,7 +75,7 @@ pm_gradient = PatternMatcher([ (UPat(Ops.CONTIGUOUS_BACKWARD), lambda ctx: (ctx.contiguous(),)), (UPat(Ops.RESHAPE, name="ret"), lambda ctx, ret: (ctx.reshape(ret.src[0].shape), None)), (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: - (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(i for i,(s,n) in enumerate(zip(ret.src[0].shape, ret.shape)) if s!=n)) + (ctx.cast(sum_acc_dtype(ctx.dtype))._rop(Ops.ADD, tuple(range(len(ret.marg)))) .reshape(ret.src[0].shape).cast(ctx.dtype), None)), (UPat(Ops.PAD, name="ret"), lambda ctx, ret: (ctx.shrink(tuple([(p[0], s+p[0]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), (UPat(Ops.SHRINK, name="ret"), lambda ctx, ret: (ctx.pad(tuple([(p[0], s-p[0]-p[1]) for s,p in zip(ret.src[0].shape, ret.marg)])), None, None)), diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 2562e3c330..cc36b3f392 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -123,9 +123,16 @@ class MovementMixin: # for each dimension, check either dim is 1, or it does not change if not all(s == ns or s == 1 for s, ns in zip(shape, new_shape)): raise ValueError(f"cannot broadcast {self.shape} to {new_shape=}") - reshaped = self.reshape(shape) - ret = reshaped._mop(Ops.EXPAND, arg=new_shape) - return reshaped if ret.shape == reshaped.shape else ret + # EXPAND only adds dims on the left. squeeze 1s that need expanding, EXPAND on left, permute back. + n_left = len(new_shape) - len(self.shape) + expand_at = tuple(i for i, s in enumerate(self.shape) if resolve(s == 1, default=False) and resolve(new_shape[n_left+i] != 1)) + kept = tuple(i for i in range(len(self.shape)) if i not in expand_at) + squeezed = self.reshape(tuple(self.shape[i] for i in kept)) + expanded = squeezed._mop(Ops.EXPAND, arg=new_shape[:n_left] + tuple(new_shape[n_left+i] for i in expand_at)) + # expanded shape = [left] + [expand_at dims] + [kept dims], permute to new_shape + perm = tuple(range(n_left)) + tuple( + n_left + (expand_at.index(i) if i in expand_at else len(expand_at) + kept.index(i)) for i in range(len(self.shape))) + return expanded.permute(perm) def expand(self, shape, *args) -> Self: """ diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index 7c69a1bc52..043fe2a7c8 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -134,7 +134,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO case Ops.SHRINK: rngs = tuple(a if off == 0 else a+off for a,(off,_) in zip(rngs, arg)) case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg)) case Ops.FLIP: rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg)) - case Ops.EXPAND: rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg)) + case Ops.EXPAND: rngs = rngs[len(arg):] case Ops.PAD: # NOTE: the .where(r-s, i) is not inside the graph_rewrite so that `convert_pad_to_where_to_keep_behavior_local` # wraps the pad with only the newly added valid @@ -248,7 +248,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do. # NOTE: this doesn't actually always end a range, but this is why convs are realized, so for now we need it if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): - ending_ranges[x] += list(UOp.sink(*[ro for ri, ro in zip(rngs, out_rngs) if ri is not ro]).ranges.keys()) + ending_ranges[x] += list(UOp.sink(*out_rngs[:len(x.marg)]).ranges.keys()) # REDUCE creates ranges for the axes it is reducing if x.op is Ops.REDUCE and x.arg[1]: diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 79cd0bee41..fbad0f8b2b 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -84,9 +84,8 @@ def reshape_multi(root:UOp, multi:UOp): return multi.src[0].reshape(new_shape).multi(new_axis) def expand_multi(root:UOp, multi:UOp): - if multi.axis is None: new_shape = root.marg - else: new_shape = tuple(multi.src[0].shape[multi.axis] if a == multi.axis else s for a,s in enumerate(root.marg)) - return multi.src[0].expand(new_shape).multi(multi.axis) + new_axis = None if multi.axis is None else multi.axis + len(root.marg) + return multi.src[0]._mop(Ops.EXPAND, arg=root.marg).multi(new_axis) def pad_multi(root:UOp, multi:UOp): assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}" diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index f697bdb6cb..f5a2e93745 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -322,9 +322,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if prod(ps) != prod(self.marg): raise ValueError(f"bad reshape: {ps} -> {self.marg}") return self.marg case Ops.EXPAND: - if len(ps) != len(self.marg) or not all(s==ns or (s==1 and ns>=0) for s,ns in zip(ps, self.marg)): - raise ValueError(f"bad expand: {ps} -> {self.marg}") - return self.marg + return tuple(self.marg) + ps case Ops.PERMUTE: if sorted(self.marg) != list(range(len(ps))): raise ValueError(f"invalid permutation {self.marg} of len {len(ps)}") return tuple(ps[i] for i in self.marg) @@ -548,7 +546,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): ret = UOp.vectorize(*stk) else: ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) - return ret.reshape((1,)*len(shape)).expand(shape) if shape is not None and shape != () and ret.shape != shape else ret + return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret @staticmethod def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs) @@ -637,6 +635,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.shape[new_axis] % len(self.device) != 0: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards") return new_axis if self.op is Ops.PERMUTE: return self.marg.index(src_axis) if src_axis is not None else None + if self.op is Ops.EXPAND: return src_axis + len(self.marg) if src_axis is not None else None return src_axis def _unshard(self, axis:int) -> UOp: @@ -688,7 +687,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def _mop(self, op:Ops, arg) -> UOp: # early NOOP - if op in {Ops.SHRINK, Ops.PAD, Ops.EXPAND} and len(arg) == 0: + if op is Ops.EXPAND and len(arg) == 0: return self + if op in {Ops.SHRINK, Ops.PAD} and len(arg) == 0: assert len(self.shape) == 0, "0 len arg only valid on zero length shape" return self match op: