From e694d7f2226c176b4cacea3210a2bd03adae5cdd Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 18 May 2026 23:14:12 -0400 Subject: [PATCH] more deviceless const prerequisites [pr] (#16256) * more deviceless const prerequisites [pr] * remove that * arange.contiguous -> arange.clone in tests arange will become deviceless const soon, update tests where it needs to be a buffer --- test/backend/test_multitensor.py | 16 ++++++------ test/backend/test_schedule.py | 12 ++++----- test/backend/test_tensor.py | 10 -------- test/null/test_schedule.py | 42 ++++++++++++++++---------------- tinygrad/mixin/__init__.py | 12 ++++++--- tinygrad/tensor.py | 3 ++- 6 files changed, 45 insertions(+), 50 deletions(-) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index bb0321e4c2..c68ac473ca 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -1174,23 +1174,23 @@ class TestMultiBufferView(unittest.TestCase): self._check(ref, a, lambda t: t[3]) def test_shrink_2d(self): - ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize() - a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(6*4).reshape(6, 4).clone().realize() + a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.shrink(((1, 4), None))) def test_reshape_then_shrink(self): - ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize() - a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(8*6).reshape(8, 6).clone().realize() + a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1]) def test_chained_shrink(self): - ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize() - a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize() + ref = Tensor.arange(10*8).reshape(10, 8).clone().realize() + a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices_2, axis=1).realize() self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None))) def test_4_devices(self): - ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize() - a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize() + ref = Tensor.arange(8*12).reshape(8, 12).clone().realize() + a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize() out = a[5].contiguous() linear, var_vals = out.linear_with_vars() if all(hasattr(Device[d].allocator, "_offset") for d in out.device): diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index cf327cc06c..6629a1479f 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -966,7 +966,7 @@ class TestSchedule(unittest.TestCase): def test_arange_index_contiguous(self): Tensor.manual_seed(0) x = Tensor.randn(5, 2).realize() - a = Tensor.arange(10).contiguous() + a = Tensor.arange(10).clone() out = (x + a[2]).sum() run_linear(*check_schedule(out, 2)) np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6) @@ -982,7 +982,7 @@ class TestSchedule(unittest.TestCase): def test_user_contiguous(self): Tensor.manual_seed(0) x = Tensor.randn(5, 2).realize() - a = (Tensor.arange(10)+1).contiguous() + a = (Tensor.arange(10)+1).clone() out = (x + a[2]).sum() run_linear(*check_schedule(out, 2)) np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6) @@ -1008,7 +1008,7 @@ class TestSchedule(unittest.TestCase): def test_fuse_assign_contiguous(self): x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize() a = Tensor.arange(8).reshape(4, 2) - run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2)) + run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.clone()), 2)) np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]]) def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True) @@ -1053,7 +1053,7 @@ class TestSchedule(unittest.TestCase): def test_no_extra_contiguous_on_setitem_assign_back(self): # pattern: contiguous copy, advanced setitem, assign back (e.g. torch backend _view_write) - base = Tensor.arange(16).reshape(4, 4).contiguous() + base = Tensor.arange(16).reshape(4, 4).clone() flat_base = base.reshape(16).contiguous() idx = Tensor([1,2,5,6], dtype=dtypes.int32) flat_base[idx] = Tensor([99,99,99,99]) @@ -1253,7 +1253,7 @@ class TestView(unittest.TestCase): # x collapses along with its children def test_parent_view_collapses(self): a = Tensor([1, 2]) - b = Tensor.arange(3).contiguous() + b = Tensor.arange(3).clone() bv = b.pad(((0, 2),))[-2:] # this becomes a late a*0 late_mul = a*bv @@ -1270,7 +1270,7 @@ class TestView(unittest.TestCase): # as long as one child realizes, x does not collapse def test_parent_multiple_children_no_collapse(self): a = Tensor([1, 2]) - b = Tensor.arange(3).contiguous() + b = Tensor.arange(3).clone() bv = b.pad(((0, 2),))[-2:] late_mul = a*bv other_child = b+2 diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index b666b02799..3249a6a3fe 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -36,16 +36,6 @@ class TestTinygrad(unittest.TestCase): self.assertTrue(t.uop.has_buffer_identity()) np.testing.assert_equal(t.numpy(), 2.0) - def test_to_deviceless_const(self): - t = Tensor(UOp.const(dtypes.float, 2.0)) - self.assertIs(t.to(f"{Device.DEFAULT}:1"), t) - self.assertIs(t.to_(f"{Device.DEFAULT}:1"), t) - - def test_shard_deviceless_const(self): - t = Tensor(UOp.const(dtypes.float, 2.0)) - self.assertIs(t.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")), t) - self.assertIs(t.shard_((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")), t) - def test_plus_equals(self): a = Tensor.randn(10,10) b = Tensor.randn(10,10) diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index dad8dbdc7d..00324bc328 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -1216,17 +1216,17 @@ class TestFusionOp(unittest.TestCase): class TestBufferView(unittest.TestCase): def test_shrink_contiguous_is_buffer_view(self): # simple 1D shrink of a realized buffer should be BUFFER_VIEW, not a copy kernel - a = Tensor.arange(100).contiguous().realize() + a = Tensor.arange(100).clone().realize() b = a.shrink(((10, 50),)).contiguous() run_linear(*check_schedule(b, 0)) def test_shrink_2d_contiguous_is_buffer_view(self): - a = Tensor.arange(100).reshape(10,10).contiguous().realize() + a = Tensor.arange(100).reshape(10,10).clone().realize() b = a.shrink(((1, 5),None)).contiguous() run_linear(*check_schedule(b, 0)) def test_chained_shrink_is_buffer_view(self): - a = Tensor.arange(1000).contiguous().realize() + a = Tensor.arange(1000).clone().realize() b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous() run_linear(*check_schedule(b, 0)) @@ -1234,96 +1234,96 @@ class TestBufferView(unittest.TestCase): # indexing a non-shard axis of a realized sharded tensor should be BUFFER_VIEW on each device, not copy kernels # this is the flat_llama pattern: weight[layer_idx] where weight is (n_layers, out, dim) sharded on axis=1 devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_shrink_2d_non_shard_axis_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.shrink(((1, 4), None)).contiguous(), 0)) def test_shrink_shard_axis_0_multi(self): # shrinking a middle dim is not contiguous per shard, so this needs copy kernels devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=0).realize() + a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=0).realize() run_linear(*check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2)) def test_reshape_then_shrink_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0)) def test_permute_then_shrink_multi(self): # permute makes per-shard view non-contiguous, needs copy kernels devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*6*2).reshape(4, 6, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2)) def test_multi_buffer_view_4_devices(self): devices = tuple(f"NULL:{i}" for i in range(4)) - a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a[5].contiguous(), 0)) def test_chained_shrink_multi(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0)) # negative tests: these should NOT become BUFFER_VIEW (non-contiguous per shard) def test_expand_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 1, 2).contiguous().shard(devices, axis=2).realize() + a = Tensor.arange(4*2).reshape(4, 1, 2).clone().shard(devices, axis=2).realize() run_linear(*check_schedule(a.expand(4, 3, 2).contiguous(), 2)) def test_pad_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2)) def test_flip_multi_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize() + a = Tensor.arange(4*2).reshape(4, 2).clone().shard(devices, axis=1).realize() run_linear(*check_schedule(a.flip(0).contiguous(), 2)) def test_replicated_reshape_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).contiguous().to(devices).realize() + a = Tensor.arange(24).clone().to(devices).realize() run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0)) def test_replicated_shrink_is_buffer_view(self): # DP pattern: replicated weight[layer_idx] devices = ("NULL:1", "NULL:2") - a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize() + a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_replicated_chained_mops_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(100).contiguous().to(devices).realize() + a = Tensor.arange(100).clone().to(devices).realize() run_linear(*check_schedule(a.reshape(10, 10).shrink(((2, 7), None)).contiguous(), 0)) def test_replicated_shard_none_is_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).contiguous().shard(devices, axis=None).realize() + a = Tensor.arange(24).clone().shard(devices, axis=None).realize() run_linear(*check_schedule(a.reshape(4, 6).contiguous(), 0)) def test_replicated_4_devices_is_buffer_view(self): devices = tuple(f"NULL:{i}" for i in range(4)) - a = Tensor.arange(8*10).reshape(8, 10).contiguous().to(devices).realize() + a = Tensor.arange(8*10).reshape(8, 10).clone().to(devices).realize() run_linear(*check_schedule(a[3].contiguous(), 0)) def test_replicated_expand_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(12).reshape(4, 1, 3).contiguous().to(devices).realize() + a = Tensor.arange(12).reshape(4, 1, 3).clone().to(devices).realize() run_linear(*check_schedule(a.expand(4, 3, 3).contiguous(), 2)) def test_replicated_permute_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize() + a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize() run_linear(*check_schedule(a.permute(1, 0).contiguous(), 2)) def test_replicated_flip_not_buffer_view(self): devices = ("NULL:1", "NULL:2") - a = Tensor.arange(24).reshape(4, 6).contiguous().to(devices).realize() + a = Tensor.arange(24).reshape(4, 6).clone().to(devices).realize() run_linear(*check_schedule(a.flip(0).contiguous(), 2)) class TestInvalidTensor(unittest.TestCase): diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 10c4f0689d..3984627cbe 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -935,7 +935,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin): print(t.gather(1, Tensor([[0, 0], [1, 0]])).numpy()) ``` """ - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") if index.ndim != self.ndim: raise RuntimeError(f"self.ndim must equal index.ndim, {self.ndim=}, {index.ndim=}") dim = self._resolve_dim(dim) assert all(s >= i for d,(s,i) in enumerate(zip(self.shape, index.shape)) if d != dim), "requires self.shape[d] >= index.shape[d] for all d != dim" @@ -975,8 +976,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin): return x.cast(self.dtype) def _pre_scatter(self, dim:int, index:Self, src:Self) -> tuple[Self, Self]: - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") - if src.device != self.device: raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if src.device is not None and self.device is not None and src.device != self.device: + raise RuntimeError(f"expected src and self on the same device, {src.device=}, {self.device=}") dim = self._resolve_dim(dim) assert index.ndim == self.ndim == src.ndim, f"self.ndim, index.ndim and src.ndim must all equal, {self.ndim=} {index.ndim=} {src.ndim=}" assert all((d == dim or self_ >= index_) and src_ >= index_ for d,(self_,index_,src_) in enumerate(zip(self.shape, index.shape, src.shape))), \ @@ -1393,7 +1396,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin): ``` """ assert 0.0 <= label_smoothing <= 1.0, "label_smoothing must be in [0.0, 1.0]" - if Y.device != self.device: raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}") + if Y.device is not None and self.device is not None and Y.device != self.device: + raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}") log_probs = self.log_softmax() loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.ones_like(dtype=dtypes.bool) y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index aac0bb28f1..909e308599 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -887,7 +887,8 @@ class Tensor(OpMixin): match index: case Tensor(): if not dtypes.is_int(index.dtype): raise IndexError(f"index dtype {index.dtype} is not supported") - if index.device != self.device: raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") + if index.device is not None and self.device is not None and index.device != self.device: + raise RuntimeError(f"expected index and self on the same device, {index.device=}, {self.device=}") assert isinstance(size, int), "size must be an int" index = (index < 0).where(index+size, index) # treat negative index values case list() | tuple():