`shape` dialect: add some ops

[silvas]

Summary

shape dialect: add some ops

• add to_extent_tensor
  • rename create_shape to from_extent_tensor for symmetry
• add split_at 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 in general.

[jpienaar]

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

316

Instead of drop back and take back, how about an op that mirrors concat (well mirrors the binary concat) like split_at which produces two shapes.

[silvas]

Address comments.

[silvas]

Update description

[jpienaar]

Nice, I think one more round and it is good

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

189

What happens for unranked and or dynamic tensors?

193

With your other rev, can these now be IndexTensor? (can also be done in follow up)

335

Add behavior for unranked here and below

[silvas]

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

189

it return the shape. Notice that the return type is not static shaped. For an unranked tensor it will return a tensor<?xi32> (or index someday).

193

Yeah, I want to do it in a follow-up. Just being consistent for now.

335

Whether something is unranked or not is a purely static concept, so it's totally orthogonal to an execution example like this.

[jpienaar]

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

189

Thus far FromExtentTensorOp has as input a tensor with known values and produces a fixed shape shape (e.g., there is no special numeric values that correspond to unknown). But another problem is that ToExtentTensor now needs to convert from a shape which may be an error to a tensor of ints. The intention is to have all shape ops be side-effect free and I don't see how that can be maintained for this op.

What I had though of here was something like:

shape.if_static(%s : shape) {
  ^bb(%t: tensor<?xi32>):
     ... now use as 
} else {
}

so that you only convert when it is safe,

[silvas]

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

189

There is no such thing as a tensor with unknown or dynamic or even partially specified shape at runtime.

One can consider the abstract !shape.shape having a runtime backing representation which represents a dataflow lattice value (such as tracking extent upper bounds to allow buffer allocation before a data-dependent shape op (like "unique") completes), but that would require specialized ops that are aware of the underlying backing representation for the abstract !shape.shape in order to construct such a value, since the choice of lattice is arbitrary.

There's only a small number of ops on !shape.shape that reify the concept of the shape being a lattice value (since there just isn't much you can abstractly do with a lattice, by definition). shape.join is one of them (and the complementary "most general shape" is the only other one I can think of). All other ops should be defined as if they are runtime computations on concrete shapes, and any particular lattice that a user intends to use will need to substitute in appropriate transfer functions either via dataflow analysis statically, or reified in the ir during lowering (in their own dialect), or literally have the runtime manifestation of !shape.shape have virtual methods they can override to substitute the transfer functions for the primitives.

As far as side effects in the error case, you can't dodge them at least how you've defined them here, at least in my use case. The result of the shape.concat/shape.slice_at/etc. is used to broadcast a tensor to a particular concrete runtime shape. If the !shape.shape is an error, then now you've just passed the error onto the "broadcast this tensor to this shape op". In general you could pass the resulting shape to "{numpy,tf,torch}.zeros(shape)" which to your point should be pure.

I feel strongly that we need a dialect that can be used to model runtime shape computations. If you want this dialect to be used purely for abstract dataflow then that's okay and I can create a different dialect.

I think this error handling situation needs more thought. Neither tf dialect nor xla_hlo dialects (or any other dialects I'm aware of that would use this) have any actual IR manifestation for what happens in error cases. E.g. most "matmul" ops that I'm aware of claim to be side effect free. Hmm....

[jpienaar]

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

189

I'm not sure I follow, shape here is a effectively an error monad, the ops operating on it propagates the error (as mentioned this could end up merging multiple errors together), but the ops are side effect free and can be folded, else many of these ops are effectively asserts that can fail (e.g., they are adding constraints on the number of elements, the maximum extent of dimensions, ...). Once lowered to a runtime executable form, then those errors would need to be considered. E.g., a tf.assert or some such needs to be inserted before spanning the gap to tf.zeros or what not that operates on tensor of ints. There are ops in those platforms which can handle those and so in lowering handle it - there are actual runtime ops there with cancellation or abort mechanisms.

The goal of this dialect is to describe shape computations for multiple different consumers at analysis, compilation or for lowering to runtime. I don't see there being a single consumer/runtime here as the way that different consumers/runtimes could handle computations, failures, error reporting, would be different and keeping it abstract means the lowerings to those targets can consider what they need (e.g., perhaps we are in "assume there is no errors" mode and all checks can be elided, or a target is massively distributed and you want error reporting to be totally ordered).

These are defined by runtime behavior, it is just one of the "runtime" is symbolic analysis (e.g., runtime behavior includes how to create a constraint programming problem). Being able to reuse these in multiple contexts is important, else we end up defining the shapes of ops for each use case. For the runtime case with all the values known, you just need 1D tensor of ints, and you could reuse normal ops to implement the computations (e.g., you could use tf.slice or xla_hlo.mul or linalg), the result of these then just gets consumed by runtime specific functions to allocate/deallocate buffers, so you don't need any special types or ops at that point. So once you are using tensor of ints you are free to use arbitrary ops and you've jumped to a new domain, but given bridging these two could require handling errors, I want the lowering for a use to handle that.

So specified as runtime behavior on shapes makes sense to me, and it makes sense to have this op to convert from !shape.shape to tensor<?xindex> in this dialect, but not for use in a shape functions (else you are just using tensor<?x index> here as a cast to appease the type system).

I might be wrong though, I was trying to sketch multiple cases in my head (and refining the comment every now and again :)) and while I can see cases where you have benefit and approaches to use such an op, I don't know how many of them are really real. But one thing that is true, if we have such an if instead, you get the conversion and if we find it is not needed, it amounts to replacing the if with a .ToTensorExtend and folding the else [which could also be the lowering pattern towards runtime you use and that would give you the exact same shape function you have today]

[silvas]

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

189

We discussed this offline.

There was agreement that an op like ToExtentTensor needs to somehow be guarded against the case that the !shape.shape is an error. For now, the op will abort if the shape is an error.

We also were envisioning different use cases. I don't think that ToExtentTensor would ever be used inside a shape transfer function, which is what Jacques was thinking about. But my use case in a downstream dialect for lowering of batch matmul doesn't arise in a shape transfer function. Shape computations are needed to broadcast the input args to their common shape, which requires passing the extent tensor for the broadcasted shapes to a broadcast_to(%t, %extent_tensor) op before passing it to the lower level batch matmul op.

mlir/include/mlir/Dialect/Shape/IR/Shape.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SHAPE_IR_SHAPE_H
 #define MLIR_SHAPE_IR_SHAPE_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffects.h"
 
 namespace mlir {
 namespace shape {
 
 namespace ShapeTypes {
 enum Kind {
   Component = Type::FIRST_SHAPE_TYPE,

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

 //===- Shape.td - Shape operations definition --------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This is the operation definition file for Shape dialect operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SHAPE_OPS
 #define SHAPE_OPS
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffects.td"
 
 // TODO(jpienaar): Move to base.
 def AnyShaped: ShapedContainerType<[AnyType], IsShapedTypePred, "shaped">;
 
 //===----------------------------------------------------------------------===//
 // Shape Inference dialect definitions
 //===----------------------------------------------------------------------===//
def Shape_ConstantOp : Shape_Op<"constant", []> {
   let results = (outs Shape_ShapeOrSizeType:$result);
 
   // TODO: Move this to main so that all shape ops implement these.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   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 = [{
     An operation that computes the least general shape of input operands. This
     effectively asserts that corresponding static dimensions are equal. The
     behavior is to match each element of the `shape.type` and propagate the most
     restrictive information, returning an invalid shape if there are
     contradictory requirements. E.g., using pseudo code
let description = [{
 
     Note: This is intended for testing and debugging only.
   }];
 
   let arguments = (ins Shape_ShapeOrSizeType:$input);
   let results =  (outs Shape_ShapeOrSizeType:$output);
 }
 
+def Shape_SplitAtOp : Shape_Op<"split_at",
+    [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Splits a shape at a given index.";
+  let description = [{
+    Splits a shape at a given dimension `index`, returning two shapes.
+    If `index` is negative, it is treated as indexing from the back of the
+    shape. This negative-handling behavior is important when handling unranked
+    shapes, where the positive index is not necessarily knowable due to a
+    dynamic number of leading dimensions.
+
+    Examples:
+    - split_at([4,5,6], index=0) -> [], [4,5,6]
+    - split_at([4,5,6], index=1) -> [4], [5,6]
+    - split_at([4,5,6], index=2) -> [4,5], [6]
+    - split_at([4,5,6], index=3) -> [4,5,6], []
+    - split_at([4,5,6], index=4) -> error
+    - split_at([4,5,6], index=-1) -> [4,5], [6]
+    - split_at([4,5,6], index=-2) -> [4], [5,6]
+    - split_at([4,5,6], index=-3) -> [], [4,5,6]
+    - split_at([4,5,6], index=-4) -> error
+
+    Requires:
+    - `index` is in the range [-rank(operand),rank(operand)]
+  }];
+
+  let arguments = (ins Shape_ShapeType:$operand, I32:$index);
+  let results = (outs Shape_ShapeType:$head, Shape_ShapeType:$tail);
+}
+
+def Shape_ConcatOp : Shape_Op<"concat",
+    [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  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

mlir/lib/Dialect/Shape/CMakeLists.txt

 add_mlir_dialect_library(MLIRShape
   IR/Shape.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Shape
 
   DEPENDS
   MLIRShapeOpsIncGen
   )
 target_link_libraries(MLIRShape
   PUBLIC
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffects
   LLVMSupport
   )

mlir/lib/Dialect/Shape/IR/Shape.cpp

if (parser.parseColonType(type))
     return failure();
 
   // Add the attribute type to the list.
   return parser.addTypeToList(type, result.types);
 }
 
 static LogicalResult verify(ConstantOp &op) { return success(); }
 
+//===----------------------------------------------------------------------===//
+// SplitAtOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult SplitAtOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    ArrayRef<NamedAttribute> attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto shapeType = ShapeType::get(context);
+  inferredReturnTypes.push_back(shapeType);
+  inferredReturnTypes.push_back(shapeType);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// ConcatOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ConcatOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    ArrayRef<NamedAttribute> attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto shapeType = ShapeType::get(context);
+  inferredReturnTypes.push_back(shapeType);
+  return success();
+}
+
 namespace mlir {
 namespace shape {
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Shape/IR/ShapeOps.cpp.inc"
 
 } // namespace shape
 } // namespace mlir