Merge pull request #249 from NVIDIA/add_alpha

Fix add ops with alpha multiplier

Author: narendasan

core/conversion/converters/impl/element_wise.cpp

@@ -79,13 +79,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto scalar = args[2].unwrapToScalar().to<float>();
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer = ctx->net->addScale(
-                          *other, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          other,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       other = scaleLayer->getOutput(0);
                     }
 
@@ -107,13 +108,14 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     auto scalar = args[2].unwrapToScalar().to<float>();
 
                     if (1 != scalar) {
-                      auto scaleW = Weights(ctx, scalar);
-                      auto unuse = Weights();
-                      // IScaleLayer assert shift, scale and power to have
-                      // the same dtype
-                      auto scaleLayer = ctx->net->addScale(
-                          *other, nvinfer1::ScaleMode::kUNIFORM, unuse.data, scaleW.data, unuse.data);
-                      TRTORCH_CHECK(scaleLayer, "Unable to create scale layer from node: " << *n);
+                      auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+                      auto scaleLayer = add_elementwise(
+                          ctx,
+                          nvinfer1::ElementWiseOperation::kPROD,
+                          other,
+                          alphaTensor,
+                          util::node_info(n) + std::string("_AlphaMultiplier"));
+                      TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
                       other = scaleLayer->getOutput(0);
                     }
 

tests/core/converters/test_element_wise.cpp

@@ -52,6 +52,30 @@ TEST(Converters, ATenAddConvertsCorrectly) {
   pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3});
 }
 
+TEST(Converters, ATenAddWithAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : float = prim::Constant[value=3.2]()
+        %3 : Tensor = aten::add(%0, %1, %2)
+        return (%3))IR";
+  pointwise_test_helper(graph, false);
+  pointwise_test_helper(graph, false, false, {3, 4}, {4});
+  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});
+}
+
+TEST(Converters, ATenAddImplicitWithAlphaConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %2 : float = prim::Constant[value=7.6]()
+        %3 : Tensor = aten::add_(%0, %1, %2)
+        return (%3))IR";
+  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});
+}
+
 TEST(Converters, ATenSubConvertsCorrectly) {
   const auto graph = R"IR(
       graph(%0 : Tensor, %1 : Tensor):
@@ -118,4 +142,4 @@ TEST(Converters, ATenAddWithScalarConvertsCorrectly) {
         %3 : Tensor = aten::add(%0, %scalar, %2)
         return (%3))IR";
   pointwise_test_helper(graph, true);
-}
+}