[MLIR][Shape] Allow `shape.get_extent` to operate on extent tensors

[frgossen]

Summary

shape.get_extent now accepts extent tensors tensor<?xindex> as an argument.

Depends On D84157

[herhut]

mlir/test/Dialect/Shape/canonicalize.mlir

257

Should these also be tested on tensor types?

[frgossen]

Test on extent tensor instead of shapes

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

def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let results = (outs IndexTensor:$result);
 
   let assemblyFormat = "attr-dict $input `:` type($result)";
 
   let hasFolder = 1;
 }
 
 def Shape_GetExtentOp : Shape_Op<"get_extent", [NoSideEffect]> {
-  let summary = "Gets the specified extent from a shape";
+  let summary = "Gets the specified extent from a shape or extent tensor";
   let description = [{
-    Gets the extent indexed by `dim` from `shape`.
-    If the shape is an error, it returns an error size.
+    Gets the extent indexed by `dim` from the `shape` operand. If the shape is
+    an error then it returns an error size.
   }];
-  let arguments = (ins
-    Shape_ShapeType:$shape,
-    Shape_SizeType:$dim
-  );
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,
+                       Shape_SizeType:$dim);
   let results = (outs Shape_SizeType:$extent);
-  let assemblyFormat = "$shape `,` $dim attr-dict";
+  let assemblyFormat = "$shape `,` $dim `:` type($shape)  attr-dict";
 
   let builders = [
     // Builder that allows passing a constant dimension as a simple integer.
     OpBuilder<"OpBuilder &builder, OperationState &result, Value shape, "
               "int64_t dim">
   ];
 
   let extraClassDeclaration = [{

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // `shape_of` operation.
 // CHECK-LABEL: @get_extent_shape_of
 // CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
 func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : !shape.size)
     -> !shape.size {
   // CHECK: %[[RESULT:.*]] = dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
   // CHECK: return %[[RESULT]] : index
   %shape = shape.shape_of %arg : tensor<2x3xf32>
-  %result = shape.get_extent %shape, %idx
+  %result = shape.get_extent %shape, %idx : !shape.shape
   return %result : !shape.size
 }
 
 // -----
 
 // Express `get_extent` as `std.extract_element` when it relies directly on the
 // outcome of a `from_extent_tensor` operation.
 // CHECK-LABEL: @get_extent_from_extent_tensor
 // CHECK-SAME: (%[[EXTENTS:.*]]: tensor<?xindex>, %[[IDX:.*]]: index) -> index
 func @get_extent_from_extent_tensor(%extents : tensor<?xindex>,
                                     %idx : !shape.size) -> !shape.size {
   // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor<?xindex>
   // CHECK: return %[[RESULT]] : index
   %shape = shape.from_extent_tensor %extents : tensor<?xindex>
-  %result = shape.get_extent %shape, %idx
+  %result = shape.get_extent %shape, %idx : !shape.shape
   return %result : !shape.size
 }
 
 // -----
 
 // Lower `any` to its first operand.
 // CHECK-LABEL: @any_of_three
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>) -> tensor<?xindex>

mlir/test/Dialect/Shape/canonicalize.mlir

 }
 
 // -----
 
 // Basic folding.
 // CHECK-LABEL: func @basic
 func @basic() -> !shape.size {
   // CHECK: shape.const_size 2
-  %0 = shape.const_shape [0, 1, 2] : !shape.shape
+  %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
   %c2 = shape.const_size 2
-  %1 = shape.get_extent %0, %c2
+  %1 = shape.get_extent %0, %c2 : tensor<?xindex>
   return %1 : !shape.size
 }
 
 // -----
 
 // Should not fold.
 // CHECK-LABEL: func @out_of_bounds
 func @out_of_bounds() -> !shape.size {
   // CHECK: shape.const_shape
   // CHECK: shape.get_extent
-  %0 = shape.const_shape [0, 1, 2] : !shape.shape
+  %0 = shape.const_shape [0, 1, 2] : tensor<?xindex>
   %c3 = shape.const_size 3
-  %1 = shape.get_extent %0, %c3
+  %1 = shape.get_extent %0, %c3 : tensor<?xindex>
   return %1 : !shape.size
 }
 
 // -----
 
 // Should not fold.
 // CHECK-LABEL: func @not_const
-func @not_const(%arg0: !shape.shape) -> !shape.size {
+func @not_const(%arg0: tensor<?xindex>) -> !shape.size {
   // CHECK: shape.get_extent
   %c3 = shape.const_size 3
-  %0 = shape.get_extent %arg0, %c3
+  %0 = shape.get_extent %arg0, %c3 : tensor<?xindex>
   return %0 : !shape.size
 }
 
 
 // -----
 // cstr_eq with non-constant but known equal shapes can be removed.
 // CHECK-LABEL: func @f
 func @f(%arg0 : !shape.shape) {

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
 }
+
+func @get_extent_on_shape(%arg : !shape.shape) -> !shape.size {
+  %c0 = shape.const_size 0
+  %result = shape.get_extent %arg, %c0 : !shape.shape
+  return %result : !shape.size
+}
+
+func @get_extent_on_extent_tensor(%arg : tensor<?xindex>) -> !shape.size {
+  %c0 = shape.const_size 0
+  %result = shape.get_extent %arg, %c0 : tensor<?xindex>
+  return %result : !shape.size
+}