Skip to content

Commit

Permalink
fix: fix bugs in aten::to
Browse files Browse the repository at this point in the history
Signed-off-by: Bo Wang <bowa@nvidia.com>
  • Loading branch information
bowang007 committed Apr 1, 2022
1 parent 474e7bf commit 2ecd187
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 17 deletions.
17 changes: 17 additions & 0 deletions core/conversion/converters/impl/cast.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,23 @@ auto cast_registrations TORCHTRT_UNUSED =
return true;
}})
.pattern(
{"aten::to.device(Tensor(a) self, Device device, int dtype, bool non_blocking=False, bool copy=False, int? memory_format=None) -> (Tensor(a))",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
// what this function does is basically the same with the previous one, however, we cannot lower this
// signature to previous one because this will incur the device issues when we run Torchscript module in
// later shape analysis phase of fallback
auto self = args[0].ITensorOrFreeze(ctx);
auto output_dtype = args[2].unwrapToScalar().to<int64_t>();

auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(output_dtype));

auto casted_itensor = castITensor(ctx, self, trt_dtype);
auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor);
LOG_DEBUG("[aten::to.device] Output tensor shape: " << output->getDimensions());

return true;
}})
.pattern(
{"aten::to.other(Tensor self, Tensor other, bool non_blocking=False, bool copy=False, int? memory_format=None) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
Expand Down
21 changes: 4 additions & 17 deletions core/lowering/passes/reduce_to.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,14 @@ namespace lowering {
namespace passes {

void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph) {
std::string to_device_pattern = R"IR(
graph(%x, %device, %dtype, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %device, %dtype, %nb, %copy, %format)
return (%out))IR";
std::string to_dtype_pattern = R"IR(
graph(%x, %device, %dtype, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %dtype, %nb, %copy, %format)
return (%out))IR";
std::string to_dtype_layout_pattern = R"IR(
graph(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format)
graph(%x, %dtype, %layout, %device, %pm, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %dtype, %layout, %device, %pm, %nb, %copy, %format)
return (%out))IR";

std::string to_dtype_multi_input_pattern = R"IR(
graph(%x, %device, %dtype, %layout, %pm, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %dtype, %nb, %copy, %format)
graph(%x, %dtype, %layout, %device, %pm, %nb, %copy, %format):
%out : Tensor = aten::to(%x, %device, %dtype, %nb, %copy, %format)
return (%out))IR";

std::string to_type_as_pattern = R"IR(
Expand All @@ -38,11 +30,6 @@ void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph) {
%out : Tensor = aten::to(%input, %other, %5, %5, %6)
return (%out))IR";

// replace aten::to.device with aten::to.dtype
torch::jit::SubgraphRewriter map_aten_device_to_dtype;
map_aten_device_to_dtype.RegisterRewritePattern(to_device_pattern, to_dtype_pattern);
map_aten_device_to_dtype.runOnGraph(graph);

// replace aten::to.dtype_layout with aten::to.dtype
torch::jit::SubgraphRewriter map_aten_dtype_layout;
map_aten_dtype_layout.RegisterRewritePattern(to_dtype_layout_pattern, to_dtype_multi_input_pattern);
Expand Down

0 comments on commit 2ecd187

Please sign in to comment.