Centralizing Partitioning State #1263
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
narendasan
added
the
WIP
github-actions
bot
added
component: api [C++]
|
@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.
|
narendasan
commented
Aug 14, 2022
| 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"; |
narendasan
commented
Aug 14, 2022
core/compiler.cpp
Outdated
| @@ -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_( | |||
There was a problem hiding this comment.
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
There was a problem hiding this comment.
@bowang007 thoughts here?
There was a problem hiding this comment.
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]>
narendasan
commented
Aug 15, 2022
core/compiler.cpp
Outdated
| @@ -219,19 +219,16 @@ void AddIfBlockToGraph( | |||
| return; | |||
| } | |||
|
|
|||
| GraphAndMapping ConstructFallbackGraph( | |||
| GraphAndMapping ConstructFallbackGraph_( | |||
There was a problem hiding this comment.
Have a more distinguishing name for this maybe?
Went through just now, I think the general architecture is very clear and comprehensive, here are some thoughts: |
I want to be able to enable features like #1257 |
|
@narendasan If we store SegmentedBlocks 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. |
7 tasks
Where do we put sub blocks now? |
Signed-off-by: Bo Wang <[email protected]>
Signed-off-by: Bo Wang <[email protected]>
Signed-off-by: Bo Wang <[email protected]>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 guidelinesSigned-off-by: Bo Wang <[email protected]>
There was a problem hiding this comment.
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 guidelinesSigned-off-by: Bo Wang <[email protected]>
Signed-off-by: Dheeraj Peri <[email protected]>
narendasan
commented
Sep 21, 2022
Signed-off-by: Dheeraj Peri <[email protected]>
peri044
removed
the
WIP
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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 toConversionCtxcentralizes 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.
Checklist: