[mlir] Add shape.with_shape op

[jpienaar]

Summary

This is an operation that can returns a new ValueShape with a different shape. Useful for composing shape function calls and reusing existing shape transfer functions.

Just adding the op in this change.

[herhut]

Adding the op SGTM.

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

420

Could we keep the op name and its textual form related? That makes it easier to match the two.

423

Can you be more precise what match means? Is the passed shape used verbatim (like a set_shape operation in a way)?

431

Nit: Drop one where.

447

I don't understand this comment. I think it answers my question from above but it is hard to parse.

[frgossen]

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

424

What does it men when shapes are non-conformant?

[frgossen]

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

423

It is only updated if the operand is a ValueShape. Otherwise, a new ValueShape is created?

440

Duplicate %%

mlir/test/Dialect/Shape/ops.mlir

158

Why shape_equal_shapes? This is essentially a join, is it?

164

Why shape_with_shape? The problem might just be that I don't fully understand the concept of ShapeValue.

[jpienaar]

Fixing typos, adding sentence to ValueShape type to describe conformance, adding more comments

[jpienaar]

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

420

Sorry, I had considered renaming it at the end and then decided it wasnt truly a refine (even though I found the naming better) but seems I forgot to rename back.

423

It is not really set: a new ValueShape is produced with the given shape, so effectively it is a copy and set.

423

A new ValueShape is always created, it is an SSA value.

424

shape_of(value) is equal to/conforms to shape at execution time. So the shape captured along with the value is consistent effectively.

447

Tried to clarify.

mlir/test/Dialect/Shape/ops.mlir

158

Correct, this is just function to invoke below, I just chose a simple one. I added a comment.

164

I was following the convention in this file op shape.X gets tested in func shape_X.

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

def Shape_ValueShapeType : DialectType<ShapeDialect,
     CPred<"$_self.isa<::mlir::shape::ValueShapeType>()">, "value shape">,
     BuildableType<"::mlir::shape::ValueShapeType::get($_builder.getContext())">
 {
   let typeDescription = [{
     `shape.value_shape` represents the value produced by an operation (this
     corresponds to `Value` in the compiler) and a shape. Conceptually this is a
     tuple of a value (potentially unknown) and `shape.type`. The value and shape
     can either or both be unknown. If both the `value` and `shape` are known,
-    then the shape of `value` is conformant with `shape`.
+    then the shape of `value` is conformant with `shape`. That is, the shape of
+    the value conforms to the shape of the ValueShape, so that if we have
+    ValueShape `(value, shape)` then `join(shape_of(value), shape)` would be
+    error free and in particular it means that if both are statically known,
+    then they are equal.
   }];
 }
 
 def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
 def Shape_ExtentTensorType : 1DTensorOf<[Index]>;
 
 def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
                                                Shape_ExtentTensorType],
                                               "shape or extent tensor">;
 
 #endif // SHAPE_BASE_TD

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

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let results = (outs Index:$result);
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
+def Shape_WithOp : Shape_Op<"with_shape", [NoSideEffect]> {
+  let summary = "Returns ValueShape with given shape";
+  let description = [{
+    Returns ValueShape with the shape updated to match the shape operand. That
+    is a new ValueShape tuple is created with value equal to `operand`'s
+    value and shape equal to `shape`. If the ValueShape and given `shape` are
+    non-conformant, then the returned ValueShape will represent an error of
+    this mismatch. Similarly if either inputs are in an error state, then an
+    error is popagated.
+
+    Usage:
+      %0 = shape.with_shape %1, %2 : tensor<...>, !shape.shape
+
+    This is used, for example, where one combines shape function calculations
+    and/or call one shape function from another. E.g.,
+
+    ```mlir
+    func @shape_foobah(%a: !shape.value_shape,
+                       %b: !shape.value_shape,
+                       %c: !shape.value_shape) -> !shape.shape {
+      %0 = call @shape_foo(%a, %b) :
+        (!shape.value_shape, !shape.value_shape) -> !shape.shape
+      %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
+      %2 = call @shape_bah(%c, %1) :
+        (!shape.value_shape, !shape.value_shape) -> !shape.shape
+      return %2 : !shape.shape
+    }
+    ```
+
+    This op need not be a refinement of the shape. In non-error cases the input
+    ValueShape's value and shape are conformant and so too for the output, but
+    the result may be less specified than `operand`'s shape as `shape` is
+    merely used to construct the new ValueShape. If join behavior is desired
+    then a join op should be used.
+  }];
+
+  let arguments = (ins AnyTypeOf<[AnyShaped, Shape_ValueShapeType]>:$operand,
+                       Shape_ShapeType:$shape);
+  let results = (outs Shape_ValueShapeType:$result);
+
+  let assemblyFormat = "operands attr-dict `:` type($operand) `,` type($shape)";
+}
+
 def Shape_YieldOp : Shape_Op<"yield",
     [HasParent<"ReduceOp">,
      NoSideEffect,
      ReturnLike,
      Terminator]> {
   let summary = "Returns the value to parent op";
 
   let arguments = (ins Variadic<AnyType>:$operands);

mlir/test/Dialect/Shape/ops.mlir

func @shape_eq_on_tensors(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {
   %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
   return %result : i1
 }
 
 func @shape_eq_on_mixed(%a : tensor<?xindex>, %b : !shape.shape) -> i1 {
   %result = shape.shape_eq %a, %b : tensor<?xindex>, !shape.shape
   return %result : i1
 }
+
+// Testing nvoking shape function from another. shape_equal_shapes is merely
+// a trivial helper function to invoke elsewhere.
+func @shape_equal_shapes(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
+  %0 = shape.shape_of %a : !shape.value_shape
+  %1 = shape.shape_of %b : !shape.value_shape
+  %2 = "shape.join"(%0, %1) : (!shape.shape, !shape.shape) -> !shape.shape
+  return %2 : !shape.shape
+}
+func @shape_with_shape(%a : !shape.value_shape, %b : !shape.value_shape) -> !shape.shape {
+  %0 = shape.shape_of %a : !shape.value_shape
+  %1 = shape.with_shape %b, %0 : !shape.value_shape, !shape.shape
+  %2 = call @shape_equal_shapes(%a, %1) : (!shape.value_shape, !shape.value_shape) -> !shape.shape
+  return %2 : !shape.shape
+}