[MLIR] Fix operand type in `from_extent_tensor` in the shape dialect

The operand of `from_extent_tensor` is now of the same index type as the result
type of the inverse operation `to_extent_tensor`.

Differential Revision: https://reviews.llvm.org/D80283

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 0d300d3..460f5be 100644 (file)
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -163,7 +163,7 @@ def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
     extents match the values of the elements.
   }];
 
-  let arguments = (ins I32Tensor:$input);
+  let arguments = (ins IndexTensor:$input);
   let results = (outs Shape_ShapeType:$result);
 }