Skip to content

[mlir][NFC] Refactor eraseState to take constant time #121670

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
Jan 6, 2025
Merged

[mlir][NFC] Refactor eraseState to take constant time #121670

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
23 changes: 9 additions & 14 deletions mlir/include/mlir/Analysis/DataFlowFramework.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,9 +332,11 @@ class DataFlowSolver {
/// does not exist.
template <typename StateT, typename AnchorT>
const StateT *lookupState(AnchorT anchor) const {
auto it =
analysisStates.find({LatticeAnchor(anchor), TypeID::get<StateT>()});
if (it == analysisStates.end())
const auto &mapIt = analysisStates.find(LatticeAnchor(anchor));
if (mapIt == analysisStates.end())
return nullptr;
auto it = mapIt->second.find(TypeID::get<StateT>());
if (it == mapIt->second.end())
return nullptr;
return static_cast<const StateT *>(it->second.get());
}
Expand All @@ -343,11 +345,7 @@ class DataFlowSolver {
template <typename AnchorT>
void eraseState(AnchorT anchor) {
LatticeAnchor la(anchor);

for (auto it = analysisStates.begin(); it != analysisStates.end(); ++it) {
if (it->first.first == la)
analysisStates.erase(it);
}
analysisStates.erase(LatticeAnchor(anchor));
}

// Erase all analysis states
Expand Down Expand Up @@ -426,7 +424,8 @@ class DataFlowSolver {

/// A type-erased map of lattice anchors to associated analysis states for
/// first-class lattice anchors.
DenseMap<std::pair<LatticeAnchor, TypeID>, std::unique_ptr<AnalysisState>>
DenseMap<LatticeAnchor, DenseMap<TypeID, std::unique_ptr<AnalysisState>>,
DenseMapInfo<LatticeAnchor::ParentTy>>
analysisStates;

/// Allow the base child analysis class to access the internals of the solver.
Expand Down Expand Up @@ -643,7 +642,7 @@ AnalysisT *DataFlowSolver::load(Args &&...args) {
template <typename StateT, typename AnchorT>
StateT *DataFlowSolver::getOrCreateState(AnchorT anchor) {
std::unique_ptr<AnalysisState> &state =
analysisStates[{LatticeAnchor(anchor), TypeID::get<StateT>()}];
analysisStates[LatticeAnchor(anchor)][TypeID::get<StateT>()];
if (!state) {
state = std::unique_ptr<StateT>(new StateT(anchor));
#if LLVM_ENABLE_ABI_BREAKING_CHECKS
Expand Down Expand Up @@ -689,10 +688,6 @@ struct DenseMapInfo<mlir::ProgramPoint> {
}
};

template <>
struct DenseMapInfo<mlir::LatticeAnchor>
: public DenseMapInfo<mlir::LatticeAnchor::ParentTy> {};

// Allow llvm::cast style functions.
template <typename To>
struct CastInfo<To, mlir::LatticeAnchor>
Expand Down
Loading