Add ceil_mode support in TorchToStablehlo lowering for AtenMaxPool1dWithIndicesOp

Author: zahidwx

lib/Conversion/TorchToStablehlo/Pooling.cpp

@@ -132,6 +132,29 @@ LogicalResult ConvertAtenOp<AtenMaxPool1dWithIndicesOp>::matchAndRewrite(
   stablehloPadding[stablehloPadding.size() - 1] = padding[0];
   stablehloPadding[stablehloPadding.size() - 2] = padding[0];
 
+  if (ceilMode) {
+    // Match PyTorch output shape with extra padding. See
+    // https://github.com/pytorch/pytorch/blob/c5de6ff079e3e5b453d6ff5190c90f02db458928/aten/src/ATen/native/Pool.h#L79
+    const int64_t inputSize = inputShape[inputRank - 1];
+    const int64_t numerator =
+        (inputSize + 2 * padding[0] - dilation[0] * (kernelSize[0] - 1) - 1);
+    const int64_t floor_output_size = (numerator) / stride[0] + 1;
+    const int64_t adj = (stride[0] - 1);
+    int64_t ceil_output_size = std::ceil((numerator + adj) / stride[0]) + 1;
+
+    // Ensure last pooling starts inside input
+    if ((ceil_output_size - 1) * stride[0] >= inputSize + padding[0]) {
+      ceil_output_size--;
+    }
+
+    // Add extra padding to make output size same as torch
+    if (ceil_output_size > floor_output_size) {
+      const int64_t sizeDiff = ceil_output_size - floor_output_size;
+      const int64_t extraPadding = sizeDiff * stride[0];
+      stablehloPadding[stablehloPadding.size() - 1] += extraPadding;
+    }
+  }
+
   Value initVal = createInitialValueForAtenPoolingOp(op, inputElemTy, rewriter);
 
   auto windowDimensions = rewriter.getDenseI64ArrayAttr(stablehloKernelSize);

projects/pt1/e2e_testing/xfail_sets.py

@@ -561,7 +561,6 @@
     "ElementwiseRemainderTensorModule_Int_Float_NegativeDivisor_basic",
     "ElementwiseRemainderTensorModule_Int_NegativeDividend_basic",
     "ElementwiseRemainderTensorModule_Int_NegativeDivisor_basic",
-    "MaxPool1dWithIndicesCeilModeModule_basic",
     "MaxPool1dCeilModeTrueModule_basic",
     "MaxPool1dStaticCeilModeTrueModule_basic",
     "MaxUnpool3dModulePad0_basic",

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

@@ -217,13 +217,13 @@ def __init__(self):
     )
     def forward(self, x):
         return torch.ops.aten.max_pool1d_with_indices(
-            x, kernel_size=[6], stride=[2], padding=[3], dilation=2, ceil_mode=True
+            x, kernel_size=[4], stride=[2], padding=[2], dilation=2, ceil_mode=True
         )
 
 
 @register_test_case(module_factory=lambda: MaxPool1dWithIndicesCeilModeModule())
 def MaxPool1dWithIndicesCeilModeModule_basic(module, tu: TestUtils):
-    module.forward(tu.rand(1, 64, 112, low=-1))
+    module.forward(tu.rand(3, 25, 37, low=-1))
 
 
 # ==============================================================================