[Mlir-commits] [mlir] 2c06199 - [MLIR][Shape] Add `shape.rank` operation

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-06-25T08:26:00Z
New Revision: 2c061998b55707b1b6f3a587ccc155fbbebd0a4e

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

LOG: [MLIR][Shape] Add `shape.rank` operation

Add `shape.rank` operation to the shape dialect.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
    mlir/test/Dialect/Shape/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 49714c96b1cd..0785a40ae513 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -169,6 +169,18 @@ def Shape_FromExtentTensorOp : Shape_Op<"from_extent_tensor", [NoSideEffect]> {
   let assemblyFormat = "attr-dict $input `:` type($input)";
 }
 
+def Shape_RankOp : Shape_Op<"rank", [NoSideEffect]> {
+  let summary = "Gets the rank of a shape";
+  let description = [{
+    Returns the rank of the shape, i.e. the number of extents.
+  }];
+
+  let arguments = (ins Shape_ShapeType:$shape);
+  let results = (outs Shape_SizeType:$rank);
+
+  let assemblyFormat = "attr-dict $shape";
+}
+
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
   let description = [{

diff  --git a/mlir/test/Dialect/Shape/ops.mlir b/mlir/test/Dialect/Shape/ops.mlir
index 2e40211e5a63..18953efa05f1 100644
--- a/mlir/test/Dialect/Shape/ops.mlir
+++ b/mlir/test/Dialect/Shape/ops.mlir
@@ -111,3 +111,8 @@ func @test_from_extent_tensor(%arg: tensor<?xindex>) -> !shape.shape {
   %0 = shape.from_extent_tensor %arg : tensor<?xindex>
   return %0 : !shape.shape
 }
+
+func @rank(%shape : !shape.shape) -> !shape.size {
+  %rank = shape.rank %shape
+  return %rank : !shape.size
+}