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     auto *blockBeforeAssuming = rewriter.getInsertionBlock();
172     auto *assumingBlock = op.getBody();
173     auto initPosition = rewriter.getInsertionPoint();
174     auto *blockAfterAssuming =
175         rewriter.splitBlock(blockBeforeAssuming, initPosition);
176 
177     // Remove the AssumingOp and AssumingYieldOp.
178     auto &yieldOp = assumingBlock->back();
179     rewriter.inlineRegionBefore(op.doRegion(), blockAfterAssuming);
180     rewriter.replaceOp(op, yieldOp.getOperands());
181     rewriter.eraseOp(&yieldOp);
182 
183     // Merge blocks together as there was no branching behavior from the
184     // AssumingOp.
185     rewriter.mergeBlocks(assumingBlock, blockBeforeAssuming);
186     rewriter.mergeBlocks(blockAfterAssuming, blockBeforeAssuming);
187     return success();
188   }
189 };
190 } // namespace
191 
192 void AssumingOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
193                                              MLIRContext *context) {
194   // If taking a passing witness, inline region
195   patterns.insert<AssumingWithTrue>(context);
196 }
197 
198 //===----------------------------------------------------------------------===//
199 // AssumingAllOp
200 //===----------------------------------------------------------------------===//
201 OpFoldResult AssumingAllOp::fold(ArrayRef<Attribute> operands) {
202   // Iterate in reverse to first handle all constant operands. They are
203   // guaranteed to be the tail of the inputs because this is commutative.
204   for (int idx = operands.size() - 1; idx >= 0; idx--) {
205     Attribute a = operands[idx];
206     // Cannot fold if any inputs are not constant;
207     if (!a)
208       return nullptr;
209 
210     // We do not need to keep statically known values after handling them in
211     // this method.
212     getOperation()->eraseOperand(idx);
213 
214     // Always false if any input is statically known false
215     if (!a.cast<BoolAttr>().getValue())
216       return a;
217   }
218   // If this is reached, all inputs were statically known passing.
219   return BoolAttr::get(true, getContext());
220 }
221 
222 static LogicalResult verify(AssumingAllOp op) {
223   // Ensure that AssumingAllOp contains at least one operand
224   if (op.getNumOperands() == 0)
225     return op.emitOpError("no operands specified");
226 
227   return success();
228 }
229 
230 //===----------------------------------------------------------------------===//
231 // BroadcastOp
232 //===----------------------------------------------------------------------===//
233 
234 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
235   if (!operands[0] || !operands[1])
236     return nullptr;
237   auto lhsShape = llvm::to_vector<6>(
238       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
239   auto rhsShape = llvm::to_vector<6>(
240       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
241   SmallVector<int64_t, 6> resultShape;
242   // If the shapes are not compatible, we can't fold it.
243   // TODO: Fold to an "error".
244   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
245     return nullptr;
246   Builder builder(getContext());
247   return builder.getIndexTensorAttr(resultShape);
248 }
249 
250 //===----------------------------------------------------------------------===//
251 // ConcatOp
252 //===----------------------------------------------------------------------===//
253 
254 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
255   if (!operands[0] || !operands[1])
256     return nullptr;
257   auto lhsShape = llvm::to_vector<6>(
258       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
259   auto rhsShape = llvm::to_vector<6>(
260       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
261   SmallVector<int64_t, 6> resultShape;
262   resultShape.append(lhsShape.begin(), lhsShape.end());
263   resultShape.append(rhsShape.begin(), rhsShape.end());
264   Builder builder(getContext());
265   return builder.getIndexTensorAttr(resultShape);
266 }
267 
268 //===----------------------------------------------------------------------===//
269 // ConstShapeOp
270 //===----------------------------------------------------------------------===//
271 
272 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
273   p << "shape.const_shape ";
274   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
275   p << "[";
276   interleaveComma(op.shape().getValues<int64_t>(), p,
277                   [&](int64_t i) { p << i; });
278   p << "]";
279 }
280 
281 static ParseResult parseConstShapeOp(OpAsmParser &parser,
282                                      OperationState &result) {
283   if (parser.parseOptionalAttrDict(result.attributes))
284     return failure();
285   // We piggy-back on ArrayAttr parsing, though we don't internally store the
286   // shape as an ArrayAttr.
287   // TODO: Implement custom parser and maybe make syntax a bit more concise.
288   Attribute extentsRaw;
289   NamedAttrList dummy;
290   if (parser.parseAttribute(extentsRaw, "dummy", dummy))
291     return failure();
292   auto extentsArray = extentsRaw.dyn_cast<ArrayAttr>();
293   if (!extentsArray)
294     return failure();
295   SmallVector<int64_t, 6> ints;
296   for (Attribute extent : extentsArray) {
297     IntegerAttr attr = extent.dyn_cast<IntegerAttr>();
298     if (!attr)
299       return failure();
300     ints.push_back(attr.getInt());
301   }
302   Builder &builder = parser.getBuilder();
303   result.addAttribute("shape", builder.getIndexTensorAttr(ints));
304 
305   result.types.push_back(ShapeType::get(builder.getContext()));
306   return success();
307 }
308 
309 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shapeAttr(); }
310 
311 //===----------------------------------------------------------------------===//
312 // CstrBroadcastableOp
313 //===----------------------------------------------------------------------===//
314 
315 void CstrBroadcastableOp::getCanonicalizationPatterns(
316     OwningRewritePatternList &patterns, MLIRContext *context) {
317   // If inputs are equal, return passing witness
318   patterns.insert<CstrBroadcastableEqOps>(context);
319 }
320 
321 OpFoldResult CstrBroadcastableOp::fold(ArrayRef<Attribute> operands) {
322   if (!operands[0] || !operands[1])
323     return nullptr;
324   auto lhsShape = llvm::to_vector<6>(
325       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
326   auto rhsShape = llvm::to_vector<6>(
327       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
328   SmallVector<int64_t, 6> resultShape;
329   if (OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
330     return BoolAttr::get(true, getContext());
331 
332   // Because a failing witness result here represents an eventual assertion
333   // failure, we do not replace it with a constant witness.
334   return nullptr;
335 }
336 
337 //===----------------------------------------------------------------------===//
338 // CstrEqOp
339 //===----------------------------------------------------------------------===//
340 
341 void CstrEqOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
342                                            MLIRContext *context) {
343   // If inputs are equal, return passing witness
344   patterns.insert<CstrEqEqOps>(context);
345 }
346 
347 OpFoldResult CstrEqOp::fold(ArrayRef<Attribute> operands) {
348   if (llvm::all_of(operands,
349                    [&](Attribute a) { return a && a == operands[0]; }))
350     return BoolAttr::get(true, getContext());
351 
352   // Because a failing witness result here represents an eventual assertion
353   // failure, we do not try to replace it with a constant witness. Similarly, we
354   // cannot if there are any non-const inputs.
355   return nullptr;
356 }
357 
358 //===----------------------------------------------------------------------===//
359 // ConstSizeOp
360 //===----------------------------------------------------------------------===//
361 
362 OpFoldResult ConstSizeOp::fold(ArrayRef<Attribute>) { return valueAttr(); }
363 
364 void ConstSizeOp::getAsmResultNames(
365     llvm::function_ref<void(Value, StringRef)> setNameFn) {
366   SmallString<4> buffer;
367   llvm::raw_svector_ostream os(buffer);
368   os << "c" << value();
369   setNameFn(getResult(), os.str());
370 }
371 
372 //===----------------------------------------------------------------------===//
373 // ConstWitnessOp
374 //===----------------------------------------------------------------------===//
375 
376 OpFoldResult ConstWitnessOp::fold(ArrayRef<Attribute>) { return passingAttr(); }
377 
378 //===----------------------------------------------------------------------===//
379 // IndexToSizeOp
380 //===----------------------------------------------------------------------===//
381 
382 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
383   // Constant values of both types, `shape.size` and `index`, are represented as
384   // `IntegerAttr`s which makes constant folding simple.
385   if (Attribute arg = operands[0])
386     return arg;
387   return {};
388 }
389 
390 //===----------------------------------------------------------------------===//
391 // FromExtentsOp
392 //===----------------------------------------------------------------------===//
393 
394 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
395   if (llvm::any_of(operands, [](Attribute a) { return !a; }))
396     return nullptr;
397   SmallVector<int64_t, 6> extents;
398   for (auto attr : operands)
399     extents.push_back(attr.cast<IntegerAttr>().getInt());
400   Builder builder(getContext());
401   return builder.getIndexTensorAttr(extents);
402 }
403 
404 //===----------------------------------------------------------------------===//
405 // GetExtentOp
406 //===----------------------------------------------------------------------===//
407 
408 Optional<int64_t> GetExtentOp::getConstantDim() {
409   if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>()) {
410     return constSizeOp.value().getLimitedValue();
411   }
412   return llvm::None;
413 }
414 
415 OpFoldResult GetExtentOp::fold(ArrayRef<Attribute> operands) {
416   auto elements = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
417   if (!elements)
418     return nullptr;
419   Optional<int64_t> dim = getConstantDim();
420   if (!dim.hasValue())
421     return nullptr;
422   if (dim.getValue() >= elements.getNumElements())
423     return nullptr;
424   return elements.getValue({(uint64_t)dim.getValue()});
425 }
426 
427 void GetExtentOp::build(OpBuilder &builder, OperationState &result, Value shape,
428                         int64_t dim) {
429   auto loc = result.location;
430   auto dimAttr = builder.getIndexAttr(dim);
431   Value dimValue = builder.create<ConstSizeOp>(loc, dimAttr);
432   build(builder, result, shape, dimValue);
433 }
434 
435 //===----------------------------------------------------------------------===//
436 // NumElementsOp
437 //===----------------------------------------------------------------------===//
438 
439 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
440 
441   // Fold only when argument constant.
442   Attribute shape = operands[0];
443   if (!shape)
444     return {};
445 
446   APInt product(64, 1);
447   for (auto value : shape.cast<DenseIntElementsAttr>())
448     product *= value;
449   Builder builder(getContext());
450   return builder.getIndexAttr(product.getLimitedValue());
451 }
452 
453 //===----------------------------------------------------------------------===//
454 // ShapeOfOp
455 //===----------------------------------------------------------------------===//
456 
457 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
458   auto type = getOperand().getType().dyn_cast<ShapedType>();
459   if (!type || !type.hasStaticShape())
460     return nullptr;
461   Builder builder(getContext());
462   return builder.getIndexTensorAttr(type.getShape());
463 }
464 
465 //===----------------------------------------------------------------------===//
466 // SizeToIndexOp
467 //===----------------------------------------------------------------------===//
468 
469 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
470   // Constant values of both types, `shape.size` and `index`, are represented as
471   // `IntegerAttr`s which makes constant folding simple.
472   if (Attribute arg = operands[0])
473     return arg;
474   return {};
475 }
476 
477 //===----------------------------------------------------------------------===//
478 // YieldOp
479 //===----------------------------------------------------------------------===//
480 
481 static LogicalResult verify(YieldOp op) {
482   auto *parentOp = op.getParentOp();
483   auto results = parentOp->getResults();
484   auto operands = op.getOperands();
485 
486   if (parentOp->getNumResults() != op.getNumOperands())
487     return op.emitOpError() << "number of operands does not match number of "
488                                "results of its parent";
489   for (auto e : llvm::zip(results, operands))
490     if (std::get<0>(e).getType() != std::get<1>(e).getType())
491       return op.emitOpError()
492              << "types mismatch between yield op and its parent";
493 
494   return success();
495 }
496 
497 //===----------------------------------------------------------------------===//
498 // SplitAtOp
499 //===----------------------------------------------------------------------===//
500 
501 LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
502                               SmallVectorImpl<OpFoldResult> &results) {
503   if (!operands[0] || !operands[1])
504     return failure();
505   auto shapeVec = llvm::to_vector<6>(
506       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
507   auto shape = llvm::makeArrayRef(shapeVec);
508   auto splitPoint = operands[1].cast<IntegerAttr>().getInt();
509   // Verify that the split point is in the correct range.
510   // TODO: Constant fold to an "error".
511   int64_t rank = shape.size();
512   if (!(-rank <= splitPoint && splitPoint <= rank))
513     return failure();
514   if (splitPoint < 0)
515     splitPoint += shape.size();
516   Builder builder(operands[0].getContext());
517   results.push_back(builder.getIndexTensorAttr(shape.take_front(splitPoint)));
518   results.push_back(builder.getIndexTensorAttr(shape.drop_front(splitPoint)));
519   return success();
520 }
521 
522 //===----------------------------------------------------------------------===//
523 // ToExtentTensorOp
524 //===----------------------------------------------------------------------===//
525 
526 OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
527   if (!operands[0])
528     return nullptr;
529   Builder builder(getContext());
530   auto shape = llvm::to_vector<6>(
531       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
532   auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
533                                     builder.getIndexType());
534   return DenseIntElementsAttr::get(type, shape);
535 }
536 
537 //===----------------------------------------------------------------------===//
538 // ReduceOp
539 //===----------------------------------------------------------------------===//
540 
541 void ReduceOp::build(OpBuilder &builder, OperationState &result, Value shape,
542                      ValueRange initVals) {
543   result.addOperands(shape);
544   result.addOperands(initVals);
545 
546   Region *bodyRegion = result.addRegion();
547   bodyRegion->push_back(new Block);
548   Block &bodyBlock = bodyRegion->front();
549   bodyBlock.addArgument(builder.getIndexType());
550   bodyBlock.addArgument(SizeType::get(builder.getContext()));
551 
552   for (Type initValType : initVals.getTypes()) {
553     bodyBlock.addArgument(initValType);
554     result.addTypes(initValType);
555   }
556 }
557 
558 static LogicalResult verify(ReduceOp op) {
559   // Verify block arg types.
560   Block &block = op.region().front();
561 
562   auto blockArgsCount = op.initVals().size() + 2;
563   if (block.getNumArguments() != blockArgsCount)
564     return op.emitOpError() << "ReduceOp body is expected to have "
565                             << blockArgsCount << " arguments";
566 
567   if (block.getArgument(0).getType() != IndexType::get(op.getContext()))
568     return op.emitOpError(
569         "argument 0 of ReduceOp body is expected to be of IndexType");
570 
571   if (block.getArgument(1).getType() != SizeType::get(op.getContext()))
572     return op.emitOpError(
573         "argument 1 of ReduceOp body is expected to be of SizeType");
574 
575   for (auto type : llvm::enumerate(op.initVals()))
576     if (block.getArgument(type.index() + 2).getType() != type.value().getType())
577       return op.emitOpError()
578              << "type mismatch between argument " << type.index() + 2
579              << " of ReduceOp body and initial value " << type.index();
580   return success();
581 }
582 
583 static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) {
584   auto *ctx = parser.getBuilder().getContext();
585   // Parse operands.
586   SmallVector<OpAsmParser::OperandType, 3> operands;
587   if (parser.parseOperandList(operands, /*requiredOperandCount=*/-1,
588                               OpAsmParser::Delimiter::Paren) ||
589       parser.parseOptionalArrowTypeList(result.types))
590     return failure();
591 
592   // Resolve operands.
593   auto initVals = llvm::makeArrayRef(operands).drop_front();
594   if (parser.resolveOperand(operands.front(), ShapeType::get(ctx),
595                             result.operands) ||
596       parser.resolveOperands(initVals, result.types, parser.getNameLoc(),
597                              result.operands))
598     return failure();
599 
600   // Parse the body.
601   Region *body = result.addRegion();
602   if (parser.parseRegion(*body, /*args=*/{}, /*argTypes=*/{}))
603     return failure();
604 
605   // Parse attributes.
606   if (parser.parseOptionalAttrDict(result.attributes))
607     return failure();
608 
609   return success();
610 }
611 
612 static void print(OpAsmPrinter &p, ReduceOp op) {
613   p << op.getOperationName() << '(' << op.shape() << ", " << op.initVals()
614     << ") ";
615   p.printOptionalArrowTypeList(op.getResultTypes());
616   p.printRegion(op.region());
617   p.printOptionalAttrDict(op.getAttrs());
618 }
619 
620 namespace mlir {
621 namespace shape {
622 
623 #define GET_OP_CLASSES
624 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
625 
626 } // namespace shape
627 } // namespace mlir
628