[Mlir-commits] [mlir] [MLIR][OpenACC] Track compute region canonicalization updates (PR #221031)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 3 12:47:07 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Mehdi Amini (joker-eph)
<details>
<summary>Changes</summary>
The compute-region canonicalizations mutate block arguments, operands, and the operand segment attribute. Route these updates through modifyOpInPlace so pattern-rewriter listeners observe every mutation.
Assisted-by: Codex
---
Full diff: https://github.com/llvm/llvm-project/pull/221031.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp (+60-41)
- (modified) mlir/test/Dialect/OpenACC/compute-region-canonicalize.mlir (+30-2)
``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
index 441e59ca0cd49..7db3fefd01e24 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACCCG.cpp
@@ -78,15 +78,18 @@ struct RemoveEmptyKernelEnvironment
}
};
-static void updateComputeRegionInputOperandSegments(ComputeRegionOp op,
- PatternRewriter &rewriter,
- size_t numInput) {
+// Capture `hasStream` before erasing inputs because segment metadata is only
+// repaired at the end of the modification transaction.
+static void setComputeRegionInputOperandSegments(ComputeRegionOp op,
+ PatternRewriter &rewriter,
+ size_t numInput,
+ bool hasStream) {
const size_t numLaunch = op.getLaunchArgs().size();
op->setInherentAttr(
rewriter.getStringAttr(ComputeRegionOp::getOperandSegmentSizeAttr()),
rewriter.getDenseI32ArrayAttr({static_cast<int32_t>(numLaunch),
static_cast<int32_t>(numInput),
- op.getStream() ? 1 : 0}));
+ hasStream ? 1 : 0}));
}
struct ComputeRegionRemoveDuplicateArgs
@@ -101,33 +104,42 @@ struct ComputeRegionRemoveDuplicateArgs
assert(body->getNumArguments() == numLaunch + numInput &&
"region args mismatch");
- bool mergedAny = false;
- while (true) {
- bool merged = false;
- for (size_t j = 1; j < numInput && !merged; ++j) {
- for (size_t i = 0; i < j; ++i) {
- if (op->getOperand(static_cast<unsigned>(numLaunch + i)) !=
- op->getOperand(static_cast<unsigned>(numLaunch + j)))
- continue;
- unsigned keepIdx = static_cast<unsigned>(numLaunch + i);
- unsigned dropIdx = static_cast<unsigned>(numLaunch + j);
- rewriter.replaceAllUsesWith(body->getArgument(dropIdx),
- body->getArgument(keepIdx));
- body->eraseArgument(dropIdx);
- op->eraseOperand(dropIdx);
- --numInput;
- merged = true;
- mergedAny = true;
+ bool hasDuplicate = false;
+ for (size_t j = 1; j < numInput && !hasDuplicate; ++j)
+ for (size_t i = 0; i < j; ++i)
+ if (op->getOperand(static_cast<unsigned>(numLaunch + i)) ==
+ op->getOperand(static_cast<unsigned>(numLaunch + j))) {
+ hasDuplicate = true;
break;
}
- }
- if (!merged)
- break;
- }
-
- if (!mergedAny)
+ if (!hasDuplicate)
return failure();
- updateComputeRegionInputOperandSegments(op, rewriter, numInput);
+
+ const bool hasStream = static_cast<bool>(op.getStream());
+ rewriter.modifyOpInPlace(op, [&] {
+ while (true) {
+ bool merged = false;
+ for (size_t j = 1; j < numInput && !merged; ++j) {
+ for (size_t i = 0; i < j; ++i) {
+ if (op->getOperand(static_cast<unsigned>(numLaunch + i)) !=
+ op->getOperand(static_cast<unsigned>(numLaunch + j)))
+ continue;
+ unsigned keepIdx = static_cast<unsigned>(numLaunch + i);
+ unsigned dropIdx = static_cast<unsigned>(numLaunch + j);
+ rewriter.replaceAllUsesWith(body->getArgument(dropIdx),
+ body->getArgument(keepIdx));
+ body->eraseArgument(dropIdx);
+ op->eraseOperand(dropIdx);
+ --numInput;
+ merged = true;
+ break;
+ }
+ }
+ if (!merged)
+ break;
+ }
+ setComputeRegionInputOperandSegments(op, rewriter, numInput, hasStream);
+ });
return success();
}
};
@@ -144,21 +156,28 @@ struct ComputeRegionRemoveUnusedArgs
assert(body->getNumArguments() == numLaunch + numInput &&
"region args mismatch");
- bool changed = false;
- for (size_t k = numLaunch; k < numLaunch + numInput;) {
- if (!body->getArgument(static_cast<unsigned>(k)).use_empty()) {
- ++k;
- continue;
+ bool hasUnused = false;
+ for (size_t k = numLaunch; k < numLaunch + numInput; ++k)
+ if (body->getArgument(static_cast<unsigned>(k)).use_empty()) {
+ hasUnused = true;
+ break;
}
- body->eraseArgument(static_cast<unsigned>(k));
- op->eraseOperand(static_cast<unsigned>(k));
- --numInput;
- changed = true;
- }
-
- if (!changed)
+ if (!hasUnused)
return failure();
- updateComputeRegionInputOperandSegments(op, rewriter, numInput);
+
+ const bool hasStream = static_cast<bool>(op.getStream());
+ rewriter.modifyOpInPlace(op, [&] {
+ for (size_t k = numLaunch; k < numLaunch + numInput;) {
+ if (!body->getArgument(static_cast<unsigned>(k)).use_empty()) {
+ ++k;
+ continue;
+ }
+ body->eraseArgument(static_cast<unsigned>(k));
+ op->eraseOperand(static_cast<unsigned>(k));
+ --numInput;
+ }
+ setComputeRegionInputOperandSegments(op, rewriter, numInput, hasStream);
+ });
return success();
}
};
diff --git a/mlir/test/Dialect/OpenACC/compute-region-canonicalize.mlir b/mlir/test/Dialect/OpenACC/compute-region-canonicalize.mlir
index cb3a3c6199bac..24d74a381b2ae 100644
--- a/mlir/test/Dialect/OpenACC/compute-region-canonicalize.mlir
+++ b/mlir/test/Dialect/OpenACC/compute-region-canonicalize.mlir
@@ -1,7 +1,5 @@
// RUN: mlir-opt -canonicalize -split-input-file %s | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
// -----
// CHECK-LABEL: func @merge_duplicate_ins
@@ -23,6 +21,36 @@ func.func @merge_duplicate_ins() -> i32 {
// -----
+// CHECK-LABEL: func @merge_duplicate_ins_with_stream
+// CHECK-SAME: (%[[M:.*]]: memref<i32>, %[[STREAM:.*]]: !gpu.async.token)
+func.func @merge_duplicate_ins_with_stream(%m: memref<i32>, %stream: !gpu.async.token) {
+ // CHECK: acc.compute_region stream(%[[STREAM]] : !gpu.async.token) ins({{.*}} = %[[M]]) : (memref<i32>) {
+ acc.compute_region stream(%stream : !gpu.async.token) ins(%a = %m, %b = %m) : (memref<i32>, memref<i32>) {
+ %va = memref.load %a[] : memref<i32>
+ %vb = memref.load %b[] : memref<i32>
+ %sum = arith.addi %va, %vb : i32
+ memref.store %sum, %a[] : memref<i32>
+ acc.yield
+ } <{origin = "acc.serial"}>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @drop_unused_ins_with_stream
+// CHECK-SAME: (%[[USED:.*]]: memref<i32>, %{{.*}}: memref<i32>, %[[STREAM:.*]]: !gpu.async.token)
+func.func @drop_unused_ins_with_stream(%used: memref<i32>, %unused: memref<i32>, %stream: !gpu.async.token) {
+ // CHECK: acc.compute_region stream(%[[STREAM]] : !gpu.async.token) ins({{.*}} = %[[USED]]) : (memref<i32>) {
+ acc.compute_region stream(%stream : !gpu.async.token) ins(%a = %used, %b = %unused) : (memref<i32>, memref<i32>) {
+ %v = memref.load %a[] : memref<i32>
+ memref.store %v, %a[] : memref<i32>
+ acc.yield
+ } <{origin = "acc.serial"}>
+ return
+}
+
+// -----
+
// CHECK-LABEL: func @merge_duplicate_ins_complex_pattern
func.func @merge_duplicate_ins_complex_pattern() -> i32 {
%c0 = arith.constant 0 : i32
``````````
</details>
https://github.com/llvm/llvm-project/pull/221031
More information about the Mlir-commits
mailing list