mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-09-18 05:23:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

The smallest promotion

Browse Source
This commit is contained in:
firestar5683
2026-08-10 22:53:46 -05:00
parent d938c9daa6
commit 6bd48c28e8
582 changed files with 64651 additions and 17100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad_repo/test/backend/test_custom_kernel.py
+260 -40
View File
@@ -1,18 +1,21 @@
import unittest
from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
from tinygrad.dtype import AddrSpace, dtypes
import numpy as np
from tinygrad.dtype import AddrSpace, dtypes, Invalid
from tinygrad.uop.ops import KernelInfo, AxisType, Ops
from tinygrad.renderer.ptx import PTXRenderer
from test.helpers import assert_kernel_count
# **** kernels ****
def custom_arange_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
return C[i].store(i.cast(C.dtype)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
def custom_eye_kernel(C:UOp) -> UOp:
i = UOp.range(C.shape[0], 0)
j = UOp.range(C.shape[1], 1)
return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
return C[i, j].store((i.eq(j)).cast(C.dtype)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
A,B = A.flatten(), B.flatten()
@@ -51,13 +54,13 @@ def flip_contract_kernel(dest:UOp, src:UOp):
i = UOp.range(dest.shape[0], 0)
j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
vec = src[i, j].contract(j)
store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
def slice_sum_kernel(dest:UOp, src:UOp):
G = UOp.range(src.shape[0], 0)
G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
slice_src = src[G, :]
reg = UOp.placeholder((1,), dest.dtype.base, 0, addrspace=AddrSpace.REG)
reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
reg = reg.after(G)[0].set(0)
R = UOp.range(src.shape[1], 1, AxisType.REDUCE)
reg = reg[0].set(reg.after(R)[0] + slice_src[R], end=R)
@@ -73,12 +76,12 @@ def simple_qkv_kernel(O:UOp, Q:UOp, K:UOp, V:UOp) -> UOp:
j = UOp.range(N, 2, axis_type=AxisType.REDUCE)
k_inner = UOp.range(d, 3, axis_type=AxisType.REDUCE)
qk_acc = UOp.placeholder((1,), Q.dtype.base, 0, addrspace=AddrSpace.REG)
qk_acc = UOp.placeholder((1,), Q.dtype, 0, addrspace=AddrSpace.REG)
qk_acc = qk_acc.after(i, j)[0].set(0.0)
qk_acc = qk_acc[0].set(qk_acc.after(k_inner)[0] + Q[i, k_inner] * K[j, k_inner], end=k_inner)
qk_score = qk_acc[0] / (d ** 0.5)
out_acc = UOp.placeholder((1,), Q.dtype.base, 1, addrspace=AddrSpace.REG)
out_acc = UOp.placeholder((1,), Q.dtype, 1, addrspace=AddrSpace.REG)
out_acc = out_acc.after(i, d_out)[0].set(0.0)
out_acc = out_acc[0].set(out_acc.after(j)[0] + qk_score * V[j, d_out], end=j)
@@ -117,13 +120,20 @@ class TestCustomKernel(unittest.TestCase):
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
def test_duplicate_call_arg(self):
x = Tensor.arange(4).clone().realize()
x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
# webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
# https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])
def test_simple_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
# ugly construction to get a sharded empty tensor
c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
out = c.flatten().tolist()
assert all(x == 2 for x in out), "all 2"
@@ -132,7 +142,7 @@ class TestCustomKernel(unittest.TestCase):
# PYTHON backend explicitly checks for OOB access for wrong multi shape regression
devs = ("PYTHON:0", "PYTHON:1")
a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
assert (c == 2).all().item()
@@ -180,6 +190,12 @@ class TestCustomKernel(unittest.TestCase):
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 15)
def test_sum_outside(self):
a = Tensor([1.0, 2, 3, 4, 5])+1
tst = Tensor.empty(1)
b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
self.assertEqual(b.item(), 20)
def test_sum_int(self):
a = Tensor([1, 2, 3, 4, 5])
tst = Tensor.empty(1, dtype=a.dtype)
@@ -199,18 +215,16 @@ class TestCustomKernel(unittest.TestCase):
c = Tensor.empty(N, N)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
def test_gemm_multi(self):
devs = ("CPU:0", "CPU:1")
N = 16
a = Tensor.randn(N, N).shard_(devs, axis=0)
b = Tensor.randn(N, N).to(devs)
c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
err = (tst - (a@b)).square().max()
self.assertLess(err.item(), 1e-6)
self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
def test_gemm_backward_custom(self): self.test_gemm_backward(True)
# NOTE: grad_fxn doesn't work with pyrender
@@ -233,14 +247,9 @@ class TestCustomKernel(unittest.TestCase):
real_grad_a, real_grad_b = a.grad, b.grad
Tensor.realize(ref, real_grad_a, real_grad_b)
err = (tst - ref).square().max()
self.assertLess(err.item(), 1e-6)
err = (grad_a - real_grad_a).square().max()
self.assertLess(err.item(), 1e-6)
err = (grad_b - real_grad_b).square().max()
self.assertLess(err.item(), 1e-6)
self.assertTrue(tst.allclose(ref, atol=1e-3).item())
self.assertTrue(grad_a.allclose(real_grad_a, atol=1e-3).item())
self.assertTrue(grad_b.allclose(real_grad_b, atol=1e-3).item())
def test_simple_qkv(self):
N, d = 8, 4
@@ -253,8 +262,7 @@ class TestCustomKernel(unittest.TestCase):
O_ref = ((Q @ K.T) / (d ** 0.5)) @ V
Tensor.realize(O_custom, O_ref)
err = (O_custom - O_ref).square().max()
self.assertLess(err.item(), 1e-6)
self.assertTrue(O_custom.allclose(O_ref, atol=1e-3).item())
def test_gemm_qkv(self):
B, N, K_DIM, H_KV, REP, D = 2, 7, 6, 2, 2, 6
@@ -275,7 +283,7 @@ class TestCustomKernel(unittest.TestCase):
GlobalCounters.reset()
out.realize()
self.assertEqual(GlobalCounters.kernel_count, 5)
assert_kernel_count(5)
def test_simple_reshape(self):
a = Tensor.ones(2,3,4).realize()
@@ -285,7 +293,7 @@ class TestCustomKernel(unittest.TestCase):
GlobalCounters.reset()
c.realize()
assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}"
self.assertEqual(GlobalCounters.kernel_count, 3)
assert_kernel_count(2)
def test_multi_after_schedule_order(self):
"""Test correct scheduling order when custom_kernel has multiple outputs.
@@ -328,6 +336,31 @@ class TestCustomKernel(unittest.TestCase):
if prg.op is not Ops.PROGRAM: continue
self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
def test_multi_invalids_custom_kernel_no_copy(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
c.realize()
assert_kernel_count(len(devs))
self.assertTrue((c == 2).all().item())
def test_partial_invalid_store_keeps_uncovered_reads(self):
x = Tensor([10., 20., 30., 40.])
after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])
def test_multi_after_invalid_store_dep_removed(self):
x = Tensor.empty(4).uop
self.assertEqual(Tensor(x.after(x.store(5), x.store(Invalid))).tolist(), [5]*4)
def test_expand_view_invalid_assign_keeps_uncovered_reads(self):
x = Tensor([[10., 11., 12., 13.], [20., 21., 22., 23.], [30., 31., 32., 33.], [40., 41., 42., 43.]]).realize()
v = x[:1, :].expand(4, 4)
v.assign(Tensor.invalids(4, 4, dtype=dtypes.float))
self.assertEqual(v.contiguous().tolist(), [[10., 11., 12., 13.]]*4)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
def test_invalids_into_custom_kernel_with_beam(self):
a = Tensor.full((4, 4), 3.).contiguous()
@@ -371,17 +404,15 @@ class TestCustomKernel(unittest.TestCase):
y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
if use_custom:
z = Tensor.empty_like(x)
z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
else: z = y.T.T+1
GlobalCounters.reset()
z.realize()
self.assertEqual(GlobalCounters.kernel_count, 2)
assert_kernel_count(2)
self.assertEqual(z.tolist(), x.add(2).tolist())
@unittest.expectedFailure
def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)
@unittest.expectedFailure
def test_sliced_buffer_function(self):
x = Tensor.arange(32).reshape(8, 4).clone().realize()
from tinygrad import function
@@ -392,24 +423,213 @@ class TestCustomKernel(unittest.TestCase):
GlobalCounters.reset()
y = run(x[0]).realize()
# it's copying the input and the output
self.assertEqual(GlobalCounters.kernel_count, 1)
# TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
self.assertEqual(y.tolist(), [1, 2, 3, 4])
@Context(DEV="CPU")
def test_simple_from_source(self):
a = Tensor([0., 1., 2.]).realize()
src = "void test_src(float* restrict a) { a[0] = 1.0; }"
a = Tensor.arange(4).clone().realize()
src = "void test_src(int* restrict a) { a[0] = 1; }"
# TODO: it currently requires a compiler for Ops.BINARY
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(A:UOp) -> UOp:
def custom_src_kernel(A:UOp, B:UOp) -> UOp:
sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
a = Tensor.custom_kernel(a.reshape(2, 2).clone(), a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [[1, 1], [2, 3]])
a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [1., 1., 2.])
@Context(DEV="CPU")
def test_simple_from_source_alt(self):
a = Tensor.arange(4).clone().realize()
src = "void copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }"
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(out:UOp, inp:UOp) -> UOp:
sink = UOp.sink(out, inp, arg=KernelInfo(name="copy"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
out = Tensor.custom_kernel(Tensor.empty_like(a), a+1, fxn=custom_src_kernel)[0]
GlobalCounters.reset()
out.realize()
assert_kernel_count(2)
self.assertEqual(out.tolist(), [1, 2, 3, 4])
@unittest.skip("this shouldn't be expected to work")
def test_inplace_transpose(self):
def custom_assign_row_max_kernel(A:UOp) -> UOp:
row = UOp.range(A.shape[0], 0)
col = UOp.range(A.shape[1], 1)
return A[row, col].store(A[row].max(axis=0)).end(col).end(row).sink(arg=KernelInfo(name=f"assign_row_max_{A.numel()}"))
a = Tensor.arange(4).clone().realize()
a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_assign_row_max_kernel)[0]
self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
self.assertEqual(a.shape, (2, 2))
class TestCustomKernelInput(unittest.TestCase):
def _test_mop(self, mop_fxn, max_kernels):
# default: input is BUFFER
x = mop_fxn(Tensor.arange(32).clone("CPU").realize())
y = Tensor.custom_kernel(Tensor.empty_like(x), x, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y.realize()
kernel_count = GlobalCounters.kernel_count
self.assertEqual(y.tolist(), x.add(1).tolist())
self.assertLessEqual(kernel_count, max_kernels)
# same test with @function, input is PARAM
from tinygrad import function
x0 = Tensor.arange(32).clone("CPU").realize()
@function(precompile=True)
def run(a:Tensor) -> Tensor:
xv = mop_fxn(a)
y = Tensor.invalids(*xv.shape, dtype=xv.dtype, device=a.device)
return Tensor.custom_kernel(y, xv, fxn=custom_add_one_kernel)[0]
GlobalCounters.reset()
y = run(x0).realize()
kernel_count = GlobalCounters.kernel_count
self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
self.assertLessEqual(kernel_count, max_kernels)
def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=2)
def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
def test_2d_shrink(self): self._test_mop(lambda x: x.reshape(4, 8)[:, 2:6], max_kernels=3)
def test_expand(self): self._test_mop(lambda x: x.reshape(16, 2)[:, :1].expand(16, 2), max_kernels=3)
class TestUnshardIndex(unittest.TestCase):
"""Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
shard. Two ownership patterns must work:
contiguous: idx = rng*shard_sz + local (thread rng owns [rng*shard_sz, ...))
strided: idx = rng + ir*shard_sz (thread rng owns {rng, rng+shard_sz, ...})
"""
def _run(self, kernel, shape=(8, 8)):
c = Tensor.empty(*shape)
out = Tensor.custom_kernel(c, fxn=kernel)[0]
try: return out.numpy()
except RuntimeError as e:
if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
self.skipTest("PTX does not support dynamic register indexing")
raise
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_contiguous_fragment_index(self):
# thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
# This is the pre-existing case that index_multi always handled.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
# 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
out = self._run(kernel, (64, 8))
assert out.shape == (64, 8)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_strided_fragment_index(self):
# thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
# strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
# The contiguous check (idx - rng*shard_sz) fails; the strided check
# (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
# 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
out = self._run(kernel, (64, 8))
assert out.shape == (64, 8)
def test_fragment_index_cannot_shard(self):
# thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
# the contiguous nor the strided ownership pattern, so index_multi must raise.
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
ir = UOp.range(8, 1, AxisType.LOOP)
j = UOp.range(8, 2, AxisType.LOOP)
frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
self._run(kernel, (64, 8))
def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
c = Tensor.empty(*out_shape)
out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
try: return out.numpy()
except RuntimeError as e:
if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
testcase.skipTest("PTX does not support dynamic register indexing")
raise
class TestUnshardAlu(unittest.TestCase):
"""Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.
An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
same sharding: peel the UNSHARD, keep the layout
scalar: broadcast to every shard
whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
"""
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_alu_scalar_broadcast(self):
# scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
# 8 values per thread, 8 threads -> 64-value full view
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(1.5)) * 2.0
return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
out = _run_fragment_kernel(self, kernel, (64,))
np.testing.assert_allclose(out, 3.0)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_alu_whole_value_subview(self):
# UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
def kernel(C:UOp, A:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(0.0)) + A
return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
a = Tensor(np.arange(64, dtype=np.float32))
out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
class TestUnshardStore(unittest.TestCase):
"""Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).
Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
"""
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_store_unshard_value(self):
# single-axis: 8 threads each own 8 values of the 64-value output tile
def kernel(C:UOp) -> UOp:
ty = UOp.range(8, 0, AxisType.LOCAL)
frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
v = frag.after(frag.store(0.0)) + 2.5
return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
out = _run_fragment_kernel(self, kernel, (64,))
np.testing.assert_allclose(out, 2.5)
@unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
def test_store_unshard_value_2axis(self):
# two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
# (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
def kernel(C:UOp, A:UOp) -> UOp:
ty = UOp.range(4, 0, AxisType.LOCAL)
tx = UOp.range(2, 1, AxisType.LOCAL)
frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
v = frag.after(frag.store(0.0)) + A
return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
class TestUOpReduce(unittest.TestCase):
def test_uop_sum(self):