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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
