332 lines · cpp
1//===- Dialect.cpp - Dialect implementation -------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/IR/Dialect.h"10#include "mlir/IR/BuiltinDialect.h"11#include "mlir/IR/Diagnostics.h"12#include "mlir/IR/DialectImplementation.h"13#include "mlir/IR/DialectInterface.h"14#include "mlir/IR/DialectRegistry.h"15#include "mlir/IR/ExtensibleDialect.h"16#include "mlir/IR/MLIRContext.h"17#include "mlir/IR/Operation.h"18#include "mlir/Support/TypeID.h"19#include "llvm/ADT/MapVector.h"20#include "llvm/ADT/SmallVectorExtras.h"21#include "llvm/ADT/Twine.h"22#include "llvm/Support/DebugLog.h"23#include "llvm/Support/Regex.h"24#include <memory>25 26#define DEBUG_TYPE "dialect"27 28using namespace mlir;29using namespace detail;30 31//===----------------------------------------------------------------------===//32// Dialect33//===----------------------------------------------------------------------===//34 35Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)36 : name(name), dialectID(id), context(context) {37 assert(isValidNamespace(name) && "invalid dialect namespace");38}39 40Dialect::~Dialect() = default;41 42/// Verify an attribute from this dialect on the argument at 'argIndex' for43/// the region at 'regionIndex' on the given operation. Returns failure if44/// the verification failed, success otherwise. This hook may optionally be45/// invoked from any operation containing a region.46LogicalResult Dialect::verifyRegionArgAttribute(Operation *, unsigned, unsigned,47 NamedAttribute) {48 return success();49}50 51/// Verify an attribute from this dialect on the result at 'resultIndex' for52/// the region at 'regionIndex' on the given operation. Returns failure if53/// the verification failed, success otherwise. This hook may optionally be54/// invoked from any operation containing a region.55LogicalResult Dialect::verifyRegionResultAttribute(Operation *, unsigned,56 unsigned, NamedAttribute) {57 return success();58}59 60/// Parse an attribute registered to this dialect.61Attribute Dialect::parseAttribute(DialectAsmParser &parser, Type type) const {62 parser.emitError(parser.getNameLoc())63 << "dialect '" << getNamespace()64 << "' provides no attribute parsing hook";65 return Attribute();66}67 68/// Parse a type registered to this dialect.69Type Dialect::parseType(DialectAsmParser &parser) const {70 // If this dialect allows unknown types, then represent this with OpaqueType.71 if (allowsUnknownTypes()) {72 StringAttr ns = StringAttr::get(getContext(), getNamespace());73 return OpaqueType::get(ns, parser.getFullSymbolSpec());74 }75 76 parser.emitError(parser.getNameLoc())77 << "dialect '" << getNamespace() << "' provides no type parsing hook";78 return Type();79}80 81std::optional<Dialect::ParseOpHook>82Dialect::getParseOperationHook(StringRef opName) const {83 return std::nullopt;84}85 86llvm::unique_function<void(Operation *, OpAsmPrinter &printer)>87Dialect::getOperationPrinter(Operation *op) const {88 assert(op->getDialect() == this &&89 "Dialect hook invoked on non-dialect owned operation");90 return nullptr;91}92 93/// Utility function that returns if the given string is a valid dialect94/// namespace95bool Dialect::isValidNamespace(StringRef str) {96 llvm::Regex dialectNameRegex("^[a-zA-Z_][a-zA-Z_0-9\\$]*$");97 return dialectNameRegex.match(str);98}99 100/// Register a set of dialect interfaces with this dialect instance.101void Dialect::addInterface(std::unique_ptr<DialectInterface> interface) {102 // Handle the case where the models resolve a promised interface.103 handleAdditionOfUndefinedPromisedInterface(getTypeID(), interface->getID());104 105 auto it = registeredInterfaces.try_emplace(interface->getID(),106 std::move(interface));107 if (!it.second)108 LDBG() << "repeated interface registration for dialect " << getNamespace();109}110 111//===----------------------------------------------------------------------===//112// Dialect Interface113//===----------------------------------------------------------------------===//114 115DialectInterface::~DialectInterface() = default;116 117MLIRContext *DialectInterface::getContext() const {118 return dialect->getContext();119}120 121DialectInterfaceCollectionBase::DialectInterfaceCollectionBase(122 MLIRContext *ctx, TypeID interfaceKind, StringRef interfaceName) {123 for (auto *dialect : ctx->getLoadedDialects()) {124#ifndef NDEBUG125 dialect->handleUseOfUndefinedPromisedInterface(126 dialect->getTypeID(), interfaceKind, interfaceName);127#endif128 if (auto *interface = dialect->getRegisteredInterface(interfaceKind)) {129 interfaces.insert(interface);130 orderedInterfaces.push_back(interface);131 }132 }133}134 135DialectInterfaceCollectionBase::~DialectInterfaceCollectionBase() = default;136 137/// Get the interface for the dialect of given operation, or null if one138/// is not registered.139const DialectInterface *140DialectInterfaceCollectionBase::getInterfaceFor(Operation *op) const {141 return getInterfaceFor(op->getDialect());142}143 144//===----------------------------------------------------------------------===//145// DialectExtension146//===----------------------------------------------------------------------===//147 148DialectExtensionBase::~DialectExtensionBase() = default;149 150void dialect_extension_detail::handleUseOfUndefinedPromisedInterface(151 Dialect &dialect, TypeID interfaceRequestorID, TypeID interfaceID,152 StringRef interfaceName) {153 dialect.handleUseOfUndefinedPromisedInterface(interfaceRequestorID,154 interfaceID, interfaceName);155}156 157void dialect_extension_detail::handleAdditionOfUndefinedPromisedInterface(158 Dialect &dialect, TypeID interfaceRequestorID, TypeID interfaceID) {159 dialect.handleAdditionOfUndefinedPromisedInterface(interfaceRequestorID,160 interfaceID);161}162 163bool dialect_extension_detail::hasPromisedInterface(Dialect &dialect,164 TypeID interfaceRequestorID,165 TypeID interfaceID) {166 return dialect.hasPromisedInterface(interfaceRequestorID, interfaceID);167}168 169//===----------------------------------------------------------------------===//170// DialectRegistry171//===----------------------------------------------------------------------===//172 173namespace {174template <typename Fn>175void applyExtensionsFn(176 Fn &&applyExtension,177 const llvm::MapVector<TypeID, std::unique_ptr<DialectExtensionBase>>178 &extensions) {179 // Note: Additional extensions may be added while applying an extension.180 // The iterators will be invalidated if extensions are added so we'll keep181 // a copy of the extensions for ourselves.182 183 const auto extractExtension =184 [](const auto &entry) -> DialectExtensionBase * {185 return entry.second.get();186 };187 188 auto startIt = extensions.begin(), endIt = extensions.end();189 size_t count = 0;190 while (startIt != endIt) {191 count += endIt - startIt;192 193 // Grab the subset of extensions we'll apply in this iteration.194 const auto subset =195 llvm::map_to_vector(llvm::make_range(startIt, endIt), extractExtension);196 197 for (const auto *ext : subset)198 applyExtension(*ext);199 200 // Book-keep for the next iteration.201 startIt = extensions.begin() + count;202 endIt = extensions.end();203 }204}205} // namespace206 207DialectRegistry::DialectRegistry() { insert<BuiltinDialect>(); }208 209DialectAllocatorFunctionRef210DialectRegistry::getDialectAllocator(StringRef name) const {211 auto it = registry.find(name);212 if (it == registry.end())213 return nullptr;214 return it->second.second;215}216 217void DialectRegistry::insert(TypeID typeID, StringRef name,218 const DialectAllocatorFunction &ctor) {219 auto inserted = registry.insert(220 std::make_pair(std::string(name), std::make_pair(typeID, ctor)));221 if (!inserted.second && inserted.first->second.first != typeID) {222 llvm::report_fatal_error(223 "Trying to register different dialects for the same namespace: " +224 name);225 }226}227 228void DialectRegistry::insertDynamic(229 StringRef name, const DynamicDialectPopulationFunction &ctor) {230 // This TypeID marks dynamic dialects. We cannot give a TypeID for the231 // dialect yet, since the TypeID of a dynamic dialect is defined at its232 // construction.233 TypeID typeID = TypeID::get<void>();234 235 // Create the dialect, and then call ctor, which allocates its components.236 auto constructor = [nameStr = name.str(), ctor](MLIRContext *ctx) {237 auto *dynDialect = ctx->getOrLoadDynamicDialect(238 nameStr, [ctx, ctor](DynamicDialect *dialect) { ctor(ctx, dialect); });239 assert(dynDialect && "Dynamic dialect creation unexpectedly failed");240 return dynDialect;241 };242 243 insert(typeID, name, constructor);244}245 246void DialectRegistry::applyExtensions(Dialect *dialect) const {247 MLIRContext *ctx = dialect->getContext();248 StringRef dialectName = dialect->getNamespace();249 250 // Functor used to try to apply the given extension.251 auto applyExtension = [&](const DialectExtensionBase &extension) {252 ArrayRef<StringRef> dialectNames = extension.getRequiredDialects();253 // An empty set is equivalent to always invoke.254 if (dialectNames.empty()) {255 extension.apply(ctx, dialect);256 return;257 }258 259 // Handle the simple case of a single dialect name. In this case, the260 // required dialect should be the current dialect.261 if (dialectNames.size() == 1) {262 if (dialectNames.front() == dialectName)263 extension.apply(ctx, dialect);264 return;265 }266 267 // Otherwise, check to see if this extension requires this dialect.268 const StringRef *nameIt = llvm::find(dialectNames, dialectName);269 if (nameIt == dialectNames.end())270 return;271 272 // If it does, ensure that all of the other required dialects have been273 // loaded.274 SmallVector<Dialect *> requiredDialects;275 requiredDialects.reserve(dialectNames.size());276 for (auto it = dialectNames.begin(), e = dialectNames.end(); it != e;277 ++it) {278 // The current dialect is known to be loaded.279 if (it == nameIt) {280 requiredDialects.push_back(dialect);281 continue;282 }283 // Otherwise, check if it is loaded.284 Dialect *loadedDialect = ctx->getLoadedDialect(*it);285 if (!loadedDialect)286 return;287 requiredDialects.push_back(loadedDialect);288 }289 extension.apply(ctx, requiredDialects);290 };291 292 applyExtensionsFn(applyExtension, extensions);293}294 295void DialectRegistry::applyExtensions(MLIRContext *ctx) const {296 // Functor used to try to apply the given extension.297 auto applyExtension = [&](const DialectExtensionBase &extension) {298 ArrayRef<StringRef> dialectNames = extension.getRequiredDialects();299 if (dialectNames.empty()) {300 auto loadedDialects = ctx->getLoadedDialects();301 extension.apply(ctx, loadedDialects);302 return;303 }304 305 // Check to see if all of the dialects for this extension are loaded.306 SmallVector<Dialect *> requiredDialects;307 requiredDialects.reserve(dialectNames.size());308 for (StringRef dialectName : dialectNames) {309 Dialect *loadedDialect = ctx->getLoadedDialect(dialectName);310 if (!loadedDialect)311 return;312 requiredDialects.push_back(loadedDialect);313 }314 extension.apply(ctx, requiredDialects);315 };316 317 applyExtensionsFn(applyExtension, extensions);318}319 320bool DialectRegistry::isSubsetOf(const DialectRegistry &rhs) const {321 // Check that all extension keys are present in 'rhs'.322 const auto hasExtension = [&](const auto &key) {323 return rhs.extensions.contains(key);324 };325 if (!llvm::all_of(make_first_range(extensions), hasExtension))326 return false;327 328 // Check that the current dialects fully overlap with the dialects in 'rhs'.329 return llvm::all_of(330 registry, [&](const auto &it) { return rhs.registry.count(it.first); });331}332