Skip to content
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

Fix arange #839

Merged
merged 2 commits into from
Feb 2, 2022
Merged

Fix arange #839

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
20e5d41
fix(//core/conversion/evaluators): Change how schemas are handled
narendasan Jan 31, 2022
096a5d4
refactor(//core/partitioning): Refactor partitioning
narendasan Jan 31, 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: 8 additions & 12 deletions core/conversion/evaluators/aten.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -620,22 +620,19 @@ auto aten_registrations TORCHTRT_UNUSED =
{"aten::tensor(t[] data, *, int? dtype=None, Device? device=None, bool requires_grad=False) -> (Tensor)"})})
.evaluator({c10::Symbol::fromQualString("aten::arange"),
[](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
int input_size = n->inputs().size();
int scalar_count = 0;
for (int i = 0; i < input_size; i++) {
if (args.at(n->input(i)).IValue()->isScalar()) {
scalar_count += 1;
}
}
if (scalar_count == 1) {
auto schema = n->maybeSchema();
TORCHTRT_CHECK(schema, "Unable to get schema for node: " << *n);
auto name = schema->operator_name();

if (c10::toString(name) == "aten::arange") {
if (args.at(n->input(0)).IValue()->isInt()) {
int end_scalar = args.at(n->input(0)).unwrapToInt();
return torch::arange(end_scalar);
} else if (args.at(n->input(0)).IValue()->isDouble()) {
float end_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
return torch::arange(end_scalar);
}
} else if (scalar_count == 2) {
} else if (c10::toString(name) == "aten::arange.start") {
if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble()) {
float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
Expand All @@ -645,7 +642,7 @@ auto aten_registrations TORCHTRT_UNUSED =
int end_scalar = args.at(n->input(1)).unwrapToInt();
return torch::arange(start_scalar, end_scalar);
}
} else if (scalar_count == 3) {
} else if (c10::toString(name) == "aten::arange.start_step") {
if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble() ||
args.at(n->input(2)).IValue()->isDouble()) {
float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
Expand All @@ -659,8 +656,7 @@ auto aten_registrations TORCHTRT_UNUSED =
return torch::arange(start_scalar, end_scalar, step_scalar);
}
} else {
TORCHTRT_THROW_ERROR(
"Invalid input argument size for aten::arange, input argument size: " << input_size);
TORCHTRT_THROW_ERROR("Unsupported aten::arange variant: " << name);
}
return {};
},
Expand Down
227 changes: 115 additions & 112 deletions core/partitioning/SegmentedBlock.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,113 +1,116 @@
#pragma once

#include <ostream>
#include <vector>

#include "NvInfer.h"
#include "core/ir/ir.h"
#include "core/partitioning/PartitionInfo.h"
#include "torch/csrc/jit/ir/ir.h"

namespace torch_tensorrt {
namespace core {
namespace partitioning {

struct SegmentedBlock {
public:
enum SegmentedBlockTarget {
kTorch,
kTensorRT,
};

static std::string target_to_str(SegmentedBlockTarget t) {
if (t == SegmentedBlockTarget::kTorch) {
return "Torch";
} else {
return "TensorRT";
}
}

using BlockID = uint64_t;

SegmentedBlock() = default;
SegmentedBlock(SegmentedBlockTarget blk_target) : target_(blk_target), g_(std::make_shared<torch::jit::Graph>()) {}
SegmentedBlock(SegmentedBlockTarget blk_target, const std::vector<torch::jit::Node*>& nodes);
SegmentedBlock(SegmentedBlockTarget blk_target, std::shared_ptr<torch::jit::Graph> g) : target_(blk_target), g_(g) {}
SegmentedBlock(BlockID id, SegmentedBlockTarget blk_target, const std::vector<torch::jit::Node*>& nodes);

torch::jit::Value* getOrAddInputForValue(torch::jit::Value* v);
torch::jit::Node* cloneNode(torch::jit::Node* node);
void appendNode(torch::jit::Node* n) {
cloneNode(n);
}
void registerOutput(torch::jit::Value* raw_output);
torch::jit::graph_node_list nodes() {
return g_->nodes();
}
const std::vector<torch::jit::Node*>& raw_nodes() const {
return nodes_;
}
torch::jit::Block* block() {
return g_->block();
}
std::shared_ptr<torch::jit::Graph>& g() {
return g_;
}
void update_graph(std::shared_ptr<torch::jit::Graph> new_g) {
g_ = new_g;
}
c10::ArrayRef<torch::jit::Value*> inputs() {
return g_->inputs();
}
c10::ArrayRef<torch::jit::Value*> outputs() {
return g_->outputs();
}
const std::vector<torch::jit::Value*>& raw_inputs() const {
return inputs_;
}
const std::vector<torch::jit::Value*>& raw_outputs() const {
return outputs_;
}
void eraseInput(size_t i);
void eraseOutput(size_t i);
bool contain_raw_value(torch::jit::Value* input) {
return old_to_new_.count(input);
}
void register_inshapes(std::vector<ir::Input>& in_shapes) {
in_shapes_ = in_shapes;
}
const std::vector<ir::Input>& in_shapes() const {
return in_shapes_;
}
void register_intypes(std::vector<at::ScalarType>& in_types) {
in_types_ = in_types;
}
const std::vector<at::ScalarType>& in_types() const {
return in_types_;
}
void update_target(SegmentedBlockTarget new_target) {
target_ = new_target;
}
enum SegmentedBlockTarget target() {
return target_;
}

friend std::ostream& operator<<(std::ostream& os, const SegmentedBlock& b);

private:
BlockID id_;
SegmentedBlockTarget target_;
std::vector<ir::Input> in_shapes_;
std::vector<at::ScalarType> in_types_;
std::vector<torch::jit::Value*> inputs_;
std::vector<torch::jit::Value*> outputs_;
std::vector<torch::jit::Node*> nodes_;
std::shared_ptr<torch::jit::Graph> g_;
std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_;
};

std::ostream& operator<<(std::ostream& os, const SegmentedBlock::SegmentedBlockTarget& t);

} // namespace partitioning
} // namespace core
#pragma once

#include <ostream>
#include <vector>

#include "NvInfer.h"
#include "core/ir/ir.h"
#include "core/partitioning/PartitionInfo.h"
#include "torch/csrc/jit/ir/ir.h"

namespace torch_tensorrt {
namespace core {
namespace partitioning {

struct SegmentedBlock {
public:
enum SegmentedBlockTarget {
kTorch,
kTensorRT,
};

static std::string target_to_str(SegmentedBlockTarget t) {
if (t == SegmentedBlockTarget::kTorch) {
return "Torch";
} else {
return "TensorRT";
}
}

using BlockID = uint64_t;

SegmentedBlock() = default;
SegmentedBlock(SegmentedBlockTarget blk_target) : target_(blk_target), g_(std::make_shared<torch::jit::Graph>()) {}
SegmentedBlock(SegmentedBlockTarget blk_target, const std::vector<torch::jit::Node*>& nodes);
SegmentedBlock(SegmentedBlockTarget blk_target, std::shared_ptr<torch::jit::Graph> g) : target_(blk_target), g_(g) {}
SegmentedBlock(BlockID id, SegmentedBlockTarget blk_target, const std::vector<torch::jit::Node*>& nodes);

torch::jit::Value* getOrAddInputForValue(torch::jit::Value* v);
torch::jit::Node* cloneNode(torch::jit::Node* node);
void appendNode(torch::jit::Node* n) {
cloneNode(n);
}
void registerOutput(torch::jit::Value* raw_output);
torch::jit::graph_node_list nodes() {
return g_->nodes();
}
const std::vector<torch::jit::Node*>& raw_nodes() const {
return nodes_;
}
torch::jit::Block* block() {
return g_->block();
}
std::shared_ptr<torch::jit::Graph>& g() {
return g_;
}
void update_graph(std::shared_ptr<torch::jit::Graph> new_g) {
g_ = new_g;
}
c10::ArrayRef<torch::jit::Value*> inputs() {
return g_->inputs();
}
c10::ArrayRef<torch::jit::Value*> outputs() {
return g_->outputs();
}
const std::vector<torch::jit::Value*>& raw_inputs() const {
return inputs_;
}
const std::vector<torch::jit::Value*>& raw_outputs() const {
return outputs_;
}
void eraseInput(size_t i);
void eraseOutput(size_t i);
bool contain_raw_value(torch::jit::Value* input) {
return old_to_new_.count(input);
}
void register_inshapes(std::vector<ir::Input>& in_shapes) {
in_shapes_ = in_shapes;
}
const std::vector<ir::Input>& in_shapes() const {
return in_shapes_;
}
void register_intypes(std::vector<at::ScalarType>& in_types) {
in_types_ = in_types;
}
const std::vector<at::ScalarType>& in_types() const {
return in_types_;
}
void update_id(BlockID new_id) {
id_ = new_id;
}
void update_target(SegmentedBlockTarget new_target) {
target_ = new_target;
}
enum SegmentedBlockTarget target() {
return target_;
}

friend std::ostream& operator<<(std::ostream& os, const SegmentedBlock& b);

private:
BlockID id_;
SegmentedBlockTarget target_;
std::vector<ir::Input> in_shapes_;
std::vector<at::ScalarType> in_types_;
std::vector<torch::jit::Value*> inputs_;
std::vector<torch::jit::Value*> outputs_;
std::vector<torch::jit::Node*> nodes_;
std::shared_ptr<torch::jit::Graph> g_;
std::unordered_map<torch::jit::Value*, torch::jit::Value*> old_to_new_;
};

std::ostream& operator<<(std::ostream& os, const SegmentedBlock::SegmentedBlockTarget& t);

} // namespace partitioning
} // namespace core
} // namespace torch_tensorrt
Loading