[Mlir-commits] [mlir] [mlir][linalg] Fix unused output cycle removal (PR #236751)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 05:19:40 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Akun (AkunCube)
<details>
<summary>Changes</summary>
Remove dead `linalg.generic` output computations without replacing multi-result operations with a single value.
Fixes #<!-- -->236365.
---
Full diff: https://github.com/llvm/llvm-project/pull/236751.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp (+39-25)
- (modified) mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir (+76)
``````````diff
diff --git a/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp b/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
index 506e20ab436f8f..0c809f3e43292b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp
@@ -6,9 +6,12 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Transforms/RegionUtils.h"
using namespace mlir;
using namespace mlir::linalg;
@@ -333,48 +336,59 @@ struct RemoveUnusedCycleInGenericOp : public OpRewritePattern<GenericOp> {
return failure();
bool hasRemovedCycles = false;
+ auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
// Iterate over output operands and remove any unused cycles.
- for (const auto &outputOpOperand :
- llvm::enumerate(genericOp.getDpsInits())) {
+ for (auto [outputIndex, _] : llvm::enumerate(genericOp.getDpsInits())) {
// Check that result from out operand is dead.
- Value result = genericOp.getResult(outputOpOperand.index());
- if (!result.use_empty())
+ if (!genericOp.getResult(outputIndex).use_empty())
continue;
- // Check that outputArg has one use in cycle.
- BlockArgument outputArg =
- genericOp.getRegionOutputArgs()[outputOpOperand.index()];
- if (!outputArg.hasOneUse())
+ BlockArgument outputArg = genericOp.getRegionOutputArgs()[outputIndex];
+ Value yieldedValue = yieldOp.getOperand(outputIndex);
+ if (yieldedValue == outputArg)
continue;
- // Check cycle has at most one use.
- Operation *cycleOp = *outputArg.user_begin();
- if (!cycleOp->hasOneUse())
+ SliceOptions options;
+ options.filter = [&](Operation *op) {
+ return genericOp->isProperAncestor(op) && op != yieldOp.getOperation();
+ };
+ llvm::SetVector<Operation *> forwardSlice;
+ getForwardSlice(outputArg, &forwardSlice, options);
+
+ // Only rewrite a value computed from this output block argument.
+ if (!forwardSlice.contains(yieldedValue.getDefiningOp()))
continue;
- // Check that the cycleUser is a yield.
- Operation *cycleUserOp = *cycleOp->user_begin();
- if (!isa<linalg::YieldOp>(cycleUserOp))
+ // Every operation reached from the output argument must be removable
+ // once the corresponding yield operand is replaced.
+ if (!llvm::all_of(forwardSlice, wouldOpBeTriviallyDead))
continue;
- // Check that argIndex matches yieldIndex, else data is being used.
- if (cycleUserOp->getOperand(outputOpOperand.index()) !=
- cycleOp->getResult(0))
+ // A different yield operand may use the output argument directly or
+ // use a value computed from it.
+ bool usedByOtherYield = false;
+ for (auto [yieldIndex, value] : llvm::enumerate(yieldOp.getValues())) {
+ if (yieldIndex == outputIndex)
+ continue;
+ if (value == outputArg ||
+ forwardSlice.contains(value.getDefiningOp())) {
+ usedByOtherYield = true;
+ break;
+ }
+ }
+ if (usedByOtherYield)
continue;
- // Directly replace the cycle with the blockArg such that
- // Deduplicate pattern can eliminate it along with unused yield.
- rewriter.replaceOp(cycleOp, outputArg);
+ rewriter.modifyOpInPlace(yieldOp, [&, index = outputIndex]() {
+ yieldOp.setOperand(index, outputArg);
+ });
+ eliminateTriviallyDeadOps(rewriter, genericOp.getRegion());
rewriter.modifyOpInPlace(genericOp, [] {});
hasRemovedCycles = true;
}
- if (hasRemovedCycles) {
- return success();
- }
-
- return failure();
+ return success(hasRemovedCycles);
}
};
diff --git a/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir b/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
index 4dc10a5fde1874..fe846871857c5e 100644
--- a/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
+++ b/mlir/test/Dialect/Linalg/erase-unused-operands-and-results.mlir
@@ -422,6 +422,82 @@ func.func @drop_dead_results_with_unused_cycles(%arg0 : tensor<?x?x?xf32>) -> (t
// -----
+// A dead output may be computed by a multi-result op.
+
+#map = affine_map<(d0) -> (d0)>
+func.func @drop_dead_multi_result_output(%init0: tensor<1xi32>, %init1: tensor<1xi32>, %rhs: i32) -> tensor<1xi32> {
+ %results:2 = linalg.generic {
+ indexing_maps = [#map, #map],
+ iterator_types = ["parallel"]}
+ outs(%init0, %init1 : tensor<1xi32>, tensor<1xi32>) {
+ ^bb0(%out0: i32, %out1: i32):
+ %diff, %borrow = arith.subui_extended %out0, %rhs : i32, i1
+ linalg.yield %diff, %out1 : i32, i32
+ } -> (tensor<1xi32>, tensor<1xi32>)
+ return %results#1 : tensor<1xi32>
+}
+
+// CHECK-LABEL: func.func @drop_dead_multi_result_output(
+// CHECK-SAME: %[[INIT0:[^:]+]]: tensor<1xi32>, %[[INIT1:[^:]+]]: tensor<1xi32>
+// CHECK: linalg.generic
+// CHECK-SAME: outs(%[[INIT1]] : tensor<1xi32>)
+// CHECK-NOT: arith.subui_extended
+// CHECK: linalg.yield
+
+// -----
+
+// Drop a dead output and its chain of computations.
+
+#map = affine_map<(d0) -> (d0)>
+func.func @drop_dead_output_computation_chain(%init0: tensor<1xi32>, %init1: tensor<1xi32>, %rhs: i32) -> tensor<1xi32> {
+ %results:2 = linalg.generic {
+ indexing_maps = [#map, #map],
+ iterator_types = ["parallel"]}
+ outs(%init0, %init1 : tensor<1xi32>, tensor<1xi32>) {
+ ^bb0(%out0: i32, %out1: i32):
+ %sum = arith.addi %out0, %rhs : i32
+ %product = arith.muli %sum, %rhs : i32
+ linalg.yield %product, %out1 : i32, i32
+ } -> (tensor<1xi32>, tensor<1xi32>)
+ return %results#1 : tensor<1xi32>
+}
+
+// CHECK-LABEL: func.func @drop_dead_output_computation_chain(
+// CHECK-SAME: %[[INIT0:[^:]+]]: tensor<1xi32>, %[[INIT1:[^:]+]]: tensor<1xi32>
+// CHECK: linalg.generic
+// CHECK-SAME: outs(%[[INIT1]] : tensor<1xi32>)
+// CHECK-NOT: arith.addi
+// CHECK-NOT: arith.muli
+// CHECK: linalg.yield
+
+// -----
+
+// Preserve a dead output whose block argument feeds a side effect.
+
+#map = affine_map<(d0) -> (d0)>
+func.func @preserve_effectful_output_computation(%init0: tensor<1xi32>, %init1: tensor<1xi32>, %rhs: i32, %mem: memref<1xi32>) -> tensor<1xi32> {
+ %c0 = arith.constant 0 : index
+ %results:2 = linalg.generic {
+ indexing_maps = [#map, #map],
+ iterator_types = ["parallel"]}
+ outs(%init0, %init1 : tensor<1xi32>, tensor<1xi32>) {
+ ^bb0(%out0: i32, %out1: i32):
+ %sum = arith.addi %out0, %rhs : i32
+ memref.store %out0, %mem[%c0] : memref<1xi32>
+ linalg.yield %sum, %out1 : i32, i32
+ } -> (tensor<1xi32>, tensor<1xi32>)
+ return %results#1 : tensor<1xi32>
+}
+
+// CHECK-LABEL: func.func @preserve_effectful_output_computation(
+// CHECK: %[[RESULTS:.+]]:2 = linalg.generic
+// CHECK: %[[SUM:.+]] = arith.addi %[[OUT0:[^,]+]],
+// CHECK: memref.store %[[OUT0]],
+// CHECK: linalg.yield %[[SUM]],
+// CHECK: return %[[RESULTS]]#1
+
+// -----
+
// Drop only the results not used by others.
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
``````````
</details>
https://github.com/llvm/llvm-project/pull/236751
More information about the Mlir-commits
mailing list