[Relay] Add new IR pass CombineParallelDense

docs/api/python/relay/transform.rst

@@ -46,6 +46,8 @@ tvm.relay.transform
 
 .. autofunction:: tvm.relay.transform.CombineParallelConv2D
 
+.. autofunction:: tvm.relay.transform.CombineParallelDense
+
 .. autofunction:: tvm.relay.transform.AlterOpLayout
 
 .. autofunction:: tvm.relay.transform.Legalize

include/tvm/relay/transform.h

@@ -482,6 +482,17 @@ TVM_DLL Pass EliminateCommonSubexpr(PackedFunc fskip = nullptr);
  */
 TVM_DLL Pass CombineParallelConv2D(uint64_t min_num_branches = 3);
 
+/*!
+ * \brief Combine parallel dense ops into a single batch_matmul if the
+ * number of branches of this dense operator is not less than
+ * `min_num_branch`.
+ *
+ * \param min_num_branches The minimun number of branches.
+ *
+ * \return The pass.
+ */
+TVM_DLL Pass CombineParallelDense(uint64_t min_num_branches = 3);
+
 /*!
  * \brief Backward fold axis scaling into weights of conv/dense operators.
  *

python/tvm/relay/transform.py

@@ -138,6 +138,7 @@ def build_config(opt_level=2,
                 "CanonicalizeCast": 3,
                 "EliminateCommonSubexpr": 3,
                 "CombineParallelConv2D": 4,
+                "CombineParallelDense": 4
             }
 
     fallback_device : int, str, or tvm.TVMContext, optional
@@ -400,6 +401,33 @@ def CombineParallelConv2D(min_num_branches=3):
     return _transform.CombineParallelConv2D(min_num_branches)
 
 
+def CombineParallelDense(min_num_branches=3):
+    """Combine multiple dense operators into one. For example:
+
+             data
+          /       \
+     dense (2,2)  dense (2,2)
+
+    Would become:
+
+             data
+              |
+        batch_matmul (2,2,2)
+
+    Parameters
+    ----------
+    min_num_branches : int
+        The minimum number of required parallel branches for performing this
+        optimization.
+
+    Returns
+    -------
+    ret: tvm.relay.Pass
+        The registered pass that combines parallel dense operators.
+    """
+    return _transform.CombineParallelDense(min_num_branches)
+
+
 def AlterOpLayout():
     """Alternate the layouts of operators or replace primitive operators with
     other expressions.

src/relay/backend/build_module.cc

@@ -299,6 +299,7 @@ class RelayBuildModule : public runtime::ModuleNode {
     });
     pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
     pass_seqs.push_back(transform::CombineParallelConv2D(3));
+    pass_seqs.push_back(transform::CombineParallelDense(3));
     pass_seqs.push_back(transform::FoldConstant());
     pass_seqs.push_back(transform::FoldScaleAxis());
     pass_seqs.push_back(transform::CanonicalizeCast());