[Mlir-commits] [mlir] 4cd9237 - [MLIR][Shape] Expose extent tensor type builder

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-08-05T09:42:57Z
New Revision: 4cd923784e9079384792e0aed38d56809d6a4f9a

URL: https://github.com/llvm/llvm-project/commit/4cd923784e9079384792e0aed38d56809d6a4f9a
DIFF: https://github.com/llvm/llvm-project/commit/4cd923784e9079384792e0aed38d56809d6a4f9a.diff

LOG: [MLIR][Shape] Expose extent tensor type builder

The extent tensor type is a `tensor<?xindex>` that is used in the shape dialect.
To facilitate the use of this type when working with the shape dialect, we
expose the helper function for its construction.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/Shape.h
    mlir/lib/Dialect/Shape/IR/Shape.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/Shape.h b/mlir/include/mlir/Dialect/Shape/IR/Shape.h
index 62e4e0c4511f..ca1e0668d070 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/Shape.h
+++ b/mlir/include/mlir/Dialect/Shape/IR/Shape.h
@@ -26,6 +26,9 @@ class PatternRewriter;
 
 namespace shape {
 
+/// Alias type for extent tensors.
+RankedTensorType getExtentTensorType(MLIRContext *ctx);
+
 namespace ShapeTypes {
 enum Kind {
   Component = Type::FIRST_SHAPE_TYPE,

diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 02fe7b8129f7..be4c3c721572 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -24,7 +24,7 @@ namespace {
 #include "ShapeCanonicalization.inc"
 }
 
-static RankedTensorType getExtentTensorType(MLIRContext *ctx) {
+RankedTensorType shape::getExtentTensorType(MLIRContext *ctx) {
   return RankedTensorType::get({ShapedType::kDynamicSize}, IndexType::get(ctx));
 }
 
@@ -713,12 +713,9 @@ OpFoldResult ShapeOfOp::fold(ArrayRef<Attribute>) {
 }
 
 void ShapeOfOp::build(OpBuilder &builder, OperationState &result, Value arg) {
-  if (arg.getType().isa<ShapedType>()) {
-    auto type = RankedTensorType::get({ShapedType::kDynamicSize},
-                                      builder.getIndexType());
-    return ShapeOfOp::build(builder, result, type, arg);
-  }
-  auto type = ShapeType::get(builder.getContext());
+  Type type = arg.getType().isa<ShapedType>()
+                  ? (Type)getExtentTensorType(builder.getContext())
+                  : (Type)builder.getType<ShapeType>();
   return ShapeOfOp::build(builder, result, type, arg);
 }