diff --git a/python/tvm/relay/op/contrib/tensorrt.py b/python/tvm/relay/op/contrib/tensorrt.py index 3bd737e6e0fd..f24d366e598a 100644 --- a/python/tvm/relay/op/contrib/tensorrt.py +++ b/python/tvm/relay/op/contrib/tensorrt.py @@ -24,8 +24,11 @@ from tvm.ir import Op from tvm.relay import transform from tvm.relay.build_module import bind_params_by_name +from tvm.relay.dataflow_pattern import is_op, wildcard from tvm.relay.expr import Call, Constant, GlobalVar, Tuple, TupleGetItem, Var from tvm.relay.expr_functor import ExprMutator, ExprVisitor +from tvm.relay.op.contrib.register import register_pattern_table +from tvm.relay.op.transform import split logger = logging.getLogger("TensorRT") supported_types = ["float32", "float16"] @@ -103,6 +106,7 @@ def partition_for_tensorrt( max_workspace_size=1 << 30, use_fp16=False, use_uint8=False, + use_patterns=False, ): """Partition the graph greedily offloading supported operators to TensorRT. @@ -133,6 +137,9 @@ def partition_for_tensorrt( lower runtime, or if no low-precision implementation exists. use_uint8: Optional[bool] Allows, TRT to automatically convert FP32 inputs to UINT8. + use_patterns: Optional[bool] + Switches to use pattern-based op suppot by applying MergeCompsite and InlineComposites + passes. Returns ------- mod_and_config : Tuple[Module, Dict[str, Any]] @@ -161,32 +168,74 @@ def partition_for_tensorrt( if params: mod["main"] = bind_params_by_name(mod["main"], params) - seq = tvm.transform.Sequential( - [ - transform.InferType(), - RemoveDropoutPass(), - transform.RemoveUnusedFunctions(), - transform.ConvertLayout( - { - "nn.conv1d": ["NCW", "default"], - "nn.conv2d": ["NCHW", "default"], - "nn.conv3d": ["NCDHW", "default"], - "nn.conv2d_transpose": ["NCHW", "default"], - } - ), - transform.FoldConstant(), - transform.AnnotateTarget("tensorrt"), - transform.MergeCompilerRegions(), - transform.PartitionGraph(), - transform.InferType(), - ] - ) + + seq = get_pass_order(use_patterns) with tvm.transform.PassContext(opt_level=3, config={"relay.ext.tensorrt.options": config}): mod = seq(mod) mod = prune_tensorrt_subgraphs(mod) return mod, config +def get_pass_order(use_patterns): + """ + Get the pass ordering based on using predicates or patterns. + + Parameters + ---------- + use_patterns: Bool + True if pass needs to work with op patterns + Returns + ---------- + ret : Sequential + Pass object + """ + return ( + tvm.transform.Sequential( + [ + transform.InferType(), + RemoveDropoutPass(), + transform.RemoveUnusedFunctions(), + transform.ConvertLayout( + { + "nn.conv1d": ["NCW", "default"], + "nn.conv2d": ["NCHW", "default"], + "nn.conv3d": ["NCDHW", "default"], + "nn.conv2d_transpose": ["NCHW", "default"], + } + ), + transform.FoldConstant(), + transform.MergeComposite(pattern_table()), + transform.AnnotateTarget("tensorrt"), + transform.MergeCompilerRegions(), + transform.PartitionGraph(), + transform.InlineComposites("tensorrt"), + transform.InferType(), + ] + ) + if use_patterns + else tvm.transform.Sequential( + [ + transform.InferType(), + RemoveDropoutPass(), + transform.RemoveUnusedFunctions(), + transform.ConvertLayout( + { + "nn.conv1d": ["NCW", "default"], + "nn.conv2d": ["NCHW", "default"], + "nn.conv3d": ["NCDHW", "default"], + "nn.conv2d_transpose": ["NCHW", "default"], + } + ), + transform.FoldConstant(), + transform.AnnotateTarget("tensorrt"), + transform.MergeCompilerRegions(), + transform.PartitionGraph(), + transform.InferType(), + ] + ) + ) + + def check_dynamism(args, op_name): """ Check for dynamism inside any of the args in the op. @@ -914,6 +963,114 @@ def conv3d_transpose_annotate_fn(expr): # pylint: disable=unused-variable return True +def unary_op_pattern(op): + """Matches unary operation""" + pattern = is_op(op)(wildcard()) + return pattern + + +def binary_op_pattern(op): + """Matches binary operation""" + pattern = is_op(op)(wildcard(), wildcard()) + return pattern + + +@register_pattern_table("tensorrt") +def pattern_table(): + """Get the Tensorrt compiler pattern table for supported ops.""" + + return [ + ("tensorrt.nn.conv3d", binary_op_pattern("nn.conv3d"), conv3d_annotate_fn), + ("tensorrt.nn.conv2d", binary_op_pattern("nn.conv2d"), conv2d_annotate_fn), + ("tensorrt.nn.conv1d", binary_op_pattern("nn.conv1d"), conv1d_annotate_fn), + ( + "tensorrt.nn.conv2d_transpose", + binary_op_pattern("nn.conv2d_transpose"), + conv2d_transpose_annotate_fn, + ), + ("tensorrt.squeeze", binary_op_pattern("squeeze"), squeeze_annotate_fn), + ("tensorrt.add", binary_op_pattern("add"), add_annotate_fn), + ("tensorrt.nn.dense", unary_op_pattern("nn.dense"), dense_annotate_fn), + ("tensorrt.bias_add", binary_op_pattern("nn.bias_add"), bias_add_annotate_fn), + ( + "tensorrt.nn.batch_matmul", + binary_op_pattern("nn.batch_matmul"), + batch_matmul_annotate_fn, + ), + ("tensorrt.divide", binary_op_pattern("divide")), + ("tensorrt.multiply", binary_op_pattern("multiply")), + ("tensorrt.split", unary_op_pattern("split")), + ("tensorrt.reshape", unary_op_pattern("reshape")), + ("tensorrt.nn.relu", unary_op_pattern("nn.relu")), + ( + "tensorrt.nn.leaky_relu", + unary_op_pattern("nn.leaky_relu"), + trt_version_annotate_fn((5, 1, 5)), + ), + ("tensorrt.nn.pad", unary_op_pattern("nn.pad")), + ("tensorrt.sigmoid", unary_op_pattern("sigmoid")), + ("tensorrt.tanh", unary_op_pattern("tanh")), + ("tensorrt.exp", unary_op_pattern("exp")), + ("tensorrt.log", unary_op_pattern("log")), + ("tensorrt.sqrt", unary_op_pattern("sqrt")), + ("tensorrt.abs", unary_op_pattern("abs")), + ("tensorrt.power", unary_op_pattern("power")), + ("tensorrt.negative", unary_op_pattern("negative")), + ("tensorrt.nn.batch_flatten", unary_op_pattern("nn.batch_flatten")), + ("tensorrt.sin", unary_op_pattern("sin"), trt_version_annotate_fn((5, 1, 5))), + ("tensorrt.clip", unary_op_pattern("clip")), + ("tensorrt.cos", unary_op_pattern("cos"), trt_version_annotate_fn((5, 1, 5))), + ("tensorrt.atan", unary_op_pattern("atan"), trt_version_annotate_fn((5, 1, 5))), + ("tensorrt.ceil", unary_op_pattern("ceil"), trt_version_annotate_fn((5, 1, 5))), + ("tensorrt.floor", unary_op_pattern("floor")), + ("tensorrt.erf", unary_op_pattern("erf"), trt_version_annotate_fn((7, 0, 0))), + ("tensorrt.sum", unary_op_pattern("sum"), reduce_annotate_fn), + ("tensorrt.prod", unary_op_pattern("prod"), reduce_annotate_fn), + ("tensorrt.max", unary_op_pattern("max"), reduce_annotate_fn), + ("tensorrt.min", unary_op_pattern("min"), reduce_annotate_fn), + ("tensorrt.max", unary_op_pattern("max"), reduce_annotate_fn), + ("tensorrt.concatenate", unary_op_pattern("concatenate"), concatenate_annotate_fn), + ("tensorrt.expand_dims", unary_op_pattern("expand_dims"), expand_dims_annotate_fn), + ( + "tensorrt.layout_transform", + unary_op_pattern("layout_transform"), + layout_transform_annotate_fn, + ), + ("tensorrt.transpose", unary_op_pattern("transpose"), transpose_annotate_fn), + ("tensorrt.reshape", unary_op_pattern("reshape"), reshape_annotate_fn), + ("tensorrt.split", unary_op_pattern("split"), split), + ("tensorrt.nn.pad", unary_op_pattern("nn.pad"), pad_annotate_fn), + ("tensorrt.strided_slice", unary_op_pattern("strided_slice"), strided_slice_annotate_fn), + ( + "tensorrt.nn.adaptive_avg_pool2d", + unary_op_pattern("nn.adaptive_avg_pool2d"), + adaptive_avg_pool2d_annotate_fn, + ), + ("tensorrt.nn.max_pool3d", unary_op_pattern("nn.max_pool3d"), max_pool_3d_annotate_fn), + ("tensorrt.nn.avg_pool3d", unary_op_pattern("nn.avg_pool3d"), avg_pool_3d_annotate_fn), + ( + "tensorrt.nn.conv3d_transpose", + unary_op_pattern("nn.conv3d_transpose"), + conv3d_transpose_annotate_fn, + ), + ("tensorrt.nn.softmax", unary_op_pattern("nn.softmax"), softmax_annotate_fn), + ("tensorrt.nn.layer_norm", unary_op_pattern("nn.layer_norm"), layer_norm_annotate_fn), + ("tensorrt.nn.max_pool2d", unary_op_pattern("nn.max_pool2d"), max_pool_2d_annotate_fn), + ("tensorrt.nn.avg_pool2d", unary_op_pattern("nn.avg_pool2d"), avg_pool_2d_annotate_fn), + ("tensorrt.nn.max_pool3d", unary_op_pattern("nn.max_pool3d"), max_pool_3d_annotate_fn), + ( + "tensorrt.nn.global_max_pool2d", + unary_op_pattern("nn.global_max_pool2d"), + global_max_pool_2d_annotate_fn, + ), + ( + "tensorrt.nn.global_avg_pool2d", + unary_op_pattern("nn.global_avg_pool2d"), + global_avg_pool_2d_annotate_fn, + ), + ] + + class IsComputeIntensiveGraph(ExprVisitor): """ Visits the Graph recursively and checks if it contains compute heavy ops like convolutions and diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py index 99c61c5bd96f..e4ee14b62941 100644 --- a/python/tvm/relay/transform/transform.py +++ b/python/tvm/relay/transform/transform.py @@ -789,6 +789,24 @@ def Inline(): return _ffi_api.Inline() +def InlineComposites(target): + """Perform inlining on the given Relay IR module. The functions originate + from the MergeComposite pass based on an input pattern table will fold back + to main. Currently, this is used for the TRT BYOC which expects a single + primitive function to operate on. + + Parameters + ---------- + target: str + The byoc target for which ops need to fold back to primitive function. + Returns + ------- + ret: tvm.transform.Pass + The registered pass that performs inlining for a Relay IR module. + """ + return _ffi_api.InlineComposites(target) + + def gradient(expr, mod=None, mode="higher_order"): """ Transform the input function, diff --git a/src/relay/transforms/inline_composites.cc b/src/relay/transforms/inline_composites.cc new file mode 100644 index 000000000000..63e7d078b0c5 --- /dev/null +++ b/src/relay/transforms/inline_composites.cc @@ -0,0 +1,119 @@ +/* + * 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 src/relay/transforms/inline_composites.cc + * \brief Undo the partioned graphs originate from merge composite. + */ +#include +#include +#include + +#include "../analysis/call_graph.h" +#include "../op/call/call.h" + +using namespace tvm::runtime; + +namespace tvm { + +namespace relay { + +class CompositeInliner : public MixedModeMutator { + public: + explicit CompositeInliner(CallGraphEntry* cur_node, CallGraphNode* call_graph) + : cur_node_(cur_node), call_graph_(call_graph) {} + + Expr Rewrite_(const CallNode* call_node) { + Call vanilla_call = GetAnyCall(call_node); + const auto* function_node = vanilla_call->op.as(); + + if (function_node) { + Array new_args; + new_args.reserve(vanilla_call->args.size()); + for (auto arg : vanilla_call->args) { + new_args.push_back(VisitExpr(arg)); + } + + Map bind_map; + for (size_t i = 0; i < new_args.size(); i++) { + bind_map.Set(function_node->params[i], new_args[i]); + } + + // Attrs need to be empty at this point to avoid propagating Composite and + // PartitionedFromPattern that fiddling TRT code gen for registered ops. + return Bind(function_node->body, bind_map); + } + + return MixedModeMutator::VisitExpr_(call_node); + } + + Function Inline(const Function& func) { + return WithFields(func, func->params, VisitExpr(func->body)); + } + + private: + /*! + * \brief The current call graph entry that is being handled. Each entry + * contains a global function. + */ + CallGraphEntry* cur_node_; + /*! \brief The call graph that is used for global function lookup. */ + const CallGraphNode* call_graph_; +}; + +IRModule InlineComposites(const IRModule& module, runtime::String target) { + CallGraph cg(module); + auto topo = cg->TopologicalOrder(); + std::reverse(topo.begin(), topo.end()); + std::unordered_set original_entry; + ICHECK(target.defined()); + for (auto* it : topo) { + auto base_func = module->Lookup(it->GetNameHint()); + + if (!base_func->GetAttr(attr::kCompiler).defined() && + base_func->GetAttr(attr::kCompiler) != target) { + continue; + } + + if (it->GetNameHint() != "main") { + if (const auto* fn = base_func.as()) { + auto func = GetRef(fn); + auto new_func = CompositeInliner(it, cg.operator->()).Inline(func); + cg->module->Update(it->GetGlobalVar(), new_func); + } + } + } + return module; +} + +namespace transform { + +Pass InlineComposites(runtime::String target) { + runtime::TypedPackedFunc pass_func = + [=](IRModule m, PassContext pc) { return relay::InlineComposites(m, target); }; + return CreateModulePass(pass_func, 0, "InlineComposites", {}); +} + +TVM_REGISTER_GLOBAL("relay._transform.InlineComposites").set_body_typed(InlineComposites); + +} // namespace transform + +} // namespace relay + +} // namespace tvm diff --git a/tests/python/relay/test_pass_inline_composites.py b/tests/python/relay/test_pass_inline_composites.py new file mode 100644 index 000000000000..54fc08c87918 --- /dev/null +++ b/tests/python/relay/test_pass_inline_composites.py @@ -0,0 +1,165 @@ +# 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, missing-docstring, too-many-statements +"""Unit tests for inline composites.""" +import pytest +import tvm +from tvm import relay, tir +from tvm.relay.dataflow_pattern import TupleGetItemPattern, is_op, wildcard +from tvm.relay.testing import run_opt_pass + +""" +The inline composite pass is designed to inline multiple kernel generated through +the merge composite composite pass. The underlying idea is to inline N kernels +produced from merge composite based on a given set of pattern into a single IR module. +Also, clears Composite and PartionedFromPatterns that infer with certain BYOC implementations + +For example suppose we have the graph: + + a b + \ / + add + | + relu + +Merge composite will wrap each standalone op to it's own function, while setting Composite and +PartitionedFromPattern attrs. + +Relay IR after merge composite pass when registering each op as a standalone pattern: +fn (%a: Tensor[(10, 10), float32], %b: Tensor[(10, 10), float32]) -> Tensor[(10, 10), float32] { + %0 = fn (%FunctionVar_0_01: Tensor[(10, 10), float32], %FunctionVar_0_1: Tensor[(10, 10), float32], PartitionedFromPattern="add_", Composite="add") -> Tensor[(10, 10), float32] { + add(%FunctionVar_0_01, %FunctionVar_0_1) /* ty=Tensor[(10, 10), float32] */ + }; + %1 = %0(%a, %b) /* ty=Tensor[(10, 10), float32] */; + %2 = fn (%FunctionVar_0_0: Tensor[(10, 10), float32], PartitionedFromPattern="nn.relu_", Composite="nn.relu") -> Tensor[(10, 10), float32] { + nn.relu(%FunctionVar_0_0) /* ty=Tensor[(10, 10), float32] */ + }; + %2(%1) /* ty=Tensor[(10, 10), float32] */ +} + +Relay IR after inline composites pass: +fn (%a: Tensor[(10, 10), float32], %b: Tensor[(10, 10), float32]) -> Tensor[(10, 10), float32] { + %0 = add(%a, %b) /* ty=Tensor[(10, 10), float32] */; + nn.relu(%0) /* ty=Tensor[(10, 10), float32] */ +} + +One convenient use of this pass is to use Pattern-based operator support to move away +from the original operator predicates, and inline them into a single primitive function to offload it +to an external BYOC backend, such as TensorRT. +""" + + +def make_add_relu_pattern(): + r"""Create a pattern to match the following graph. + + add + | + relu + """ + add_node = wildcard() + wildcard() + r = is_op("nn.relu")(add_node) + return r + + +def make_relu_pattern(): + r"""Create a pattern to match the following graph + a + | + relu + | + """ + pattern = is_op("nn.relu")(wildcard()) + return pattern + + +def make_add_pattern(): + r"""Create a pattern to match the following graph + a b + \ / + add + | + """ + pattern = is_op("add")(wildcard(), wildcard()) + return pattern + + +def check_success_composite_pass(func): + return func.body.op.attrs["Composite"] is not None + + +def check_result(pattern_table, expected_graph, import_prelude=False): + """Utility function to check inline composites results.""" + result = run_opt_pass( + expected_graph, relay.transform.MergeComposite(pattern_table), import_prelude=import_prelude + ) + assert check_success_composite_pass( + result + ), "Merge Composite pass didn't produced partioned from Pattern" + result = run_opt_pass( + expected_graph, relay.transform.InlineComposites(target=""), import_prelude=import_prelude + ) + assert not relay.analysis.free_vars(result), "Found free vars in the result graph: {0}".format( + str(result) + ) + expected = run_opt_pass(expected_graph, relay.transform.InferType()) + assert tvm.ir.structural_equal( + result, expected, map_free_vars=True + ), "Graph mismatch: output vs. expected\n{0}\n=====\n{1}".format(str(result), str(expected)) + + +def test_single_op_registry(): + r"""Test inline composite pass is correctly inline the post-merge composite graph. + + We could expect the patterns `make_add_pattern` and `make_relu_pattern` to be inlined + into a single func instead of an single func per registered pattern. + + """ + pattern_table = [("add", make_add_pattern()), ("nn.relu", make_relu_pattern())] + + def expected(): + in_1 = relay.var("in_1", shape=(10, 10)) + in_2 = relay.var("in_2", shape=(10, 10)) + add_node = relay.add(in_1, in_2) + relu_node = relay.nn.relu(add_node) + add_relu = relay.Function([in_1, in_2], relu_node) + return add_relu + + check_result(pattern_table, expected()) + + +def test_mix_fused_and_single_op(): + r"""Test inline composite pass is correctly inline the merge composite result""" + pattern_table = [("add_relu", make_add_relu_pattern()), ("nn.relu", make_relu_pattern())] + + def expected(): + a = relay.var("a", shape=(10, 10)) + b = relay.var("b", shape=(10, 10)) + + # add_relu function + in_1 = relay.var("in_1", shape=(10, 10)) + in_2 = relay.var("in_2", shape=(10, 10)) + add_node = relay.add(in_1, in_2) + relu_node = relay.nn.relu(add_node) + relu_nd = relay.nn.relu(relu_node) + add_relu = relay.Function([in_1, in_2], relu_nd) + return add_relu + + check_result(pattern_table, expected()) + + +if __name__ == "__main__": + pytest.main()