Skip to content

MLIR: Add ReplaceOpAction into DialectConversion and PatternMatch #709

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

Jump to bottom
Open
wants to merge 1 commit into
base: aie-public
Choose a base branch
from
Open

MLIR: Add ReplaceOpAction into DialectConversion and PatternMatch #709

wants to merge 1 commit into from
+58 −1

Conversation

@mgehre-amd
Copy link
Collaborator

@mgehre-amd mgehre-amd commented Nov 12, 2025 •
edited
Loading

Add action to enable listener for RewriterBase::replaceOp, RewriterBase::replaceAllUsesWith and friends.

@mgehre-amd mgehre-amd force-pushed the matthias.replace_op_action_public branch from a84528d to a1da90b Compare November 13, 2025 07:51
@mgehre-amd mgehre-amd marked this pull request as ready for review November 13, 2025 07:51
cferry-AMD
cferry-AMD reviewed Nov 13, 2025
View reviewed changes
mlir/include/mlir/IR/PatternMatch.h
/// Find uses of `from` and replace them with `to`. Also notify the listener
/// about every in-place op modification (for every use that was replaced).
void replaceAllUsesWith(Value from, Value to) {
if (auto *fromOp = from.getDefiningOp())
Copy link
Contributor

@cferry-AMD cferry-AMD Nov 13, 2025

Choose a reason for hiding this comment

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

Also notify the listener
  /// about every in-place op modification (for every use that was replaced).

Was this not up-to-date? Or is an action something else than the listener?

Copy link
Contributor

@jorickert jorickert Nov 14, 2025

Choose a reason for hiding this comment

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

My main question is also whats the difference between an action and a listener

Copy link
Collaborator Author

@mgehre-amd mgehre-amd Nov 14, 2025

Choose a reason for hiding this comment

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

The RewriterBase supports attaching listeners, and the MLIRContext supports action handlers. In this PR, we use the action handlers on the MLIRContext, which allow us to observe across all passes without modifying each.

cferry-AMD
cferry-AMD reviewed Nov 13, 2025
View reviewed changes
cferry-AMD
cferry-AMD reviewed Nov 13, 2025
View reviewed changes
Copy link
Contributor

@cferry-AMD cferry-AMD left a comment

Choose a reason for hiding this comment

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

replaceAllUsesWith(ValueRange, ValueRange) should also be modified, shouldn't it?

@mgehre-amd
Copy link
Collaborator Author

mgehre-amd commented Nov 14, 2025

replaceAllUsesWith(ValueRange, ValueRange)
this calls the replaceAllUsesWith(Value from, Value to), which calls the action here

void replaceAllUsesWith(ValueRange from, ValueRange to) {
    assert(from.size() == to.size() && "incorrect number of replacements");
    for (auto it : llvm::zip(from, to))
      replaceAllUsesWith(std::get<0>(it), std::get<1>(it));
  }

cferry-AMD
cferry-AMD approved these changes Nov 14, 2025
View reviewed changes
Copy link
Contributor

@cferry-AMD cferry-AMD left a comment

Choose a reason for hiding this comment

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

Thanks for the claraifications, that looks good.

jorickert
jorickert approved these changes Nov 17, 2025
View reviewed changes
mlir/lib/IR/PatternMatch.cpp

void ReplaceOpAction::print(raw_ostream &os) const {
OpPrintingFlags flags;
flags.elideLargeElementsAttrs(10);
Copy link
Contributor

@jorickert jorickert Nov 17, 2025

Choose a reason for hiding this comment

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

I assume we do not have access to the compiler flags at this point?

mlir/lib/IR/PatternMatch.cpp
Comment on lines +145 to +152
for (auto r : replacement) {
if (!first)
os << ", ";
os << "`";
r.print(os, flags);
os << "`";
first = false;
}
Copy link
Contributor

@jorickert jorickert Nov 17, 2025

Choose a reason for hiding this comment

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

nit: This could have been written with llvm::interleave, but does not make a big difference

roberteg16
roberteg16 approved these changes Nov 17, 2025
View reviewed changes
Copy link
Contributor

@roberteg16 roberteg16 left a comment

Choose a reason for hiding this comment

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

Thanks!

mlir/include/mlir/IR/PatternMatch.h
void replaceAllUsesWith(Value from, Value to) {
if (auto *fromOp = from.getDefiningOp())
getContext()->executeAction<ReplaceOpAction>(
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
Copy link
Contributor

@roberteg16 roberteg16 Nov 17, 2025

Choose a reason for hiding this comment

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

Suggested change
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
/*actionFn=*/[]() {}, ArrayRef<IRUnit>{fromOp}, to);

mlir/lib/IR/PatternMatch.cpp
Comment on lines +144 to +152
bool first = true;
for (auto r : replacement) {
if (!first)
os << ", ";
os << "`";
r.print(os, flags);
os << "`";
first = false;
}
Copy link
Contributor

@roberteg16 roberteg16 Nov 17, 2025

Choose a reason for hiding this comment

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

Can you use llvm:interleaveComma here?

mlir/lib/IR/PatternMatch.cpp
if (replace) {
if (auto *fromOp = from.getDefiningOp())
getContext()->executeAction<ReplaceOpAction>(
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
Copy link
Contributor

@roberteg16 roberteg16 Nov 17, 2025

Choose a reason for hiding this comment

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

Suggested change
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
/*actionFn=*/[]() {}, ArrayRef<IRUnit>{fromOp}, to);

mlir/include/mlir/IR/PatternMatch.h
#include "mlir/IR/BuiltinOps.h"
#include "llvm/ADT/FunctionExtras.h"
#include "llvm/Support/TypeName.h"
#include <mlir/IR/PatternMatchAction.h>
Copy link
Contributor

@roberteg16 roberteg16 Nov 17, 2025

Choose a reason for hiding this comment

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

I believe mlir project prefers "" for internal headers

Suggested change
#include <mlir/IR/PatternMatchAction.h>
#include "mlir/IR/PatternMatchAction.h"

Copy link
Contributor

@cferry-AMD cferry-AMD Nov 17, 2025

Choose a reason for hiding this comment

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

Thanks Robert, TIL :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cferry-AMD cferry-AMD cferry-AMD approved these changes

@jorickert jorickert jorickert approved these changes

@roberteg16 roberteg16 roberteg16 approved these changes

@josel-amd josel-amd josel-amd approved these changes

@ehsan-toosi ehsan-toosi Awaiting requested review from ehsan-toosi ehsan-toosi is a code owner

@ttjost ttjost Awaiting requested review from ttjost ttjost is a code owner

@ljfitz ljfitz Awaiting requested review from ljfitz ljfitz is a code owner

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@mgehre-amd @jorickert @roberteg16 @cferry-AMD @josel-amd