let results = (outs Shape_SizeType:$result);
}
+def Shape_NumElementsOp : Shape_Op<"num_elements", [
+ NoSideEffect,
+ DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+
+ let summary = "Returns the number of elements for a given shape";
+ let description = [{
+ Returns the number of elements for a given shape which is the product of its
+ dimensions.
+ }];
+
+ let arguments = (ins Shape_ShapeType:$shape);
+ let results = (outs Shape_SizeType:$result);
+
+ let assemblyFormat = "attr-dict $shape";
+
+ let hasFolder = 1;
+}
+
def Shape_ReduceOp : Shape_Op<"reduce", []> {
let summary = "Returns an expression reduced over a shape";
let description = [{
}
//===----------------------------------------------------------------------===//
+// NumElementsOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {
+
+ // Fold only when argument constant.
+ Attribute shape = operands[0];
+ if (!shape)
+ return {};
+
+ APInt product(64, 1);
+ for (auto value : shape.cast<DenseIntElementsAttr>())
+ product *= value;
+ Builder builder(getContext());
+ return builder.getIndexAttr(product.getLimitedValue());
+}
+
+LogicalResult NumElementsOp::inferReturnTypes(
+ MLIRContext *context, Optional<Location> location, ValueRange operands,
+ DictionaryAttr attributes, RegionRange regions,
+ SmallVectorImpl<Type> &inferredReturnTypes) {
+ inferredReturnTypes.push_back(SizeType::get(context));
+ return success();
+}
+
+//===----------------------------------------------------------------------===//
// ShapeOfOp
//===----------------------------------------------------------------------===//
}
// -----
+// Fold number of elements computation.
+// CHECK-LABEL: func @num_elements
+func @num_elements() -> !shape.size {
+ // CHECK-NOT: shape.const_shape
+ %shape = shape.const_shape [4, 5, 6]
+ // CHECK-NOT: shape.num_elements
+ %num_elements = shape.num_elements %shape
+ // CHECK: %[[NUM:.*]] = shape.const_size 120
+ // CHECK-NEXT: return %[[NUM]] : !shape.size
+ return %num_elements : !shape.size
+}
+
+// -----
+// No folding.
+// CHECK-LABEL: func @nonfoldable_num_elements
+func @nonfoldable_num_elements(%shape : !shape.shape) -> !shape.size {
+ // CHECK-NOT: shape.const_{{.*}}
+ %num_elements = shape.num_elements %shape
+ return %num_elements : !shape.size
+}
+
+// -----
// Canonicalization of shape.get_extent