477 lines · cpp
1//===- TestDialect.cpp - MLIR Dialect for Testing -------------------------===//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 "TestDialect.h"10#include "TestOps.h"11#include "TestTypes.h"12#include "mlir/Bytecode/BytecodeImplementation.h"13#include "mlir/Dialect/Arith/IR/Arith.h"14#include "mlir/Dialect/Func/IR/FuncOps.h"15#include "mlir/Dialect/Tensor/IR/Tensor.h"16#include "mlir/IR/AsmState.h"17#include "mlir/IR/BuiltinAttributes.h"18#include "mlir/IR/BuiltinOps.h"19#include "mlir/IR/Diagnostics.h"20#include "mlir/IR/ExtensibleDialect.h"21#include "mlir/IR/MLIRContext.h"22#include "mlir/IR/ODSSupport.h"23#include "mlir/IR/OperationSupport.h"24#include "mlir/IR/PatternMatch.h"25#include "mlir/IR/TypeUtilities.h"26#include "mlir/IR/Verifier.h"27#include "mlir/Interfaces/CallInterfaces.h"28#include "mlir/Interfaces/FunctionImplementation.h"29#include "mlir/Interfaces/InferIntRangeInterface.h"30#include "mlir/Support/LLVM.h"31#include "mlir/Transforms/FoldUtils.h"32#include "mlir/Transforms/InliningUtils.h"33#include "llvm/ADT/STLFunctionalExtras.h"34#include "llvm/ADT/SmallString.h"35#include "llvm/ADT/StringExtras.h"36#include "llvm/ADT/StringSwitch.h"37#include "llvm/Support/Base64.h"38#include "llvm/Support/Casting.h"39 40#include "mlir/Dialect/DLTI/DLTI.h"41#include "mlir/Interfaces/FoldInterfaces.h"42#include "mlir/Reducer/ReductionPatternInterface.h"43#include "mlir/Transforms/InliningUtils.h"44#include <cstdint>45#include <numeric>46#include <optional>47 48// Include this before the using namespace lines below to test that we don't49// have namespace dependencies.50#include "TestOpsDialect.cpp.inc"51 52using namespace mlir;53using namespace test;54 55//===----------------------------------------------------------------------===//56// PropertiesWithCustomPrint57//===----------------------------------------------------------------------===//58 59LogicalResult60test::setPropertiesFromAttribute(PropertiesWithCustomPrint &prop,61 Attribute attr,62 function_ref<InFlightDiagnostic()> emitError) {63 DictionaryAttr dict = dyn_cast<DictionaryAttr>(attr);64 if (!dict) {65 emitError() << "expected DictionaryAttr to set TestProperties";66 return failure();67 }68 auto label = dict.getAs<mlir::StringAttr>("label");69 if (!label) {70 emitError() << "expected StringAttr for key `label`";71 return failure();72 }73 auto valueAttr = dict.getAs<IntegerAttr>("value");74 if (!valueAttr) {75 emitError() << "expected IntegerAttr for key `value`";76 return failure();77 }78 79 prop.label = std::make_shared<std::string>(label.getValue());80 prop.value = valueAttr.getValue().getSExtValue();81 return success();82}83 84DictionaryAttr85test::getPropertiesAsAttribute(MLIRContext *ctx,86 const PropertiesWithCustomPrint &prop) {87 SmallVector<NamedAttribute> attrs;88 Builder b{ctx};89 attrs.push_back(b.getNamedAttr("label", b.getStringAttr(*prop.label)));90 attrs.push_back(b.getNamedAttr("value", b.getI32IntegerAttr(prop.value)));91 return b.getDictionaryAttr(attrs);92}93 94llvm::hash_code test::computeHash(const PropertiesWithCustomPrint &prop) {95 return llvm::hash_combine(prop.value, StringRef(*prop.label));96}97 98void test::customPrintProperties(OpAsmPrinter &p,99 const PropertiesWithCustomPrint &prop) {100 p.printKeywordOrString(*prop.label);101 p << " is " << prop.value;102}103 104ParseResult test::customParseProperties(OpAsmParser &parser,105 PropertiesWithCustomPrint &prop) {106 std::string label;107 if (parser.parseKeywordOrString(&label) || parser.parseKeyword("is") ||108 parser.parseInteger(prop.value))109 return failure();110 prop.label = std::make_shared<std::string>(std::move(label));111 return success();112}113 114//===----------------------------------------------------------------------===//115// MyPropStruct116//===----------------------------------------------------------------------===//117 118Attribute MyPropStruct::asAttribute(MLIRContext *ctx) const {119 return StringAttr::get(ctx, content);120}121 122LogicalResult123MyPropStruct::setFromAttr(MyPropStruct &prop, Attribute attr,124 function_ref<InFlightDiagnostic()> emitError) {125 StringAttr strAttr = dyn_cast<StringAttr>(attr);126 if (!strAttr) {127 emitError() << "Expect StringAttr but got " << attr;128 return failure();129 }130 prop.content = strAttr.getValue();131 return success();132}133 134llvm::hash_code MyPropStruct::hash() const {135 return hash_value(StringRef(content));136}137 138LogicalResult test::readFromMlirBytecode(DialectBytecodeReader &reader,139 MyPropStruct &prop) {140 StringRef str;141 if (failed(reader.readString(str)))142 return failure();143 prop.content = str.str();144 return success();145}146 147void test::writeToMlirBytecode(DialectBytecodeWriter &writer,148 MyPropStruct &prop) {149 writer.writeOwnedString(prop.content);150}151 152//===----------------------------------------------------------------------===//153// VersionedProperties154//===----------------------------------------------------------------------===//155 156LogicalResult157test::setPropertiesFromAttribute(VersionedProperties &prop, Attribute attr,158 function_ref<InFlightDiagnostic()> emitError) {159 DictionaryAttr dict = dyn_cast<DictionaryAttr>(attr);160 if (!dict) {161 emitError() << "expected DictionaryAttr to set VersionedProperties";162 return failure();163 }164 auto value1Attr = dict.getAs<IntegerAttr>("value1");165 if (!value1Attr) {166 emitError() << "expected IntegerAttr for key `value1`";167 return failure();168 }169 auto value2Attr = dict.getAs<IntegerAttr>("value2");170 if (!value2Attr) {171 emitError() << "expected IntegerAttr for key `value2`";172 return failure();173 }174 175 prop.value1 = value1Attr.getValue().getSExtValue();176 prop.value2 = value2Attr.getValue().getSExtValue();177 return success();178}179 180DictionaryAttr test::getPropertiesAsAttribute(MLIRContext *ctx,181 const VersionedProperties &prop) {182 SmallVector<NamedAttribute> attrs;183 Builder b{ctx};184 attrs.push_back(b.getNamedAttr("value1", b.getI32IntegerAttr(prop.value1)));185 attrs.push_back(b.getNamedAttr("value2", b.getI32IntegerAttr(prop.value2)));186 return b.getDictionaryAttr(attrs);187}188 189llvm::hash_code test::computeHash(const VersionedProperties &prop) {190 return llvm::hash_combine(prop.value1, prop.value2);191}192 193void test::customPrintProperties(OpAsmPrinter &p,194 const VersionedProperties &prop) {195 p << prop.value1 << " | " << prop.value2;196}197 198ParseResult test::customParseProperties(OpAsmParser &parser,199 VersionedProperties &prop) {200 if (parser.parseInteger(prop.value1) || parser.parseVerticalBar() ||201 parser.parseInteger(prop.value2))202 return failure();203 return success();204}205 206//===----------------------------------------------------------------------===//207// Bytecode Support208//===----------------------------------------------------------------------===//209 210LogicalResult test::readFromMlirBytecode(DialectBytecodeReader &reader,211 MutableArrayRef<int64_t> prop) {212 uint64_t size;213 if (failed(reader.readVarInt(size)))214 return failure();215 if (size != prop.size())216 return reader.emitError("array size mismach when reading properties: ")217 << size << " vs expected " << prop.size();218 for (auto &elt : prop) {219 uint64_t value;220 if (failed(reader.readVarInt(value)))221 return failure();222 elt = value;223 }224 return success();225}226 227void test::writeToMlirBytecode(DialectBytecodeWriter &writer,228 ArrayRef<int64_t> prop) {229 writer.writeVarInt(prop.size());230 for (auto elt : prop)231 writer.writeVarInt(elt);232}233 234//===----------------------------------------------------------------------===//235// Dynamic operations236//===----------------------------------------------------------------------===//237 238static std::unique_ptr<DynamicOpDefinition>239getDynamicGenericOp(TestDialect *dialect) {240 return DynamicOpDefinition::get(241 "dynamic_generic", dialect, [](Operation *op) { return success(); },242 [](Operation *op) { return success(); });243}244 245static std::unique_ptr<DynamicOpDefinition>246getDynamicOneOperandTwoResultsOp(TestDialect *dialect) {247 return DynamicOpDefinition::get(248 "dynamic_one_operand_two_results", dialect,249 [](Operation *op) {250 if (op->getNumOperands() != 1) {251 op->emitOpError()252 << "expected 1 operand, but had " << op->getNumOperands();253 return failure();254 }255 if (op->getNumResults() != 2) {256 op->emitOpError()257 << "expected 2 results, but had " << op->getNumResults();258 return failure();259 }260 return success();261 },262 [](Operation *op) { return success(); });263}264 265static std::unique_ptr<DynamicOpDefinition>266getDynamicCustomParserPrinterOp(TestDialect *dialect) {267 auto verifier = [](Operation *op) {268 if (op->getNumOperands() == 0 && op->getNumResults() == 0)269 return success();270 op->emitError() << "operation should have no operands and no results";271 return failure();272 };273 auto regionVerifier = [](Operation *op) { return success(); };274 275 auto parser = [](OpAsmParser &parser, OperationState &state) {276 return parser.parseKeyword("custom_keyword");277 };278 279 auto printer = [](Operation *op, OpAsmPrinter &printer, llvm::StringRef) {280 printer << op->getName() << " custom_keyword";281 };282 283 return DynamicOpDefinition::get("dynamic_custom_parser_printer", dialect,284 verifier, regionVerifier, parser, printer);285}286 287//===----------------------------------------------------------------------===//288// TestDialect289//===----------------------------------------------------------------------===//290 291void test::registerTestDialect(DialectRegistry ®istry) {292 registry.insert<TestDialect>();293}294 295void test::testSideEffectOpGetEffect(296 Operation *op,297 SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>>298 &effects) {299 auto effectsAttr = op->getAttrOfType<AffineMapAttr>("effect_parameter");300 if (!effectsAttr)301 return;302 303 effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);304}305 306// This is the implementation of a dialect fallback for `TestEffectOpInterface`.307struct TestOpEffectInterfaceFallback308 : public TestEffectOpInterface::FallbackModel<309 TestOpEffectInterfaceFallback> {310 static bool classof(Operation *op) {311 bool isSupportedOp =312 op->getName().getStringRef() == "test.unregistered_side_effect_op";313 assert(isSupportedOp && "Unexpected dispatch");314 return isSupportedOp;315 }316 317 void318 getEffects(Operation *op,319 SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>>320 &effects) const {321 testSideEffectOpGetEffect(op, effects);322 }323};324 325void TestDialect::initialize() {326 registerAttributes();327 registerTypes();328 registerOpsSyntax();329 addOperations<ManualCppOpWithFold>();330 registerTestDialectOperations(this);331 registerDynamicOp(getDynamicGenericOp(this));332 registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));333 registerDynamicOp(getDynamicCustomParserPrinterOp(this));334 registerInterfaces();335 allowUnknownOperations();336 337 // Instantiate our fallback op interface that we'll use on specific338 // unregistered op.339 fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;340}341 342TestDialect::~TestDialect() {343 delete static_cast<TestOpEffectInterfaceFallback *>(344 fallbackEffectOpInterfaces);345}346 347Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,348 Type type, Location loc) {349 return TestOpConstant::create(builder, loc, type, value);350}351 352void *TestDialect::getRegisteredInterfaceForOp(TypeID typeID,353 OperationName opName) {354 if (opName.getIdentifier() == "test.unregistered_side_effect_op" &&355 typeID == TypeID::get<TestEffectOpInterface>())356 return fallbackEffectOpInterfaces;357 return nullptr;358}359 360LogicalResult TestDialect::verifyOperationAttribute(Operation *op,361 NamedAttribute namedAttr) {362 if (namedAttr.getName() == "test.invalid_attr")363 return op->emitError() << "invalid to use 'test.invalid_attr'";364 return success();365}366 367LogicalResult TestDialect::verifyRegionArgAttribute(Operation *op,368 unsigned regionIndex,369 unsigned argIndex,370 NamedAttribute namedAttr) {371 if (namedAttr.getName() == "test.invalid_attr")372 return op->emitError() << "invalid to use 'test.invalid_attr'";373 return success();374}375 376LogicalResult377TestDialect::verifyRegionResultAttribute(Operation *op, unsigned regionIndex,378 unsigned resultIndex,379 NamedAttribute namedAttr) {380 if (namedAttr.getName() == "test.invalid_attr")381 return op->emitError() << "invalid to use 'test.invalid_attr'";382 return success();383}384 385std::optional<Dialect::ParseOpHook>386TestDialect::getParseOperationHook(StringRef opName) const {387 if (opName == "test.dialect_custom_printer") {388 return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {389 return parser.parseKeyword("custom_format");390 }};391 }392 if (opName == "test.dialect_custom_format_fallback") {393 return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {394 return parser.parseKeyword("custom_format_fallback");395 }};396 }397 if (opName == "test.dialect_custom_printer.with.dot") {398 return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {399 return ParseResult::success();400 }};401 }402 return std::nullopt;403}404 405llvm::unique_function<void(Operation *, OpAsmPrinter &)>406TestDialect::getOperationPrinter(Operation *op) const {407 StringRef opName = op->getName().getStringRef();408 if (opName == "test.dialect_custom_printer") {409 return [](Operation *op, OpAsmPrinter &printer) {410 printer.getStream() << " custom_format";411 };412 }413 if (opName == "test.dialect_custom_format_fallback") {414 return [](Operation *op, OpAsmPrinter &printer) {415 printer.getStream() << " custom_format_fallback";416 };417 }418 return {};419}420 421static LogicalResult422dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,423 PatternRewriter &rewriter) {424 rewriter.replaceOpWithNewOp<arith::ConstantOp>(425 op, rewriter.getI32IntegerAttr(42));426 return success();427}428 429void TestDialect::getCanonicalizationPatterns(430 RewritePatternSet &results) const {431 results.add(&dialectCanonicalizationPattern);432}433 434//===----------------------------------------------------------------------===//435// TestCallWithSegmentsOp436//===----------------------------------------------------------------------===//437// The op `test.call_with_segments` models a call-like operation whose operands438// are divided into 3 variadic segments: `prefix`, `args`, and `suffix`.439// Only the middle segment represents the actual call arguments. The op uses440// the AttrSizedOperandSegments trait, so we can derive segment boundaries from441// the generated `operandSegmentSizes` attribute. We provide custom helpers to442// expose the logical call arguments as both a read-only range and a mutable443// range bound to the proper segment so that insertion/erasure updates the444// attribute automatically.445 446// Segment layout indices in the DenseI32ArrayAttr: [prefix, args, suffix].447static constexpr unsigned kTestCallWithSegmentsArgsSegIndex = 1;448 449Operation::operand_range CallWithSegmentsOp::getArgOperands() {450 // Leverage generated getters for segment sizes: slice between prefix and451 // suffix using current operand list.452 return getOperation()->getOperands().slice(getPrefix().size(),453 getArgs().size());454}455 456MutableOperandRange CallWithSegmentsOp::getArgOperandsMutable() {457 Operation *op = getOperation();458 459 // Obtain the canonical segment size attribute name for this op.460 auto segName =461 CallWithSegmentsOp::getOperandSegmentSizesAttrName(op->getName());462 auto sizesAttr = op->getAttrOfType<DenseI32ArrayAttr>(segName);463 assert(sizesAttr && "missing operandSegmentSizes attribute on op");464 465 // Compute the start and length of the args segment from the prefix size and466 // args size stored in the attribute.467 auto sizes = sizesAttr.asArrayRef();468 unsigned start = static_cast<unsigned>(sizes[0]); // prefix size469 unsigned len = static_cast<unsigned>(sizes[1]); // args size470 471 NamedAttribute segNamed(segName, sizesAttr);472 MutableOperandRange::OperandSegment binding{kTestCallWithSegmentsArgsSegIndex,473 segNamed};474 475 return MutableOperandRange(op, start, len, {binding});476}477