nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

Merge branch 'master' into dsp_search_merged

Browse Source
This commit is contained in:
George Hotz
2025-04-02 16:45:20 +08:00
committed by GitHub
parent 4496cc6e61 f72a87fd0e
commit 7e4ac744ac
6 changed files with 8 additions and 242 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/actions/setup-tinygrad/action.yml
+1 -1
View File
@@ -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'
extra/onnx.py
-1
View File
@@ -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]
test/test_linearizer.py
+1 -224
View File
@@ -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()
test/test_ops.py
+1 -1
View File
@@ -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])
tinygrad/codegen/kernel.py
+4 -14
View File
@@ -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:
tinygrad/renderer/__init__.py
+1 -1
View File
@@ -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