Add a combine batch_matmul pass (#5791)

* Add a combine batch_matmul pass

Contrary what you might expect, this doesn't share as much code with
the combine dense pass as it does with the combine 2d conv pass.
This is because it concatenates the "output feature" dimensions.

* fix docstring

include/tvm/relay/transform.h

@@ -228,6 +228,17 @@ TVM_DLL Pass CombineParallelConv2D(uint64_t min_num_branches = 3);
  */
 TVM_DLL Pass CombineParallelDense(uint64_t min_num_branches = 3);
 
+/*!
+ * \brief Combine parallel batch_matmul 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 CombineParallelBatchMatmul(uint64_t min_num_branches = 3);
+
 /*!
  * \brief Backward fold axis scaling into weights of conv/dense operators.
  *

python/tvm/relay/transform/transform.py

@@ -58,6 +58,7 @@ def build_config(opt_level=2,
                 "EliminateCommonSubexpr": 3,
                 "CombineParallelConv2D": 4,
                 "CombineParallelDense": 4,
+                "CombineParallelBatchMatmul": 4,
                 "FastMath": 4
             }
 
@@ -307,6 +308,39 @@ def CombineParallelDense(min_num_branches=3):
     """
     return _ffi_api.CombineParallelDense(min_num_branches)
 
+def CombineParallelBatchMatmul(min_num_branches=3):
+    """Combine multiple batch matmul operators into one. For example:
+
+    .. code-block
+                             data (1, 2, 3)
+                         /                  \
+        batch_matmul(data, (1, 4, 3))    batch_matmul(data, (1, 5, 3))
+            |                                |
+        elemwise/bcast (1, 2, 4)         elemwise/bcast (1, 2, 5)
+
+    Would become:
+
+    .. code-block
+
+                data (1, 2, 3)
+                |
+            batch_matmul(data, (1, 4+5, 3))
+                |
+            elemwise/bcast (1 ,2, 4+5)
+
+    Parameters
+    ----------
+    min_num_branches : int
+        The minimum number of required parallel branches for performing this
+        optimization.
+
+    Returns
+    -------
+    ret: tvm.transform.Pass
+        The registered pass that combines parallel dense operators.
+    """
+    return _ffi_api.CombineParallelBatchMatmul(min_num_branches)
+
 
 def AlterOpLayout():
     """Alternate the layouts of operators or replace primitive operators with

src/relay/backend/build_module.cc

@@ -278,6 +278,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::CombineParallelBatchMatmul(3));
     pass_seqs.push_back(transform::FoldConstant());
     pass_seqs.push_back(transform::FoldScaleAxis());
     pass_seqs.push_back(transform::CanonicalizeCast());

src/relay/transforms/combine_parallel_batch_matmul.cc

@@ -0,0 +1,160 @@
+/*
+ * 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.
+ */
+
+/*!
+ *
+ * \file combine_parallel_batch_matmul.cc
+ * \brief Combine parallel batch matmuls into a single one.
+ *
+ * This pass replaces batch_matmul that share the same lhs node with a
+ * single batch matmul.Elemwise and broadcast ops following batch_matmul are also
+ * combined if possible.
+ *
+ * This prevents launching multiple kernels in networks with multiple
+ * convolution branches, such as Inception block.
+ */
+
+#include <tvm/relay/analysis.h>
+#include <tvm/relay/attrs/nn.h>
+#include <tvm/relay/attrs/transform.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/relay/op_attr_types.h>
+#include <tvm/relay/transform.h>
+
+#include <unordered_map>
+#include <unordered_set>
+
+#include "./combine_parallel_op.h"
+#include "./expr_subst.h"
+#include "pattern_util.h"
+
+namespace tvm {
+namespace relay {
+
+class ParallelBatchMatmulCombiner : public ParallelOpCombiner {
+ public:
+  explicit ParallelBatchMatmulCombiner(uint64_t min_num_branches)
+      : ParallelOpCombiner("nn.batch_matmul", min_num_branches) {}
+
+ protected:
+  bool IsSupportedOp(const CallNode* n) { return true; }
+
+  bool CanOpsBeCombined(const CallNode* a, const CallNode* b) {
+    StructuralEqual eq;
+    const auto* rhs_a = a->args[1]->type_as<TensorTypeNode>();
+    const auto* rhs_b = b->args[1]->type_as<TensorTypeNode>();
+    const auto* restype_a = a->type_as<TensorTypeNode>();
+    const auto* restype_b = b->type_as<TensorTypeNode>();
+    // shape[2] is the contraction axis and automatically consistent
+    // if it were valid batch_matmul ops
+    auto res = eq(rhs_a->dtype, rhs_b->dtype) && eq(restype_a->dtype, restype_b->dtype) &&
+               (rhs_a->shape.size() == 3) && (rhs_b->shape.size() == 3) &&
+               eq(rhs_a->shape[0], rhs_b->shape[0]);
+    return res;
+  }
+
+  Call MakeCombinedOp(const Group& branches) {
+    const Op& batch_matmul = Op::Get("nn.batch_matmul");
+    Expr data = branches[0][0]->args[0];
+
+    Array<Expr> weights;
+    for (const auto& branch : branches) {
+      auto batch_matmul = branch[0];
+      weights.push_back(batch_matmul->args[1]);
+    }
+    Expr new_weight = MakeConcatenate(Tuple(weights), 1);
+    return Call(batch_matmul, {data, new_weight}, {}, {});
+  }
+
+  bool IsArgCompatible(const CallNode* a, const CallNode* b, size_t index) { return true; }
+
+  Call MakeCombinedCallFromFollowingOps(const Expr& data, const Group& branches, size_t depth,
+                                        size_t parent_index) {
+    Array<Expr> new_args;
+    const CallNode* call = branches[0][depth];
+
+    for (size_t i = 0; i < call->args.size(); i++) {
+      if (i == parent_index) {
+        new_args.push_back(data);
+        continue;
+      }
+
+      Array<Expr> tuple;
+      for (const auto& branch : branches) {
+        tuple.push_back(branch[depth]->args[i]);
+      }
+
+      auto concat = MakeConcatenate(Tuple(tuple), -1);
+      new_args.push_back(std::move(concat));
+    }
+
+    return Call(call->op, new_args, call->attrs, {});
+  }
+
+  void UpdateGroupOutput(const Expr& data, const Group& branches, size_t depth,
+                         ExprSubstMap* subst_map) {
+    int64_t index = 0;
+
+    for (const auto& branch : branches) {
+      const CallNode* batch_matmul = branch[0];
+      auto feature_dim = batch_matmul->args[1]->type_as<TensorTypeNode>()->shape[1];
+      auto fpp = tir::as_const_int(feature_dim);
+      int64_t features = *fpp;
+      std::vector<int64_t> begin;
+      std::vector<int64_t> end;
+      for (size_t i = 0; i < 2; i++) {
+        begin.push_back(0);
+        end.push_back(-1);
+      }
+      begin.push_back(index);
+      index += features;
+      end.push_back(features);
+      std::vector<int64_t> strides(begin.size(), 1);
+      std::vector<int64_t> ndarray_shape = {static_cast<int64_t>(begin.size())};
+      Constant begin_const = MakeConstantTensor(DataType::Int(64), ndarray_shape, begin);
+      Constant end_const = MakeConstantTensor(DataType::Int(64), ndarray_shape, end);
+      Constant strides_const = MakeConstantTensor(DataType::Int(64), ndarray_shape, strides);
+      auto slice = MakeStridedSlice(data, begin_const, end_const, strides_const, "size");
+      subst_map->insert({GetRef<Expr>(branch[depth]), slice});
+    }
+  }
+};
+
+/*! \brief Combine parallel batch_matmul if number of branches >= min_num_branches */
+Expr CombineParallelBatchMatmul(const Expr& expr, uint64_t min_num_branches) {
+  return ParallelBatchMatmulCombiner(min_num_branches).Combine(expr);
+}
+
+namespace transform {
+
+Pass CombineParallelBatchMatmul(uint64_t min_num_branches) {
+  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
+      [=](Function f, IRModule m, PassContext pc) {
+        return Downcast<Function>(CombineParallelBatchMatmul(f, min_num_branches));
+      };
+  return CreateFunctionPass(pass_func, 4, "CombineParallelBatchMatmul", {"InferType"});
+}
+
+TVM_REGISTER_GLOBAL("relay._transform.CombineParallelBatchMatmul")
+    .set_body_typed(CombineParallelBatchMatmul);
+
+}  // namespace transform
+
+}  // namespace relay
+}  // namespace tvm

tests/python/relay/test_pass_combine_parallel_batch_matmul.py

@@ -0,0 +1,146 @@
+# 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.
+# pylint: disable=invalid-name,too-many-locals,too-many-arguments,missing-module-docstring
+
+import tvm
+from tvm import relay
+from tvm.relay import transform
+
+
+def run_opt_pass(expr, opt_pass):
+    "runs the opt_pass on the expr of a function the function"
+    assert isinstance(opt_pass, tvm.transform.Pass)
+    mod = tvm.IRModule.from_expr(expr)
+    mod = opt_pass(mod)
+    return mod["main"]
+
+def test_combine_parallel_batch_matmul():
+    """Simple testcase."""
+    def before(x, w1, w2, w3):
+        args = [x, w1, w2, w3]
+        y1 = relay.nn.batch_matmul(x, w1)
+        y2 = relay.nn.batch_matmul(x, w2)
+        y3 = relay.nn.batch_matmul(x, w3)
+        y = relay.Tuple((y1, y2, y3))
+        return relay.Function(args, y)
+
+    def expected(x, w1, w2, w3):
+        # use a fixed order of args so alpha equal check can pass
+        s1 = w1.type_annotation.shape[1]
+        s2 = w2.type_annotation.shape[1]
+        s3 = w3.type_annotation.shape[1]
+        args = [x, w1, w2, w3]
+        w = relay.concatenate((w1, w2, w3), axis=1)
+        y = relay.nn.batch_matmul(x, w)
+        y1 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, 0], "int64"),
+                                 end=relay.const([-1, -1, s1], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y2 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, s1], "int64"),
+                                 end=relay.const([-1, -1, s2], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y3 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, s1+s2], "int64"),
+                                 end=relay.const([-1, -1, s3], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y = relay.Tuple((y1, y2, y3))
+        return relay.Function(args, y)
+
+    def check(b, i, j, k):
+        x = relay.var("x", shape=(b, i, k))
+        w1 = relay.var("w1", shape=(b, j, k))
+        w2 = relay.var("w2", shape=(b, j, k))
+        w3 = relay.var("w3", shape=(b, j, k))
+
+        y_before = before(x, w1, w2, w3)
+        y = run_opt_pass(y_before,
+                         transform.CombineParallelBatchMatmul(min_num_branches=2))
+        y_expected = expected(x, w1, w2, w3)
+        y_expected = run_opt_pass(y_expected, transform.InferType())
+        tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
+
+    check(2, 3, 5, 4)
+    check(1, 100, 200, 300)
+
+def test_combine_parallel_batch_matmul_biasadd():
+    """Simple testcase with bias"""
+    def before(x, w1, w2, w3, b1, b2, b3):
+        args = [x, w1, w2, w3, b1, b2, b3]
+        y1 = relay.nn.batch_matmul(x, w1)
+        y2 = relay.nn.batch_matmul(x, w2)
+        y3 = relay.nn.batch_matmul(x, w3)
+        y1 = relay.add(y1, b1)
+        y2 = relay.add(y2, b2)
+        y3 = relay.add(y3, b3)
+        y = relay.Tuple((y1, y2, y3))
+        return relay.Function(args, y)
+
+    def expected(x, w1, w2, w3, b1, b2, b3):
+        # use a fixed order of args so alpha equal check can pass
+        s1 = w1.type_annotation.shape[1]
+        s2 = w2.type_annotation.shape[1]
+        s3 = w3.type_annotation.shape[1]
+        args = [x, w1, w2, w3, b1, b2, b3]
+        w = relay.concatenate((w1, w2, w3), axis=1)
+        b = relay.concatenate((b1, b2, b3), axis=-1)
+        y = relay.nn.batch_matmul(x, w)
+        y = relay.add(y, b)
+        y1 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, 0], "int64"),
+                                 end=relay.const([-1, -1, s1], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y2 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, s1], "int64"),
+                                 end=relay.const([-1, -1, s2], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y3 = relay.strided_slice(y,
+                                 begin=relay.const([0, 0, s1+s2], "int64"),
+                                 end=relay.const([-1, -1, s3], "int64"),
+                                 strides=relay.const([1, 1, 1], 'int64'),
+                                 slice_mode="size")
+        y = relay.Tuple((y1, y2, y3))
+        return relay.Function(args, y)
+
+    def check(b, i, j, k):
+        x = relay.var("x", shape=(b, i, k))
+        w1 = relay.var("w1", shape=(b, j, k))
+        w2 = relay.var("w2", shape=(b, j, k))
+        w3 = relay.var("w3", shape=(b, j, k))
+        b1 = relay.var("b1", shape=(j,))
+        b2 = relay.var("b2", shape=(j,))
+        b3 = relay.var("b3", shape=(j,))
+
+        y_before = before(x, w1, w2, w3, b1, b2, b3)
+        y = run_opt_pass(y_before,
+                         transform.CombineParallelBatchMatmul(min_num_branches=2))
+        y_expected = expected(x, w1, w2, w3, b1, b2, b3)
+        y_expected = run_opt_pass(y_expected, transform.InferType())
+        tvm.ir.assert_structural_equal(y, y_expected, map_free_vars=True)
+
+    check(2, 3, 5, 4)
+    check(1, 100, 200, 300)
+
+
+if __name__ == "__main__":
+    test_combine_parallel_batch_matmul()
+    test_combine_parallel_batch_matmul_biasadd()