[MLIR][TORCH] Add E2E support for max_pool2d_with_indices op

This commit adds lowering of `max_pool2d_with_indices` op.

Signed-Off By: Vivek Khandelwal <vivek@nod-labs.com>

Author: vivekkhandelwal1

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

@@ -2938,6 +2938,35 @@ def Torch_AtenMaxPool2dOp : Torch_Op<"aten.max_pool2d", [
   }];
 }
 
+def Torch_AtenMaxPool2dWithIndicesOp : Torch_Op<"aten.max_pool2d_with_indices", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::max_pool2d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self,
+    ListOfTorchIntType:$kernel_size,
+    ListOfTorchIntType:$stride,
+    ListOfTorchIntType:$padding,
+    ListOfTorchIntType:$dilation,
+    Torch_BoolType:$ceil_mode
+  );
+  let results = (outs
+    AnyTorchTensorType:$result0,
+    AnyTorchTensorType:$result1
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult AtenMaxPool2dWithIndicesOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 6, 2);
+    }
+    void AtenMaxPool2dWithIndicesOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 6, 2);
+    }
+  }];
+}
+
 def Torch_AtenSoftmaxIntOp : Torch_Op<"aten.softmax.int", [
     AllowsTypeRefinement,
     HasValueSemantics,

lib/Conversion/TorchToLinalg/Linear.cpp

@@ -513,8 +513,8 @@ class ConvertAtenConvolutionOp : public OpConversionPattern<AtenConvolutionOp> {
     SmallVector<int64_t, 4> paddingIncludingNC = {0, 0};
     paddingIncludingNC.insert(paddingIncludingNC.end(), paddingInts.begin(),
                               paddingInts.end());
-    Value paddedInput = torch_to_linalg::getPaddedTensor(op, rewriter, input,
-                                                         paddingIncludingNC);
+    Value paddedInput = torch_to_linalg::getZeroPaddedTensor(
+        op, rewriter, input, paddingIncludingNC);
 
     SmallVector<Value> paddingIntValues =
         getAsConstantIntValues(rewriter, loc, paddingInts);

lib/Conversion/TorchToLinalg/Pooling.cpp

@@ -55,9 +55,9 @@ Value torch_to_linalg::getPaddedTensor(
 // Helper function to get the padding tensor given the padding int values.
 // It's assumed that the padding on the low end and high end are the same,
 // and that zero padding is required.
-Value torch_to_linalg::getPaddedTensor(Operation *op, OpBuilder &b,
-                                       Value &input,
-                                       SmallVectorImpl<int64_t> &paddingInts) {
+Value torch_to_linalg::getZeroPaddedTensor(
+    Operation *op, OpBuilder &b, Value &input,
+    SmallVectorImpl<int64_t> &paddingInts) {
   assert(input.getType().isa<RankedTensorType>() &&
          "input must be RankedTensorType");
   Location loc = op->getLoc();

lib/Conversion/TorchToLinalg/Utils.cpp

@@ -21,8 +21,8 @@ Value getPaddedTensor(Operation *op, OpBuilder &b, Value &input,
 // Helper function to get the padding tensor given the padding int values.
 // It's assumed that the padding on the low end and high end are the same,
 // and that zero padding is required.
-Value getPaddedTensor(Operation *op, OpBuilder &b, Value &input,
-                      SmallVectorImpl<int64_t> &paddingInts);
+Value getZeroPaddedTensor(Operation *op, OpBuilder &b, Value &input,
+                          SmallVectorImpl<int64_t> &paddingInts);
 
 // Helper function to caculate the output tensor dims for convolution-like ops.
 // Along each dim:

lib/Conversion/TorchToLinalg/Utils.h

@@ -669,6 +669,21 @@ ChangeResult TypeAnalyzer::visitOperation(
     return changed;
   }
 
+  if (isa<AtenMaxPool2dWithIndicesOp>(op)) {
+    auto self = operands[0]->getValue();
+    auto result0Knowledge =
+        ValueKnowledge::getNotNonePessimisticValueState(op->getContext());
+    result0Knowledge.dtype = self.dtype;
+    auto result1Knowledge =
+        ValueKnowledge::getNotNonePessimisticValueState(op->getContext());
+    result1Knowledge.dtype =
+        IntegerType::get(op->getContext(), 64, IntegerType::Signed);
+    ;
+    auto changed = incorporateKnowledge(op->getResult(0), result0Knowledge);
+    changed |= incorporateKnowledge(op->getResult(1), result1Knowledge);
+    return changed;
+  }
+
   if (auto arange = dyn_cast<AtenArangeOp>(op)) {
     return visitAtenArangeOp(arange);
   }

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -1572,6 +1572,11 @@ module {
     }
     return %none : !torch.none
   }
+  func @"__torch_mlir_shape_fn.aten.max_pool2d_with_indices"(%arg0: !torch.list<int>, %arg1: !torch.list<int>, %arg2: !torch.list<int>, %arg3: !torch.list<int>, %arg4: !torch.list<int>, %arg5: !torch.bool) -> !torch.tuple<list<int>, list<int>> {
+    %0 = call @__torch__.torch_mlir.dialects.torch.importer.jit_ir.build_tools.upstream_shape_helpers.max_pool2d(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) : (!torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.list<int>, !torch.bool) -> !torch.list<int>
+    %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 @"__torch_mlir_shape_fn.aten.adaptive_avg_pool2d"(%arg0: !torch.list<int>, %arg1: !torch.list<int>) -> !torch.list<int> {
     %0 = call @__torch__.torch_mlir.dialects.torch.importer.jit_ir.build_tools.upstream_shape_helpers.adaptive_avg_pool2d(%arg0, %arg1) : (!torch.list<int>, !torch.list<int>) -> !torch.list<int>
     return %0 : !torch.list<int>

lib/Dialect/Torch/Transforms/ShapeLibrary.cpp

@@ -563,6 +563,10 @@ def aten〇resize_(self: List[int], size: List[int], memory_format: Optional[int
 def aten〇max_pool2d(self: List[int], kernel_size: List[int], stride: List[int] = (), padding: List[int] = (0, 0), dilation: List[int] = (1, 1), ceil_mode: bool = False) -> List[int]:
     return upstream_shape_helpers.max_pool2d(self, kernel_size, stride, padding, dilation, ceil_mode)
 
+def aten〇max_pool2d_with_indices(self: List[int], kernel_size: List[int], stride: List[int] = (), padding: List[int] = (0, 0), dilation: List[int] = (1, 1), ceil_mode: bool = False) -> Tuple[List[int], List[int]]:
+    maxpool2d = indices = upstream_shape_helpers.max_pool2d(self, kernel_size, stride, padding, dilation, ceil_mode)
+    return maxpool2d, indices
+
 def aten〇adaptive_avg_pool2d(self: List[int], output_size: List[int]) -> List[int]:
     return upstream_shape_helpers.adaptive_avg_pool2d(self, output_size)
 

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

@@ -325,6 +325,9 @@ def emit_with_mutating_variants(key, **kwargs):
     emit(
         "aten::max_pool2d : (Tensor, int[], int[], int[], int[], bool) -> (Tensor)"
     )
+    emit(
+        "aten::max_pool2d_with_indices : (Tensor, int[], int[], int[], int[], bool) -> (Tensor, Tensor)"
+    )
     emit(
         "aten::softmax.int : (Tensor, int, int?) -> (Tensor)"
     )

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

@@ -35,3 +35,4 @@ def register_all_tests():
     from . import rng
     from . import cast
     from . import index_put
+    from . import pooling