Add automatic type promotion to element-wise ops

Adds automatic type promotion to match the default torch-script behavior for element-wise ops. Debug messages added for the type mismatch and cast.

Messages written to log.
```
DEBUG: [Torch-TensorRT] - Type mismatch for inputs in element-wise operation %3 : Tensor = aten::add(%0, %1, %2): Int32, Float32
DEBUG: [Torch-TensorRT] - Element-wise op type promotion adding cast from Int32 to Float32 for layer %3 : Tensor = aten::add(%0, %1, %2)
```
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

Signed-off-by: Michael Feliz <[email protected]>

Author: mfeliz-cruise

core/conversion/converters/converter_util.cpp

@@ -59,6 +59,14 @@ nvinfer1::ITensor* addUnpadding(
   }
 }
 
+nvinfer1::DataType promote_types(nvinfer1::DataType type_a, nvinfer1::DataType type_b){
+  auto torch_type_a = util::TRTDataTypeToScalarType(type_a);
+  auto torch_type_b = util::TRTDataTypeToScalarType(type_b);
+  auto promo_type = at::promote_types(torch_type_a, torch_type_b);
+  auto trt_promo_type = util::ScalarTypeToTRTDataType(promo_type);
+  return trt_promo_type;
+}
+
 nvinfer1::ILayer* add_elementwise(
     ConversionCtx* ctx,
     nvinfer1::ElementWiseOperation op,
@@ -71,6 +79,20 @@ nvinfer1::ILayer* add_elementwise(
     std::swap(self, other);
     swapSelfOther = true;
   }
+
+  if(self->getType() != other->getType()){
+    LOG_DEBUG("Type mismatch for inputs in element-wise operation " << name << ": " << self->getType() << ", " << other->getType());
+    auto promo_type = promote_types(self->getType(), other->getType());
+    if(self->getType() != promo_type){
+      LOG_DEBUG("Element-wise op type promotion adding cast from " << self->getType() << " to " << promo_type << " for layer " << name);
+      self = castITensor(ctx, self, promo_type);
+    }
+    if(other->getType() != promo_type){
+      LOG_DEBUG("Element-wise op type promotion adding cast from " << other->getType() << " to " << promo_type << " for layer " << name);
+      other = castITensor(ctx, other, promo_type);
+    }
+  }
+
   auto selfDim = util::toVec(self->getDimensions());
   auto otherDim = util::toVec(other->getDimensions());
   if (selfDim.size() != otherDim.size()) {

core/conversion/converters/impl/element_wise.cpp

@@ -54,10 +54,10 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                     // Should implement self + alpha * other
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto other = args[1].ITensorOrFreeze(ctx);
-                    auto scalar = args[2].unwrapToScalar().to<float>();
+                    auto scalar = args[2].unwrapToScalar();
 
-                    if (1 != scalar) {
-                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    if (1 != scalar.to<float>()) {
+                      auto alphaTensor = impl::scalar_to_tensor(ctx, scalar);
                       auto scaleLayer = add_elementwise(
                           ctx,
                           nvinfer1::ElementWiseOperation::kPROD,
@@ -83,10 +83,10 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                     // Should implement self + alpha * other
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto other = args[1].ITensorOrFreeze(ctx);
-                    auto scalar = args[2].unwrapToScalar().to<float>();
+                    auto scalar = args[2].unwrapToScalar();
 
-                    if (1 != scalar) {
-                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    if (1 != scalar.to<float>()) {
+                      auto alphaTensor = impl::scalar_to_tensor(ctx, scalar);
                       auto scaleLayer = add_elementwise(
                           ctx,
                           nvinfer1::ElementWiseOperation::kPROD,
@@ -257,12 +257,11 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     // Should implement other - alpha * self
                     auto self = args[0].ITensorOrFreeze(ctx);
-                    auto otherScalar = args[1].unwrapToScalar().to<float>();
-                    auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
-                    auto scalar = args[2].unwrapToScalar().to<float>();
+                    auto other = impl::scalar_to_tensor(ctx, args[1].unwrapToScalar());
+                    auto scalar = args[2].unwrapToScalar();
 
-                    if (1 != scalar) {
-                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    if (1 != scalar.to<float>()) {
+                      auto alphaTensor = impl::scalar_to_tensor(ctx, scalar);
                       auto scaleLayer = add_elementwise(
                           ctx,
                           nvinfer1::ElementWiseOperation::kPROD,
@@ -287,10 +286,10 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                     // Should implement other - alpha * self
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto other = args[1].ITensorOrFreeze(ctx);
-                    auto scalar = args[2].unwrapToScalar().to<float>();
+                    auto scalar = args[2].unwrapToScalar();
 
-                    if (1 != scalar) {
-                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                    if (1 != scalar.to<float>()) {
+                      auto alphaTensor = impl::scalar_to_tensor(ctx, scalar);
                       auto scaleLayer = add_elementwise(
                           ctx,
                           nvinfer1::ElementWiseOperation::kPROD,

tests/core/conversion/converters/test_element_wise.cpp

@@ -12,26 +12,31 @@ void pointwise_test_helper(
     std::vector<int64_t> shape1 = {5},
     std::vector<int64_t> shape2 = {5},
     bool negative_input = false,
-    bool int_tensors = false) {
+    at::ScalarType type1 = at::kFloat,
+    at::ScalarType type2 = at::kFloat) {
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph_ir, g.get());
 
   // singleInput case is enabled when elementwise operation is performed
   // with an input and a constant embedded in graph
   std::vector<at::Tensor> torch_inputs;
-  if (negative_input) {
-    torch_inputs.push_back(at::randint(-5, 5, shape1, {at::kCUDA}));
-  } else {
-    torch_inputs.push_back(at::randint(1, 5, shape1, {at::kCUDA}));
+  int first_min = negative_input ? -5 : 1;
+  int first_max = 5;
+  int second_min = 1;
+  int second_max = 5;
+  if(type1 == at::kBool){
+    first_min = 0;
+    first_max = 1;
   }
-  if (!singleInput) {
-    torch_inputs.push_back(at::randint(1, 5, shape2, {at::kCUDA}));
+  if(type2 == at::kBool){
+    second_min = 0;
+    second_max = 1;
   }
-  if(int_tensors){
-    for(size_t i = 0UL; i < torch_inputs.size(); ++i){
-      torch_inputs[i] = torch_inputs[i].to(at::kInt);
-    }
+  torch_inputs.push_back(at::randint(first_min, first_max, shape1, at::TensorOptions(at::kCUDA).dtype(type1)));
+  if (!singleInput) {
+    torch_inputs.push_back(at::randint(second_min, second_max, shape2, at::TensorOptions(at::kCUDA).dtype(type2)));
   }
+
   auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
   auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, torch_inputs);
 
@@ -62,6 +67,13 @@ TEST(Converters, ATenAddConvertsCorrectly) {
   pointwise_test_helper(graph, false, false, {4}, {3, 4});
   pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
   pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kInt, at::kInt);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kInt, at::kFloat);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kBool, at::kInt);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kBool, at::kFloat);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kInt, at::kBool);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kFloat, at::kBool);
 }
 
 TEST(Converters, ATenAddWithAlphaConvertsCorrectly) {
@@ -86,6 +98,17 @@ TEST(Converters, ATenAddImplicitWithAlphaConvertsCorrectly) {
   pointwise_test_helper(graph, false);
   pointwise_test_helper(graph, false, false, {3, 4}, {4});
   pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
+  pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kFloat, at::kInt);
+}
+
+TEST(Converters, ATenAddImplicitWithIntAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : int = prim::Constant[value=42]()
+        %3 : Tensor = aten::add_(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false, false, {2, 2}, {2, 2}, false, at::kInt, at::kInt);
+  pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kInt, at::kInt);
 }
 
 TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
@@ -138,7 +161,7 @@ TEST(Converters, ATenMulWithIntScalarConvertsCorrectly) {
         %scalar : int = prim::Constant[value=2]()
         %1 : Tensor = aten::mul(%0, %scalar)
         return (%1))IR";
-  pointwise_test_helper(graph, true, false, {5}, {5}, false, true);
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
 }
 
 TEST(Converters, ATenDivConvertsCorrectly) {
@@ -151,6 +174,8 @@ TEST(Converters, ATenDivConvertsCorrectly) {
   pointwise_test_helper(graph, false, false, {4}, {3, 4});
   pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3});
   pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kFloat, at::kInt);
+  pointwise_test_helper(graph, false, false, {4, 3}, {3, 4, 3}, false, at::kInt, at::kFloat);
 }
 
 TEST(Converters, ATenDivWithScalarConvertsCorrectly) {
@@ -295,6 +320,8 @@ TEST(Converters, ATenRsubWithTensorConvertsCorrectly) {
   pointwise_test_helper(graph, false, false, {3, 4}, {4});
   pointwise_test_helper(graph, false, false, {4}, {3, 4});
   pointwise_test_helper(graph, false, true, {4, 3, 3, 3}, {4, 3, 3, 3});
+  pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kFloat);
+  pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kInt);
 }
 
 TEST(Converters, ATenRsubWithScalarConvertsCorrectly) {
@@ -307,6 +334,46 @@ TEST(Converters, ATenRsubWithScalarConvertsCorrectly) {
   pointwise_test_helper(graph, true, false, {4, 3, 3, 3});
 }
 
+TEST(Converters, ATenRsubWithIntScalarConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=2]()
+        %scalar : int = prim::Constant[value=8]()
+        %3 : Tensor = aten::rsub(%0, %scalar, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, true, false, {4, 3, 3, 3}, {}, false, at::kInt);
+}
+
+TEST(Converters, ATenClipMinConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=1.5]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clip(%x.1, %2, %3)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClipMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=3.5]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clip(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClipMinMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=3.5]()
+          %3 : float = prim::Constant[value=1.5]()
+          %4 : Tensor = aten::clip(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
 TEST(Converters, ATenClampMinConvertsCorrectly) {
   const auto graph = R"IR(
   graph(%x.1 : Tensor):
@@ -337,6 +404,36 @@ TEST(Converters, ATenClampMinMaxConvertsCorrectly) {
   pointwise_test_helper(graph, true);
 }
 
+TEST(Converters, ATenClampIntMinConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : int = prim::Constant[value=1]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clamp(%x.1, %2, %3)
+          return (%4))IR";
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
+}
+
+TEST(Converters, ATenClampIntMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : int = prim::Constant[value=3]()
+          %3 : None = prim::Constant()
+          %4 : Tensor = aten::clamp(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
+}
+
+TEST(Converters, ATenClampIntMinMaxConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : int = prim::Constant[value=3]()
+          %3 : int = prim::Constant[value=1]()
+          %4 : Tensor = aten::clamp(%x.1, %3, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt);
+}
+
 TEST(Converters, ATenClampMinimumConvertsCorrectly) {
   const auto graph = R"IR(
   graph(%x.1 : Tensor):
@@ -487,4 +584,4 @@ TEST(Converters, ATenRemainderWithScalarConvertsCorrectly) {
   auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
 
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
-}
+}