[Mlir-commits] [mlir] 24debf5 - [MLIR][Shape] Lower `shape.rank`

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-06-25T08:44:06Z
New Revision: 24debf5a76e03cfeab6df2530151e3061eb110f2

URL: https://github.com/llvm/llvm-project/commit/24debf5a76e03cfeab6df2530151e3061eb110f2
DIFF: https://github.com/llvm/llvm-project/commit/24debf5a76e03cfeab6df2530151e3061eb110f2.diff

LOG: [MLIR][Shape] Lower `shape.rank`

Lower `shape.rank` to standard dialect.
A shape's size is the same as the extent of the first and only dimension of the
`tensor<?xindex>` it is represented by.

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

Added: 
    

Modified: 
    mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
    mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
index 6a02bdc2c286..5fd9be0bd73a 100644
--- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
+++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
@@ -90,6 +90,20 @@ class ConstSizeOpConverter : public OpConversionPattern<ConstSizeOp> {
   }
 };
 
+class RankOpConverter : public OpConversionPattern<shape::RankOp> {
+public:
+  using OpConversionPattern<shape::RankOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::RankOp::Adaptor transformed(operands);
+    rewriter.replaceOpWithNewOp<DimOp>(op.getOperation(), transformed.shape(),
+                                       0);
+    return success();
+  }
+};
+
 /// Type conversions.
 class ShapeTypeConverter : public TypeConverter {
 public:
@@ -147,6 +161,7 @@ void mlir::populateShapeToStandardConversionPatterns(
       BinaryOpConversion<AddOp, AddIOp>,
       BinaryOpConversion<MulOp, MulIOp>,
       ConstSizeOpConverter,
+      RankOpConverter,
       ShapeOfOpConversion>(ctx);
   // clang-format on
 }

diff  --git a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
index bfe3c2b599c5..a9b4bf701909 100644
--- a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
+++ b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
@@ -86,7 +86,6 @@ func @size_const() -> !shape.size {
 }
 // CHECK: %[[C1:.*]] = constant 1 : index
 // CHECK: return %[[C1]] : index
-
 // -----
 
 // Lower `shape_of` for statically shaped tensor.
@@ -115,3 +114,16 @@ func @shape_of_dyn(%arg : tensor<1x5x?xf32>) {
   %shape = shape.shape_of %arg : tensor<1x5x?xf32>
   return
 }
+
+// -----
+
+// Convert `rank` to `dim` of the first dimension.
+// CHECK-LABEL: @rank
+// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
+func @rank(%shape : !shape.shape) -> !shape.size {
+  // CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  // CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
+  // CHECK-DAG: return %[[RESULT]] : index
+  %rank = shape.rank %shape
+  return %rank : !shape.size
+}