[Torch] support adaptive_max_pool1d when return_indices equals False (#…

…3783)
llvm · Oct 11, 2024 · 7b11dfc · 7b11dfc
1 parent 8787970
commit 7b11dfc
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -7368,10 +7368,19 @@ class DecomposeAtenAdaptiveMaxPool1dOp
         loc, Torch::ListType::get(Torch::IntType::get(context)),
         ValueRange{constantOne});
 
-    rewriter.replaceOpWithNewOp<AtenMaxPool1dWithIndicesOp>(
-        op, op.getType(0), op.getType(1), input, kernelSizeList, strideList,
-        paddingSizeList, dialationList,
-        /*ceil_mode=*/constantFalse);
+    if (op.getResult(1).use_empty()) {
+      auto maxPool = rewriter.create<AtenMaxPool1dOp>(
+          loc, op.getType(0), input, kernelSizeList, strideList,
+          paddingSizeList, dialationList,
+          /*ceil_mode=*/constantFalse);
+      rewriter.replaceOp(op, {maxPool.getResult(), Value()});
+    } else {
+      auto maxPool = rewriter.create<AtenMaxPool1dWithIndicesOp>(
+          loc, op.getType(0), op.getType(1), input, kernelSizeList, strideList,
+          paddingSizeList, dialationList,
+          /*ceil_mode=*/constantFalse);
+      rewriter.replaceOp(op, maxPool.getResults());
+    }
     return success();
   }
 };