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