From ecff91fb27596e75c0f6495d2a94fbac21a3883d Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Wed, 30 Nov 2022 18:09:37 -0800 Subject: [PATCH] [fix]Disambiguate cast layer names Add mechanism to add a prefix to cast layer names. This allows us to avoid name conflict errors in cases where we perform the same cast on a tensor twice either internally in ops or in cases where there is both an int32 cast and a long cast which is automatically truncated to int32. Fixes # (issue) Please delete options that are not relevant and/or add your own. - Bug fix (non-breaking change which fixes an issue) - New feature (non-breaking change which adds functionality) - Breaking change (fix or feature that would cause existing functionality to not work as expected) - This change requires a documentation update - [ ] 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 --- core/conversion/converters/converter_util.cpp | 9 +++++- core/conversion/converters/converter_util.h | 6 +++- core/conversion/converters/impl/cast.cpp | 8 +++--- .../core/conversion/converters/test_cast.cpp | 28 +++++++++++++++++++ 4 files changed, 45 insertions(+), 6 deletions(-) diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp index 36356d9db6..7c116b6c2c 100644 --- a/core/conversion/converters/converter_util.cpp +++ b/core/conversion/converters/converter_util.cpp @@ -205,7 +205,11 @@ nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor return id_out_tensor; } -nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype) { +nvinfer1::ITensor* castITensor( + ConversionCtx* ctx, + nvinfer1::ITensor* tensor, + nvinfer1::DataType dtype, + const std::string& layer_name_prefix) { if (tensor->getType() != dtype) { std::ostringstream tensor_id; tensor_id << reinterpret_cast(tensor); @@ -219,6 +223,9 @@ nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nv LOG_DEBUG(ctx->logger, "Casting ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype); std::stringstream ss; + if (layer_name_prefix.size()) { + ss << layer_name_prefix << " "; + } ss << "[Cast ITensor " << tensor_id.str() << " from " << tensor->getType() << " to " << dtype << "]"; id_layer->setName(ss.str().c_str()); return casted_tensor; diff --git a/core/conversion/converters/converter_util.h b/core/conversion/converters/converter_util.h index 245b93f498..c5e4e4eebc 100644 --- a/core/conversion/converters/converter_util.h +++ b/core/conversion/converters/converter_util.h @@ -56,7 +56,11 @@ nvinfer1::ITensor* add_abs( nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& name); // If an ITensor is of a type not dtype, add an Identity layer to cast it to dtype -nvinfer1::ITensor* castITensor(ConversionCtx* ctx, nvinfer1::ITensor* tensor, nvinfer1::DataType dtype); +nvinfer1::ITensor* castITensor( + ConversionCtx* ctx, + nvinfer1::ITensor* tensor, + nvinfer1::DataType dtype, + const std::string& layer_name_prefix = ""); // Freeze an at::Tensor in a IConstant layer nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t, const std::string& name = std::string()); diff --git a/core/conversion/converters/impl/cast.cpp b/core/conversion/converters/impl/cast.cpp index 04323d171a..7e20144123 100644 --- a/core/conversion/converters/impl/cast.cpp +++ b/core/conversion/converters/impl/cast.cpp @@ -26,7 +26,7 @@ auto cast_registrations TORCHTRT_UNUSED = } else { trt_dtype = util::ScalarTypeToTRTDataType(static_cast(output_dtype)); } - auto casted_itensor = castITensor(ctx, self, trt_dtype); + auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n)); auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor); LOG_DEBUG("[aten::to.dtype] Output tensor shape: " << output->getDimensions()); @@ -48,7 +48,7 @@ auto cast_registrations TORCHTRT_UNUSED = } else { trt_dtype = util::ScalarTypeToTRTDataType(static_cast(output_dtype)); } - auto casted_itensor = castITensor(ctx, self, trt_dtype); + auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n)); auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor); LOG_DEBUG("[aten::to.device] Output tensor shape: " << output->getDimensions()); @@ -59,7 +59,7 @@ auto cast_registrations TORCHTRT_UNUSED = [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { auto self = args[0].ITensorOrFreeze(ctx); nvinfer1::DataType other_dtype = args[1].ITensorOrFreeze(ctx)->getType(); - auto casted_itensor = castITensor(ctx, self, other_dtype); + auto casted_itensor = castITensor(ctx, self, other_dtype, util::node_info(n)); auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor); LOG_DEBUG("[aten::to.other] Output tensor shape: " << output->getDimensions()); @@ -77,7 +77,7 @@ auto cast_registrations TORCHTRT_UNUSED = auto output_dtype = args[2].unwrapToScalar().to(); auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast(output_dtype)); - auto casted_itensor = castITensor(ctx, self, trt_dtype); + auto casted_itensor = castITensor(ctx, self, trt_dtype, util::node_info(n)); auto output = ctx->AssociateValueAndTensor(n->outputs()[0], casted_itensor); LOG_DEBUG("[aten::to.prim_Device] Output tensor shape: " << output->getDimensions()); diff --git a/tests/core/conversion/converters/test_cast.cpp b/tests/core/conversion/converters/test_cast.cpp index f36cadd04b..601eefa2bf 100644 --- a/tests/core/conversion/converters/test_cast.cpp +++ b/tests/core/conversion/converters/test_cast.cpp @@ -163,6 +163,34 @@ TEST(Converters, ATenToSingleConvertsCorrectly) { ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt, 2e-6)); } +TEST(Converters, ATenToDuplicateConvertsCorrectly) { + const auto graph = R"IR( + graph(%y.1 : Tensor): + %4 : int = prim::Constant[value=3]() + %5 : bool = prim::Constant[value=0]() + %6 : None = prim::Constant() + %y0.1 : Tensor = aten::to(%y.1, %4, %5, %5, %6) + %y0.2 : Tensor = aten::to(%y.1, %4, %5, %5, %6) + return (%y0.1, %y0.2))IR"; + + auto g = std::make_shared(); + torch::jit::parseIR(graph, &*g); + + auto in = at::randint(1, 10, {3}, {at::kCUDA}); + + auto jit_in = at::clone(in); + auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {jit_in}); + + auto trt_in = at::clone(in); + params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); + auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {trt_in}); + for (size_t i = 0UL; i < jit_results.size(); ++i) { + ASSERT_TRUE(jit_results[i].scalar_type() == trt_results[i].scalar_type()); + ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[i], trt_results[i], 2e-6)); + } +} + TEST(Converters, ATenTypeAsConvertsCorrectly) { const auto graph = R"IR( graph(%0 : Tensor,