[MLIR][Shape] Add `shape.shape_eq` operation

[frgossen]

Summary

Add shape.shape_eq operation to the shape dialect.
The operation allows to test shapes for equality.

[silvas]

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

141

Should it be SameOperandsAndResultType?

146

What about equality of errors?

152

Having two separate ops almost seems preferable (and is easier to isa/dyn_cast as well):

shape.shape_eq %lhs, %rhs
shape.size_eq %lhs, %rhs

vs

shape.eq %lhs, %rhs : !shape.shape
shape.eq %lhs, %rhs : !shape.size

More redundancy in the latter.

[silvas]

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

152

Also, it is much easier to builder.create with with two separate ops, since with two separate ops the return type can be inferred.

builder.create<shape::ShapeEqOp>(loc, lhs, rhs)
builder.create<shape::SizeEqOp>(loc, lhs, rhs)

vs

builder.create<shape::EqOp>(loc, builder.getType<shape::ShapeType>(), lhs, rhs)
builder.create<shape::EqOp>(loc, builder.getType<shape::SizeType>(), lhs, rhs)

[silvas]

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

141

D'oh, sorry. It's i1 result! My bad!

How is the error case going to be propagated then?

[silvas]

(we could also use shape.eq to be the one for shapes, which is a bit shorter than shape.shape_eq; I personally like being a bit more explicit but it's up to you what feels best namingwise)

[silvas]

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

152

D'oh, ignore that last comment. It returns i1.

[frgossen]

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

146

We could leave it undefined for now (and mention that in the documentation ofc).
Do you have any preferences? @tpopp , @herhut

152

Sounds good to me. I will split it into two.

[frgossen]

Move shape.size_eq to other CL

[silvas]

LGTM modulo getting feedback for others on how to define the error cases.

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

146

@jpienaar thoughts?

For now, we can say it is undefined. (I don't know what "unknown" means in this context)

[frgossen]

Update documentation

[herhut]

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

145

This entire unknown is a bit of a mess. I think we have two different concerns that we are trying to model here. One is the representation of partially constant values for constant folding and analysis. That is where the unknown comes in. On the other hand we have the runtime behavior, which does not know the concept of unknown.

Maybe we should separate the two issues and only describe the runtime semantics for the operations in the docs. Have an additional description of the special encoding that we use for constant folding.

That would fix these issues. @jpienaar WDYT?

If any of the inputs is an error, then the result is undefined works for me.

[jpienaar]

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

145

You are correct we do. I think there is more than that: there are operations that express constraints and operations that perform computations. And for the former this is useful but for the latter not. Having it expressed consistently is probably less confusion.

This operation being added here is an operation on shapes but express no constraint. And as such does not need to operate in constraint domain or make sense where errors or partial known values are, e.g., this op should be operating on tensor<?xindex>. For expressing constraint you'd have shape.join %0, %1 followed by feeding into a assuming op (currently they only take witness, but the valid bit of shape serves the purpose of witness, so operationally they make sense there to me). While if this operation only makes sense where no errors are possible, then it should only operate there and so take tensor<?xindex> instead.

146

As mentioned in the dialect doc:

all functions that return a shape and take shapes as input, return the invalid shape if one of its operands is an invalid shape

This op cannot be executed in any situation where one of the shapes can be an error. It only makes sense post verification/inside assume. As such, I don't know if if this op provides anything which equality check on tensor<?xindex> doesn't.

152

Also, it is much easier to builder.create with with two separate ops, since with two separate ops the return type can be inferred

Return type can be inferred in either case. As long as one has the requirement that that operands and result type match.

More redundancy in the latter.

Textual form is meant for debugging only. There is no more or less redundancy in the actual IR. This is still an IR and not a programming language.

[frgossen]

Reflect comments and email discussion

[frgossen]

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

146

I could not find an equality check on tensors in std (at least not with a lowering).
Having this operation in the shape dialect also keeps it a bit more closed (@herhut).

[frgossen]

Clarify documentation

[silvas]

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

112

Let's wait before we discuss this more before we add support for both types in the same op. I don't think we have consensus on the direction for this.

[herhut]

I think this is ready to land after yesterday's discussion.

[frgossen]

Correct linebreaks

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

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`.
   }];
 }
 
-def Shape_ShapeOrSizeType: AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
+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_ConstSizeOp : Shape_Op<"const_size", [
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, int64_t value">
   ];
 
   let assemblyFormat = "$value attr-dict";
   let hasFolder = 1;
 }
 
+def Shape_ShapeEqOp : Shape_Op<"shape_eq", [Commutative, NoSideEffect]> {
+  let summary = "Returns whether the input shapes or extent tensors are equal";
+  let description = [{
+    Takes two shape or extent tensor operands and determines whether they are
+    equal. When extent tensors are compared to shapes they are regarded as their
+    equivalent non-error shapes. Error shapes can be tested for equality like
+    any other shape value, meaning that the error value is equal to itself.
+  }];
+
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
+                       Shape_ShapeOrExtentTensorType:$rhs);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
+}
+
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
   let summary = "Creates a shape from extents";
   let description = [{
     Creates a shape from multiple SSA values representing the extents of
     the shape.
 
     ```mlir
     // Rank 2 shape.

mlir/test/Dialect/Shape/ops.mlir

func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
   %0 = shape.from_extent_tensor %arg : tensor<?xindex>
   return %0 : !shape.shape
 }
 
 func @rank(%shape : !shape.shape) -> !shape.size {
   %rank = shape.rank %shape
   return %rank : !shape.size
 }
+
+func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
+  %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
+  return %result : i1
+}
+
+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
+}