From 7ee377096172efb2ac0079d02d833eb5a79aff4e Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 7 Aug 2025 10:32:34 -0700 Subject: [PATCH] FUSE_ARANGE=1 (#11427) * FUSE_ARANGE=1 * fix test --------- Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com> --- .github/workflows/test.yml | 2 +- test/test_const_folding.py | 7 +++---- test/test_nn.py | 4 ++-- test/test_schedule.py | 8 ++++---- tinygrad/helpers.py | 2 +- 5 files changed, 11 insertions(+), 12 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 065f9dbab3..5f0f8ea93d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -624,7 +624,7 @@ jobs: - name: Test LLVM=1 DEVECTORIZE=0 for model run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py - name: Test CPU=1 DEVECTORIZE=0 - run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure" + run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure" testwebgpu: name: Linux (WebGPU) diff --git a/test/test_const_folding.py b/test/test_const_folding.py index 187a4b5142..60c36b0e56 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -139,10 +139,9 @@ class TestBitcastConstFolding(unittest.TestCase): class TestIndexingConstFolding(unittest.TestCase): def test_scalar_index(self): t = Tensor.arange(16).float().reshape(1,1,4,4).realize() - # TODO: fold these - _check_ast_count(2, t[:,:,Tensor(1),:]) - _check_ast_count(2, t[:,:,Tensor(1)+2,:]) - _check_ast_count(2, t[:,:,Tensor(1),Tensor(0)]) + _check_ast_count(1, t[:,:,Tensor(1),:]) + _check_ast_count(1, t[:,:,Tensor(1)+2,:]) + _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)]) @unittest.expectedFailure def test_const_tensor_index(self): diff --git a/test/test_nn.py b/test/test_nn.py index 42e7ea3632..9f5e70d509 100755 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -401,7 +401,7 @@ class TestNN(unittest.TestCase): torch_z = torch_layer(torch_x) np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8) - def test_embedding_one_kernel(self, ops=41410, kcount=3): + def test_embedding_one_kernel(self, ops=612000, kcount=2): GlobalCounters.reset() layer = Embedding(20, 30) layer.weight = Tensor.zeros_like(layer.weight).contiguous() @@ -409,7 +409,7 @@ class TestNN(unittest.TestCase): [12, 19, 8, 1]]) result = layer(a) schedule = result.schedule() - self.assertEqual(kcount, len([item for item in schedule if item.ast.op is Ops.SINK]), "first run realizes weight and embedding") + self.assertEqual(len([item for item in schedule if item.ast.op is Ops.SINK]), kcount, "first run realizes weight and embedding") run_schedule(schedule) b = Tensor([[1, 2, 3], diff --git a/test/test_schedule.py b/test/test_schedule.py index 39f70161ef..87530de80e 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -70,7 +70,7 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float, **kwargs): def schedule_graph_rewrite(big_sink:UOp): return get_kernelize_map(big_sink)[big_sink] class TestSchedule(unittest.TestCase): - def test_arange_avgpool2d(self, kcount=2): + def test_arange_avgpool2d(self, kcount=1): x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32) t = x.avg_pool2d(padding=1) sched = t.schedule() @@ -1028,14 +1028,14 @@ class TestSchedule(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(4, 32).realize() out = x.argmin(-1) - run_schedule(check_schedule(out, 3)) + 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, 3)) + 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): @@ -1613,7 +1613,7 @@ class TestSchedule(unittest.TestCase): Tensor.manual_seed(0) x = Tensor.randn(10, 20).realize() out = x.argmax(1) - run_schedule(check_schedule(out, 3)) # TODO: push a reduceop through a reshape + 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) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 90914c7d22..1bce0c41a9 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -131,7 +131,7 @@ JIT_BATCH_SIZE = ContextVar("JIT_BATCH_SIZE", 32) WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1) USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0) TRANSCENDENTAL, TC_SEARCH_OVER_SHAPE, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("TC_SEARCH_OVER_SHAPE", 1), ContextVar("NOLOCALS", 0) -FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 0), ContextVar("FUSE_CONV_BW", 0) +FUSE_ARANGE, FUSE_CONV_BW = ContextVar("FUSE_ARANGE", 1), ContextVar("FUSE_CONV_BW", 0) SPLIT_REDUCEOP, NO_MEMORY_PLANNER, RING = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("RING", 1) PICKLE_BUFFERS, PROFILE, LRU = ContextVar("PICKLE_BUFFERS", 1), ContextVar("PROFILE", getenv("VIZ")), ContextVar("LRU", 1) CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0), ContextVar("DEVECTORIZE", 1)