onnx.MaxPool add atenMaxPool1d lowering support (#3452)

fixes #3422
llvm · Jun 13, 2024 · 919b599 · 919b599
1 parent 39d882f
commit 919b599
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
@@ -565,15 +565,15 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
         Value cstCeilMode =
             rewriter.create<Torch::ConstantBoolOp>(binder.getLoc(), ceilMode);
 
-        if (rank == 3)
-          return rewriter.notifyMatchFailure(binder.op,
-                                             "Unimplemented: AtenMaxPool1dOp");
-
         if (binder.op->getNumResults() == 2) {
           Torch::ValueTensorType resultTypeIndices;
           if (binder.tensorResultTypeAtIndex(resultTypeIndices, 1))
             return failure();
 
+          if (rank == 3)
+            return rewriter.notifyMatchFailure(
+                binder.op, "Unimplemented: AtenMaxPool1dWithIndicesOp");
+
           if (rank == 4) {
             rewriter.replaceOpWithNewOp<Torch::AtenMaxPool2dWithIndicesOp>(
                 binder.op, resultTypeOut, resultTypeIndices, operand,
@@ -589,6 +589,12 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
             return success();
           }
         } else {
+          if (rank == 3) {
+            rewriter.replaceOpWithNewOp<Torch::AtenMaxPool1dOp>(
+                binder.op, resultTypeOut, operand, kernelSizeList, stridesList,
+                paddingList, dilationsList, cstCeilMode);
+            return success();
+          }
           if (rank == 4) {
             rewriter.replaceOpWithNewOp<Torch::AtenMaxPool2dOp>(
                 binder.op, resultTypeOut, operand, kernelSizeList, stridesList,

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -2418,10 +2418,7 @@
     "LogSoftmaxBackwardModule_basic",
     "MaskedScatterStaticBasic_basic",
     "MaxPool1dCeilModeTrueModule_basic",
-    "MaxPool1dEmptyStrideStaticModule_basic",
     "MaxPool1dModule_basic",
-    "MaxPool1dStaticCeilModeTrueModule_basic",
-    "MaxPool1dStaticModule_basic",
     "MaxPool2dCeilModeTrueModule_basic",
     "MaxPool2dModule_basic",
     "MaxPool2dWithIndicesAllOnesModule_basic",