579 lines · cpp
1//===- OpDefinitionsGen.cpp - IRDL op definitions generator ---------------===//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// OpDefinitionsGen uses the description of operations to generate IRDL10// definitions for ops.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/IRDL/IR/IRDL.h"15#include "mlir/IR/Attributes.h"16#include "mlir/IR/Builders.h"17#include "mlir/IR/BuiltinOps.h"18#include "mlir/IR/Diagnostics.h"19#include "mlir/IR/Dialect.h"20#include "mlir/IR/MLIRContext.h"21#include "mlir/TableGen/AttrOrTypeDef.h"22#include "mlir/TableGen/GenInfo.h"23#include "mlir/TableGen/GenNameParser.h"24#include "mlir/TableGen/Interfaces.h"25#include "mlir/TableGen/Operator.h"26#include "llvm/ADT/StringExtras.h"27#include "llvm/Support/CommandLine.h"28#include "llvm/Support/InitLLVM.h"29#include "llvm/Support/raw_ostream.h"30#include "llvm/TableGen/Main.h"31#include "llvm/TableGen/Record.h"32#include "llvm/TableGen/TableGenBackend.h"33 34using namespace llvm;35using namespace mlir;36using tblgen::NamedTypeConstraint;37 38static llvm::cl::OptionCategory dialectGenCat("Options for -gen-irdl-dialect");39static llvm::cl::opt<std::string>40 selectedDialect("dialect", llvm::cl::desc("The dialect to gen for"),41 llvm::cl::cat(dialectGenCat), llvm::cl::Required);42 43static Value createPredicate(OpBuilder &builder, tblgen::Pred pred) {44 MLIRContext *ctx = builder.getContext();45 46 if (pred.isCombined()) {47 auto combiner = pred.getDef().getValueAsDef("kind")->getName();48 if (combiner == "PredCombinerAnd" || combiner == "PredCombinerOr") {49 std::vector<Value> constraints;50 for (auto *child : pred.getDef().getValueAsListOfDefs("children")) {51 constraints.push_back(createPredicate(builder, tblgen::Pred(child)));52 }53 if (combiner == "PredCombinerAnd") {54 auto op =55 irdl::AllOfOp::create(builder, UnknownLoc::get(ctx), constraints);56 return op.getOutput();57 }58 auto op =59 irdl::AnyOfOp::create(builder, UnknownLoc::get(ctx), constraints);60 return op.getOutput();61 }62 }63 64 std::string condition = pred.getCondition();65 // Build a CPredOp to match the C constraint built.66 irdl::CPredOp op = irdl::CPredOp::create(builder, UnknownLoc::get(ctx),67 StringAttr::get(ctx, condition));68 return op;69}70 71static Value typeToConstraint(OpBuilder &builder, Type type) {72 MLIRContext *ctx = builder.getContext();73 auto op =74 irdl::IsOp::create(builder, UnknownLoc::get(ctx), TypeAttr::get(type));75 return op.getOutput();76}77 78static Value baseToConstraint(OpBuilder &builder, StringRef baseClass) {79 MLIRContext *ctx = builder.getContext();80 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx),81 StringAttr::get(ctx, baseClass));82 return op.getOutput();83}84 85static std::optional<Type> recordToType(MLIRContext *ctx,86 const Record &predRec) {87 if (predRec.isSubClassOf("I")) {88 auto width = predRec.getValueAsInt("bitwidth");89 return IntegerType::get(ctx, width, IntegerType::Signless);90 }91 92 if (predRec.isSubClassOf("SI")) {93 auto width = predRec.getValueAsInt("bitwidth");94 return IntegerType::get(ctx, width, IntegerType::Signed);95 }96 97 if (predRec.isSubClassOf("UI")) {98 auto width = predRec.getValueAsInt("bitwidth");99 return IntegerType::get(ctx, width, IntegerType::Unsigned);100 }101 102 // Index type103 if (predRec.getName() == "Index") {104 return IndexType::get(ctx);105 }106 107 // Float types108 if (predRec.isSubClassOf("F")) {109 auto width = predRec.getValueAsInt("bitwidth");110 switch (width) {111 case 16:112 return Float16Type::get(ctx);113 case 32:114 return Float32Type::get(ctx);115 case 64:116 return Float64Type::get(ctx);117 case 80:118 return Float80Type::get(ctx);119 case 128:120 return Float128Type::get(ctx);121 }122 }123 124 if (predRec.getName() == "NoneType") {125 return NoneType::get(ctx);126 }127 128 if (predRec.getName() == "BF16") {129 return BFloat16Type::get(ctx);130 }131 132 if (predRec.getName() == "TF32") {133 return FloatTF32Type::get(ctx);134 }135 136 if (predRec.getName() == "F8E4M3FN") {137 return Float8E4M3FNType::get(ctx);138 }139 140 if (predRec.getName() == "F8E5M2") {141 return Float8E5M2Type::get(ctx);142 }143 144 if (predRec.getName() == "F8E4M3") {145 return Float8E4M3Type::get(ctx);146 }147 148 if (predRec.getName() == "F8E4M3FNUZ") {149 return Float8E4M3FNUZType::get(ctx);150 }151 152 if (predRec.getName() == "F8E4M3B11FNUZ") {153 return Float8E4M3B11FNUZType::get(ctx);154 }155 156 if (predRec.getName() == "F8E5M2FNUZ") {157 return Float8E5M2FNUZType::get(ctx);158 }159 160 if (predRec.getName() == "F8E3M4") {161 return Float8E3M4Type::get(ctx);162 }163 164 if (predRec.isSubClassOf("Complex")) {165 const Record *elementRec = predRec.getValueAsDef("elementType");166 auto elementType = recordToType(ctx, *elementRec);167 if (elementType.has_value()) {168 return ComplexType::get(elementType.value());169 }170 }171 172 return std::nullopt;173}174 175static Value createTypeConstraint(OpBuilder &builder,176 tblgen::Constraint constraint) {177 MLIRContext *ctx = builder.getContext();178 const Record &predRec = constraint.getDef();179 180 if (predRec.isSubClassOf("Variadic") || predRec.isSubClassOf("Optional"))181 return createTypeConstraint(builder, predRec.getValueAsDef("baseType"));182 183 if (predRec.getName() == "AnyType") {184 auto op = irdl::AnyOp::create(builder, UnknownLoc::get(ctx));185 return op.getOutput();186 }187 188 if (predRec.isSubClassOf("TypeDef")) {189 auto dialect = predRec.getValueAsDef("dialect")->getValueAsString("name");190 if (dialect == selectedDialect) {191 std::string combined = ("!" + predRec.getValueAsString("mnemonic")).str();192 SmallVector<FlatSymbolRefAttr> nested = {193 SymbolRefAttr::get(ctx, combined)};194 auto typeSymbol = SymbolRefAttr::get(ctx, dialect, nested);195 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx), typeSymbol);196 return op.getOutput();197 }198 std::string typeName = ("!" + predRec.getValueAsString("typeName")).str();199 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx),200 StringAttr::get(ctx, typeName));201 return op.getOutput();202 }203 204 if (predRec.isSubClassOf("AnyTypeOf")) {205 std::vector<Value> constraints;206 for (const Record *child : predRec.getValueAsListOfDefs("allowedTypes")) {207 constraints.push_back(208 createTypeConstraint(builder, tblgen::Constraint(child)));209 }210 auto op = irdl::AnyOfOp::create(builder, UnknownLoc::get(ctx), constraints);211 return op.getOutput();212 }213 214 if (predRec.isSubClassOf("AllOfType")) {215 std::vector<Value> constraints;216 for (const Record *child : predRec.getValueAsListOfDefs("allowedTypes")) {217 constraints.push_back(218 createTypeConstraint(builder, tblgen::Constraint(child)));219 }220 auto op = irdl::AllOfOp::create(builder, UnknownLoc::get(ctx), constraints);221 return op.getOutput();222 }223 224 // Integer types225 if (predRec.getName() == "AnyInteger") {226 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx),227 StringAttr::get(ctx, "!builtin.integer"));228 return op.getOutput();229 }230 231 if (predRec.isSubClassOf("AnyI")) {232 auto width = predRec.getValueAsInt("bitwidth");233 std::vector<Value> types = {234 typeToConstraint(builder,235 IntegerType::get(ctx, width, IntegerType::Signless)),236 typeToConstraint(builder,237 IntegerType::get(ctx, width, IntegerType::Signed)),238 typeToConstraint(builder,239 IntegerType::get(ctx, width, IntegerType::Unsigned))};240 auto op = irdl::AnyOfOp::create(builder, UnknownLoc::get(ctx), types);241 return op.getOutput();242 }243 244 auto type = recordToType(ctx, predRec);245 246 if (type.has_value()) {247 return typeToConstraint(builder, type.value());248 }249 250 // Confined type251 if (predRec.isSubClassOf("ConfinedType")) {252 std::vector<Value> constraints;253 constraints.push_back(createTypeConstraint(254 builder, tblgen::Constraint(predRec.getValueAsDef("baseType"))));255 for (const Record *child : predRec.getValueAsListOfDefs("predicateList")) {256 constraints.push_back(createPredicate(builder, tblgen::Pred(child)));257 }258 auto op = irdl::AllOfOp::create(builder, UnknownLoc::get(ctx), constraints);259 return op.getOutput();260 }261 262 return createPredicate(builder, constraint.getPredicate());263}264 265static Value createAttrConstraint(OpBuilder &builder,266 tblgen::Constraint constraint) {267 MLIRContext *ctx = builder.getContext();268 const Record &predRec = constraint.getDef();269 270 if (predRec.isSubClassOf("DefaultValuedAttr") ||271 predRec.isSubClassOf("DefaultValuedOptionalAttr") ||272 predRec.isSubClassOf("OptionalAttr")) {273 return createAttrConstraint(builder, predRec.getValueAsDef("baseAttr"));274 }275 276 if (predRec.isSubClassOf("ConfinedAttr")) {277 std::vector<Value> constraints;278 constraints.push_back(createAttrConstraint(279 builder, tblgen::Constraint(predRec.getValueAsDef("baseAttr"))));280 for (const Record *child :281 predRec.getValueAsListOfDefs("attrConstraints")) {282 constraints.push_back(createPredicate(283 builder, tblgen::Pred(child->getValueAsDef("predicate"))));284 }285 auto op = irdl::AllOfOp::create(builder, UnknownLoc::get(ctx), constraints);286 return op.getOutput();287 }288 289 if (predRec.isSubClassOf("AnyAttrOf")) {290 std::vector<Value> constraints;291 for (const Record *child :292 predRec.getValueAsListOfDefs("allowedAttributes")) {293 constraints.push_back(294 createAttrConstraint(builder, tblgen::Constraint(child)));295 }296 auto op = irdl::AnyOfOp::create(builder, UnknownLoc::get(ctx), constraints);297 return op.getOutput();298 }299 300 if (predRec.getName() == "AnyAttr") {301 auto op = irdl::AnyOp::create(builder, UnknownLoc::get(ctx));302 return op.getOutput();303 }304 305 if (predRec.isSubClassOf("AnyIntegerAttrBase") ||306 predRec.isSubClassOf("SignlessIntegerAttrBase") ||307 predRec.isSubClassOf("SignedIntegerAttrBase") ||308 predRec.isSubClassOf("UnsignedIntegerAttrBase") ||309 predRec.isSubClassOf("BoolAttr")) {310 return baseToConstraint(builder, "!builtin.integer");311 }312 313 if (predRec.isSubClassOf("FloatAttrBase")) {314 return baseToConstraint(builder, "!builtin.float");315 }316 317 if (predRec.isSubClassOf("StringBasedAttr")) {318 return baseToConstraint(builder, "!builtin.string");319 }320 321 if (predRec.getName() == "UnitAttr") {322 auto op =323 irdl::IsOp::create(builder, UnknownLoc::get(ctx), UnitAttr::get(ctx));324 return op.getOutput();325 }326 327 if (predRec.isSubClassOf("AttrDef")) {328 auto dialect = predRec.getValueAsDef("dialect")->getValueAsString("name");329 if (dialect == selectedDialect) {330 std::string combined = ("#" + predRec.getValueAsString("mnemonic")).str();331 SmallVector<FlatSymbolRefAttr> nested = {SymbolRefAttr::get(ctx, combined)332 333 };334 auto typeSymbol = SymbolRefAttr::get(ctx, dialect, nested);335 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx), typeSymbol);336 return op.getOutput();337 }338 std::string typeName = ("#" + predRec.getValueAsString("attrName")).str();339 auto op = irdl::BaseOp::create(builder, UnknownLoc::get(ctx),340 StringAttr::get(ctx, typeName));341 return op.getOutput();342 }343 344 return createPredicate(builder, constraint.getPredicate());345}346 347static Value createRegionConstraint(OpBuilder &builder,348 tblgen::Region constraint) {349 MLIRContext *ctx = builder.getContext();350 const Record &predRec = constraint.getDef();351 352 if (predRec.getName() == "AnyRegion") {353 ValueRange entryBlockArgs = {};354 auto op =355 irdl::RegionOp::create(builder, UnknownLoc::get(ctx), entryBlockArgs);356 return op.getResult();357 }358 359 if (predRec.isSubClassOf("SizedRegion")) {360 ValueRange entryBlockArgs = {};361 auto ty = IntegerType::get(ctx, 32);362 auto op = irdl::RegionOp::create(363 builder, UnknownLoc::get(ctx), entryBlockArgs,364 IntegerAttr::get(ty, predRec.getValueAsInt("blocks")));365 return op.getResult();366 }367 368 return createPredicate(builder, constraint.getPredicate());369}370 371/// Returns the name of the operation without the dialect prefix.372static StringRef getOperatorName(tblgen::Operator &tblgenOp) {373 StringRef opName = tblgenOp.getDef().getValueAsString("opName");374 return opName;375}376 377/// Returns the name of the type without the dialect prefix.378static StringRef getTypeName(tblgen::TypeDef &tblgenType) {379 StringRef opName = tblgenType.getDef()->getValueAsString("mnemonic");380 return opName;381}382 383/// Returns the name of the attr without the dialect prefix.384static StringRef getAttrName(tblgen::AttrDef &tblgenType) {385 StringRef opName = tblgenType.getDef()->getValueAsString("mnemonic");386 return opName;387}388 389/// Extract an operation to IRDL.390static irdl::OperationOp createIRDLOperation(OpBuilder &builder,391 tblgen::Operator &tblgenOp) {392 MLIRContext *ctx = builder.getContext();393 StringRef opName = getOperatorName(tblgenOp);394 395 irdl::OperationOp op = irdl::OperationOp::create(396 builder, UnknownLoc::get(ctx), StringAttr::get(ctx, opName));397 398 // Add the block in the region.399 Block &opBlock = op.getBody().emplaceBlock();400 OpBuilder consBuilder = OpBuilder::atBlockBegin(&opBlock);401 402 SmallDenseSet<StringRef> usedNames;403 for (auto &namedCons : tblgenOp.getOperands())404 usedNames.insert(namedCons.name);405 for (auto &namedCons : tblgenOp.getResults())406 usedNames.insert(namedCons.name);407 for (auto &namedReg : tblgenOp.getRegions())408 usedNames.insert(namedReg.name);409 410 size_t generateCounter = 0;411 auto generateName = [&](StringRef prefix) -> StringAttr {412 SmallString<16> candidate;413 do {414 candidate.clear();415 raw_svector_ostream candidateStream(candidate);416 candidateStream << prefix << generateCounter;417 generateCounter++;418 } while (usedNames.contains(candidate));419 return StringAttr::get(ctx, candidate);420 };421 auto normalizeName = [&](StringRef name) -> StringAttr {422 if (name == "")423 return generateName("unnamed");424 return StringAttr::get(ctx, name);425 };426 427 auto getValues = [&](tblgen::Operator::const_value_range namedCons) {428 SmallVector<Value> operands;429 SmallVector<Attribute> names;430 SmallVector<irdl::VariadicityAttr> variadicity;431 432 for (const NamedTypeConstraint &namedCons : namedCons) {433 auto operand = createTypeConstraint(consBuilder, namedCons.constraint);434 operands.push_back(operand);435 436 names.push_back(normalizeName(namedCons.name));437 438 irdl::VariadicityAttr var;439 if (namedCons.isOptional())440 var = consBuilder.getAttr<irdl::VariadicityAttr>(441 irdl::Variadicity::optional);442 else if (namedCons.isVariadic())443 var = consBuilder.getAttr<irdl::VariadicityAttr>(444 irdl::Variadicity::variadic);445 else446 var = consBuilder.getAttr<irdl::VariadicityAttr>(447 irdl::Variadicity::single);448 449 variadicity.push_back(var);450 }451 return std::make_tuple(operands, names, variadicity);452 };453 454 auto [operands, operandNames, operandVariadicity] =455 getValues(tblgenOp.getOperands());456 auto [results, resultNames, resultVariadicity] =457 getValues(tblgenOp.getResults());458 459 SmallVector<Value> attributes;460 SmallVector<Attribute> attrNames;461 for (auto namedAttr : tblgenOp.getAttributes()) {462 if (namedAttr.attr.isOptional())463 continue;464 attributes.push_back(createAttrConstraint(consBuilder, namedAttr.attr));465 attrNames.push_back(StringAttr::get(ctx, namedAttr.name));466 }467 468 SmallVector<Value> regions;469 SmallVector<Attribute> regionNames;470 for (auto namedRegion : tblgenOp.getRegions()) {471 regions.push_back(472 createRegionConstraint(consBuilder, namedRegion.constraint));473 regionNames.push_back(normalizeName(namedRegion.name));474 }475 476 // Create the operands and results operations.477 if (!operands.empty())478 irdl::OperandsOp::create(consBuilder, UnknownLoc::get(ctx), operands,479 ArrayAttr::get(ctx, operandNames),480 operandVariadicity);481 if (!results.empty())482 irdl::ResultsOp::create(consBuilder, UnknownLoc::get(ctx), results,483 ArrayAttr::get(ctx, resultNames),484 resultVariadicity);485 if (!attributes.empty())486 irdl::AttributesOp::create(consBuilder, UnknownLoc::get(ctx), attributes,487 ArrayAttr::get(ctx, attrNames));488 if (!regions.empty())489 irdl::RegionsOp::create(consBuilder, UnknownLoc::get(ctx), regions,490 ArrayAttr::get(ctx, regionNames));491 492 return op;493}494 495static irdl::TypeOp createIRDLType(OpBuilder &builder,496 tblgen::TypeDef &tblgenType) {497 MLIRContext *ctx = builder.getContext();498 StringRef typeName = getTypeName(tblgenType);499 std::string combined = ("!" + typeName).str();500 501 irdl::TypeOp op = irdl::TypeOp::create(builder, UnknownLoc::get(ctx),502 StringAttr::get(ctx, combined));503 504 op.getBody().emplaceBlock();505 506 return op;507}508 509static irdl::AttributeOp createIRDLAttr(OpBuilder &builder,510 tblgen::AttrDef &tblgenAttr) {511 MLIRContext *ctx = builder.getContext();512 StringRef attrName = getAttrName(tblgenAttr);513 std::string combined = ("#" + attrName).str();514 515 irdl::AttributeOp op = irdl::AttributeOp::create(516 builder, UnknownLoc::get(ctx), StringAttr::get(ctx, combined));517 518 op.getBody().emplaceBlock();519 520 return op;521}522 523static irdl::DialectOp createIRDLDialect(OpBuilder &builder) {524 MLIRContext *ctx = builder.getContext();525 return irdl::DialectOp::create(builder, UnknownLoc::get(ctx),526 StringAttr::get(ctx, selectedDialect));527}528 529static bool emitDialectIRDLDefs(const RecordKeeper &records, raw_ostream &os) {530 // Initialize.531 MLIRContext ctx;532 ctx.getOrLoadDialect<irdl::IRDLDialect>();533 OpBuilder builder(&ctx);534 535 // Create a module op and set it as the insertion point.536 OwningOpRef<ModuleOp> module =537 ModuleOp::create(builder, UnknownLoc::get(&ctx));538 builder = builder.atBlockBegin(module->getBody());539 // Create the dialect and insert it.540 irdl::DialectOp dialect = createIRDLDialect(builder);541 // Set insertion point to start of DialectOp.542 builder = builder.atBlockBegin(&dialect.getBody().emplaceBlock());543 544 for (const Record *type :545 records.getAllDerivedDefinitionsIfDefined("TypeDef")) {546 tblgen::TypeDef tblgenType(type);547 if (tblgenType.getDialect().getName() != selectedDialect)548 continue;549 createIRDLType(builder, tblgenType);550 }551 552 for (const Record *attr :553 records.getAllDerivedDefinitionsIfDefined("AttrDef")) {554 tblgen::AttrDef tblgenAttr(attr);555 if (tblgenAttr.getDialect().getName() != selectedDialect)556 continue;557 createIRDLAttr(builder, tblgenAttr);558 }559 560 for (const Record *def : records.getAllDerivedDefinitionsIfDefined("Op")) {561 tblgen::Operator tblgenOp(def);562 if (tblgenOp.getDialectName() != selectedDialect)563 continue;564 565 createIRDLOperation(builder, tblgenOp);566 }567 568 // Print the module.569 module->print(os);570 571 return false;572}573 574static mlir::GenRegistration575 genOpDefs("gen-dialect-irdl-defs", "Generate IRDL dialect definitions",576 [](const RecordKeeper &records, raw_ostream &os) {577 return emitDialectIRDLDefs(records, os);578 });579