torch.aten.max pool2d with indices (#125)

* Update to latest xilinx/llvm-project feature/backport_ea1_ops * Decompose aten.max_pool2d_with_indices
Xilinx · Jul 21, 2023 · 2174481 · 2174481
1 parent 90856f5
commit 2174481
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 7 deletions.
diff --git a/e2e_testing/xfail_sets.py b/e2e_testing/xfail_sets.py
@@ -1275,16 +1275,11 @@
 
     # 'tensor.empty' op incorrect number of dynamic sizes, has 1, expected 0
     "BatchNorm1DStaticShapeModule_basic",
+    "ResNet18StaticModule_basic",
 
     # Dynamic shape, has extra unsupported broadcast ops
     "Matmul_3d",
 
-    # failed to legalize operation 'torch.aten.max_pool2d_with_indices
-    "MaxPool2dEmptyStrideStaticModule_basic",
-    "MaxPool2dStaticCeilModeTrueModule_basic",
-    "MaxPool2dStaticModule_basic",
-    "ResNet18StaticModule_basic",
-
     # Unimplemented operator 'aten._index_put_impl_.hacked_twin'
     "IndexPutImpl1DFloatNonAccumulateModule_basic",
     "IndexPutImpl1DIntNonAccumulateModule_basic",

diff --git a/externals/llvm-project b/externals/llvm-project
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -4628,6 +4628,31 @@ class DecomposeAtenFakeQuantizePerTensorAffineCachemaskOp
 };
 } // namespace
 
+namespace {
+// Decompose aten.max_pool2d_with_indices
+// into aten.max_pool2d
+// when the second result is unused.
+class DecomposeAtenMaxPool2dWithIndicesOp
+    : public OpRewritePattern<AtenMaxPool2dWithIndicesOp> {
+public:
+  using OpRewritePattern<
+      AtenMaxPool2dWithIndicesOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenMaxPool2dWithIndicesOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op->getResult(1).use_empty())
+      return failure();
+
+    auto newOp = rewriter.create<AtenMaxPool2dOp>(
+        op.getLoc(), op->getResult(0).getType(), op.getSelf(), op.getKernelSize(),
+        op.getStride(), op.getPadding(), op.getDilation(), op.getCeilMode());
+
+    rewriter.replaceAllUsesWith(op->getResult(0), newOp);
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+} // namespace
+
 namespace {
 class DecomposeComplexOpsPass
     : public DecomposeComplexOpsBase<DecomposeComplexOpsPass> {
@@ -4803,6 +4828,8 @@ class DecomposeComplexOpsPass
     addPatternIfTargetOpIsIllegal<DecomposePrimsSumOp>(patterns);
     addPatternIfTargetOpIsIllegal<
         DecomposeAtenFakeQuantizePerTensorAffineCachemaskOp>(patterns);
+    addPatternIfTargetOpIsIllegal<
+        DecomposeAtenMaxPool2dWithIndicesOp>(patterns);
 
     GreedyRewriteConfig config;
     config.useTopDownTraversal = true;