Merge pull request #320 from NVIDIA/neg_index_sum

Support -1 indexing for aten::sum

Author: narendasan

core/conversion/converters/impl/reduce.cpp

@@ -80,11 +80,23 @@ auto reduce_registrations TRTORCH_UNUSED =
              [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                auto in_tensor = args[0].ITensorOrFreeze(ctx);
                auto dims = args[1].unwrapToIntList();
-               LOG_DEBUG("Dim to reduce:" << util::toDims(dims)); // Some abuse of toDim but just for debug info
+               c10::List<int64_t> calculated_dims;
+               auto in_dims = util::toVec(in_tensor->getDimensions());
+               LOG_DEBUG("InDims " << in_dims); // Some abuse of toDim but just for debug info
+               LOG_DEBUG(
+                   "Dim to reduce(original):" << util::toDims(dims)); // Some abuse of toDim but just for debug info
+               for (size_t i = 0; i < dims.size(); i++) {
+                 auto dim_val = dims[i] < 0 ? (in_dims.size() + dims[i]) : dims[i];
+                 calculated_dims.push_back(dim_val);
+               }
+
+               LOG_DEBUG(
+                   "Dim to reduce(converted):"
+                   << util::toDims(calculated_dims)); // Some abuse of toDim but just for debug info
 
                uint32_t axis_mask = 0;
-               for (size_t d = 0; d < dims.size(); d++) {
-                 axis_mask |= 1 << dims[d];
+               for (size_t d = 0; d < calculated_dims.size(); d++) {
+                 axis_mask |= 1 << calculated_dims[d];
                }
                LOG_DEBUG("Axis Mask" << std::bitset<32>(axis_mask));
 

tests/core/conversion/converters/test_reduce.cpp

@@ -134,6 +134,58 @@ converts_keepdims_correctly(mean, Mean);
 
 #undef converts_keepdims_correctly
 
+TEST(Converters, ATenSumDimNegOneIndexConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-1]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=0]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegOneIndexKeepDimsConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-1]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=1]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegIndexConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-2]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=0]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
+TEST(Converters, ATenSumDimNegIndexKeepDimsConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%0 : Tensor):
+      %1 : int = prim::Constant[value=-2]()
+      %2 : int[] = prim::ListConstruct(%1)
+      %3 : bool = prim::Constant[value=1]()
+      %4 : None = prim::Constant()
+      %5 : Tensor = aten::sum(%0, %2, %3, %4)
+      return (%5))IR";
+  auto in = at::randint(-5, 5, {4, 4, 4}, at::kCUDA);
+  test_body(graph, in);
+}
+
 TEST(Converters, ATenProdDimConvertsCorrectly) {
   const auto graph = R"IR(
     graph(%0 : Tensor):