[Unity] Add pass for combining parallel matmul

include/tvm/relax/dataflow_matcher.h

@@ -58,6 +58,16 @@ Optional<Map<DFPattern, Expr>> ExtractMatchedExpr(
 TVM_DLL Optional<Map<DFPattern, Var>> MatchGraph(const PatternContext& ctx,
                                                  const DataflowBlock& dfb);
 
+/**
+ * \brief Rewrite a function with the given pattern and the rewriter function.
+ * \param ctx The pattern constraint context under which rewriting takes place.
+ * \param rewriter The function to be called on a successful matching for rewriting.
+    Given the map of patterns and corresponding variables (bound variables or parameters),
+    it should return a map that specifies new values for matched bound variables.
+ * \param f The function to rewrite
+ * \return The rewritten or the input function, depending on the pattern matching result.
+ */
+TVM_DLL Function RewriteBindings(const PatternContext& ctx, PackedFunc rewriter, Function f);
 }  // namespace relax
 }  // namespace tvm
 

include/tvm/relax/dataflow_pattern.h

@@ -248,14 +248,12 @@ class PatternContext : public ObjectRef {
   /*! \brief Get the constraint context object on the top of the stack */
   TVM_DLL static Optional<PatternContext> Current();
 
-  class Internal;
-
- private:
   /*! \brief The RAII-like entry of a constraint context scope */
-  TVM_DLL void EnterWithScope();
+  TVM_DLL void EnterWithScope() const;
   /*! \brief The RAII-like exit of a constraint context scope */
-  TVM_DLL void ExitWithScope();
-  friend class Internal;
+  TVM_DLL void ExitWithScope() const;
+
+ private:
   friend class With<PatternContext>;
 };
 

python/tvm/relax/transform/transform.py

@@ -934,6 +934,22 @@ def SplitCallTIRByPattern(patterns, fcodegen) -> tvm.ir.transform.Pass:
     return _ffi_api.SplitCallTIRByPattern(patterns, fcodegen)  # type: ignore
 
 
+def CombineParallelMatmul():
+    """Combine multiple matmul operators sharing the same LHS matrix into one,
+    followed by slicing. When all matmul branches in a tree have the same set of fused ops,
+    the fused ops are applied to the combined matmul output before slicing.
+
+    Currently, only a limited set of fused ops is supported. It includes bias add,
+    relu, gelu, and silu activation.
+
+    Returns
+    -------
+    ret : tvm.transform.Pass
+        The corresponding pass.
+    """
+    return _ffi_api.CombineParallelMatmul()  # type: ignore
+
+
 def _wrap_class_function_pass(pass_cls, pass_info):
     """Wrap a python class as function pass."""
 

src/relax/ir/dataflow_matcher.cc

@@ -22,6 +22,7 @@
  * \brief The dataflow pattern matcher for Relax.
  */
 
+#include <tvm/node/structural_equal.h>
 #include <tvm/relax/analysis.h>
 #include <tvm/relax/dataflow_matcher.h>
 #include <tvm/relax/dataflow_pattern.h>
@@ -791,7 +792,7 @@ class PatternRewriter : ExprMutator {
       : ctx_(ctx), rewriter_func_(rewriter_func), params_(params) {}
 
   template <typename PatternType>
-  static Expr Run(PatternType pat, PackedFunc rewriter_func, Function f) {
+  static Function Run(PatternType pat, PackedFunc rewriter_func, Function f) {
     std::unordered_set<const VarNode*> params;
     for (const auto& p : f->params) {
       params.insert(p.get());
@@ -868,15 +869,17 @@ class PatternRewriter : ExprMutator {
 
       std::unordered_set<const VarNode*> emitted_vars;
 
+      bool changed = false;
       for (size_t i = 0; i < block->bindings.size(); ++i) {
         const auto& binding = block->bindings[i];
         if (auto var_bind = binding.as<VarBindingNode>()) {
-          if (replacements.count(var_bind->var)) {
-            auto new_val = replacements[var_bind->var];
+          if (auto new_val = replacements.Get(var_bind->var).value_or(var_bind->value);
+              !StructuralEqual()(var_bind->value, new_val)) {
             Array<Binding> pending_bindings(block->bindings.begin() + i + 1, block->bindings.end());
             // Make sure there is no unbound variable used in the new value before it is emitted
             EmitUsedVars(new_val, pending_bindings, &emitted_vars);
             this->ReEmitBinding(var_bind, builder_->Normalize(new_val));
+            changed = true;
           } else if (!emitted_vars.count(var_bind->var.get())) {
             this->VisitBinding(binding);
             emitted_vars.insert(var_bind->var.get());
@@ -885,7 +888,11 @@ class PatternRewriter : ExprMutator {
           this->VisitBinding(binding);
         }
       }
-      return RewriteDataflowBlockFixedPoint(builder_->EndBlock());
+
+      auto new_block = builder_->EndBlock();
+
+      if (!changed) return new_block;
+      return RewriteDataflowBlockFixedPoint(new_block);
     }
     return block;
   }
@@ -909,15 +916,16 @@ class PatternRewriter : ExprMutator {
   std::unordered_map<const Object*, Expr> memo_;
 };
 
+Function RewriteBindings(const PatternContext& ctx, PackedFunc rewriter, Function f) {
+  return PatternRewriter::Run(ctx, rewriter, f);
+}
+
 TVM_REGISTER_GLOBAL("relax.dpl.rewrite_call")
     .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
       return PatternRewriter::Run(pat, rewriter, f);
     });
 
-TVM_REGISTER_GLOBAL("relax.dpl.rewrite_bindings")
-    .set_body_typed([](const PatternContext& ctx, PackedFunc rewriter, Function f) {
-      return PatternRewriter::Run(ctx, rewriter, f);
-    });
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite_bindings").set_body_typed(RewriteBindings);
 
 }  // namespace relax
 }  // namespace tvm

src/relax/ir/dataflow_pattern.cc

@@ -412,9 +412,9 @@ PatternContext::PatternContext(bool incremental) {
   data_ = std::move(n);
 }
 
-void PatternContext::EnterWithScope() { pattern_ctx_stack().push(*this); }
+void PatternContext::EnterWithScope() const { pattern_ctx_stack().push(*this); }
 
-void PatternContext::ExitWithScope() {
+void PatternContext::ExitWithScope() const {
   ICHECK(pattern_ctx_stack().top().same_as(*this));
   pattern_ctx_stack().pop();
 }
@@ -610,15 +610,13 @@ TVM_REGISTER_GLOBAL("relax.dpl.current_context").set_body_typed([] {
   return PatternContext::Current();
 });
 
-class PatternContext::Internal {
- public:
-  static void EnterScope(PatternContext pass_ctx) { pass_ctx.EnterWithScope(); }
-  static void ExitScope(PatternContext pass_ctx) { pass_ctx.ExitWithScope(); }
-};
-
-TVM_REGISTER_GLOBAL("relax.dpl.enter_context").set_body_typed(PatternContext::Internal::EnterScope);
+TVM_REGISTER_GLOBAL("relax.dpl.enter_context").set_body_typed([](const PatternContext& ctx) {
+  ctx.EnterWithScope();
+});
 
-TVM_REGISTER_GLOBAL("relax.dpl.exit_context").set_body_typed(PatternContext::Internal::ExitScope);
+TVM_REGISTER_GLOBAL("relax.dpl.exit_context").set_body_typed([](const PatternContext& ctx) {
+  ctx.ExitWithScope();
+});
 
 }  // namespace relax
 }  // namespace tvm