[Mlir-commits] [mlir] 483f278 - Use index for shape.to_extent_tensor

[Sean Silva]

Author: Sean Silva
Date: 2020-04-14T15:09:18-07:00
New Revision: 483f2783ae013f0a07e11d16cf8f6d00fa552704

URL: https://github.com/llvm/llvm-project/commit/483f2783ae013f0a07e11d16cf8f6d00fa552704
DIFF: https://github.com/llvm/llvm-project/commit/483f2783ae013f0a07e11d16cf8f6d00fa552704.diff

LOG: 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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
    mlir/include/mlir/IR/OpBase.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 77516b9f29cd..795fb353768b 100644
--- 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 56c958d16a8f..e552279e8a8e 100644
--- 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]>;