Skip to content

[mlir][Canonicalize] Copy ParallelOp Attributes in Single-iteration #145852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+42 −0
Open

[mlir][Canonicalize] Copy ParallelOp Attributes in Single-iteration #145852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions mlir/lib/Dialect/SCF/IR/SCF.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
});
}
Copy link
Collaborator

@joker-eph joker-eph Jun 26, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. This looks costly to me to do this one attribute at a time. It'll incur a lot of round-trip to the uniquer
  2. These a discardable attributes, we can't propagate them blindly through transformations: this is unsafe (unless proven otherwise, which seems hard to me here).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have an attribute that I wanted to preserve with this transformation.
How do you suggest I do that?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And why would it round trip through the uniquer? It isnt creating a new attribute, just transfering it. It should be cheap.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mmarjieh You may find this presentation from EuroLLVM 2025 conference interesting.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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.

// 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(),
Expand Down
32 changes: 32 additions & 0 deletions mlir/test/Dialect/SCF/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading