diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td --- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td +++ b/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