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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "TestDialect.h"
10 #include "TestAttributes.h"
11 #include "TestInterfaces.h"
12 #include "TestTypes.h"
13 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
14 #include "mlir/Dialect/DLTI/DLTI.h"
15 #include "mlir/Dialect/Func/IR/FuncOps.h"
16 #include "mlir/Dialect/Tensor/IR/Tensor.h"
17 #include "mlir/IR/BuiltinOps.h"
18 #include "mlir/IR/DialectImplementation.h"
19 #include "mlir/IR/ExtensibleDialect.h"
20 #include "mlir/IR/PatternMatch.h"
21 #include "mlir/IR/TypeUtilities.h"
22 #include "mlir/IR/Verifier.h"
23 #include "mlir/Reducer/ReductionPatternInterface.h"
24 #include "mlir/Transforms/FoldUtils.h"
25 #include "mlir/Transforms/InliningUtils.h"
26 #include "llvm/ADT/StringExtras.h"
27 #include "llvm/ADT/StringSwitch.h"
28 
29 // Include this before the using namespace lines below to
30 // test that we don't have namespace dependencies.
31 #include "TestOpsDialect.cpp.inc"
32 
33 using namespace mlir;
34 using namespace test;
35 
36 void test::registerTestDialect(DialectRegistry &registry) {
37   registry.insert<TestDialect>();
38 }
39 
40 //===----------------------------------------------------------------------===//
41 // TestDialect Interfaces
42 //===----------------------------------------------------------------------===//
43 
44 namespace {
45 
46 /// Testing the correctness of some traits.
47 static_assert(
48     llvm::is_detected<OpTrait::has_implicit_terminator_t,
49                       SingleBlockImplicitTerminatorOp>::value,
50     "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");
51 static_assert(OpTrait::hasSingleBlockImplicitTerminator<
52                   SingleBlockImplicitTerminatorOp>::value,
53               "hasSingleBlockImplicitTerminator does not match "
54               "SingleBlockImplicitTerminatorOp");
55 
56 // Test support for interacting with the AsmPrinter.
57 struct TestOpAsmInterface : public OpAsmDialectInterface {
58   using OpAsmDialectInterface::OpAsmDialectInterface;
59 
60   AliasResult getAlias(Attribute attr, raw_ostream &os) const final {
61     StringAttr strAttr = attr.dyn_cast<StringAttr>();
62     if (!strAttr)
63       return AliasResult::NoAlias;
64 
65     // Check the contents of the string attribute to see what the test alias
66     // should be named.
67     Optional<StringRef> aliasName =
68         StringSwitch<Optional<StringRef>>(strAttr.getValue())
69             .Case("alias_test:dot_in_name", StringRef("test.alias"))
70             .Case("alias_test:trailing_digit", StringRef("test_alias0"))
71             .Case("alias_test:prefixed_digit", StringRef("0_test_alias"))
72             .Case("alias_test:sanitize_conflict_a",
73                   StringRef("test_alias_conflict0"))
74             .Case("alias_test:sanitize_conflict_b",
75                   StringRef("test_alias_conflict0_"))
76             .Case("alias_test:tensor_encoding", StringRef("test_encoding"))
77             .Default(llvm::None);
78     if (!aliasName)
79       return AliasResult::NoAlias;
80 
81     os << *aliasName;
82     return AliasResult::FinalAlias;
83   }
84 
85   AliasResult getAlias(Type type, raw_ostream &os) const final {
86     if (auto tupleType = type.dyn_cast<TupleType>()) {
87       if (tupleType.size() > 0 &&
88           llvm::all_of(tupleType.getTypes(), [](Type elemType) {
89             return elemType.isa<SimpleAType>();
90           })) {
91         os << "test_tuple";
92         return AliasResult::FinalAlias;
93       }
94     }
95     if (auto intType = type.dyn_cast<TestIntegerType>()) {
96       if (intType.getSignedness() ==
97               TestIntegerType::SignednessSemantics::Unsigned &&
98           intType.getWidth() == 8) {
99         os << "test_ui8";
100         return AliasResult::FinalAlias;
101       }
102     }
103     return AliasResult::NoAlias;
104   }
105 };
106 
107 struct TestDialectFoldInterface : public DialectFoldInterface {
108   using DialectFoldInterface::DialectFoldInterface;
109 
110   /// Registered hook to check if the given region, which is attached to an
111   /// operation that is *not* isolated from above, should be used when
112   /// materializing constants.
113   bool shouldMaterializeInto(Region *region) const final {
114     // If this is a one region operation, then insert into it.
115     return isa<OneRegionOp>(region->getParentOp());
116   }
117 };
118 
119 /// This class defines the interface for handling inlining with standard
120 /// operations.
121 struct TestInlinerInterface : public DialectInlinerInterface {
122   using DialectInlinerInterface::DialectInlinerInterface;
123 
124   //===--------------------------------------------------------------------===//
125   // Analysis Hooks
126   //===--------------------------------------------------------------------===//
127 
128   bool isLegalToInline(Operation *call, Operation *callable,
129                        bool wouldBeCloned) const final {
130     // Don't allow inlining calls that are marked `noinline`.
131     return !call->hasAttr("noinline");
132   }
133   bool isLegalToInline(Region *, Region *, bool,
134                        BlockAndValueMapping &) const final {
135     // Inlining into test dialect regions is legal.
136     return true;
137   }
138   bool isLegalToInline(Operation *, Region *, bool,
139                        BlockAndValueMapping &) const final {
140     return true;
141   }
142 
143   bool shouldAnalyzeRecursively(Operation *op) const final {
144     // Analyze recursively if this is not a functional region operation, it
145     // froms a separate functional scope.
146     return !isa<FunctionalRegionOp>(op);
147   }
148 
149   //===--------------------------------------------------------------------===//
150   // Transformation Hooks
151   //===--------------------------------------------------------------------===//
152 
153   /// Handle the given inlined terminator by replacing it with a new operation
154   /// as necessary.
155   void handleTerminator(Operation *op,
156                         ArrayRef<Value> valuesToRepl) const final {
157     // Only handle "test.return" here.
158     auto returnOp = dyn_cast<TestReturnOp>(op);
159     if (!returnOp)
160       return;
161 
162     // Replace the values directly with the return operands.
163     assert(returnOp.getNumOperands() == valuesToRepl.size());
164     for (const auto &it : llvm::enumerate(returnOp.getOperands()))
165       valuesToRepl[it.index()].replaceAllUsesWith(it.value());
166   }
167 
168   /// Attempt to materialize a conversion for a type mismatch between a call
169   /// from this dialect, and a callable region. This method should generate an
170   /// operation that takes 'input' as the only operand, and produces a single
171   /// result of 'resultType'. If a conversion can not be generated, nullptr
172   /// should be returned.
173   Operation *materializeCallConversion(OpBuilder &builder, Value input,
174                                        Type resultType,
175                                        Location conversionLoc) const final {
176     // Only allow conversion for i16/i32 types.
177     if (!(resultType.isSignlessInteger(16) ||
178           resultType.isSignlessInteger(32)) ||
179         !(input.getType().isSignlessInteger(16) ||
180           input.getType().isSignlessInteger(32)))
181       return nullptr;
182     return builder.create<TestCastOp>(conversionLoc, resultType, input);
183   }
184 
185   void processInlinedCallBlocks(
186       Operation *call,
187       iterator_range<Region::iterator> inlinedBlocks) const final {
188     if (!isa<ConversionCallOp>(call))
189       return;
190 
191     // Set attributed on all ops in the inlined blocks.
192     for (Block &block : inlinedBlocks) {
193       block.walk([&](Operation *op) {
194         op->setAttr("inlined_conversion", UnitAttr::get(call->getContext()));
195       });
196     }
197   }
198 };
199 
200 struct TestReductionPatternInterface : public DialectReductionPatternInterface {
201 public:
202   TestReductionPatternInterface(Dialect *dialect)
203       : DialectReductionPatternInterface(dialect) {}
204 
205   void populateReductionPatterns(RewritePatternSet &patterns) const final {
206     populateTestReductionPatterns(patterns);
207   }
208 };
209 
210 } // namespace
211 
212 //===----------------------------------------------------------------------===//
213 // Dynamic operations
214 //===----------------------------------------------------------------------===//
215 
216 std::unique_ptr<DynamicOpDefinition> getDynamicGenericOp(TestDialect *dialect) {
217   return DynamicOpDefinition::get(
218       "dynamic_generic", dialect, [](Operation *op) { return success(); },
219       [](Operation *op) { return success(); });
220 }
221 
222 std::unique_ptr<DynamicOpDefinition>
223 getDynamicOneOperandTwoResultsOp(TestDialect *dialect) {
224   return DynamicOpDefinition::get(
225       "dynamic_one_operand_two_results", dialect,
226       [](Operation *op) {
227         if (op->getNumOperands() != 1) {
228           op->emitOpError()
229               << "expected 1 operand, but had " << op->getNumOperands();
230           return failure();
231         }
232         if (op->getNumResults() != 2) {
233           op->emitOpError()
234               << "expected 2 results, but had " << op->getNumResults();
235           return failure();
236         }
237         return success();
238       },
239       [](Operation *op) { return success(); });
240 }
241 
242 std::unique_ptr<DynamicOpDefinition>
243 getDynamicCustomParserPrinterOp(TestDialect *dialect) {
244   auto verifier = [](Operation *op) {
245     if (op->getNumOperands() == 0 && op->getNumResults() == 0)
246       return success();
247     op->emitError() << "operation should have no operands and no results";
248     return failure();
249   };
250   auto regionVerifier = [](Operation *op) { return success(); };
251 
252   auto parser = [](OpAsmParser &parser, OperationState &state) {
253     return parser.parseKeyword("custom_keyword");
254   };
255 
256   auto printer = [](Operation *op, OpAsmPrinter &printer, llvm::StringRef) {
257     printer << op->getName() << " custom_keyword";
258   };
259 
260   return DynamicOpDefinition::get("dynamic_custom_parser_printer", dialect,
261                                   verifier, regionVerifier, parser, printer);
262 }
263 
264 //===----------------------------------------------------------------------===//
265 // TestDialect
266 //===----------------------------------------------------------------------===//
267 
268 static void testSideEffectOpGetEffect(
269     Operation *op,
270     SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>> &effects);
271 
272 // This is the implementation of a dialect fallback for `TestEffectOpInterface`.
273 struct TestOpEffectInterfaceFallback
274     : public TestEffectOpInterface::FallbackModel<
275           TestOpEffectInterfaceFallback> {
276   static bool classof(Operation *op) {
277     bool isSupportedOp =
278         op->getName().getStringRef() == "test.unregistered_side_effect_op";
279     assert(isSupportedOp && "Unexpected dispatch");
280     return isSupportedOp;
281   }
282 
283   void
284   getEffects(Operation *op,
285              SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>>
286                  &effects) const {
287     testSideEffectOpGetEffect(op, effects);
288   }
289 };
290 
291 void TestDialect::initialize() {
292   registerAttributes();
293   registerTypes();
294   addOperations<
295 #define GET_OP_LIST
296 #include "TestOps.cpp.inc"
297       >();
298   registerDynamicOp(getDynamicGenericOp(this));
299   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
300   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
301 
302   addInterfaces<TestOpAsmInterface, TestDialectFoldInterface,
303                 TestInlinerInterface, TestReductionPatternInterface>();
304   allowUnknownOperations();
305 
306   // Instantiate our fallback op interface that we'll use on specific
307   // unregistered op.
308   fallbackEffectOpInterfaces = new TestOpEffectInterfaceFallback;
309 }
310 TestDialect::~TestDialect() {
311   delete static_cast<TestOpEffectInterfaceFallback *>(
312       fallbackEffectOpInterfaces);
313 }
314 
315 Operation *TestDialect::materializeConstant(OpBuilder &builder, Attribute value,
316                                             Type type, Location loc) {
317   return builder.create<TestOpConstant>(loc, type, value);
318 }
319 
320 ::mlir::LogicalResult FormatInferType2Op::inferReturnTypes(
321     ::mlir::MLIRContext *context, ::llvm::Optional<::mlir::Location> location,
322     ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
323     ::mlir::RegionRange regions,
324     ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
325   inferredReturnTypes.assign({::mlir::IntegerType::get(context, 16)});
326   return ::mlir::success();
327 }
328 
329 void *TestDialect::getRegisteredInterfaceForOp(TypeID typeID,
330                                                OperationName opName) {
331   if (opName.getIdentifier() == "test.unregistered_side_effect_op" &&
332       typeID == TypeID::get<TestEffectOpInterface>())
333     return fallbackEffectOpInterfaces;
334   return nullptr;
335 }
336 
337 LogicalResult TestDialect::verifyOperationAttribute(Operation *op,
338                                                     NamedAttribute namedAttr) {
339   if (namedAttr.getName() == "test.invalid_attr")
340     return op->emitError() << "invalid to use 'test.invalid_attr'";
341   return success();
342 }
343 
344 LogicalResult TestDialect::verifyRegionArgAttribute(Operation *op,
345                                                     unsigned regionIndex,
346                                                     unsigned argIndex,
347                                                     NamedAttribute namedAttr) {
348   if (namedAttr.getName() == "test.invalid_attr")
349     return op->emitError() << "invalid to use 'test.invalid_attr'";
350   return success();
351 }
352 
353 LogicalResult
354 TestDialect::verifyRegionResultAttribute(Operation *op, unsigned regionIndex,
355                                          unsigned resultIndex,
356                                          NamedAttribute namedAttr) {
357   if (namedAttr.getName() == "test.invalid_attr")
358     return op->emitError() << "invalid to use 'test.invalid_attr'";
359   return success();
360 }
361 
362 Optional<Dialect::ParseOpHook>
363 TestDialect::getParseOperationHook(StringRef opName) const {
364   if (opName == "test.dialect_custom_printer") {
365     return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
366       return parser.parseKeyword("custom_format");
367     }};
368   }
369   if (opName == "test.dialect_custom_format_fallback") {
370     return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
371       return parser.parseKeyword("custom_format_fallback");
372     }};
373   }
374   return None;
375 }
376 
377 llvm::unique_function<void(Operation *, OpAsmPrinter &)>
378 TestDialect::getOperationPrinter(Operation *op) const {
379   StringRef opName = op->getName().getStringRef();
380   if (opName == "test.dialect_custom_printer") {
381     return [](Operation *op, OpAsmPrinter &printer) {
382       printer.getStream() << " custom_format";
383     };
384   }
385   if (opName == "test.dialect_custom_format_fallback") {
386     return [](Operation *op, OpAsmPrinter &printer) {
387       printer.getStream() << " custom_format_fallback";
388     };
389   }
390   return {};
391 }
392 
393 //===----------------------------------------------------------------------===//
394 // TestBranchOp
395 //===----------------------------------------------------------------------===//
396 
397 SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
398   assert(index == 0 && "invalid successor index");
399   return SuccessorOperands(getTargetOperandsMutable());
400 }
401 
402 //===----------------------------------------------------------------------===//
403 // TestProducingBranchOp
404 //===----------------------------------------------------------------------===//
405 
406 SuccessorOperands TestProducingBranchOp::getSuccessorOperands(unsigned index) {
407   assert(index <= 1 && "invalid successor index");
408   if (index == 1)
409     return SuccessorOperands(getFirstOperandsMutable());
410   return SuccessorOperands(getSecondOperandsMutable());
411 }
412 
413 //===----------------------------------------------------------------------===//
414 // TestProducingBranchOp
415 //===----------------------------------------------------------------------===//
416 
417 SuccessorOperands TestInternalBranchOp::getSuccessorOperands(unsigned index) {
418   assert(index <= 1 && "invalid successor index");
419   if (index == 0)
420     return SuccessorOperands(0, getSuccessOperandsMutable());
421   return SuccessorOperands(1, getErrorOperandsMutable());
422 }
423 
424 //===----------------------------------------------------------------------===//
425 // TestDialectCanonicalizerOp
426 //===----------------------------------------------------------------------===//
427 
428 static LogicalResult
429 dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,
430                                PatternRewriter &rewriter) {
431   rewriter.replaceOpWithNewOp<arith::ConstantOp>(
432       op, rewriter.getI32IntegerAttr(42));
433   return success();
434 }
435 
436 void TestDialect::getCanonicalizationPatterns(
437     RewritePatternSet &results) const {
438   results.add(&dialectCanonicalizationPattern);
439 }
440 
441 //===----------------------------------------------------------------------===//
442 // TestCallOp
443 //===----------------------------------------------------------------------===//
444 
445 LogicalResult TestCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
446   // Check that the callee attribute was specified.
447   auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
448   if (!fnAttr)
449     return emitOpError("requires a 'callee' symbol reference attribute");
450   if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))
451     return emitOpError() << "'" << fnAttr.getValue()
452                          << "' does not reference a valid function";
453   return success();
454 }
455 
456 //===----------------------------------------------------------------------===//
457 // TestFoldToCallOp
458 //===----------------------------------------------------------------------===//
459 
460 namespace {
461 struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
462   using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
463 
464   LogicalResult matchAndRewrite(FoldToCallOp op,
465                                 PatternRewriter &rewriter) const override {
466     rewriter.replaceOpWithNewOp<func::CallOp>(op, TypeRange(),
467                                               op.getCalleeAttr(), ValueRange());
468     return success();
469   }
470 };
471 } // namespace
472 
473 void FoldToCallOp::getCanonicalizationPatterns(RewritePatternSet &results,
474                                                MLIRContext *context) {
475   results.add<FoldToCallOpPattern>(context);
476 }
477 
478 //===----------------------------------------------------------------------===//
479 // Test Format* operations
480 //===----------------------------------------------------------------------===//
481 
482 //===----------------------------------------------------------------------===//
483 // Parsing
484 
485 static ParseResult parseCustomOptionalOperand(
486     OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &optOperand) {
487   if (succeeded(parser.parseOptionalLParen())) {
488     optOperand.emplace();
489     if (parser.parseOperand(*optOperand) || parser.parseRParen())
490       return failure();
491   }
492   return success();
493 }
494 
495 static ParseResult parseCustomDirectiveOperands(
496     OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,
497     Optional<OpAsmParser::UnresolvedOperand> &optOperand,
498     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands) {
499   if (parser.parseOperand(operand))
500     return failure();
501   if (succeeded(parser.parseOptionalComma())) {
502     optOperand.emplace();
503     if (parser.parseOperand(*optOperand))
504       return failure();
505   }
506   if (parser.parseArrow() || parser.parseLParen() ||
507       parser.parseOperandList(varOperands) || parser.parseRParen())
508     return failure();
509   return success();
510 }
511 static ParseResult
512 parseCustomDirectiveResults(OpAsmParser &parser, Type &operandType,
513                             Type &optOperandType,
514                             SmallVectorImpl<Type> &varOperandTypes) {
515   if (parser.parseColon())
516     return failure();
517 
518   if (parser.parseType(operandType))
519     return failure();
520   if (succeeded(parser.parseOptionalComma())) {
521     if (parser.parseType(optOperandType))
522       return failure();
523   }
524   if (parser.parseArrow() || parser.parseLParen() ||
525       parser.parseTypeList(varOperandTypes) || parser.parseRParen())
526     return failure();
527   return success();
528 }
529 static ParseResult
530 parseCustomDirectiveWithTypeRefs(OpAsmParser &parser, Type operandType,
531                                  Type optOperandType,
532                                  const SmallVectorImpl<Type> &varOperandTypes) {
533   if (parser.parseKeyword("type_refs_capture"))
534     return failure();
535 
536   Type operandType2, optOperandType2;
537   SmallVector<Type, 1> varOperandTypes2;
538   if (parseCustomDirectiveResults(parser, operandType2, optOperandType2,
539                                   varOperandTypes2))
540     return failure();
541 
542   if (operandType != operandType2 || optOperandType != optOperandType2 ||
543       varOperandTypes != varOperandTypes2)
544     return failure();
545 
546   return success();
547 }
548 static ParseResult parseCustomDirectiveOperandsAndTypes(
549     OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,
550     Optional<OpAsmParser::UnresolvedOperand> &optOperand,
551     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands,
552     Type &operandType, Type &optOperandType,
553     SmallVectorImpl<Type> &varOperandTypes) {
554   if (parseCustomDirectiveOperands(parser, operand, optOperand, varOperands) ||
555       parseCustomDirectiveResults(parser, operandType, optOperandType,
556                                   varOperandTypes))
557     return failure();
558   return success();
559 }
560 static ParseResult parseCustomDirectiveRegions(
561     OpAsmParser &parser, Region &region,
562     SmallVectorImpl<std::unique_ptr<Region>> &varRegions) {
563   if (parser.parseRegion(region))
564     return failure();
565   if (failed(parser.parseOptionalComma()))
566     return success();
567   std::unique_ptr<Region> varRegion = std::make_unique<Region>();
568   if (parser.parseRegion(*varRegion))
569     return failure();
570   varRegions.emplace_back(std::move(varRegion));
571   return success();
572 }
573 static ParseResult
574 parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,
575                                SmallVectorImpl<Block *> &varSuccessors) {
576   if (parser.parseSuccessor(successor))
577     return failure();
578   if (failed(parser.parseOptionalComma()))
579     return success();
580   Block *varSuccessor;
581   if (parser.parseSuccessor(varSuccessor))
582     return failure();
583   varSuccessors.append(2, varSuccessor);
584   return success();
585 }
586 static ParseResult parseCustomDirectiveAttributes(OpAsmParser &parser,
587                                                   IntegerAttr &attr,
588                                                   IntegerAttr &optAttr) {
589   if (parser.parseAttribute(attr))
590     return failure();
591   if (succeeded(parser.parseOptionalComma())) {
592     if (parser.parseAttribute(optAttr))
593       return failure();
594   }
595   return success();
596 }
597 
598 static ParseResult parseCustomDirectiveAttrDict(OpAsmParser &parser,
599                                                 NamedAttrList &attrs) {
600   return parser.parseOptionalAttrDict(attrs);
601 }
602 static ParseResult parseCustomDirectiveOptionalOperandRef(
603     OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &optOperand) {
604   int64_t operandCount = 0;
605   if (parser.parseInteger(operandCount))
606     return failure();
607   bool expectedOptionalOperand = operandCount == 0;
608   return success(expectedOptionalOperand != optOperand.hasValue());
609 }
610 
611 //===----------------------------------------------------------------------===//
612 // Printing
613 
614 static void printCustomOptionalOperand(OpAsmPrinter &printer, Operation *,
615                                        Value optOperand) {
616   if (optOperand)
617     printer << "(" << optOperand << ") ";
618 }
619 
620 static void printCustomDirectiveOperands(OpAsmPrinter &printer, Operation *,
621                                          Value operand, Value optOperand,
622                                          OperandRange varOperands) {
623   printer << operand;
624   if (optOperand)
625     printer << ", " << optOperand;
626   printer << " -> (" << varOperands << ")";
627 }
628 static void printCustomDirectiveResults(OpAsmPrinter &printer, Operation *,
629                                         Type operandType, Type optOperandType,
630                                         TypeRange varOperandTypes) {
631   printer << " : " << operandType;
632   if (optOperandType)
633     printer << ", " << optOperandType;
634   printer << " -> (" << varOperandTypes << ")";
635 }
636 static void printCustomDirectiveWithTypeRefs(OpAsmPrinter &printer,
637                                              Operation *op, Type operandType,
638                                              Type optOperandType,
639                                              TypeRange varOperandTypes) {
640   printer << " type_refs_capture ";
641   printCustomDirectiveResults(printer, op, operandType, optOperandType,
642                               varOperandTypes);
643 }
644 static void printCustomDirectiveOperandsAndTypes(
645     OpAsmPrinter &printer, Operation *op, Value operand, Value optOperand,
646     OperandRange varOperands, Type operandType, Type optOperandType,
647     TypeRange varOperandTypes) {
648   printCustomDirectiveOperands(printer, op, operand, optOperand, varOperands);
649   printCustomDirectiveResults(printer, op, operandType, optOperandType,
650                               varOperandTypes);
651 }
652 static void printCustomDirectiveRegions(OpAsmPrinter &printer, Operation *,
653                                         Region &region,
654                                         MutableArrayRef<Region> varRegions) {
655   printer.printRegion(region);
656   if (!varRegions.empty()) {
657     printer << ", ";
658     for (Region &region : varRegions)
659       printer.printRegion(region);
660   }
661 }
662 static void printCustomDirectiveSuccessors(OpAsmPrinter &printer, Operation *,
663                                            Block *successor,
664                                            SuccessorRange varSuccessors) {
665   printer << successor;
666   if (!varSuccessors.empty())
667     printer << ", " << varSuccessors.front();
668 }
669 static void printCustomDirectiveAttributes(OpAsmPrinter &printer, Operation *,
670                                            Attribute attribute,
671                                            Attribute optAttribute) {
672   printer << attribute;
673   if (optAttribute)
674     printer << ", " << optAttribute;
675 }
676 
677 static void printCustomDirectiveAttrDict(OpAsmPrinter &printer, Operation *op,
678                                          DictionaryAttr attrs) {
679   printer.printOptionalAttrDict(attrs.getValue());
680 }
681 
682 static void printCustomDirectiveOptionalOperandRef(OpAsmPrinter &printer,
683                                                    Operation *op,
684                                                    Value optOperand) {
685   printer << (optOperand ? "1" : "0");
686 }
687 
688 //===----------------------------------------------------------------------===//
689 // Test IsolatedRegionOp - parse passthrough region arguments.
690 //===----------------------------------------------------------------------===//
691 
692 ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
693                                     OperationState &result) {
694   // Parse the input operand.
695   OpAsmParser::Argument argInfo;
696   argInfo.type = parser.getBuilder().getIndexType();
697   if (parser.parseOperand(argInfo.ssaName) ||
698       parser.resolveOperand(argInfo.ssaName, argInfo.type, result.operands))
699     return failure();
700 
701   // Parse the body region, and reuse the operand info as the argument info.
702   Region *body = result.addRegion();
703   return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);
704 }
705 
706 void IsolatedRegionOp::print(OpAsmPrinter &p) {
707   p << "test.isolated_region ";
708   p.printOperand(getOperand());
709   p.shadowRegionArgs(getRegion(), getOperand());
710   p << ' ';
711   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
712 }
713 
714 //===----------------------------------------------------------------------===//
715 // Test SSACFGRegionOp
716 //===----------------------------------------------------------------------===//
717 
718 RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {
719   return RegionKind::SSACFG;
720 }
721 
722 //===----------------------------------------------------------------------===//
723 // Test GraphRegionOp
724 //===----------------------------------------------------------------------===//
725 
726 RegionKind GraphRegionOp::getRegionKind(unsigned index) {
727   return RegionKind::Graph;
728 }
729 
730 //===----------------------------------------------------------------------===//
731 // Test AffineScopeOp
732 //===----------------------------------------------------------------------===//
733 
734 ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
735   // Parse the body region, and reuse the operand info as the argument info.
736   Region *body = result.addRegion();
737   return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
738 }
739 
740 void AffineScopeOp::print(OpAsmPrinter &p) {
741   p << "test.affine_scope ";
742   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
743 }
744 
745 //===----------------------------------------------------------------------===//
746 // Test parser.
747 //===----------------------------------------------------------------------===//
748 
749 ParseResult ParseIntegerLiteralOp::parse(OpAsmParser &parser,
750                                          OperationState &result) {
751   if (parser.parseOptionalColon())
752     return success();
753   uint64_t numResults;
754   if (parser.parseInteger(numResults))
755     return failure();
756 
757   IndexType type = parser.getBuilder().getIndexType();
758   for (unsigned i = 0; i < numResults; ++i)
759     result.addTypes(type);
760   return success();
761 }
762 
763 void ParseIntegerLiteralOp::print(OpAsmPrinter &p) {
764   if (unsigned numResults = getNumResults())
765     p << " : " << numResults;
766 }
767 
768 ParseResult ParseWrappedKeywordOp::parse(OpAsmParser &parser,
769                                          OperationState &result) {
770   StringRef keyword;
771   if (parser.parseKeyword(&keyword))
772     return failure();
773   result.addAttribute("keyword", parser.getBuilder().getStringAttr(keyword));
774   return success();
775 }
776 
777 void ParseWrappedKeywordOp::print(OpAsmPrinter &p) { p << " " << getKeyword(); }
778 
779 //===----------------------------------------------------------------------===//
780 // Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
781 
782 ParseResult WrappingRegionOp::parse(OpAsmParser &parser,
783                                     OperationState &result) {
784   if (parser.parseKeyword("wraps"))
785     return failure();
786 
787   // Parse the wrapped op in a region
788   Region &body = *result.addRegion();
789   body.push_back(new Block);
790   Block &block = body.back();
791   Operation *wrappedOp = parser.parseGenericOperation(&block, block.begin());
792   if (!wrappedOp)
793     return failure();
794 
795   // Create a return terminator in the inner region, pass as operand to the
796   // terminator the returned values from the wrapped operation.
797   SmallVector<Value, 8> returnOperands(wrappedOp->getResults());
798   OpBuilder builder(parser.getContext());
799   builder.setInsertionPointToEnd(&block);
800   builder.create<TestReturnOp>(wrappedOp->getLoc(), returnOperands);
801 
802   // Get the results type for the wrapping op from the terminator operands.
803   Operation &returnOp = body.back().back();
804   result.types.append(returnOp.operand_type_begin(),
805                       returnOp.operand_type_end());
806 
807   // Use the location of the wrapped op for the "test.wrapping_region" op.
808   result.location = wrappedOp->getLoc();
809 
810   return success();
811 }
812 
813 void WrappingRegionOp::print(OpAsmPrinter &p) {
814   p << " wraps ";
815   p.printGenericOp(&getRegion().front().front());
816 }
817 
818 //===----------------------------------------------------------------------===//
819 // Test PrettyPrintedRegionOp -  exercising the following parser APIs
820 //   parseGenericOperationAfterOpName
821 //   parseCustomOperationName
822 //===----------------------------------------------------------------------===//
823 
824 ParseResult PrettyPrintedRegionOp::parse(OpAsmParser &parser,
825                                          OperationState &result) {
826 
827   SMLoc loc = parser.getCurrentLocation();
828   Location currLocation = parser.getEncodedSourceLoc(loc);
829 
830   // Parse the operands.
831   SmallVector<OpAsmParser::UnresolvedOperand, 2> operands;
832   if (parser.parseOperandList(operands))
833     return failure();
834 
835   // Check if we are parsing the pretty-printed version
836   //  test.pretty_printed_region start <inner-op> end : <functional-type>
837   // Else fallback to parsing the "non pretty-printed" version.
838   if (!succeeded(parser.parseOptionalKeyword("start")))
839     return parser.parseGenericOperationAfterOpName(
840         result, llvm::makeArrayRef(operands));
841 
842   FailureOr<OperationName> parseOpNameInfo = parser.parseCustomOperationName();
843   if (failed(parseOpNameInfo))
844     return failure();
845 
846   StringAttr innerOpName = parseOpNameInfo->getIdentifier();
847 
848   FunctionType opFntype;
849   Optional<Location> explicitLoc;
850   if (parser.parseKeyword("end") || parser.parseColon() ||
851       parser.parseType(opFntype) ||
852       parser.parseOptionalLocationSpecifier(explicitLoc))
853     return failure();
854 
855   // If location of the op is explicitly provided, then use it; Else use
856   // the parser's current location.
857   Location opLoc = explicitLoc.getValueOr(currLocation);
858 
859   // Derive the SSA-values for op's operands.
860   if (parser.resolveOperands(operands, opFntype.getInputs(), loc,
861                              result.operands))
862     return failure();
863 
864   // Add a region for op.
865   Region &region = *result.addRegion();
866 
867   // Create a basic-block inside op's region.
868   Block &block = region.emplaceBlock();
869 
870   // Create and insert an "inner-op" operation in the block.
871   // Just for testing purposes, we can assume that inner op is a binary op with
872   // result and operand types all same as the test-op's first operand.
873   Type innerOpType = opFntype.getInput(0);
874   Value lhs = block.addArgument(innerOpType, opLoc);
875   Value rhs = block.addArgument(innerOpType, opLoc);
876 
877   OpBuilder builder(parser.getBuilder().getContext());
878   builder.setInsertionPointToStart(&block);
879 
880   Operation *innerOp =
881       builder.create(opLoc, innerOpName, /*operands=*/{lhs, rhs}, innerOpType);
882 
883   // Insert a return statement in the block returning the inner-op's result.
884   builder.create<TestReturnOp>(innerOp->getLoc(), innerOp->getResults());
885 
886   // Populate the op operation-state with result-type and location.
887   result.addTypes(opFntype.getResults());
888   result.location = innerOp->getLoc();
889 
890   return success();
891 }
892 
893 void PrettyPrintedRegionOp::print(OpAsmPrinter &p) {
894   p << ' ';
895   p.printOperands(getOperands());
896 
897   Operation &innerOp = getRegion().front().front();
898   // Assuming that region has a single non-terminator inner-op, if the inner-op
899   // meets some criteria (which in this case is a simple one  based on the name
900   // of inner-op), then we can print the entire region in a succinct way.
901   // Here we assume that the prototype of "special.op" can be trivially derived
902   // while parsing it back.
903   if (innerOp.getName().getStringRef().equals("special.op")) {
904     p << " start special.op end";
905   } else {
906     p << " (";
907     p.printRegion(getRegion());
908     p << ")";
909   }
910 
911   p << " : ";
912   p.printFunctionalType(*this);
913 }
914 
915 //===----------------------------------------------------------------------===//
916 // Test PolyForOp - parse list of region arguments.
917 //===----------------------------------------------------------------------===//
918 
919 ParseResult PolyForOp::parse(OpAsmParser &parser, OperationState &result) {
920   SmallVector<OpAsmParser::Argument, 4> ivsInfo;
921   // Parse list of region arguments without a delimiter.
922   if (parser.parseArgumentList(ivsInfo, OpAsmParser::Delimiter::None))
923     return failure();
924 
925   // Parse the body region.
926   Region *body = result.addRegion();
927   for (auto &iv : ivsInfo)
928     iv.type = parser.getBuilder().getIndexType();
929   return parser.parseRegion(*body, ivsInfo);
930 }
931 
932 void PolyForOp::print(OpAsmPrinter &p) { p.printGenericOp(*this); }
933 
934 void PolyForOp::getAsmBlockArgumentNames(Region &region,
935                                          OpAsmSetValueNameFn setNameFn) {
936   auto arrayAttr = getOperation()->getAttrOfType<ArrayAttr>("arg_names");
937   if (!arrayAttr)
938     return;
939   auto args = getRegion().front().getArguments();
940   auto e = std::min(arrayAttr.size(), args.size());
941   for (unsigned i = 0; i < e; ++i) {
942     if (auto strAttr = arrayAttr[i].dyn_cast<StringAttr>())
943       setNameFn(args[i], strAttr.getValue());
944   }
945 }
946 
947 //===----------------------------------------------------------------------===//
948 // Test removing op with inner ops.
949 //===----------------------------------------------------------------------===//
950 
951 namespace {
952 struct TestRemoveOpWithInnerOps
953     : public OpRewritePattern<TestOpWithRegionPattern> {
954   using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;
955 
956   void initialize() { setDebugName("TestRemoveOpWithInnerOps"); }
957 
958   LogicalResult matchAndRewrite(TestOpWithRegionPattern op,
959                                 PatternRewriter &rewriter) const override {
960     rewriter.eraseOp(op);
961     return success();
962   }
963 };
964 } // namespace
965 
966 void TestOpWithRegionPattern::getCanonicalizationPatterns(
967     RewritePatternSet &results, MLIRContext *context) {
968   results.add<TestRemoveOpWithInnerOps>(context);
969 }
970 
971 OpFoldResult TestOpWithRegionFold::fold(ArrayRef<Attribute> operands) {
972   return getOperand();
973 }
974 
975 OpFoldResult TestOpConstant::fold(ArrayRef<Attribute> operands) {
976   return getValue();
977 }
978 
979 LogicalResult TestOpWithVariadicResultsAndFolder::fold(
980     ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
981   for (Value input : this->getOperands()) {
982     results.push_back(input);
983   }
984   return success();
985 }
986 
987 OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
988   assert(operands.size() == 1);
989   if (operands.front()) {
990     (*this)->setAttr("attr", operands.front());
991     return getResult();
992   }
993   return {};
994 }
995 
996 OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
997   return getOperand();
998 }
999 
1000 LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
1001     MLIRContext *, Optional<Location> location, ValueRange operands,
1002     DictionaryAttr attributes, RegionRange regions,
1003     SmallVectorImpl<Type> &inferredReturnTypes) {
1004   if (operands[0].getType() != operands[1].getType()) {
1005     return emitOptionalError(location, "operand type mismatch ",
1006                              operands[0].getType(), " vs ",
1007                              operands[1].getType());
1008   }
1009   inferredReturnTypes.assign({operands[0].getType()});
1010   return success();
1011 }
1012 
1013 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
1014     MLIRContext *context, Optional<Location> location, ValueShapeRange operands,
1015     DictionaryAttr attributes, RegionRange regions,
1016     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1017   // Create return type consisting of the last element of the first operand.
1018   auto operandType = operands.front().getType();
1019   auto sval = operandType.dyn_cast<ShapedType>();
1020   if (!sval) {
1021     return emitOptionalError(location, "only shaped type operands allowed");
1022   }
1023   int64_t dim =
1024       sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;
1025   auto type = IntegerType::get(context, 17);
1026   inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type));
1027   return success();
1028 }
1029 
1030 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
1031     OpBuilder &builder, ValueRange operands,
1032     llvm::SmallVectorImpl<Value> &shapes) {
1033   shapes = SmallVector<Value, 1>{
1034       builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
1035   return success();
1036 }
1037 
1038 LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
1039     OpBuilder &builder, ValueRange operands,
1040     llvm::SmallVectorImpl<Value> &shapes) {
1041   Location loc = getLoc();
1042   shapes.reserve(operands.size());
1043   for (Value operand : llvm::reverse(operands)) {
1044     auto rank = operand.getType().cast<RankedTensorType>().getRank();
1045     auto currShape = llvm::to_vector<4>(
1046         llvm::map_range(llvm::seq<int64_t>(0, rank), [&](int64_t dim) -> Value {
1047           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
1048         }));
1049     shapes.push_back(builder.create<tensor::FromElementsOp>(
1050         getLoc(), RankedTensorType::get({rank}, builder.getIndexType()),
1051         currShape));
1052   }
1053   return success();
1054 }
1055 
1056 LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
1057     OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
1058   Location loc = getLoc();
1059   shapes.reserve(getNumOperands());
1060   for (Value operand : llvm::reverse(getOperands())) {
1061     auto currShape = llvm::to_vector<4>(llvm::map_range(
1062         llvm::seq<int64_t>(
1063             0, operand.getType().cast<RankedTensorType>().getRank()),
1064         [&](int64_t dim) -> Value {
1065           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
1066         }));
1067     shapes.emplace_back(std::move(currShape));
1068   }
1069   return success();
1070 }
1071 
1072 //===----------------------------------------------------------------------===//
1073 // Test SideEffect interfaces
1074 //===----------------------------------------------------------------------===//
1075 
1076 namespace {
1077 /// A test resource for side effects.
1078 struct TestResource : public SideEffects::Resource::Base<TestResource> {
1079   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestResource)
1080 
1081   StringRef getName() final { return "<Test>"; }
1082 };
1083 } // namespace
1084 
1085 static void testSideEffectOpGetEffect(
1086     Operation *op,
1087     SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>>
1088         &effects) {
1089   auto effectsAttr = op->getAttrOfType<AffineMapAttr>("effect_parameter");
1090   if (!effectsAttr)
1091     return;
1092 
1093   effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);
1094 }
1095 
1096 void SideEffectOp::getEffects(
1097     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
1098   // Check for an effects attribute on the op instance.
1099   ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");
1100   if (!effectsAttr)
1101     return;
1102 
1103   // If there is one, it is an array of dictionary attributes that hold
1104   // information on the effects of this operation.
1105   for (Attribute element : effectsAttr) {
1106     DictionaryAttr effectElement = element.cast<DictionaryAttr>();
1107 
1108     // Get the specific memory effect.
1109     MemoryEffects::Effect *effect =
1110         StringSwitch<MemoryEffects::Effect *>(
1111             effectElement.get("effect").cast<StringAttr>().getValue())
1112             .Case("allocate", MemoryEffects::Allocate::get())
1113             .Case("free", MemoryEffects::Free::get())
1114             .Case("read", MemoryEffects::Read::get())
1115             .Case("write", MemoryEffects::Write::get());
1116 
1117     // Check for a non-default resource to use.
1118     SideEffects::Resource *resource = SideEffects::DefaultResource::get();
1119     if (effectElement.get("test_resource"))
1120       resource = TestResource::get();
1121 
1122     // Check for a result to affect.
1123     if (effectElement.get("on_result"))
1124       effects.emplace_back(effect, getResult(), resource);
1125     else if (Attribute ref = effectElement.get("on_reference"))
1126       effects.emplace_back(effect, ref.cast<SymbolRefAttr>(), resource);
1127     else
1128       effects.emplace_back(effect, resource);
1129   }
1130 }
1131 
1132 void SideEffectOp::getEffects(
1133     SmallVectorImpl<TestEffects::EffectInstance> &effects) {
1134   testSideEffectOpGetEffect(getOperation(), effects);
1135 }
1136 
1137 //===----------------------------------------------------------------------===//
1138 // StringAttrPrettyNameOp
1139 //===----------------------------------------------------------------------===//
1140 
1141 // This op has fancy handling of its SSA result name.
1142 ParseResult StringAttrPrettyNameOp::parse(OpAsmParser &parser,
1143                                           OperationState &result) {
1144   // Add the result types.
1145   for (size_t i = 0, e = parser.getNumResults(); i != e; ++i)
1146     result.addTypes(parser.getBuilder().getIntegerType(32));
1147 
1148   if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
1149     return failure();
1150 
1151   // If the attribute dictionary contains no 'names' attribute, infer it from
1152   // the SSA name (if specified).
1153   bool hadNames = llvm::any_of(result.attributes, [](NamedAttribute attr) {
1154     return attr.getName() == "names";
1155   });
1156 
1157   // If there was no name specified, check to see if there was a useful name
1158   // specified in the asm file.
1159   if (hadNames || parser.getNumResults() == 0)
1160     return success();
1161 
1162   SmallVector<StringRef, 4> names;
1163   auto *context = result.getContext();
1164 
1165   for (size_t i = 0, e = parser.getNumResults(); i != e; ++i) {
1166     auto resultName = parser.getResultName(i);
1167     StringRef nameStr;
1168     if (!resultName.first.empty() && !isdigit(resultName.first[0]))
1169       nameStr = resultName.first;
1170 
1171     names.push_back(nameStr);
1172   }
1173 
1174   auto namesAttr = parser.getBuilder().getStrArrayAttr(names);
1175   result.attributes.push_back({StringAttr::get(context, "names"), namesAttr});
1176   return success();
1177 }
1178 
1179 void StringAttrPrettyNameOp::print(OpAsmPrinter &p) {
1180   // Note that we only need to print the "name" attribute if the asmprinter
1181   // result name disagrees with it.  This can happen in strange cases, e.g.
1182   // when there are conflicts.
1183   bool namesDisagree = getNames().size() != getNumResults();
1184 
1185   SmallString<32> resultNameStr;
1186   for (size_t i = 0, e = getNumResults(); i != e && !namesDisagree; ++i) {
1187     resultNameStr.clear();
1188     llvm::raw_svector_ostream tmpStream(resultNameStr);
1189     p.printOperand(getResult(i), tmpStream);
1190 
1191     auto expectedName = getNames()[i].dyn_cast<StringAttr>();
1192     if (!expectedName ||
1193         tmpStream.str().drop_front() != expectedName.getValue()) {
1194       namesDisagree = true;
1195     }
1196   }
1197 
1198   if (namesDisagree)
1199     p.printOptionalAttrDictWithKeyword((*this)->getAttrs());
1200   else
1201     p.printOptionalAttrDictWithKeyword((*this)->getAttrs(), {"names"});
1202 }
1203 
1204 // We set the SSA name in the asm syntax to the contents of the name
1205 // attribute.
1206 void StringAttrPrettyNameOp::getAsmResultNames(
1207     function_ref<void(Value, StringRef)> setNameFn) {
1208 
1209   auto value = getNames();
1210   for (size_t i = 0, e = value.size(); i != e; ++i)
1211     if (auto str = value[i].dyn_cast<StringAttr>())
1212       if (!str.getValue().empty())
1213         setNameFn(getResult(i), str.getValue());
1214 }
1215 
1216 void CustomResultsNameOp::getAsmResultNames(
1217     function_ref<void(Value, StringRef)> setNameFn) {
1218   ArrayAttr value = getNames();
1219   for (size_t i = 0, e = value.size(); i != e; ++i)
1220     if (auto str = value[i].dyn_cast<StringAttr>())
1221       if (!str.getValue().empty())
1222         setNameFn(getResult(i), str.getValue());
1223 }
1224 
1225 //===----------------------------------------------------------------------===//
1226 // ResultTypeWithTraitOp
1227 //===----------------------------------------------------------------------===//
1228 
1229 LogicalResult ResultTypeWithTraitOp::verify() {
1230   if ((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
1231     return success();
1232   return emitError("result type should have trait 'TestTypeTrait'");
1233 }
1234 
1235 //===----------------------------------------------------------------------===//
1236 // AttrWithTraitOp
1237 //===----------------------------------------------------------------------===//
1238 
1239 LogicalResult AttrWithTraitOp::verify() {
1240   if (getAttr().hasTrait<AttributeTrait::TestAttrTrait>())
1241     return success();
1242   return emitError("'attr' attribute should have trait 'TestAttrTrait'");
1243 }
1244 
1245 //===----------------------------------------------------------------------===//
1246 // RegionIfOp
1247 //===----------------------------------------------------------------------===//
1248 
1249 void RegionIfOp::print(OpAsmPrinter &p) {
1250   p << " ";
1251   p.printOperands(getOperands());
1252   p << ": " << getOperandTypes();
1253   p.printArrowTypeList(getResultTypes());
1254   p << " then ";
1255   p.printRegion(getThenRegion(),
1256                 /*printEntryBlockArgs=*/true,
1257                 /*printBlockTerminators=*/true);
1258   p << " else ";
1259   p.printRegion(getElseRegion(),
1260                 /*printEntryBlockArgs=*/true,
1261                 /*printBlockTerminators=*/true);
1262   p << " join ";
1263   p.printRegion(getJoinRegion(),
1264                 /*printEntryBlockArgs=*/true,
1265                 /*printBlockTerminators=*/true);
1266 }
1267 
1268 ParseResult RegionIfOp::parse(OpAsmParser &parser, OperationState &result) {
1269   SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfos;
1270   SmallVector<Type, 2> operandTypes;
1271 
1272   result.regions.reserve(3);
1273   Region *thenRegion = result.addRegion();
1274   Region *elseRegion = result.addRegion();
1275   Region *joinRegion = result.addRegion();
1276 
1277   // Parse operand, type and arrow type lists.
1278   if (parser.parseOperandList(operandInfos) ||
1279       parser.parseColonTypeList(operandTypes) ||
1280       parser.parseArrowTypeList(result.types))
1281     return failure();
1282 
1283   // Parse all attached regions.
1284   if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
1285       parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
1286       parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
1287     return failure();
1288 
1289   return parser.resolveOperands(operandInfos, operandTypes,
1290                                 parser.getCurrentLocation(), result.operands);
1291 }
1292 
1293 OperandRange RegionIfOp::getSuccessorEntryOperands(unsigned index) {
1294   assert(index < 2 && "invalid region index");
1295   return getOperands();
1296 }
1297 
1298 void RegionIfOp::getSuccessorRegions(
1299     Optional<unsigned> index, ArrayRef<Attribute> operands,
1300     SmallVectorImpl<RegionSuccessor> &regions) {
1301   // We always branch to the join region.
1302   if (index.hasValue()) {
1303     if (index.getValue() < 2)
1304       regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));
1305     else
1306       regions.push_back(RegionSuccessor(getResults()));
1307     return;
1308   }
1309 
1310   // The then and else regions are the entry regions of this op.
1311   regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
1312   regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
1313 }
1314 
1315 void RegionIfOp::getRegionInvocationBounds(
1316     ArrayRef<Attribute> operands,
1317     SmallVectorImpl<InvocationBounds> &invocationBounds) {
1318   // Each region is invoked at most once.
1319   invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});
1320 }
1321 
1322 //===----------------------------------------------------------------------===//
1323 // AnyCondOp
1324 //===----------------------------------------------------------------------===//
1325 
1326 void AnyCondOp::getSuccessorRegions(Optional<unsigned> index,
1327                                     ArrayRef<Attribute> operands,
1328                                     SmallVectorImpl<RegionSuccessor> &regions) {
1329   // The parent op branches into the only region, and the region branches back
1330   // to the parent op.
1331   if (index)
1332     regions.emplace_back(&getRegion());
1333   else
1334     regions.emplace_back(getResults());
1335 }
1336 
1337 void AnyCondOp::getRegionInvocationBounds(
1338     ArrayRef<Attribute> operands,
1339     SmallVectorImpl<InvocationBounds> &invocationBounds) {
1340   invocationBounds.emplace_back(1, 1);
1341 }
1342 
1343 //===----------------------------------------------------------------------===//
1344 // SingleNoTerminatorCustomAsmOp
1345 //===----------------------------------------------------------------------===//
1346 
1347 ParseResult SingleNoTerminatorCustomAsmOp::parse(OpAsmParser &parser,
1348                                                  OperationState &state) {
1349   Region *body = state.addRegion();
1350   if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
1351     return failure();
1352   return success();
1353 }
1354 
1355 void SingleNoTerminatorCustomAsmOp::print(OpAsmPrinter &printer) {
1356   printer.printRegion(
1357       getRegion(), /*printEntryBlockArgs=*/false,
1358       // This op has a single block without terminators. But explicitly mark
1359       // as not printing block terminators for testing.
1360       /*printBlockTerminators=*/false);
1361 }
1362 
1363 //===----------------------------------------------------------------------===//
1364 // TestVerifiersOp
1365 //===----------------------------------------------------------------------===//
1366 
1367 LogicalResult TestVerifiersOp::verify() {
1368   if (!getRegion().hasOneBlock())
1369     return emitOpError("`hasOneBlock` trait hasn't been verified");
1370 
1371   Operation *definingOp = getInput().getDefiningOp();
1372   if (definingOp && failed(mlir::verify(definingOp)))
1373     return emitOpError("operand hasn't been verified");
1374 
1375   emitRemark("success run of verifier");
1376 
1377   return success();
1378 }
1379 
1380 LogicalResult TestVerifiersOp::verifyRegions() {
1381   if (!getRegion().hasOneBlock())
1382     return emitOpError("`hasOneBlock` trait hasn't been verified");
1383 
1384   for (Block &block : getRegion())
1385     for (Operation &op : block)
1386       if (failed(mlir::verify(&op)))
1387         return emitOpError("nested op hasn't been verified");
1388 
1389   emitRemark("success run of region verifier");
1390 
1391   return success();
1392 }
1393 
1394 #include "TestOpEnums.cpp.inc"
1395 #include "TestOpInterfaces.cpp.inc"
1396 #include "TestOpStructs.cpp.inc"
1397 #include "TestTypeInterfaces.cpp.inc"
1398 
1399 #define GET_OP_CLASSES
1400 #include "TestOps.cpp.inc"
1401