[BYOC][TRT] Add DFPattern support for TRT backend by mikepapadim · Pull Request #10759 · apache/tvm

mikepapadim · 2022-03-24T12:14:21Z

This PR adds DFPattern support for the TRT backend without removing the existing predicate registry.

Adds and extends the following:

In tensorrt.py: Add a pattern_table for all the supported ops and consumes the pre-existing op_registry checks
Adds an additional pass as unmerge_composites.cc. This is required for the TRT backend as it expects a single primitive function to work with, while the MergeComposite and PartitionGraph will produce a single function for each Composite pattern.
Adds test_inline_composites.py which tests the newly introduced pass.

Both the pattern-based and predicate-based pass sequences produce syntactically equivalent IRModules.
This is to ensure backwards compatibility."

Original Pass orderding:

  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(),
        ]
    )

Pass ordering with MergeComposites and UnmergeComposites:

  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.MergeComposite(pattern_table()),                                   <-------- Change #1
            transform.AnnotateTarget("tensorrt"),
            transform.MergeCompilerRegions(),
            transform.PartitionGraph(),
            transform.InlineComposites("tensorrt"),                                     <-------- Change #2
            transform.InferType(),
        ]
    )

@mbs-octoml @mbaret @masahi

… to Collage

mbaret

I think this pass needs its own unit tests so it can be tested outside of the TRT partitioning flow.

mbaret · 2022-03-24T13:45:01Z


 @_register_external_dynamic_check_func("nn.batch_matmul")
-def batch_matmul_annotate_fn(expr):
+def batch_matmul_annotate_fn(expr):  # pylint: disable=unused-variable


Is this change required?

mbaret · 2022-03-24T13:47:25Z

    for tup in pattern_table:
        if len(tup) == 2:
            pattern_name, pattern = tup
-            check = lambda extract: True


Could you explain this change?

missmatch in black autoformat it

mbaret · 2022-03-24T14:18:02Z

+
+namespace relay {
+
+class Unmerger : ExprMutator {


MixedModeMutator is now preferred where possible.

mbaret · 2022-03-24T14:21:39Z

+      Function gv = GetRef<Function>(function_var_node);
+      const auto* fn = gv.as<FunctionNode>();


Is this needed - it looks like we already start with the FunctionNode?

mbaret · 2022-03-24T14:22:01Z

+
+  Expr VisitExpr_(const CallNode* call_node) final {
+    Call vanilla_call = GetAnyCall(call_node);
+    const auto* global_var_node = vanilla_call->op.as<GlobalVarNode>();


mbaret · 2022-03-24T14:26:58Z

+      // Attrs need to be empty at this point to avoid propagating Composite and
+      // PartitionedFromPattern that fiddling TRT code gen for registered ops.
+      auto func = Function(fn->params, fn->body, fn->ret_type, fn->type_params, {});
+      return Bind(func->body, bind_map);


Not sure I understand this, why can't we just do

return Bind(fn->body, bind_map);

mbaret · 2022-03-24T14:29:25Z

+
+    if (!base_func->GetAttr<String>(attr::kCompiler).defined() &&
+        base_func->GetAttr<String>(attr::kCompiler) != target) {
+      return module;


I think it'd be better to continue; here rather than return, otherwise it seems if any partitioning for a different target has taken place, this will bail out.

mbaret · 2022-03-24T14:30:51Z

+ */
+
+/*!
+ * \file src/relay/transforms/unmerge_composites.cc


Personal preference for the name here would be either InlineComposite or RemoveComposite, not a huge deal though, if no one else agrees we can keep it as Unmerge.

InlineComposite makes sense, I will rename it

mbaret · 2022-03-24T14:32:50Z

+
+/*!
+ * \file src/relay/transforms/unmerge_composites.cc
+ * \brief Undo the partioned graphs originate from merge composite.


I think 'Inline composite functions for a given target' describes this a bit better.

mikepapadim · 2022-03-25T10:03:57Z

PTAL

mbaret · 2022-03-25T15:22:18Z

I still think this needs a couple of simple unit tests to confirm the behaviour. Also ping @mbs-octoml if you want to take a quick look.

mikepapadim · 2022-03-25T15:47:29Z

@mbs-octoml

mikepapadim · 2022-03-31T12:15:05Z

@mbaret PTAL. Under test_inline_composites.py are a couple o unit-tests for testing the new pass.

mbaret · 2022-03-31T13:57:51Z

+    print("merge composite reusult")
+    print(result)
+    print("---------------------")


We should probably omit the prints.

mbaret · 2022-03-31T13:59:46Z

+    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)
+        add_relu = relay.Function([in_1, in_2], relu_node)
+        return add_relu


I think this is the same as before() (given a and b aren't used). If all we really want to test is that doing InlineComposites undoes MergeComposite, we can probably just test that the result is equal to the input.

mbaret · 2022-03-31T14:00:38Z

+    """Utility function to check inline composites results."""
+    result = run_opt_pass(
+        graph, relay.transform.MergeComposite(pattern_table), import_prelude=import_prelude
+    )


I think we should put some form of check here just to confirm that a composite function has been created (so we know MergeComposite didn't just skip everything if for instance there was a pattern error).

mbaret · 2022-03-31T14:02:09Z

+       relu                            relu
+
+    """
+    pattern_table = [("add", make_add_pattern()), ("nn.relu", make_relu_pattern())]


This doesn't seem to match the description above.

mbaret · 2022-03-31T14:03:13Z

+"""
+
+
+def make_conv_bias_relu_pattern():


This pattern doesn't seem to be used, I think either add a test for it or remove.

mbaret

lgtm

This PR adds DFPattern support for the TRT backend without removing the existing predicate registry. Adds and extends the following: In tensorrt.py: Add a pattern_table for all the supported ops and consumes the pre-existing op_registry checks Adds an additional pass as unmerge_composites.cc. This is required for the TRT backend as it expects a single primitive function to work with, while the MergeComposite and PartitionGraph will produce a single function for each Composite pattern. Adds test_inline_composites.py which tests the newly introduced pass. Both the pattern-based and predicate-based pass sequences produce syntactically equivalent IRModules. This is to ensure backwards compatibility."

Michalis Papapdimitriou added 4 commits March 23, 2022 06:47

[TRT] Move ops from predicates to DFPattern-based to cater transition…

8c425a1

… to Collage

Register more patterns for TRT

548b5de

Add more patterns

4b10225

[WIP] Pre-cleanup on checks

067fc87

mikepapadim force-pushed the pattern_trt branch 2 times, most recently from a4f84ce to 53b0c48 Compare March 24, 2022 12:48

mbaret requested changes Mar 24, 2022

View reviewed changes

More cleanup and op support added for TRT patterns

a4d95c7

mikepapadim force-pushed the pattern_trt branch from 53b0c48 to a4d95c7 Compare March 24, 2022 16:50

Address PR comments

a3dc554

mikepapadim force-pushed the pattern_trt branch from a961ef8 to a3dc554 Compare March 25, 2022 10:08

mbs-octoml mentioned this pull request Mar 28, 2022

Collage RFC apache/tvm-rfcs#62

Merged

Add unit tests for inline composites pass

2a6060f

mbaret requested changes Mar 31, 2022

View reviewed changes

Address PR comments

6c39b81

mbaret approved these changes Apr 4, 2022

View reviewed changes

mbaret merged commit 98580a2 into apache:main Apr 4, 2022

mbs-octoml mentioned this pull request Jun 27, 2022

[BYOC] InlineCompilerFunctions helper pass #11923

Merged

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

		Function gv = GetRef<Function>(function_var_node);
		const auto* fn = gv.as<FunctionNode>();

Conversation

mikepapadim commented Mar 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mbaret left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mikepapadim commented Mar 25, 2022

Uh oh!

mbaret commented Mar 25, 2022

Uh oh!

mikepapadim commented Mar 25, 2022

Uh oh!

mikepapadim commented Mar 31, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mbaret left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mikepapadim commented Mar 24, 2022 •

edited

Loading