vinx13 commented on code in PR #14161:
URL: https://github.com/apache/tvm/pull/14161#discussion_r1123570753
##########
tests/python/unittest/test_tir_schedule_cache_read_write.py:
##########
@@ -1336,5 +1388,38 @@ def test_cache_write_allocate_const():
verify_trace_roundtrip(sch=sch, mod=cache_write_allocate_const)
+def test_reindex_cache_read():
+ sch = tir.Schedule(elementwise, debug_mask="all")
+ sch.reindex_cache_read("C", 0, "shared", lambda i, j: (j, i // 2, i % 2))
+ tvm.ir.assert_structural_equal(elementwise_reindex_cache_read,
sch.mod["main"])
+ verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+ sch = tir.Schedule(elementwise, debug_mask="all")
+ with pytest.raises(tvm.tir.ScheduleError):
+ sch.reindex_cache_read(
+ "C",
+ 0,
+ "shared",
+ lambda i, j: j * 2,
+ )
+
+
+def test_reindex_cache_write():
+ sch = tir.Schedule(elementwise, debug_mask="all")
+ sch.reindex_cache_write("B", 0, "shared", lambda i, j: (j, i))
+ print(sch.mod["main"].show())
+ tvm.ir.assert_structural_equal(elementwise_reindex_cache_write,
sch.mod["main"])
+ verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+ sch = tir.Schedule(elementwise, debug_mask="all")
Review Comment:
ditto
##########
tests/python/unittest/test_tir_schedule_cache_read_write.py:
##########
@@ -1336,5 +1388,38 @@ def test_cache_write_allocate_const():
verify_trace_roundtrip(sch=sch, mod=cache_write_allocate_const)
+def test_reindex_cache_read():
+ sch = tir.Schedule(elementwise, debug_mask="all")
+ sch.reindex_cache_read("C", 0, "shared", lambda i, j: (j, i // 2, i % 2))
+ tvm.ir.assert_structural_equal(elementwise_reindex_cache_read,
sch.mod["main"])
+ verify_trace_roundtrip(sch=sch, mod=elementwise)
+
+ sch = tir.Schedule(elementwise, debug_mask="all")
Review Comment:
nit: separate to another test case
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
