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

[frgossen]

Summary

Add shape.rank operation to the shape dialect.

Depends On D82022

[silvas]

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

172

Can we call this "rank", which is the terminology usually used for this? Otherwise, it's quite confusing especially since we have !shape.size which is a totally different thing.

181

can you also add a constant folder?

[frgossen]

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

172

I see that the name can be confusing.
@pifon2a and I discussed the naming and we though that "rank" would be very confusing as the rank of the argument itself (the shape) is always 1.
What do you think about get_size as a name?

181

Will do that ofc.
Thought that should go in a separate CL.

[frgossen]

Update name

[frgossen]

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

181

See https://reviews.llvm.org/D82075

[frgossen]

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

181

See https://reviews.llvm.org/D82076

[silvas]

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

172

I would agree with you if the argument was an extent tensor. But the argument is a !shape.shape. "rank" is unambiguous when applied to something of !shape.shape type.

"size" is super ambiguous: Just based on the name it sounds like it could be any of these ops:

• shape.num_elements
• shape.get_rank
• shape.get_extent
• shape.const_size

Just as we have been very consistent to use "extent" to mean the size of a dimension for clarity, we should really use the term "rank" here as that has no other possible meaning in this context.

[frgossen]

Rename operation

[silvas]

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

175

nit: just to be precise and self-contained, add "The rank is the number of extents in a shape"

[frgossen]

Address nit

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

def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   }];
 
   let arguments = (ins IndexTensor:$input);
   let results = (outs Shape_ShapeType:$result);
 
   let assemblyFormat = "attr-dict $input `:` type($input)";
 }
 
+def Shape_Size : Shape_Op<"size", [NoSideEffect]> {
+  let summary = "Gets the size of a shape";
+  let description = [{
+    Returns the size of the shape, i.e. the number of extents.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$shape);
+  let results = (outs Shape_SizeType:$size);
+
+  let assemblyFormat = "attr-dict $shape";
+}
+
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   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.
 
     If the shape represents an error, this op's behavior is undefined.

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 @shape_size(%shape : !shape.shape) -> !shape.size {
+  %size = shape.size %shape
+  return %size : !shape.size
+}