--- /dev/null
+/*
+ * 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
--- /dev/null
+# 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()
projects / platform / upstream / tvm.git / commitdiff