Enable lowering to linalg

Signed-off-by: Zahid Wakeel <[email protected]>

Author: zahidwx

lib/Conversion/TorchToLinalg/Pooling.cpp

@@ -225,6 +225,10 @@ template <> struct DimensionTraits<AtenMaxPool1dOp> {
   static_assert(Dim == Dim);
 };
 
+template <>
+struct DimensionTraits<AtenMaxPool1dWithIndicesOp>
+    : DimensionTraits<AtenMaxPool1dOp> {};
+
 template <> struct DimensionTraits<AtenMaxPool2dOp> {
   static constexpr int64_t Dim = 2;
   // unused const variable warning suppression:
@@ -250,7 +254,8 @@ class ConvertAtenMaxPoolOp : public OpConversionPattern<OpTy> {
   using OpConversionPattern<OpTy>::OpConversionPattern;
 
   static const bool withIndices =
-      llvm::is_one_of<OpTy, AtenMaxPool2dWithIndicesOp,
+      llvm::is_one_of<OpTy, AtenMaxPool1dWithIndicesOp,
+                      AtenMaxPool2dWithIndicesOp,
                       AtenMaxPool3dWithIndicesOp>::value;
 
 private:
@@ -1687,8 +1692,11 @@ void mlir::torch::torch_to_linalg::populatePoolingPatternsAndLegality(
   patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool2dOp>>(typeConverter, context);
   patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool3dOp>>(typeConverter, context);
 
+  target.addIllegalOp<AtenMaxPool1dWithIndicesOp>();
   target.addIllegalOp<AtenMaxPool2dWithIndicesOp>();
   target.addIllegalOp<AtenMaxPool3dWithIndicesOp>();
+  patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool1dWithIndicesOp>>(typeConverter,
+                                                                 context);
   patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool2dWithIndicesOp>>(typeConverter,
                                                                  context);
   patterns.add<ConvertAtenMaxPoolOp<AtenMaxPool3dWithIndicesOp>>(typeConverter,

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

@@ -180,6 +180,55 @@ def AdaptiveAvgPool2dUnitOutputSizeDynamicModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class MaxPool1dWithIndicesModule(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.max_pool1d_with_indices(
+            x, kernel_size=[6], stride=[2], padding=[3], dilation=2, ceil_mode=False
+        )
+
+
+@register_test_case(module_factory=lambda: MaxPool1dWithIndicesModule())
+def MaxPool1dWithIndicesModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(1, 64, 112, low=-1))
+
+
+class MaxPool1dWithIndicesCeilModeModule(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.max_pool1d_with_indices(
+            x, kernel_size=[6], stride=[2], padding=[3], 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))
+
+
+# ==============================================================================
+
+
 class MaxPool1dModule(torch.nn.Module):
 
     def __init__(self):