179 lines · cpp
1//===- Context.cpp --------------------------------------------------------===//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/Tools/PDLL/ODS/Context.h"10#include "mlir/Tools/PDLL/ODS/Constraint.h"11#include "mlir/Tools/PDLL/ODS/Dialect.h"12#include "mlir/Tools/PDLL/ODS/Operation.h"13#include "llvm/Support/ScopedPrinter.h"14#include "llvm/Support/raw_ostream.h"15 16using namespace mlir;17using namespace mlir::pdll::ods;18 19//===----------------------------------------------------------------------===//20// Context21//===----------------------------------------------------------------------===//22 23Context::Context() = default;24Context::~Context() = default;25 26const AttributeConstraint &27Context::insertAttributeConstraint(StringRef name, StringRef summary,28 StringRef cppClass) {29 std::unique_ptr<AttributeConstraint> &constraint = attributeConstraints[name];30 if (!constraint) {31 constraint.reset(new AttributeConstraint(name, summary, cppClass));32 } else {33 assert(constraint->getCppClass() == cppClass &&34 constraint->getSummary() == summary &&35 "constraint with the same name was already registered with a "36 "different class");37 }38 return *constraint;39}40 41const TypeConstraint &Context::insertTypeConstraint(StringRef name,42 StringRef summary,43 StringRef cppClass) {44 std::unique_ptr<TypeConstraint> &constraint = typeConstraints[name];45 if (!constraint)46 constraint.reset(new TypeConstraint(name, summary, cppClass));47 return *constraint;48}49 50Dialect &Context::insertDialect(StringRef name) {51 std::unique_ptr<Dialect> &dialect = dialects[name];52 if (!dialect)53 dialect.reset(new Dialect(name));54 return *dialect;55}56 57const Dialect *Context::lookupDialect(StringRef name) const {58 auto it = dialects.find(name);59 return it == dialects.end() ? nullptr : &*it->second;60}61 62std::pair<Operation *, bool>63Context::insertOperation(StringRef name, StringRef summary, StringRef desc,64 StringRef nativeClassName,65 bool supportsResultTypeInferrence, SMLoc loc) {66 std::pair<StringRef, StringRef> dialectAndName = name.split('.');67 return insertDialect(dialectAndName.first)68 .insertOperation(name, summary, desc, nativeClassName,69 supportsResultTypeInferrence, loc);70}71 72const Operation *Context::lookupOperation(StringRef name) const {73 std::pair<StringRef, StringRef> dialectAndName = name.split('.');74 if (const Dialect *dialect = lookupDialect(dialectAndName.first))75 return dialect->lookupOperation(name);76 return nullptr;77}78 79template <typename T>80static SmallVector<T *>81sortMapByName(const llvm::StringMap<std::unique_ptr<T>> &map) {82 SmallVector<T *> storage;83 for (auto &entry : map)84 storage.push_back(entry.second.get());85 llvm::sort(storage, [](const auto &lhs, const auto &rhs) {86 return lhs->getName() < rhs->getName();87 });88 return storage;89}90 91void Context::print(raw_ostream &os) const {92 auto printVariableLengthCst = [&](StringRef cst, VariableLengthKind kind) {93 switch (kind) {94 case VariableLengthKind::Optional:95 os << "Optional<" << cst << ">";96 break;97 case VariableLengthKind::Single:98 os << cst;99 break;100 case VariableLengthKind::Variadic:101 os << "Variadic<" << cst << ">";102 break;103 }104 };105 106 llvm::ScopedPrinter printer(os);107 llvm::DictScope odsScope(printer, "ODSContext");108 for (const Dialect *dialect : sortMapByName(dialects)) {109 printer.startLine() << "Dialect `" << dialect->getName() << "` {\n";110 printer.indent();111 112 for (const Operation *op : sortMapByName(dialect->getOperations())) {113 printer.startLine() << "Operation `" << op->getName() << "` {\n";114 printer.indent();115 116 // Attributes.117 ArrayRef<Attribute> attributes = op->getAttributes();118 if (!attributes.empty()) {119 printer.startLine() << "Attributes { ";120 llvm::interleaveComma(attributes, os, [&](const Attribute &attr) {121 os << attr.getName() << " : ";122 123 auto kind = attr.isOptional() ? VariableLengthKind::Optional124 : VariableLengthKind::Single;125 printVariableLengthCst(attr.getConstraint().getDemangledName(), kind);126 });127 os << " }\n";128 }129 130 // Operands.131 ArrayRef<OperandOrResult> operands = op->getOperands();132 if (!operands.empty()) {133 printer.startLine() << "Operands { ";134 llvm::interleaveComma(135 operands, os, [&](const OperandOrResult &operand) {136 os << operand.getName() << " : ";137 printVariableLengthCst(operand.getConstraint().getDemangledName(),138 operand.getVariableLengthKind());139 });140 os << " }\n";141 }142 143 // Results.144 ArrayRef<OperandOrResult> results = op->getResults();145 if (!results.empty()) {146 printer.startLine() << "Results { ";147 llvm::interleaveComma(results, os, [&](const OperandOrResult &result) {148 os << result.getName() << " : ";149 printVariableLengthCst(result.getConstraint().getDemangledName(),150 result.getVariableLengthKind());151 });152 os << " }\n";153 }154 155 printer.objectEnd();156 }157 printer.objectEnd();158 }159 for (const AttributeConstraint *cst : sortMapByName(attributeConstraints)) {160 printer.startLine() << "AttributeConstraint `" << cst->getDemangledName()161 << "` {\n";162 printer.indent();163 164 printer.startLine() << "Summary: " << cst->getSummary() << "\n";165 printer.startLine() << "CppClass: " << cst->getCppClass() << "\n";166 printer.objectEnd();167 }168 for (const TypeConstraint *cst : sortMapByName(typeConstraints)) {169 printer.startLine() << "TypeConstraint `" << cst->getDemangledName()170 << "` {\n";171 printer.indent();172 173 printer.startLine() << "Summary: " << cst->getSummary() << "\n";174 printer.startLine() << "CppClass: " << cst->getCppClass() << "\n";175 printer.objectEnd();176 }177 printer.objectEnd();178}179