[Relax][Pass] FuseOps (tlc-pack#15)

* [Relax][Pass] FuseOps * Fix 2 bugs
MasterJH5574 · Apr 13, 2022 · 523a815 · 523a815
1 parent acd1327
commit 523a815
Show file tree

Hide file tree

Showing 3 changed files with 642 additions and 552 deletions.
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
@@ -139,7 +139,12 @@ TVM_DLL Pass LayoutRewrite();
 TVM_DLL Pass FoldConstant();
 
 /*!
- * \brief Fuse operators in an expr to a larger operator according to some rules.
+ * \brief This pass groups bindings in a dataflow block of Relaxfunctions and generate a new grouped
+ * Relax function for each group, according to the fusion algorithm described in the pass
+ * implementation. By grouping bindings into new Relax functions, we substitute the bindings in the
+ * function being manipulated into function calls to the new grouped function.
+ *
+ * A follow-up pass named "FuseTIR" will generate a TIR PrimFunc for each grouped function.
  * \param fuse_opt_level The level of fuse optimization.
  *        -1 indicates that the level will be inferred from pass context.
  * \return The Pass.

diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py
@@ -183,7 +183,12 @@ def LayoutRewrite() -> tvm.ir.transform.Pass:
 
 
 def FuseOps(fuse_opt_level=-1) -> tvm.ir.transform.Pass:
-    """Fuse operators in an expr to a larger operator according to some rules.
+    """This pass groups bindings in a dataflow block of Relax functions and generate a new grouped
+    Relax function for each group, according to the fusion algorithm described in the pass
+    implementation. By grouping bindings into new Relax functions, we substitute the bindings in
+    the function being manipulated into function calls to the new grouped function.
+
+    A follow-up pass named "FuseTIR" will generate a TIR PrimFunc for each grouped function.
 
     Parameters
     ----------