1 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
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 "mlir/Dialect/Shape/IR/Shape.h"
10 
11 #include "mlir/Dialect/Traits.h"
12 #include "mlir/IR/Builders.h"
13 #include "mlir/IR/DialectImplementation.h"
14 #include "mlir/IR/PatternMatch.h"
15 #include "mlir/IR/StandardTypes.h"
16 #include "llvm/ADT/SmallString.h"
17 #include "llvm/Support/raw_ostream.h"
18 
19 using namespace mlir;
20 using namespace mlir::shape;
21 
22 namespace {
23 #include "ShapeCanonicalization.inc"
24 }
25 
26 ShapeDialect::ShapeDialect(MLIRContext *context)
27     : Dialect(getDialectNamespace(), context) {
28   addOperations<
29 #define GET_OP_LIST
30 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
31       >();
32   addTypes<ComponentType, ElementType, ShapeType, SizeType, ValueShapeType,
33            WitnessType>();
34   // Allow unknown operations during prototyping and testing. As the dialect is
35   // still evolving it makes it simple to start with an unregistered ops and
36   // try different variants before actually defining the op.
37   allowUnknownOperations();
38 }
39 
40 Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
41                                              Attribute value, Type type,
42                                              Location loc) {
43   if (auto shapeType = type.dyn_cast<ShapeType>())
44     return builder.create<ConstShapeOp>(loc, type,
45                                         value.cast<DenseIntElementsAttr>());
46   if (auto sizeType = type.dyn_cast<SizeType>())
47     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
48   if (auto witnessType = type.dyn_cast<WitnessType>())
49     return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
50   return nullptr;
51 }
52 
53 /// Parse a type registered to this dialect.
54 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
55   StringRef keyword;
56   if (parser.parseKeyword(&keyword))
57     return Type();
58 
59   if (keyword == "component")
60     return ComponentType::get(getContext());
61   if (keyword == "element")
62     return ElementType::get(getContext());
63   if (keyword == "shape")
64     return ShapeType::get(getContext());
65   if (keyword == "size")
66     return SizeType::get(getContext());
67   if (keyword == "value_shape")
68     return ValueShapeType::get(getContext());
69   if (keyword == "witness")
70     return WitnessType::get(getContext());
71 
72   parser.emitError(parser.getNameLoc(), "unknown shape type: ") << keyword;
73   return Type();
74 }
75 
76 /// Print a type registered to this dialect.
77 void ShapeDialect::printType(Type type, DialectAsmPrinter &os) const {
78   switch (type.getKind()) {
79   case ShapeTypes::Component:
80     os << "component";
81     return;
82   case ShapeTypes::Element:
83     os << "element";
84     return;
85   case ShapeTypes::Size:
86     os << "size";
87     return;
88   case ShapeTypes::Shape:
89     os << "shape";
90     return;
91   case ShapeTypes::ValueShape:
92     os << "value_shape";
93     return;
94   case ShapeTypes::Witness:
95     os << "witness";
96     return;
97   default:
98     llvm_unreachable("unexpected 'shape' type kind");
99   }
100 }
101 
102 //===----------------------------------------------------------------------===//
103 // AnyOp
104 //===----------------------------------------------------------------------===//
105 
106 // TODO: Canonicalization should be implemented for shapes that can be
107 // determined through mixtures of the known dimensions of the inputs.
108 OpFoldResult AnyOp::fold(ArrayRef<Attribute> operands) {
109   // Only the last operand is checked because AnyOp is commutative.
110   if (operands.back())
111     return operands.back();
112 
113   return nullptr;
114 }
115 
116 //===----------------------------------------------------------------------===//
117 // AssumingOp
118 //===----------------------------------------------------------------------===//
119 
120 static ParseResult parseAssumingOp(OpAsmParser &parser,
121                                    OperationState &result) {
122   result.regions.reserve(1);
123   Region *doRegion = result.addRegion();
124 
125   auto &builder = parser.getBuilder();
126   OpAsmParser::OperandType cond;
127   if (parser.parseOperand(cond) ||
128       parser.resolveOperand(cond, builder.getType<WitnessType>(),
129                             result.operands))
130     return failure();
131 
132   // Parse optional results type list.
133   if (parser.parseOptionalArrowTypeList(result.types))
134     return failure();
135 
136   // Parse the region and add a terminator if elided.
137   if (parser.parseRegion(*doRegion, /*arguments=*/{}, /*argTypes=*/{}))
138     return failure();
139   AssumingOp::ensureTerminator(*doRegion, parser.getBuilder(), result.location);
140 
141   // Parse the optional attribute list.
142   if (parser.parseOptionalAttrDict(result.attributes))
143     return failure();
144   return success();
145 }
146 
147 static void print(OpAsmPrinter &p, AssumingOp op) {
148   bool yieldsResults = !op.results().empty();
149 
150   p << AssumingOp::getOperationName() << " " << op.witness();
151   if (yieldsResults) {
152     p << " -> (" << op.getResultTypes() << ")";
153   }
154   p.printRegion(op.doRegion(),
155                 /*printEntryBlockArgs=*/false,
156                 /*printBlockTerminators=*/yieldsResults);
157   p.printOptionalAttrDict(op.getAttrs());
158 }
159 
160 namespace {
161 // Removes AssumingOp with a passing witness and inlines the region.
162 struct AssumingWithTrue : public OpRewritePattern<AssumingOp> {
163   using OpRewritePattern<AssumingOp>::OpRewritePattern;
164 
165   LogicalResult matchAndRewrite(AssumingOp op,
166                                 PatternRewriter &rewriter) const override {
167     auto witness = op.witness().getDefiningOp<ConstWitnessOp>();
168     if (!witness || !witness.passingAttr())
169       return failure();
170 
171     AssumingOp::inlineRegionIntoParent(op, rewriter);
172     return success();
173   }
174 };
175 } // namespace
176 
177 void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
178                                              MLIRContext *context) {
179   // If taking a passing witness, inline region.
180   patterns.insert<AssumingWithTrue>(context);
181 }
182 
183 void AssumingOp::inlineRegionIntoParent(AssumingOp &op,
184                                         PatternRewriter &rewriter) {
185   auto *blockBeforeAssuming = rewriter.getInsertionBlock();
186   auto *assumingBlock = op.getBody();
187   auto initPosition = rewriter.getInsertionPoint();
188   auto *blockAfterAssuming =
189       rewriter.splitBlock(blockBeforeAssuming, initPosition);
190 
191   // Remove the AssumingOp and AssumingYieldOp.
192   auto &yieldOp = assumingBlock->back();
193   rewriter.inlineRegionBefore(op.doRegion(), blockAfterAssuming);
194   rewriter.replaceOp(op, yieldOp.getOperands());
195   rewriter.eraseOp(&yieldOp);
196 
197   // Merge blocks together as there was no branching behavior from the
198   // AssumingOp.
199   rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
200   rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
201 }
202 
203 //===----------------------------------------------------------------------===//
204 // AssumingAllOp
205 //===----------------------------------------------------------------------===//
206 OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
207   // Iterate in reverse to first handle all constant operands. They are
208   // guaranteed to be the tail of the inputs because this is commutative.
209   for (int idx = operands.size() - 1; idx >= 0; idx--) {
210     Attribute a = operands[idx];
211     // Cannot fold if any inputs are not constant;
212     if (!a)
213       return nullptr;
214 
215     // We do not need to keep statically known values after handling them in
216     // this method.
217     getOperation()->eraseOperand(idx);
218 
219     // Always false if any input is statically known false
220     if (!a.cast<BoolAttr>().getValue())
221       return a;
222   }
223   // If this is reached, all inputs were statically known passing.
224   return BoolAttr::get(true, getContext());
225 }
226 
227 static LogicalResult verify(AssumingAllOp op) {
228   // Ensure that AssumingAllOp contains at least one operand
229   if (op.getNumOperands() == 0)
230     return op.emitOpError("no operands specified");
231 
232   return success();
233 }
234 
235 //===----------------------------------------------------------------------===//
236 // BroadcastOp
237 //===----------------------------------------------------------------------===//
238 
239 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
240   if (!operands[0] || !operands[1])
241     return nullptr;
242   auto lhsShape = llvm::to_vector<6>(
243       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
244   auto rhsShape = llvm::to_vector<6>(
245       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
246   SmallVector<int64_t, 6> resultShape;
247   // If the shapes are not compatible, we can't fold it.
248   // TODO: Fold to an "error".
249   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
250     return nullptr;
251   Builder builder(getContext());
252   return builder.getIndexTensorAttr(resultShape);
253 }
254 
255 //===----------------------------------------------------------------------===//
256 // ConcatOp
257 //===----------------------------------------------------------------------===//
258 
259 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
260   if (!operands[0] || !operands[1])
261     return nullptr;
262   auto lhsShape = llvm::to_vector<6>(
263       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
264   auto rhsShape = llvm::to_vector<6>(
265       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
266   SmallVector<int64_t, 6> resultShape;
267   resultShape.append(lhsShape.begin(), lhsShape.end());
268   resultShape.append(rhsShape.begin(), rhsShape.end());
269   Builder builder(getContext());
270   return builder.getIndexTensorAttr(resultShape);
271 }
272 
273 //===----------------------------------------------------------------------===//
274 // ConstShapeOp
275 //===----------------------------------------------------------------------===//
276 
277 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
278   p << "shape.const_shape ";
279   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
280   p << "[";
281   interleaveComma(op.shape().getValues<int64_t>(), p,
282                   [&](int64_t i) { p << i; });
283   p << "]";
284 }
285 
286 static ParseResult parseConstShapeOp(OpAsmParser &parser,
287                                      OperationState &result) {
288   if (parser.parseOptionalAttrDict(result.attributes))
289     return failure();
290   // We piggy-back on ArrayAttr parsing, though we don't internally store the
291   // shape as an ArrayAttr.
292   // TODO: Implement custom parser and maybe make syntax a bit more concise.
293   Attribute extentsRaw;
294   NamedAttrList dummy;
295   if (parser.parseAttribute(extentsRaw, "dummy", dummy))
296     return failure();
297   auto extentsArray = extentsRaw.dyn_cast<ArrayAttr>();
298   if (!extentsArray)
299     return failure();
300   SmallVector<int64_t, 6> ints;
301   for (Attribute extent : extentsArray) {
302     IntegerAttr attr = extent.dyn_cast<IntegerAttr>();
303     if (!attr)
304       return failure();
305     ints.push_back(attr.getInt());
306   }
307   Builder &builder = parser.getBuilder();
308   result.addAttribute("shape", builder.getIndexTensorAttr(ints));
309 
310   result.types.push_back(ShapeType::get(builder.getContext()));
311   return success();
312 }
313 
314 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
315 
316 //===----------------------------------------------------------------------===//
317 // CstrBroadcastableOp
318 //===----------------------------------------------------------------------===//
319 
320 void CstrBroadcastableOp::getCanonicalizationPatterns(
321     OwningRewritePatternList &patterns, MLIRContext *context) {
322   // If inputs are equal, return passing witness
323   patterns.insert<CstrBroadcastableEqOps>(context);
324 }
325 
326 OpFoldResult CstrBroadcastableOp::fold(ArrayRef<Attribute> operands) {
327   if (!operands[0] || !operands[1])
328     return nullptr;
329   auto lhsShape = llvm::to_vector<6>(
330       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
331   auto rhsShape = llvm::to_vector<6>(
332       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
333   SmallVector<int64_t, 6> resultShape;
334   if (OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
335     return BoolAttr::get(true, getContext());
336 
337   // Because a failing witness result here represents an eventual assertion
338   // failure, we do not replace it with a constant witness.
339   return nullptr;
340 }
341 
342 //===----------------------------------------------------------------------===//
343 // CstrEqOp
344 //===----------------------------------------------------------------------===//
345 
346 void CstrEqOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
347                                            MLIRContext *context) {
348   // If inputs are equal, return passing witness
349   patterns.insert<CstrEqEqOps>(context);
350 }
351 
352 OpFoldResult CstrEqOp::fold(ArrayRef<Attribute> operands) {
353   if (llvm::all_of(operands,
354                    [&](Attribute a) { return a && a == operands[0]; }))
355     return BoolAttr::get(true, getContext());
356 
357   // Because a failing witness result here represents an eventual assertion
358   // failure, we do not try to replace it with a constant witness. Similarly, we
359   // cannot if there are any non-const inputs.
360   return nullptr;
361 }
362 
363 //===----------------------------------------------------------------------===//
364 // ConstSizeOp
365 //===----------------------------------------------------------------------===//
366 
367 void ConstSizeOp::build(OpBuilder &builder, OperationState &result,
368                         int64_t value) {
369   build(builder, result, builder.getIndexAttr(value));
370 }
371 
372 OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
373 
374 void ConstSizeOp::getAsmResultNames(
375     llvm::function_ref<void(Value, StringRef)> setNameFn) {
376   SmallString<4> buffer;
377   llvm::raw_svector_ostream os(buffer);
378   os << "c" << value();
379   setNameFn(getResult(), os.str());
380 }
381 
382 //===----------------------------------------------------------------------===//
383 // ConstWitnessOp
384 //===----------------------------------------------------------------------===//
385 
386 OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
387 
388 //===----------------------------------------------------------------------===//
389 // IndexToSizeOp
390 //===----------------------------------------------------------------------===//
391 
392 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
393   // Constant values of both types, `shape.size` and `index`, are represented as
394   // `IntegerAttr`s which makes constant folding simple.
395   if (Attribute arg = operands[0])
396     return arg;
397   return {};
398 }
399 
400 void IndexToSizeOp::getCanonicalizationPatterns(
401     OwningRewritePatternList &patterns, MLIRContext *context) {
402   patterns.insert<SizeToIndexToSizeCanonicalization>(context);
403 }
404 
405 //===----------------------------------------------------------------------===//
406 // FromExtentsOp
407 //===----------------------------------------------------------------------===//
408 
409 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
410   if (llvm::any_of(operands, [](Attribute a) { return !a; }))
411     return nullptr;
412   SmallVector<int64_t, 6> extents;
413   for (auto attr : operands)
414     extents.push_back(attr.cast<IntegerAttr>().getInt());
415   Builder builder(getContext());
416   return builder.getIndexTensorAttr(extents);
417 }
418 
419 //===----------------------------------------------------------------------===//
420 // GetExtentOp
421 //===----------------------------------------------------------------------===//
422 
423 Optional<int64_t> GetExtentOp::getConstantDim() {
424   if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>()) {
425     return constSizeOp.value().getLimitedValue();
426   }
427   return llvm::None;
428 }
429 
430 OpFoldResult GetExtentOp::fold(ArrayRef<Attribute> operands) {
431   auto elements = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
432   if (!elements)
433     return nullptr;
434   Optional<int64_t> dim = getConstantDim();
435   if (!dim.hasValue())
436     return nullptr;
437   if (dim.getValue() >= elements.getNumElements())
438     return nullptr;
439   return elements.getValue({(uint64_t)dim.getValue()});
440 }
441 
442 void GetExtentOp::build(OpBuilder &builder, OperationState &result, Value shape,
443                         int64_t dim) {
444   auto loc = result.location;
445   auto dimAttr = builder.getIndexAttr(dim);
446   Value dimValue = builder.create<ConstSizeOp>(loc, dimAttr);
447   build(builder, result, shape, dimValue);
448 }
449 
450 //===----------------------------------------------------------------------===//
451 // RankOp
452 //===----------------------------------------------------------------------===//
453 
454 OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
455   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
456   if (!shape)
457     return {};
458   int64_t rank = shape.getNumElements();
459   Builder builder(getContext());
460   return builder.getIndexAttr(rank);
461 }
462 
463 /// Evaluate the `rank` operation for shapes of ranked tensors at compile time.
464 /// Constant folding fails in cases where only the rank is constant, not the
465 /// shape itself.
466 /// This canonicalization matches `shape.rank(shape.shape_of(%ranked_tensor))`.
467 ///
468 /// Example:
469 ///
470 /// %shape = shape.shape_of %ranked_tensor : tensor<1x2x?xf32>
471 /// %rank = shape.rank %shape
472 ///
473 /// becomes
474 ///
475 /// %rank = shape.const_size 3
476 
477 namespace {
478 struct RankShapeOfCanonicalizationPattern : public OpRewritePattern<RankOp> {
479   using OpRewritePattern<RankOp>::OpRewritePattern;
480 
481   LogicalResult matchAndRewrite(RankOp op,
482                                 PatternRewriter &rewriter) const override {
483     auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>();
484     if (!shapeOfOp)
485       return failure();
486     auto rankedTensorType =
487         shapeOfOp.arg().getType().dyn_cast<RankedTensorType>();
488     if (!rankedTensorType)
489       return failure();
490     int64_t rank = rankedTensorType.getRank();
491     rewriter.replaceOpWithNewOp<ConstSizeOp>(op.getOperation(), rank);
492     return success();
493   }
494 };
495 } // namespace
496 
497 void RankOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
498                                          MLIRContext *context) {
499   patterns.insert<RankShapeOfCanonicalizationPattern>(context);
500 }
501 
502 //===----------------------------------------------------------------------===//
503 // NumElementsOp
504 //===----------------------------------------------------------------------===//
505 
506 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
507 
508   // Fold only when argument constant.
509   Attribute shape = operands[0];
510   if (!shape)
511     return {};
512 
513   APInt product(64, 1);
514   for (auto value : shape.cast<DenseIntElementsAttr>())
515     product *= value;
516   Builder builder(getContext());
517   return builder.getIndexAttr(product.getLimitedValue());
518 }
519 
520 //===----------------------------------------------------------------------===//
521 // ShapeOfOp
522 //===----------------------------------------------------------------------===//
523 
524 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
525   auto type = getOperand().getType().dyn_cast<ShapedType>();
526   if (!type || !type.hasStaticShape())
527     return nullptr;
528   Builder builder(getContext());
529   return builder.getIndexTensorAttr(type.getShape());
530 }
531 
532 //===----------------------------------------------------------------------===//
533 // SizeToIndexOp
534 //===----------------------------------------------------------------------===//
535 
536 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
537   // Constant values of both types, `shape.size` and `index`, are represented as
538   // `IntegerAttr`s which makes constant folding simple.
539   if (Attribute arg = operands[0])
540     return arg;
541   return {};
542 }
543 
544 void SizeToIndexOp::getCanonicalizationPatterns(
545     OwningRewritePatternList &patterns, MLIRContext *context) {
546   patterns.insert<IndexToSizeToIndexCanonicalization>(context);
547 }
548 
549 //===----------------------------------------------------------------------===//
550 // YieldOp
551 //===----------------------------------------------------------------------===//
552 
553 static LogicalResult verify(YieldOp op) {
554   auto *parentOp = op.getParentOp();
555   auto results = parentOp->getResults();
556   auto operands = op.getOperands();
557 
558   if (parentOp->getNumResults() != op.getNumOperands())
559     return op.emitOpError() << "number of operands does not match number of "
560                                "results of its parent";
561   for (auto e : llvm::zip(results, operands))
562     if (std::get<0>(e).getType() != std::get<1>(e).getType())
563       return op.emitOpError()
564              << "types mismatch between yield op and its parent";
565 
566   return success();
567 }
568 
569 //===----------------------------------------------------------------------===//
570 // SplitAtOp
571 //===----------------------------------------------------------------------===//
572 
573 LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
574                               SmallVectorImpl<OpFoldResult> &results) {
575   if (!operands[0] || !operands[1])
576     return failure();
577   auto shapeVec = llvm::to_vector<6>(
578       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
579   auto shape = llvm::makeArrayRef(shapeVec);
580   auto splitPoint = operands[1].cast<IntegerAttr>().getInt();
581   // Verify that the split point is in the correct range.
582   // TODO: Constant fold to an "error".
583   int64_t rank = shape.size();
584   if (!(-rank <= splitPoint && splitPoint <= rank))
585     return failure();
586   if (splitPoint < 0)
587     splitPoint += shape.size();
588   Builder builder(operands[0].getContext());
589   results.push_back(builder.getIndexTensorAttr(shape.take_front(splitPoint)));
590   results.push_back(builder.getIndexTensorAttr(shape.drop_front(splitPoint)));
591   return success();
592 }
593 
594 //===----------------------------------------------------------------------===//
595 // ToExtentTensorOp
596 //===----------------------------------------------------------------------===//
597 
598 OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
599   if (!operands[0])
600     return nullptr;
601   Builder builder(getContext());
602   auto shape = llvm::to_vector<6>(
603       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
604   auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
605                                     builder.getIndexType());
606   return DenseIntElementsAttr::get(type, shape);
607 }
608 
609 //===----------------------------------------------------------------------===//
610 // ReduceOp
611 //===----------------------------------------------------------------------===//
612 
613 void ReduceOp::build(OpBuilder &builder, OperationState &result, Value shape,
614                      ValueRange initVals) {
615   result.addOperands(shape);
616   result.addOperands(initVals);
617 
618   Region *bodyRegion = result.addRegion();
619   bodyRegion->push_back(new Block);
620   Block &bodyBlock = bodyRegion->front();
621   bodyBlock.addArgument(builder.getIndexType());
622   bodyBlock.addArgument(SizeType::get(builder.getContext()));
623 
624   for (Type initValType : initVals.getTypes()) {
625     bodyBlock.addArgument(initValType);
626     result.addTypes(initValType);
627   }
628 }
629 
630 static LogicalResult verify(ReduceOp op) {
631   // Verify block arg types.
632   Block &block = op.region().front();
633 
634   auto blockArgsCount = op.initVals().size() + 2;
635   if (block.getNumArguments() != blockArgsCount)
636     return op.emitOpError() << "ReduceOp body is expected to have "
637                             << blockArgsCount << " arguments";
638 
639   if (block.getArgument(0).getType() != IndexType::get(op.getContext()))
640     return op.emitOpError(
641         "argument 0 of ReduceOp body is expected to be of IndexType");
642 
643   if (block.getArgument(1).getType() != SizeType::get(op.getContext()))
644     return op.emitOpError(
645         "argument 1 of ReduceOp body is expected to be of SizeType");
646 
647   for (auto type : llvm::enumerate(op.initVals()))
648     if (block.getArgument(type.index() + 2).getType() != type.value().getType())
649       return op.emitOpError()
650              << "type mismatch between argument " << type.index() + 2
651              << " of ReduceOp body and initial value " << type.index();
652   return success();
653 }
654 
655 static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) {
656   auto *ctx = parser.getBuilder().getContext();
657   // Parse operands.
658   SmallVector<OpAsmParser::OperandType, 3> operands;
659   if (parser.parseOperandList(operands, /*requiredOperandCount=*/-1,
660                               OpAsmParser::Delimiter::Paren) ||
661       parser.parseOptionalArrowTypeList(result.types))
662     return failure();
663 
664   // Resolve operands.
665   auto initVals = llvm::makeArrayRef(operands).drop_front();
666   if (parser.resolveOperand(operands.front(), ShapeType::get(ctx),
667                             result.operands) ||
668       parser.resolveOperands(initVals, result.types, parser.getNameLoc(),
669                              result.operands))
670     return failure();
671 
672   // Parse the body.
673   Region *body = result.addRegion();
674   if (parser.parseRegion(*body, /*args=*/{}, /*argTypes=*/{}))
675     return failure();
676 
677   // Parse attributes.
678   if (parser.parseOptionalAttrDict(result.attributes))
679     return failure();
680 
681   return success();
682 }
683 
684 static void print(OpAsmPrinter &p, ReduceOp op) {
685   p << op.getOperationName() << '(' << op.shape() << ", " << op.initVals()
686     << ") ";
687   p.printOptionalArrowTypeList(op.getResultTypes());
688   p.printRegion(op.region());
689   p.printOptionalAttrDict(op.getAttrs());
690 }
691 
692 namespace mlir {
693 namespace shape {
694 
695 #define GET_OP_CLASSES
696 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
697 
698 } // namespace shape
699 } // namespace mlir
700