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

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

MLIR: Add ReplaceOpAction into DialectConversion and PatternMatch #709

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
4 changes: 4 additions & 0 deletions mlir/include/mlir/IR/PatternMatch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#include "mlir/IR/BuiltinOps.h"
#include "llvm/ADT/FunctionExtras.h"
#include "llvm/Support/TypeName.h"
#include <mlir/IR/PatternMatchAction.h>
#include <optional>

using llvm::SmallPtrSetImpl;
Expand Down Expand Up @@ -652,6 +653,9 @@ class RewriterBase : public OpBuilder {
/// 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.

getContext()->executeAction<ReplaceOpAction>(
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
Operation *op = operand.getOwner();
modifyOpInPlace(op, [&]() { operand.set(to); });
Expand Down
20 changes: 20 additions & 0 deletions mlir/include/mlir/IR/PatternMatchAction.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#ifndef MLIR_IR_PATTERNMATCHACTION_H
#define MLIR_IR_PATTERNMATCHACTION_H

#include "mlir/IR/Action.h"

namespace mlir {
struct ReplaceOpAction : public tracing::ActionImpl<ReplaceOpAction> {
using Base = tracing::ActionImpl<ReplaceOpAction>;
ReplaceOpAction(ArrayRef<IRUnit> irUnits, ValueRange replacement);
static constexpr StringLiteral tag = "op-replacement";
void print(raw_ostream &os) const override;

Operation *getOp() const;

public:
ValueRange replacement;
};
} // namespace mlir

#endif
35 changes: 34 additions & 1 deletion mlir/lib/IR/PatternMatch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,35 @@ void RewriterBase::replaceAllOpUsesWith(Operation *from, Operation *to) {
replaceAllUsesWith(from->getResults(), to->getResults());
}

ReplaceOpAction::ReplaceOpAction(ArrayRef<IRUnit> irUnits,
ValueRange replacement)
: Base(irUnits), replacement(replacement) {
assert(irUnits.size() == 1);
assert(irUnits[0]);
assert(isa<Operation *>(irUnits[0]));
}

void ReplaceOpAction::print(raw_ostream &os) const {
OpPrintingFlags flags;
flags.elideLargeElementsAttrs(10);
os << "`" << tag << "` replacing operation `";
getOp()->print(os, flags);
os << "` by ";
bool first = true;
for (auto r : replacement) {
if (!first)
os << ", ";
os << "`";
r.print(os, flags);
os << "`";
first = false;
}
}

Operation *ReplaceOpAction::getOp() const {
return cast<Operation *>(getContextIRUnits()[0]);
}

/// This method replaces the results of the operation with the specified list of
/// values. The number of provided values must match the number of results of
/// the operation. The replaced op is erased.
Expand Down Expand Up @@ -265,8 +294,12 @@ void RewriterBase::replaceUsesWithIf(Value from, Value to,
bool allReplaced = true;
for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
bool replace = functor(operand);
if (replace)
if (replace) {
if (auto *fromOp = from.getDefiningOp())
getContext()->executeAction<ReplaceOpAction>(
[]() {}, ArrayRef<IRUnit>{fromOp}, to);
modifyOpInPlace(operand.getOwner(), [&]() { operand.set(to); });
}
allReplaced &= replace;
}
if (allUsesReplaced)
Expand Down