[mlir][Canonicalize] Copy ParallelOp Attributes in Single-iteration #145852
Conversation
|
@llvm/pr-subscribers-mlir-scf @llvm/pr-subscribers-mlir Author: Michael Marjieh (mmarjieh) ChangesFull diff: https://github.com/llvm/llvm-project/pull/145852.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 79012dbd32f80..7992e561895d8 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3133,6 +3133,16 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
newSteps, op.getInitVals(), nullptr);
// Erase the empty block that was inserted by the builder.
rewriter.eraseBlock(newOp.getBody());
+
+ // The new ParallelOp needs to keep all attributes from the old one, except
+ // for "operandSegmentSizes" which will be outdated.
+ for (const auto &namedAttr : op->getAttrs()) {
+ if (namedAttr.getName() == ParallelOp::getOperandSegmentSizeAttr())
+ continue;
+ rewriter.modifyOpInPlace(newOp, [&]() {
+ newOp->setAttr(namedAttr.getName(), namedAttr.getValue());
+ });
+ }
// Clone the loop body and remap the block arguments of the collapsed loops
// (inlining does not support a cancellable block argument mapping).
rewriter.cloneRegionBefore(op.getRegion(), newOp.getRegion(),
diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir
index 5e32a3a78c032..41f608f8f0f30 100644
--- a/mlir/test/Dialect/SCF/canonicalize.mlir
+++ b/mlir/test/Dialect/SCF/canonicalize.mlir
@@ -94,6 +94,38 @@ func.func @single_iteration_reduce(%A: index, %B: index) -> (index, index) {
// -----
+func.func @single_iteration_with_attributes(%A: memref<?x?x?xi32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c3 = arith.constant 3 : index
+ %c6 = arith.constant 6 : index
+ %c7 = arith.constant 7 : index
+ %c10 = arith.constant 10 : index
+ scf.parallel (%i0, %i1, %i2) = (%c0, %c3, %c7) to (%c1, %c6, %c10) step (%c1, %c2, %c3) {
+ %c42 = arith.constant 42 : i32
+ memref.store %c42, %A[%i0, %i1, %i2] : memref<?x?x?xi32>
+ scf.reduce
+ } {some_attr}
+ return
+}
+
+// CHECK-LABEL: func @single_iteration_with_attributes(
+// CHECK-SAME: [[ARG0:%.*]]: memref<?x?x?xi32>) {
+// CHECK-DAG: [[C42:%.*]] = arith.constant 42 : i32
+// CHECK-DAG: [[C7:%.*]] = arith.constant 7 : index
+// CHECK-DAG: [[C6:%.*]] = arith.constant 6 : index
+// CHECK-DAG: [[C3:%.*]] = arith.constant 3 : index
+// CHECK-DAG: [[C2:%.*]] = arith.constant 2 : index
+// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index
+// CHECK: scf.parallel ([[V0:%.*]]) = ([[C3]]) to ([[C6]]) step ([[C2]]) {
+// CHECK: memref.store [[C42]], [[ARG0]]{{\[}}[[C0]], [[V0]], [[C7]]] : memref<?x?x?xi32>
+// CHECK: scf.reduce
+// CHECK: } {some_attr}
+// CHECK: return
+
+// -----
+
func.func @nested_parallel(%0: memref<?x?x?xf64>) -> memref<?x?x?xf64> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
|
| rewriter.modifyOpInPlace(newOp, [&]() { | ||
| newOp->setAttr(namedAttr.getName(), namedAttr.getValue()); | ||
| }); | ||
| } |
There was a problem hiding this comment.
- This looks costly to me to do this one attribute at a time. It'll incur a lot of round-trip to the uniquer
- These a discardable attributes, we can't propagate them blindly through transformations: this is unsafe (unless proven otherwise, which seems hard to me here).
There was a problem hiding this comment.
I have an attribute that I wanted to preserve with this transformation.
How do you suggest I do that?
There was a problem hiding this comment.
I don't have a suggestion for this, because this is a fundamental issue in MLIR that is hard to resolve. I tried to look into this a while back (see https://discourse.llvm.org/t/rfc-implicit-propagation-of-dialect-attributes-best-effort/2657 for example) but we never found a solution.
For now "discardable attributes" are meant to be discarded outside of your own transformation that will understand them and preserve them.
There was a problem hiding this comment.
Sure its discardable, but it can also propagate it. It doesnt know/doesnt care. If the attribute is not valid anymore, then that is the problem of something that is adding the attribute downstream. I think it is a pragmatic solution to keep the attributes here without any gaurantee that the attribute semantics are preserved on the rewrite.
There was a problem hiding this comment.
And why would it round trip through the uniquer? It isnt creating a new attribute, just transfering it. It should be cheap.
There was a problem hiding this comment.
Sure its discardable, but it can also propagate it. It doesnt know/doesnt care. If the attribute is not valid anymore, then that is the problem of something that is adding the attribute downstream. I think it is a pragmatic solution to keep the attributes here without any gaurantee that the attribute semantics are preserved on the rewrite.
I already explained to you why this is incorrect, I won't get into it again, please refer to past discussions.
And why would it round trip through the uniquer? It isnt creating a new attribute, just transfering it. It should be cheap.
This is iterating through the dictionary of attributes of the original op, and for each of them it'll call newOp->setAttr(). This means that for each attribute it will take the existing incrementally built dictionary on the new op, unpack it to a list of NamedAttribute, append the new NamedAttribute, and build a new dictionary through the uniquer: as many times as there are attributes to be added.
There was a problem hiding this comment.
@mmarjieh you will not get a resolution of this upstream. It will pull you into a discussion hole. Downstream in IREE we ended up using a listener based approach that adds back attributes that get dropped due to upstream https://github.com/iree-org/iree/blob/57d61720368c1f23bab13463cb52d0d785dd82f0/compiler/src/iree/compiler/Codegen/Common/ConfigTrackingCanonicalizer.cpp#L92 . Thats the best way forward for you
There was a problem hiding this comment.
Side comment: in another downstream project I used regions to model custom op attributes https://github.com/intel/mlir-extensions/blob/main/docs/rfcs/RegionDialect.md, which is more likely to survive unrelated passes.
There was a problem hiding this comment.
@mmarjieh You may find this presentation from EuroLLVM 2025 conference interesting.
There was a problem hiding this comment.
@joker-eph @Hardcode84 @matthias-springer @MaheshRavishankar
Thanks everyone that replied. Bottom line, I can't rely on attributes to transfer data between different passes since every target has their own attributes and we can't simply propagate them without understanding if the semantics stays the same after the transformation,
I also went for a solution that doesn't rely on attributes.
No description provided.