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   if (opName == "test.dialect_custom_printer.with.dot") {
375     return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
376       return ParseResult::success();
377     }};
378   }
379   return None;
380 }
381 
382 llvm::unique_function<void(Operation *, OpAsmPrinter &)>
383 TestDialect::getOperationPrinter(Operation *op) const {
384   StringRef opName = op->getName().getStringRef();
385   if (opName == "test.dialect_custom_printer") {
386     return [](Operation *op, OpAsmPrinter &printer) {
387       printer.getStream() << " custom_format";
388     };
389   }
390   if (opName == "test.dialect_custom_format_fallback") {
391     return [](Operation *op, OpAsmPrinter &printer) {
392       printer.getStream() << " custom_format_fallback";
393     };
394   }
395   return {};
396 }
397 
398 //===----------------------------------------------------------------------===//
399 // TestBranchOp
400 //===----------------------------------------------------------------------===//
401 
402 SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
403   assert(index == 0 && "invalid successor index");
404   return SuccessorOperands(getTargetOperandsMutable());
405 }
406 
407 //===----------------------------------------------------------------------===//
408 // TestProducingBranchOp
409 //===----------------------------------------------------------------------===//
410 
411 SuccessorOperands TestProducingBranchOp::getSuccessorOperands(unsigned index) {
412   assert(index <= 1 && "invalid successor index");
413   if (index == 1)
414     return SuccessorOperands(getFirstOperandsMutable());
415   return SuccessorOperands(getSecondOperandsMutable());
416 }
417 
418 //===----------------------------------------------------------------------===//
419 // TestProducingBranchOp
420 //===----------------------------------------------------------------------===//
421 
422 SuccessorOperands TestInternalBranchOp::getSuccessorOperands(unsigned index) {
423   assert(index <= 1 && "invalid successor index");
424   if (index == 0)
425     return SuccessorOperands(0, getSuccessOperandsMutable());
426   return SuccessorOperands(1, getErrorOperandsMutable());
427 }
428 
429 //===----------------------------------------------------------------------===//
430 // TestDialectCanonicalizerOp
431 //===----------------------------------------------------------------------===//
432 
433 static LogicalResult
434 dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,
435                                PatternRewriter &rewriter) {
436   rewriter.replaceOpWithNewOp<arith::ConstantOp>(
437       op, rewriter.getI32IntegerAttr(42));
438   return success();
439 }
440 
441 void TestDialect::getCanonicalizationPatterns(
442     RewritePatternSet &results) const {
443   results.add(&dialectCanonicalizationPattern);
444 }
445 
446 //===----------------------------------------------------------------------===//
447 // TestCallOp
448 //===----------------------------------------------------------------------===//
449 
450 LogicalResult TestCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
451   // Check that the callee attribute was specified.
452   auto fnAttr = (*this)->getAttrOfType<FlatSymbolRefAttr>("callee");
453   if (!fnAttr)
454     return emitOpError("requires a 'callee' symbol reference attribute");
455   if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))
456     return emitOpError() << "'" << fnAttr.getValue()
457                          << "' does not reference a valid function";
458   return success();
459 }
460 
461 //===----------------------------------------------------------------------===//
462 // TestFoldToCallOp
463 //===----------------------------------------------------------------------===//
464 
465 namespace {
466 struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
467   using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
468 
469   LogicalResult matchAndRewrite(FoldToCallOp op,
470                                 PatternRewriter &rewriter) const override {
471     rewriter.replaceOpWithNewOp<func::CallOp>(op, TypeRange(),
472                                               op.getCalleeAttr(), ValueRange());
473     return success();
474   }
475 };
476 } // namespace
477 
478 void FoldToCallOp::getCanonicalizationPatterns(RewritePatternSet &results,
479                                                MLIRContext *context) {
480   results.add<FoldToCallOpPattern>(context);
481 }
482 
483 //===----------------------------------------------------------------------===//
484 // Test Format* operations
485 //===----------------------------------------------------------------------===//
486 
487 //===----------------------------------------------------------------------===//
488 // Parsing
489 
490 static ParseResult parseCustomOptionalOperand(
491     OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &optOperand) {
492   if (succeeded(parser.parseOptionalLParen())) {
493     optOperand.emplace();
494     if (parser.parseOperand(*optOperand) || parser.parseRParen())
495       return failure();
496   }
497   return success();
498 }
499 
500 static ParseResult parseCustomDirectiveOperands(
501     OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,
502     Optional<OpAsmParser::UnresolvedOperand> &optOperand,
503     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands) {
504   if (parser.parseOperand(operand))
505     return failure();
506   if (succeeded(parser.parseOptionalComma())) {
507     optOperand.emplace();
508     if (parser.parseOperand(*optOperand))
509       return failure();
510   }
511   if (parser.parseArrow() || parser.parseLParen() ||
512       parser.parseOperandList(varOperands) || parser.parseRParen())
513     return failure();
514   return success();
515 }
516 static ParseResult
517 parseCustomDirectiveResults(OpAsmParser &parser, Type &operandType,
518                             Type &optOperandType,
519                             SmallVectorImpl<Type> &varOperandTypes) {
520   if (parser.parseColon())
521     return failure();
522 
523   if (parser.parseType(operandType))
524     return failure();
525   if (succeeded(parser.parseOptionalComma())) {
526     if (parser.parseType(optOperandType))
527       return failure();
528   }
529   if (parser.parseArrow() || parser.parseLParen() ||
530       parser.parseTypeList(varOperandTypes) || parser.parseRParen())
531     return failure();
532   return success();
533 }
534 static ParseResult
535 parseCustomDirectiveWithTypeRefs(OpAsmParser &parser, Type operandType,
536                                  Type optOperandType,
537                                  const SmallVectorImpl<Type> &varOperandTypes) {
538   if (parser.parseKeyword("type_refs_capture"))
539     return failure();
540 
541   Type operandType2, optOperandType2;
542   SmallVector<Type, 1> varOperandTypes2;
543   if (parseCustomDirectiveResults(parser, operandType2, optOperandType2,
544                                   varOperandTypes2))
545     return failure();
546 
547   if (operandType != operandType2 || optOperandType != optOperandType2 ||
548       varOperandTypes != varOperandTypes2)
549     return failure();
550 
551   return success();
552 }
553 static ParseResult parseCustomDirectiveOperandsAndTypes(
554     OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,
555     Optional<OpAsmParser::UnresolvedOperand> &optOperand,
556     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands,
557     Type &operandType, Type &optOperandType,
558     SmallVectorImpl<Type> &varOperandTypes) {
559   if (parseCustomDirectiveOperands(parser, operand, optOperand, varOperands) ||
560       parseCustomDirectiveResults(parser, operandType, optOperandType,
561                                   varOperandTypes))
562     return failure();
563   return success();
564 }
565 static ParseResult parseCustomDirectiveRegions(
566     OpAsmParser &parser, Region &region,
567     SmallVectorImpl<std::unique_ptr<Region>> &varRegions) {
568   if (parser.parseRegion(region))
569     return failure();
570   if (failed(parser.parseOptionalComma()))
571     return success();
572   std::unique_ptr<Region> varRegion = std::make_unique<Region>();
573   if (parser.parseRegion(*varRegion))
574     return failure();
575   varRegions.emplace_back(std::move(varRegion));
576   return success();
577 }
578 static ParseResult
579 parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,
580                                SmallVectorImpl<Block *> &varSuccessors) {
581   if (parser.parseSuccessor(successor))
582     return failure();
583   if (failed(parser.parseOptionalComma()))
584     return success();
585   Block *varSuccessor;
586   if (parser.parseSuccessor(varSuccessor))
587     return failure();
588   varSuccessors.append(2, varSuccessor);
589   return success();
590 }
591 static ParseResult parseCustomDirectiveAttributes(OpAsmParser &parser,
592                                                   IntegerAttr &attr,
593                                                   IntegerAttr &optAttr) {
594   if (parser.parseAttribute(attr))
595     return failure();
596   if (succeeded(parser.parseOptionalComma())) {
597     if (parser.parseAttribute(optAttr))
598       return failure();
599   }
600   return success();
601 }
602 
603 static ParseResult parseCustomDirectiveAttrDict(OpAsmParser &parser,
604                                                 NamedAttrList &attrs) {
605   return parser.parseOptionalAttrDict(attrs);
606 }
607 static ParseResult parseCustomDirectiveOptionalOperandRef(
608     OpAsmParser &parser, Optional<OpAsmParser::UnresolvedOperand> &optOperand) {
609   int64_t operandCount = 0;
610   if (parser.parseInteger(operandCount))
611     return failure();
612   bool expectedOptionalOperand = operandCount == 0;
613   return success(expectedOptionalOperand != optOperand.hasValue());
614 }
615 
616 //===----------------------------------------------------------------------===//
617 // Printing
618 
619 static void printCustomOptionalOperand(OpAsmPrinter &printer, Operation *,
620                                        Value optOperand) {
621   if (optOperand)
622     printer << "(" << optOperand << ") ";
623 }
624 
625 static void printCustomDirectiveOperands(OpAsmPrinter &printer, Operation *,
626                                          Value operand, Value optOperand,
627                                          OperandRange varOperands) {
628   printer << operand;
629   if (optOperand)
630     printer << ", " << optOperand;
631   printer << " -> (" << varOperands << ")";
632 }
633 static void printCustomDirectiveResults(OpAsmPrinter &printer, Operation *,
634                                         Type operandType, Type optOperandType,
635                                         TypeRange varOperandTypes) {
636   printer << " : " << operandType;
637   if (optOperandType)
638     printer << ", " << optOperandType;
639   printer << " -> (" << varOperandTypes << ")";
640 }
641 static void printCustomDirectiveWithTypeRefs(OpAsmPrinter &printer,
642                                              Operation *op, Type operandType,
643                                              Type optOperandType,
644                                              TypeRange varOperandTypes) {
645   printer << " type_refs_capture ";
646   printCustomDirectiveResults(printer, op, operandType, optOperandType,
647                               varOperandTypes);
648 }
649 static void printCustomDirectiveOperandsAndTypes(
650     OpAsmPrinter &printer, Operation *op, Value operand, Value optOperand,
651     OperandRange varOperands, Type operandType, Type optOperandType,
652     TypeRange varOperandTypes) {
653   printCustomDirectiveOperands(printer, op, operand, optOperand, varOperands);
654   printCustomDirectiveResults(printer, op, operandType, optOperandType,
655                               varOperandTypes);
656 }
657 static void printCustomDirectiveRegions(OpAsmPrinter &printer, Operation *,
658                                         Region &region,
659                                         MutableArrayRef<Region> varRegions) {
660   printer.printRegion(region);
661   if (!varRegions.empty()) {
662     printer << ", ";
663     for (Region &region : varRegions)
664       printer.printRegion(region);
665   }
666 }
667 static void printCustomDirectiveSuccessors(OpAsmPrinter &printer, Operation *,
668                                            Block *successor,
669                                            SuccessorRange varSuccessors) {
670   printer << successor;
671   if (!varSuccessors.empty())
672     printer << ", " << varSuccessors.front();
673 }
674 static void printCustomDirectiveAttributes(OpAsmPrinter &printer, Operation *,
675                                            Attribute attribute,
676                                            Attribute optAttribute) {
677   printer << attribute;
678   if (optAttribute)
679     printer << ", " << optAttribute;
680 }
681 
682 static void printCustomDirectiveAttrDict(OpAsmPrinter &printer, Operation *op,
683                                          DictionaryAttr attrs) {
684   printer.printOptionalAttrDict(attrs.getValue());
685 }
686 
687 static void printCustomDirectiveOptionalOperandRef(OpAsmPrinter &printer,
688                                                    Operation *op,
689                                                    Value optOperand) {
690   printer << (optOperand ? "1" : "0");
691 }
692 
693 //===----------------------------------------------------------------------===//
694 // Test IsolatedRegionOp - parse passthrough region arguments.
695 //===----------------------------------------------------------------------===//
696 
697 ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
698                                     OperationState &result) {
699   // Parse the input operand.
700   OpAsmParser::Argument argInfo;
701   argInfo.type = parser.getBuilder().getIndexType();
702   if (parser.parseOperand(argInfo.ssaName) ||
703       parser.resolveOperand(argInfo.ssaName, argInfo.type, result.operands))
704     return failure();
705 
706   // Parse the body region, and reuse the operand info as the argument info.
707   Region *body = result.addRegion();
708   return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);
709 }
710 
711 void IsolatedRegionOp::print(OpAsmPrinter &p) {
712   p << "test.isolated_region ";
713   p.printOperand(getOperand());
714   p.shadowRegionArgs(getRegion(), getOperand());
715   p << ' ';
716   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
717 }
718 
719 //===----------------------------------------------------------------------===//
720 // Test SSACFGRegionOp
721 //===----------------------------------------------------------------------===//
722 
723 RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {
724   return RegionKind::SSACFG;
725 }
726 
727 //===----------------------------------------------------------------------===//
728 // Test GraphRegionOp
729 //===----------------------------------------------------------------------===//
730 
731 RegionKind GraphRegionOp::getRegionKind(unsigned index) {
732   return RegionKind::Graph;
733 }
734 
735 //===----------------------------------------------------------------------===//
736 // Test AffineScopeOp
737 //===----------------------------------------------------------------------===//
738 
739 ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
740   // Parse the body region, and reuse the operand info as the argument info.
741   Region *body = result.addRegion();
742   return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
743 }
744 
745 void AffineScopeOp::print(OpAsmPrinter &p) {
746   p << "test.affine_scope ";
747   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
748 }
749 
750 //===----------------------------------------------------------------------===//
751 // Test parser.
752 //===----------------------------------------------------------------------===//
753 
754 ParseResult ParseIntegerLiteralOp::parse(OpAsmParser &parser,
755                                          OperationState &result) {
756   if (parser.parseOptionalColon())
757     return success();
758   uint64_t numResults;
759   if (parser.parseInteger(numResults))
760     return failure();
761 
762   IndexType type = parser.getBuilder().getIndexType();
763   for (unsigned i = 0; i < numResults; ++i)
764     result.addTypes(type);
765   return success();
766 }
767 
768 void ParseIntegerLiteralOp::print(OpAsmPrinter &p) {
769   if (unsigned numResults = getNumResults())
770     p << " : " << numResults;
771 }
772 
773 ParseResult ParseWrappedKeywordOp::parse(OpAsmParser &parser,
774                                          OperationState &result) {
775   StringRef keyword;
776   if (parser.parseKeyword(&keyword))
777     return failure();
778   result.addAttribute("keyword", parser.getBuilder().getStringAttr(keyword));
779   return success();
780 }
781 
782 void ParseWrappedKeywordOp::print(OpAsmPrinter &p) { p << " " << getKeyword(); }
783 
784 //===----------------------------------------------------------------------===//
785 // Test WrapRegionOp - wrapping op exercising `parseGenericOperation()`.
786 
787 ParseResult WrappingRegionOp::parse(OpAsmParser &parser,
788                                     OperationState &result) {
789   if (parser.parseKeyword("wraps"))
790     return failure();
791 
792   // Parse the wrapped op in a region
793   Region &body = *result.addRegion();
794   body.push_back(new Block);
795   Block &block = body.back();
796   Operation *wrappedOp = parser.parseGenericOperation(&block, block.begin());
797   if (!wrappedOp)
798     return failure();
799 
800   // Create a return terminator in the inner region, pass as operand to the
801   // terminator the returned values from the wrapped operation.
802   SmallVector<Value, 8> returnOperands(wrappedOp->getResults());
803   OpBuilder builder(parser.getContext());
804   builder.setInsertionPointToEnd(&block);
805   builder.create<TestReturnOp>(wrappedOp->getLoc(), returnOperands);
806 
807   // Get the results type for the wrapping op from the terminator operands.
808   Operation &returnOp = body.back().back();
809   result.types.append(returnOp.operand_type_begin(),
810                       returnOp.operand_type_end());
811 
812   // Use the location of the wrapped op for the "test.wrapping_region" op.
813   result.location = wrappedOp->getLoc();
814 
815   return success();
816 }
817 
818 void WrappingRegionOp::print(OpAsmPrinter &p) {
819   p << " wraps ";
820   p.printGenericOp(&getRegion().front().front());
821 }
822 
823 //===----------------------------------------------------------------------===//
824 // Test PrettyPrintedRegionOp -  exercising the following parser APIs
825 //   parseGenericOperationAfterOpName
826 //   parseCustomOperationName
827 //===----------------------------------------------------------------------===//
828 
829 ParseResult PrettyPrintedRegionOp::parse(OpAsmParser &parser,
830                                          OperationState &result) {
831 
832   SMLoc loc = parser.getCurrentLocation();
833   Location currLocation = parser.getEncodedSourceLoc(loc);
834 
835   // Parse the operands.
836   SmallVector<OpAsmParser::UnresolvedOperand, 2> operands;
837   if (parser.parseOperandList(operands))
838     return failure();
839 
840   // Check if we are parsing the pretty-printed version
841   //  test.pretty_printed_region start <inner-op> end : <functional-type>
842   // Else fallback to parsing the "non pretty-printed" version.
843   if (!succeeded(parser.parseOptionalKeyword("start")))
844     return parser.parseGenericOperationAfterOpName(
845         result, llvm::makeArrayRef(operands));
846 
847   FailureOr<OperationName> parseOpNameInfo = parser.parseCustomOperationName();
848   if (failed(parseOpNameInfo))
849     return failure();
850 
851   StringAttr innerOpName = parseOpNameInfo->getIdentifier();
852 
853   FunctionType opFntype;
854   Optional<Location> explicitLoc;
855   if (parser.parseKeyword("end") || parser.parseColon() ||
856       parser.parseType(opFntype) ||
857       parser.parseOptionalLocationSpecifier(explicitLoc))
858     return failure();
859 
860   // If location of the op is explicitly provided, then use it; Else use
861   // the parser's current location.
862   Location opLoc = explicitLoc.getValueOr(currLocation);
863 
864   // Derive the SSA-values for op's operands.
865   if (parser.resolveOperands(operands, opFntype.getInputs(), loc,
866                              result.operands))
867     return failure();
868 
869   // Add a region for op.
870   Region &region = *result.addRegion();
871 
872   // Create a basic-block inside op's region.
873   Block &block = region.emplaceBlock();
874 
875   // Create and insert an "inner-op" operation in the block.
876   // Just for testing purposes, we can assume that inner op is a binary op with
877   // result and operand types all same as the test-op's first operand.
878   Type innerOpType = opFntype.getInput(0);
879   Value lhs = block.addArgument(innerOpType, opLoc);
880   Value rhs = block.addArgument(innerOpType, opLoc);
881 
882   OpBuilder builder(parser.getBuilder().getContext());
883   builder.setInsertionPointToStart(&block);
884 
885   Operation *innerOp =
886       builder.create(opLoc, innerOpName, /*operands=*/{lhs, rhs}, innerOpType);
887 
888   // Insert a return statement in the block returning the inner-op's result.
889   builder.create<TestReturnOp>(innerOp->getLoc(), innerOp->getResults());
890 
891   // Populate the op operation-state with result-type and location.
892   result.addTypes(opFntype.getResults());
893   result.location = innerOp->getLoc();
894 
895   return success();
896 }
897 
898 void PrettyPrintedRegionOp::print(OpAsmPrinter &p) {
899   p << ' ';
900   p.printOperands(getOperands());
901 
902   Operation &innerOp = getRegion().front().front();
903   // Assuming that region has a single non-terminator inner-op, if the inner-op
904   // meets some criteria (which in this case is a simple one  based on the name
905   // of inner-op), then we can print the entire region in a succinct way.
906   // Here we assume that the prototype of "special.op" can be trivially derived
907   // while parsing it back.
908   if (innerOp.getName().getStringRef().equals("special.op")) {
909     p << " start special.op end";
910   } else {
911     p << " (";
912     p.printRegion(getRegion());
913     p << ")";
914   }
915 
916   p << " : ";
917   p.printFunctionalType(*this);
918 }
919 
920 //===----------------------------------------------------------------------===//
921 // Test PolyForOp - parse list of region arguments.
922 //===----------------------------------------------------------------------===//
923 
924 ParseResult PolyForOp::parse(OpAsmParser &parser, OperationState &result) {
925   SmallVector<OpAsmParser::Argument, 4> ivsInfo;
926   // Parse list of region arguments without a delimiter.
927   if (parser.parseArgumentList(ivsInfo, OpAsmParser::Delimiter::None))
928     return failure();
929 
930   // Parse the body region.
931   Region *body = result.addRegion();
932   for (auto &iv : ivsInfo)
933     iv.type = parser.getBuilder().getIndexType();
934   return parser.parseRegion(*body, ivsInfo);
935 }
936 
937 void PolyForOp::print(OpAsmPrinter &p) { p.printGenericOp(*this); }
938 
939 void PolyForOp::getAsmBlockArgumentNames(Region &region,
940                                          OpAsmSetValueNameFn setNameFn) {
941   auto arrayAttr = getOperation()->getAttrOfType<ArrayAttr>("arg_names");
942   if (!arrayAttr)
943     return;
944   auto args = getRegion().front().getArguments();
945   auto e = std::min(arrayAttr.size(), args.size());
946   for (unsigned i = 0; i < e; ++i) {
947     if (auto strAttr = arrayAttr[i].dyn_cast<StringAttr>())
948       setNameFn(args[i], strAttr.getValue());
949   }
950 }
951 
952 //===----------------------------------------------------------------------===//
953 // Test removing op with inner ops.
954 //===----------------------------------------------------------------------===//
955 
956 namespace {
957 struct TestRemoveOpWithInnerOps
958     : public OpRewritePattern<TestOpWithRegionPattern> {
959   using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;
960 
961   void initialize() { setDebugName("TestRemoveOpWithInnerOps"); }
962 
963   LogicalResult matchAndRewrite(TestOpWithRegionPattern op,
964                                 PatternRewriter &rewriter) const override {
965     rewriter.eraseOp(op);
966     return success();
967   }
968 };
969 } // namespace
970 
971 void TestOpWithRegionPattern::getCanonicalizationPatterns(
972     RewritePatternSet &results, MLIRContext *context) {
973   results.add<TestRemoveOpWithInnerOps>(context);
974 }
975 
976 OpFoldResult TestOpWithRegionFold::fold(ArrayRef<Attribute> operands) {
977   return getOperand();
978 }
979 
980 OpFoldResult TestOpConstant::fold(ArrayRef<Attribute> operands) {
981   return getValue();
982 }
983 
984 LogicalResult TestOpWithVariadicResultsAndFolder::fold(
985     ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
986   for (Value input : this->getOperands()) {
987     results.push_back(input);
988   }
989   return success();
990 }
991 
992 OpFoldResult TestOpInPlaceFold::fold(ArrayRef<Attribute> operands) {
993   assert(operands.size() == 1);
994   if (operands.front()) {
995     (*this)->setAttr("attr", operands.front());
996     return getResult();
997   }
998   return {};
999 }
1000 
1001 OpFoldResult TestPassthroughFold::fold(ArrayRef<Attribute> operands) {
1002   return getOperand();
1003 }
1004 
1005 LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
1006     MLIRContext *, Optional<Location> location, ValueRange operands,
1007     DictionaryAttr attributes, RegionRange regions,
1008     SmallVectorImpl<Type> &inferredReturnTypes) {
1009   if (operands[0].getType() != operands[1].getType()) {
1010     return emitOptionalError(location, "operand type mismatch ",
1011                              operands[0].getType(), " vs ",
1012                              operands[1].getType());
1013   }
1014   inferredReturnTypes.assign({operands[0].getType()});
1015   return success();
1016 }
1017 
1018 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
1019     MLIRContext *context, Optional<Location> location, ValueShapeRange operands,
1020     DictionaryAttr attributes, RegionRange regions,
1021     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
1022   // Create return type consisting of the last element of the first operand.
1023   auto operandType = operands.front().getType();
1024   auto sval = operandType.dyn_cast<ShapedType>();
1025   if (!sval) {
1026     return emitOptionalError(location, "only shaped type operands allowed");
1027   }
1028   int64_t dim =
1029       sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;
1030   auto type = IntegerType::get(context, 17);
1031   inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type));
1032   return success();
1033 }
1034 
1035 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
1036     OpBuilder &builder, ValueRange operands,
1037     llvm::SmallVectorImpl<Value> &shapes) {
1038   shapes = SmallVector<Value, 1>{
1039       builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
1040   return success();
1041 }
1042 
1043 LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
1044     OpBuilder &builder, ValueRange operands,
1045     llvm::SmallVectorImpl<Value> &shapes) {
1046   Location loc = getLoc();
1047   shapes.reserve(operands.size());
1048   for (Value operand : llvm::reverse(operands)) {
1049     auto rank = operand.getType().cast<RankedTensorType>().getRank();
1050     auto currShape = llvm::to_vector<4>(
1051         llvm::map_range(llvm::seq<int64_t>(0, rank), [&](int64_t dim) -> Value {
1052           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
1053         }));
1054     shapes.push_back(builder.create<tensor::FromElementsOp>(
1055         getLoc(), RankedTensorType::get({rank}, builder.getIndexType()),
1056         currShape));
1057   }
1058   return success();
1059 }
1060 
1061 LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
1062     OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
1063   Location loc = getLoc();
1064   shapes.reserve(getNumOperands());
1065   for (Value operand : llvm::reverse(getOperands())) {
1066     auto currShape = llvm::to_vector<4>(llvm::map_range(
1067         llvm::seq<int64_t>(
1068             0, operand.getType().cast<RankedTensorType>().getRank()),
1069         [&](int64_t dim) -> Value {
1070           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
1071         }));
1072     shapes.emplace_back(std::move(currShape));
1073   }
1074   return success();
1075 }
1076 
1077 //===----------------------------------------------------------------------===//
1078 // Test SideEffect interfaces
1079 //===----------------------------------------------------------------------===//
1080 
1081 namespace {
1082 /// A test resource for side effects.
1083 struct TestResource : public SideEffects::Resource::Base<TestResource> {
1084   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestResource)
1085 
1086   StringRef getName() final { return "<Test>"; }
1087 };
1088 } // namespace
1089 
1090 static void testSideEffectOpGetEffect(
1091     Operation *op,
1092     SmallVectorImpl<SideEffects::EffectInstance<TestEffects::Effect>>
1093         &effects) {
1094   auto effectsAttr = op->getAttrOfType<AffineMapAttr>("effect_parameter");
1095   if (!effectsAttr)
1096     return;
1097 
1098   effects.emplace_back(TestEffects::Concrete::get(), effectsAttr);
1099 }
1100 
1101 void SideEffectOp::getEffects(
1102     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
1103   // Check for an effects attribute on the op instance.
1104   ArrayAttr effectsAttr = (*this)->getAttrOfType<ArrayAttr>("effects");
1105   if (!effectsAttr)
1106     return;
1107 
1108   // If there is one, it is an array of dictionary attributes that hold
1109   // information on the effects of this operation.
1110   for (Attribute element : effectsAttr) {
1111     DictionaryAttr effectElement = element.cast<DictionaryAttr>();
1112 
1113     // Get the specific memory effect.
1114     MemoryEffects::Effect *effect =
1115         StringSwitch<MemoryEffects::Effect *>(
1116             effectElement.get("effect").cast<StringAttr>().getValue())
1117             .Case("allocate", MemoryEffects::Allocate::get())
1118             .Case("free", MemoryEffects::Free::get())
1119             .Case("read", MemoryEffects::Read::get())
1120             .Case("write", MemoryEffects::Write::get());
1121 
1122     // Check for a non-default resource to use.
1123     SideEffects::Resource *resource = SideEffects::DefaultResource::get();
1124     if (effectElement.get("test_resource"))
1125       resource = TestResource::get();
1126 
1127     // Check for a result to affect.
1128     if (effectElement.get("on_result"))
1129       effects.emplace_back(effect, getResult(), resource);
1130     else if (Attribute ref = effectElement.get("on_reference"))
1131       effects.emplace_back(effect, ref.cast<SymbolRefAttr>(), resource);
1132     else
1133       effects.emplace_back(effect, resource);
1134   }
1135 }
1136 
1137 void SideEffectOp::getEffects(
1138     SmallVectorImpl<TestEffects::EffectInstance> &effects) {
1139   testSideEffectOpGetEffect(getOperation(), effects);
1140 }
1141 
1142 //===----------------------------------------------------------------------===//
1143 // StringAttrPrettyNameOp
1144 //===----------------------------------------------------------------------===//
1145 
1146 // This op has fancy handling of its SSA result name.
1147 ParseResult StringAttrPrettyNameOp::parse(OpAsmParser &parser,
1148                                           OperationState &result) {
1149   // Add the result types.
1150   for (size_t i = 0, e = parser.getNumResults(); i != e; ++i)
1151     result.addTypes(parser.getBuilder().getIntegerType(32));
1152 
1153   if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
1154     return failure();
1155 
1156   // If the attribute dictionary contains no 'names' attribute, infer it from
1157   // the SSA name (if specified).
1158   bool hadNames = llvm::any_of(result.attributes, [](NamedAttribute attr) {
1159     return attr.getName() == "names";
1160   });
1161 
1162   // If there was no name specified, check to see if there was a useful name
1163   // specified in the asm file.
1164   if (hadNames || parser.getNumResults() == 0)
1165     return success();
1166 
1167   SmallVector<StringRef, 4> names;
1168   auto *context = result.getContext();
1169 
1170   for (size_t i = 0, e = parser.getNumResults(); i != e; ++i) {
1171     auto resultName = parser.getResultName(i);
1172     StringRef nameStr;
1173     if (!resultName.first.empty() && !isdigit(resultName.first[0]))
1174       nameStr = resultName.first;
1175 
1176     names.push_back(nameStr);
1177   }
1178 
1179   auto namesAttr = parser.getBuilder().getStrArrayAttr(names);
1180   result.attributes.push_back({StringAttr::get(context, "names"), namesAttr});
1181   return success();
1182 }
1183 
1184 void StringAttrPrettyNameOp::print(OpAsmPrinter &p) {
1185   // Note that we only need to print the "name" attribute if the asmprinter
1186   // result name disagrees with it.  This can happen in strange cases, e.g.
1187   // when there are conflicts.
1188   bool namesDisagree = getNames().size() != getNumResults();
1189 
1190   SmallString<32> resultNameStr;
1191   for (size_t i = 0, e = getNumResults(); i != e && !namesDisagree; ++i) {
1192     resultNameStr.clear();
1193     llvm::raw_svector_ostream tmpStream(resultNameStr);
1194     p.printOperand(getResult(i), tmpStream);
1195 
1196     auto expectedName = getNames()[i].dyn_cast<StringAttr>();
1197     if (!expectedName ||
1198         tmpStream.str().drop_front() != expectedName.getValue()) {
1199       namesDisagree = true;
1200     }
1201   }
1202 
1203   if (namesDisagree)
1204     p.printOptionalAttrDictWithKeyword((*this)->getAttrs());
1205   else
1206     p.printOptionalAttrDictWithKeyword((*this)->getAttrs(), {"names"});
1207 }
1208 
1209 // We set the SSA name in the asm syntax to the contents of the name
1210 // attribute.
1211 void StringAttrPrettyNameOp::getAsmResultNames(
1212     function_ref<void(Value, StringRef)> setNameFn) {
1213 
1214   auto value = getNames();
1215   for (size_t i = 0, e = value.size(); i != e; ++i)
1216     if (auto str = value[i].dyn_cast<StringAttr>())
1217       if (!str.getValue().empty())
1218         setNameFn(getResult(i), str.getValue());
1219 }
1220 
1221 void CustomResultsNameOp::getAsmResultNames(
1222     function_ref<void(Value, StringRef)> setNameFn) {
1223   ArrayAttr value = getNames();
1224   for (size_t i = 0, e = value.size(); i != e; ++i)
1225     if (auto str = value[i].dyn_cast<StringAttr>())
1226       if (!str.getValue().empty())
1227         setNameFn(getResult(i), str.getValue());
1228 }
1229 
1230 //===----------------------------------------------------------------------===//
1231 // ResultTypeWithTraitOp
1232 //===----------------------------------------------------------------------===//
1233 
1234 LogicalResult ResultTypeWithTraitOp::verify() {
1235   if ((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
1236     return success();
1237   return emitError("result type should have trait 'TestTypeTrait'");
1238 }
1239 
1240 //===----------------------------------------------------------------------===//
1241 // AttrWithTraitOp
1242 //===----------------------------------------------------------------------===//
1243 
1244 LogicalResult AttrWithTraitOp::verify() {
1245   if (getAttr().hasTrait<AttributeTrait::TestAttrTrait>())
1246     return success();
1247   return emitError("'attr' attribute should have trait 'TestAttrTrait'");
1248 }
1249 
1250 //===----------------------------------------------------------------------===//
1251 // RegionIfOp
1252 //===----------------------------------------------------------------------===//
1253 
1254 void RegionIfOp::print(OpAsmPrinter &p) {
1255   p << " ";
1256   p.printOperands(getOperands());
1257   p << ": " << getOperandTypes();
1258   p.printArrowTypeList(getResultTypes());
1259   p << " then ";
1260   p.printRegion(getThenRegion(),
1261                 /*printEntryBlockArgs=*/true,
1262                 /*printBlockTerminators=*/true);
1263   p << " else ";
1264   p.printRegion(getElseRegion(),
1265                 /*printEntryBlockArgs=*/true,
1266                 /*printBlockTerminators=*/true);
1267   p << " join ";
1268   p.printRegion(getJoinRegion(),
1269                 /*printEntryBlockArgs=*/true,
1270                 /*printBlockTerminators=*/true);
1271 }
1272 
1273 ParseResult RegionIfOp::parse(OpAsmParser &parser, OperationState &result) {
1274   SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfos;
1275   SmallVector<Type, 2> operandTypes;
1276 
1277   result.regions.reserve(3);
1278   Region *thenRegion = result.addRegion();
1279   Region *elseRegion = result.addRegion();
1280   Region *joinRegion = result.addRegion();
1281 
1282   // Parse operand, type and arrow type lists.
1283   if (parser.parseOperandList(operandInfos) ||
1284       parser.parseColonTypeList(operandTypes) ||
1285       parser.parseArrowTypeList(result.types))
1286     return failure();
1287 
1288   // Parse all attached regions.
1289   if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
1290       parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
1291       parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
1292     return failure();
1293 
1294   return parser.resolveOperands(operandInfos, operandTypes,
1295                                 parser.getCurrentLocation(), result.operands);
1296 }
1297 
1298 OperandRange RegionIfOp::getSuccessorEntryOperands(unsigned index) {
1299   assert(index < 2 && "invalid region index");
1300   return getOperands();
1301 }
1302 
1303 void RegionIfOp::getSuccessorRegions(
1304     Optional<unsigned> index, ArrayRef<Attribute> operands,
1305     SmallVectorImpl<RegionSuccessor> &regions) {
1306   // We always branch to the join region.
1307   if (index.hasValue()) {
1308     if (index.getValue() < 2)
1309       regions.push_back(RegionSuccessor(&getJoinRegion(), getJoinArgs()));
1310     else
1311       regions.push_back(RegionSuccessor(getResults()));
1312     return;
1313   }
1314 
1315   // The then and else regions are the entry regions of this op.
1316   regions.push_back(RegionSuccessor(&getThenRegion(), getThenArgs()));
1317   regions.push_back(RegionSuccessor(&getElseRegion(), getElseArgs()));
1318 }
1319 
1320 void RegionIfOp::getRegionInvocationBounds(
1321     ArrayRef<Attribute> operands,
1322     SmallVectorImpl<InvocationBounds> &invocationBounds) {
1323   // Each region is invoked at most once.
1324   invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});
1325 }
1326 
1327 //===----------------------------------------------------------------------===//
1328 // AnyCondOp
1329 //===----------------------------------------------------------------------===//
1330 
1331 void AnyCondOp::getSuccessorRegions(Optional<unsigned> index,
1332                                     ArrayRef<Attribute> operands,
1333                                     SmallVectorImpl<RegionSuccessor> &regions) {
1334   // The parent op branches into the only region, and the region branches back
1335   // to the parent op.
1336   if (index)
1337     regions.emplace_back(&getRegion());
1338   else
1339     regions.emplace_back(getResults());
1340 }
1341 
1342 void AnyCondOp::getRegionInvocationBounds(
1343     ArrayRef<Attribute> operands,
1344     SmallVectorImpl<InvocationBounds> &invocationBounds) {
1345   invocationBounds.emplace_back(1, 1);
1346 }
1347 
1348 //===----------------------------------------------------------------------===//
1349 // SingleNoTerminatorCustomAsmOp
1350 //===----------------------------------------------------------------------===//
1351 
1352 ParseResult SingleNoTerminatorCustomAsmOp::parse(OpAsmParser &parser,
1353                                                  OperationState &state) {
1354   Region *body = state.addRegion();
1355   if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
1356     return failure();
1357   return success();
1358 }
1359 
1360 void SingleNoTerminatorCustomAsmOp::print(OpAsmPrinter &printer) {
1361   printer.printRegion(
1362       getRegion(), /*printEntryBlockArgs=*/false,
1363       // This op has a single block without terminators. But explicitly mark
1364       // as not printing block terminators for testing.
1365       /*printBlockTerminators=*/false);
1366 }
1367 
1368 //===----------------------------------------------------------------------===//
1369 // TestVerifiersOp
1370 //===----------------------------------------------------------------------===//
1371 
1372 LogicalResult TestVerifiersOp::verify() {
1373   if (!getRegion().hasOneBlock())
1374     return emitOpError("`hasOneBlock` trait hasn't been verified");
1375 
1376   Operation *definingOp = getInput().getDefiningOp();
1377   if (definingOp && failed(mlir::verify(definingOp)))
1378     return emitOpError("operand hasn't been verified");
1379 
1380   emitRemark("success run of verifier");
1381 
1382   return success();
1383 }
1384 
1385 LogicalResult TestVerifiersOp::verifyRegions() {
1386   if (!getRegion().hasOneBlock())
1387     return emitOpError("`hasOneBlock` trait hasn't been verified");
1388 
1389   for (Block &block : getRegion())
1390     for (Operation &op : block)
1391       if (failed(mlir::verify(&op)))
1392         return emitOpError("nested op hasn't been verified");
1393 
1394   emitRemark("success run of region verifier");
1395 
1396   return success();
1397 }
1398 
1399 #include "TestOpEnums.cpp.inc"
1400 #include "TestOpInterfaces.cpp.inc"
1401 #include "TestOpStructs.cpp.inc"
1402 #include "TestTypeInterfaces.cpp.inc"
1403 
1404 #define GET_OP_CLASSES
1405 #include "TestOps.cpp.inc"
1406