From 0b58ff679cecf5c15a41dbde291cfaf410002881 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 29 Aug 2026 19:18:11 -0400 Subject: [PATCH] cannot decomp long Variable (#17842) raise instead of silent error. also fix negative WEBGPU variable --- test/backend/test_tensor_variable.py | 8 ++++---- tinygrad/codegen/decomp/dtype.py | 10 +++++++--- tinygrad/runtime/ops_webgpu.py | 2 +- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/test/backend/test_tensor_variable.py b/test/backend/test_tensor_variable.py index efab3b0aab..834d8c998e 100644 --- a/test/backend/test_tensor_variable.py +++ b/test/backend/test_tensor_variable.py @@ -41,10 +41,8 @@ class TestTensorVariable(unittest.TestCase): self.assertEqual(Tensor(2**40).item(), 2**40) self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35+3)).item(), 2**35+3) - @unittest.expectedFailure - def test_long_variable_emulated(self): - # the long decomp splits the variable into two int PARAMs that share its name, so both bind the full value and truncate - with Context(EMULATED_DTYPES="long"): self.assertEqual(Tensor(Variable("c", 0, 2**40).bind(2**35+3)).item(), 2**35+3) + def test_long_variable_emulated_raises(self): + with Context(EMULATED_DTYPES="long"), self.assertRaises(RuntimeError): Tensor(Variable("c", 0, 2**40).bind(2**35+3)).item() def test_variable_tensor_dtype_arg(self): vv = Variable("a", 1, 10).bind(2) @@ -59,6 +57,8 @@ class TestTensorVariable(unittest.TestCase): # bound variables in an expression are fine self.assertEqual(Tensor(Variable("u", 1, 10).bind(2) + 1).item(), 3) + def test_negative_variable_on_device(self): self.assertEqual(Tensor(Variable("n", -10, 10).bind(-3)).clone().item(), -3) + def test_shrink_beyond_buffer_variable(self): # TODO: shrink by a variable whose vmax exceeds the dim should fail at build, today only CHECK_OOB=1 rejects it t = Tensor.ones(3).contiguous()[:Variable("a", 1, 10).bind(5)] diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index cde69c6f75..f02e781d06 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -1,5 +1,5 @@ from dataclasses import replace -from tinygrad.dtype import dtypes, DType, truncate +from tinygrad.dtype import dtypes, DType, AddrSpace, truncate from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES from tinygrad.uop import GroupOp from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite @@ -80,6 +80,11 @@ def l2i(op: Ops, dt: DType, *uops:UOp): case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1) case _: raise NotImplementedError(f"long decomposition of {op} unsupported") +def l2i_define(x:UOp) -> UOp: + # cannot decomp a Variable + if x.addrspace == AddrSpace.ALU: raise RuntimeError(f"long decomposition of variable {x.arg.name} unsupported") + return UOp(x.op, arg=replace(x.arg, dtype=l2i_dt[x.dtype], size=None if x.arg.size is None else x.arg.size*2), tag=x.tag) + def split_l2i(ctx:dict, op: Ops, dt: DType, *uops:UOp): # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself. # both word halves of a node ask for the same split, so ctx memos it for the pass @@ -140,8 +145,7 @@ def f2f_store(st, idx, val, fr:DType, to:DType): pm_long_decomp: PatternMatcher = PatternMatcher([ # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))), - (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), name="x"), lambda x: - UOp(x.op, arg=replace(x.arg, dtype=l2i_dt[x.dtype], size=None if x.arg.size is None else x.arg.size*2), tag=x.tag)), + (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), name="x"), l2i_define), (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None), (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val: diff --git a/tinygrad/runtime/ops_webgpu.py b/tinygrad/runtime/ops_webgpu.py index dc3466e4ac..de2103f281 100644 --- a/tinygrad/runtime/ops_webgpu.py +++ b/tinygrad/runtime/ops_webgpu.py @@ -201,7 +201,7 @@ class WebGpuDevice(Compiled): def create_uniform(self, val:int|float) -> webgpu.WGPUBuffer: buf = webgpu.wgpuDeviceCreateBuffer(self.device_res, webgpu.WGPUBufferDescriptor(size=4, usage=webgpu.WGPUBufferUsage_Uniform | webgpu.WGPUBufferUsage_CopyDst)) - self.write_buffer(buf, val.to_bytes(4, "little") if isinstance(val, int) else struct.pack(' webgpu.WGPUBuffer: size = webgpu.wgpuBufferGetSize(buf)