[core IR] Add decomps for aten.sum and aten.squeeze variants (#110645)

Stephen Jia · facebook-github-bot · commit 382eeca89b06 · 2023-10-06T11:44:40.000-07:00
Summary:

## Context

Both `aten.sum` and `aten.squeeze` have a "most generic" variant in the form of `aten.sum.dim_IntList` and `aten.squeeze.dims` respectively. Add decompositions for other non generic variants of these operators to express them using the most generic variant.

Note that to register these decomps, the reference implementation under `_refs` had to be removed as registered decompositions.

Test Plan: Github CI + Meta Internal CI

Reviewed By: manuelcandales

Differential Revision: D49965952
diff --git a/test/expect/HasDecompTest.test_aten_core_operators.expect b/test/expect/HasDecompTest.test_aten_core_operators.expect
@@ -462,19 +462,15 @@ aten::split_with_sizes
 aten::sqrt
 aten::sqrt.out
 aten::sqrt_
-aten::squeeze
-aten::squeeze.dim
 aten::squeeze.dims
 aten::sub.Scalar
 aten::sub.Scalar_out
 aten::sub.Tensor
 aten::sub.out
 aten::sub_.Scalar
 aten::sub_.Tensor
-aten::sum
 aten::sum.IntList_out
 aten::sum.dim_IntList
-aten::sum.out
 aten::tan
 aten::tan.out
 aten::tan_
diff --git a/torch/_decomp/__init__.py b/torch/_decomp/__init__.py
@@ -404,9 +404,13 @@ def core_aten_decompositions() -> Dict[torch._ops.OperatorBase, Callable]:
             aten.special_log_ndtr,
             aten.special_xlog1py,
             aten.split.Tensor,
+            aten.squeeze.default,
+            aten.squeeze.dim,
             aten.std,
             aten.std_mean,
             aten.stack,
+            aten.sum.default,
+            aten.sum.out,
             aten.t,
             aten.tanh_backward,
             aten.threshold,
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
@@ -4211,6 +4211,27 @@ def floor_divide(self, other):
     return torch.div(self, other, rounding_mode="floor")
 
 
+@register_decomposition([aten.sum.default, aten.sum.out])
+def sum_default(
+    self: Tensor,
+    *,
+    dtype: Optional[torch.dtype] = None,
+    out: Optional[Tensor] = None,
+) -> Tensor:
+    if out is None:
+        return aten.sum.dim_IntList(self, [], dtype=dtype)
+    else:
+        return aten.sum.IntList_out(self, [], dtype=dtype, out=out)
+
+
+@register_decomposition([aten.squeeze.default, aten.squeeze.dim])
+def squeeze_default(self: Tensor, dim: Optional[int] = None):
+    if dim is None:
+        return aten.squeeze.dims(self, list(range(self.dim())))
+    else:
+        return aten.squeeze.dims(self, [dim])
+
+
 register_inplace(aten.addbmm_, aten.addbmm)
 register_inplace(aten.addmm_, aten.addmm)
 register_inplace(aten.addmv_, aten.addmv)
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -73,6 +73,8 @@
     aten.clamp_min,
     aten.glu,  # inductor lowers this directly
     aten.split.Tensor,  # inductor lowers this directly
+    aten.squeeze,  # inductor lowers this directly
+    aten.sum,  # inductor lowers this directly
     aten.unbind,  # inductor lowers this directly
 ]
 
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -2266,7 +2266,7 @@ def any(
     return result
 
 
-@register_decomposition(aten.sum)
+@register_decomposition([aten.sum.dim_IntList, aten.sum.IntList_out])
 def sum(
     a: TensorLikeType,
     dim: Union[Optional[int], Optional[List[int]]] = None,
@@ -3987,7 +3987,7 @@ def index_select(x: TensorLike, dim: int, index: TensorLike):
     return x[idx]
 
 
-@register_decomposition(aten.squeeze)
+@register_decomposition(aten.squeeze.dims)
 def squeeze(a: TensorLikeType, dim: Optional[DimsType] = None) -> TensorLikeType:
     if dim is None:
         dims = tuple(idx for idx, size in enumerate(a.shape) if size == 1)
