Centralizing Partitioning State #1263

narendasan · 2022-08-14T17:11:48Z

Description

It is incredibly hard to track the state of the partitioning process because the pipeline relies on a number of shared lookup tables that are modified in place during the course of partitioning.

PartitioningCtx, similar to ConversionCtx centralizes key information like decisions about node executors, user settings, the in progress graph etc. so that this information is uniformly managed and queryable at any point during the pipelines execution. This means it is hard to client code to alter state in unpredictable ways and it also reduces the amount of arguments that need to be passed to each phase.

Addresses #949

Type of change

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

New feature (non-breaking change which adds functionality)

Checklist:

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

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

narendasan · 2022-08-14T17:19:02Z

@bowang007 This should be roughly equivalent to our current system. There are a couple improvements on this that I'd like to do that I want your input on.

I renamed a bunch of stuff to hopefully help make it clearer to new devs, let me know if you disagree about anything.
I want every node in the global_fallback_map (now called node_executor_decision_map) and make an explicit decision on if it will run in torch or tensorrt (this can be changed like if a block would run in trt but is not large enough). That way in debugging or otherwise I can just ask about a particular node and get back what it will run on. Right now I am having trouble doing this because there is an assumption in the system that if the node is not in the map it must run in tensorrt.
There is still work to be done on being able to easily dump state at any given time.
Are there other data sources that we might want to have managed by the context?

narendasan · 2022-08-14T17:21:54Z

core/partitioning/partitioningctx/PartitioningCtx.cpp

+std::ostream& operator<<(std::ostream& os, const NodeExecutorDecision& format) {
+  switch (format) {
+    case NodeExecutorDecision::kUNSUPPORTED:
+      return os << "to run torch due to lack of converter support";


"run in torch"

narendasan · 2022-08-14T17:28:29Z

core/compiler.cpp

@@ -272,8 +269,8 @@ GraphAndMapping ConstructFallbackGraph(
        // convert the 2 blocks in prim::if and get the converted graph with mappings
        std::vector<GraphAndMapping> graph_and_mappings;
        for (auto cur_block : if_node->blocks()) {
-          graph_and_mappings.push_back(
-              ConstructFallbackGraph(new_mod, cur_block, example_tensor_map, cfg, static_params, fallback_nodes));
+          graph_and_mappings.push_back(ConstructFallbackGraph_(


Not wildly pressing but is there a way to do all the partitioning beforehand then go through and compile specific blocks? Having them mixed is not as easy to debug

@bowang007 thoughts here?

Perhaps what we do is recursively partition, then recursively compile the final graph. Not sure if graph stitching can handle this right now.

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

narendasan · 2022-08-15T15:54:11Z

core/compiler.cpp

@@ -219,19 +219,16 @@ void AddIfBlockToGraph(
  return;
 }

-GraphAndMapping ConstructFallbackGraph(
+GraphAndMapping ConstructFallbackGraph_(


Have a more distinguishing name for this maybe?

bowang007 · 2022-08-16T01:13:22Z

@bowang007 This should be roughly equivalent to our current system. There are a couple improvements on this that I'd like to do that I want your input on.

I renamed a bunch of stuff to hopefully help make it clearer to new devs, let me know if you disagree about anything.

I want every node in the global_fallback_map (now called node_executor_decision_map) and make an explicit decision on if it will run in torch or tensorrt (this can be changed like if a block would run in trt but is not large enough). That way in debugging or otherwise I can just ask about a particular node and get back what it will run on. Right now I am having trouble doing this because there is an assumption in the system that if the node is not in the map it must run in tensorrt.

There is still work to be done on being able to easily dump state at any given time.

Are there other data sources that we might want to have managed by the context?

Went through just now, I think the general architecture is very clear and comprehensive, here are some thoughts:
2. We can have a map for every nodes in the graph, logic might be a little bit different but shouldn't be to complicated.
3. more details on this?

narendasan · 2022-08-16T04:17:20Z

@bowang007 This should be roughly equivalent to our current system. There are a couple improvements on this that I'd like to do that I want your input on.

I renamed a bunch of stuff to hopefully help make it clearer to new devs, let me know if you disagree about anything.

I want every node in the global_fallback_map (now called node_executor_decision_map) and make an explicit decision on if it will run in torch or tensorrt (this can be changed like if a block would run in trt but is not large enough). That way in debugging or otherwise I can just ask about a particular node and get back what it will run on. Right now I am having trouble doing this because there is an assumption in the system that if the node is not in the map it must run in tensorrt.

There is still work to be done on being able to easily dump state at any given time.

Are there other data sources that we might want to have managed by the context?

Went through just now, I think the general architecture is very clear and comprehensive, here are some thoughts: 2. We can have a map for every nodes in the graph, logic might be a little bit different but shouldn't be to complicated. 3. more details on this?

I want to be able to enable features like #1257
which would be able to show a user info like inputs to a block, node assignments, shape calculation results etc.

bowang007 · 2022-08-19T23:58:26Z

@narendasan If we store SegmentedBlocks

TensorRT/core/partitioning/partitioningctx/PartitioningCtx.h

Line 51 in 1a4084b

PartitionedGraph blocks;

in PartitioningCtx, what about the sub-blocks' segmentedBlocks? Looks like it could get messed up if there is some node like If::node which contains sub-blocks.

narendasan · 2022-08-25T19:13:14Z

@narendasan If we store SegmentedBlocks

TensorRT/core/partitioning/partitioningctx/PartitioningCtx.h

Line 51 in 1a4084b

PartitionedGraph blocks;

in PartitioningCtx, what about the sub-blocks' segmentedBlocks? Looks like it could get messed up if there is some node like If::node which contains sub-blocks.

Where do we put sub blocks now?

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Bo Wang <[email protected]>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Bo Wang <[email protected]>

github-actions

There are some changes that do not conform to C++ style guidelines:

diff --git a/home/runner/work/TensorRT/TensorRT/core/compiler.cpp b/tmp/changes.txt
index 178d3c4..8599be8 100644
--- a/home/runner/work/TensorRT/TensorRT/core/compiler.cpp
+++ b/tmp/changes.txt
@@ -142,7 +142,7 @@ partitioning::GraphAndMapping BuildHybridGraph(

  partitioning::Partition(&partitioning_ctx, collection_input_ivalues_map);

-  for (auto &partitioned_block : partitioning_ctx.partitioned_blocks) {
+  for (auto& partitioned_block : partitioning_ctx.partitioned_blocks) {
    partitioning::PartitionedGraph& segmented_blocks = partitioned_block.second;

    for (auto& seg_block : segmented_blocks) {
diff --git a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioning.cpp b/tmp/changes.txt
index 1a596a2..0d18f37 100644
--- a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioning.cpp
+++ b/tmp/changes.txt
@@ -73,7 +73,6 @@ void SetExplicitFallbackNodes(PartitioningCtx* ctx, torch::jit::Block* block) {
      // Set the rest nodes to TensorRt
      ctx->setNodeExecutorDecision(n, NodeExecutorDecision::kCONVERT);
    }
-
  }
  return;
}
@@ -236,9 +235,12 @@ void resolveTRTNonTensorInputs(PartitioningCtx* ctx, torch::jit::Block* block) {
        }
      }
      if (!inputs_to_resolve.empty()) {
-        std::vector<torch::jit::Node*> dependency_nodes = getDependencyNodes(inputs_to_resolve, cur_partitioned_block[i]);
+        std::vector<torch::jit::Node*> dependency_nodes =
+            getDependencyNodes(inputs_to_resolve, cur_partitioned_block[i]);
        dependency_nodes.insert(
-            dependency_nodes.end(), cur_partitioned_block[i].raw_nodes().begin(), cur_partitioned_block[i].raw_nodes().end());
+            dependency_nodes.end(),
+            cur_partitioned_block[i].raw_nodes().begin(),
+            cur_partitioned_block[i].raw_nodes().end());
        cur_partitioned_block[i] = SegmentedBlock(SegmentedBlock::kTensorRT, dependency_nodes);
      }
    }
@@ -339,7 +341,6 @@ void SegmentGraph(PartitioningCtx* ctx, torch::jit::Block* block) {

  std::vector<torch::jit::Node*> in_prog_trt_blk_nodes, in_prog_pyt_blk_nodes;
  for (const auto n : nodes) {
-
    // Skip constant nodes as they are resources for both kinds of modules
    if (n->kind() == torch::jit::prim::Constant) {
      continue;
@@ -438,7 +439,6 @@ void Partition(PartitioningCtx* ctx, ExampleIValues& example_tensor_map) {

  // Go through all the blocks to do the partitioning
  for (torch::jit::Block* block : ctx->original_blocks) {
-
    // Find all the fallback nodes and build execution decision LUT for all nodes
    SetNodeExecutorLUT(ctx, block);

@@ -453,22 +453,17 @@ void Partition(PartitioningCtx* ctx, ExampleIValues& example_tensor_map) {
    LOG_DEBUG("Registering input/output torch::jit::Value for segmented graphs");
    registerSegmentsOutputs(ctx, block);

-    for (auto &i : ctx->partitioned_blocks[block]) {
+    for (auto& i : ctx->partitioned_blocks[block]) {
      LOG_DEBUG(i);
    }

    // run shape analysis on each segmented block
    runShapeAnalysis(ctx, block, example_tensor_map);
-
  }

-
-
-//  for (uint64_t i = 0; i < ctx->blocks.size(); i++) {
-//    ctx->blocks[i].update_id(i);
-//  }
-
-
+  //  for (uint64_t i = 0; i < ctx->blocks.size(); i++) {
+  //    ctx->blocks[i].update_id(i);
+  //  }
}

} // namespace partitioning
diff --git a/home/runner/work/TensorRT/TensorRT/core/partitioning/stitching.cpp b/tmp/changes.txt
index f8a9633..42cfe9d 100644
--- a/home/runner/work/TensorRT/TensorRT/core/partitioning/stitching.cpp
+++ b/tmp/changes.txt
@@ -97,7 +97,6 @@ void AddIfBlockToGraph(
  return;
}

-
GraphAndMapping Stitch(PartitioningCtx* ctx, torch::jit::Block* block) {
  auto new_g = std::make_shared<torch::jit::Graph>();

@@ -146,8 +145,7 @@ GraphAndMapping Stitch(PartitioningCtx* ctx, torch::jit::Block* block) {
    }
  }
  return {new_g, old_to_new_g};
-
-}
-}
-}
}
+} // namespace partitioning
+} // namespace core
+} // namespace torch_tensorrt
diff --git a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioningctx/PartitioningCtx.cpp b/tmp/changes.txt
index 4b8368d..3c1db64 100644
--- a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioningctx/PartitioningCtx.cpp
+++ b/tmp/changes.txt
@@ -38,7 +38,7 @@ void PartitioningCtx::setNodeExecutorDecision(torch::jit::Node* n, NodeExecutorD
  // NOTE: This is this way due to partitioning.cpp L#134 I dont know if this is what we should do.

  auto result = node_executor_decision_map[n] = decision;
-  return ;
+  return;
}

bool PartitioningCtx::shouldNodeRunInTorch(torch::jit::Node* n) {
ERROR: Some files do not conform to style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Bo Wang <[email protected]>

github-actions

There are some changes that do not conform to C++ style guidelines:

diff --git a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioning.cpp b/tmp/changes.txt
index 86cfd6d..53b9f22 100644
--- a/home/runner/work/TensorRT/TensorRT/core/partitioning/partitioning.cpp
+++ b/tmp/changes.txt
@@ -454,7 +454,6 @@ void Partition(PartitioningCtx* ctx, ExampleIValues& example_tensor_map) {
    LOG_DEBUG("Registering input/output torch::jit::Value for segmented graphs");
    RegisterSegmentsOutputs(ctx, block);

-
    // run shape analysis on each segmented block
    RunShapeAnalysis(ctx, block, example_tensor_map);
  }
ERROR: Some files do not conform to style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Bo Wang <[email protected]>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

Signed-off-by: Dheeraj Peri <[email protected]>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

core/partitioning/BUILD

core/partitioning/partitioning.cpp

Signed-off-by: Dheeraj Peri <[email protected]>

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

narendasan added 3 commits August 12, 2022 11:15

feat(aten::hardsigmoid): Unpack hardsigmoid

d00240e

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

chore: reorg in prep for state centralizing

c9e504d

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

refactor(//core/partitioning): Centralizing partitioning state

85f8cc3

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

narendasan added the WIP Work is in progress, pull request should not be merged yet label Aug 14, 2022

facebook-github-bot added the cla signed label Aug 14, 2022

github-actions bot added component: api [C++] Issues re: C++ API component: core Issues re: The core compiler component: lowering Issues re: The lowering / preprocessing passes component: partitioning component: tests Issues re: Tests labels Aug 14, 2022

github-actions bot requested review from andi4191, bowang007 and peri044 August 14, 2022 17:12

github-actions bot approved these changes Aug 14, 2022

View reviewed changes

narendasan added the release: v1.3 Tagged to be included in v1.3 label Aug 14, 2022

narendasan commented Aug 14, 2022

View reviewed changes

chore: update usage in py api

1a4084b

Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>

github-actions bot added the component: api [Python] Issues re: Python API label Aug 14, 2022

github-actions bot approved these changes Aug 14, 2022

View reviewed changes

narendasan mentioned this pull request Aug 15, 2022

✨[Feature] Perform constant folding of operations in the graph if nodes can be evaluated at compile time #1266

Closed

narendasan commented Aug 15, 2022

View reviewed changes

narendasan mentioned this pull request Aug 24, 2022

[feat] Add dependency awareness to torch-trt partitioning #1304

Merged

7 tasks

fix: Update new info name

4cc3143

github-actions bot approved these changes Aug 26, 2022

View reviewed changes

bowang007 added 2 commits August 30, 2022 15:38

refactor: refactor the NodeExecutor logic

c3082f5

Signed-off-by: Bo Wang <[email protected]>

fix: fix typo

a165811

Signed-off-by: Bo Wang <[email protected]>

github-actions bot approved these changes Aug 30, 2022

View reviewed changes

bowang007 mentioned this pull request Sep 6, 2022

❓ [Question] How do you solve the error: Expected Tensor but got Uninitialized? #1282

Closed

refactor: extract stitching phase out of compiler.cpp

ff22707

Signed-off-by: Bo Wang <[email protected]>

github-actions bot requested changes Sep 7, 2022

View reviewed changes

github-actions bot approved these changes Sep 7, 2022

View reviewed changes

fix: fix bugs found when running tests

77f4c09

Signed-off-by: Bo Wang <[email protected]>

github-actions bot requested changes Sep 8, 2022

View reviewed changes

github-actions bot approved these changes Sep 8, 2022

View reviewed changes

fix: change the shouldNodeRunInTorch logic

34366b3

Signed-off-by: Bo Wang <[email protected]>

github-actions bot approved these changes Sep 8, 2022

View reviewed changes

chore: Rebase with master

6db9ffb

Signed-off-by: Dheeraj Peri <[email protected]>

github-actions bot approved these changes Sep 21, 2022

View reviewed changes

narendasan commented Sep 21, 2022

View reviewed changes

core/partitioning/BUILD Show resolved Hide resolved

core/partitioning/partitioning.cpp Outdated Show resolved Hide resolved

chore: address review comments

d053b4d

Signed-off-by: Dheeraj Peri <[email protected]>

peri044 removed the WIP Work is in progress, pull request should not be merged yet label Sep 22, 2022

github-actions bot approved these changes Sep 22, 2022

View reviewed changes

peri044 marked this pull request as ready for review September 22, 2022 17:51

peri044 merged commit 24172f0 into master Sep 22, 2022

github-actions bot approved these changes Sep 22, 2022

View reviewed changes

bowang007 deleted the partitioning_ctx branch December 7, 2022 04:14

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Centralizing Partitioning State #1263

Centralizing Partitioning State #1263

narendasan commented Aug 14, 2022

github-actions bot left a comment

github-actions bot left a comment

narendasan commented Aug 14, 2022

narendasan Aug 14, 2022

narendasan Aug 14, 2022

narendasan Aug 15, 2022

narendasan Aug 15, 2022

github-actions bot left a comment

github-actions bot left a comment

narendasan Aug 15, 2022

bowang007 commented Aug 16, 2022 •

edited

Loading

narendasan commented Aug 16, 2022

bowang007 commented Aug 19, 2022

narendasan commented Aug 25, 2022

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

github-actions bot left a comment

Centralizing Partitioning State #1263

Centralizing Partitioning State #1263

Conversation

narendasan commented Aug 14, 2022

Description

Type of change

Checklist:

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

narendasan commented Aug 14, 2022

narendasan Aug 14, 2022

Choose a reason for hiding this comment

narendasan Aug 14, 2022

Choose a reason for hiding this comment

narendasan Aug 15, 2022

Choose a reason for hiding this comment

narendasan Aug 15, 2022

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

narendasan Aug 15, 2022

Choose a reason for hiding this comment

bowang007 commented Aug 16, 2022 • edited Loading

narendasan commented Aug 16, 2022

bowang007 commented Aug 19, 2022

narendasan commented Aug 25, 2022

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

bowang007 commented Aug 16, 2022 •

edited

Loading