596 lines · cpp
1//===- ExtensibleDialect.cpp - Extensible dialect ---------------*- C++ -*-===//2//3// This file is licensed 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/ExtensibleDialect.h"10#include "mlir/IR/AttributeSupport.h"11#include "mlir/IR/DialectImplementation.h"12#include "mlir/IR/OperationSupport.h"13#include "mlir/IR/StorageUniquerSupport.h"14#include "llvm/Support/InterleavedRange.h"15 16using namespace mlir;17 18//===----------------------------------------------------------------------===//19// Dynamic types and attributes shared functions20//===----------------------------------------------------------------------===//21 22/// Default parser for dynamic attribute or type parameters.23/// Parse in the format '(<>)?' or '<attr (,attr)*>'.24static LogicalResult25typeOrAttrParser(AsmParser &parser, SmallVectorImpl<Attribute> &parsedParams) {26 // No parameters27 if (parser.parseOptionalLess() || !parser.parseOptionalGreater())28 return success();29 30 Attribute attr;31 if (parser.parseAttribute(attr))32 return failure();33 parsedParams.push_back(attr);34 35 while (parser.parseOptionalGreater()) {36 Attribute attr;37 if (parser.parseComma() || parser.parseAttribute(attr))38 return failure();39 parsedParams.push_back(attr);40 }41 42 return success();43}44 45/// Default printer for dynamic attribute or type parameters.46/// Print in the format '(<>)?' or '<attr (,attr)*>'.47static void typeOrAttrPrinter(AsmPrinter &printer, ArrayRef<Attribute> params) {48 if (params.empty())49 return;50 51 printer << "<" << llvm::interleaved(params) << ">";52}53 54//===----------------------------------------------------------------------===//55// Dynamic type56//===----------------------------------------------------------------------===//57 58std::unique_ptr<DynamicTypeDefinition>59DynamicTypeDefinition::get(StringRef name, ExtensibleDialect *dialect,60 VerifierFn &&verifier) {61 return DynamicTypeDefinition::get(name, dialect, std::move(verifier),62 typeOrAttrParser, typeOrAttrPrinter);63}64 65std::unique_ptr<DynamicTypeDefinition>66DynamicTypeDefinition::get(StringRef name, ExtensibleDialect *dialect,67 VerifierFn &&verifier, ParserFn &&parser,68 PrinterFn &&printer) {69 return std::unique_ptr<DynamicTypeDefinition>(70 new DynamicTypeDefinition(name, dialect, std::move(verifier),71 std::move(parser), std::move(printer)));72}73 74DynamicTypeDefinition::DynamicTypeDefinition(StringRef nameRef,75 ExtensibleDialect *dialect,76 VerifierFn &&verifier,77 ParserFn &&parser,78 PrinterFn &&printer)79 : name(nameRef), dialect(dialect), verifier(std::move(verifier)),80 parser(std::move(parser)), printer(std::move(printer)),81 ctx(dialect->getContext()) {}82 83DynamicTypeDefinition::DynamicTypeDefinition(ExtensibleDialect *dialect,84 StringRef nameRef)85 : name(nameRef), dialect(dialect), ctx(dialect->getContext()) {}86 87void DynamicTypeDefinition::registerInTypeUniquer() {88 detail::TypeUniquer::registerType<DynamicType>(&getContext(), getTypeID());89}90 91namespace mlir {92namespace detail {93/// Storage of DynamicType.94/// Contains a pointer to the type definition and type parameters.95struct DynamicTypeStorage : public TypeStorage {96 97 using KeyTy = std::pair<DynamicTypeDefinition *, ArrayRef<Attribute>>;98 99 explicit DynamicTypeStorage(DynamicTypeDefinition *typeDef,100 ArrayRef<Attribute> params)101 : typeDef(typeDef), params(params) {}102 103 bool operator==(const KeyTy &key) const {104 return typeDef == key.first && params == key.second;105 }106 107 static llvm::hash_code hashKey(const KeyTy &key) {108 return llvm::hash_value(key);109 }110 111 static DynamicTypeStorage *construct(TypeStorageAllocator &alloc,112 const KeyTy &key) {113 return new (alloc.allocate<DynamicTypeStorage>())114 DynamicTypeStorage(key.first, alloc.copyInto(key.second));115 }116 117 /// Definition of the type.118 DynamicTypeDefinition *typeDef;119 120 /// The type parameters.121 ArrayRef<Attribute> params;122};123} // namespace detail124} // namespace mlir125 126DynamicType DynamicType::get(DynamicTypeDefinition *typeDef,127 ArrayRef<Attribute> params) {128 auto &ctx = typeDef->getContext();129 auto emitError = detail::getDefaultDiagnosticEmitFn(&ctx);130 assert(succeeded(typeDef->verify(emitError, params)));131 return detail::TypeUniquer::getWithTypeID<DynamicType>(132 &ctx, typeDef->getTypeID(), typeDef, params);133}134 135DynamicType136DynamicType::getChecked(function_ref<InFlightDiagnostic()> emitError,137 DynamicTypeDefinition *typeDef,138 ArrayRef<Attribute> params) {139 if (failed(typeDef->verify(emitError, params)))140 return {};141 auto &ctx = typeDef->getContext();142 return detail::TypeUniquer::getWithTypeID<DynamicType>(143 &ctx, typeDef->getTypeID(), typeDef, params);144}145 146DynamicTypeDefinition *DynamicType::getTypeDef() { return getImpl()->typeDef; }147 148ArrayRef<Attribute> DynamicType::getParams() { return getImpl()->params; }149 150bool DynamicType::classof(Type type) {151 return type.hasTrait<TypeTrait::IsDynamicType>();152}153 154ParseResult DynamicType::parse(AsmParser &parser,155 DynamicTypeDefinition *typeDef,156 DynamicType &parsedType) {157 SmallVector<Attribute> params;158 if (failed(typeDef->parser(parser, params)))159 return failure();160 parsedType = parser.getChecked<DynamicType>(typeDef, params);161 if (!parsedType)162 return failure();163 return success();164}165 166void DynamicType::print(AsmPrinter &printer) {167 printer << getTypeDef()->getName();168 getTypeDef()->printer(printer, getParams());169}170 171//===----------------------------------------------------------------------===//172// Dynamic attribute173//===----------------------------------------------------------------------===//174 175std::unique_ptr<DynamicAttrDefinition>176DynamicAttrDefinition::get(StringRef name, ExtensibleDialect *dialect,177 VerifierFn &&verifier) {178 return DynamicAttrDefinition::get(name, dialect, std::move(verifier),179 typeOrAttrParser, typeOrAttrPrinter);180}181 182std::unique_ptr<DynamicAttrDefinition>183DynamicAttrDefinition::get(StringRef name, ExtensibleDialect *dialect,184 VerifierFn &&verifier, ParserFn &&parser,185 PrinterFn &&printer) {186 return std::unique_ptr<DynamicAttrDefinition>(187 new DynamicAttrDefinition(name, dialect, std::move(verifier),188 std::move(parser), std::move(printer)));189}190 191DynamicAttrDefinition::DynamicAttrDefinition(StringRef nameRef,192 ExtensibleDialect *dialect,193 VerifierFn &&verifier,194 ParserFn &&parser,195 PrinterFn &&printer)196 : name(nameRef), dialect(dialect), verifier(std::move(verifier)),197 parser(std::move(parser)), printer(std::move(printer)),198 ctx(dialect->getContext()) {}199 200DynamicAttrDefinition::DynamicAttrDefinition(ExtensibleDialect *dialect,201 StringRef nameRef)202 : name(nameRef), dialect(dialect), ctx(dialect->getContext()) {}203 204void DynamicAttrDefinition::registerInAttrUniquer() {205 detail::AttributeUniquer::registerAttribute<DynamicAttr>(&getContext(),206 getTypeID());207}208 209namespace mlir {210namespace detail {211/// Storage of DynamicAttr.212/// Contains a pointer to the attribute definition and attribute parameters.213struct DynamicAttrStorage : public AttributeStorage {214 using KeyTy = std::pair<DynamicAttrDefinition *, ArrayRef<Attribute>>;215 216 explicit DynamicAttrStorage(DynamicAttrDefinition *attrDef,217 ArrayRef<Attribute> params)218 : attrDef(attrDef), params(params) {}219 220 bool operator==(const KeyTy &key) const {221 return attrDef == key.first && params == key.second;222 }223 224 static llvm::hash_code hashKey(const KeyTy &key) {225 return llvm::hash_value(key);226 }227 228 static DynamicAttrStorage *construct(AttributeStorageAllocator &alloc,229 const KeyTy &key) {230 return new (alloc.allocate<DynamicAttrStorage>())231 DynamicAttrStorage(key.first, alloc.copyInto(key.second));232 }233 234 /// Definition of the type.235 DynamicAttrDefinition *attrDef;236 237 /// The type parameters.238 ArrayRef<Attribute> params;239};240} // namespace detail241} // namespace mlir242 243DynamicAttr DynamicAttr::get(DynamicAttrDefinition *attrDef,244 ArrayRef<Attribute> params) {245 auto &ctx = attrDef->getContext();246 return detail::AttributeUniquer::getWithTypeID<DynamicAttr>(247 &ctx, attrDef->getTypeID(), attrDef, params);248}249 250DynamicAttr251DynamicAttr::getChecked(function_ref<InFlightDiagnostic()> emitError,252 DynamicAttrDefinition *attrDef,253 ArrayRef<Attribute> params) {254 if (failed(attrDef->verify(emitError, params)))255 return {};256 return get(attrDef, params);257}258 259DynamicAttrDefinition *DynamicAttr::getAttrDef() { return getImpl()->attrDef; }260 261ArrayRef<Attribute> DynamicAttr::getParams() { return getImpl()->params; }262 263bool DynamicAttr::classof(Attribute attr) {264 return attr.hasTrait<AttributeTrait::IsDynamicAttr>();265}266 267ParseResult DynamicAttr::parse(AsmParser &parser,268 DynamicAttrDefinition *attrDef,269 DynamicAttr &parsedAttr) {270 SmallVector<Attribute> params;271 if (failed(attrDef->parser(parser, params)))272 return failure();273 parsedAttr = parser.getChecked<DynamicAttr>(attrDef, params);274 if (!parsedAttr)275 return failure();276 return success();277}278 279void DynamicAttr::print(AsmPrinter &printer) {280 printer << getAttrDef()->getName();281 getAttrDef()->printer(printer, getParams());282}283 284//===----------------------------------------------------------------------===//285// Dynamic operation286//===----------------------------------------------------------------------===//287 288DynamicOpDefinition::DynamicOpDefinition(289 StringRef name, ExtensibleDialect *dialect,290 OperationName::VerifyInvariantsFn &&verifyFn,291 OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,292 OperationName::ParseAssemblyFn &&parseFn,293 OperationName::PrintAssemblyFn &&printFn,294 OperationName::FoldHookFn &&foldHookFn,295 GetCanonicalizationPatternsFn &&getCanonicalizationPatternsFn,296 OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrsFn)297 : Impl(StringAttr::get(dialect->getContext(),298 (dialect->getNamespace() + "." + name).str()),299 dialect, dialect->allocateTypeID(),300 /*interfaceMap=*/detail::InterfaceMap()),301 verifyFn(std::move(verifyFn)), verifyRegionFn(std::move(verifyRegionFn)),302 parseFn(std::move(parseFn)), printFn(std::move(printFn)),303 foldHookFn(std::move(foldHookFn)),304 getCanonicalizationPatternsFn(std::move(getCanonicalizationPatternsFn)),305 populateDefaultAttrsFn(std::move(populateDefaultAttrsFn)) {306 typeID = dialect->allocateTypeID();307}308 309std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(310 StringRef name, ExtensibleDialect *dialect,311 OperationName::VerifyInvariantsFn &&verifyFn,312 OperationName::VerifyRegionInvariantsFn &&verifyRegionFn) {313 auto parseFn = [](OpAsmParser &parser, OperationState &result) {314 return parser.emitError(315 parser.getCurrentLocation(),316 "dynamic operation do not define any parser function");317 };318 319 auto printFn = [](Operation *op, OpAsmPrinter &printer, StringRef) {320 printer.printGenericOp(op);321 };322 323 return DynamicOpDefinition::get(name, dialect, std::move(verifyFn),324 std::move(verifyRegionFn), std::move(parseFn),325 std::move(printFn));326}327 328std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(329 StringRef name, ExtensibleDialect *dialect,330 OperationName::VerifyInvariantsFn &&verifyFn,331 OperationName::VerifyRegionInvariantsFn &&verifyRegionFn,332 OperationName::ParseAssemblyFn &&parseFn,333 OperationName::PrintAssemblyFn &&printFn) {334 auto foldHookFn = [](Operation *op, ArrayRef<Attribute> operands,335 SmallVectorImpl<OpFoldResult> &results) {336 return failure();337 };338 339 auto getCanonicalizationPatternsFn = [](RewritePatternSet &, MLIRContext *) {340 };341 342 auto populateDefaultAttrsFn = [](const OperationName &, NamedAttrList &) {};343 344 return DynamicOpDefinition::get(name, dialect, std::move(verifyFn),345 std::move(verifyRegionFn), std::move(parseFn),346 std::move(printFn), std::move(foldHookFn),347 std::move(getCanonicalizationPatternsFn),348 std::move(populateDefaultAttrsFn));349}350 351std::unique_ptr<DynamicOpDefinition> DynamicOpDefinition::get(352 StringRef name, ExtensibleDialect *dialect,353 OperationName::VerifyInvariantsFn &&verifyFn,354 OperationName::VerifyInvariantsFn &&verifyRegionFn,355 OperationName::ParseAssemblyFn &&parseFn,356 OperationName::PrintAssemblyFn &&printFn,357 OperationName::FoldHookFn &&foldHookFn,358 GetCanonicalizationPatternsFn &&getCanonicalizationPatternsFn,359 OperationName::PopulateDefaultAttrsFn &&populateDefaultAttrsFn) {360 return std::unique_ptr<DynamicOpDefinition>(new DynamicOpDefinition(361 name, dialect, std::move(verifyFn), std::move(verifyRegionFn),362 std::move(parseFn), std::move(printFn), std::move(foldHookFn),363 std::move(getCanonicalizationPatternsFn),364 std::move(populateDefaultAttrsFn)));365}366 367//===----------------------------------------------------------------------===//368// Extensible dialect369//===----------------------------------------------------------------------===//370 371namespace {372/// Interface that can only be implemented by extensible dialects.373/// The interface is used to check if a dialect is extensible or not.374class IsExtensibleDialect : public DialectInterface::Base<IsExtensibleDialect> {375public:376 IsExtensibleDialect(Dialect *dialect) : Base(dialect) {}377 378 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IsExtensibleDialect)379};380} // namespace381 382ExtensibleDialect::ExtensibleDialect(StringRef name, MLIRContext *ctx,383 TypeID typeID)384 : Dialect(name, ctx, typeID) {385 addInterfaces<IsExtensibleDialect>();386}387 388void ExtensibleDialect::registerDynamicType(389 std::unique_ptr<DynamicTypeDefinition> &&type) {390 DynamicTypeDefinition *typePtr = type.get();391 TypeID typeID = type->getTypeID();392 StringRef name = type->getName();393 ExtensibleDialect *dialect = type->getDialect();394 395 assert(dialect == this &&396 "trying to register a dynamic type in the wrong dialect");397 398 // If a type with the same name is already defined, fail.399 auto registered = dynTypes.try_emplace(typeID, std::move(type)).second;400 (void)registered;401 assert(registered && "type TypeID was not unique");402 403 registered = nameToDynTypes.insert({name, typePtr}).second;404 (void)registered;405 assert(registered &&406 "Trying to create a new dynamic type with an existing name");407 408 // The StringAttr allocates the type name StringRef for the duration of the409 // MLIR context.410 MLIRContext *ctx = getContext();411 auto nameAttr =412 StringAttr::get(ctx, getNamespace() + "." + typePtr->getName());413 414 auto abstractType = AbstractType::get(415 *dialect, DynamicAttr::getInterfaceMap(), DynamicType::getHasTraitFn(),416 DynamicType::getWalkImmediateSubElementsFn(),417 DynamicType::getReplaceImmediateSubElementsFn(), typeID, nameAttr);418 419 /// Add the type to the dialect and the type uniquer.420 addType(typeID, std::move(abstractType));421 typePtr->registerInTypeUniquer();422}423 424void ExtensibleDialect::registerDynamicAttr(425 std::unique_ptr<DynamicAttrDefinition> &&attr) {426 auto *attrPtr = attr.get();427 auto typeID = attr->getTypeID();428 auto name = attr->getName();429 auto *dialect = attr->getDialect();430 431 assert(dialect == this &&432 "trying to register a dynamic attribute in the wrong dialect");433 434 // If an attribute with the same name is already defined, fail.435 auto registered = dynAttrs.try_emplace(typeID, std::move(attr)).second;436 (void)registered;437 assert(registered && "attribute TypeID was not unique");438 439 registered = nameToDynAttrs.insert({name, attrPtr}).second;440 (void)registered;441 assert(registered &&442 "Trying to create a new dynamic attribute with an existing name");443 444 // The StringAttr allocates the attribute name StringRef for the duration of445 // the MLIR context.446 MLIRContext *ctx = getContext();447 auto nameAttr =448 StringAttr::get(ctx, getNamespace() + "." + attrPtr->getName());449 450 auto abstractAttr = AbstractAttribute::get(451 *dialect, DynamicAttr::getInterfaceMap(), DynamicAttr::getHasTraitFn(),452 DynamicAttr::getWalkImmediateSubElementsFn(),453 DynamicAttr::getReplaceImmediateSubElementsFn(), typeID, nameAttr);454 455 /// Add the type to the dialect and the type uniquer.456 addAttribute(typeID, std::move(abstractAttr));457 attrPtr->registerInAttrUniquer();458}459 460void ExtensibleDialect::registerDynamicOp(461 std::unique_ptr<DynamicOpDefinition> &&op) {462 assert(op->dialect == this &&463 "trying to register a dynamic op in the wrong dialect");464 RegisteredOperationName::insert(std::move(op), /*attrNames=*/{});465}466 467bool ExtensibleDialect::classof(const Dialect *dialect) {468 return const_cast<Dialect *>(dialect)469 ->getRegisteredInterface<IsExtensibleDialect>();470}471 472OptionalParseResult ExtensibleDialect::parseOptionalDynamicType(473 StringRef typeName, AsmParser &parser, Type &resultType) const {474 DynamicTypeDefinition *typeDef = lookupTypeDefinition(typeName);475 if (!typeDef)476 return std::nullopt;477 478 DynamicType dynType;479 if (DynamicType::parse(parser, typeDef, dynType))480 return failure();481 resultType = dynType;482 return success();483}484 485LogicalResult ExtensibleDialect::printIfDynamicType(Type type,486 AsmPrinter &printer) {487 if (auto dynType = llvm::dyn_cast<DynamicType>(type)) {488 dynType.print(printer);489 return success();490 }491 return failure();492}493 494OptionalParseResult ExtensibleDialect::parseOptionalDynamicAttr(495 StringRef attrName, AsmParser &parser, Attribute &resultAttr) const {496 DynamicAttrDefinition *attrDef = lookupAttrDefinition(attrName);497 if (!attrDef)498 return std::nullopt;499 500 DynamicAttr dynAttr;501 if (DynamicAttr::parse(parser, attrDef, dynAttr))502 return failure();503 resultAttr = dynAttr;504 return success();505}506 507LogicalResult ExtensibleDialect::printIfDynamicAttr(Attribute attribute,508 AsmPrinter &printer) {509 if (auto dynAttr = llvm::dyn_cast<DynamicAttr>(attribute)) {510 dynAttr.print(printer);511 return success();512 }513 return failure();514}515 516//===----------------------------------------------------------------------===//517// Dynamic dialect518//===----------------------------------------------------------------------===//519 520namespace {521/// Interface that can only be implemented by extensible dialects.522/// The interface is used to check if a dialect is extensible or not.523class IsDynamicDialect : public DialectInterface::Base<IsDynamicDialect> {524public:525 IsDynamicDialect(Dialect *dialect) : Base(dialect) {}526 527 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IsDynamicDialect)528};529} // namespace530 531DynamicDialect::DynamicDialect(StringRef name, MLIRContext *ctx)532 : SelfOwningTypeID(),533 ExtensibleDialect(name, ctx, SelfOwningTypeID::getTypeID()) {534 addInterfaces<IsDynamicDialect>();535}536 537bool DynamicDialect::classof(const Dialect *dialect) {538 return const_cast<Dialect *>(dialect)539 ->getRegisteredInterface<IsDynamicDialect>();540}541 542Type DynamicDialect::parseType(DialectAsmParser &parser) const {543 auto loc = parser.getCurrentLocation();544 StringRef typeTag;545 if (failed(parser.parseKeyword(&typeTag)))546 return Type();547 548 {549 Type dynType;550 auto parseResult = parseOptionalDynamicType(typeTag, parser, dynType);551 if (parseResult.has_value()) {552 if (succeeded(parseResult.value()))553 return dynType;554 return Type();555 }556 }557 558 parser.emitError(loc, "expected dynamic type");559 return Type();560}561 562void DynamicDialect::printType(Type type, DialectAsmPrinter &printer) const {563 auto wasDynamic = printIfDynamicType(type, printer);564 (void)wasDynamic;565 assert(succeeded(wasDynamic) &&566 "non-dynamic type defined in dynamic dialect");567}568 569Attribute DynamicDialect::parseAttribute(DialectAsmParser &parser,570 Type type) const {571 auto loc = parser.getCurrentLocation();572 StringRef typeTag;573 if (failed(parser.parseKeyword(&typeTag)))574 return Attribute();575 576 {577 Attribute dynAttr;578 auto parseResult = parseOptionalDynamicAttr(typeTag, parser, dynAttr);579 if (parseResult.has_value()) {580 if (succeeded(parseResult.value()))581 return dynAttr;582 return Attribute();583 }584 }585 586 parser.emitError(loc, "expected dynamic attribute");587 return Attribute();588}589void DynamicDialect::printAttribute(Attribute attr,590 DialectAsmPrinter &printer) const {591 auto wasDynamic = printIfDynamicAttr(attr, printer);592 (void)wasDynamic;593 assert(succeeded(wasDynamic) &&594 "non-dynamic attribute defined in dynamic dialect");595}596