Skip to content

Commit

Permalink
[mlir][Transforms][NFC] CSE: Add C++ entry point
Browse files Browse the repository at this point in the history 
* All IR modifications are done with a rewriter.
* The new C++ entry point takes a `RewriterBase &`, which may have a listener attached to it.

This revision is useful because it allows users to run CSE and track IR modifications via a listener that can be attached to the rewriter.

This is a reupload. The original CL was reverted (9979417) due to a memory leak. The memory leak is unrelated to this change and fixed with D154185.

Differential Revision: https://reviews.llvm.org/D145226
  • Loading branch information
@matthias-springer
matthias-springer committed Jul 3, 2023
1 parent c487fe3 commit b9bdff4
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 30 deletions.
6 changes: 6 additions & 0 deletions mlir/include/mlir/IR/PatternMatch.h
View file
Expand Up @@ -627,6 +627,12 @@ class RewriterBase : public OpBuilder {
/// in-place operation modification is about to happen.
void replaceUsesWithIf(Value from, Value to,
function_ref<bool(OpOperand &)> functor);
void replaceUsesWithIf(ValueRange from, ValueRange to,
function_ref<bool(OpOperand &)> functor) {
assert(from.size() == to.size() && "incorrect number of replacements");
for (auto it : llvm::zip(from, to))
replaceUsesWithIf(std::get<0>(it), std::get<1>(it), functor);
}

/// Find uses of `from` and replace them with `to` except if the user is
/// `exceptedUser`. It also marks every modified uses and notifies the
Expand Down
32 changes: 32 additions & 0 deletions mlir/include/mlir/Transforms/CSE.h
View file
@@ -0,0 +1,32 @@
//===- CSE.h - Common Subexpression Elimination -----------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file declares methods for eliminating common subexpressions.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_TRANSFORMS_CSE_H_
#define MLIR_TRANSFORMS_CSE_H_

namespace mlir {

class DominanceInfo;
class Operation;
class RewriterBase;

/// Eliminate common subexpressions within the given operation. This transform
/// looks for and deduplicates equivalent operations.
///
/// `changed` indicates whether the IR was modified or not.
void eliminateCommonSubExpressions(RewriterBase &rewriter,
DominanceInfo &domInfo, Operation *op,
bool *changed = nullptr);

} // namespace mlir

#endif // MLIR_TRANSFORMS_CSE_H_
120 changes: 90 additions & 30 deletions mlir/lib/Transforms/CSE.cpp
View file
Expand Up @@ -11,11 +11,13 @@
//
//===----------------------------------------------------------------------===//

#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/CSE.h"

#include "mlir/IR/Dominance.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/Passes.h"
#include "llvm/ADT/DenseMapInfo.h"
#include "llvm/ADT/Hashing.h"
#include "llvm/ADT/ScopedHashTable.h"
Expand Down Expand Up @@ -56,7 +58,18 @@ struct SimpleOperationInfo : public llvm::DenseMapInfo<Operation *> {

namespace {
/// Simple common sub-expression elimination.
struct CSE : public impl::CSEBase<CSE> {
class CSEDriver {
public:
CSEDriver(RewriterBase &rewriter, DominanceInfo *domInfo)
: rewriter(rewriter), domInfo(domInfo) {}

/// Simplify all operations within the given op.
void simplify(Operation *op, bool *changed = nullptr);

int64_t getNumCSE() const { return numCSE; }
int64_t getNumDCE() const { return numDCE; }

private:
/// Shared implementation of operation elimination and scoped map definitions.
using AllocatorTy = llvm::RecyclingAllocator<
llvm::BumpPtrAllocator,
Expand Down Expand Up @@ -94,39 +107,59 @@ struct CSE : public impl::CSEBase<CSE> {
void simplifyBlock(ScopedMapTy &knownValues, Block *bb, bool hasSSADominance);
void simplifyRegion(ScopedMapTy &knownValues, Region &region);

void runOnOperation() override;

private:
void replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
Operation *existing, bool hasSSADominance);

/// Check if there is side-effecting operations other than the given effect
/// between the two operations.
bool hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp);

/// A rewriter for modifying the IR.
RewriterBase &rewriter;

/// Operations marked as dead and to be erased.
std::vector<Operation *> opsToErase;
DominanceInfo *domInfo = nullptr;
MemEffectsCache memEffectsCache;

// Various statistics.
int64_t numCSE = 0;
int64_t numDCE = 0;
};
} // namespace

void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
Operation *existing, bool hasSSADominance) {
void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
Operation *existing,
bool hasSSADominance) {
// If we find one then replace all uses of the current operation with the
// existing one and mark it for deletion. We can only replace an operand in
// an operation if it has not been visited yet.
if (hasSSADominance) {
// If the region has SSA dominance, then we are guaranteed to have not
// visited any use of the current operation.
op->replaceAllUsesWith(existing);
if (auto *rewriteListener =
dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
rewriteListener->notifyOperationReplaced(op, existing);
// Replace all uses, but do not remote the operation yet. This does not
// notify the listener because the original op is not erased.
rewriter.replaceAllUsesWith(op->getResults(), existing->getResults());
opsToErase.push_back(op);
} else {
// When the region does not have SSA dominance, we need to check if we
// have visited a use before replacing any use.
op->replaceUsesWithIf(existing->getResults(), [&](OpOperand &operand) {
auto wasVisited = [&](OpOperand &operand) {
return !knownValues.count(operand.getOwner());
});
};
if (auto *rewriteListener =
dyn_cast_if_present<RewriterBase::Listener>(rewriter.getListener()))
for (Value v : op->getResults())
if (all_of(v.getUses(), wasVisited))
rewriteListener->notifyOperationReplaced(op, existing);

// Replace all uses, but do not remote the operation yet. This does not
// notify the listener because the original op is not erased.
rewriter.replaceUsesWithIf(op->getResults(), existing->getResults(),
wasVisited);

// There may be some remaining uses of the operation.
if (op->use_empty())
Expand All @@ -142,7 +175,8 @@ void CSE::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
++numCSE;
}

bool CSE::hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp) {
bool CSEDriver::hasOtherSideEffectingOpInBetween(Operation *fromOp,
Operation *toOp) {
assert(fromOp->getBlock() == toOp->getBlock());
assert(
isa<MemoryEffectOpInterface>(fromOp) &&
Expand Down Expand Up @@ -183,8 +217,9 @@ bool CSE::hasOtherSideEffectingOpInBetween(Operation *fromOp, Operation *toOp) {
}

/// Attempt to eliminate a redundant operation.
LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
bool hasSSADominance) {
LogicalResult CSEDriver::simplifyOperation(ScopedMapTy &knownValues,
Operation *op,
bool hasSSADominance) {
// Don't simplify terminator operations.
if (op->hasTrait<OpTrait::IsTerminator>())
return failure();
Expand Down Expand Up @@ -240,8 +275,8 @@ LogicalResult CSE::simplifyOperation(ScopedMapTy &knownValues, Operation *op,
return failure();
}

void CSE::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
bool hasSSADominance) {
void CSEDriver::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
bool hasSSADominance) {
for (auto &op : *bb) {
// Most operations don't have regions, so fast path that case.
if (op.getNumRegions() != 0) {
Expand All @@ -267,7 +302,7 @@ void CSE::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
memEffectsCache.clear();
}

void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
void CSEDriver::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
// If the region is empty there is nothing to do.
if (region.empty())
return;
Expand Down Expand Up @@ -322,29 +357,54 @@ void CSE::simplifyRegion(ScopedMapTy &knownValues, Region &region) {
}
}

void CSE::runOnOperation() {
/// A scoped hash table of defining operations within a region.
void CSEDriver::simplify(Operation *op, bool *changed) {
/// Simplify all regions.
ScopedMapTy knownValues;

domInfo = &getAnalysis<DominanceInfo>();
Operation *rootOp = getOperation();

for (auto &region : rootOp->getRegions())
for (auto &region : op->getRegions())
simplifyRegion(knownValues, region);

// If no operations were erased, then we mark all analyses as preserved.
if (opsToErase.empty())
return markAllAnalysesPreserved();

/// Erase any operations that were marked as dead during simplification.
for (auto *op : opsToErase)
op->erase();
opsToErase.clear();
rewriter.eraseOp(op);
if (changed)
*changed = !opsToErase.empty();

// Note: CSE does currently not remove ops with regions, so DominanceInfo
// does not have to be invalidated.
}

void mlir::eliminateCommonSubExpressions(RewriterBase &rewriter,
DominanceInfo &domInfo, Operation *op,
bool *changed) {
CSEDriver driver(rewriter, &domInfo);
driver.simplify(op, changed);
}

namespace {
/// CSE pass.
struct CSE : public impl::CSEBase<CSE> {
void runOnOperation() override;
};
} // namespace

void CSE::runOnOperation() {
// Simplify the IR.
IRRewriter rewriter(&getContext());
CSEDriver driver(rewriter, &getAnalysis<DominanceInfo>());
bool changed = false;
driver.simplify(getOperation(), &changed);

// Set statistics.
numCSE = driver.getNumCSE();
numDCE = driver.getNumDCE();

// If there was no change to the IR, we mark all analyses as preserved.
if (!changed)
return markAllAnalysesPreserved();

// We currently don't remove region operations, so mark dominance as
// preserved.
markAnalysesPreserved<DominanceInfo, PostDominanceInfo>();
domInfo = nullptr;
}

std::unique_ptr<Pass> mlir::createCSEPass() { return std::make_unique<CSE>(); }

0 comments on commit b9bdff4

Please sign in to comment.