Add support for dim=None to AtenMeanDimOp

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -3879,10 +3879,10 @@ def Torch_AtenMeanDimOp : Torch_Op<"aten.mean.dim", [
     HasValueSemantics,
     ReadOnly
   ]> {
-  let summary = "Generated op for `aten::mean.dim : (Tensor, int[], bool, int?) -> (Tensor)`";
+  let summary = "Generated op for `aten::mean.dim : (Tensor, int[]?, bool, int?) -> (Tensor)`";
   let arguments = (ins
     AnyTorchTensorType:$self,
-    AnyTorchListOfTorchIntType:$dim,
+    AnyTorchOptionalListOfTorchIntType:$dim,
     Torch_BoolType:$keepdim,
     AnyTorchOptionalIntType:$dtype
   );

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -1053,10 +1053,11 @@ class DecomposeAtenMeanDimOp : public OpRewritePattern<AtenMeanDimOp> {
           op, "only floating-point type is supported");
     }
 
-    auto dimListConstruct = dimList.getDefiningOp<PrimListConstructOp>();
-    if (!dimListConstruct) {
+    SmallVector<Value> dimListElements;
+    if (!getListConstructElements(dimList, dimListElements) &&
+        !dimList.getType().isa<Torch::NoneType>()) {
       return rewriter.notifyMatchFailure(
-          op, "expect dimList to be constructed from list construct");
+          op, "expected `dim` to be `None` or constructed from list construct");
     }
 
     // Compute sum along dimensions specified in `dimList`.
@@ -1066,12 +1067,12 @@ class DecomposeAtenMeanDimOp : public OpRewritePattern<AtenMeanDimOp> {
     // `productDimSize` is product of sizes of dimensions to be reduced.
     Value productDimSize;
     // Case: Reduce along all dims.
-    if (dimListConstruct.elements().empty() && inputRank != 0) {
+    if (dimListElements.empty() && inputRank != 0) {
       productDimSize = rewriter.create<AtenNumelOp>(loc, input);
     } else {
       productDimSize = rewriter.create<Torch::ConstantIntOp>(
           loc, rewriter.getI64IntegerAttr(1));
-      for (Value dim : dimListConstruct.elements()) {
+      for (Value dim : dimListElements) {
         Value dimSize = rewriter.create<AtenSizeIntOp>(loc, input, dim);
         productDimSize =
             rewriter.create<AtenMulIntOp>(loc, productDimSize, dimSize);

lib/Dialect/Torch/Transforms/ShapeLibrary.cpp

@@ -5678,11 +5678,19 @@ module {
     %1 = torch.prim.TupleConstruct %0, %0 : !torch.list<int>, !torch.list<int> -> !torch.tuple<list<int>, list<int>>
     return %1 : !torch.tuple<list<int>, list<int>>
   }
-  func.func @"__torch_mlir_shape_fn.aten.mean.dim"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.bool, %arg3: !torch.optional<int>) -> !torch.list<int> {
+  func.func @"__torch_mlir_shape_fn.aten.mean.dim"(%arg0: !torch.list<int>, %arg1: !torch.optional<list<int>>, %arg2: !torch.bool, %arg3: !torch.optional<int>) -> !torch.list<int> {
     %true = torch.constant.bool true
+    %none = torch.constant.none
     %int0 = torch.constant.int 0
-    %0 = torch.aten.len.t %arg1 : !torch.list<int> -> !torch.int
-    %1 = torch.aten.eq.int %0, %int0 : !torch.int, !torch.int -> !torch.bool
+    %0 = torch.aten.__is__ %arg1, %none : !torch.optional<list<int>>, !torch.none -> !torch.bool
+    %1 = torch.prim.If %0 -> (!torch.bool) {
+      torch.prim.If.yield %true : !torch.bool
+    } else {
+      %5 = torch.prim.unchecked_cast %arg1 : !torch.optional<list<int>> -> !torch.list<int>
+      %6 = torch.aten.len.t %5 : !torch.list<int> -> !torch.int
+      %7 = torch.aten.eq.int %6, %int0 : !torch.int, !torch.int -> !torch.bool
+      torch.prim.If.yield %7 : !torch.bool
+    }
     %2 = torch.prim.If %1 -> (!torch.list<int>) {
       %5 = torch.aten.len.t %arg0 : !torch.list<int> -> !torch.int
       %6 = torch.prim.ListConstruct  : () -> !torch.list<int>
@@ -5693,7 +5701,8 @@ module {
       } : (!torch.int, !torch.bool) -> ()
       torch.prim.If.yield %6 : !torch.list<int>
     } else {
-      torch.prim.If.yield %arg1 : !torch.list<int>
+      %5 = torch.prim.unchecked_cast %arg1 : !torch.optional<list<int>> -> !torch.list<int>
+      torch.prim.If.yield %5 : !torch.list<int>
     }
     %3 = torch.derefine %arg3 : !torch.optional<int> to !torch.any
     %4 = call @__torch__.torch.jit._shape_functions.mean_dim(%arg0, %2, %arg2, %3) : (!torch.list<int>, !torch.list<int>, !torch.bool, !torch.any) -> !torch.list<int>

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/shape_lib_gen.py

@@ -537,8 +537,8 @@ def aten〇max〇dim(self: List[int], dim: int, keepdim: bool = False) -> Tuple[
     reduced_shape = _reduce_along_dim(self, dim, keepdim)
     return reduced_shape, reduced_shape
 
-def aten〇mean〇dim(self: List[int], dim: List[int], keepdim: bool = False, dtype: Optional[int] = None) -> List[int]:
-    if len(dim)==0:
+def aten〇mean〇dim(self: List[int], dim: Optional[List[int]], keepdim: bool = False, dtype: Optional[int] = None) -> List[int]:
+    if dim is None or len(dim)==0:
         dim = list(range(len(self)))
     return upstream_shape_functions.mean_dim(self, dim, keepdim, dtype)
 

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -376,7 +376,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::cumsum : (Tensor, int, int?) -> (Tensor)")
     emit("aten::floor_divide.Scalar : (Tensor, Scalar) -> (Tensor)")
     emit("aten::logsumexp : (Tensor, int[], bool) -> (Tensor)")
-    emit("aten::mean.dim : (Tensor, int[], bool, int?) -> (Tensor)")
+    emit("aten::mean.dim : (Tensor, int[]?, bool, int?) -> (Tensor)")
     emit("aten::__and__.Tensor : (Tensor, Tensor) -> (Tensor)")
     emit("aten::_softmax : (Tensor, int, bool) -> (Tensor)")
     emit("aten::mean : (Tensor, int?) -> (Tensor)")

python/torch_mlir_e2e_test/test_suite/stats.py

@@ -202,6 +202,25 @@ def MeanDimEmptyDimModule_basic(module, tu: TestUtils):
 
 # ==============================================================================
 
+class MeanDimNoneDimModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1, -1], torch.float32, True),
+    ])
+    def forward(self, x):
+        return torch.ops.aten.mean(x, dim=None)
+
+
+@register_test_case(module_factory=lambda: MeanDimNoneDimModule())
+def MeanDimNoneDimModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4, 5))
+
+# ==============================================================================
+
 class VarUnbiasedModule(torch.nn.Module):
     def __init__(self):
         super().__init__()