[feat] Add dependency awareness to torch-trt partitioning

[mfeliz-cruise]

Adds a heuristic to torch-trt partitioning's segmentation to avoid materializing segments until we hit a dependency of that segment. This can significantly reduce the number of segments/engines in cases where the linear traversal of torchscipt nodes would otherwise produce alternating torch and TRT segments which are not dependent on each-other

Fixes # (issue)

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)

• New feature (non-breaking change which adds functionality)

• Breaking change (fix or feature that would cause existing functionality to not work as expected)

• This change requires a documentation update

• My code follows the style guidelines of this project (You can use the linters)

• I have performed a self-review of my own code

• I have commented my code, particularly in hard-to-understand areas and hacks

• I have made corresponding changes to the documentation

• I have added tests to verify my fix or my feature

• New and existing unit tests pass locally with my changes

• I have added the relevant labels to my PR in so that relevant reviewers are notified

[mfeliz-cruise]

I think this can be handled more simply in future by directly partitioning on the dependency graph. This would require updating the min_block_size logic, but would remove the need to merge segments after the initial partition.

[narendasan]

@mfeliz-cruise we are currently working on a major restructuring of the partitioning phase to hopefully bring it closer to other design patterns in the project and make it easier to debug the state and develop new features (#1263). Could you try rebasing this work on that branch and point the PR to merge into partitioning_ctx?

[bowang007]

@mfeliz-cruise we are currently working on a major restructuring of the partitioning phase to hopefully bring it closer to other design patterns in the project and make it easier to debug the state and develop new features (#1263). Could you try rebasing this work on that branch and point the PR to merge into partitioning_ctx?

Looks like now there are still some errors on that branch.

[mfeliz-cruise]

@mfeliz-cruise we are currently working on a major restructuring of the partitioning phase to hopefully bring it closer to other design patterns in the project and make it easier to debug the state and develop new features (#1263). Could you try rebasing this work on that branch and point the PR to merge into partitioning_ctx?

Looks like now there are still some errors on that branch.

I'll hold off for now until it stabilizes.

[mfeliz-cruise]

@mfeliz-cruise we are currently working on a major restructuring of the partitioning phase to hopefully bring it closer to other design patterns in the project and make it easier to debug the state and develop new features (#1263). Could you try rebasing this work on that branch and point the PR to merge into partitioning_ctx?

Looks like now there are still some errors on that branch.

I'll hold off for now until it stabilizes.

I've rebased and should be ready for review.

[peri044]

Hello @mfeliz-cruise
I went through this PR and observed the test cases. I ran them and understand the final graphs you are trying to achieve post segmentation. However, I'm a little unclear on the code logic. Can you explain the heuristic and logic in a write-up ( with also some references to the modified test cases) ? Since this is advanced segmentation/merging, it would be nice for this write-up to serve as a reference for future.

[mfeliz-cruise]

Sure @peri044, do you have a standard place you put this kind of documentation or should I just expand the PR description?

[narendasan]

We keep documentation for contributors on the implementation of partitioning here: https://pytorch.org/TensorRT/contributors/partitioning.html#partitioning

[mfeliz-cruise]

We keep documentation for contributors on the implementation of partitioning here: https://pytorch.org/TensorRT/contributors/partitioning.html#partitioning

I've taken a first pass at documenting this in docsrc/contributors/partitioning.rst

[peri044]

• How does the merge adjacent segments work ? What's the criteria to merge ?

[mfeliz-cruise]

• How does the merge adjacent segments work ? What's the criteria to merge ?
  I added some more about this in partitioning.rst. Let me know if you have more questions @peri044.
  https://github.com/pytorch/TensorRT/pull/1304/files#diff-a9f595cc75ff499ecbaedbe818d92ad9543b68b00243b8da7f250f72e7f12cdcR239

[peri044]

Hello @mfeliz-cruise ,
I removed this snippet and ran the tests and they ran successfully. I'm wondering what is the usecase behind this snippet ?

From my understanding this is written for in-place ops.
For eg: In the sample graph

%2 = aten::cat(%1)
%2 = aten::append(%2, %3)
%4 = aten::relu
For n = aten::append, use.user would be %2 (aten::cat output) which would not be isAfter(n) correct ?
Do you have an example in mind which uses this ?

[mfeliz-cruise]

Hello @mfeliz-cruise , I removed this snippet and ran the tests and they ran successfully. I'm wondering what is the usecase behind this snippet ?

From my understanding this is written for in-place ops. For eg: In the sample graph

%2 = aten::cat(%1) %2 = aten::append(%2, %3) %4 = aten::relu For n = aten::append, use.user would be %2 (aten::cat output) which would not be isAfter(n) correct ? Do you have an example in mind which uses this ?

It would be a case like this #1018 where we have an op that modifies its input without producing the modified value:
= aten::_set_item(%out_dict.1, %3, %x.1)
%z.1 : Tensor = aten::__getitem__(%out_dict.1, %3)

Here %out_dict.1 is modified by _set_item and we should recognize that this makes aten::_getitem__ a dependent of the set. If we only look at the node outputs here we would not identify this relationship.