[PATCH] D76817: `shape` dialect: add some ops

[Sean Silva via Phabricator via llvm-commits]

silvas created this revision.
silvas added a reviewer: jpienaar.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, rriddle, mehdi_amini.
Herald added a project: LLVM.

- add `to_extent_tensor`
  - rename `create_shape` to `from_extent_tensor` for symmetry
- add `drop_back`, `take_back`, and `concat` ops for basic shape manipulations

This set of ops is inspired by the requirements of lowering a dynamic-shape-aware batch matmul op. For such an op, the "matrix" dimensions aren't subject to broadcasting but the others are, and so we need to slice, broadcast, and reconstruct the final output shape. Furthermore, the actual broadcasting op used downstream uses a tensor of extents as its preferred shape interface for the actual op that does the broadcasting.

However, this functionality is quite general. It's obvious that `to_extent_tensor` is needed long-term to support many common patterns that involve computations on shapes. We can evolve the shape manipulation ops introduced here. The specific choices made here took into consideration the potentially unranked nature of the !shape.shape type, which means that a simple listing of dimensions to extract isn't possible for an op like `drop_back`.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D76817

Files:
  mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td


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
@@ -168,17 +168,30 @@
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def Shape_CreateShapeOp : Shape_Op<"create_shape", []> {
-  let summary = "Creates a shape descriptor from a tensor";
+def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
+  let summary = "Creates a shape from a tensor of extents";
   let description = [{
-    Creates a shape from a 1D integral tensor. The rank equals the number of
-    elements in the tensor, and extent matches the values of the elements.
+    Creates a shape from a 1D integral tensor of extents. The rank of the
+    resulting shape equals the number of elements in the tensor, and the
+    extents match the values of the elements.
   }];
 
   let arguments = (ins I32Tensor:$input);
   let results = (outs Shape_ShapeType:$result);
 }
 
+def Shape_ToExtentTensorOp : Shape_Op<"to_tensor", []> {
+  let summary = "Creates a dimension tensor from a shape";
+  let description = [{
+    Converts a shape to a 1D integral tensor of extents. The number of elements
+    in the tensor equals the rank of the shape, and the elements equal the
+    extents of the shape.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$input);
+  let results = (outs I32Tensor:$result);
+}
+
 def Shape_JoinOp : Shape_Op<"join", []> {
   let summary = "Returns the least general shape.size of its operands";
   let description = [{
@@ -299,4 +312,60 @@
   let results =  (outs Shape_ShapeOrSizeType:$output);
 }
 
+def Shape_DropBackOp : Shape_Op<"drop_back", []> {
+  let summary = "Drops dimensions off the back of a shape.";
+  let description = [{
+    Drops `amount` dimensions off the back of a shape, resulting in a shape
+    with equal or lesser rank.
+
+    Examples:
+    drop_back([4,5,6], amount=1) -> [4,5]
+    drop_back([4,5,6], amount=3) -> []
+    drop_back([4,5,6], amount=4) -> error
+
+
+    Requires (otherwise result is an error shape):
+    - amount >= 0
+    - amount <= rank(operand)
+  }];
+
+  let arguments = (ins Shape_ShapeType:$operand, I32:$amount);
+  let results = (outs Shape_ShapeType:$result);
+}
+def Shape_TakeBackOp : Shape_Op<"take_back", []> {
+  let summary = "Takes dimensions off the back of a shape.";
+  let description = [{
+    Takes the last `amount` dimensions from `operand`, resulting in a shape
+    with equal or lesser rank.
+
+    Examples:
+    take_back([5,6,7], amount=1) -> [7]
+    take_back([5,6,7], amount=3) -> [5,6,7]
+    take_back([5,6,7], amount=6) -> error
+
+    Requires:
+    - amount >= 0
+    - amount <= rank(operand)
+  }];
+
+  let arguments = (ins Shape_ShapeType:$operand, I32:$amount);
+  let results = (outs Shape_ShapeType:$result);
+}
+
+def Shape_ConcatOp : Shape_Op<"concat", []> {
+  let summary = "Concatenates two shapes.";
+  let description = [{
+    Creates a shape whose dimensions consist of first the dimensions from `lhs`
+    followed by the dimensions of `rhs`.
+
+    Example:
+    concat([2,3], [4,5]) -> [2,3,4,5]
+    concat([], []) -> []
+    concat([], [4,5,6]) -> [4,5,6]
+  }];
+
+  let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs);
+  let results = (outs Shape_ShapeType:$result);
+}
+
 #endif // SHAPE_OPS


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D76817.252721.patch
Type: text/x-patch
Size: 3415 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200326/0b592f78/attachment.bin>