[GitHub] [tvm] junrushao commented on a diff in pull request #13093: [MetaSchedule] Fix AddReadReuse for multiple consumer block case

[GitBox <gi...@apache.org>]

junrushao commented on code in PR #13093:
URL: https://github.com/apache/tvm/pull/13093#discussion_r1072582257


##########
tests/python/unittest/test_meta_schedule_schedule_rule_mlt.py:
##########
@@ -623,6 +623,29 @@ def cpu_conv2d_nhwc(
     )
 
 
+def test_cache_read_specify_consumer():
+    A, B, C = te_workload.matmul(512, 512, 512)
+    mod = te.create_prim_func([A, B, C + A])
+
+    space = generate_design_space(
+        kind="cuda",
+        mod=mod,
+        target=Target("nvidia/geforce-rtx-3080"),
+        types=ms.schedule_rule.MultiLevelTiling,
+    )
+
+    residual_block = """
+        for i0, i1 in T.grid(512, 512):
+            with T.block("T_add"):
+                ax0, ax1 = T.axis.remap("SS", [i0, i1])
+                T.reads(C[ax0, ax1], A[ax0, ax1])
+                T.writes(T_add[ax0, ax1])
+                T_add[ax0, ax1] = C[ax0, ax1] + A[ax0, ax1]
+    """
+
+    assert residual_block in space[0].mod.script()

Review Comment:
   Hey a late nitpick: looks like there could be some way around to make it not using string equality. I just tweaked it a bit to using structural equality :-)



-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org