diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 5dc965ca88..0ff635bc3b 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -368,5 +368,16 @@ class TestUOpRender(unittest.TestCase): u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) self.assertEqual(u.render(), "{0,1,2}") +class TestContiguousViewOffset(unittest.TestCase): + def _check(self, u, expected): self.assertEqual(u.contiguous_view_offset(), expected) + + def test_simple(self): self._check(UOp.empty(10), 0) + def test_shrink(self): self._check(UOp.empty(10)[1:8], 1) + def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7) + def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1) + def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None) + def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None) + def test_strided(self): self._check(UOp.empty(4)[::2], None) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 56e3e511fc..8438e2f147 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -829,17 +829,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # by relevant CL runtimes at time of writing. if any(d.startswith(("WEBGPU", "CL")) for d in ((self.device,) if isinstance(self.device, str) else self.device)): return None - numel = self.numel() - out = graph_rewrite(self.flatten().index(UOp.range(numel, 0)), pm_mops+symbolic, name="contiguous_view_offset") - if out.op is not Ops.INDEX: return None - if len(out.src) == 1: return 0 if resolve(numel == 1, False) else None - idx, has_range = out.src[1], False - if idx.op is Ops.RANGE: return 0 - if idx.op is Ops.ADD and idx.src[0].op is Ops.RANGE: idx, has_range = idx.src[1], True - if idx.op is Ops.CONST and (has_range or resolve(numel == 1, False)): - if not isinstance(idx.arg, int): return None # masked/padded regions produce InvalidType - return idx.arg - return None + idx = self.flatten().index(UOp.range(self.numel(), 0)) + out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") + return out.arg if out.op is Ops.CONST and isinstance(out.arg, int) else None def has_buffer_identity(self): """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain).""" @@ -1719,6 +1711,14 @@ def do_unbind(ctx:dict[Variable, int], x:UOp): return v pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)]) +# ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset +pm_contiguous_view_offset = PatternMatcher([ + (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(dtypes.index, 0)), + (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(dtypes.index, 0)), + (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c), + (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None), +]) + # *** what was symbolic.py *** sint = int|UOp