Merge pull request #1513 from mfeliz-cruise/michael.feliz/cast_name_conflict

[fix]Disambiguate cast layer names

Author: peri044

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<int*>(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;

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());

core/conversion/converters/impl/cast.cpp

@@ -26,7 +26,7 @@ auto cast_registrations TORCHTRT_UNUSED =
                } else {
                  trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(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<at::ScalarType>(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<int64_t>();
                auto trt_dtype = util::ScalarTypeToTRTDataType(static_cast<at::ScalarType>(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());
 

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::Graph>();
+  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,