[Fix][DLight] Decline matmul schedules with unsupported consumer chains

python/tvm/s_tir/dlight/gpu/matmul.py

@@ -95,7 +95,8 @@ def auto_inline_consumers(
 def auto_inline_consumer_chain(
     sch: s_tir.Schedule,
     block: s_tir.schedule.SBlockRV,
-):
+) -> bool:
+    """Try to inline the consumer chain into the cache-write stage."""
     auto_inline_consumers(sch, block)
     remaining_consumers = sch.get_consumers(block)
 
@@ -106,13 +107,15 @@ def auto_inline_consumer_chain(
             for p in sch.get_producers(c):
                 if sch.get(p) != sch.get(block):
                     auto_inline_producers(sch, p)
-                    sch.compute_inline(p)
+                    try:
+                        sch.compute_inline(p)
+                    except s_tir.ScheduleError:
+                        return False
 
-        # Try inlining into the cache-write stage again, this time it should succeed.
+        # Retry after inlining the other producers.  The chain may still be unsupported.
         auto_inline_consumers(sch, block)
 
-    msg = "There are some consumers of the cache-write stage that are not properly inlined."
-    assert len(sch.get_consumers(block)) == 0, msg
+    return len(sch.get_consumers(block)) == 0
 
 
 class IterKind(Enum):
@@ -458,7 +461,8 @@ def tensorize_block(block: s_tir.schedule.SBlockRV, intrin: str):
         *_, i, j, k = sch.get_loops(main_block)
         sch.tensorize(i, intrin_group["compute"])
 
-        auto_inline_consumer_chain(sch, C_s2g)
+        if not auto_inline_consumer_chain(sch, C_s2g):
+            return None
         fused = sch.fuse(*sch.get_loops(C_s2g)[-2:])
         _, _tz, _ty, _tx, vec = sch.split(fused, [None, tz_len, ty_len, warp_size, vector_size])
         sch.bind(_tz, "threadIdx.z")
@@ -658,7 +662,8 @@ def _cooperative_fetch(index, vec_len):
         else:
             auto_inline_producers(sch, main_block)
 
-        auto_inline_consumer_chain(sch, l2g)
+        if not auto_inline_consumer_chain(sch, l2g):
+            return None
 
         sch.decompose_reduction(main_block, ko)
         return sch

tests/python/s_tir/dlight/test_gpu_matmul_consumer_chain.py

@@ -0,0 +1,151 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Matmul must decline consumer chains it cannot inline (issue #20615)."""
+
+import pytest
+
+import tvm_tirx
+import tvm_tirx.testing
+from tvm import s_tir
+from tvm.s_tir import dlight as dl
+from tvm.s_tir import script as Ts
+from tvm.s_tir.dlight.gpu.matmul import auto_inline_consumer_chain
+from tvm_tirx.ir import assert_structural_equal
+from tvm_tirx.script import tirx as T
+from tvm_tirx.target import Target
+
+
+@pytest.fixture(params=[False, True], ids=["bias_relu", "bias_relu_pad"])
+def conv_consumer_chain(request):
+    pad = int(request.param)
+
+    @Ts.function
+    def before(
+        data: T.Tensor((1, 8, 9, 17), "float32"),
+        weight: T.Tensor((4, 8, 2, 2), "float32"),
+        bias: T.Tensor((1, 4, 1, 1), "float32"),
+        output: T.Tensor((1, 4, 8 + pad, 16 + pad), "float32"),
+    ):
+        T.func_attr({"tirx.noalias": True})
+        padded_input = Ts.sblock_alloc_tensor((1, 8, 9, 17))
+        conv = Ts.sblock_alloc_tensor((1, 4, 8, 16))
+        add = Ts.sblock_alloc_tensor((1, 4, 8, 16))
+        relu = Ts.sblock_alloc_tensor((1, 4, 8, 16))
+
+        for n, c, h, w in T.grid(1, 8, 9, 17):
+            with Ts.sblock("copy_input"):
+                vn, vc, vh, vw = Ts.axis.remap("SSSS", [n, c, h, w])
+                padded_input[vn, vc, vh, vw] = data[vn, vc, vh, vw]
+
+        for n, f, y, x, rc, ry, rx in T.grid(1, 4, 8, 16, 8, 2, 2):
+            with Ts.sblock("conv2d_nchw"):
+                vn, vf, vy, vx, vrc, vry, vrx = Ts.axis.remap("SSSSRRR", [n, f, y, x, rc, ry, rx])
+                with Ts.init():
+                    conv[vn, vf, vy, vx] = T.float32(0)
+                conv[vn, vf, vy, vx] += (
+                    padded_input[vn, vrc, vy + vry, vx + vrx] * weight[vf, vrc, vry, vrx]
+                )
+
+        for n, f, y, x in T.grid(1, 4, 8, 16):
+            with Ts.sblock("bias_add"):
+                vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
+                add[vn, vf, vy, vx] = conv[vn, vf, vy, vx] + bias[vn, vf, 0, 0]
+
+        for n, f, y, x in T.grid(1, 4, 8, 16):
+            with Ts.sblock("relu"):
+                vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
+                relu[vn, vf, vy, vx] = T.max(add[vn, vf, vy, vx], T.float32(0))
+
+        for n, f, y, x in T.grid(1, 4, 8 + pad, 16 + pad):
+            with Ts.sblock("pad_output"):
+                vn, vf, vy, vx = Ts.axis.remap("SSSS", [n, f, y, x])
+                output[vn, vf, vy, vx] = T.if_then_else(
+                    vy < 8 and vx < 16, relu[vn, vf, vy, vx], T.float32(0)
+                )
+
+    return before, bool(pad)
+
+
+def test_matmul_consumer_chain(conv_consumer_chain):
+    before, padded = conv_consumer_chain
+    original = tvm_tirx.IRModule({"main": before})
+    snapshot = tvm_tirx.ir.load_json(tvm_tirx.ir.save_json(original))
+    target = Target("cuda")
+    schedule = dl.gpu.Matmul().apply(before, target, False)
+    assert (schedule is None) == padded
+    assert_structural_equal(original, snapshot)
+
+    with target:
+        matmul_only = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(original)
+        if padded:
+            assert_structural_equal(matmul_only, snapshot)
+        else:
+            assert matmul_only["main"].attrs["tirx.is_scheduled"] == 1
+
+
+def test_matmul_consumer_chain_following_rules(conv_consumer_chain):
+    before, padded = conv_consumer_chain
+    following_rules = [
+        dl.gpu.GEMV(),
+        dl.gpu.Reduction(),
+        dl.gpu.GeneralReduction(),
+        dl.gpu.Fallback(),
+    ]
+    with Target("cuda"):
+        actual = dl.ApplyDefaultSchedule(dl.gpu.Matmul(), *following_rules)(
+            tvm_tirx.IRModule({"main": before})
+        )
+        if padded:
+            expected = dl.ApplyDefaultSchedule(*following_rules)(
+                tvm_tirx.IRModule({"main": before})
+            )
+            assert_structural_equal(actual, expected)
+    assert actual["main"].attrs["tirx.is_scheduled"] == 1
+
+
+def test_consumer_chain_non_inlineable_producer():
+    @Ts.function
+    def before(
+        a: T.Tensor((16, 16), "float32"),
+        b: T.Tensor((16, 16), "float32"),
+        output: T.Tensor((16,), "float32"),
+    ):
+        lhs = Ts.sblock_alloc_tensor((16,))
+        rhs = Ts.sblock_alloc_tensor((16,))
+        for i, k in T.grid(16, 16):
+            with Ts.sblock("lhs"):
+                vi, vk = Ts.axis.remap("SR", [i, k])
+                with Ts.init():
+                    lhs[vi] = T.float32(0)
+                lhs[vi] += a[vi, vk]
+        for i, k in T.grid(16, 16):
+            with Ts.sblock("rhs"):
+                vi, vk = Ts.axis.remap("SR", [i, k])
+                with Ts.init():
+                    rhs[vi] = T.float32(0)
+                rhs[vi] += b[vi, vk]
+        for i in range(16):
+            with Ts.sblock("add"):
+                vi = Ts.axis.spatial(16, i)
+                output[vi] = lhs[vi] + rhs[vi]
+
+    sch = s_tir.Schedule(before)
+    assert not auto_inline_consumer_chain(sch, sch.get_sblock("lhs"))
+
+
+if __name__ == "__main__":
+    tvm_tirx.testing.main()