[Mlir-commits] [mlir] a97940d - [MLIR][Shape] Limit `shape.rank` lowering to its extent tensor variant

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-07-30T11:43:08Z
New Revision: a97940d4e0ed0507eb9beba7f22e34d62449a58f

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

LOG: [MLIR][Shape] Limit `shape.rank` lowering to its extent tensor variant

When lowering to the standard dialect, we currently support only the extent
tensor variant of the shape.rank operation. This change lets the conversion
pattern fail in a well-defined manner.

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

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 41d4d90b33d3..debdedc89e87 100644
--- a/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
+++ b/mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp
@@ -231,6 +231,10 @@ class RankOpConverter : public OpConversionPattern<shape::RankOp> {
 LogicalResult
 RankOpConverter::matchAndRewrite(shape::RankOp op, ArrayRef<Value> operands,
                                  ConversionPatternRewriter &rewriter) const {
+  // For now, this lowering supports only error-free types.
+  if (op.getType().isa<SizeType>())
+    return failure();
+
   shape::RankOp::Adaptor transformed(operands);
   rewriter.replaceOpWithNewOp<DimOp>(op, transformed.shape(), 0);
   return success();

diff  --git a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
index 3f19de9c52f0..b0fb5bac9071 100644
--- a/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
+++ b/mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir
@@ -90,6 +90,16 @@ func @get_extent(%shape : tensor<?xindex>, %idx : !shape.size) -> !shape.size {
 
 // -----
 
+// Don't lower `rank` if type is not error-free.
+// CHECK-LABEL: @rank
+func @rank(%shape : !shape.shape) {
+  // CHECK: shape.rank
+  %rank = shape.rank %shape : !shape.shape -> !shape.size
+  return
+}
+
+// -----
+
 // Express `get_extent` as `std.dim` when it relies directly on the outcome of a
 // `shape_of` operation.
 // CHECK-LABEL: @get_extent_shape_of