[Fix][Relax] Skip constant LHS in CombineParallelMatmul

Author: Nanmur

src/relax/transform/combine_parallel_matmul.cc

@@ -295,7 +295,8 @@ std::vector<BranchInfo> GetBranchInfo(Function f) {
       if (!match) return;
 
       auto matmul_call = match.value()[matmul_pat].as_or_throw<Call>();
-      auto matmul_lhs = matmul_call->args[0].as_or_throw<Var>();
+      const auto* matmul_lhs = matmul_call->args[0].as<VarNode>();
+      if (matmul_lhs == nullptr) return;
 
       std::optional<int> bias_dim = std::nullopt;
       std::optional<std::string> activation = std::nullopt;
@@ -310,7 +311,7 @@ std::vector<BranchInfo> GetBranchInfo(Function f) {
         }
       }
 
-      if (auto it = groups.find(matmul_lhs.get()); it != groups.end()) {
+      if (auto it = groups.find(matmul_lhs); it != groups.end()) {
         // Create a new branch in the existing parallel matmul subtree, and
         // invalidate bias and activation information when needed.
         BranchInfo* branch = &it->second;
@@ -326,7 +327,7 @@ std::vector<BranchInfo> GetBranchInfo(Function f) {
         }
       } else {
         // Create a new subgraph with one matmul
-        groups[matmul_lhs.get()] = {1, bias_dim, activation};
+        groups[matmul_lhs] = {1, bias_dim, activation};
       }
     });
 

tests/python/relax/test_transform_combine_parallel_matmul.py

@@ -743,5 +743,18 @@ def before(
     tvm_tirx.ir.assert_structural_equal(after, before)
 
 
+def test_skip_matmul_with_constant_lhs():
+    @R.function(private=True)
+    def before(x: R.Tensor((2, 2), "float32")) -> R.Tensor((2, 2), "float32"):
+        with R.dataflow():
+            out = R.matmul(R.const([[1.0, 2.0], [3.0, 4.0]], "float32"), x)
+            R.output(out)
+        return out
+
+    after = CombineParallelMatmul()(tvm_tirx.IRModule.from_expr(before))["main"]
+
+    tvm_tirx.ir.assert_structural_equal(after, before)
+
+
 if __name__ == "__main__":
     tvm_tirx.testing.main()