Skip to content

feat: rewriting param to a Constant if it's a introduced input #1298

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Nov 13, 2022
Merged

feat: rewriting param to a Constant if it's a introduced input #1298

Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
05bf80c
feat: rewriting param to a Constant if it's a introduced input
bowang007 Aug 20, 2022
ab977f5
fix: deal with edge cases when introduced value is Tensor with gradient
bowang007 Sep 9, 2022
f2ef0eb
refactor: refactor RewriteInputsWithParams() to a lowering pass
bowang007 Sep 14, 2022
9fecda0
chore: rebase on master branch
bowang007 Oct 13, 2022
318281d
test: add test case for rewrite input with params pass
bowang007 Oct 14, 2022
ce63baa
chore: apply linting
bowang007 Oct 14, 2022
b09994d
Merge branch 'master' into param_input
bowang007 Nov 10, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions core/compiler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

#include "torch/csrc/jit/frontend/function_schema_parser.h"
#include "torch/csrc/jit/ir/ir.h"
#include "torch/csrc/jit/ir/constants.h"
#include "torch/csrc/jit/ir/ir_views.h"
#include "torch/csrc/jit/passes/graph_fuser.h"
#include "torch/csrc/jit/passes/loop_unrolling.h"
Expand All @@ -28,6 +29,22 @@
namespace torch_tensorrt {
namespace core {

void RewriteInputsWithParams(std::shared_ptr<torch::jit::Graph> g, std::vector<torch::jit::IValue> params) {
auto input_size = g->inputs().size();
auto param_it = params.rbegin();
for (int i = input_size - 1; i >= 0; --i) {
if (g->inputs()[i]->type() != c10::TensorType::get() && g->inputs()[i]->type()->kind() != torch::jit::TypeKind::TupleType &&
g->inputs()[i]->type()->kind() != torch::jit::TypeKind::ListType && param_it != params.rend()) {
auto new_constant = torch::jit::tryInsertConstant(*g, *param_it);
++param_it;
if (new_constant) {
g->inputs()[i]->replaceAllUsesWith(*new_constant);
g->eraseInput(i);
}
}
}
}

void AddEngineToGraph(
torch::jit::script::Module mod,
std::shared_ptr<torch::jit::Graph>& g,
Expand Down Expand Up @@ -434,6 +451,9 @@ torch::jit::Module CompileGraph(const torch::jit::Module& mod, CompileSpec cfg)
(!(cfg.lower_info.forced_fallback_modules.size() == 0 &&
cfg.partition_info.forced_fallback_operators.size() == 0 && isBlockConvertible) ||
outputIsCollection)) {
if (!static_params.empty()) {
RewriteInputsWithParams(g, params);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should this be a lowering pass?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It could be a lowering pass. However, I'm not so sure about what happens with conversion stage if the input are stored in params. Any thought on this? @narendasan

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

At this point static params is usually empty, it is just legacy behavior we still exploit for test so it should be able to handle it.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cool let me change it to a lowering pass.

}
std::unordered_map<torch::jit::Node*, int> fallback_nodes;
auto collection_input_ivalues_map =
partitioning::generateRandomInputs(cfg.convert_info.collection_input_spec_map, first_use_types);
Expand Down