[Torch] Add support for aten.any.dims (#4200)

* Added lowering to Linalg-on-Tensors
* Added test to
projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

Author: mmanzoorTT

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

@@ -10940,6 +10940,32 @@ def Torch_AtenAnyDimOp : Torch_Op<"aten.any.dim", [
   }];
 }
 
+def Torch_AtenAnyDimsOp : Torch_Op<"aten.any.dims", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::any.dims : (Tensor, int[]?, bool) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    AnyTorchOptionalListOfTorchIntType:$dim,
+    Torch_BoolType:$keepdim
+  );
+  let results = (outs
+    AnyTorchOptionalTensorType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenAnyDimsOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 3, 1);
+    }
+    void AtenAnyDimsOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 3, 1);
+    }
+  }];
+  let hasFolder = 1;
+}
+
 def Torch_AtenArangeOp : Torch_Op<"aten.arange", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Conversion/TorchToLinalg/Reduction.cpp

@@ -337,7 +337,7 @@ static Value createInitElementForReduceOp(OpBuilder &b, Location loc,
     return b.create<arith::ConstantOp>(loc, b.getBoolAttr(true));
   }
 
-  if (isa<AtenAnyOp>(op)) {
+  if (isa<AtenAnyOp, AtenAnyDimsOp>(op)) {
     return b.create<arith::ConstantOp>(loc, b.getBoolAttr(false));
   }
 
@@ -434,7 +434,7 @@ static Value createLinalgPayloadForReduceOp(OpBuilder &b, Location loc,
     Value result = payloadArgs[1];
     Value self = convertScalarToDtype(b, loc, elem, resultElementType);
     return b.create<arith::AndIOp>(loc, self, result);
-  } else if (isa<AtenAnyOp>(op)) {
+  } else if (isa<AtenAnyOp, AtenAnyDimsOp>(op)) {
     Value elem = payloadArgs[0];
     Value result = payloadArgs[1];
     Value self = convertScalarToDtype(b, loc, elem, resultElementType);
@@ -532,6 +532,9 @@ class ConvertReductionOp : public ConversionPattern {
     if (auto allOp = dyn_cast<AtenAllDimOp>(op))
       return computeReductionOpInfoForDimVariantOp(allOp, operands, rewriter);
 
+    if (auto anyOp = dyn_cast<AtenAnyDimsOp>(op))
+      return computeReductionOpInfoForDimVariantOp(anyOp, operands, rewriter);
+
     return rewriter.notifyMatchFailure(op, "not a supported reduce op");
   }
 
@@ -709,6 +712,7 @@ void mlir::torch::torch_to_linalg::populateReductionPatternsAndLegality(
   patterns.add<ConvertAtenMinMaxDimOp<AtenMinDimOp>>(typeConverter, context);
   target.addIllegalOp<AtenSumOp>();
   target.addIllegalOp<AtenAnyOp>();
+  target.addIllegalOp<AtenAnyDimsOp>();
   target.addIllegalOp<AtenAllOp>();
   target.addIllegalOp<AtenSumDimIntListOp>();
   target.addIllegalOp<AtenProdOp>();

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -1120,6 +1120,22 @@ OpFoldResult AtenDimOp::fold(FoldAdaptor adaptor) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// AtenAnyDimsOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult AtenAnyDimsOp::fold(FoldAdaptor adaptor) {
+  auto resultType = dyn_cast<ValueTensorType>(getResult().getType());
+  auto resultShape = resultType.toBuiltinTensor().getShape();
+  auto inputType = dyn_cast<ValueTensorType>(getOperand(0).getType());
+  auto inputShape = inputType.toBuiltinTensor().getShape();
+  if ((inputType.getDtype() == resultType.getDtype()) &&
+      (inputShape == resultShape)) {
+    return getSelf();
+  }
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // AtenLenTOp
 //===----------------------------------------------------------------------===//

lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

@@ -7504,6 +7504,12 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    %1 = call @__torch__.torch.jit._shape_functions.argmax(%arg0, %0, %arg2) : (!torch.list<int>, !torch.optional<int>, !torch.bool) -> !torch.list<int>\n"
 "    return %1 : !torch.list<int>\n"
 "  }\n"
+"  func.func @\"__torch_mlir_shape_fn.aten.any.dims\"(%arg0: !torch.list<int>, %arg1: !torch.optional<list<int>>, %arg2: !torch.bool) -> !torch.list<int> {\n"
+"    %none = torch.constant.none\n"
+"    %0 = torch.derefine %none : !torch.none to !torch.any\n"
+"    %1 = call @__torch__.torch.jit._shape_functions.sum_mean_dim(%arg0, %arg1, %arg2, %0) : (!torch.list<int>, !torch.optional<list<int>>, !torch.bool, !torch.any) -> !torch.list<int>\n"
+"    return %1 : !torch.list<int>\n"
+"  }\n"
 "  func.func @\"__torch_mlir_shape_fn.aten.all.dim\"(%arg0: !torch.list<int>, %arg1: !torch.int, %arg2: !torch.bool) -> !torch.list<int> {\n"
 "    %0 = torch.derefine %arg1 : !torch.int to !torch.optional<int>\n"
 "    %1 = call @__torch__.torch.jit._shape_functions.argmax(%arg0, %0, %arg2) : (!torch.list<int>, !torch.optional<int>, !torch.bool) -> !torch.list<int>\n"
@@ -15420,6 +15426,18 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
 "    }\n"
 "    return %2 : !torch.int\n"
 "  }\n"
+"  func.func @\"__torch_mlir_dtype_fn.aten.any.dims\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.optional<list<int>>, %arg2: !torch.bool) -> !torch.int {\n"
+"    %int11 = torch.constant.int 11\n"
+"    %int0 = torch.constant.int 0\n"
+"    %0:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple<int, int> -> !torch.int, !torch.int\n"
+"    %1 = torch.aten.eq.int %0#1, %int0 : !torch.int, !torch.int -> !torch.bool\n"
+"    %2 = torch.prim.If %1 -> (!torch.int) {\n"
+"      torch.prim.If.yield %0#1 : !torch.int\n"
+"    } else {\n"
+"      torch.prim.If.yield %int11 : !torch.int\n"
+"    }\n"
+"    return %2 : !torch.int\n"
+"  }\n"
 "  func.func @\"__torch_mlir_dtype_fn.aten.all.dim\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.int, %arg2: !torch.bool) -> !torch.int {\n"
 "    %int11 = torch.constant.int 11\n"
 "    %int0 = torch.constant.int 0\n"

projects/pt1/e2e_testing/xfail_sets.py

@@ -824,6 +824,7 @@
     "RandnLikeDtypeModule_basic",
     "RandnLikeModule_basic",
     "RandnModule_basic",
+    "ReduceAnyDimsFloatModule_basic",
     "ReflectionPad1dModule2dInput_Right",
     "ReflectionPad1dModule2dInput_basic",
     "ReflectionPad1dModule3dInput_Left",
@@ -2691,6 +2692,7 @@
     "PermuteNegativeIndexModule_basic",
     # Failure - incorrect numerics
     "ReduceAnyDimFloatModule_basic",
+    "ReduceAnyDimsFloatModule_basic",
     "AvgPool2dDivisorOverrideModule_basic",
     "BroadcastDynamicDimModule_basic",
     "ElementwiseAtan2TensorIntModule_basic",
@@ -3793,6 +3795,7 @@
     "RandnLikeModule_basic",
     "RandnModule_basic",
     "ReduceAllDimEmpty_basic",
+    "ReduceAnyDimsFloatModule_basic",
     "ReduceFrobeniusNormComplexModule_basic",
     "ReduceL1NormComplexModule_basic",
     "ReduceL2NormComplexModule_basic",

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

@@ -778,6 +778,9 @@ def aten〇one_hot〡shape(self: List[int], num_classes: int = -1) -> List[int]:
 def aten〇any〇dim〡shape(self: List[int], dim: int, keepdim: bool = False) -> List[int]:
     return upstream_shape_functions.argmax(self, dim, keepdim)
 
+def aten〇any〇dims〡shape(self: List[int], dim: Optional[List[int]] = None, keepdim: bool = False) -> List[int]:
+    return upstream_shape_functions.sum_mean_dim(self, dim, keepdim, None)
+
 def aten〇all〇dim〡shape(self: List[int], dim: int, keepdim: bool = False) -> List[int]:
     return upstream_shape_functions.argmax(self, dim, keepdim)
 
@@ -5249,6 +5252,13 @@ def aten〇any〇dim〡dtype(self_rank_dtype: Tuple[int, int], dim: int, keepdim
         return self_dtype
     return torch.bool
 
+@check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1))
+def aten〇any〇dims〡dtype(self_rank_dtype: Tuple[int, int], dim: Optional[List[int]] = None, keepdim: bool = False) -> int:
+    self_rank, self_dtype = self_rank_dtype
+    if self_dtype == torch.uint8:
+        return self_dtype
+    return torch.bool
+
 @check_dtype_function(_check_tensors_with_the_same_dtype(num_of_tensors=1, dim=0))
 def aten〇all〇dim〡dtype(self_rank_dtype: Tuple[int, int], dim: int, keepdim: bool = False) -> int:
     self_rank, self_dtype = self_rank_dtype

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/torch_ods_gen.py

@@ -842,6 +842,7 @@ def emit_with_mutating_variants(key, **kwargs):
     emit("aten::all.dim : (Tensor, int, bool) -> (Tensor)")
     emit("aten::any : (Tensor) -> (Tensor)")
     emit("aten::any.dim : (Tensor, int, bool) -> (Tensor)")
+    emit("aten::any.dims : (Tensor, int[]?, bool) -> (Tensor)", has_folder=True)
     emit("aten::arange : (Scalar, int?, int?, Device?, bool?) -> (Tensor)")
     emit(
         "aten::arange.start : (Scalar, Scalar, int?, int?, Device?, bool?) -> (Tensor)"

projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

@@ -302,6 +302,26 @@ def ReduceAnyDimFloatModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(3, 4, 5))
 
 
+class ReduceAnyDimsFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, a):
+        return torch.ops.aten.any(a, dim=[0, 1])
+
+
+@register_test_case(module_factory=lambda: ReduceAnyDimsFloatModule())
+def ReduceAnyDimsFloatModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(3, 4, 5))
+
+
 # ==============================================================================