mawei/onepilot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update251203 (#233)

Browse Source
This commit is contained in:
carrot
2025-12-03 10:28:27 +09:00
committed by GitHub
parent d6899edd97
commit c5ebcbcb97
347 changed files with 8678 additions and 13489 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad_repo/test/test_schedule.py
+129 -217
View File
@@ -12,9 +12,9 @@ from tinygrad import nn, dtypes, Device, Tensor
from tinygrad.device import is_dtype_supported
from tinygrad.dtype import DType, ImageDType
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
from tinygrad.uop.symbolic import symbolic_simple
from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp, RANGEIFY
from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
@@ -33,6 +33,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
# test lowering all the ScheduleItems to ExecItems
kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
if kernel_cnt != allowed:
if RANGEIFY: return sched # allow different kernel count, TODO: fix the asserts
print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
if DEBUG >= 3:
for i,s in enumerate(sched):
@@ -41,6 +42,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
raise KernelCountException(f"{kernel_cnt} != {allowed}")
return sched
def expect_rangeify_fails(fxn): return (unittest.expectedFailure if RANGEIFY else (lambda f:f))(fxn)
def _realize_weights(m):
for p in nn.state.get_parameters(m): p.realize()
@@ -111,6 +114,7 @@ class TestSchedule(unittest.TestCase):
self.assertListEqual(a.tolist(), [[15]])
@unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch")
@expect_rangeify_fails
def test_error_on_device_mismatch(self):
a = Tensor.empty(10)
b = Tensor.empty(10, device="CPU")
@@ -118,11 +122,12 @@ class TestSchedule(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1)
@unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch")
@expect_rangeify_fails
def test_error_on_device_mismatch_alt(self):
a = Tensor.empty(10)
b = Tensor.empty((1,), device="CPU").expand(10).contiguous()
c = a+b
with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1)
with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2 if RANGEIFY else 1)
@unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1")
@unittest.skip("CAST_AFTER_EXPAND is not supported")
@@ -140,6 +145,7 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])
@unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@unittest.skipIf(RANGEIFY, "rangeify doesn't implement input buffer limiting")
def test_add_chain_buffers(self):
N = 31
with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -198,9 +204,10 @@ class TestSchedule(unittest.TestCase):
def test_simplify_padded_const(self):
a = Tensor.empty(1022).cummax(axis=0)
sched = check_schedule(a, 5)
ast = sched[0].ast
self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)
check_schedule(a, 5)
# TODO: what is this testing?
#ast = sched[0].ast
#self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)
def test_basic_binop_fusion(self):
a = Tensor.empty(10)
@@ -278,7 +285,7 @@ class TestSchedule(unittest.TestCase):
a = Tensor.empty(10,10,10)
b = Tensor.empty(10,10,1)
c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
with self.assertRaises(KernelCountException): check_schedule(c, 1)
check_schedule(c, 2)
def test_allow_push_permutes(self):
a = Tensor.randn(10,10,10).realize()
@@ -316,7 +323,7 @@ class TestSchedule(unittest.TestCase):
b = Tensor.empty(10)
c = a+b
d = a.reshape(10,1)+b.reshape(10,1)
with self.assertRaises(KernelCountException): check_schedule(d, 0, [c])
check_schedule(d, 1, [c])
# failing in new lazy
def test_cache_binaryop_transpose(self):
@@ -324,7 +331,7 @@ class TestSchedule(unittest.TestCase):
b = Tensor.empty(10,10)
c = (a.T*b.T).T #.contiguous()
d = a*b
with self.assertRaises(KernelCountException): check_schedule(d, 0, [c])
check_schedule(d, 1, [c])
def test_cache_two_reduceops(self):
a = Tensor.empty(10)
@@ -339,7 +346,7 @@ class TestSchedule(unittest.TestCase):
r1 = (x - r0).sum(axis=0).div(2)
out = r0 + r1
schedule = check_schedule(out, 2)
reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS]
reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
assert len(reduceops) == 2
def test_cache_reduce_multiple_children(self):
@@ -349,9 +356,9 @@ class TestSchedule(unittest.TestCase):
r1 = (x - r0).sum(axis=0).div(2)
out0 = r0 + y
out1 = r1 + y
schedule = check_schedule([out0, out1], 4)
reduceops = [x for si in schedule for x in si.ast.toposort() if x.op is Ops.REDUCE_AXIS]
assert len(reduceops) == 2
schedule = check_schedule([out0, out1], 2 if RANGEIFY else 4)
reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
assert len(reduceops) == (3 if RANGEIFY else 2)
def test_div_collapse_buffer(self):
a = Tensor.full((4,), 4.0).contiguous().realize()
@@ -394,6 +401,7 @@ class TestSchedule(unittest.TestCase):
# a and b share the same underlying device memory
self.assertIs(a.uop.realized, b.uop.realized)
@expect_rangeify_fails
def test_clone_doesnt_dedup(self):
src = Tensor.ones(4).contiguous().realize()
a = src.clone()
@@ -417,6 +425,11 @@ class TestSchedule(unittest.TestCase):
b = Tensor.full((4, 4), 1.).contiguous().realize()
check_schedule([a+b, a+b], 1)
def test_const_realize(self):
t = Tensor.ones(2)
check_schedule(t[0], 0)
check_schedule(t[1], 0)
def test_fold_double_unary(self):
y = Tensor.empty(2)
out = y.sum(keepdim=True).sqrt().neg()
@@ -558,7 +571,7 @@ class TestSchedule(unittest.TestCase):
c = a+b
d = a.reshape(10,1)+b.reshape(10,1)
out = c.sum() + d.sum()
with self.assertRaises(KernelCountException): check_schedule(out, 1)
check_schedule(out, 2)
def test_children_dont_push(self):
a = Tensor.empty(10, 10, 1)
@@ -569,6 +582,7 @@ class TestSchedule(unittest.TestCase):
check_schedule(f, 2)
# failing in new lazy
@unittest.skip("always fusing elementwise")
def test_dont_fuse_binops_with_children(self):
a = Tensor.empty(10)
b = Tensor.empty(10)
@@ -576,8 +590,8 @@ class TestSchedule(unittest.TestCase):
keep_me = a+b
e = keep_me.sum() # noqa: F841 give keep_me a child (NOTE: BinaryOps won't be a child since it will instant fuse)
d = keep_me+c
with self.assertRaises(KernelCountException): check_schedule(d, 2)
with self.assertRaises(KernelCountException): check_schedule(keep_me, 0, [d])
check_schedule(d, 2)
check_schedule(keep_me, 0, [d])
#@unittest.skip("failing in old lazy")
def test_permute_breaks_fusion(self):
@@ -627,7 +641,8 @@ class TestSchedule(unittest.TestCase):
x = x.image_conv2d(w3, b3)
# NOOP, 3 convs, contiguous
with self.assertRaises(KernelCountException): check_schedule(x, 5)
#check_schedule(x, 5)
check_schedule(x, 8)
def test_image_conv_fusion_minimal(self):
b1 = Tensor.empty(16)
@@ -682,6 +697,7 @@ class TestSchedule(unittest.TestCase):
c = (a.sum(2).contiguous() + b).contiguous()
check_schedule(c, 2)
@expect_rangeify_fails
def test_kernelize(self):
a = Tensor.empty(10)
b = Tensor.empty(10)
@@ -689,12 +705,14 @@ class TestSchedule(unittest.TestCase):
d = c+2
check_schedule(d, 2)
@expect_rangeify_fails
def test_kernelize_view(self):
a = Tensor.empty(4,1)
b = a*2
c = b.kernelize()+Tensor.empty(4,4)
check_schedule(c, 2)
@expect_rangeify_fails
def test_kernelize_diamond(self):
a = Tensor([0]).realize()
prev_a = (a+1).contiguous()
@@ -703,6 +721,7 @@ class TestSchedule(unittest.TestCase):
assert prev_a.uop in a.uop.src, "contiguous usage must run before assign"
self.assertEqual((prev_a+a*3).item(), 1+2*3)
@expect_rangeify_fails
def test_multioutput_ast(self):
a = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
b = Tensor.zeros(1, dtype=dtypes.int).contiguous().realize().uop
@@ -714,6 +733,7 @@ class TestSchedule(unittest.TestCase):
self.assertEqual(b.buffer.numpy(), [12])
# unlike schedule, kernelize can be called multiple times on a Tensor
@expect_rangeify_fails
def test_double_kerenlize(self):
a = Tensor.empty(10)
b = Tensor.empty(10)
@@ -722,6 +742,7 @@ class TestSchedule(unittest.TestCase):
e = c.kernelize()+d.kernelize()
check_schedule(e, 3)
@expect_rangeify_fails
def test_kernelize_bw(self):
a = Tensor.full((3,), 2.0, requires_grad=True).contiguous()
b = Tensor.full((3,), 3.0, requires_grad=True).contiguous()
@@ -732,6 +753,7 @@ class TestSchedule(unittest.TestCase):
self.assertEqual(z.item(), 18.0)
self.assertEqual(z.grad.item(), 1.0)
@expect_rangeify_fails
def test_kernelize_bw_view(self):
a = Tensor.full((3,1), 2.0, requires_grad=True).contiguous()
b = Tensor.full((3,1), 3.0, requires_grad=True).contiguous()
@@ -784,6 +806,13 @@ class TestSchedule(unittest.TestCase):
out = x + 1
check_schedule(out, 0, filter_sink=False)
def test_zero_size_assign(self):
f = Tensor.full((2,), 0.).contiguous().realize()
a = f.shrink_to((0,))
a.assign(Tensor.ones_like(a))
check_schedule(a, 0)
self.assertEqual(a.tolist(), [])
def test_reduce_permute_nofuse(self):
x = Tensor.empty(32, 32, 32)
y = Tensor.empty(32, 32)
@@ -888,26 +917,24 @@ class TestSchedule(unittest.TestCase):
out = x.contiguous() + y.contiguous()
check_schedule(out, 2, filter_sink=False)
@unittest.expectedFailure
def test_reduce_same_size(self):
Tensor.manual_seed(0)
a = Tensor.randn(4, 4).realize()
out0 = a.sum() + 2
out1 = a.sum() + 4
out2 = out0 * out1
run_schedule(check_schedule([out0, out1, out2], 1))
run_schedule(check_schedule([out0, out1, out2], 1 if RANGEIFY else 4))
np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
@unittest.expectedFailure
def test_reduce_multiple_paths(self):
Tensor.manual_seed(0)
a = Tensor.randn(4, 4).realize()
out0 = a.sum().exp2()
# out1 has two paths to a.sum()
out1 = a.sum() + out0
run_schedule(check_schedule([out0, out1], 1))
run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)
@@ -983,7 +1010,6 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(e.numpy(), e_np:=b.numpy() + out0_np, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(out1.numpy(), r_np + e_np[0][0][0], atol=1e-4, rtol=1e-4)
# changed by multireduce
def test_reduce_expand_child(self):
Tensor.manual_seed(0)
a = Tensor.randn((32, 32, 32)).realize()
@@ -995,13 +1021,12 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy(), atol=1e-4, rtol=1e-4)
@unittest.expectedFailure
def test_reduce_shrink_child(self):
a = Tensor.empty(100, 100)
b = Tensor.empty(10,)
c = a.sum() + b[0]
d = a.sum() + 2
check_schedule([c, d], 1)
check_schedule([c, d], 1 if RANGEIFY else 3)
def test_reduce_multiple_paths_midshrink(self):
a = Tensor.empty(4, 4)
@@ -1024,20 +1049,6 @@ class TestSchedule(unittest.TestCase):
run_schedule(check_schedule(out, 2))
np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
def test_argmin_multireduce_fusion(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmin(-1)
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
def test_argmax_multireduce_fusion(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmax(-1)
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
def test_scaled_dot_product_attention_multireduce_fusion(self):
Tensor.manual_seed(0)
q = Tensor.randn(32,8,16,8).realize()
@@ -1050,6 +1061,14 @@ class TestSchedule(unittest.TestCase):
compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
with Context(FUSE_ATTENTION=1):
out = Tensor.scaled_dot_product_attention(q,k,v)
run_schedule(check_schedule(out, 1))
if getenv("CHECK", 1):
import torch
compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
def test_ugly_reduceop_pairing(self):
Tensor.manual_seed(0)
a = Tensor.randn(4, 32).realize()
@@ -1171,13 +1190,14 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@expect_rangeify_fails
def test_softmax_upcast(self):
# input half, softmax in float
Tensor.manual_seed(0)
x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
out = x.softmax(dtype=dtypes.float)
sched = out.schedule()
self.assertEqual(len(sched), 3)
self.assertEqual(len(sched), 2 if RANGEIFY else 3)
self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
# input float, softmax in float
@@ -1194,7 +1214,6 @@ class TestSchedule(unittest.TestCase):
x.softmax().sum().backward()
run_schedule(check_schedule(x.grad, 4))
# changed by: multireduce spec
def test_layernorm_onelayer_fusion(self):
Tensor.manual_seed(0)
layer = nn.LayerNorm([10, 10])
@@ -1308,6 +1327,7 @@ class TestSchedule(unittest.TestCase):
with Context(FUSE_CONV_BW=1): check_schedule(opt.schedule_step(), 14)
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@expect_rangeify_fails
def test_prefer_half_buffer(self):
x = Tensor.ones(4).contiguous().realize()
# y = Tensor.ones(4).contiguous().realize()
@@ -1425,7 +1445,6 @@ class TestSchedule(unittest.TestCase):
run_schedule(schedule)
np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)
# changed by: multireduce spec
# pattern in test_transformer
def test_partial_fuse1(self):
Tensor.manual_seed(0)
@@ -1438,7 +1457,6 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)
# changed by: multireduce spec
# pattern in conv
def test_partial_fuse2(self):
Tensor.manual_seed(0)
@@ -1451,9 +1469,7 @@ class TestSchedule(unittest.TestCase):
np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(d.numpy(), b.numpy().sum()-(a.numpy().sum()+2), atol=1e-4, rtol=1e-4)
# changed by: multireduce spec
# pattern in adam
@unittest.expectedFailure
def test_partial_fuse3(self):
Tensor.manual_seed(0)
a = Tensor.randn(16, 16).realize()
@@ -1463,14 +1479,12 @@ class TestSchedule(unittest.TestCase):
e = c * d
f = b.sum() - e
# run_schedule(check_schedule([c, d, e, f], 1))
run_schedule(check_schedule([c, d, e, f], 2))
run_schedule(check_schedule([c, d, e, f], 2 if RANGEIFY else 5))
np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(f.numpy(), b.numpy().sum() - e_np, atol=1e-4, rtol=1e-4)
# changed by: multireduce spec
@unittest.expectedFailure
def test_partial_fuse4(self):
Tensor.manual_seed(0)
a = Tensor.randn(16, 16).realize()
@@ -1480,7 +1494,7 @@ class TestSchedule(unittest.TestCase):
e = c * d
f = (b - d).sum() - e
# run_schedule(check_schedule([c, d, e, f], 1))
run_schedule(check_schedule([c, d, e, f], 3))
run_schedule(check_schedule([c, d, e, f], 5))
np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1615,11 +1629,11 @@ class TestSchedule(unittest.TestCase):
out = x.argmax(1)
run_schedule(check_schedule(out, 2))
def test_conv2d(self): _test_conv2d(7)
def test_conv2d_fused(self): _test_conv2d(5, FUSE_CONV_BW=1)
def test_conv2d(self): _test_conv2d(4 if RANGEIFY else 7)
def test_conv2d_fused(self): _test_conv2d(4 if RANGEIFY else 5, FUSE_CONV_BW=1)
@unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong")
def test_conv2d_half(self): _test_conv2d(7, dtype=dtypes.half)
def test_conv2d_half(self): _test_conv2d(4 if RANGEIFY else 7, dtype=dtypes.half)
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail")
@unittest.expectedFailure
@@ -1646,7 +1660,8 @@ class TestSchedule(unittest.TestCase):
constv = Tensor.empty(2, 2).uop.const_like(10).contiguous()
check_schedule(constv, 1)
@unittest.skipIf(Device.DEFAULT != "GPU", "image only supported on GPU")
@unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
@expect_rangeify_fails
def test_image_matmul(self):
with Context(IMAGE=2):
x = Tensor.randn((9, 9)).realize()
@@ -1682,6 +1697,7 @@ class TestSchedule(unittest.TestCase):
def test_late_fusion_post_expand(self):
self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)
@expect_rangeify_fails
def test_cast_padded_view(self):
a = Tensor.arange(4).reshape(1, 4)
casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
@@ -1711,6 +1727,7 @@ class TestSchedule(unittest.TestCase):
self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
@given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
@expect_rangeify_fails
def test_cast_padded_const(self, dt1, dt2):
assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
@@ -1720,53 +1737,41 @@ class TestSchedule(unittest.TestCase):
run_schedule(check_schedule(realized_const_view, 1))
np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]])
class TestIndexing(unittest.TestCase):
def check_schedule(self, xt:Tensor|list[Tensor], cnt:int):
with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):
lst = [xt] if isinstance(xt, Tensor) else xt
s = Tensor.schedule(*lst)
lowered = [x[1] for x in lower_schedule(s.copy())]
kernels = [ei for ei in list(lowered) if isinstance(ei.prg, CompiledRunner)]
if FUSE_ARANGE and len(kernels) != cnt:
raise KernelCountException(f"{len(kernels)} != {cnt}")
for ei in lowered: ei.run(do_update_stats=True)
return s
def test_simple_indexing(self):
X = Tensor.randn(10, 10).realize()
idxs = Tensor([0, 2]).realize()
xt = X[idxs]
self.check_schedule(xt, 2)
run_schedule(check_schedule(xt, 2))
np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])
def test_simple_indexing_alt(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[[1, 2], [-1, 2]]
self.check_schedule(xt, 1)
run_schedule(check_schedule(xt, 1))
np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]])
def test_advanced_indexing(self):
X = Tensor.arange(10)+1
xt = X[[0, -1]]
self.check_schedule(xt, 1)
run_schedule(check_schedule(xt, 1))
np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]])
def test_advanced_indexing_alt(self):
X = Tensor.arange(6).reshape(3, 2)+1
xt = X[[Tensor([2]), Tensor([1])]]
self.check_schedule(xt, 3)
run_schedule(check_schedule(xt, 3))
np.testing.assert_equal(xt.numpy(), 6)
def test_advanced_simple_indexing_combined(self):
X = Tensor.arange(16).reshape(4, 4)
xt = X[1:2, [-1, 2]]
self.check_schedule(xt, 1)
run_schedule(check_schedule(xt, 1))
def test_push_through_reshape(self):
Tensor.manual_seed(0)
x = Tensor.randn(10, 20).realize()
out = x.argmax(1)
self.check_schedule(out, 2)
run_schedule(check_schedule(out, 2))
np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1))
def test_arange_push_through_expand(self):
@@ -1774,35 +1779,35 @@ class TestIndexing(unittest.TestCase):
a = Tensor.arange(4,)
b = Tensor.randn(4, 4).realize()
out = (a+b).sum()
self.check_schedule(out, 1)
run_schedule(check_schedule(out, 1))
np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5)
def test_argmin(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmin(-1)
self.check_schedule(out, 2)
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
def test_argmax(self):
Tensor.manual_seed(0)
x = Tensor.randn(4, 32).realize()
out = x.argmax(-1)
self.check_schedule(out, 2)
run_schedule(check_schedule(out, 2))
np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
def test_arange_transposed(self):
Tensor.manual_seed(0)
x = Tensor.randint(4, 1).realize()
a = ((Tensor.arange(4,)*x).T).sum()
self.check_schedule(a, 1)
run_schedule(check_schedule(a, 1))
np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())
def test_div_padded_arange(self):
x = Tensor.full((2,2), 16)
y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
out = y.sum(axis=1)
with Context(FUSE_ARANGE=1): run_schedule(check_schedule(out, 2))
run_schedule(check_schedule(out, 2))
self.assertListEqual(out.tolist(), [0, 12, 4, 0])
def test_arange_transposed_descendants(self):
@@ -1811,7 +1816,7 @@ class TestIndexing(unittest.TestCase):
a = (Tensor.arange(4,)*x).T
b = Tensor.randint(4, 4).realize()
out = (a+b).sum()
self.check_schedule(out, 1)
run_schedule(check_schedule(out, 1))
np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum())
def test_arange_index(self):
@@ -1819,7 +1824,7 @@ class TestIndexing(unittest.TestCase):
x = Tensor.randn(5, 2).realize()
a = Tensor.arange(10)
out = (x + a[2]).sum()
self.check_schedule(out, 1)
run_schedule(check_schedule(out, 1))
np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
def test_arange_index_shrink(self):
@@ -1828,14 +1833,14 @@ class TestIndexing(unittest.TestCase):
x = Tensor.randn(11).realize()
a = Tensor.arange(22)
out = (x + a[:11]).sum()
self.check_schedule(out, 1)
check_schedule(out, 1)
def test_arange_index_contiguous(self):
Tensor.manual_seed(0)
x = Tensor.randn(5, 2).realize()
a = Tensor.arange(10).contiguous()
out = (x + a[2]).sum()
self.check_schedule(out, 3)
run_schedule(check_schedule(out, 3))
np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
def test_arange_index_child(self):
@@ -1843,62 +1848,24 @@ class TestIndexing(unittest.TestCase):
x = Tensor.randn(5, 2).realize()
a = Tensor.arange(10)+1
out = (x + a[2]).sum()
self.check_schedule(out, 1)
run_schedule(check_schedule(out, 1))
np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
def test_arange_index_contiguous_child(self):
def test_user_contiguous(self):
Tensor.manual_seed(0)
x = Tensor.randn(5, 2).realize()
a = (Tensor.arange(10)+1).contiguous()
out = (x + a[2]).sum()
self.check_schedule(out, 3)
run_schedule(check_schedule(out, 3))
np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
def test_arange_childless_base(self):
a = Tensor.arange(4)
self.check_schedule(a, 1)
np.testing.assert_equal(a.numpy(), np.arange(4))
def test_arange_childless_view(self):
a = Tensor.arange(4).reshape(2, 2)
a[0] = 4
np.testing.assert_equal(a.numpy(), [[4, 4], [2, 3]])
def test_arange_group_childless_base(self):
Tensor.manual_seed(0)
x = Tensor.randint(4).realize()
a = Tensor.arange(4)+x
self.check_schedule(a, 1)
np.testing.assert_equal(a.numpy(), np.arange(4)+x.numpy())
def test_arange_group_childless_view(self):
Tensor.manual_seed(0)
x = Tensor.ones(4).contiguous().realize()
a = Tensor.arange(4)+x
a[0] = 6
np.testing.assert_equal(a.numpy(), [6., 2., 3., 4.])
@unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
def test_arange_view_op(self):
a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous()
sched = self.check_schedule(a, 1)
sched = run_schedule(check_schedule(a, 1))
self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW)
np.testing.assert_equal(a.numpy(), [[4, 5]])
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
def test_arange_shrink_copy(self):
a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).to("CPU")
sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
self.assertIs(sched[-1].ast.op, Ops.COPY)
np.testing.assert_equal(a.numpy(), [[4, 5]])
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
def test_arange_expand_copy(self):
a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU")
sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
self.assertIs(sched[2].ast.op, Ops.COPY)
np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]])
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
def test_precompute_freqs_cis(self):
from extra.models.llama import precompute_freqs_cis
@@ -1914,23 +1881,33 @@ class TestIndexing(unittest.TestCase):
def test_fuse_assign_contiguous(self):
x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
a = Tensor.arange(8).reshape(4, 2)
self.check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2)
run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
def test_assign_non_contiguous(self):
x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
y = Tensor.randint(4, 2)
a = Tensor.arange(8).reshape(4, 2)+y
x.shrink((None, (0, 2))).assign(a).realize()
xref = np.zeros((4, 4), dtype=int)
xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
def test_assign_non_contiguous(self, alt=False):
x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
xref = x.numpy()
if alt:
y = Tensor.randint(2, 4).contiguous().realize()
a = Tensor.arange(8).reshape(2, 4)+y
tst = x.shrink(((0, 2), None)).assign(a).realize()
xref[:2, :] = np.arange(8).reshape(2, 4)+y.numpy()
else:
y = Tensor.randint(4, 2).contiguous().realize()
a = Tensor.arange(8).reshape(4, 2)+y
tst = x.shrink((None, (0, 2))).assign(a).realize()
xref[:, :2] = np.arange(8).reshape(4, 2)+y.numpy()
np.testing.assert_equal(x.numpy(), xref)
if RANGEIFY > 0:
# NOTE: this is a bug on non rangeify
np.testing.assert_equal(tst.numpy(), a.numpy())
def test_sparse_categorical_crossentropy_simple(self):
X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
Y = Tensor([1, 2]).realize()
loss = X.sparse_categorical_crossentropy(Y)
self.check_schedule(loss, 4)
run_schedule(check_schedule(loss, 4))
np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Validation error on WebGPU")
@@ -1942,28 +1919,21 @@ class TestIndexing(unittest.TestCase):
yt = Tensor.randn(BS, 10).realize()
with Context(SPLIT_REDUCEOP=0):
loss = yt.sparse_categorical_crossentropy(Y_train[samples])
self.check_schedule(loss, 6)
run_schedule(check_schedule(loss, 6))
loss_fused = loss.numpy()
loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@unittest.expectedFailure
def test_arange_fuse_grouped_children(self):
X = Tensor.randn(4, 4).realize()
r = (X+Tensor.arange(16).reshape(4, 4)).sum()
out0 = r+2
out1 = r+3
self.check_schedule([out0, out1], 1)
run_schedule(check_schedule([out0, out1], 1 if RANGEIFY else 3))
r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
def test_dont_fold_arange_contiguous_view(self):
X = Tensor.randn(4, 4).realize()
r = (X+Tensor.arange(16).reshape(4, 4).contiguous()).sum(1, keepdim=True)
self.check_schedule([r], 2)
np.testing.assert_allclose(r.numpy(), (X.numpy()+np.arange(16).reshape(4, 4)).sum(1, keepdims=True), atol=1e-5, rtol=1e-6)
@unittest.skip("multi output isn't supported")
def test_multiview_arange_children(self):
X = Tensor.randn(2,3,4,4).numpy()
@@ -2093,6 +2063,7 @@ class TestView(unittest.TestCase):
run_schedule(sched)
np.testing.assert_equal(b.numpy(), 0)
@expect_rangeify_fails
def test_mask_dim_1(self):
# mask out dim = 1 works too
a = Tensor.rand(10, 10).realize()
@@ -2119,6 +2090,7 @@ class TestView(unittest.TestCase):
# a*VIEW(x), where VIEW(x) = 0
# x collapses along with its children
@unittest.skipIf(RANGEIFY, "this only fails if you run all of TestSchedule, some global tensor map bug?")
def test_parent_view_collapses(self):
a = Tensor([1, 2])
b = Tensor.arange(3).contiguous()
@@ -2136,6 +2108,7 @@ class TestView(unittest.TestCase):
# a*VIEW(x), where VIEW(x) = 0
# x+2
# as long as one child realizes, x does not collapse
@expect_rangeify_fails
def test_parent_multiple_children_no_collapse(self):
a = Tensor([1, 2])
b = Tensor.arange(3).contiguous()
@@ -2200,84 +2173,6 @@ class TestSimplifier(unittest.TestCase):
assert UPat(Ops.CONST, arg=False).match(sink, {}), f"expected {sink} to collapse to a const False"
assert sink.shape == a.shape
tensor_const_pm = PatternMatcher([
(UPat(Ops.CONST, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)),)), lambda: True),
(UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR, src=(UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),)))), UPat(Ops.CONST))), lambda: True),
])
class TestConst(unittest.TestCase):
# ** part 1: basic functionality of a tensor directly created from CONST
def test_tensor_const(self):
a = Tensor(1)
print(a.uop)
self.assertTrue(tensor_const_pm.rewrite(a.uop))
def test_tensor_variable(self):
vv = UOp.variable("a", 0, 10).bind(1)
a = Tensor(vv)
print(a.uop)
self.assertTrue(tensor_const_pm.rewrite(a.uop))
def test_const_schedule(self):
a = Tensor.ones((4, 4))
sched = a.schedule()
self.assertEqual(len(sched), 0)
def test_const_contiguous_schedule(self):
# this ends up in the big graph
a = Tensor.ones((4,)).contiguous()
sched = a.schedule()
self.assertEqual(len(sched), 1)
# ** part 2: scheduler behavior when const folding happens later
def test_const_folding_no_realize(self):
a = Tensor([1, 2, 3, 4])*0
sched = a.schedule()
self.assertEqual(len(sched), 0)
def test_src_const_folding(self):
with Context(TRACK_MATCH_STATS=0):
a = Tensor.full((4,), 1).contiguous().realize()
b = Tensor.full((4,), 2).contiguous().realize()
mul0 = a*0
add = b+mul0
sched = add.schedule()
self.assertEqual(len(sched), 0)
# b+0 and b share the same underlying device memory
self.assertIs(add.uop.buffer, b.uop.buffer)
self.assertListEqual(add.tolist(), [2, 2, 2, 2])
def test_src_masked_const_folding(self):
with Context(TRACK_MATCH_STATS=0):
a = Tensor.full((4,), 1).contiguous().realize()
b = Tensor.full((6,), 2).contiguous().realize()
mul0 = a*0
add = b+mul0.pad((1, 1), value=2)
sched = add.schedule()
self.assertEqual(len(sched), 1)
run_schedule(sched)
# add gets assigned to a new buffer
self.assertIsNot(add.uop.base.realized, b.uop.base.realized)
self.assertListEqual(add.tolist(), [4, 2, 2, 2, 2, 4])
# ** part 3: Tensor variable bindings
#@unittest.expectedFailure # TODO: should schedule assert if you try to realize a Variable?
def test_var_schedule(self):
vv = UOp.variable("a", 0, 10).bind(1)
a = Tensor(vv)
sched = a.schedule()
self.assertEqual(len(sched), 0)
def test_add_tvar(self):
vv = UOp.variable("a", 0, 10).bind(1)
a = Tensor(vv)+2
sched, var_vals = a.schedule_with_vars()
self.assertEqual(len(sched), 1)
run_schedule(sched, var_vals)
self.assertEqual(a.tolist(), 3)
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
class TestCopyFolding(unittest.TestCase):
def test_const_copy_is_free(self):
@@ -2291,6 +2186,7 @@ class TestCopyFolding(unittest.TestCase):
b = (a*zeros).to("CPU")
run_schedule(check_schedule(b, 0, filter_sink=False))
self.assertListEqual(b.tolist(), [0, 0, 0])
self.assertEqual(b.device, "CPU")
def test_alu_after_copy(self):
a = Tensor.ones((4,)).to("CPU")
@@ -2299,6 +2195,12 @@ class TestCopyFolding(unittest.TestCase):
add.kernelize()
assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
def test_alu_before_copy(self):
buf = Tensor.ones(1).contiguous().realize()
a = buf+1
b = a.to("CPU")
self.assertListEqual(b.tolist(), [2.])
def test_copy_to_same_device(self):
a = Tensor.empty(4).uop
b = a.copy_to_device(a.device)
@@ -2345,6 +2247,7 @@ class TestCopyFolding(unittest.TestCase):
b.realize()
self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])
@expect_rangeify_fails
def test_permute_on_disk(self):
with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_buffer())
a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
@@ -2491,6 +2394,7 @@ class TestUOpBecome(unittest.TestCase):
self.assertEqual(add.uop.shape, (8, 2))
assert add.uop is not add.uop.base
@expect_rangeify_fails
def test_new_flat_buffer(self):
a = Tensor.empty(4,)
b = Tensor.empty(4,)
@@ -2516,6 +2420,7 @@ class TestUOpBecome(unittest.TestCase):
z = (img*x) / y
check_schedule(z, 1)
@expect_rangeify_fails
def test_become_existing_buffer(self):
a = Tensor.empty(4, 4)
b = a*1
@@ -2543,6 +2448,7 @@ class TestUOpBecome(unittest.TestCase):
check_schedule(b, 0)
assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER)
@expect_rangeify_fails
def test_become_const_in_view(self):
# if we shrink the base down to a size 0, only the VIEW becomes CONST, base is unchanged.
add = Tensor.empty(2, 2)+Tensor.empty(2, 2)
@@ -2560,6 +2466,7 @@ class TestUOpBecome(unittest.TestCase):
assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {})
# tensors can become another realized tensor source
@expect_rangeify_fails
def test_become_existing_buf_simple(self):
a = Tensor.empty(4, 4)
b = a+0
@@ -2568,12 +2475,14 @@ class TestUOpBecome(unittest.TestCase):
self.assertIs(a.uop, b.uop)
# they can also chain other movement ops on top of the tensor source
@expect_rangeify_fails
def test_become_existing_buf_view(self):
a = Tensor.empty(4, 4)
b = a.permute((1, 0))+0
check_schedule(b, 0)
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).st)
@expect_rangeify_fails
def test_become_existing_buf_view_alt(self):
a = Tensor.empty(4, 4)
b = a.permute((1, 0)).reshape((8, 2))+0
@@ -2581,6 +2490,7 @@ class TestUOpBecome(unittest.TestCase):
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)
# they can also have other base parents that simplified, in that case we just backtrack to the chained mops
@expect_rangeify_fails
def test_become_existing_buf_complex(self):
a = Tensor.empty(4, 4)
b = (a.permute((1, 0))+0).reshape((8, 2))+0
@@ -2588,6 +2498,7 @@ class TestUOpBecome(unittest.TestCase):
self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st)
assert b.uop.base.op is Ops.BUFFER
@expect_rangeify_fails
def test_become_multiple_choices(self):
a = Tensor.empty(16)
b = (a.reshape(1, 1, 4, 1, 4)+0).reshape(1, 1, 4, 4).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
@@ -2599,6 +2510,7 @@ class TestUOpBecome(unittest.TestCase):
assert b.uop is c.uop
assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {})
@expect_rangeify_fails
def test_setitem_becomes_subbuffer(self):
a = Tensor.full((4,), 2.).contiguous().realize()
b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0))