copy is already contiguous (#8886)

This commit is contained in:
chenyu
2025-02-04 17:53:33 -05:00
committed by GitHub
parent 4c28235bd1
commit 48349efdc1
4 changed files with 4 additions and 8 deletions

View File

@@ -219,9 +219,7 @@ class TestMultiConstFolding(unittest.TestCase):
t = Tensor.arange(16).float().realize().to(ds)
# non const folding case creates one ast on each shard
# NOTE: there's extra contiguous kernels here since it's realizing both the CONTIGUOUS and its parent COPY
# why does multi call contiguous on a COPY?
_check_ast_count(7, t + 1)
_check_ast_count(4, t + 1)
_check_ast_count(4, 1 + t)
_check_ast_count(4, t * 2)
_check_ast_count(4, 2 * t)

View File

@@ -774,10 +774,9 @@ class TestMultiTensor(unittest.TestCase):
zeros = Tensor.zeros(3).realize()
b = a.to(devices_2)*zeros.to(devices_2)
sched = b.schedule()
self.assertEqual(len(sched), 8)
self.assertEqual(len(sched), 6)
# notably, only two copies (for the arange) - vs 4 copies if we didn't fold the const copy
self.assertEqual(len([x for x in sched if any(u.op is Ops.COPY for u in x.ast.toposort)]), 2)
# all these kernels are just because multi calls contiguous on every single shard
run_schedule(sched)
self.assertListEqual(b.tolist(), [0, 0, 0])