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