[BYOC][TRT] Add DFPattern support for TRT backend

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

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,

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 <tvm/relay/expr.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/relay/transform.h>
+
+#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<FunctionNode>();
+
+    if (function_node) {
+      Array<Expr> new_args;
+      new_args.reserve(vanilla_call->args.size());
+      for (auto arg : vanilla_call->args) {
+        new_args.push_back(VisitExpr(arg));
+      }
+
+      Map<Var, Expr> 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<CallGraphEntry*> original_entry;
+  ICHECK(target.defined());
+  for (auto* it : topo) {
+    auto base_func = module->Lookup(it->GetNameHint());
+
+    if (!base_func->GetAttr<String>(attr::kCompiler).defined() &&
+        base_func->GetAttr<String>(attr::kCompiler) != target) {
+      continue;
+    }
+
+    if (it->GetNameHint() != "main") {
+      if (const auto* fn = base_func.as<FunctionNode>()) {
+        auto func = GetRef<Function>(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<IRModule(IRModule, PassContext)> 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