geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

FUSE_ARANGE=1 (#11427)

Browse Source 
* FUSE_ARANGE=1

* fix test

---------

Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
chenyu
2025-08-07 13:32:34 -04:00
committed by GitHub
co-authored by George Hotz
parent 4dfcfb1ae5
commit 7ee3770961
5 changed files with 11 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test.yml
+1 -1
View File
@@ -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)
test/test_const_folding.py
+3 -4
View File
@@ -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):
test/test_nn.py
+2 -2
View File
@@ -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],
test/test_schedule.py
+4 -4
View File
@@ -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)
tinygrad/helpers.py
+1 -1
View File
@@ -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)