[PATCH] D80281: [MLIR] Add `num_elements` to the shape dialect
Frederik Gossen via Phabricator via llvm-commits
llvm-commits at lists.llvm.org
Wed May 20 02:40:06 PDT 2020
frgossen created this revision.
Herald added subscribers: llvm-commits, jurahul, Kayjukh, grosul1, Joonsoo, stephenneuendorffer, liufengdb, aartbik, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: jpienaar.
Herald added a reviewer: silvas.
Herald added a project: LLVM.
frgossen added a child revision: D80283: [MLIR] Fix operand type in `from_extent_tensor` in the shape dialect.
frgossen added reviewers: pifon2a, herhut.
The operation `num_elements` determines the number of elements for a given
shape.
That is the product of its dimensions.
A tensor of the given shape holds this many elements.
Depends On D80280 <https://reviews.llvm.org/D80280>
Repository:
rG LLVM Github Monorepo
https://reviews.llvm.org/D80281
Files:
mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
mlir/lib/Dialect/Shape/IR/Shape.cpp
mlir/test/Dialect/Shape/canonicalize.mlir
Index: mlir/test/Dialect/Shape/canonicalize.mlir
===================================================================
--- mlir/test/Dialect/Shape/canonicalize.mlir
+++ mlir/test/Dialect/Shape/canonicalize.mlir
@@ -139,3 +139,25 @@
%cs = shape.index_to_size %ci
return %cs : !shape.size
}
+
+// -----
+// 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
+}
Index: mlir/lib/Dialect/Shape/IR/Shape.cpp
===================================================================
--- mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -252,6 +252,32 @@
return success();
}
+//===----------------------------------------------------------------------===//
+// 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
//===----------------------------------------------------------------------===//
Index: mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
===================================================================
--- mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -231,6 +231,26 @@
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.
+ A tensor of the given shape will hold this many elements.
+ Expects a `shape.shape` and returns a `shape.size` value.
+ }];
+
+ 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 = [{
-------------- next part --------------
A non-text attachment was scrubbed...
Name: D80281.265172.patch
Type: text/x-patch
Size: 3339 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200520/18ed41ac/attachment-0001.bin>
More information about the llvm-commits
mailing list