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' diff --git a/extra/onnx.py b/extra/onnx.py index d653f40d3e..b427b2796e 100644 --- a/extra/onnx.py +++ b/extra/onnx.py @@ -725,7 +725,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: in_chans = ret.shape[1] 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/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]) diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 78f8deb767..aaa67ae1a6 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: @@ -714,11 +710,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() @@ -727,7 +718,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 111fa563c9..ff1ffcca99 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