add aten.sum.default (#2807)

Summary:

The operator `aten.sum.dim_IntList` could take an empty list as the parameter for `dims`. We modify `vulkan_graph_builder.py` to accommodate the empty list.

Moreover, the op `aten.sum.default` is implemented as a [decomposition](https://www.internalfb.com/code/fbsource/[96e496f9db8f92967b4394bd4f60e39ab916740b]/xplat/caffe2/torch/_decomp/decompositions.py?lines=4676) into `aten.sum.dim_IntList` with empty `dims`. So we will support `aten.sum.default` with the changes.

Context: `torch.sum(x, ())` and `torch.sum(x)` are two ways to compute the sum of all elements in tensor `x`.

Differential Revision: D55630993

Author: copyrightly

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -48,6 +48,7 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.max_pool2d_with_indices.default,
             # Sum
             exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.aten.sum.default,
             # Other
             operator.getitem,
         ]

backends/vulkan/runtime/graph/ops/impl/Sum.cpp

@@ -123,10 +123,17 @@ void add_sum_dim_IntList(
   auto dims_to_sum = graph.get_val(opt_dim).toIntList();
   int64_t in_dim = in_tensor.sizes().size();
 
-  for (const auto& dim : dims_to_sum) {
-    // Normalize (negative) dim into range [0, self.dim() - 1]
-    int64_t dim_normalized = normalize(dim, in_dim);
-    dims_set.insert(dim_normalized);
+  if (dims_to_sum.empty()) {
+    // If dim is not specified, reduce over all dims
+    for (int64_t i = 0; i < in_dim; ++i) {
+      dims_set.insert(i);
+    }
+  } else {
+    for (const auto& dim : dims_to_sum) {
+      // Normalize (negative) dim into range [0, self.dim() - 1]
+      int64_t dim_normalized = normalize(dim, in_dim);
+      dims_set.insert(dim_normalized);
+    }
   }
 
   // Reduce the higher dimensionalities first, otherwise when keepdim is

backends/vulkan/serialization/vulkan_graph_builder.py

@@ -178,7 +178,11 @@ def create_tensor_value(self, spec: TensorSpec, constant_id: int = -1) -> int:
 
     def create_scalar_list_value(self, arg: List[_ScalarType]) -> int:
         new_id = len(self.values)
-        if isinstance(arg[0], bool):
+        if len(arg) == 0:
+            self.values.append(
+                vk_graph_schema.VkValue(vk_graph_schema.IntList(items=[]))
+            )
+        elif isinstance(arg[0], bool):
             self.values.append(
                 vk_graph_schema.VkValue(
                     vk_graph_schema.BoolList(items=[cast(bool, e) for e in arg])
@@ -227,7 +231,9 @@ def get_or_create_value_for(self, arg: _Argument):
             return self.create_scalar_value(arg)
         elif isinstance(arg, TensorSpec):
             return self.create_tensor_value(arg)
-        elif isinstance(arg, list) and isinstance(arg[0], _ScalarType):
+        elif isinstance(arg, list) and (
+            len(arg) == 0 or isinstance(arg[0], _ScalarType)
+        ):
             # pyre-ignore[6]
             return self.create_scalar_list_value(arg)
         elif isinstance(arg, list) and isinstance(arg[0], Node):

backends/vulkan/test/test_vulkan_delegate.py

@@ -496,3 +496,23 @@ def forward(self, x):
             sample_inputs,
             memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
         )
+
+    def test_vulkan_backend_sum(self):
+        class SumModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            # test both torch.sum(x, ()) and torch.sum(x)
+            def forward(self, x):
+                x = torch.sum(x, (), keepdim=True)
+                x = torch.sum(x)
+                return x
+
+        module = SumModule()
+        sample_inputs = (torch.rand(size=(3, 2, 7, 5), dtype=torch.float32),)
+
+        self.lower_module_and_test_output(
+            module,
+            sample_inputs,
+            memory_layouts=[vk_graph_schema.VkMemoryLayout.TENSOR_CHANNELS_PACKED],
+        )