From 2c061998b55707b1b6f3a587ccc155fbbebd0a4e Mon Sep 17 00:00:00 2001
From: Frederik Gossen <frgossen@google.com>
Date: Thu, 25 Jun 2020 08:25:06 +0000
Subject: [PATCH] [MLIR][Shape] Add `shape.rank` operation

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

Differential Revision: https://reviews.llvm.org/D82028
---
 mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td | 12 ++++++++++++
 mlir/test/Dialect/Shape/ops.mlir               |  5 +++++
 2 files changed, 17 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 49714c9..0785a40 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 2e40211..18953ef 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
+}
-- 
2.7.4