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   // Allow unknown operations during prototyping and testing. As the dialect is
29   // still evolving it makes it simple to start with an unregistered ops and
30   // try different variants before actually defining the op.
31   allowUnknownOperations();
32 }
33 
34 Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
35                                              Attribute value, Type type,
36                                              Location loc) {
37   if (auto shapeType = type.dyn_cast<ShapeType>()) {
38     return builder.create<ConstShapeOp>(loc, type,
39                                         value.cast<DenseIntElementsAttr>());
40   }
41   if (auto sizeType = type.dyn_cast<SizeType>()) {
42     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
43   }
44   return nullptr;
45 }
46 
47 /// Parse a type registered to this dialect.
48 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
49   StringRef keyword;
50   if (parser.parseKeyword(&keyword))
51     return Type();
52 
53   if (keyword == "component")
54     return ComponentType::get(getContext());
55   if (keyword == "element")
56     return ElementType::get(getContext());
57   if (keyword == "shape")
58     return ShapeType::get(getContext());
59   if (keyword == "size")
60     return SizeType::get(getContext());
61   if (keyword == "value_shape")
62     return ValueShapeType::get(getContext());
63 
64   parser.emitError(parser.getNameLoc(), "unknown shape type: ") << keyword;
65   return Type();
66 }
67 
68 /// Print a type registered to this dialect.
69 void ShapeDialect::printType(Type type, DialectAsmPrinter &os) const {
70   switch (type.getKind()) {
71   case ShapeTypes::Component:
72     os << "component";
73     return;
74   case ShapeTypes::Element:
75     os << "element";
76     return;
77   case ShapeTypes::Size:
78     os << "size";
79     return;
80   case ShapeTypes::Shape:
81     os << "shape";
82     return;
83   case ShapeTypes::ValueShape:
84     os << "value_shape";
85     return;
86   default:
87     llvm_unreachable("unexpected 'shape' type kind");
88   }
89 }
90 
91 //===----------------------------------------------------------------------===//
92 // BroadcastOp
93 //===----------------------------------------------------------------------===//
94 
95 LogicalResult BroadcastOp::inferReturnTypes(
96     MLIRContext *context, Optional<Location> location, ValueRange operands,
97     ArrayRef<NamedAttribute> attributes, RegionRange regions,
98     SmallVectorImpl<Type> &inferredReturnTypes) {
99   inferredReturnTypes.push_back(ShapeType::get(context));
100   return success();
101 }
102 
103 OpFoldResult BroadcastOp::fold(ArrayRef<Attribute> operands) {
104   if (!operands[0] || !operands[1])
105     return nullptr;
106   auto lhsShape = llvm::to_vector<6>(
107       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
108   auto rhsShape = llvm::to_vector<6>(
109       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
110   SmallVector<int64_t, 6> resultShape;
111   // If the shapes are not compatible, we can't fold it.
112   // TODO: Fold to an "error".
113   if (!OpTrait::util::getBroadcastedShape(lhsShape, rhsShape, resultShape))
114     return nullptr;
115   Builder builder(getContext());
116   return builder.getI64TensorAttr(resultShape);
117 }
118 
119 //===----------------------------------------------------------------------===//
120 // ConstShapeOp
121 //===----------------------------------------------------------------------===//
122 
123 static void print(OpAsmPrinter &p, ConstShapeOp &op) {
124   p << "shape.const_shape ";
125   p.printOptionalAttrDict(op.getAttrs(), /*elidedAttrs=*/{"shape"});
126   p << "[";
127   interleaveComma(op.shape().getValues<int64_t>(), p,
128                   [&](int64_t i) { p << i; });
129   p << "]";
130 }
131 
132 static ParseResult parseConstShapeOp(OpAsmParser &parser,
133                                      OperationState &result) {
134   if (parser.parseOptionalAttrDict(result.attributes))
135     return failure();
136   // We piggy-back on ArrayAttr parsing, though we don't internally store the
137   // shape as an ArrayAttr.
138   // TODO: Implement custom parser and maybe make syntax a bit more concise.
139   Attribute extentsRaw;
140   SmallVector<NamedAttribute, 6> dummy;
141   if (parser.parseAttribute(extentsRaw, "dummy", dummy))
142     return failure();
143   auto extentsArray = extentsRaw.dyn_cast<ArrayAttr>();
144   if (!extentsArray)
145     return failure();
146   SmallVector<int64_t, 6> ints;
147   for (Attribute extent : extentsArray) {
148     IntegerAttr attr = extent.dyn_cast<IntegerAttr>();
149     if (!attr)
150       return failure();
151     ints.push_back(attr.getInt());
152   }
153   Builder &builder = parser.getBuilder();
154   result.addAttribute("shape", builder.getI64TensorAttr(ints));
155 
156   result.types.push_back(ShapeType::get(builder.getContext()));
157   return success();
158 }
159 
160 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shape(); }
161 
162 LogicalResult ConstShapeOp::inferReturnTypes(
163     MLIRContext *context, Optional<Location> location, ValueRange operands,
164     ArrayRef<NamedAttribute> attributes, RegionRange regions,
165     SmallVectorImpl<Type> &inferredReturnTypes) {
166   inferredReturnTypes.push_back(ShapeType::get(context));
167   return success();
168 }
169 
170 //===----------------------------------------------------------------------===//
171 // ConstSizeOp
172 //===----------------------------------------------------------------------===//
173 
174 LogicalResult ConstSizeOp::inferReturnTypes(
175     MLIRContext *context, Optional<Location> location, ValueRange operands,
176     ArrayRef<NamedAttribute> attributes, RegionRange regions,
177     SmallVectorImpl<Type> &inferredReturnTypes) {
178   inferredReturnTypes.push_back(SizeType::get(context));
179   return success();
180 }
181 
182 //===----------------------------------------------------------------------===//
183 // ShapeOfOp
184 //===----------------------------------------------------------------------===//
185 
186 LogicalResult ShapeOfOp::inferReturnTypes(
187     MLIRContext *context, Optional<Location> location, ValueRange operands,
188     ArrayRef<NamedAttribute> attributes, RegionRange regions,
189     SmallVectorImpl<Type> &inferredReturnTypes) {
190   inferredReturnTypes.push_back(ShapeType::get(context));
191   return success();
192 }
193 
194 OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
195   auto type = getOperand().getType().dyn_cast<ShapedType>();
196   if (!type || !type.hasStaticShape())
197     return nullptr;
198   Builder builder(getContext());
199   return builder.getI64TensorAttr(type.getShape());
200 }
201 
202 //===----------------------------------------------------------------------===//
203 // SplitAtOp
204 //===----------------------------------------------------------------------===//
205 
206 LogicalResult SplitAtOp::inferReturnTypes(
207     MLIRContext *context, Optional<Location> location, ValueRange operands,
208     ArrayRef<NamedAttribute> attributes, RegionRange regions,
209     SmallVectorImpl<Type> &inferredReturnTypes) {
210   auto shapeType = ShapeType::get(context);
211   inferredReturnTypes.push_back(shapeType);
212   inferredReturnTypes.push_back(shapeType);
213   return success();
214 }
215 
216 LogicalResult SplitAtOp::fold(ArrayRef<Attribute> operands,
217                               SmallVectorImpl<OpFoldResult> &results) {
218   if (!operands[0] || !operands[1])
219     return failure();
220   auto shapeVec = llvm::to_vector<6>(
221       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
222   auto shape = llvm::makeArrayRef(shapeVec);
223   auto splitPoint = operands[1].cast<IntegerAttr>().getInt();
224   // Verify that the split point is in the correct range.
225   // TODO: Constant fold to an "error".
226   int64_t rank = shape.size();
227   if (!(-rank <= splitPoint && splitPoint <= rank))
228     return failure();
229   if (splitPoint < 0)
230     splitPoint += shape.size();
231   Builder builder(operands[0].getContext());
232   results.push_back(builder.getI64TensorAttr(shape.take_front(splitPoint)));
233   results.push_back(builder.getI64TensorAttr(shape.drop_front(splitPoint)));
234   return success();
235 }
236 
237 //===----------------------------------------------------------------------===//
238 // ConcatOp
239 //===----------------------------------------------------------------------===//
240 
241 LogicalResult ConcatOp::inferReturnTypes(
242     MLIRContext *context, Optional<Location> location, ValueRange operands,
243     ArrayRef<NamedAttribute> attributes, RegionRange regions,
244     SmallVectorImpl<Type> &inferredReturnTypes) {
245   auto shapeType = ShapeType::get(context);
246   inferredReturnTypes.push_back(shapeType);
247   return success();
248 }
249 
250 OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
251   if (!operands[0] || !operands[1])
252     return nullptr;
253   auto lhsShape = llvm::to_vector<6>(
254       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
255   auto rhsShape = llvm::to_vector<6>(
256       operands[1].cast<DenseIntElementsAttr>().getValues<int64_t>());
257   SmallVector<int64_t, 6> resultShape;
258   resultShape.append(lhsShape.begin(), lhsShape.end());
259   resultShape.append(rhsShape.begin(), rhsShape.end());
260   Builder builder(getContext());
261   return builder.getI64TensorAttr(resultShape);
262 }
263 
264 //===----------------------------------------------------------------------===//
265 // ToExtentTensorOp
266 //===----------------------------------------------------------------------===//
267 
268 OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
269   if (!operands[0])
270     return nullptr;
271   Builder builder(getContext());
272   auto shape = llvm::to_vector<6>(
273       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
274   auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
275                                     builder.getIndexType());
276   return DenseIntElementsAttr::get(type, shape);
277 }
278 
279 namespace mlir {
280 namespace shape {
281 
282 #define GET_OP_CLASSES
283 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
284 
285 } // namespace shape
286 } // namespace mlir
287