[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

137

Should it be SameOperandsAndResultType?

142

What about equality of errors?

148

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

148

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

137

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

148

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

[frgossen]

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

142

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

148

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

142

@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

141

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

141

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.

142

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.

148

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

142

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 ↗

(On Diff #277430)

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/ShapeOps.td

def Shape_ConstSizeOp : Shape_Op<"const_size", [
 
   let arguments = (ins IndexAttr:$value);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "$value attr-dict";
   let hasFolder = 1;
 }
 
+def Shape_EqOp : Shape_Op<"eq", [Commutative, NoSideEffect, SameTypeOperands]> {
+  let summary = "Returns whether the input shapes or sizes are equal";
+  let description = [{
+    Takes either two shape operands or two size operands and determines whether
+    they are equal.
+    Shape operands are equal if their rank and all extents are equal.
+  }];
+
+  let arguments = (ins Shape_ShapeOrSizeType:$lhs, Shape_ShapeOrSizeType:$rhs);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs `:` type($lhs) attr-dict";
+}
+
 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_to_extent_tensor(%arg: !shape.shape) -> tensor<3xindex> {
   %0 = shape.to_extent_tensor %arg : tensor<3xindex>
   return %0 : tensor<3xindex>
 }
 
 func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
   %0 = shape.from_extent_tensor %arg : tensor<?xindex>
   return %0 : !shape.shape
 }
+
+func @eq_shape(%a : !shape.shape, %b : !shape.shape) -> i1 {
+  %result = shape.eq %a, %b : !shape.shape
+  return %result : i1
+}
+
+func @eq_size(%a : !shape.size, %b : !shape.size) -> i1 {
+  %result = shape.eq %a, %b : !shape.size
+  return %result : i1
+}