Skip to content

[mlir][Transforms] GreedyPatternRewriteDriver: Hash ops separately #78312

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

Merged
merged 1 commit into from
Feb 1, 2024
Merged

[mlir][Transforms] GreedyPatternRewriteDriver: Hash ops separately #78312

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: 2 additions & 2 deletions mlir/include/mlir/IR/OperationSupport.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1316,10 +1316,10 @@ LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE();
//===----------------------------------------------------------------------===//

/// A unique fingerprint for a specific operation, and all of it's internal
/// operations.
/// operations (if `includeNested` is set).
class OperationFingerPrint {
public:
OperationFingerPrint(Operation *topOp);
OperationFingerPrint(Operation *topOp, bool includeNested = true);
OperationFingerPrint(const OperationFingerPrint &) = default;
OperationFingerPrint &operator=(const OperationFingerPrint &) = default;

Expand Down
15 changes: 11 additions & 4 deletions mlir/lib/IR/OperationSupport.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -911,11 +911,12 @@ static void addDataToHash(llvm::SHA1 &hasher, const T &data) {
ArrayRef<uint8_t>(reinterpret_cast<const uint8_t *>(&data), sizeof(T)));
}

OperationFingerPrint::OperationFingerPrint(Operation *topOp) {
OperationFingerPrint::OperationFingerPrint(Operation *topOp,
bool includeNested) {
llvm::SHA1 hasher;

// Hash each of the operations based upon their mutable bits:
topOp->walk([&](Operation *op) {
// Helper function that hashes an operation based on its mutable bits:
auto addOperationToHash = [&](Operation *op) {
// - Operation pointer
addDataToHash(hasher, op);
// - Parent operation pointer (to take into account the nesting structure)
Expand Down Expand Up @@ -944,6 +945,12 @@ OperationFingerPrint::OperationFingerPrint(Operation *topOp) {
// - Result types
for (Type t : op->getResultTypes())
addDataToHash(hasher, t);
});
};

if (includeNested)
topOp->walk(addOperationToHash);
else
addOperationToHash(topOp);

hash = hasher.result();
}
31 changes: 16 additions & 15 deletions mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,9 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
void computeFingerPrints(Operation *topLevel) {
this->topLevel = topLevel;
this->topLevelFingerPrint.emplace(topLevel);
topLevel->walk([&](Operation *op) { fingerprints.try_emplace(op, op); });
topLevel->walk([&](Operation *op) {
fingerprints.try_emplace(op, op, /*includeNested=*/false);
});
}

/// Clear all finger prints.
Expand Down Expand Up @@ -95,7 +97,8 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {
// API.) Finger print computation does may not crash if a new op was
// created at the same memory location. (But then the finger print should
// have changed.)
if (it.second != OperationFingerPrint(it.first)) {
if (it.second !=
OperationFingerPrint(it.first, /*includeNested=*/false)) {
// Note: Run "mlir-opt -debug" to see which pattern is broken.
llvm::report_fatal_error("operation finger print changed");
}
Expand Down Expand Up @@ -125,25 +128,23 @@ struct ExpensiveChecks : public RewriterBase::ForwardingListener {

protected:
/// Invalidate the finger print of the given op, i.e., remove it from the map.
void invalidateFingerPrint(Operation *op) {
// Invalidate all finger prints until the top level.
while (op && op != topLevel) {
fingerprints.erase(op);
op = op->getParentOp();
}
void invalidateFingerPrint(Operation *op) { fingerprints.erase(op); }

void notifyBlockRemoved(Block *block) override {
RewriterBase::ForwardingListener::notifyBlockRemoved(block);

// The block structure (number of blocks, types of block arguments, etc.)
// is part of the fingerprint of the parent op.
// TODO: The parent op fingerprint should also be invalidated when modifying
// the block arguments of a block, but we do not have a
// `notifyBlockModified` callback yet.
invalidateFingerPrint(block->getParentOp());
}

void notifyOperationInserted(Operation *op,
OpBuilder::InsertPoint previous) override {
RewriterBase::ForwardingListener::notifyOperationInserted(op, previous);
// Invalidate the finger print of the op that owns the block into which the
// op was inserted into.
invalidateFingerPrint(op->getParentOp());

// Also invalidate the finger print of the op that owns the block from which
// the op was moved from. (Only applicable if the op was moved.)
if (previous.isSet())
invalidateFingerPrint(previous.getBlock()->getParentOp());
}

void notifyOperationModified(Operation *op) override {
Expand Down