134 lines
4.7 KiB
C++
134 lines
4.7 KiB
C++
//===- SymCache.h - Declare Symbol Cache ------------------------*- 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 a Symbol Cache.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#ifndef MLIR_SUPPORT_SYMCACHE_H
|
|
#define MLIR_SUPPORT_SYMCACHE_H
|
|
|
|
#include "mlir/IR/SymbolTable.h"
|
|
#include "llvm/ADT/iterator.h"
|
|
#include "llvm/Support/Casting.h"
|
|
|
|
namespace mlir {
|
|
|
|
/// Base symbol cache class to allow for cache lookup through a pointer to some
|
|
/// abstract cache. A symbol cache stores lookup tables to make manipulating and
|
|
/// working with the IR more efficient.
|
|
class SymbolCacheBase {
|
|
public:
|
|
virtual ~SymbolCacheBase();
|
|
|
|
/// Defines 'op' as associated with the 'symbol' in the cache.
|
|
virtual void addDefinition(mlir::Attribute symbol, mlir::Operation *op) = 0;
|
|
|
|
/// Adds the symbol-defining 'op' to the cache.
|
|
void addSymbol(mlir::SymbolOpInterface op) {
|
|
addDefinition(op.getNameAttr(), op);
|
|
}
|
|
|
|
/// Populate the symbol cache with all symbol-defining operations within the
|
|
/// 'top' operation.
|
|
void addDefinitions(mlir::Operation *top);
|
|
|
|
/// Lookup a definition for 'symbol' in the cache.
|
|
virtual mlir::Operation *getDefinition(mlir::Attribute symbol) const = 0;
|
|
|
|
/// Lookup a definition for 'symbol' in the cache.
|
|
mlir::Operation *getDefinition(mlir::FlatSymbolRefAttr symbol) const {
|
|
return getDefinition(symbol.getAttr());
|
|
}
|
|
|
|
/// Iterator support through a pointer to some abstract cache.
|
|
/// The implementing cache must provide an iterator that carries values on the
|
|
/// form of <mlir::Attribute, mlir::Operation*>.
|
|
using CacheItem = std::pair<mlir::Attribute, mlir::Operation *>;
|
|
struct CacheIteratorImpl {
|
|
virtual ~CacheIteratorImpl() {}
|
|
virtual void operator++() = 0;
|
|
virtual CacheItem operator*() = 0;
|
|
virtual bool operator==(CacheIteratorImpl *other) = 0;
|
|
};
|
|
|
|
struct Iterator
|
|
: public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
|
|
CacheItem> {
|
|
Iterator(std::unique_ptr<CacheIteratorImpl> &&impl)
|
|
: impl(std::move(impl)) {}
|
|
CacheItem operator*() const { return **impl; }
|
|
using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
|
|
CacheItem>::operator++;
|
|
bool operator==(const Iterator &other) const {
|
|
return *impl == other.impl.get();
|
|
}
|
|
void operator++() { impl->operator++(); }
|
|
|
|
private:
|
|
std::unique_ptr<CacheIteratorImpl> impl;
|
|
};
|
|
virtual Iterator begin() = 0;
|
|
virtual Iterator end() = 0;
|
|
};
|
|
|
|
/// Default symbol cache implementation; stores associations between names
|
|
/// (StringAttr's) to mlir::Operation's.
|
|
/// Adding/getting definitions from the symbol cache is not
|
|
/// thread safe. If this is required, synchronizing cache acccess should be
|
|
/// ensured by the caller.
|
|
class SymbolCache : public SymbolCacheBase {
|
|
public:
|
|
/// In the building phase, add symbols.
|
|
void addDefinition(mlir::Attribute key, mlir::Operation *op) override {
|
|
symbolCache.try_emplace(key, op);
|
|
}
|
|
|
|
// Pull in getDefinition(mlir::FlatSymbolRefAttr symbol)
|
|
using SymbolCacheBase::getDefinition;
|
|
mlir::Operation *getDefinition(mlir::Attribute attr) const override {
|
|
auto it = symbolCache.find(attr);
|
|
if (it == symbolCache.end())
|
|
return nullptr;
|
|
return it->second;
|
|
}
|
|
|
|
protected:
|
|
/// This stores a lookup table from symbol attribute to the operation
|
|
/// that defines it.
|
|
llvm::DenseMap<mlir::Attribute, mlir::Operation *> symbolCache;
|
|
|
|
private:
|
|
/// Iterator support: A simple mapping between decltype(symbolCache)::iterator
|
|
/// to SymbolCacheBase::Iterator.
|
|
using Iterator = decltype(symbolCache)::iterator;
|
|
struct SymbolCacheIteratorImpl : public CacheIteratorImpl {
|
|
SymbolCacheIteratorImpl(Iterator it) : it(it) {}
|
|
CacheItem operator*() override { return {it->getFirst(), it->getSecond()}; }
|
|
void operator++() override { it++; }
|
|
bool operator==(CacheIteratorImpl *other) override {
|
|
return it == static_cast<SymbolCacheIteratorImpl *>(other)->it;
|
|
}
|
|
Iterator it;
|
|
};
|
|
|
|
public:
|
|
SymbolCacheBase::Iterator begin() override {
|
|
return SymbolCacheBase::Iterator(
|
|
std::make_unique<SymbolCacheIteratorImpl>(symbolCache.begin()));
|
|
}
|
|
SymbolCacheBase::Iterator end() override {
|
|
return SymbolCacheBase::Iterator(
|
|
std::make_unique<SymbolCacheIteratorImpl>(symbolCache.end()));
|
|
}
|
|
};
|
|
|
|
} // namespace mlir
|
|
|
|
#endif // MLIR_SUPPORT_SYMCACHE_H
|