Use index for shape.to_extent_tensor

Summary:
Also,
- add IndexTensor to OpBase.td
- fix typo in the op name. It was mistakenly `to_tensor` instead of
  `to_extent_tensor`.

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

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 77516b9..795fb35 100644 (file)
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -181,7 +181,7 @@ def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", []> {
   let results = (outs Shape_ShapeType:$result);
 }
 
-def Shape_ToExtentTensorOp : Shape_Op<"to_tensor", []> {
+def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", []> {
   let summary = "Creates a dimension tensor from a shape";
   // TODO: Think more about the error situation. Perhaps factor out the
   // error detection into a separate op so downstream consumers can control
@@ -197,7 +197,7 @@ def Shape_ToExtentTensorOp : Shape_Op<"to_tensor", []> {
   }];
 
   let arguments = (ins Shape_ShapeType:$input);
-  let results = (outs I32Tensor:$result);
+  let results = (outs IndexTensor:$result);
 }
 
 def Shape_JoinOp : Shape_Op<"join", []> {

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 56c958d..e552279 100644 (file)
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -582,6 +582,7 @@ def I8Tensor   : TensorOf<[I8]>;
 def I16Tensor  : TensorOf<[I16]>;
 def I32Tensor  : TensorOf<[I32]>;
 def I64Tensor  : TensorOf<[I64]>;
+def IndexTensor: TensorOf<[Index]>;
 
 def BF16Tensor : TensorOf<[BF16]>;
 def F16Tensor  : TensorOf<[F16]>;