From cc9bf8ccbc0b7eb0e3b8510d475fa56263ef8cab Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Tue, 10 Feb 2026 13:35:17 +0800 Subject: [PATCH] move more to null/unit tests (#14658) * move more to null tests * move test_gc * no test fusion op --- .github/workflows/test.yml | 2 +- test/{ => null}/test_gc.py | 1 - test/null/test_schedule.py | 1000 +++++++++++++++++++++++++++++++++++- test/test_fusion_op.py | 68 --- test/test_schedule.py | 964 +--------------------------------- 5 files changed, 1015 insertions(+), 1020 deletions(-) rename test/{ => null}/test_gc.py (98%) delete mode 100644 test/test_fusion_op.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4f22319e03..a07cb1cde2 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -739,7 +739,7 @@ jobs: DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add - name: Run pytest (cuda) # skip multitensor because it's slow - run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20 + run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore=test/null --ignore test/test_multitensor.py --durations=20 - name: Run TestOps.test_add with PMA run: VIZ=-1 PMA=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add - name: Run process replay tests diff --git a/test/test_gc.py b/test/null/test_gc.py similarity index 98% rename from test/test_gc.py rename to test/null/test_gc.py index 3f3c8129b7..2865ca2b5e 100644 --- a/test/test_gc.py +++ b/test/null/test_gc.py @@ -72,7 +72,6 @@ class TestGC(unittest.TestCase): ys = y.schedule() del x run_schedule(ys) - np.testing.assert_equal(y.numpy(), np.full((256,), 2)) self.assertEqual(bufs_allocated()-init, 1) del y self.assertEqual(bufs_allocated()-init, 0) diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 035ba198f6..b31238fdd6 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -1,8 +1,9 @@ # schedule tests that pass on NULL backend (no copyout needed) -import unittest -from tinygrad import Tensor -from tinygrad.uop.ops import UOp -from tinygrad.helpers import DEBUG, Context +import unittest, time +from tinygrad import nn, dtypes, Device, Tensor +from tinygrad.device import is_dtype_supported +from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat +from tinygrad.helpers import DEBUG, GlobalCounters, Context from tinygrad.engine.realize import CompiledRunner, run_schedule class KernelCountException(Exception): pass @@ -26,6 +27,9 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te raise KernelCountException(f"{kernel_cnt} != {allowed}") return sched +def _realize_weights(m): + for p in nn.state.get_parameters(m): p.realize() + class TestBufferUOp(unittest.TestCase): # BUFFER has a ShapeTracker of shape=(n,) and stride=(1,) def test_buffer_has_buffer(self): @@ -139,5 +143,993 @@ class TestSimpleSchedule(unittest.TestCase): a2 = a.reshape(16,1,1) self.assertEqual(len(Tensor.schedule(a1, a2)), 1) +class TestSchedule(unittest.TestCase): + @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") + def test_error_on_device_mismatch(self): + a = Tensor.empty(10) + b = Tensor.empty(10, device="CPU") + c = a+b + with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1) + + @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") + 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, 2) + + def test_rand(self): + x = Tensor.rand(32) + check_schedule(x, 1, [Tensor._device_rng_counters[x.device]]) + + def test_rand_recompute_arange(self): + x = Tensor.rand(32) + check_schedule(x, 1, [Tensor._device_rng_counters[x.device]]) + + def test_empty_is_not_realized(self): + a = Tensor.empty(10) + child = a+2 + assert not a.uop.is_realized + child.realize() + assert a.uop.is_realized + + def test_realize_view_of_realized_has_empty_schedule(self): + # views of realized buffers produce an empty schedule + t = Tensor.zeros((3, 3)).contiguous().realize() + v = t[1] # view - is_realized but not has_buffer_identity + assert v.uop.is_realized + sched, _ = Tensor.schedule_with_vars(v) + self.assertEqual(len(sched), 0) + + # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated. + def test_childless_empty_never_allocates(self): + a = Tensor.empty(10) + a.realize() + assert not a.uop.is_realized + + def test_simplify_padded_const(self): + a, _ = Tensor.empty(1022).cummax(axis=0) + check_schedule(a, 3) + + def test_basic_binop_fusion(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = Tensor.empty(10) + d = a+b+c + check_schedule(d, 1) + + def test_basic_binop_fusion_deep(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = Tensor.empty(10) + d = Tensor.empty(10) + e = a+b+c+d + check_schedule(e, 1) + + def test_mulacc_fusion(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = (a*b).sum() + check_schedule(c, 1) + + def test_mulacc_relu_fusion(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = (a*b).sum().relu() + check_schedule(c, 1) + + def test_binop_reshape_fusion(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = Tensor.empty(5,2) + d = (a+b).reshape(5,2)+c + check_schedule(d, 1) + + def test_binop_permute_fusion(self): + a = Tensor.empty(2,5) + b = Tensor.empty(2,5) + c = Tensor.empty(5,2) + d = (a+b).permute(1,0)+c + check_schedule(d, 1) + + def test_constants_are_embedded(self): + a = Tensor.empty(3,3) * 2 + check_schedule(a, 1, filter_sink=False) + + def tests_constants_are_folded(self): + a = Tensor(2) + check_schedule(a, 0) + + def test_binop_elu_fusion(self): + a = Tensor.empty(10) + b = a.elu() + check_schedule(b, 1) + + def test_binop_reshape_reduce_fusion(self): + a = Tensor.empty(100) + b = Tensor.empty(100) + c = (a+b).reshape(10, 10).sum(axis=0, keepdim=True) + check_schedule(c, 1) + + def test_reduce_reshape_binop_fusion(self): + a = Tensor.empty(10,10) + b = Tensor.empty(10) + c = a.sum(axis=0) + b + check_schedule(c, 1) + + def test_reduce_permute_binop_fusion(self): + a = Tensor.empty(10,10,10) + b = Tensor.empty(10,10,1) + c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b + check_schedule(c, 1) + + def test_binop_early_reshape_reduce_fusion(self): + a = Tensor.empty(100) + b = Tensor.empty(100) + c = Tensor.empty(10,10) + d = ((a+b).reshape(10,10) + c).sum(axis=0) + check_schedule(d, 1) + + def test_diamond_folded(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = Tensor.empty(10) + d = Tensor.empty(10) + ab = a+b + e = (ab+c) + (ab+d) + check_schedule(e, 1) + + def test_cache_binaryop(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = a+b + d = a+b + check_schedule(d, 0, [c]) + + # failing in new lazy + def test_cache_binaryop_reshaped(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = a+b + d = a.reshape(10,1)+b.reshape(10,1) + check_schedule(d, 1, [c]) + + # failing in new lazy + def test_cache_binaryop_transpose(self): + a = Tensor.empty(10,10) + b = Tensor.empty(10,10) + c = (a.T*b.T).T #.contiguous() + d = a*b + check_schedule(d, 1, [c]) + + def test_cache_two_reduceops(self): + a = Tensor.empty(10) + b = a.sum() + c = a.sum() + bc = b+c + check_schedule(bc, 1) + + def test_cache_reduce_parent(self): + x = Tensor.empty(32) + r0 = x.mean(axis=0, keepdim=True) + 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 in {Ops.REDUCE_AXIS, Ops.REDUCE}] + assert len(reduceops) == 2 + + def test_cache_reduce_multiple_children(self): + x = Tensor.empty(32) + y = Tensor.empty(4, 4) + r0 = x.mean(axis=0, keepdim=True) + r1 = (x - r0).sum(axis=0).div(2) + out0 = r0 + y + out1 = r1 + y + schedule = check_schedule([out0, out1], 3) + reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}] + self.assertEqual(len(reduceops), 2) # why is RANGEIFY different? + + def test_dedup_assign(self): + a = Tensor.ones(4).contiguous().realize() + b = Tensor.full((4,), 2.).contiguous() + first = a.assign(b) + second = a.assign(b) + check_schedule([first, second], 2) # TODO: 1? + + def test_no_dedup_empty(self): + a = Tensor.empty((4,)) + b = Tensor.empty((4,)) + # NOTE: empty does not have any schedule + check_schedule([a, b], 0, filter_sink=False) + self.assertIsNot(a.uop.buffer, b.uop.buffer) + + def test_dedup_outputs(self): + a = Tensor.full((4, 4), 1.).contiguous().realize() + 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() + check_schedule(out, 1) + + #@unittest.skip("may want to reconsider this") + def test_fold_batchnorm(self): + with Tensor.train(): + img = Tensor.empty(1,32,4,4) + bn = nn.BatchNorm2d(32, track_running_stats=False) + out = bn(img) + check_schedule(out, 3) + + def test_fold_conv_batchnorm_notrain(self): + with Tensor.train(False): + img = Tensor.empty(1,3,8,8) + c1 = nn.Conv2d(3,32,3) + bn = nn.BatchNorm2d(32, track_running_stats=True) + out = bn(c1(img)).relu() + check_schedule(out, 1, [c1.weight, c1.bias]) + + def test_fold_conv_batchnorm_notrain_no_running_stats(self): + with Tensor.train(False): + img = Tensor.empty(1,3,8,8) + c1 = nn.Conv2d(3,32,3) + bn = nn.BatchNorm2d(32, track_running_stats=False) + out = bn(c1(img)).relu() + check_schedule(out, 4, [c1.weight, c1.bias]) + + def test_fold_conv_batchnorm(self): + with Tensor.train(): + img = Tensor.empty(1,3,8,8) + c1 = nn.Conv2d(3,32,3) + bn = nn.BatchNorm2d(32, track_running_stats=False) + out = bn(c1(img)).relu() + check_schedule(out, 4, [c1.weight, c1.bias]) + + @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong") + def test_fold_conv_batchnorm_optim(self): + # this is too high + for optim, cnt in [(nn.optim.Adam, 27), (nn.optim.SGD, 7)]: + with self.subTest(optim=optim.__name__): + with Tensor.train(): + img = Tensor.ones(1,3,4,4) + c1 = nn.Conv2d(3,32,3) + bn = nn.BatchNorm2d(32, track_running_stats=False) + _realize_weights([c1, bn]) + opt = optim(nn.state.get_parameters([c1, bn])) + img_bn = bn(c1(img)).elu().sum() + opt.zero_grad() + img_bn.backward() + check_schedule(opt.schedule_step(), cnt) + + def test_fold_batchnorm_backward(self): + with Tensor.train(): + x = Tensor.empty((2, 16, 8, 8)).contiguous() + bn = nn.BatchNorm2d(16) + bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True + fw = bn(x).contiguous_backward().relu().contiguous() + fw.sum().backward() + # TODO: this is too many + check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9) + + def test_fold_conv_relu(self): + c1 = nn.Conv2d(3,16,3) + # run + img = Tensor.ones(2,3,64,64) + out = c1(img).relu() + check_schedule(out, 1, [c1.weight, c1.bias]) + + def test_fold_conv_relu_alt(self): + img = Tensor.ones(1,4,8,8) + c1 = nn.Conv2d(4, 4, kernel_size=3) + c2 = nn.Conv2d(4, 4, kernel_size=3) + img_conv = img.sequential([c1, Tensor.relu, c2, Tensor.relu]) + check_schedule(img_conv, 2, [*nn.state.get_parameters(c1), *nn.state.get_parameters(c2), img]) + + def test_fold_conv_relu_nobias(self): + img = Tensor.ones(1,4,8,8) + c1 = nn.Conv2d(4, 4, kernel_size=3, bias=False) + c2 = nn.Conv2d(4, 4, kernel_size=3, bias=False) + out = img.sequential([c1, Tensor.relu, c2, Tensor.relu]) + check_schedule(out, 2, [c1.weight, c2.weight, img]) + + def test_fold_conv_elu(self): + c1 = nn.Conv2d(3,16,3) + # run + img = Tensor.rand(2,3,64,64) + out = c1(img).elu() + check_schedule(out, 1, [c1.weight, c1.bias, img]) + + def test_fold_conv_elu_alt(self): + img = Tensor.ones(1,4,8,8).contiguous() + c1 = nn.Conv2d(4, 4, kernel_size=3) + c2 = nn.Conv2d(4, 4, kernel_size=3) + img_conv = img.sequential([c1, Tensor.elu, c2, Tensor.elu]) + check_schedule(img_conv, 2, [*nn.state.get_parameters(c1), *nn.state.get_parameters(c2), img]) + + def test_two_sum(self): + img = Tensor.empty(64,64) + x = (img.sum(0) + img.sum(1)) + out = x.relu() + check_schedule(out, 1) + + def test_push_permute_through_reshape(self): + a = Tensor.empty(16,16) + b = Tensor.empty(16,16) + c = (a+b).reshape(4,4,4,4).permute(2,3,0,1).contiguous() + check_schedule(c, 1) + + #@unittest.skip("failing in old lazy") + def test_push_permute_through_reshape_alt(self): + a = Tensor.empty(4,4,4,4) + b = Tensor.empty(4,4,4,4) + c = (a+b).reshape(16,16).permute(1,0).contiguous() + check_schedule(c, 1) + + def test_no_binop_rerun(self): + a = Tensor.empty(16) + b = Tensor.empty(16) + c = a+b + d = (a+b).reshape(16,1) + check_schedule(d, 0, [c]) + + @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + def test_multi_permute_should_collapse(self): + a = Tensor.empty(4,4,4,4) + b = Tensor.empty(16) + c = a.sum((0,1)).cast(dtypes.float16).permute(1,0).reshape(4,4,1).permute(1,0,2).reshape(16) + b + check_schedule(c, 1) + + def test_fancy_reshape_fusion(self): + a = Tensor.empty(10) + b = Tensor.empty(10) + c = a+b + d = a.reshape(10,1)+b.reshape(10,1) + out = c.sum() + d.sum() + check_schedule(out, 1) + + def test_children_dont_push(self): + a = Tensor.empty(10, 10, 1) + b = Tensor.empty(10, 10, 1) + d = (a+b).expand(10, 10, 10) + e = (a+b).permute(2,1,0) + f = d+e + check_schedule(f, 1) + + # 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) + c = Tensor.empty(10) + 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 + check_schedule(d, 2) + check_schedule(keep_me, 0, [d]) + + #@unittest.skip("failing in old lazy") + def test_permute_breaks_fusion(self): + a = Tensor.empty(10, 10, 10) + b = Tensor.empty(10, 10) + c = (a.sum(axis=2) + b).permute(1,0) + d = c.permute(1,0) + check_schedule(d, 1) + + def test_some_permute_fusion(self): + a = Tensor.empty(8192, 16) + b = Tensor.empty(1, 16) + d = (a.T + b.expand(8192, 16).T) + c = a + b.expand(8192, 16) + e = d.T + check_schedule(c, 1) + check_schedule(e, 1) + + def test_shrink_fuse(self): + a = Tensor.empty(8192, 16) + b = Tensor.empty(8192, 16) + c = a * b + d = Tensor.empty(1, 16) + e = c[0] * d + check_schedule(e, 1) + + def test_expand_fuse(self): + a = Tensor.empty(1, 16) + b = Tensor.empty(1, 16) + c = a * b + d = Tensor.empty(8192, 16) + e = c * d + check_schedule(e, 1) + + # this is the failing case in openpilot...it's very simple like this + def test_image_conv_fusion(self): + w1 = Tensor.empty(16, 16, 1, 1) + b1 = Tensor.empty(16) + w2 = Tensor.empty(16, 16, 1, 1) + b2 = Tensor.empty(16) + w3 = Tensor.empty(16, 16, 1, 1) + b3 = Tensor.empty(16) + + x = Tensor.empty(1, 16, 32, 32) + x = base = x.image_conv2d(w1, b1) + x = x.image_conv2d(w2, b2) + base + x = x.image_conv2d(w3, b3) + + # NOOP, 3 convs, contiguous + #check_schedule(x, 5) + check_schedule(x, 7) + + def test_image_conv_fusion_minimal(self): + b1 = Tensor.empty(16) + b2 = Tensor.empty(16) + def p(x): return x.permute(1,0).contiguous().reshape(32,16,1).expand(32,16,16).sum(axis=2).permute(1,0) + + x = Tensor.empty(16, 32) + x = base = p(x) + b1.reshape(16,1) + x = p(x) + x = x + b2.reshape(16,1) + x = x + base + del base + x = p(x) + check_schedule(x, 4) + + def test_image_conv_fusion_more_minimal(self): + b1 = Tensor.empty(16) + def p(x): return x.permute(1,0).contiguous().reshape(32,16,1).expand(32,16,16).sum(axis=2).permute(1,0) + + x = Tensor.empty(16, 32) + x = base = p(x) + b1.reshape(16,1) + x = p(x) + del base + check_schedule(x, 3) + + def test_contiguous_while_contiguous(self): + x = Tensor.empty(1, 64, 32, 32) + out = x.contiguous() + check_schedule(out, 0, filter_sink=False) + + def test_contiguous_while_not_contiguous(self): + x = Tensor.empty(1, 64, 32, 32) + out = x.permute(0,2,3,1).contiguous() + check_schedule(out, 1, filter_sink=False) + + def test_fold_with_contiguous(self): + a = Tensor.randn(16, 16, 16).realize() + b = Tensor.randn(16, 16).realize() + c = (a.sum(2).contiguous() + b).contiguous() + check_schedule(c, 2) + + def _alu_from_tensor(self, t:Tensor): + s = [s for s in t.schedule() if s.ast.op is Ops.SINK] + self.assertEqual(len(s), 1) + return [u.op for u in s[0].ast.toposort() if u.op in GroupOp.ALU] + + def test_2_pow_is_exp2(self): + t = 2.0 ** Tensor([1.0, 2.0, 3.0]) + self.assertEqual(self._alu_from_tensor(t), [Ops.EXP2]) + + def test_pow_05_is_sqrt(self): + t = Tensor([1.0, 2.0, 3.0]) ** 0.5 + self.assertEqual(self._alu_from_tensor(t), [Ops.SQRT]) + + def test_pow_neg_05_is_rsqrt(self): + t = Tensor([1.0, 2.0, 3.0]) ** -0.5 + self.assertEqual(self._alu_from_tensor(t), [Ops.RECIPROCAL, Ops.SQRT]) + + def test_pow_2_has_1_mul(self): + t = Tensor([1.0, 2.0, 3.0]) ** Tensor(2.0) + self.assertEqual(self._alu_from_tensor(t), [Ops.MUL]) + + def test_pow_8_has_3_muls(self): + t = Tensor([1.0, 2.0, 3.0]) ** 8 + self.assertEqual(self._alu_from_tensor(t), [Ops.MUL, Ops.MUL, Ops.MUL]) + + def test_pow_const_tensor_to_zero(self): + x = Tensor([1,2,3,4]) + out = x ** Tensor(0.0) + # NOTE: this is UOp.const(0) + UOp.const(1) + check_schedule(out, 0) + + def test_zero_size(self): + x = Tensor.empty(2, 3, 0) + out = x + 1 + check_schedule(out, 0, filter_sink=False) + + def test_reduce_permute_nofuse(self): + x = Tensor.empty(32, 32, 32) + y = Tensor.empty(32, 32) + out = x.sum(axis=2).T+y + check_schedule(out, 1) + + def test_two_elus_sum(self): + x = Tensor.empty(32, 32) + y = Tensor.empty(32, 32) + out = x.sum(1).relu().elu() + y.sum(1).relu().elu() + check_schedule(out, 1) + + def test_multistage_reduce(self): + x = Tensor.empty(32, 32, 32) + out = x.sum(2).relu().sum(1) + check_schedule(out, 1) + + def test_multistage_reduce_fork(self): + x = Tensor.empty(32, 32, 32) + x = x.sum(2) + out2 = x + 1 + out = x.relu().sum(1) + out2[0] + check_schedule(out, 2) + + def test_contiguous_add(self): + x = Tensor.empty(32) + y = Tensor.empty(32) + z = Tensor.empty(32) + out = (x+y).contiguous()+z + check_schedule(out, 2) + + def test_double_sum_ref(self): + x = Tensor.empty(32, 32, 32) + x = x.sum(2) + out = x + x[:, 4] + check_schedule(out, 2) + + def test_reduce_shrink(self): + x = Tensor.empty(32, 32) + y = Tensor.empty(16) + x = x.sum(1) + x = x[:16] + out = x + y + check_schedule(out, 1) + + def test_const_no_recompute(self): + x = Tensor(2) + Tensor(2) + y = Tensor(2) + Tensor(2) + out = x.contiguous() + y.contiguous() + check_schedule(out, 2, filter_sink=False) + + 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], 2) # TODO: 1? + + def test_reduce_multiple_paths_midshrink(self): + a = Tensor.empty(4, 4) + r = a.sum(axis=1) + out0 = r.exp2() + out1 = out0[0] + out0 + check_schedule([r, out0, out1], 3) + + def test_reduce_shrink_output(self): + a = Tensor.empty(4, 4) + r = a.sum(keepdim=True) + out0 = r.exp2() + out1 = out0[0] + Tensor.empty(1, ) + check_schedule([r, out0, out1], 3) + + @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") + 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(sched[0].bufs[0].dtype, dtypes.float) + + # input float, softmax in float + Tensor.manual_seed(0) + x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.float).realize() + out = x.softmax(dtype=dtypes.float) + sched = out.schedule() + self.assertEqual(len(sched), 3) + self.assertEqual(sched[0].bufs[0].dtype, dtypes.float) + + def test_softmax_backward(self): + Tensor.manual_seed(0) + x = Tensor.randn(4, 12, 64, 64, requires_grad=True).realize() + x.softmax().sum().backward() + run_schedule(check_schedule(x.grad, 4)) + + def test_scaled_dot_product_attention_fusion(self): + x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4)) + out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m) + check_schedule(out, 4) + + def test_scaled_dot_product_attention_causal_fusion(self): + x, y, z = (Tensor.empty(32, 8, 16, 16) for _ in range(3)) + out = Tensor.scaled_dot_product_attention(x, y, z, is_causal=True) + check_schedule(out, 4) + + def test_adam_step_fusion(self): + with Tensor.train(): + x = Tensor.empty(4, 64, 32) + layer = nn.Linear(32, 32*4) + _realize_weights(layer) + opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4) + layer(x).relu().sum().backward() + check_schedule(opt.schedule_step(), 19) + + def test_adam_conv_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,4,4) + c1 = nn.Conv2d(3,32,3) + _realize_weights(c1) + opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4) + opt.zero_grad() + c1(img).relu().sum().backward() + check_schedule(opt.schedule_step(), 19) + + def test_adam_2convs_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,4,4) + c1 = nn.Conv2d(3,16,3,bias=False) + c2 = nn.Conv2d(16,32,2,bias=False) + _realize_weights([c1, c2]) + opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4) + opt.zero_grad() + c2(c1(img).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 21) + + def test_sgd_conv_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,4,4) + c1 = nn.Conv2d(3,32,3) + _realize_weights(c1) + opt = nn.optim.SGD(nn.state.get_parameters(c1)) + opt.zero_grad() + c1(img).relu().sum().backward() + check_schedule(opt.schedule_step(), 5) # TODO: 3? + + def test_sgd_2convs_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,4,4) + c1 = nn.Conv2d(3,16,3,bias=False) + c2 = nn.Conv2d(16,32,2,bias=False) + _realize_weights([c1, c2]) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2])) + opt.zero_grad() + c2(c1(img).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 7) + + @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong") + def test_fold_2convs_sgd_nesterov_momentum_wd(self): + with Tensor.train(): + img = Tensor.empty(2,3,4,4) + c1 = nn.Conv2d(3,16,3,bias=False) + c2 = nn.Conv2d(16,32,2,bias=False) + _realize_weights([c1, c2]) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1) + opt.zero_grad() + c2(c1(img).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 13) + + def test_sgd_4convs_fuse(self): + with Tensor.train(): + img = Tensor.empty(2,3,16,16) + c1 = nn.Conv2d(3,4,3,bias=False) + c2 = nn.Conv2d(4,8,3,bias=False) + c3 = nn.Conv2d(8,16,3,bias=False) + c4 = nn.Conv2d(16,32,3,bias=False) + _realize_weights([c1, c2, c3, c4]) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) + opt.zero_grad() + c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 15) + + def test_sgd_4convs_fuse_conv_bw(self): + with Tensor.train(): + img = Tensor.empty(2,3,16,16) + c1 = nn.Conv2d(3,4,3,bias=False) + c2 = nn.Conv2d(4,8,3,bias=False) + c3 = nn.Conv2d(8,16,3,bias=False) + c4 = nn.Conv2d(16,32,3,bias=False) + _realize_weights([c1, c2, c3, c4]) + opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) + opt.zero_grad() + c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() + check_schedule(opt.schedule_step(), 15) + + def test_reduce_simple_chase(self): + a = Tensor.empty(4, 4, 4) + r = a.sum(0) + 6 + b = r.sum(0) * 4 + c = r.sum(1) * 2 + check_schedule([b, c], 3) + + def test_push_permute_chase(self): + a = Tensor.empty(4, 4, 4) + b = Tensor.empty(4, 4) + r = a.sum(2) + b + d = r.T * 4 + e = r * d + check_schedule([d, e], 3) + + def test_push_shrink_chase(self): + a = Tensor.empty(16, 16) + b = Tensor.empty(4) + c = Tensor.empty(16, ) + r = a.sum(1) + c + d = r[:4] * b + check_schedule(d, 1) + + def test_midreduce_nochase(self): + a = Tensor.empty(16, 16) + b = (a.sum(0) + a.max(1)) + 2 + check_schedule(b, 1) + + def test_bitcast_fuses(self): + x = Tensor.empty(1, dtype=dtypes.float32) + a = x.exp2().bitcast(dtypes.int32) + b = x.bitcast(dtypes.int32) + check_schedule(a+b, 1) # this should fuse when it makes sense + + def test_reduceop_reshape_dont_push(self): + Tensor.manual_seed(0) + x = Tensor.randn(10, 20).realize() + out = x.argmax(1) + run_schedule(check_schedule(out, 2)) + + def test_resnet_conv2d(self): + x = Tensor.empty(1, 8, 32, 32) + w1 = Tensor.empty(8, 8, 3, 3) + w2 = Tensor.empty(8, 8, 1, 1) + out = x.conv2d(w1).conv2d(w2) + check_schedule(out, 2) + + def test_schedule_mem_used(self): + base = GlobalCounters.mem_used + Tensor.ones(256).contiguous().realize() + Tensor.ones(5, 5).contiguous().schedule() + self.assertEqual(GlobalCounters.mem_used-base, 0) + + def test_const_schedule(self): + constv = Tensor.empty(2, 2).uop.const_like(10) + check_schedule(constv, 0) + + def test_const_schedule_contig(self): + constv = Tensor.empty(2, 2).uop.const_like(10).contiguous() + check_schedule(constv, 1) + + def test_advanced_simple_indexing_combined(self): + X = Tensor.arange(16).reshape(4, 4) + xt = X[1:2, [-1, 2]] + check_schedule(xt, 1) + + def test_arange_index_shrink(self): + Tensor.manual_seed(0) + with Context(TRACK_MATCH_STATS=0): + x = Tensor.randn(11).realize() + a = Tensor.arange(22) + out = (x + a[:11]).sum() + check_schedule(out, 1) + + def test_fuse_arange_avg_pool2d_ceil_mode(self): + x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True) + sched = check_schedule(x, 1) + self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1) + + def test_fuse_arange_pad_circular_mode_bw(self): + x = Tensor.empty(1,1,5,5,5) + out = x.pad((1,2,3,5,1,2), mode="circular") + g = out.sum().gradient(x)[0] + sched = check_schedule(g, 1) + self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0) + + def test_resnet_block(self): + with Tensor.train(False): + in_planes, planes = 64, 64 + conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + bn1 = nn.BatchNorm2d(planes) + conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, stride=1, bias=False) + bn2 = nn.BatchNorm2d(planes) + x = Tensor.empty(1, 64, 32, 32) + out = bn1(conv1(x)).relu() + out = bn2(conv2(out)) + out = (out + x).relu() + run_schedule(check_schedule(out, 2, [conv1.weight, conv2.weight])) + +class TestSwizzle(unittest.TestCase): + def test_softmax_one_kernel(self): + Tensor.manual_seed(0) + with Context(DEBUG=0, TRACK_MATCH_STATS=0): + a = Tensor.randn(32, 32).realize() + t = a.softmax() + check_schedule(t, 3) # TODO: 1? + + def test_argmax_one_kernel(self): + Tensor.manual_seed(0) + with Context(DEBUG=0, TRACK_MATCH_STATS=0): + a = Tensor.randn(10, 20).realize() + t = a.argmax(0) + check_schedule(t, 2) # TODO: 1? + +class TestView(unittest.TestCase): + def test_zero_size_alt(self): + a = Tensor.empty(135, 0, 9) + b = a.pad(((0, 0), (0, 0), (18, 0))) + check_schedule(b, 0) + +class TestUOpBecome(unittest.TestCase): + # the simplest case, if we create a new BUFFER for this tensor UOp + def test_new_buffer(self): + a = Tensor.empty(4, 4) + b = Tensor.empty(4, 4) + add = a+b + check_schedule(add, 1) + # NOTE: realized base is always a flat buffer + assert UPat(Ops.BUFFER).match(add.uop.base, {}) + # the Tensor UOp can optionally stack a VIEW on top of the BUFFER, in this case to preserve the (4, 4) shape of the tensor + assert add.uop is not add.uop.base + self.assertEqual(add.uop.size, 16) + self.assertEqual(add.uop.shape, (4, 4)) + + def test_new_buffer_view(self): + a = Tensor.empty(4, 4) + b = Tensor.empty(4, 4) + add = (a+b).reshape(8, 2) + check_schedule(add, 1) + assert UPat(Ops.BUFFER).match(add.uop.base, {}) + # the shape is preserverd in the becomes_map. + self.assertEqual(add.uop.shape, (8, 2)) + assert add.uop is not add.uop.base + + def test_new_flat_buffer(self): + a = Tensor.empty(4,) + b = Tensor.empty(4,) + add = a+b + check_schedule(add, 1) + # BUFFER already has a shape (4,), this tensor just becomes a contiguous BUFFER + assert UPat(Ops.BUFFER).match(add.uop.base, {}) + + # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer + + def test_reorder_expand(self): + a = Tensor.empty(4, 1) + b = a.expand(4, 4).reciprocal() + check_schedule(b, 1) + self.assertEqual(b.uop.base.buffer.size, 4) + self.assertEqual(b.uop.shape, (4, 4)) + + def test_reorder_expand_alt(self): + x = Tensor.empty(4, 1) + y = Tensor.empty(4, 1) + img = Tensor.empty(4, 4) + z = (img*x) / y + check_schedule(z, 1) + + # TODO: rangeify doesn't yet cleanup this kind of re-indexing + @unittest.expectedFailure + def test_become_existing_buffer(self): + a = Tensor.empty(4, 4) + b = a*1 + assert UPat(Ops.MUL).match(b.uop, {}) # before scheduling it's a mul + check_schedule(b, 0) + self.assertIs(a.uop.base.buffer, b.uop.base.buffer) + + def test_become_buf_with_mops(self): + a = Tensor.empty(2, 4, 2) + noop = a.shrink(((1, 2), (0, 4), (0, 2))).reshape(4, 2)*1+0 + # before realizing, this tensor is base + assert noop.uop is noop.uop.base + noop.realize() + # it becomes a realized view after realize + assert noop.uop is not noop.uop.base + assert noop.uop.base.op is Ops.BUFFER + late_add = noop+2 + late_add.realize() + + def test_become_const_in_base(self): + a = Tensor.empty(4) + b = a*0 + assert UPat(Ops.MUL).match(b.uop, {}) # before scheduling it's a mul + check_schedule(b, 0) + assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER) + + def test_become_const_from_const(self): + const_add = Tensor(1)+Tensor(2) + assert UPat(Ops.ADD).match(const_add.uop, {}) + check_schedule(const_add, 0) + assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {}) + + # tensors can become another realized tensor source + @unittest.expectedFailure + def test_become_existing_buf_simple(self): + a = Tensor.empty(4, 4) + b = a+0 + check_schedule(b, 0) + assert b.uop.base.op is Ops.BUFFER + self.assertIs(a.uop, b.uop) + + # they can also chain other movement ops on top of the tensor source + @unittest.expectedFailure + 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) + + @unittest.expectedFailure + def test_become_existing_buf_view_alt(self): + a = Tensor.empty(4, 4) + b = a.permute((1, 0)).reshape((8, 2))+0 + check_schedule(b, 0) + 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 + @unittest.expectedFailure + def test_become_existing_buf_complex(self): + a = Tensor.empty(4, 4) + b = (a.permute((1, 0))+0).reshape((8, 2))+0 + check_schedule(b, 0) + self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st) + assert b.uop.base.op is Ops.BUFFER + + @unittest.expectedFailure + 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 + c = (a.reshape(1, 1, 4, 4)+0).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0 + check_schedule([b, c], 0) + from tinygrad.helpers import all_same + assert all_same([x.uop.base.realized for x in [a,b,c]]) + + def test_setitem_becomes_subbuffer(self): + a = Tensor.full((4,), 2.).contiguous().realize() + b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0)) + b.realize() + assert a.uop.is_realized + assert a.uop.buffer._base is None + assert b.uop.op_in_backward_slice_with_self(Ops.SHRINK) + assert b.uop.base is a.uop.base + +class TestFusionOp(unittest.TestCase): + def test_recursive_add(self): + st = time.perf_counter() + a = Tensor([1,2,3,4]) + for _ in range(24): a = a + a + sched = a.schedule() + sched[-1].lower() + self.assertLess(time.perf_counter()-st, 2.0) + assert len(sched[-1].prg.p.src.splitlines()) < 250 + + def test_recursive_add_cmp(self): + st = time.perf_counter() + a = Tensor([1,2,3,4]) + for _ in range(24): a = a + a + sched1 = a.schedule() + b = Tensor([1,2,3,4]) + for _ in range(24): b = b + b + sched2 = b.schedule() + c = Tensor([1,2,3,4]) + for _ in range(23): c = c + c + sched3 = c.schedule() + self.assertEqual(sched1[-1].ast, sched2[-1].ast) + with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast) + self.assertLess(time.perf_counter()-st, 2.0) + + def test_recursive_pad(self): + st = time.perf_counter() + val = 1.0 + a = Tensor(val) + for _ in range(24): a = Tensor.stack(a, a)[0] + sched = a.schedule() + self.assertEqual(len(sched), 0) + self.assertLess(time.perf_counter()-st, 2.0) + + def test_recursive_reshape(self): + st = time.perf_counter() + a = Tensor.empty(32, 32).realize() + b = Tensor.empty(16, 2).realize() + r = a.sum(1) + for _ in range(24): r = r.reshape(16, 2) + b + sched = r.schedule() + self.assertEqual(len(sched), 1) + self.assertLess(time.perf_counter()-st, 2.0) + if __name__ == '__main__': unittest.main(verbosity=2) diff --git a/test/test_fusion_op.py b/test/test_fusion_op.py deleted file mode 100644 index a01c80c7d9..0000000000 --- a/test/test_fusion_op.py +++ /dev/null @@ -1,68 +0,0 @@ -import unittest -import time -import numpy as np -from tinygrad import Tensor, dtypes -from tinygrad.engine.realize import run_schedule - -class TestFusionOp(unittest.TestCase): - def test_contiguous_add(self): - def test(contig=False): - bt = Tensor(np.arange(16), dtype=dtypes.float32).reshape(4,4) - x = bt.permute(1,0) - if contig: x = x.contiguous() - return (x.permute(1,0) + bt).data() - assert test() == test(True) - - def test_expand_fuse(self): - bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32) - out = (bt*2).expand(10,10).sum(1) - sched = out.schedule() - run_schedule(sched) - outd = out.tolist() - assert all(x == 20.0 for x in outd) - - def test_recursive_add(self): - st = time.perf_counter() - a = Tensor([1,2,3,4]) - for _ in range(24): a = a + a - sched = a.schedule() - sched[-1].lower() - self.assertLess(time.perf_counter()-st, 2.0) - assert len(sched[-1].prg.p.src.splitlines()) < 250 - - def test_recursive_add_cmp(self): - st = time.perf_counter() - a = Tensor([1,2,3,4]) - for _ in range(24): a = a + a - sched1 = a.schedule() - b = Tensor([1,2,3,4]) - for _ in range(24): b = b + b - sched2 = b.schedule() - c = Tensor([1,2,3,4]) - for _ in range(23): c = c + c - sched3 = c.schedule() - self.assertEqual(sched1[-1].ast, sched2[-1].ast) - with self.assertRaises(AssertionError): self.assertEqual(sched1[-1].ast, sched3[-1].ast) - self.assertLess(time.perf_counter()-st, 2.0) - - def test_recursive_pad(self): - st = time.perf_counter() - val = 1.0 - a = Tensor(val) - for _ in range(24): a = Tensor.stack(a, a)[0] - sched = a.schedule() - self.assertEqual(len(sched), 0) - self.assertLess(time.perf_counter()-st, 2.0) - - def test_recursive_reshape(self): - st = time.perf_counter() - a = Tensor.empty(32, 32).realize() - b = Tensor.empty(16, 2).realize() - r = a.sum(1) - for _ in range(24): r = r.reshape(16, 2) + b - sched = r.schedule() - self.assertEqual(len(sched), 1) - self.assertLess(time.perf_counter()-st, 2.0) - -if __name__ == '__main__': - unittest.main(verbosity=2) diff --git a/test/test_schedule.py b/test/test_schedule.py index b97b16783d..f3f1355ddd 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -10,7 +10,7 @@ from hypothesis import assume, given, settings, strategies as strat from tinygrad import nn, dtypes, Device, Tensor, Variable from tinygrad.device import is_dtype_supported from tinygrad.dtype import DType, ImageDType -from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat +from tinygrad.uop.ops import UOp, Ops, UPat from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp from tinygrad.engine.realize import CompiledRunner, run_schedule @@ -98,20 +98,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(a, 1)) self.assertListEqual(a.tolist(), [[15]]) - @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") - def test_error_on_device_mismatch(self): - a = Tensor.empty(10) - b = Tensor.empty(10, device="CPU") - c = a+b - with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 1) - - @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch") - 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, 2) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_expand_buffer_before_cast(self): a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2)) @@ -159,116 +145,11 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(z, 1, [x,y])) self.assertEqual(z.item(), 32) - def test_rand(self): - x = Tensor.rand(32) - check_schedule(x, 1, [Tensor._device_rng_counters[x.device]]) - - def test_rand_recompute_arange(self): - x = Tensor.rand(32) - check_schedule(x, 1, [Tensor._device_rng_counters[x.device]]) - - def test_empty_is_not_realized(self): - a = Tensor.empty(10) - child = a+2 - assert not a.uop.is_realized - child.realize() - assert a.uop.is_realized - - def test_realize_view_of_realized_has_empty_schedule(self): - # views of realized buffers produce an empty schedule - t = Tensor.zeros((3, 3)).contiguous().realize() - v = t[1] # view - is_realized but not has_buffer_identity - assert v.uop.is_realized - sched, _ = Tensor.schedule_with_vars(v) - self.assertEqual(len(sched), 0) - - # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated. - def test_childless_empty_never_allocates(self): - a = Tensor.empty(10) - a.realize() - assert not a.uop.is_realized - - def test_simplify_padded_const(self): - a, _ = Tensor.empty(1022).cummax(axis=0) - check_schedule(a, 3) - - def test_basic_binop_fusion(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = Tensor.empty(10) - d = a+b+c - check_schedule(d, 1) - - def test_basic_binop_fusion_deep(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = Tensor.empty(10) - d = Tensor.empty(10) - e = a+b+c+d - check_schedule(e, 1) - - def test_mulacc_fusion(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = (a*b).sum() - check_schedule(c, 1) - - def test_mulacc_relu_fusion(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = (a*b).sum().relu() - check_schedule(c, 1) - - def test_binop_reshape_fusion(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = Tensor.empty(5,2) - d = (a+b).reshape(5,2)+c - check_schedule(d, 1) - - def test_binop_permute_fusion(self): - a = Tensor.empty(2,5) - b = Tensor.empty(2,5) - c = Tensor.empty(5,2) - d = (a+b).permute(1,0)+c - check_schedule(d, 1) - - def test_constants_are_embedded(self): - a = Tensor.empty(3,3) * 2 - check_schedule(a, 1, filter_sink=False) - - def tests_constants_are_folded(self): - a = Tensor(2) - check_schedule(a, 0) - def test_constants_can_store(self): a = Tensor(2).contiguous() run_schedule(check_schedule(a, 1)) np.testing.assert_equal(a.numpy(), 2) - def test_binop_elu_fusion(self): - a = Tensor.empty(10) - b = a.elu() - check_schedule(b, 1) - - def test_binop_reshape_reduce_fusion(self): - a = Tensor.empty(100) - b = Tensor.empty(100) - c = (a+b).reshape(10, 10).sum(axis=0, keepdim=True) - check_schedule(c, 1) - - def test_reduce_reshape_binop_fusion(self): - a = Tensor.empty(10,10) - b = Tensor.empty(10) - c = a.sum(axis=0) + b - check_schedule(c, 1) - - def test_reduce_permute_binop_fusion(self): - a = Tensor.empty(10,10,10) - b = Tensor.empty(10,10,1) - c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b - check_schedule(c, 1) - def test_allow_push_permutes(self): a = Tensor.randn(10,10,10).realize() b = Tensor.randn(10,10,1).realize() @@ -276,72 +157,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(c, 1)) np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy()) - def test_binop_early_reshape_reduce_fusion(self): - a = Tensor.empty(100) - b = Tensor.empty(100) - c = Tensor.empty(10,10) - d = ((a+b).reshape(10,10) + c).sum(axis=0) - check_schedule(d, 1) - - def test_diamond_folded(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = Tensor.empty(10) - d = Tensor.empty(10) - ab = a+b - e = (ab+c) + (ab+d) - check_schedule(e, 1) - - def test_cache_binaryop(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = a+b - d = a+b - check_schedule(d, 0, [c]) - - # failing in new lazy - def test_cache_binaryop_reshaped(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = a+b - d = a.reshape(10,1)+b.reshape(10,1) - check_schedule(d, 1, [c]) - - # failing in new lazy - def test_cache_binaryop_transpose(self): - a = Tensor.empty(10,10) - b = Tensor.empty(10,10) - c = (a.T*b.T).T #.contiguous() - d = a*b - check_schedule(d, 1, [c]) - - def test_cache_two_reduceops(self): - a = Tensor.empty(10) - b = a.sum() - c = a.sum() - bc = b+c - check_schedule(bc, 1) - - def test_cache_reduce_parent(self): - x = Tensor.empty(32) - r0 = x.mean(axis=0, keepdim=True) - 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 in {Ops.REDUCE_AXIS, Ops.REDUCE}] - assert len(reduceops) == 2 - - def test_cache_reduce_multiple_children(self): - x = Tensor.empty(32) - y = Tensor.empty(4, 4) - r0 = x.mean(axis=0, keepdim=True) - r1 = (x - r0).sum(axis=0).div(2) - out0 = r0 + y - out1 = r1 + y - schedule = check_schedule([out0, out1], 3) - reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}] - self.assertEqual(len(reduceops), 2) # why is RANGEIFY different? - def test_div_collapse_buffer(self): a = Tensor.full((4,), 4.0).contiguous().realize() b = Tensor.full((4,), 2.0).contiguous().realize() @@ -366,13 +181,6 @@ class TestSchedule(unittest.TestCase): self.assertLessEqual(GlobalCounters.global_ops, 4*3) np.testing.assert_allclose(expr.numpy(), (a.numpy()/b.numpy())/c.numpy()) - def test_dedup_assign(self): - a = Tensor.ones(4).contiguous().realize() - b = Tensor.full((4,), 2.).contiguous() - first = a.assign(b) - second = a.assign(b) - check_schedule([first, second], 2) # TODO: 1? - # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562 # should contiguous dedup? @unittest.skip("we do the exact opposite now") @@ -393,332 +201,12 @@ class TestSchedule(unittest.TestCase): # a and b are assigned to the same device Buffer self.assertIsNot(a.uop.base.realized, b.uop.base.realized) - # EMPTY is assigned to a unique device Buffer - - def test_no_dedup_empty(self): - a = Tensor.empty((4,)) - b = Tensor.empty((4,)) - # NOTE: empty does not have any schedule - check_schedule([a, b], 0, filter_sink=False) - self.assertIsNot(a.uop.buffer, b.uop.buffer) - - def test_dedup_outputs(self): - a = Tensor.full((4, 4), 1.).contiguous().realize() - 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() - check_schedule(out, 1) - - #@unittest.skip("may want to reconsider this") - def test_fold_batchnorm(self): - with Tensor.train(): - img = Tensor.empty(1,32,4,4) - bn = nn.BatchNorm2d(32, track_running_stats=False) - out = bn(img) - check_schedule(out, 3) - - def test_fold_conv_batchnorm_notrain(self): - with Tensor.train(False): - img = Tensor.empty(1,3,8,8) - c1 = nn.Conv2d(3,32,3) - bn = nn.BatchNorm2d(32, track_running_stats=True) - out = bn(c1(img)).relu() - check_schedule(out, 1, [c1.weight, c1.bias]) - - def test_fold_conv_batchnorm_notrain_no_running_stats(self): - with Tensor.train(False): - img = Tensor.empty(1,3,8,8) - c1 = nn.Conv2d(3,32,3) - bn = nn.BatchNorm2d(32, track_running_stats=False) - out = bn(c1(img)).relu() - check_schedule(out, 4, [c1.weight, c1.bias]) - - def test_fold_conv_batchnorm(self): - with Tensor.train(): - img = Tensor.empty(1,3,8,8) - c1 = nn.Conv2d(3,32,3) - bn = nn.BatchNorm2d(32, track_running_stats=False) - out = bn(c1(img)).relu() - check_schedule(out, 4, [c1.weight, c1.bias]) - - @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong") - def test_fold_conv_batchnorm_optim(self): - # this is too high - for optim, cnt in [(nn.optim.Adam, 27), (nn.optim.SGD, 7)]: - with self.subTest(optim=optim.__name__): - with Tensor.train(): - img = Tensor.ones(1,3,4,4) - c1 = nn.Conv2d(3,32,3) - bn = nn.BatchNorm2d(32, track_running_stats=False) - _realize_weights([c1, bn]) - opt = optim(nn.state.get_parameters([c1, bn])) - img_bn = bn(c1(img)).elu().sum() - opt.zero_grad() - img_bn.backward() - check_schedule(opt.schedule_step(), cnt) - - def test_fold_batchnorm_backward(self): - with Tensor.train(): - x = Tensor.empty((2, 16, 8, 8)).contiguous() - bn = nn.BatchNorm2d(16) - bn.weight.requires_grad = bn.bias.requires_grad = x.requires_grad = True - fw = bn(x).contiguous_backward().relu().contiguous() - fw.sum().backward() - # TODO: this is too many - check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9) - - def test_fold_conv_relu(self): - c1 = nn.Conv2d(3,16,3) - - # run - img = Tensor.ones(2,3,64,64) - out = c1(img).relu() - check_schedule(out, 1, [c1.weight, c1.bias]) - - def test_fold_conv_relu_alt(self): - img = Tensor.ones(1,4,8,8) - c1 = nn.Conv2d(4, 4, kernel_size=3) - c2 = nn.Conv2d(4, 4, kernel_size=3) - img_conv = img.sequential([c1, Tensor.relu, c2, Tensor.relu]) - check_schedule(img_conv, 2, [*nn.state.get_parameters(c1), *nn.state.get_parameters(c2), img]) - - def test_fold_conv_relu_nobias(self): - img = Tensor.ones(1,4,8,8) - c1 = nn.Conv2d(4, 4, kernel_size=3, bias=False) - c2 = nn.Conv2d(4, 4, kernel_size=3, bias=False) - out = img.sequential([c1, Tensor.relu, c2, Tensor.relu]) - check_schedule(out, 2, [c1.weight, c2.weight, img]) - - def test_fold_conv_elu(self): - c1 = nn.Conv2d(3,16,3) - - # run - img = Tensor.rand(2,3,64,64) - out = c1(img).elu() - check_schedule(out, 1, [c1.weight, c1.bias, img]) - - def test_fold_conv_elu_alt(self): - img = Tensor.ones(1,4,8,8).contiguous() - c1 = nn.Conv2d(4, 4, kernel_size=3) - c2 = nn.Conv2d(4, 4, kernel_size=3) - img_conv = img.sequential([c1, Tensor.elu, c2, Tensor.elu]) - check_schedule(img_conv, 2, [*nn.state.get_parameters(c1), *nn.state.get_parameters(c2), img]) - - def test_two_sum(self): - img = Tensor.empty(64,64) - x = (img.sum(0) + img.sum(1)) - out = x.relu() - check_schedule(out, 1) - - def test_push_permute_through_reshape(self): - a = Tensor.empty(16,16) - b = Tensor.empty(16,16) - c = (a+b).reshape(4,4,4,4).permute(2,3,0,1).contiguous() - check_schedule(c, 1) - - #@unittest.skip("failing in old lazy") - def test_push_permute_through_reshape_alt(self): - a = Tensor.empty(4,4,4,4) - b = Tensor.empty(4,4,4,4) - c = (a+b).reshape(16,16).permute(1,0).contiguous() - check_schedule(c, 1) - - def test_no_binop_rerun(self): - a = Tensor.empty(16) - b = Tensor.empty(16) - c = a+b - d = (a+b).reshape(16,1) - check_schedule(d, 0, [c]) - - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") - def test_multi_permute_should_collapse(self): - a = Tensor.empty(4,4,4,4) - b = Tensor.empty(16) - c = a.sum((0,1)).cast(dtypes.float16).permute(1,0).reshape(4,4,1).permute(1,0,2).reshape(16) + b - check_schedule(c, 1) - - def test_fancy_reshape_fusion(self): - a = Tensor.empty(10) - b = Tensor.empty(10) - c = a+b - d = a.reshape(10,1)+b.reshape(10,1) - out = c.sum() + d.sum() - check_schedule(out, 1) - - def test_children_dont_push(self): - a = Tensor.empty(10, 10, 1) - b = Tensor.empty(10, 10, 1) - d = (a+b).expand(10, 10, 10) - e = (a+b).permute(2,1,0) - f = d+e - check_schedule(f, 1) - - # 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) - c = Tensor.empty(10) - 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 - check_schedule(d, 2) - check_schedule(keep_me, 0, [d]) - - #@unittest.skip("failing in old lazy") - def test_permute_breaks_fusion(self): - a = Tensor.empty(10, 10, 10) - b = Tensor.empty(10, 10) - c = (a.sum(axis=2) + b).permute(1,0) - d = c.permute(1,0) - check_schedule(d, 1) - - def test_some_permute_fusion(self): - a = Tensor.empty(8192, 16) - b = Tensor.empty(1, 16) - d = (a.T + b.expand(8192, 16).T) - c = a + b.expand(8192, 16) - e = d.T - check_schedule(c, 1) - check_schedule(e, 1) - - def test_shrink_fuse(self): - a = Tensor.empty(8192, 16) - b = Tensor.empty(8192, 16) - c = a * b - d = Tensor.empty(1, 16) - e = c[0] * d - check_schedule(e, 1) - - def test_expand_fuse(self): - a = Tensor.empty(1, 16) - b = Tensor.empty(1, 16) - c = a * b - d = Tensor.empty(8192, 16) - e = c * d - check_schedule(e, 1) - - # this is the failing case in openpilot...it's very simple like this - def test_image_conv_fusion(self): - w1 = Tensor.empty(16, 16, 1, 1) - b1 = Tensor.empty(16) - w2 = Tensor.empty(16, 16, 1, 1) - b2 = Tensor.empty(16) - w3 = Tensor.empty(16, 16, 1, 1) - b3 = Tensor.empty(16) - - x = Tensor.empty(1, 16, 32, 32) - x = base = x.image_conv2d(w1, b1) - x = x.image_conv2d(w2, b2) + base - x = x.image_conv2d(w3, b3) - - # NOOP, 3 convs, contiguous - #check_schedule(x, 5) - check_schedule(x, 7) - - def test_image_conv_fusion_minimal(self): - b1 = Tensor.empty(16) - b2 = Tensor.empty(16) - def p(x): return x.permute(1,0).contiguous().reshape(32,16,1).expand(32,16,16).sum(axis=2).permute(1,0) - - x = Tensor.empty(16, 32) - x = base = p(x) + b1.reshape(16,1) - x = p(x) - x = x + b2.reshape(16,1) - x = x + base - del base - x = p(x) - check_schedule(x, 4) - - def test_image_conv_fusion_more_minimal(self): - b1 = Tensor.empty(16) - def p(x): return x.permute(1,0).contiguous().reshape(32,16,1).expand(32,16,16).sum(axis=2).permute(1,0) - - x = Tensor.empty(16, 32) - x = base = p(x) + b1.reshape(16,1) - x = p(x) - del base - check_schedule(x, 3) - - def test_resnet_block(self): - with Tensor.train(False): - in_planes, planes = 64, 64 - conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=1, padding=1, bias=False) - bn1 = nn.BatchNorm2d(planes) - conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, stride=1, bias=False) - bn2 = nn.BatchNorm2d(planes) - x = Tensor.empty(1, 64, 32, 32) - out = bn1(conv1(x)).relu() - out = bn2(conv2(out)) - out = (out + x).relu() - run_schedule(check_schedule(out, 2, [conv1.weight, conv2.weight])) - - def test_contiguous_while_contiguous(self): - x = Tensor.empty(1, 64, 32, 32) - out = x.contiguous() - check_schedule(out, 0, filter_sink=False) - - def test_contiguous_while_not_contiguous(self): - x = Tensor.empty(1, 64, 32, 32) - out = x.permute(0,2,3,1).contiguous() - check_schedule(out, 1, filter_sink=False) - - def test_fold_with_contiguous(self): - a = Tensor.randn(16, 16, 16).realize() - b = Tensor.randn(16, 16).realize() - c = (a.sum(2).contiguous() + b).contiguous() - check_schedule(c, 2) - @unittest.skip("no longer supported") def test_double_from(self): x = Tensor([1,2,3,4]) out = x.to('python') check_schedule(out, 0, filter_sink=False) - def _alu_from_tensor(self, t:Tensor): - s = [s for s in t.schedule() if s.ast.op is Ops.SINK] - self.assertEqual(len(s), 1) - return [u.op for u in s[0].ast.toposort() if u.op in GroupOp.ALU] - - def test_2_pow_is_exp2(self): - t = 2.0 ** Tensor([1.0, 2.0, 3.0]) - self.assertEqual(self._alu_from_tensor(t), [Ops.EXP2]) - - def test_pow_05_is_sqrt(self): - t = Tensor([1.0, 2.0, 3.0]) ** 0.5 - self.assertEqual(self._alu_from_tensor(t), [Ops.SQRT]) - - def test_pow_neg_05_is_rsqrt(self): - t = Tensor([1.0, 2.0, 3.0]) ** -0.5 - self.assertEqual(self._alu_from_tensor(t), [Ops.RECIPROCAL, Ops.SQRT]) - - def test_pow_2_has_1_mul(self): - t = Tensor([1.0, 2.0, 3.0]) ** Tensor(2.0) - self.assertEqual(self._alu_from_tensor(t), [Ops.MUL]) - - def test_pow_8_has_3_muls(self): - t = Tensor([1.0, 2.0, 3.0]) ** 8 - self.assertEqual(self._alu_from_tensor(t), [Ops.MUL, Ops.MUL, Ops.MUL]) - - def test_pow_const_tensor_to_zero(self): - x = Tensor([1,2,3,4]) - out = x ** Tensor(0.0) - # NOTE: this is UOp.const(0) + UOp.const(1) - check_schedule(out, 0) - - def test_zero_size(self): - x = Tensor.empty(2, 3, 0) - 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,)) @@ -734,18 +222,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(out, 1)) self.assertEqual(out.item(), 4.) - def test_reduce_permute_nofuse(self): - x = Tensor.empty(32, 32, 32) - y = Tensor.empty(32, 32) - out = x.sum(axis=2).T+y - check_schedule(out, 1) - - def test_two_elus_sum(self): - x = Tensor.empty(32, 32) - y = Tensor.empty(32, 32) - out = x.sum(1).relu().elu() + y.sum(1).relu().elu() - check_schedule(out, 1) - @unittest.skipUnless(SPLIT_REDUCEOP, "Testing split reducop requires SPLIT_REDUCEOP") def test_preserve_multistage_reduce(self): big_enough = getenv("REDUCEOP_SPLIT_THRESHOLD", 32768) @@ -754,18 +230,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(out, 4)) np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max()) - def test_multistage_reduce(self): - x = Tensor.empty(32, 32, 32) - out = x.sum(2).relu().sum(1) - check_schedule(out, 1) - - def test_multistage_reduce_fork(self): - x = Tensor.empty(32, 32, 32) - x = x.sum(2) - out2 = x + 1 - out = x.relu().sum(1) + out2[0] - check_schedule(out, 2) - @unittest.skip("these two Tensors are the same") def test_example_matmul(self): x = Tensor.eye(64, requires_grad=True) @@ -794,27 +258,6 @@ class TestSchedule(unittest.TestCase): # NOTE: the gradient flows twice np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64))) - def test_contiguous_add(self): - x = Tensor.empty(32) - y = Tensor.empty(32) - z = Tensor.empty(32) - out = (x+y).contiguous()+z - check_schedule(out, 2) - - def test_double_sum_ref(self): - x = Tensor.empty(32, 32, 32) - x = x.sum(2) - out = x + x[:, 4] - check_schedule(out, 2) - - def test_reduce_shrink(self): - x = Tensor.empty(32, 32) - y = Tensor.empty(16) - x = x.sum(1) - x = x[:16] - out = x + y - check_schedule(out, 1) - def test_multireduce_shrink(self): Tensor.manual_seed(0) a = Tensor.randn(32, 32).realize() @@ -828,14 +271,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(out, 1)) np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4) - # broken due to const folding and two contiguous are different kernels - # NOTE: passes after delete_lazy - def test_const_no_recompute(self): - x = Tensor(2) + Tensor(2) - y = Tensor(2) + Tensor(2) - out = x.contiguous() + y.contiguous() - check_schedule(out, 2, filter_sink=False) - def test_reduce_same_size(self): Tensor.manual_seed(0) a = Tensor.randn(4, 4).realize() @@ -940,27 +375,6 @@ 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) - 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], 2) # TODO: 1? - - def test_reduce_multiple_paths_midshrink(self): - a = Tensor.empty(4, 4) - r = a.sum(axis=1) - out0 = r.exp2() - out1 = out0[0] + out0 - check_schedule([r, out0, out1], 3) - - def test_reduce_shrink_output(self): - a = Tensor.empty(4, 4) - r = a.sum(keepdim=True) - out0 = r.exp2() - out1 = out0[0] + Tensor.empty(1, ) - check_schedule([r, out0, out1], 3) - def test_std_multireduce_fusion(self): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() @@ -1116,30 +530,6 @@ class TestSchedule(unittest.TestCase): expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True) np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4) - @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") - 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(sched[0].bufs[0].dtype, dtypes.float) - - # input float, softmax in float - Tensor.manual_seed(0) - x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.float).realize() - out = x.softmax(dtype=dtypes.float) - sched = out.schedule() - self.assertEqual(len(sched), 3) - self.assertEqual(sched[0].bufs[0].dtype, dtypes.float) - - def test_softmax_backward(self): - Tensor.manual_seed(0) - x = Tensor.randn(4, 12, 64, 64, requires_grad=True).realize() - x.softmax().sum().backward() - run_schedule(check_schedule(x.grad, 4)) - def test_layernorm_onelayer_fusion(self): Tensor.manual_seed(0) layer = nn.LayerNorm([10, 10]) @@ -1153,112 +543,6 @@ class TestSchedule(unittest.TestCase): expected = y / np.sqrt((y*y).mean(layer.axis, keepdims=True) + layer.eps) np.testing.assert_allclose(out.numpy(), expected * layer.weight.numpy() + layer.bias.numpy(), atol=1e-4, rtol=1e-4) - def test_scaled_dot_product_attention_fusion(self): - x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4)) - out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m) - check_schedule(out, 4) - - def test_scaled_dot_product_attention_causal_fusion(self): - x, y, z = (Tensor.empty(32, 8, 16, 16) for _ in range(3)) - out = Tensor.scaled_dot_product_attention(x, y, z, is_causal=True) - check_schedule(out, 4) - - def test_adam_step_fusion(self): - with Tensor.train(): - x = Tensor.empty(4, 64, 32) - layer = nn.Linear(32, 32*4) - _realize_weights(layer) - opt = nn.optim.Adam(nn.state.get_parameters(layer), lr=1e-4) - layer(x).relu().sum().backward() - check_schedule(opt.schedule_step(), 19) - - def test_adam_conv_fuse(self): - with Tensor.train(): - img = Tensor.empty(2,3,4,4) - c1 = nn.Conv2d(3,32,3) - _realize_weights(c1) - opt = nn.optim.Adam(nn.state.get_parameters(c1), lr=1e-4) - opt.zero_grad() - c1(img).relu().sum().backward() - check_schedule(opt.schedule_step(), 19) - - def test_adam_2convs_fuse(self): - with Tensor.train(): - img = Tensor.empty(2,3,4,4) - c1 = nn.Conv2d(3,16,3,bias=False) - c2 = nn.Conv2d(16,32,2,bias=False) - _realize_weights([c1, c2]) - opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4) - opt.zero_grad() - c2(c1(img).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 21) - - def test_sgd_conv_fuse(self): - with Tensor.train(): - img = Tensor.empty(2,3,4,4) - c1 = nn.Conv2d(3,32,3) - _realize_weights(c1) - opt = nn.optim.SGD(nn.state.get_parameters(c1)) - opt.zero_grad() - c1(img).relu().sum().backward() - check_schedule(opt.schedule_step(), 5) # TODO: 3? - - def test_sgd_2convs_fuse(self): - with Tensor.train(): - img = Tensor.empty(2,3,4,4) - c1 = nn.Conv2d(3,16,3,bias=False) - c2 = nn.Conv2d(16,32,2,bias=False) - _realize_weights([c1, c2]) - opt = nn.optim.SGD(nn.state.get_parameters([c1, c2])) - opt.zero_grad() - c2(c1(img).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 7) - - @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong") - def test_fold_2convs_sgd_nesterov_momentum_wd(self): - with Tensor.train(): - img = Tensor.empty(2,3,4,4) - c1 = nn.Conv2d(3,16,3,bias=False) - c2 = nn.Conv2d(16,32,2,bias=False) - _realize_weights([c1, c2]) - opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1) - opt.zero_grad() - c2(c1(img).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 13) - - def test_sgd_4convs_fuse(self): - with Tensor.train(): - img = Tensor.empty(2,3,16,16) - c1 = nn.Conv2d(3,4,3,bias=False) - c2 = nn.Conv2d(4,8,3,bias=False) - c3 = nn.Conv2d(8,16,3,bias=False) - c4 = nn.Conv2d(16,32,3,bias=False) - _realize_weights([c1, c2, c3, c4]) - opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) - opt.zero_grad() - c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 15) - - def test_sgd_4convs_fuse_conv_bw(self): - with Tensor.train(): - img = Tensor.empty(2,3,16,16) - c1 = nn.Conv2d(3,4,3,bias=False) - c2 = nn.Conv2d(4,8,3,bias=False) - c3 = nn.Conv2d(8,16,3,bias=False) - c4 = nn.Conv2d(16,32,3,bias=False) - _realize_weights([c1, c2, c3, c4]) - opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4])) - opt.zero_grad() - c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward() - check_schedule(opt.schedule_step(), 15) - - def test_reduce_simple_chase(self): - a = Tensor.empty(4, 4, 4) - r = a.sum(0) + 6 - b = r.sum(0) * 4 - c = r.sum(1) * 2 - check_schedule([b, c], 3) - def test_multireduce_simple_chase(self): Tensor.manual_seed(0) a = Tensor.randn(4, 4, 4).realize() @@ -1273,14 +557,6 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(b.numpy(), np_r.sum(0) + 8, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(c.numpy(), np_r.sum(1) + 12, atol=1e-4, rtol=1e-4) - def test_push_permute_chase(self): - a = Tensor.empty(4, 4, 4) - b = Tensor.empty(4, 4) - r = a.sum(2) + b - d = r.T * 4 - e = r * d - check_schedule([d, e], 3) - def test_multireduce_push_permute_chase(self): Tensor.manual_seed(0) a = Tensor.randn(4, 4, 4).realize() @@ -1293,14 +569,6 @@ class TestSchedule(unittest.TestCase): np.testing.assert_allclose(d.numpy(), (a.numpy().sum(2) + b.numpy()).T * 4, atol=1e-4, rtol=1e-4) np.testing.assert_allclose(e.numpy(), (a.numpy().sum(2) + b.numpy()) * (d.numpy() + a.numpy()).sum(2), atol=1e-4, rtol=1e-4) - def test_push_shrink_chase(self): - a = Tensor.empty(16, 16) - b = Tensor.empty(4) - c = Tensor.empty(16, ) - r = a.sum(1) + c - d = r[:4] * b - check_schedule(d, 1) - def test_multireduce_push_shrink_chase(self): Tensor.manual_seed(0) a = Tensor.randn(16, 16).realize() @@ -1313,11 +581,6 @@ class TestSchedule(unittest.TestCase): run_schedule(schedule) np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4) - def test_midreduce_nochase(self): - a = Tensor.empty(16, 16) - b = (a.sum(0) + a.max(1)) + 2 - check_schedule(b, 1) - def test_multireduce_midreduce_nochase(self): Tensor.manual_seed(0) a = Tensor.randn(16, 16).realize() @@ -1452,18 +715,6 @@ class TestSchedule(unittest.TestCase): run_schedule(sched) np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3) - def test_fuse_arange_avg_pool2d_ceil_mode(self): - x = Tensor.avg_pool2d(Tensor.empty(1,1,6,6), kernel_size=(3,3), padding=1, stride=3, ceil_mode=True) - sched = check_schedule(x, 1) - self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 1) - - def test_fuse_arange_pad_circular_mode_bw(self): - x = Tensor.empty(1,1,5,5,5) - out = x.pad((1,2,3,5,1,2), mode="circular") - g = out.sum().gradient(x)[0] - sched = check_schedule(g, 1) - self.assertEqual(len([x for x in sched[0].ast.backward_slice_with_self if x.op is Ops.REDUCE]), 0) - # TODO like openpilot with imagef @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") def test_base_change_expand_expand(self): @@ -1500,12 +751,6 @@ class TestSchedule(unittest.TestCase): p = np.tile(p, 2) np.testing.assert_allclose(tiny_ret, p) - def test_bitcast_fuses(self): - x = Tensor.empty(1, dtype=dtypes.float32) - a = x.exp2().bitcast(dtypes.int32) - b = x.bitcast(dtypes.int32) - check_schedule(a+b, 1) # this should fuse when it makes sense - @unittest.skip("disabling subbuffer manually isn't supported anymore") def test_bitcast_disable_subbufer(self): x = cast(UOp, Tensor.empty(1, dtype=dtypes.float32).realize().uop) @@ -1514,34 +759,15 @@ class TestSchedule(unittest.TestCase): b = a.alu(Ops.ADD, b) check_schedule(b, 1) - def test_reduceop_reshape_dont_push(self): - Tensor.manual_seed(0) - x = Tensor.randn(10, 20).realize() - out = x.argmax(1) - run_schedule(check_schedule(out, 2)) - def test_conv2d(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4) def test_conv2d_fused(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4) - def test_resnet_conv2d(self): - x = Tensor.empty(1, 8, 32, 32) - w1 = Tensor.empty(8, 8, 3, 3) - w2 = Tensor.empty(8, 8, 1, 1) - out = x.conv2d(w1).conv2d(w2) - check_schedule(out, 2) - @unittest.skipUnless(is_dtype_supported(dtypes.half) and is_dtype_supported(dtypes.ulong), "need half and ulong") def test_conv2d_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half) @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half") @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail") def test_conv2d_fused_half(self): _test_conv2d(5 if SPLIT_REDUCEOP else 4, dtype=dtypes.half) - def test_schedule_mem_used(self): - base = GlobalCounters.mem_used - Tensor.ones(256).contiguous().realize() - Tensor.ones(5, 5).contiguous().schedule() - self.assertEqual(GlobalCounters.mem_used-base, 0) - @unittest.skip("TODO: this is consistently creating non reproducible failures") def test_schedule_mem_used_with_inputs(self): base = GlobalCounters.mem_used @@ -1549,14 +775,6 @@ class TestSchedule(unittest.TestCase): (x+Tensor.ones(256).contiguous()).schedule() self.assertEqual(GlobalCounters.mem_used-base, 1024) - def test_const_schedule(self): - constv = Tensor.empty(2, 2).uop.const_like(10) - check_schedule(constv, 0) - - def test_const_schedule_contig(self): - constv = Tensor.empty(2, 2).uop.const_like(10).contiguous() - check_schedule(constv, 1) - @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL") def test_image_matmul(self): with Context(IMAGE=2): @@ -1657,11 +875,6 @@ class TestSchedule(unittest.TestCase): run_schedule(check_schedule(xt, 1)) 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]] - run_schedule(check_schedule(xt, 1)) - def test_push_through_reshape(self): Tensor.manual_seed(0) x = Tensor.randn(10, 20).realize() @@ -1722,14 +935,6 @@ class TestSchedule(unittest.TestCase): 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): - Tensor.manual_seed(0) - with Context(TRACK_MATCH_STATS=0): - x = Tensor.randn(11).realize() - a = Tensor.arange(22) - out = (x + a[:11]).sum() - check_schedule(out, 1) - def test_arange_index_contiguous(self): Tensor.manual_seed(0) x = Tensor.randn(5, 2).realize() @@ -1908,20 +1113,6 @@ class TestSwizzle(unittest.TestCase): run_schedule(check_schedule(add, 1)) self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0)) - def test_softmax_one_kernel(self): - Tensor.manual_seed(0) - with Context(DEBUG=0, TRACK_MATCH_STATS=0): - a = Tensor.randn(32, 32).realize() - t = a.softmax() - check_schedule(t, 3) # TODO: 1? - - def test_argmax_one_kernel(self): - Tensor.manual_seed(0) - with Context(DEBUG=0, TRACK_MATCH_STATS=0): - a = Tensor.randn(10, 20).realize() - t = a.argmax(0) - check_schedule(t, 2) # TODO: 1? - def test_swizzle_reduceop(self): Tensor.manual_seed(0) x = Tensor.randn(4,4).realize() @@ -1992,11 +1183,6 @@ class TestView(unittest.TestCase): run_schedule(sched) np.testing.assert_equal(b.numpy(), 0) - def test_zero_size_alt(self): - a = Tensor.empty(135, 0, 9) - b = a.pad(((0, 0), (0, 0), (18, 0))) - check_schedule(b, 0) - def test_partial_mask(self): # partial masked out does not degrade into CONST a = Tensor.rand(10, 10).realize() @@ -2153,137 +1339,6 @@ class TestCopyFolding(unittest.TestCase): self.assertListEqual(b.tolist(), [[0, 2], [1, 3]]) class TestUOpBecome(unittest.TestCase): - # the simplest case, if we create a new BUFFER for this tensor UOp - def test_new_buffer(self): - a = Tensor.empty(4, 4) - b = Tensor.empty(4, 4) - add = a+b - check_schedule(add, 1) - # NOTE: realized base is always a flat buffer - assert UPat(Ops.BUFFER).match(add.uop.base, {}) - # the Tensor UOp can optionally stack a VIEW on top of the BUFFER, in this case to preserve the (4, 4) shape of the tensor - assert add.uop is not add.uop.base - self.assertEqual(add.uop.size, 16) - self.assertEqual(add.uop.shape, (4, 4)) - - def test_new_buffer_view(self): - a = Tensor.empty(4, 4) - b = Tensor.empty(4, 4) - add = (a+b).reshape(8, 2) - check_schedule(add, 1) - assert UPat(Ops.BUFFER).match(add.uop.base, {}) - # the shape is preserverd in the becomes_map. - self.assertEqual(add.uop.shape, (8, 2)) - assert add.uop is not add.uop.base - - def test_new_flat_buffer(self): - a = Tensor.empty(4,) - b = Tensor.empty(4,) - add = a+b - check_schedule(add, 1) - # BUFFER already has a shape (4,), this tensor just becomes a contiguous BUFFER - assert UPat(Ops.BUFFER).match(add.uop.base, {}) - - # sometimes we prefer to perform an op before movement ops, in this case we should stack the mops on top of the new buffer - - def test_reorder_expand(self): - a = Tensor.empty(4, 1) - b = a.expand(4, 4).reciprocal() - check_schedule(b, 1) - self.assertEqual(b.uop.base.buffer.size, 4) - self.assertEqual(b.uop.shape, (4, 4)) - - def test_reorder_expand_alt(self): - x = Tensor.empty(4, 1) - y = Tensor.empty(4, 1) - img = Tensor.empty(4, 4) - z = (img*x) / y - check_schedule(z, 1) - - # TODO: rangeify doesn't yet cleanup this kind of re-indexing - @unittest.expectedFailure - def test_become_existing_buffer(self): - a = Tensor.empty(4, 4) - b = a*1 - assert UPat(Ops.MUL).match(b.uop, {}) # before scheduling it's a mul - check_schedule(b, 0) - self.assertIs(a.uop.base.buffer, b.uop.base.buffer) - - def test_become_buf_with_mops(self): - a = Tensor.empty(2, 4, 2) - noop = a.shrink(((1, 2), (0, 4), (0, 2))).reshape(4, 2)*1+0 - # before realizing, this tensor is base - assert noop.uop is noop.uop.base - noop.realize() - # it becomes a realized view after realize - assert noop.uop is not noop.uop.base - assert noop.uop.base.op is Ops.BUFFER - late_add = noop+2 - late_add.realize() - - def test_become_const_in_base(self): - a = Tensor.empty(4) - b = a*0 - assert UPat(Ops.MUL).match(b.uop, {}) # before scheduling it's a mul - check_schedule(b, 0) - assert UPat(Ops.CONST, arg=0).match(b.uop.base, {}) # scheduling replaces the tensor uop with a VIEW(BUFFER) - - def test_become_const_from_const(self): - const_add = Tensor(1)+Tensor(2) - assert UPat(Ops.ADD).match(const_add.uop, {}) - check_schedule(const_add, 0) - assert UPat(Ops.CONST, arg=3).match(const_add.uop.base, {}) - - # tensors can become another realized tensor source - @unittest.expectedFailure - def test_become_existing_buf_simple(self): - a = Tensor.empty(4, 4) - b = a+0 - check_schedule(b, 0) - assert b.uop.base.op is Ops.BUFFER - self.assertIs(a.uop, b.uop) - - # they can also chain other movement ops on top of the tensor source - @unittest.expectedFailure - 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) - - @unittest.expectedFailure - def test_become_existing_buf_view_alt(self): - a = Tensor.empty(4, 4) - b = a.permute((1, 0)).reshape((8, 2))+0 - check_schedule(b, 0) - 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 - @unittest.expectedFailure - def test_become_existing_buf_complex(self): - a = Tensor.empty(4, 4) - b = (a.permute((1, 0))+0).reshape((8, 2))+0 - check_schedule(b, 0) - self.assertEqual(b.uop.st, a.uop.permute((1, 0)).reshape((8, 2)).st) - assert b.uop.base.op is Ops.BUFFER - - @unittest.expectedFailure - 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 - c = (a.reshape(1, 1, 4, 4)+0).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0 - check_schedule([b, c], 0) - assert all_same([x.uop.base.realized for x in [a,b,c]]) - - def test_setitem_becomes_subbuffer(self): - a = Tensor.full((4,), 2.).contiguous().realize() - b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0)) - b.realize() - assert a.uop.is_realized - assert a.uop.buffer._base is None - assert b.uop.op_in_backward_slice_with_self(Ops.SHRINK) - assert b.uop.base is a.uop.base - def test_setitem_offset(self): a = Tensor.full((16,), 0.).contiguous().realize() b = Tensor.full((16,), 1.).contiguous().realize() @@ -2291,5 +1346,22 @@ class TestUOpBecome(unittest.TestCase): b.shrink(((0,4),)).assign(a_view).realize() self.assertListEqual(b.tolist(), [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) +class TestFusionOp(unittest.TestCase): + def test_contiguous_add(self): + def test(contig=False): + bt = Tensor(np.arange(16), dtype=dtypes.float32).reshape(4,4) + x = bt.permute(1,0) + if contig: x = x.contiguous() + return (x.permute(1,0) + bt).data() + assert test() == test(True) + + def test_expand_fuse(self): + bt = Tensor(np.ones((10, 1)), dtype=dtypes.float32) + out = (bt*2).expand(10,10).sum(1) + sched = out.schedule() + run_schedule(sched) + outd = out.tolist() + assert all(x == 20.0 for x in outd) + if __name__ == '__main__': unittest.main(verbosity=2)