diff --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td --- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td +++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td @@ -169,6 +169,18 @@ 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 --- a/mlir/test/Dialect/Shape/ops.mlir +++ b/mlir/test/Dialect/Shape/ops.mlir @@ -111,3 +111,8 @@ %0 = shape.from_extent_tensor %arg : tensor return %0 : !shape.shape } + +func @rank(%shape : !shape.shape) -> !shape.size { + %rank = shape.rank %shape + return %rank : !shape.size +}