From e78e8722dca29365138c33f2b260070c2d7efb92 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 2 Apr 2025 15:31:32 +0800 Subject: [PATCH 1/4] Revert "LDS noop and spec (#9669)" (#9691) This reverts commit 870b545aceabfead4c6f61517854699d293b3de3. Co-authored-by: Ignacio Sica --- test/test_linearizer.py | 225 +--------------------------------- tinygrad/codegen/kernel.py | 18 +-- tinygrad/renderer/__init__.py | 2 +- 3 files changed, 6 insertions(+), 239 deletions(-) diff --git a/test/test_linearizer.py b/test/test_linearizer.py index c79a35bf1c..3871b88a4d 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -2218,228 +2218,5 @@ class TestKernelOpts(unittest.TestCase): ] helper_linearizer_opt(r, [x[0] for x in opts_shapes], color_sizes=[x[1] for x in opts_shapes]) -def helper_lds_allclose(opts:list[Opt], expected_bufs, N=16, M=16, K=16, dtype_in=dtypes.float, acc_dtype=dtypes.float): - with Context(DEBUG=0): a, b = Tensor.rand(M, K, dtype=dtype_in).realize(), Tensor.rand(K, N, dtype=dtype_in).realize() - realized_ast, bufs = helper_realized_ast(a.matmul(b, dtype=acc_dtype)) - k = Kernel(realized_ast) - for opt in opts: - k.apply_opt(opt) - prg = k.to_program() - CompiledRunner(replace(prg, device=Device.DEFAULT)).exec(bufs) - - atol, rtol = 1e-4, 1e-4 - if dtype_in == dtypes.half: atol, rtol = 1e-2, 1e-2 - np.testing.assert_allclose(bufs[0].numpy().reshape((M,N)), a.numpy() @ b.numpy(), atol=atol, rtol=rtol) - - local_buffers = [uop for uop in k.uops if uop.op is Ops.DEFINE_LOCAL] - assert len(local_buffers) == len(expected_bufs), f"Expected exactly {len(expected_bufs)} local buffers, got {len(local_buffers)}" - for i,(buf, sz) in enumerate(expected_bufs): - assert local_buffers[i].arg == buf, f"Expected buffer argument index {buf}, got {local_buffers[i].arg}" - expected_dtype = (acc_dtype if buf == 0 else dtype_in).ptr(sz, local=True) - assert local_buffers[i].dtype == expected_dtype, f"Expected buffer dtype {expected_dtype}, got {local_buffers[i].dtype} for {opts=}" - # TODO: check all access to the global buffer are proxied through the local buffer - -@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared") -class TestLDS(unittest.TestCase): - # lds tile size for inputs are the same size as the memory accessed by each thread inside the reduce loop - # test no reshape opt after lds? true for lds_swap - # test TC3? - - def test_lds_args(self): - realized_ast, _ = helper_realized_ast(Tensor.rand(4, 4) @ Tensor.rand(4, 4)) - k = Kernel(realized_ast) - valid_opts = [Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - for opt in valid_opts: - k = Kernel(realized_ast) - k.apply_opt(opt) - - invalid_opts = [Opt(OptOps.LDS, -1, None), - Opt(OptOps.LDS, 3, None)] - for opt in invalid_opts: - k = Kernel(realized_ast) - with self.assertRaises(KernelOptError): - k.apply_opt(opt) - - @unittest.expectedFailure - def test_lds_output_basic(self): - helper_lds_allclose(opts=[Opt(OptOps.LDS, 0, None)], expected_bufs=[(0,1)]) - - @unittest.expectedFailure - def test_lds_input_basic(self): - helper_lds_allclose(opts=[Opt(OptOps.LDS, 1, None)], expected_bufs=[(1,1)]) - helper_lds_allclose(opts=[Opt(OptOps.LDS, 2, None)], expected_bufs=[(2,1)]) - - @unittest.expectedFailure - def test_lds_multi_basic(self): - helper_lds_allclose(opts=[Opt(OptOps.LDS, 0, None), Opt(OptOps.LDS, 1, None)], expected_bufs=[(0,1),(1,1)]) - helper_lds_allclose(opts=[Opt(OptOps.LDS, 0, None), Opt(OptOps.LDS, 1, None), Opt(OptOps.LDS, 2, None)], expected_bufs=[(0,1),(1,1),(2,1)]) - - @unittest.expectedFailure - def test_lds_unroll(self): - # unroll doesn't change local output buffer size - for sz in [2,4,8]: - helper_lds_allclose(opts=[Opt(OptOps.UNROLL, 0, sz), Opt(OptOps.LDS, 0, None)], expected_bufs=[(0,1)]) - helper_lds_allclose(opts=[Opt(OptOps.UNROLL, 0, sz), Opt(OptOps.LDS, 1, None)], expected_bufs=[(1,sz)]) - helper_lds_allclose(opts=[Opt(OptOps.UNROLL, 0, sz), Opt(OptOps.LDS, 2, None)], expected_bufs=[(2,sz)]) - - @unittest.expectedFailure - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") - def test_lds_local(self): - # if only locals are applied, local buffer size for output should be prod(locals) - - basic_local_opts = [Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=basic_local_opts, expected_bufs=[(0,2),(1,2),(2,1)]) - - multi_local_opts = [Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.LOCAL, 0, 8), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=multi_local_opts, expected_bufs=[(0,16),(1,16),(2,1)]) - - multi_axis_local_opts = [Opt(OptOps.LOCAL, 1, 4), - Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=multi_axis_local_opts, expected_bufs=[(0,8),(1,2),(2,4)]) - - full_local_opts = [Opt(OptOps.LOCAL, 0, 16), - Opt(OptOps.LOCAL, 0, 16), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=full_local_opts, expected_bufs=[(0,256),(1,16),(2,16)]) - - @unittest.expectedFailure - def test_lds_upcast(self): - # if only upcasts are applied, local buffer size for output should be prod(upcast) - - basic_upcast_opts = [Opt(OptOps.UPCAST, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=basic_upcast_opts, expected_bufs=[(0,2),(1,2),(2,1)]) - - multi_upcast_opts = [Opt(OptOps.UPCAST, 0, 2), - Opt(OptOps.UPCAST, 0, 8), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=multi_upcast_opts, expected_bufs=[(0,16),(1,16),(2,1)]) - - multi_axis_upcast_opts = [Opt(OptOps.UPCAST, 1, 4), - Opt(OptOps.UPCAST, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=multi_axis_upcast_opts, expected_bufs=[(0,8),(1,2),(2,4)]) - - full_upcast_opts = [Opt(OptOps.UPCAST, 0, 16), - Opt(OptOps.UPCAST, 0, 16), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=full_upcast_opts, expected_bufs=[(0,256),(1,16),(2,16)]) - - @unittest.expectedFailure - @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") - def test_lds_tc(self): - for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if tc.dtype_in == dtypes.bfloat16 or tc.dtype_out == dtypes.bfloat16: continue - (N, M, K) = tc.dims - opts = [Opt(OptOps.TC, 0, (-1, 0)), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,N*M),(1,M*K),(2,K*N)], N=N, M=M, K=K, dtype_in=tc.dtype_in, acc_dtype=tc.dtype_out) - - opts = [Opt(OptOps.TC, 0, (-1, 0)), - Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.UPCAST, 1, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,N*M*4),(1,M*K*2),(2,K*N*2)], N=N*4, M=M*4, K=K*4, dtype_in=tc.dtype_in, acc_dtype=tc.dtype_out) - - opts = [Opt(OptOps.TC, 0, (-1, 0)), - Opt(OptOps.UNROLL, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,N*M),(1,M*K*2),(2,K*N*2)], N=N*4, M=M*4, K=K*4, dtype_in=tc.dtype_in, acc_dtype=tc.dtype_out) - - opts = [Opt(OptOps.TC, 0, (-1, 0)), - Opt(OptOps.UNROLL, 0, 2), - Opt(OptOps.UPCAST, 1, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,N*M*2),(1,M*K*2),(2,K*N*4)], N=N*4, M=M*4, K=K*4, dtype_in=tc.dtype_in, acc_dtype=tc.dtype_out) - - @unittest.expectedFailure - @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") - def test_lds_tc_padded(self): - for tc in Device[Device.DEFAULT].renderer.tensor_cores: - if tc.dtype_in == dtypes.bfloat16 or tc.dtype_out == dtypes.bfloat16: continue - (N, M, K) = tc.dims - opts = [Opt(OptOps.TC, 0, (-1, 2)), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,N*M),(1,M*K),(2,K*N)], N=N+3, M=M+3, K=K+3, dtype_in=tc.dtype_in, acc_dtype=tc.dtype_out) - - @unittest.expectedFailure - @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals") - def test_lds_full(self): - opts = [Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.UPCAST, 1, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,4),(1,2),(2,2)]) - - opts = [Opt(OptOps.LOCAL, 0, 2), - Opt(OptOps.UPCAST, 0, 4), - Opt(OptOps.LOCAL, 1, 8), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,64),(1,8),(2,8)]) - - opts = [Opt(OptOps.LOCAL, 0, 16), - Opt(OptOps.UPCAST, 1, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,16),(1,16),(2,1)]) - - opts = [Opt(OptOps.LOCAL, 0, 16), - Opt(OptOps.UPCAST, 0, 16), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,256),(1,16),(2,16)]) - - opts = [Opt(OptOps.LOCAL, 1, 16), - Opt(OptOps.UPCAST, 1, 16), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,16),(1,1),(2,16)]) - - opts = [Opt(OptOps.LOCAL, 1, 4), - Opt(OptOps.UNROLL, 0, 2), - Opt(OptOps.UPCAST, 0, 2), - Opt(OptOps.LDS, 0, None), - Opt(OptOps.LDS, 1, None), - Opt(OptOps.LDS, 2, None)] - helper_lds_allclose(opts=opts, expected_bufs=[(0,8),(1,4),(2,8)]) - -if __name__ == "__main__": +if __name__ == '__main__': unittest.main() diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 79a2c09bdc..73227c89c4 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -5,7 +5,7 @@ from collections import defaultdict from typing import Optional, cast, Final, Callable, Sequence from tinygrad.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, track_rewrites, view_left, print_uops -from tinygrad.ops import PatternMatcher, UPat +from tinygrad.ops import PatternMatcher from tinygrad.spec import type_verify, shape_spec from tinygrad.device import Device from tinygrad.renderer import Renderer, TensorCore, ProgramSpec, Opt, OptOps @@ -78,7 +78,6 @@ class Kernel: self.tensor_core_opts: Optional[TensorCoreOptions] = None self.use_tensor_cores: int = 0 self.dont_use_locals: bool = False - self.lds: list[bool] = [False] * len(self.bufs) # group simplifies self.simplify_ones() @@ -95,8 +94,8 @@ class Kernel: ret.sts = self.sts[:len(ret.bufs)+len(ret.reduceops)*2] # NOTE: must redo the local buffers with TC in beam # parameters for optimizations - ret.applied_opts, ret.group_for_reduces, ret.upcasted, ret.local_dims, ret.dont_use_locals, ret.lds = \ - self.applied_opts[:], self.group_for_reduces, self.upcasted, self.local_dims, self.dont_use_locals, self.lds + ret.applied_opts, ret.group_for_reduces, ret.upcasted, ret.local_dims, ret.dont_use_locals = \ + self.applied_opts[:], self.group_for_reduces, self.upcasted, self.local_dims, self.dont_use_locals ret.tensor_core, ret.tensor_core_opts, ret.use_tensor_cores = self.tensor_core, self.tensor_core_opts, self.use_tensor_cores return ret @@ -354,7 +353,7 @@ class Kernel: return axis = self.real_axis(opt) - if opt.op != OptOps.LDS: check(axis < len(self.full_shape), "invalid axis") + check(axis < len(self.full_shape), "invalid axis") if opt.op is OptOps.SWAP: amt = cast(int, opt.arg) # arg is an axis in the SWAPs elif opt.arg is not None: @@ -425,9 +424,6 @@ class Kernel: self.sts[i] = st.pad(((0,0),) * axis + ((0,ru),) + ((0,0),) * (len(st.shape)-axis-1)) padded = True check(padded, "nothing was padded") - elif opt.op is OptOps.LDS: - check(0 <= axis < len(self.bufs), f"invalid buffer {axis}") - self.lds = self.lds[:axis] + [True] + self.lds[axis+1:] if append_opt: self.applied_opts.append(opt) if self.simplify_ones() and self.tensor_core_opts: @@ -659,11 +655,6 @@ class Kernel: return graph_rewrite(fixup_ast(self.ast), view_left) - def apply_lds(self, ast) -> UOp: - def transform(ctx:tuple[Kernel, set[UOp]], global_access:UOp): return None - - return graph_rewrite(ast, PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), name="global_access"), transform)]), ctx=(self, set())) - # **** this is the lowerer **** @track_rewrites() @@ -672,7 +663,6 @@ class Kernel: if getenv("VIZ"): graph_rewrite(self.ast, PatternMatcher([]), name="View Base AST") modified_ast = self.get_optimized_ast(name_override) - modified_ast = self.apply_lds(modified_ast) if ast_transform is not None: modified_ast = ast_transform(self, modified_ast) if DEBUG >= 3: diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index ebdf97a6be..ef0aba0d5d 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -8,7 +8,7 @@ from tinygrad.ops import Ops, UOp, sym_infer, sint, Variable, ssimplify, GroupOp from tinygrad.dtype import DType class OptOps(Enum): - TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto(); LDS = auto() # noqa: E702 + TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702 GROUP = auto(); GROUPTOP = auto(); NOLOCALS = auto(); PADTO = auto(); SWAP = auto() # noqa: E702 def __lt__(self, x:OptOps): return self.value < x.value From bc3bfcbad441cf1af16826972376829b1f38c877 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 2 Apr 2025 04:10:34 -0400 Subject: [PATCH 2/4] update install gpuocelot (#9693) `-DCMAKE_POLICY_VERSION_MINIMUM=3.5` --- .github/actions/setup-tinygrad/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/setup-tinygrad/action.yml b/.github/actions/setup-tinygrad/action.yml index 4096b60835..506b84422a 100644 --- a/.github/actions/setup-tinygrad/action.yml +++ b/.github/actions/setup-tinygrad/action.yml @@ -175,7 +175,7 @@ runs: git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99 mkdir build cd build - cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF + cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5 ninja - name: Install gpuocelot if: inputs.cuda == 'true' From 3b8d923692ea0caf4b6b408fc34b2a2414bb7128 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 2 Apr 2025 04:15:00 -0400 Subject: [PATCH 3/4] remove skip LLVM in test_div_int (#9686) --- test/test_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_ops.py b/test/test_ops.py index 69e05419dd..9a078b74d1 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -571,7 +571,7 @@ class TestOps(unittest.TestCase): np.testing.assert_equal(x.numpy(), 2**64 - 1) # 1 // 0 is device dependent, but it should not raise Tensor([1]).idiv(1).realize() - if not (CI and (Device.DEFAULT=="LLVM" or getenv("PTX"))): # TODO: crashed in CI + if not (CI and getenv("PTX")): # TODO: crashed in PTX CI # ... because if might be in a where branch that the output is well defined t = Tensor([-1, 0, 1, 2]) np.testing.assert_equal((t > 0).where(1//t, t).numpy(), [-1, 0, 1, 0]) From f72a87fd0e2120abc92c88fa56bf909e2bd88776 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Wed, 2 Apr 2025 16:38:01 +0800 Subject: [PATCH 4/4] add proper support for Ops.IGNORE to remove store masks (#9692) * add proper support for Ops.IGNORE to remove store masks * remove useless NHWC * revert that --- extra/onnx.py | 6 ------ tinygrad/codegen/expander.py | 32 ++++++++++++++++++++++++++++++-- tinygrad/renderer/__init__.py | 3 +-- 3 files changed, 31 insertions(+), 10 deletions(-) diff --git a/extra/onnx.py b/extra/onnx.py index a2985be7c0..9939153220 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -724,8 +724,6 @@ def get_onnx_ops(): ret = _clamp_cast((x / y_scale + 0.4999999 + y_zero_point).int(), out_dtype) else: ret = _clamp_cast(((x / y_scale).round() + y_zero_point), out_dtype) - # you need both NHWC=1 DONT_GROUP_REDUCES=1 for this to work - if getenv("NHWC") and len(ret.shape) == 4: return ret.permute(0,2,3,1).contiguous().permute(0,3,1,2) return ret.contiguous() def DynamicQuantizeLinear(x: Tensor): @@ -737,10 +735,6 @@ def get_onnx_ops(): return y, scale, zero_point def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:int=1, block_size:int=0): - WEIGHT_SHIFT = 4 - if getenv("NHWC") and len(x.shape) == 4 and x.shape[2:] == (1,1) and x.shape[1]%WEIGHT_SHIFT == 0: - # DSP swizzle memory - x = x.reshape(x.shape[0], x.shape[1]//WEIGHT_SHIFT, WEIGHT_SHIFT).permute(1,0,2).contiguous().permute(1,0,2).reshape(x.shape) x_scale, x_zero_point = _prepare_quantize(x, x_scale, x_zero_point, axis, block_size) return ((x.int() - x_zero_point) * x_scale).cast(x_scale.dtype) diff --git a/tinygrad/codegen/expander.py b/tinygrad/codegen/expander.py index 8585a7e92e..9f78bc0a3f 100644 --- a/tinygrad/codegen/expander.py +++ b/tinygrad/codegen/expander.py @@ -116,9 +116,37 @@ migrate_indexing = PatternMatcher([ (UPat(Ops.STORE, name="root"), create_gate), ]) +# **** IGNORE support **** + +pm_store_ignore = PatternMatcher([ + (UPat().index(UPat(), UPat(name="mask")).store(UPat()).named("store"), + lambda store,mask: store.replace(src=(store.src[0], UOp(Ops.IGNORE, src=(store.src[1], mask)))) if store.src[1].op is not Ops.IGNORE else None), +]) + +pm_move_ignore = PatternMatcher([ + # IGNORE on SELF is nothing + (UPat(Ops.IGNORE, src=(UPat(name="x"), UPat(name="x"))), lambda x: x.const_like(True)), + # IGNORE on a CONST is nothing + (UPat(Ops.IGNORE, src=(UPat((Ops.CONST, Ops.VCONST), name="c"), UPat())), lambda c: c), + # move the IGNOREs + (UPat(Ops.IGNORE, src=(UPat((*GroupOp.ALU, Ops.CAST, Ops.VECTORIZE), name="alu"), UPat.var("mask")), name="ig"), + lambda ig,alu,mask: alu.replace(src=tuple(UOp(Ops.IGNORE, x.dtype, (x, mask)) for x in alu.src))), +]) + +pm_delete_ignore = PatternMatcher([ + # IGNORE on SELF is nothing + (UPat(Ops.IGNORE, src=(UPat(name="x"), UPat())), lambda x: x), +]) + def expand_rewrite(sink:UOp) -> UOp: # initial symbolic + migrate indexing (remove this) sink = graph_rewrite(sink, sym+migrate_indexing) - # expand - return graph_rewrite(sink, sym+expander) + # store IGNORE + sink = graph_rewrite(sink, pm_store_ignore, name="store_ignore") + + # move IGNORE + sink = graph_rewrite(sink, pm_move_ignore, name="move_ignore") + + # expand + remove surviving ignores + return graph_rewrite(sink, pm_delete_ignore+sym+expander) diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index ef0aba0d5d..2d03f8a3d0 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -111,8 +111,7 @@ class ProgramSpec: # NOTE: you have to set local_size and global_size to the base [1,1,1] outside this if u.arg[0][0] == 'i': self.local_size = None special_size = self.local_size if u.arg[0][0] == 'l' else self.global_size - assert special_size is not None - special_size[int(u.arg[0][-1])] = u.arg[1] + if special_size is not None: special_size[int(u.arg[0][-1])] = u.arg[1] self.vars = sorted(self.vars, key=lambda v: v.arg) self.outs = sorted(dedup(self.outs)) self.ins = sorted(dedup(self.ins))