1640 lines · cpp
1//===- TestOpDefs.cpp - MLIR Test Dialect Operation Hooks -----------------===//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 "mlir/Dialect/Bufferization/IR/Bufferization.h"12#include "mlir/Dialect/Tensor/IR/Tensor.h"13#include "mlir/IR/Verifier.h"14#include "mlir/Interfaces/FunctionImplementation.h"15#include "mlir/Interfaces/MemorySlotInterfaces.h"16 17using namespace mlir;18using namespace test;19 20//===----------------------------------------------------------------------===//21// OverridenSymbolVisibilityOp22//===----------------------------------------------------------------------===//23 24SymbolTable::Visibility OverriddenSymbolVisibilityOp::getVisibility() {25 return SymbolTable::Visibility::Private;26}27 28static StringLiteral getVisibilityString(SymbolTable::Visibility visibility) {29 switch (visibility) {30 case SymbolTable::Visibility::Private:31 return "private";32 case SymbolTable::Visibility::Nested:33 return "nested";34 case SymbolTable::Visibility::Public:35 return "public";36 }37}38 39void OverriddenSymbolVisibilityOp::setVisibility(40 SymbolTable::Visibility visibility) {41 42 emitOpError("cannot change visibility of symbol to ")43 << getVisibilityString(visibility);44}45 46//===----------------------------------------------------------------------===//47// TestBranchOp48//===----------------------------------------------------------------------===//49 50SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {51 assert(index == 0 && "invalid successor index");52 return SuccessorOperands(getTargetOperandsMutable());53}54 55//===----------------------------------------------------------------------===//56// TestProducingBranchOp57//===----------------------------------------------------------------------===//58 59SuccessorOperands TestProducingBranchOp::getSuccessorOperands(unsigned index) {60 assert(index <= 1 && "invalid successor index");61 if (index == 1)62 return SuccessorOperands(getFirstOperandsMutable());63 return SuccessorOperands(getSecondOperandsMutable());64}65 66//===----------------------------------------------------------------------===//67// TestInternalBranchOp68//===----------------------------------------------------------------------===//69 70SuccessorOperands TestInternalBranchOp::getSuccessorOperands(unsigned index) {71 assert(index <= 1 && "invalid successor index");72 if (index == 0)73 return SuccessorOperands(0, getSuccessOperandsMutable());74 return SuccessorOperands(1, getErrorOperandsMutable());75}76 77//===----------------------------------------------------------------------===//78// TestCallOp79//===----------------------------------------------------------------------===//80 81LogicalResult TestCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {82 // Check that the callee attribute was specified.83 auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");84 if (!fnAttr)85 return emitOpError("requires a 'callee' symbol reference attribute");86 if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))87 return emitOpError() << "'" << fnAttr.getValue()88 << "' does not reference a valid function";89 return success();90}91 92//===----------------------------------------------------------------------===//93// FoldToCallOp94//===----------------------------------------------------------------------===//95 96namespace {97struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {98 using OpRewritePattern<FoldToCallOp>::OpRewritePattern;99 100 LogicalResult matchAndRewrite(FoldToCallOp op,101 PatternRewriter &rewriter) const override {102 rewriter.replaceOpWithNewOp<func::CallOp>(op, TypeRange(),103 op.getCalleeAttr(), ValueRange());104 return success();105 }106};107} // namespace108 109void FoldToCallOp::getCanonicalizationPatterns(RewritePatternSet &results,110 MLIRContext *context) {111 results.add<FoldToCallOpPattern>(context);112}113 114//===----------------------------------------------------------------------===//115// IsolatedRegionOp - test parsing passthrough operands116//===----------------------------------------------------------------------===//117 118ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,119 OperationState &result) {120 // Parse the input operand.121 OpAsmParser::Argument argInfo;122 argInfo.type = parser.getBuilder().getIndexType();123 if (parser.parseOperand(argInfo.ssaName) ||124 parser.resolveOperand(argInfo.ssaName, argInfo.type, result.operands))125 return failure();126 127 // Parse the body region, and reuse the operand info as the argument info.128 Region *body = result.addRegion();129 return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);130}131 132void IsolatedRegionOp::print(OpAsmPrinter &p) {133 p << ' ';134 p.printOperand(getOperand());135 p.shadowRegionArgs(getRegion(), getOperand());136 p << ' ';137 p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);138}139 140//===----------------------------------------------------------------------===//141// SSACFGRegionOp142//===----------------------------------------------------------------------===//143 144RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {145 return RegionKind::SSACFG;146}147 148//===----------------------------------------------------------------------===//149// GraphRegionOp150//===----------------------------------------------------------------------===//151 152RegionKind GraphRegionOp::getRegionKind(unsigned index) {153 return RegionKind::Graph;154}155 156//===----------------------------------------------------------------------===//157// IsolatedGraphRegionOp158//===----------------------------------------------------------------------===//159 160RegionKind IsolatedGraphRegionOp::getRegionKind(unsigned index) {161 return RegionKind::Graph;162}163 164//===----------------------------------------------------------------------===//165// AffineScopeOp166//===----------------------------------------------------------------------===//167 168ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {169 // Parse the body region, and reuse the operand info as the argument info.170 Region *body = result.addRegion();171 return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});172}173 174void AffineScopeOp::print(OpAsmPrinter &p) {175 p << " ";176 p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);177}178 179//===----------------------------------------------------------------------===//180// TestRemoveOpWithInnerOps181//===----------------------------------------------------------------------===//182 183namespace {184struct TestRemoveOpWithInnerOps185 : public OpRewritePattern<TestOpWithRegionPattern> {186 using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;187 188 void initialize() { setDebugName("TestRemoveOpWithInnerOps"); }189 190 LogicalResult matchAndRewrite(TestOpWithRegionPattern op,191 PatternRewriter &rewriter) const override {192 rewriter.eraseOp(op);193 return success();194 }195};196} // namespace197 198//===----------------------------------------------------------------------===//199// TestOpWithRegionPattern200//===----------------------------------------------------------------------===//201 202void TestOpWithRegionPattern::getCanonicalizationPatterns(203 RewritePatternSet &results, MLIRContext *context) {204 results.add<TestRemoveOpWithInnerOps>(context);205}206 207//===----------------------------------------------------------------------===//208// TestOpWithRegionFold209//===----------------------------------------------------------------------===//210 211OpFoldResult TestOpWithRegionFold::fold(FoldAdaptor adaptor) {212 return getOperand();213}214 215//===----------------------------------------------------------------------===//216// TestOpConstant217//===----------------------------------------------------------------------===//218 219OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) { return getValue(); }220 221//===----------------------------------------------------------------------===//222// TestOpWithVariadicResultsAndFolder223//===----------------------------------------------------------------------===//224 225LogicalResult TestOpWithVariadicResultsAndFolder::fold(226 FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {227 for (Value input : this->getOperands()) {228 results.push_back(input);229 }230 return success();231}232 233//===----------------------------------------------------------------------===//234// TestOpInPlaceFold235//===----------------------------------------------------------------------===//236 237OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {238 // Exercise the fact that an operation created with createOrFold should be239 // allowed to access its parent block.240 assert(getOperation()->getBlock() &&241 "expected that operation is not unlinked");242 243 if (adaptor.getOp() && !getProperties().attr) {244 // The folder adds "attr" if not present.245 getProperties().attr = dyn_cast_or_null<IntegerAttr>(adaptor.getOp());246 return getResult();247 }248 return {};249}250 251//===----------------------------------------------------------------------===//252// OpWithInferTypeInterfaceOp253//===----------------------------------------------------------------------===//254 255LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(256 MLIRContext *, std::optional<Location> location, ValueRange operands,257 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,258 SmallVectorImpl<Type> &inferredReturnTypes) {259 if (operands[0].getType() != operands[1].getType()) {260 return emitOptionalError(location, "operand type mismatch ",261 operands[0].getType(), " vs ",262 operands[1].getType());263 }264 inferredReturnTypes.assign({operands[0].getType()});265 return success();266}267 268//===----------------------------------------------------------------------===//269// OpWithShapedTypeInferTypeInterfaceOp270//===----------------------------------------------------------------------===//271 272LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(273 MLIRContext *context, std::optional<Location> location,274 ValueShapeRange operands, DictionaryAttr attributes,275 OpaqueProperties properties, RegionRange regions,276 SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {277 // Create return type consisting of the last element of the first operand.278 auto operandType = operands.front().getType();279 auto sval = dyn_cast<ShapedType>(operandType);280 if (!sval)281 return emitOptionalError(location, "only shaped type operands allowed");282 int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamic;283 auto type = IntegerType::get(context, 17);284 285 Attribute encoding;286 if (auto rankedTy = dyn_cast<RankedTensorType>(sval))287 encoding = rankedTy.getEncoding();288 inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));289 return success();290}291 292LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(293 OpBuilder &builder, ValueRange operands,294 llvm::SmallVectorImpl<Value> &shapes) {295 shapes = SmallVector<Value, 1>{296 builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};297 return success();298}299 300//===----------------------------------------------------------------------===//301// OpWithResultShapeInterfaceOp302//===----------------------------------------------------------------------===//303 304LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(305 OpBuilder &builder, ValueRange operands,306 llvm::SmallVectorImpl<Value> &shapes) {307 Location loc = getLoc();308 shapes.reserve(operands.size());309 for (Value operand : llvm::reverse(operands)) {310 auto rank = cast<RankedTensorType>(operand.getType()).getRank();311 auto currShape = llvm::to_vector<4>(312 llvm::map_range(llvm::seq<int64_t>(0, rank), [&](int64_t dim) -> Value {313 return builder.createOrFold<tensor::DimOp>(loc, operand, dim);314 }));315 shapes.push_back(tensor::FromElementsOp::create(316 builder, getLoc(),317 RankedTensorType::get({rank}, builder.getIndexType()), currShape));318 }319 return success();320}321 322//===----------------------------------------------------------------------===//323// ReifyShapedTypeUsingReifyResultShapesOp324//===----------------------------------------------------------------------===//325 326LogicalResult ReifyShapedTypeUsingReifyResultShapesOp::reifyResultShapes(327 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {328 Location loc = getLoc();329 shapes.reserve(getNumOperands());330 for (Value operand : llvm::reverse(getOperands())) {331 auto tensorType = cast<RankedTensorType>(operand.getType());332 auto currShape = llvm::to_vector<4>(llvm::map_range(333 llvm::seq<int64_t>(0, tensorType.getRank()),334 [&](int64_t dim) -> OpFoldResult {335 return tensorType.isDynamicDim(dim)336 ? static_cast<OpFoldResult>(337 builder.createOrFold<tensor::DimOp>(loc, operand,338 dim))339 : static_cast<OpFoldResult>(340 builder.getIndexAttr(tensorType.getDimSize(dim)));341 }));342 shapes.emplace_back(std::move(currShape));343 }344 return success();345}346 347//===----------------------------------------------------------------------===//348// ReifyShapedTypeUsingReifyShapeOfResultOp349//===----------------------------------------------------------------------===//350 351LogicalResult ReifyShapedTypeUsingReifyShapeOfResultOp::reifyResultShapes(352 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {353 return failure();354}355 356FailureOr<SmallVector<OpFoldResult>>357ReifyShapedTypeUsingReifyShapeOfResultOp::reifyShapeOfResult(OpBuilder &builder,358 int resultIndex) {359 Location loc = getLoc();360 Value sourceOperand = getOperand(getNumOperands() - 1 - resultIndex);361 SmallVector<OpFoldResult> shape =362 tensor::getMixedSizes(builder, loc, sourceOperand);363 return shape;364}365 366//===----------------------------------------------------------------------===//367// ReifyShapedTypeUsingReifyDimOfResultOp368//===----------------------------------------------------------------------===//369 370LogicalResult ReifyShapedTypeUsingReifyDimOfResultOp::reifyResultShapes(371 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {372 return failure();373}374 375FailureOr<SmallVector<OpFoldResult>>376ReifyShapedTypeUsingReifyDimOfResultOp::reifyShapeOfResult(OpBuilder &builder,377 int resultIndex) {378 return failure();379}380 381FailureOr<OpFoldResult>382ReifyShapedTypeUsingReifyDimOfResultOp::reifyDimOfResult(OpBuilder &builder,383 int resultIndex,384 int dim) {385 Location loc = getLoc();386 Value sourceOperand = getOperand(getNumOperands() - 1 - resultIndex);387 OpFoldResult shape = tensor::getMixedSize(builder, loc, sourceOperand, dim);388 return shape;389}390 391//===----------------------------------------------------------------------===//392// UnreifableResultShapesOp393//===----------------------------------------------------------------------===//394 395LogicalResult UnreifiableResultShapesOp::reifyResultShapes(396 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {397 Location loc = getLoc();398 shapes.resize(1);399 shapes[0] = {tensor::getMixedSize(builder, loc, getOperand(), 0),400 OpFoldResult()};401 return success();402}403 404//===----------------------------------------------------------------------===//405// UnreifableResultShapeOp406//===----------------------------------------------------------------------===//407 408LogicalResult UnreifiableResultShapeOp::reifyResultShapes(409 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {410 return failure();411}412 413FailureOr<SmallVector<OpFoldResult>>414UnreifiableResultShapeOp::reifyShapeOfResult(OpBuilder &builder,415 int resultIndex) {416 SmallVector<OpFoldResult> shape = {417 tensor::getMixedSize(builder, getLoc(), getOperand(), 0), OpFoldResult()};418 return shape;419}420 421//===----------------------------------------------------------------------===//422// UnreifableResultShapeOp423//===----------------------------------------------------------------------===//424 425LogicalResult UnreifiableDimOfResultShapeOp::reifyResultShapes(426 OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {427 return failure();428}429 430FailureOr<SmallVector<OpFoldResult>>431UnreifiableDimOfResultShapeOp::reifyShapeOfResult(OpBuilder &builder,432 int resultIndex) {433 return failure();434}435 436FailureOr<OpFoldResult>437UnreifiableDimOfResultShapeOp::reifyDimOfResult(OpBuilder &builder,438 int resultIndex, int dim) {439 if (dim == 0)440 return tensor::getMixedSize(builder, getLoc(), getOperand(), 0);441 return failure();442}443 444//===----------------------------------------------------------------------===//445// SideEffectOp446//===----------------------------------------------------------------------===//447 448namespace {449/// A test resource for side effects.450struct TestResource : public SideEffects::Resource::Base<TestResource> {451 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestResource)452 453 StringRef getName() final { return "<Test>"; }454};455} // namespace456 457void SideEffectOp::getEffects(458 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {459 // Check for an effects attribute on the op instance.460 ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");461 if (!effectsAttr)462 return;463 464 for (Attribute element : effectsAttr) {465 DictionaryAttr effectElement = cast<DictionaryAttr>(element);466 467 // Get the specific memory effect.468 MemoryEffects::Effect *effect =469 StringSwitch<MemoryEffects::Effect *>(470 cast<StringAttr>(effectElement.get("effect")).getValue())471 .Case("allocate", MemoryEffects::Allocate::get())472 .Case("free", MemoryEffects::Free::get())473 .Case("read", MemoryEffects::Read::get())474 .Case("write", MemoryEffects::Write::get());475 476 // Check for a non-default resource to use.477 SideEffects::Resource *resource = SideEffects::DefaultResource::get();478 if (effectElement.get("test_resource"))479 resource = TestResource::get();480 481 // Check for a result to affect.482 if (effectElement.get("on_result"))483 effects.emplace_back(effect, getOperation()->getOpResults()[0], resource);484 else if (Attribute ref = effectElement.get("on_reference"))485 effects.emplace_back(effect, cast<SymbolRefAttr>(ref), resource);486 else487 effects.emplace_back(effect, resource);488 }489}490 491void SideEffectOp::getEffects(492 SmallVectorImpl<TestEffects::EffectInstance> &effects) {493 testSideEffectOpGetEffect(getOperation(), effects);494}495 496void SideEffectWithRegionOp::getEffects(497 SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {498 // Check for an effects attribute on the op instance.499 ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");500 if (!effectsAttr)501 return;502 503 for (Attribute element : effectsAttr) {504 DictionaryAttr effectElement = cast<DictionaryAttr>(element);505 506 // Get the specific memory effect.507 MemoryEffects::Effect *effect =508 StringSwitch<MemoryEffects::Effect *>(509 cast<StringAttr>(effectElement.get("effect")).getValue())510 .Case("allocate", MemoryEffects::Allocate::get())511 .Case("free", MemoryEffects::Free::get())512 .Case("read", MemoryEffects::Read::get())513 .Case("write", MemoryEffects::Write::get());514 515 // Check for a non-default resource to use.516 SideEffects::Resource *resource = SideEffects::DefaultResource::get();517 if (effectElement.get("test_resource"))518 resource = TestResource::get();519 520 // Check for a result to affect.521 if (effectElement.get("on_result"))522 effects.emplace_back(effect, getOperation()->getOpResults()[0], resource);523 else if (effectElement.get("on_operand"))524 effects.emplace_back(effect, &getOperation()->getOpOperands()[0],525 resource);526 else if (effectElement.get("on_argument"))527 effects.emplace_back(effect, getOperation()->getRegion(0).getArgument(0),528 resource);529 else if (Attribute ref = effectElement.get("on_reference"))530 effects.emplace_back(effect, cast<SymbolRefAttr>(ref), resource);531 else532 effects.emplace_back(effect, resource);533 }534}535 536void SideEffectWithRegionOp::getEffects(537 SmallVectorImpl<TestEffects::EffectInstance> &effects) {538 testSideEffectOpGetEffect(getOperation(), effects);539}540 541//===----------------------------------------------------------------------===//542// StringAttrPrettyNameOp543//===----------------------------------------------------------------------===//544 545// This op has fancy handling of its SSA result name.546ParseResult StringAttrPrettyNameOp::parse(OpAsmParser &parser,547 OperationState &result) {548 // Add the result types.549 for (size_t i = 0, e = parser.getNumResults(); i != e; ++i)550 result.addTypes(parser.getBuilder().getIntegerType(32));551 552 if (parser.parseOptionalAttrDictWithKeyword(result.attributes))553 return failure();554 555 // If the attribute dictionary contains no 'names' attribute, infer it from556 // the SSA name (if specified).557 bool hadNames = llvm::any_of(result.attributes, [](NamedAttribute attr) {558 return attr.getName() == "names";559 });560 561 // If there was no name specified, check to see if there was a useful name562 // specified in the asm file.563 if (hadNames || parser.getNumResults() == 0)564 return success();565 566 SmallVector<StringRef, 4> names;567 auto *context = result.getContext();568 569 for (size_t i = 0, e = parser.getNumResults(); i != e; ++i) {570 auto resultName = parser.getResultName(i);571 StringRef nameStr;572 if (!resultName.first.empty() && !isdigit(resultName.first[0]))573 nameStr = resultName.first;574 575 names.push_back(nameStr);576 }577 578 auto namesAttr = parser.getBuilder().getStrArrayAttr(names);579 result.attributes.push_back({StringAttr::get(context, "names"), namesAttr});580 return success();581}582 583void StringAttrPrettyNameOp::print(OpAsmPrinter &p) {584 // Note that we only need to print the "name" attribute if the asmprinter585 // result name disagrees with it. This can happen in strange cases, e.g.586 // when there are conflicts.587 bool namesDisagree = getNames().size() != getNumResults();588 589 SmallString<32> resultNameStr;590 for (size_t i = 0, e = getNumResults(); i != e && !namesDisagree; ++i) {591 resultNameStr.clear();592 llvm::raw_svector_ostream tmpStream(resultNameStr);593 p.printOperand(getResult(i), tmpStream);594 595 auto expectedName = dyn_cast<StringAttr>(getNames()[i]);596 if (!expectedName ||597 tmpStream.str().drop_front() != expectedName.getValue()) {598 namesDisagree = true;599 }600 }601 602 if (namesDisagree)603 p.printOptionalAttrDictWithKeyword((*this)->getAttrs());604 else605 p.printOptionalAttrDictWithKeyword((*this)->getAttrs(), {"names"});606}607 608// We set the SSA name in the asm syntax to the contents of the name609// attribute.610void StringAttrPrettyNameOp::getAsmResultNames(611 function_ref<void(Value, StringRef)> setNameFn) {612 613 auto value = getNames();614 for (size_t i = 0, e = value.size(); i != e; ++i)615 if (auto str = dyn_cast<StringAttr>(value[i]))616 if (!str.getValue().empty())617 setNameFn(getResult(i), str.getValue());618}619 620//===----------------------------------------------------------------------===//621// CustomResultsNameOp622//===----------------------------------------------------------------------===//623 624void CustomResultsNameOp::getAsmResultNames(625 function_ref<void(Value, StringRef)> setNameFn) {626 ArrayAttr value = getNames();627 for (size_t i = 0, e = value.size(); i != e; ++i)628 if (auto str = dyn_cast<StringAttr>(value[i]))629 if (!str.empty())630 setNameFn(getResult(i), str.getValue());631}632 633//===----------------------------------------------------------------------===//634// ResultNameFromTypeOp635//===----------------------------------------------------------------------===//636 637void ResultNameFromTypeOp::getAsmResultNames(638 function_ref<void(Value, StringRef)> setNameFn) {639 auto result = getResult();640 auto setResultNameFn = [&](::llvm::StringRef name) {641 setNameFn(result, name);642 };643 auto opAsmTypeInterface =644 ::mlir::cast<::mlir::OpAsmTypeInterface>(result.getType());645 opAsmTypeInterface.getAsmName(setResultNameFn);646}647 648//===----------------------------------------------------------------------===//649// BlockArgumentNameFromTypeOp650//===----------------------------------------------------------------------===//651 652void BlockArgumentNameFromTypeOp::getAsmBlockArgumentNames(653 ::mlir::Region ®ion, ::mlir::OpAsmSetValueNameFn setNameFn) {654 for (auto &block : region) {655 for (auto arg : block.getArguments()) {656 if (auto opAsmTypeInterface =657 ::mlir::dyn_cast<::mlir::OpAsmTypeInterface>(arg.getType())) {658 auto setArgNameFn = [&](StringRef name) { setNameFn(arg, name); };659 opAsmTypeInterface.getAsmName(setArgNameFn);660 }661 }662 }663}664 665//===----------------------------------------------------------------------===//666// ResultTypeWithTraitOp667//===----------------------------------------------------------------------===//668 669LogicalResult ResultTypeWithTraitOp::verify() {670 if ((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())671 return success();672 return emitError("result type should have trait 'TestTypeTrait'");673}674 675//===----------------------------------------------------------------------===//676// AttrWithTraitOp677//===----------------------------------------------------------------------===//678 679LogicalResult AttrWithTraitOp::verify() {680 if (getAttr().hasTrait<AttributeTrait::TestAttrTrait>())681 return success();682 return emitError("'attr' attribute should have trait 'TestAttrTrait'");683}684 685//===----------------------------------------------------------------------===//686// RegionIfOp687//===----------------------------------------------------------------------===//688 689void RegionIfOp::print(OpAsmPrinter &p) {690 p << " ";691 p.printOperands(getOperands());692 p << ": " << getOperandTypes();693 p.printArrowTypeList(getResultTypes());694 p << " then ";695 p.printRegion(getThenRegion(),696 /*printEntryBlockArgs=*/true,697 /*printBlockTerminators=*/true);698 p << " else ";699 p.printRegion(getElseRegion(),700 /*printEntryBlockArgs=*/true,701 /*printBlockTerminators=*/true);702 p << " join ";703 p.printRegion(getJoinRegion(),704 /*printEntryBlockArgs=*/true,705 /*printBlockTerminators=*/true);706}707 708ParseResult RegionIfOp::parse(OpAsmParser &parser, OperationState &result) {709 SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfos;710 SmallVector<Type, 2> operandTypes;711 712 result.regions.reserve(3);713 Region *thenRegion = result.addRegion();714 Region *elseRegion = result.addRegion();715 Region *joinRegion = result.addRegion();716 717 // Parse operand, type and arrow type lists.718 if (parser.parseOperandList(operandInfos) ||719 parser.parseColonTypeList(operandTypes) ||720 parser.parseArrowTypeList(result.types))721 return failure();722 723 // Parse all attached regions.724 if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||725 parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||726 parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))727 return failure();728 729 return parser.resolveOperands(operandInfos, operandTypes,730 parser.getCurrentLocation(), result.operands);731}732 733OperandRange RegionIfOp::getEntrySuccessorOperands(RegionSuccessor successor) {734 assert(llvm::is_contained({&getThenRegion(), &getElseRegion()},735 successor.getSuccessor()) &&736 "invalid region index");737 return getOperands();738}739 740void RegionIfOp::getSuccessorRegions(741 RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {742 // We always branch to the join region.743 if (!point.isParent()) {744 if (point.getTerminatorPredecessorOrNull()->getParentRegion() !=745 &getJoinRegion())746 regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));747 else748 regions.push_back(RegionSuccessor(getOperation(), getResults()));749 return;750 }751 752 // The then and else regions are the entry regions of this op.753 regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));754 regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));755}756 757void RegionIfOp::getRegionInvocationBounds(758 ArrayRef<Attribute> operands,759 SmallVectorImpl<InvocationBounds> &invocationBounds) {760 // Each region is invoked at most once.761 invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});762}763 764//===----------------------------------------------------------------------===//765// AnyCondOp766//===----------------------------------------------------------------------===//767 768void AnyCondOp::getSuccessorRegions(RegionBranchPoint point,769 SmallVectorImpl<RegionSuccessor> ®ions) {770 // The parent op branches into the only region, and the region branches back771 // to the parent op.772 if (point.isParent())773 regions.emplace_back(&getRegion());774 else775 regions.emplace_back(getOperation(), getResults());776}777 778void AnyCondOp::getRegionInvocationBounds(779 ArrayRef<Attribute> operands,780 SmallVectorImpl<InvocationBounds> &invocationBounds) {781 invocationBounds.emplace_back(1, 1);782}783 784//===----------------------------------------------------------------------===//785// SingleBlockImplicitTerminatorOp786//===----------------------------------------------------------------------===//787 788/// Testing the correctness of some traits.789static_assert(790 llvm::is_detected<OpTrait::has_implicit_terminator_t,791 SingleBlockImplicitTerminatorOp>::value,792 "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");793static_assert(OpTrait::hasSingleBlockImplicitTerminator<794 SingleBlockImplicitTerminatorOp>::value,795 "hasSingleBlockImplicitTerminator does not match "796 "SingleBlockImplicitTerminatorOp");797 798//===----------------------------------------------------------------------===//799// SingleNoTerminatorCustomAsmOp800//===----------------------------------------------------------------------===//801 802ParseResult SingleNoTerminatorCustomAsmOp::parse(OpAsmParser &parser,803 OperationState &state) {804 Region *body = state.addRegion();805 if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))806 return failure();807 return success();808}809 810void SingleNoTerminatorCustomAsmOp::print(OpAsmPrinter &printer) {811 printer.printRegion(812 getRegion(), /*printEntryBlockArgs=*/false,813 // This op has a single block without terminators. But explicitly mark814 // as not printing block terminators for testing.815 /*printBlockTerminators=*/false);816}817 818//===----------------------------------------------------------------------===//819// TestVerifiersOp820//===----------------------------------------------------------------------===//821 822LogicalResult TestVerifiersOp::verify() {823 if (!getRegion().hasOneBlock())824 return emitOpError("`hasOneBlock` trait hasn't been verified");825 826 Operation *definingOp = getInput().getDefiningOp();827 if (definingOp && failed(mlir::verify(definingOp)))828 return emitOpError("operand hasn't been verified");829 830 // Avoid using `emitRemark(msg)` since that will trigger an infinite verifier831 // loop.832 mlir::emitRemark(getLoc(), "success run of verifier");833 834 return success();835}836 837LogicalResult TestVerifiersOp::verifyRegions() {838 if (!getRegion().hasOneBlock())839 return emitOpError("`hasOneBlock` trait hasn't been verified");840 841 for (Block &block : getRegion())842 for (Operation &op : block)843 if (failed(mlir::verify(&op)))844 return emitOpError("nested op hasn't been verified");845 846 // Avoid using `emitRemark(msg)` since that will trigger an infinite verifier847 // loop.848 mlir::emitRemark(getLoc(), "success run of region verifier");849 850 return success();851}852 853//===----------------------------------------------------------------------===//854// Test InferIntRangeInterface855//===----------------------------------------------------------------------===//856 857//===----------------------------------------------------------------------===//858// TestWithBoundsOp859//===----------------------------------------------------------------------===//860 861void TestWithBoundsOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,862 SetIntRangeFn setResultRanges) {863 setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});864}865 866//===----------------------------------------------------------------------===//867// TestWithBoundsRegionOp868//===----------------------------------------------------------------------===//869 870ParseResult TestWithBoundsRegionOp::parse(OpAsmParser &parser,871 OperationState &result) {872 if (parser.parseOptionalAttrDict(result.attributes))873 return failure();874 875 // Parse the input argument876 OpAsmParser::Argument argInfo;877 if (failed(parser.parseArgument(argInfo, true)))878 return failure();879 880 // Parse the body region, and reuse the operand info as the argument info.881 Region *body = result.addRegion();882 return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/false);883}884 885void TestWithBoundsRegionOp::print(OpAsmPrinter &p) {886 p.printOptionalAttrDict((*this)->getAttrs());887 p << ' ';888 p.printRegionArgument(getRegion().getArgument(0), /*argAttrs=*/{},889 /*omitType=*/false);890 p << ' ';891 p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);892}893 894void TestWithBoundsRegionOp::inferResultRanges(895 ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {896 Value arg = getRegion().getArgument(0);897 setResultRanges(arg, {getUmin(), getUmax(), getSmin(), getSmax()});898}899 900//===----------------------------------------------------------------------===//901// TestIncrementOp902//===----------------------------------------------------------------------===//903 904void TestIncrementOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,905 SetIntRangeFn setResultRanges) {906 const ConstantIntRanges &range = argRanges[0];907 APInt one(range.umin().getBitWidth(), 1);908 setResultRanges(getResult(),909 {range.umin().uadd_sat(one), range.umax().uadd_sat(one),910 range.smin().sadd_sat(one), range.smax().sadd_sat(one)});911}912 913//===----------------------------------------------------------------------===//914// TestReflectBoundsOp915//===----------------------------------------------------------------------===//916 917void TestReflectBoundsOp::inferResultRanges(918 ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {919 const ConstantIntRanges &range = argRanges[0];920 MLIRContext *ctx = getContext();921 Builder b(ctx);922 Type sIntTy, uIntTy;923 // For plain `IntegerType`s, we can derive the appropriate signed and unsigned924 // Types for the Attributes.925 Type type = getElementTypeOrSelf(getType());926 if (auto intTy = llvm::dyn_cast<IntegerType>(type)) {927 unsigned bitwidth = intTy.getWidth();928 sIntTy = b.getIntegerType(bitwidth, /*isSigned=*/true);929 uIntTy = b.getIntegerType(bitwidth, /*isSigned=*/false);930 } else {931 sIntTy = uIntTy = type;932 }933 934 setUminAttr(b.getIntegerAttr(uIntTy, range.umin()));935 setUmaxAttr(b.getIntegerAttr(uIntTy, range.umax()));936 setSminAttr(b.getIntegerAttr(sIntTy, range.smin()));937 setSmaxAttr(b.getIntegerAttr(sIntTy, range.smax()));938 setResultRanges(getResult(), range);939}940 941//===----------------------------------------------------------------------===//942// ConversionFuncOp943//===----------------------------------------------------------------------===//944 945ParseResult ConversionFuncOp::parse(OpAsmParser &parser,946 OperationState &result) {947 auto buildFuncType =948 [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,949 function_interface_impl::VariadicFlag,950 std::string &) { return builder.getFunctionType(argTypes, results); };951 952 return function_interface_impl::parseFunctionOp(953 parser, result, /*allowVariadic=*/false,954 getFunctionTypeAttrName(result.name), buildFuncType,955 getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));956}957 958void ConversionFuncOp::print(OpAsmPrinter &p) {959 function_interface_impl::printFunctionOp(960 p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),961 getArgAttrsAttrName(), getResAttrsAttrName());962}963 964//===----------------------------------------------------------------------===//965// TestValueWithBoundsOp966//===----------------------------------------------------------------------===//967 968void TestValueWithBoundsOp::populateBoundsForIndexValue(969 Value v, ValueBoundsConstraintSet &cstr) {970 cstr.bound(v) >= getMin().getSExtValue();971 cstr.bound(v) <= getMax().getSExtValue();972}973 974//===----------------------------------------------------------------------===//975// ReifyBoundOp976//===----------------------------------------------------------------------===//977 978mlir::presburger::BoundType ReifyBoundOp::getBoundType() {979 if (getType() == "EQ")980 return mlir::presburger::BoundType::EQ;981 if (getType() == "LB")982 return mlir::presburger::BoundType::LB;983 if (getType() == "UB")984 return mlir::presburger::BoundType::UB;985 llvm_unreachable("invalid bound type");986}987 988LogicalResult ReifyBoundOp::verify() {989 if (isa<ShapedType>(getVar().getType())) {990 if (!getDim().has_value())991 return emitOpError("expected 'dim' attribute for shaped type variable");992 } else if (getVar().getType().isIndex()) {993 if (getDim().has_value())994 return emitOpError("unexpected 'dim' attribute for index variable");995 } else {996 return emitOpError("expected index-typed variable or shape type variable");997 }998 if (getConstant() && getScalable())999 return emitOpError("'scalable' and 'constant' are mutually exlusive");1000 if (getScalable() != getVscaleMin().has_value())1001 return emitOpError("expected 'vscale_min' if and only if 'scalable'");1002 if (getScalable() != getVscaleMax().has_value())1003 return emitOpError("expected 'vscale_min' if and only if 'scalable'");1004 return success();1005}1006 1007ValueBoundsConstraintSet::Variable ReifyBoundOp::getVariable() {1008 if (getDim().has_value())1009 return ValueBoundsConstraintSet::Variable(getVar(), *getDim());1010 return ValueBoundsConstraintSet::Variable(getVar());1011}1012 1013//===----------------------------------------------------------------------===//1014// CompareOp1015//===----------------------------------------------------------------------===//1016 1017ValueBoundsConstraintSet::ComparisonOperator1018CompareOp::getComparisonOperator() {1019 if (getCmp() == "EQ")1020 return ValueBoundsConstraintSet::ComparisonOperator::EQ;1021 if (getCmp() == "LT")1022 return ValueBoundsConstraintSet::ComparisonOperator::LT;1023 if (getCmp() == "LE")1024 return ValueBoundsConstraintSet::ComparisonOperator::LE;1025 if (getCmp() == "GT")1026 return ValueBoundsConstraintSet::ComparisonOperator::GT;1027 if (getCmp() == "GE")1028 return ValueBoundsConstraintSet::ComparisonOperator::GE;1029 llvm_unreachable("invalid comparison operator");1030}1031 1032mlir::ValueBoundsConstraintSet::Variable CompareOp::getLhs() {1033 if (!getLhsMap())1034 return ValueBoundsConstraintSet::Variable(getVarOperands()[0]);1035 SmallVector<Value> mapOperands(1036 getVarOperands().slice(0, getLhsMap()->getNumInputs()));1037 return ValueBoundsConstraintSet::Variable(*getLhsMap(), mapOperands);1038}1039 1040mlir::ValueBoundsConstraintSet::Variable CompareOp::getRhs() {1041 int64_t rhsOperandsBegin = getLhsMap() ? getLhsMap()->getNumInputs() : 1;1042 if (!getRhsMap())1043 return ValueBoundsConstraintSet::Variable(1044 getVarOperands()[rhsOperandsBegin]);1045 SmallVector<Value> mapOperands(1046 getVarOperands().slice(rhsOperandsBegin, getRhsMap()->getNumInputs()));1047 return ValueBoundsConstraintSet::Variable(*getRhsMap(), mapOperands);1048}1049 1050LogicalResult CompareOp::verify() {1051 if (getCompose() && (getLhsMap() || getRhsMap()))1052 return emitOpError(1053 "'compose' not supported when 'lhs_map' or 'rhs_map' is present");1054 int64_t expectedNumOperands = getLhsMap() ? getLhsMap()->getNumInputs() : 1;1055 expectedNumOperands += getRhsMap() ? getRhsMap()->getNumInputs() : 1;1056 if (getVarOperands().size() != size_t(expectedNumOperands))1057 return emitOpError("expected ")1058 << expectedNumOperands << " operands, but got "1059 << getVarOperands().size();1060 return success();1061}1062 1063//===----------------------------------------------------------------------===//1064// TestOpInPlaceSelfFold1065//===----------------------------------------------------------------------===//1066 1067OpFoldResult TestOpInPlaceSelfFold::fold(FoldAdaptor adaptor) {1068 if (!getFolded()) {1069 // The folder adds the "folded" if not present.1070 setFolded(true);1071 return getResult();1072 }1073 return {};1074}1075 1076//===----------------------------------------------------------------------===//1077// TestOpFoldWithFoldAdaptor1078//===----------------------------------------------------------------------===//1079 1080OpFoldResult TestOpFoldWithFoldAdaptor::fold(FoldAdaptor adaptor) {1081 int64_t sum = 0;1082 if (auto value = dyn_cast_or_null<IntegerAttr>(adaptor.getOp()))1083 sum += value.getValue().getSExtValue();1084 1085 for (Attribute attr : adaptor.getVariadic())1086 if (auto value = dyn_cast_or_null<IntegerAttr>(attr))1087 sum += 2 * value.getValue().getSExtValue();1088 1089 for (ArrayRef<Attribute> attrs : adaptor.getVarOfVar())1090 for (Attribute attr : attrs)1091 if (auto value = dyn_cast_or_null<IntegerAttr>(attr))1092 sum += 3 * value.getValue().getSExtValue();1093 1094 sum += 4 * std::distance(adaptor.getBody().begin(), adaptor.getBody().end());1095 1096 return IntegerAttr::get(getType(), sum);1097}1098 1099//===----------------------------------------------------------------------===//1100// OpWithInferTypeAdaptorInterfaceOp1101//===----------------------------------------------------------------------===//1102 1103LogicalResult OpWithInferTypeAdaptorInterfaceOp::inferReturnTypes(1104 MLIRContext *, std::optional<Location> location,1105 OpWithInferTypeAdaptorInterfaceOp::Adaptor adaptor,1106 SmallVectorImpl<Type> &inferredReturnTypes) {1107 if (adaptor.getX().getType() != adaptor.getY().getType()) {1108 return emitOptionalError(location, "operand type mismatch ",1109 adaptor.getX().getType(), " vs ",1110 adaptor.getY().getType());1111 }1112 inferredReturnTypes.assign({adaptor.getX().getType()});1113 return success();1114}1115 1116//===----------------------------------------------------------------------===//1117// OpWithRefineTypeInterfaceOp1118//===----------------------------------------------------------------------===//1119 1120// TODO: We should be able to only define either inferReturnType or1121// refineReturnType, currently only refineReturnType can be omitted.1122LogicalResult OpWithRefineTypeInterfaceOp::inferReturnTypes(1123 MLIRContext *context, std::optional<Location> location, ValueRange operands,1124 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,1125 SmallVectorImpl<Type> &returnTypes) {1126 returnTypes.clear();1127 return OpWithRefineTypeInterfaceOp::refineReturnTypes(1128 context, location, operands, attributes, properties, regions,1129 returnTypes);1130}1131 1132LogicalResult OpWithRefineTypeInterfaceOp::refineReturnTypes(1133 MLIRContext *, std::optional<Location> location, ValueRange operands,1134 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,1135 SmallVectorImpl<Type> &returnTypes) {1136 if (operands[0].getType() != operands[1].getType()) {1137 return emitOptionalError(location, "operand type mismatch ",1138 operands[0].getType(), " vs ",1139 operands[1].getType());1140 }1141 // TODO: Add helper to make this more concise to write.1142 if (returnTypes.empty())1143 returnTypes.resize(1, nullptr);1144 if (returnTypes[0] && returnTypes[0] != operands[0].getType())1145 return emitOptionalError(location,1146 "required first operand and result to match");1147 returnTypes[0] = operands[0].getType();1148 return success();1149}1150 1151//===----------------------------------------------------------------------===//1152// TilingNoDpsOp1153//===----------------------------------------------------------------------===//1154 1155SmallVector<Range> TilingNoDpsOp::getIterationDomain(OpBuilder &builder) {1156 return {};1157}1158 1159SmallVector<utils::IteratorType> TilingNoDpsOp::getLoopIteratorTypes() {1160 return {};1161}1162 1163FailureOr<TilingResult>1164TilingNoDpsOp::getTiledImplementation(OpBuilder &builder,1165 ArrayRef<OpFoldResult> offsets,1166 ArrayRef<OpFoldResult> sizes) {1167 return failure();1168}1169 1170LogicalResult TilingNoDpsOp::getResultTilePosition(1171 OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,1172 ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,1173 SmallVector<OpFoldResult> &resultSizes) {1174 return failure();1175}1176 1177//===----------------------------------------------------------------------===//1178// OpWithShapedTypeInferTypeAdaptorInterfaceOp1179//===----------------------------------------------------------------------===//1180 1181LogicalResult1182OpWithShapedTypeInferTypeAdaptorInterfaceOp::inferReturnTypeComponents(1183 MLIRContext *context, std::optional<Location> location,1184 OpWithShapedTypeInferTypeAdaptorInterfaceOp::Adaptor adaptor,1185 SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {1186 // Create return type consisting of the last element of the first operand.1187 auto operandType = adaptor.getOperand1().getType();1188 auto sval = dyn_cast<ShapedType>(operandType);1189 if (!sval)1190 return emitOptionalError(location, "only shaped type operands allowed");1191 int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamic;1192 auto type = IntegerType::get(context, 17);1193 1194 Attribute encoding;1195 if (auto rankedTy = dyn_cast<RankedTensorType>(sval))1196 encoding = rankedTy.getEncoding();1197 inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));1198 return success();1199}1200 1201LogicalResult1202OpWithShapedTypeInferTypeAdaptorInterfaceOp::reifyReturnTypeShapes(1203 OpBuilder &builder, ValueRange operands,1204 llvm::SmallVectorImpl<Value> &shapes) {1205 shapes = SmallVector<Value, 1>{1206 builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};1207 return success();1208}1209 1210//===----------------------------------------------------------------------===//1211// TestOpWithPropertiesAndInferredType1212//===----------------------------------------------------------------------===//1213 1214LogicalResult TestOpWithPropertiesAndInferredType::inferReturnTypes(1215 MLIRContext *context, std::optional<Location>, ValueRange operands,1216 DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,1217 SmallVectorImpl<Type> &inferredReturnTypes) {1218 1219 Adaptor adaptor(operands, attributes, properties, regions);1220 inferredReturnTypes.push_back(IntegerType::get(1221 context, adaptor.getLhs() + adaptor.getProperties().rhs));1222 return success();1223}1224 1225//===----------------------------------------------------------------------===//1226// LoopBlockOp1227//===----------------------------------------------------------------------===//1228 1229void LoopBlockOp::getSuccessorRegions(1230 RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {1231 regions.emplace_back(&getBody(), getBody().getArguments());1232 if (point.isParent())1233 return;1234 1235 regions.emplace_back(getOperation(), getOperation()->getResults());1236}1237 1238OperandRange LoopBlockOp::getEntrySuccessorOperands(RegionSuccessor successor) {1239 assert(successor.getSuccessor() == &getBody());1240 return MutableOperandRange(getInitMutable());1241}1242 1243//===----------------------------------------------------------------------===//1244// LoopBlockTerminatorOp1245//===----------------------------------------------------------------------===//1246 1247MutableOperandRange1248LoopBlockTerminatorOp::getMutableSuccessorOperands(RegionSuccessor successor) {1249 if (successor.isParent())1250 return getExitArgMutable();1251 return getNextIterArgMutable();1252}1253 1254//===----------------------------------------------------------------------===//1255// SwitchWithNoBreakOp1256//===----------------------------------------------------------------------===//1257 1258void TestNoTerminatorOp::getSuccessorRegions(1259 RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {}1260 1261//===----------------------------------------------------------------------===//1262// Test InferIntRangeInterface1263//===----------------------------------------------------------------------===//1264 1265OpFoldResult ManualCppOpWithFold::fold(ArrayRef<Attribute> attributes) {1266 // Just a simple fold for testing purposes that reads an operands constant1267 // value and returns it.1268 if (!attributes.empty())1269 return attributes.front();1270 return nullptr;1271}1272 1273//===----------------------------------------------------------------------===//1274// Tensor/Buffer Ops1275//===----------------------------------------------------------------------===//1276 1277void ReadBufferOp::getEffects(1278 SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>1279 &effects) {1280 // The buffer operand is read.1281 effects.emplace_back(MemoryEffects::Read::get(), &getBufferMutable(),1282 SideEffects::DefaultResource::get());1283 // The buffer contents are dumped.1284 effects.emplace_back(MemoryEffects::Write::get(),1285 SideEffects::DefaultResource::get());1286}1287 1288//===----------------------------------------------------------------------===//1289// Test Dataflow1290//===----------------------------------------------------------------------===//1291 1292//===----------------------------------------------------------------------===//1293// TestCallAndStoreOp1294//===----------------------------------------------------------------------===//1295 1296CallInterfaceCallable TestCallAndStoreOp::getCallableForCallee() {1297 return getCallee();1298}1299 1300void TestCallAndStoreOp::setCalleeFromCallable(CallInterfaceCallable callee) {1301 setCalleeAttr(cast<SymbolRefAttr>(callee));1302}1303 1304Operation::operand_range TestCallAndStoreOp::getArgOperands() {1305 return getCalleeOperands();1306}1307 1308MutableOperandRange TestCallAndStoreOp::getArgOperandsMutable() {1309 return getCalleeOperandsMutable();1310}1311 1312//===----------------------------------------------------------------------===//1313// TestCallOnDeviceOp1314//===----------------------------------------------------------------------===//1315 1316CallInterfaceCallable TestCallOnDeviceOp::getCallableForCallee() {1317 return getCallee();1318}1319 1320void TestCallOnDeviceOp::setCalleeFromCallable(CallInterfaceCallable callee) {1321 setCalleeAttr(cast<SymbolRefAttr>(callee));1322}1323 1324Operation::operand_range TestCallOnDeviceOp::getArgOperands() {1325 return getForwardedOperands();1326}1327 1328MutableOperandRange TestCallOnDeviceOp::getArgOperandsMutable() {1329 return getForwardedOperandsMutable();1330}1331 1332//===----------------------------------------------------------------------===//1333// TestStoreWithARegion1334//===----------------------------------------------------------------------===//1335 1336void TestStoreWithARegion::getSuccessorRegions(1337 RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {1338 if (point.isParent())1339 regions.emplace_back(&getBody(), getBody().front().getArguments());1340 else1341 regions.emplace_back(getOperation(), getOperation()->getResults());1342}1343 1344//===----------------------------------------------------------------------===//1345// TestStoreWithALoopRegion1346//===----------------------------------------------------------------------===//1347 1348void TestStoreWithALoopRegion::getSuccessorRegions(1349 RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> ®ions) {1350 // Both the operation itself and the region may be branching into the body or1351 // back into the operation itself. It is possible for the operation not to1352 // enter the body.1353 regions.emplace_back(1354 RegionSuccessor(&getBody(), getBody().front().getArguments()));1355 regions.emplace_back(getOperation(), getOperation()->getResults());1356}1357 1358//===----------------------------------------------------------------------===//1359// TestVersionedOpA1360//===----------------------------------------------------------------------===//1361 1362LogicalResult1363TestVersionedOpA::readProperties(mlir::DialectBytecodeReader &reader,1364 mlir::OperationState &state) {1365 auto &prop = state.getOrAddProperties<Properties>();1366 if (mlir::failed(reader.readAttribute(prop.dims)))1367 return mlir::failure();1368 1369 // Check if we have a version. If not, assume we are parsing the current1370 // version.1371 auto maybeVersion = reader.getDialectVersion<test::TestDialect>();1372 if (succeeded(maybeVersion)) {1373 // If version is less than 2.0, there is no additional attribute to parse.1374 // We can materialize missing properties post parsing before verification.1375 const auto *version =1376 reinterpret_cast<const TestDialectVersion *>(*maybeVersion);1377 if ((version->major_ < 2)) {1378 return success();1379 }1380 }1381 1382 if (mlir::failed(reader.readAttribute(prop.modifier)))1383 return mlir::failure();1384 return mlir::success();1385}1386 1387void TestVersionedOpA::writeProperties(mlir::DialectBytecodeWriter &writer) {1388 auto &prop = getProperties();1389 writer.writeAttribute(prop.dims);1390 1391 auto maybeVersion = writer.getDialectVersion<test::TestDialect>();1392 if (succeeded(maybeVersion)) {1393 // If version is less than 2.0, there is no additional attribute to write.1394 const auto *version =1395 reinterpret_cast<const TestDialectVersion *>(*maybeVersion);1396 if ((version->major_ < 2)) {1397 llvm::outs() << "downgrading op properties...\n";1398 return;1399 }1400 }1401 writer.writeAttribute(prop.modifier);1402}1403 1404//===----------------------------------------------------------------------===//1405// TestOpWithVersionedProperties1406//===----------------------------------------------------------------------===//1407 1408llvm::LogicalResult TestOpWithVersionedProperties::readFromMlirBytecode(1409 mlir::DialectBytecodeReader &reader, test::VersionedProperties &prop) {1410 uint64_t value1, value2 = 0;1411 if (failed(reader.readVarInt(value1)))1412 return failure();1413 1414 // Check if we have a version. If not, assume we are parsing the current1415 // version.1416 auto maybeVersion = reader.getDialectVersion<test::TestDialect>();1417 bool needToParseAnotherInt = true;1418 if (succeeded(maybeVersion)) {1419 // If version is less than 2.0, there is no additional attribute to parse.1420 // We can materialize missing properties post parsing before verification.1421 const auto *version =1422 reinterpret_cast<const TestDialectVersion *>(*maybeVersion);1423 if ((version->major_ < 2))1424 needToParseAnotherInt = false;1425 }1426 if (needToParseAnotherInt && failed(reader.readVarInt(value2)))1427 return failure();1428 1429 prop.value1 = value1;1430 prop.value2 = value2;1431 return success();1432}1433 1434void TestOpWithVersionedProperties::writeToMlirBytecode(1435 mlir::DialectBytecodeWriter &writer,1436 const test::VersionedProperties &prop) {1437 writer.writeVarInt(prop.value1);1438 writer.writeVarInt(prop.value2);1439}1440 1441//===----------------------------------------------------------------------===//1442// TestMultiSlotAlloca1443//===----------------------------------------------------------------------===//1444 1445llvm::SmallVector<MemorySlot> TestMultiSlotAlloca::getPromotableSlots() {1446 SmallVector<MemorySlot> slots;1447 for (Value result : getResults()) {1448 slots.push_back(MemorySlot{1449 result, cast<MemRefType>(result.getType()).getElementType()});1450 }1451 return slots;1452}1453 1454Value TestMultiSlotAlloca::getDefaultValue(const MemorySlot &slot,1455 OpBuilder &builder) {1456 return TestOpConstant::create(builder, getLoc(), slot.elemType,1457 builder.getI32IntegerAttr(42));1458}1459 1460void TestMultiSlotAlloca::handleBlockArgument(const MemorySlot &slot,1461 BlockArgument argument,1462 OpBuilder &builder) {1463 // Not relevant for testing.1464}1465 1466/// Creates a new TestMultiSlotAlloca operation, just without the `slot`.1467static std::optional<TestMultiSlotAlloca>1468createNewMultiAllocaWithoutSlot(const MemorySlot &slot, OpBuilder &builder,1469 TestMultiSlotAlloca oldOp) {1470 1471 if (oldOp.getNumResults() == 1) {1472 oldOp.erase();1473 return std::nullopt;1474 }1475 1476 SmallVector<Type> newTypes;1477 SmallVector<Value> remainingValues;1478 1479 for (Value oldResult : oldOp.getResults()) {1480 if (oldResult == slot.ptr)1481 continue;1482 remainingValues.push_back(oldResult);1483 newTypes.push_back(oldResult.getType());1484 }1485 1486 OpBuilder::InsertionGuard guard(builder);1487 builder.setInsertionPoint(oldOp);1488 auto replacement =1489 TestMultiSlotAlloca::create(builder, oldOp->getLoc(), newTypes);1490 for (auto [oldResult, newResult] :1491 llvm::zip_equal(remainingValues, replacement.getResults()))1492 oldResult.replaceAllUsesWith(newResult);1493 1494 oldOp.erase();1495 return replacement;1496}1497 1498std::optional<PromotableAllocationOpInterface>1499TestMultiSlotAlloca::handlePromotionComplete(const MemorySlot &slot,1500 Value defaultValue,1501 OpBuilder &builder) {1502 if (defaultValue && defaultValue.use_empty())1503 defaultValue.getDefiningOp()->erase();1504 return createNewMultiAllocaWithoutSlot(slot, builder, *this);1505}1506 1507SmallVector<DestructurableMemorySlot>1508TestMultiSlotAlloca::getDestructurableSlots() {1509 SmallVector<DestructurableMemorySlot> slots;1510 for (Value result : getResults()) {1511 auto memrefType = cast<MemRefType>(result.getType());1512 auto destructurable = dyn_cast<DestructurableTypeInterface>(memrefType);1513 if (!destructurable)1514 continue;1515 1516 std::optional<DenseMap<Attribute, Type>> destructuredType =1517 destructurable.getSubelementIndexMap();1518 if (!destructuredType)1519 continue;1520 slots.emplace_back(1521 DestructurableMemorySlot{{result, memrefType}, *destructuredType});1522 }1523 return slots;1524}1525 1526DenseMap<Attribute, MemorySlot> TestMultiSlotAlloca::destructure(1527 const DestructurableMemorySlot &slot,1528 const SmallPtrSetImpl<Attribute> &usedIndices, OpBuilder &builder,1529 SmallVectorImpl<DestructurableAllocationOpInterface> &newAllocators) {1530 OpBuilder::InsertionGuard guard(builder);1531 builder.setInsertionPointAfter(*this);1532 1533 DenseMap<Attribute, MemorySlot> slotMap;1534 1535 for (Attribute usedIndex : usedIndices) {1536 Type elemType = slot.subelementTypes.lookup(usedIndex);1537 MemRefType elemPtr = MemRefType::get({}, elemType);1538 auto subAlloca = TestMultiSlotAlloca::create(builder, getLoc(), elemPtr);1539 newAllocators.push_back(subAlloca);1540 slotMap.try_emplace<MemorySlot>(usedIndex,1541 {subAlloca.getResult(0), elemType});1542 }1543 1544 return slotMap;1545}1546 1547std::optional<DestructurableAllocationOpInterface>1548TestMultiSlotAlloca::handleDestructuringComplete(1549 const DestructurableMemorySlot &slot, OpBuilder &builder) {1550 return createNewMultiAllocaWithoutSlot(slot, builder, *this);1551}1552 1553namespace {1554/// Returns test dialect's memref layout for test dialect's tensor encoding when1555/// applicable.1556MemRefLayoutAttrInterface1557getMemRefLayoutForTensorEncoding(RankedTensorType tensorType) {1558 if (auto encoding =1559 dyn_cast<test::TestTensorEncodingAttr>(tensorType.getEncoding())) {1560 return cast<MemRefLayoutAttrInterface>(test::TestMemRefLayoutAttr::get(1561 tensorType.getContext(), encoding.getDummy()));1562 }1563 return {};1564}1565 1566/// Auxiliary bufferization function for test and builtin tensors.1567bufferization::BufferLikeType1568convertTensorToBuffer(mlir::Operation *op,1569 const bufferization::BufferizationOptions &options,1570 bufferization::TensorLikeType tensorLike) {1571 auto buffer =1572 *tensorLike.getBufferType(options, [&]() { return op->emitError(); });1573 if (auto memref = dyn_cast<MemRefType>(buffer)) {1574 // Note: For the sake of testing, we want to ensure that encoding -> layout1575 // bufferization happens. This is currently achieved manually.1576 auto layout =1577 getMemRefLayoutForTensorEncoding(cast<RankedTensorType>(tensorLike));1578 return cast<bufferization::BufferLikeType>(1579 MemRefType::get(memref.getShape(), memref.getElementType(), layout,1580 memref.getMemorySpace()));1581 }1582 return buffer;1583}1584} // namespace1585 1586::mlir::LogicalResult test::TestDummyTensorOp::bufferize(1587 ::mlir::RewriterBase &rewriter,1588 const ::mlir::bufferization::BufferizationOptions &options,1589 ::mlir::bufferization::BufferizationState &state) {1590 auto buffer =1591 mlir::bufferization::getBuffer(rewriter, getInput(), options, state);1592 if (mlir::failed(buffer))1593 return failure();1594 1595 const auto outType = getOutput().getType();1596 const auto bufferizedOutType =1597 convertTensorToBuffer(getOperation(), options, outType);1598 // replace op with memref analogy1599 auto dummyMemrefOp = test::TestDummyMemrefOp::create(1600 rewriter, getLoc(), bufferizedOutType, *buffer);1601 1602 mlir::bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),1603 dummyMemrefOp.getResult());1604 1605 return mlir::success();1606}1607 1608::mlir::LogicalResult test::TestCreateTensorOp::bufferize(1609 ::mlir::RewriterBase &rewriter,1610 const ::mlir::bufferization::BufferizationOptions &options,1611 ::mlir::bufferization::BufferizationState &state) {1612 // Note: mlir::bufferization::getBufferType() would internally call1613 // TestCreateTensorOp::getBufferType()1614 const auto bufferizedOutType =1615 mlir::bufferization::getBufferType(getOutput(), options, state);1616 if (mlir::failed(bufferizedOutType))1617 return failure();1618 1619 // replace op with memref analogy1620 auto createMemrefOp =1621 test::TestCreateMemrefOp::create(rewriter, getLoc(), *bufferizedOutType);1622 1623 mlir::bufferization::replaceOpWithBufferizedValues(1624 rewriter, getOperation(), createMemrefOp.getResult());1625 1626 return mlir::success();1627}1628 1629mlir::FailureOr<mlir::bufferization::BufferLikeType>1630test::TestCreateTensorOp::getBufferType(1631 mlir::Value value, const mlir::bufferization::BufferizationOptions &options,1632 const mlir::bufferization::BufferizationState &,1633 llvm::SmallVector<::mlir::Value> &) {1634 const auto type = dyn_cast<bufferization::TensorLikeType>(value.getType());1635 if (type == nullptr)1636 return failure();1637 1638 return convertTensorToBuffer(getOperation(), options, type);1639}1640