[MLIR][Shape] Allow `shape.rank` to operate on extent tensors

[frgossen]

Summary

Depends On D84424

[jpienaar]

mlir/lib/Dialect/Shape/IR/Shape.cpp

17

clang-tidy: error: 'mlir/include/mlir/Dialect/StandardOps/IR/Ops.h' file not found [clang-diagnostic-error]

(just remove mlir/include part)

578

This one I would relax, as I can see this happening in two steps:

%x = ... : shape // (1)
%0 = shape.rank %x : !shape.size

then we end folding to the point where (1) could be verified to be error free and canonicalized. If we made rank more permissive it would still be valid and we could canonicalize it post rather than flagging as invalid. (also this doesn't lose info and so is valid)

[frgossen]

Address comments

mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td

 
 def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
 def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
                                                Shape_ExtentTensorType],
                                               "shape or extent tensor">;
 
+def Shape_SizeOrIndexType : AnyTypeOf<[Shape_SizeType, Index], "size or index">;
+
 #endif // SHAPE_BASE_TD

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

 
 def Shape_RankOp : Shape_Op<"rank", [NoSideEffect]> {
   let summary = "Gets the rank of a shape";
   let description = [{
     Returns the rank of the shape or extent tensor, i.e. the number of extents.
   }];
 
   let arguments = (ins Shape_ShapeOrExtentTensorType:$shape);
-  let results = (outs Shape_SizeType:$rank);
+  let results = (outs Shape_SizeOrIndexType:$rank);
 
-  let assemblyFormat = "$shape `:` type($shape) attr-dict";
+  let assemblyFormat = "$shape `:` type($shape) `->` type($rank) attr-dict";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
+  let verifier = [{ return ::verify(*this); }];
 }
 
 def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [NoSideEffect]> {
   let summary = "Creates a dimension tensor from a shape";
   let description = [{
     Converts a shape to a 1D integral tensor of extents. The number of elements
     in the tensor equals the rank of the shape, and the elements equal the
     extents of the shape.

mlir/lib/Dialect/Shape/IR/Shape.cpp

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::shape;
 
 namespace {
 #include "ShapeCanonicalization.inc"
 }
Operation *ShapeDialect::materializeConstant(OpBuilder &builder,
   if (type.isa<ShapeType>() ||
       type == getExtentTensorType(builder.getContext()))
     return builder.create<ConstShapeOp>(loc, type,
                                         value.cast<DenseIntElementsAttr>());
   if (type.isa<SizeType>())
     return builder.create<ConstSizeOp>(loc, type, value.cast<IntegerAttr>());
   if (type.isa<WitnessType>())
     return builder.create<ConstWitnessOp>(loc, type, value.cast<BoolAttr>());
+  if (type.isa<IndexType>())
+    return builder.create<ConstantOp>(loc, type, value);
   return nullptr;
 }
 
 /// Parse a type registered to this dialect.
 Type ShapeDialect::parseType(DialectAsmParser &parser) const {
   StringRef keyword;
   if (parser.parseKeyword(&keyword))
     return Type();
void GetExtentOp::build(OpBuilder &builder, OperationState &result, Value shape,
   Value dimValue = builder.create<ConstSizeOp>(loc, dimAttr);
   build(builder, result, shape, dimValue);
 }
 
 //===----------------------------------------------------------------------===//
 // RankOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult RankOp::fold(ArrayRef<Attribute> operands) {
+static LogicalResult verify(shape::RankOp op) {
+  Type argTy = op.shape().getType();
+  Type resultTy = op.rank().getType();
+  if (argTy.isa<ShapeType>() && !resultTy.isa<SizeType>())
+    return op.emitOpError()
+           << "if operand is of type `shape` then the result must be of type "
+              "`size` to propagate potential errors";
+  return success();
+}
+
+OpFoldResult shape::RankOp::fold(ArrayRef<Attribute> operands) {
   auto shape = operands[0].dyn_cast_or_null<DenseIntElementsAttr>();
   if (!shape)
     return {};
   int64_t rank = shape.getNumElements();
   Builder builder(getContext());
   return builder.getIndexAttr(rank);
 }
 
 /// Evaluate the `rank` operation for shapes of ranked tensors at compile time.
 /// Constant folding fails in cases where only the rank is constant, not the
 /// shape itself.
 /// This canonicalization matches `shape.rank(shape.shape_of(%ranked_tensor))`.
 ///
 /// Example:
 ///
 /// %shape = shape.shape_of %ranked_tensor : tensor<1x2x?xf32>
 /// %rank = shape.rank %shape
 ///
 /// becomes
 ///
 /// %rank = shape.const_size 3
 
 namespace {
-struct RankShapeOfCanonicalizationPattern : public OpRewritePattern<RankOp> {
-  using OpRewritePattern<RankOp>::OpRewritePattern;
+struct RankShapeOfCanonicalizationPattern
+    : public OpRewritePattern<shape::RankOp> {
+  using OpRewritePattern<shape::RankOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(RankOp op,
+  LogicalResult matchAndRewrite(shape::RankOp op,
                                 PatternRewriter &rewriter) const override {
     auto shapeOfOp = op.shape().getDefiningOp<ShapeOfOp>();
     if (!shapeOfOp)
       return failure();
     auto rankedTensorType =
         shapeOfOp.arg().getType().dyn_cast<RankedTensorType>();
     if (!rankedTensorType)
       return failure();
+    assert(op.getType().isa<IndexType>() &&
+           "expected `rank(shape_of( ... )]` based on a shaped argument to "
+           "yield an index type");
     int64_t rank = rankedTensorType.getRank();
-    rewriter.replaceOpWithNewOp<ConstSizeOp>(op.getOperation(), rank);
+    rewriter.replaceOpWithNewOp<ConstantIndexOp>(op.getOperation(), rank);
     return success();
   }
 };
 } // namespace
 
-void RankOp::getCanonicalizationPatterns(OwningRewritePatternList &patterns,
-                                         MLIRContext *context) {
+void shape::RankOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<RankShapeOfCanonicalizationPattern>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // NumElementsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult NumElementsOp::fold(ArrayRef<Attribute> operands) {

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

func @shape_of_dyn(%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 : !shape.shape
-  return %rank : !shape.size
+func @rank(%shape : tensor<?xindex>) -> index {
+  // CHECK: %[[C0:.*]] = constant 0 : index
+  // CHECK: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
+  // CHECK: return %[[RESULT]] : index
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
 }
 
 // -----
 
 // Express `get_extent` as `std.dim` when it relies directly on the outcome of a
 // `shape_of` operation.
 // CHECK-LABEL: @get_extent_shape_of
 // CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index

mlir/test/Dialect/Shape/canonicalize.mlir

func @broadcastable_on_extent_tensors(%arg : tensor<?xindex>) {
   return
 }
 
 // -----
 
 // Fold `rank` based on constant shape.
 // CHECK-LABEL: @fold_rank
 func @fold_rank() -> !shape.size {
-  // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 5
-  // CHECK-DAG: return %[[RESULT]] : !shape.size
+  // CHECK: %[[RESULT:.*]] = shape.const_size 5
+  // CHECK: return %[[RESULT]] : !shape.size
   %shape = shape.const_shape [3, 4, 5, 6, 7] : !shape.shape
-  %rank = shape.rank %shape : !shape.shape
+  %rank = shape.rank %shape : !shape.shape -> !shape.size
   return %rank : !shape.size
 }
 
 // -----
 
 // Do not fold `rank` if shape is dynamic.
 // CHECK-LABEL: @dont_fold_rank
 // CHECK-SAME: (%[[SHAPE:.*]]: !shape.shape) -> !shape.size
 func @dont_fold_rank(%shape : !shape.shape) -> !shape.size {
-  // CHECK-DAG: %[[RESULT:.*]] = shape.rank %[[SHAPE]]
-  // CHECK-DAG: return %[[RESULT]] : !shape.size
-  %rank = shape.rank %shape : !shape.shape
+  // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]]
+  // CHECK: return %[[RESULT]] : !shape.size
+  %rank = shape.rank %shape : !shape.shape -> !shape.size
   return %rank : !shape.size
 }
 
 // -----
 
+// Fold `rank` based on constant extent tensor.
+// CHECK-LABEL: @fold_rank
+func @fold_rank() -> index {
+  // CHECK: %[[RESULT:.*]] = constant 5 : index
+  // CHECK: return %[[RESULT]] : index
+  %shape = shape.const_shape [3, 4, 5, 6, 7] : tensor<?xindex>
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
+}
+
+// -----
+
+// Do not fold `rank` for non-constant extent tensors.
+// CHECK-LABEL: @dont_fold_rank
+// CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
+func @dont_fold_rank(%shape : tensor<?xindex>) -> index {
+  // CHECK: %[[RESULT:.*]] = shape.rank %[[SHAPE]] : tensor<?xindex> -> index
+  // CHECK: return %[[RESULT]] : index
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
+}
+
+// -----
+
 // Canonicalize `rank` when shape is derived from ranked tensor.
 // CHECK-LABEL: @canonicalize_rank
-func @canonicalize_rank(%arg : tensor<1x2x?xf32>) -> !shape.size {
-  // CHECK-DAG: %[[RESULT:.*]] = shape.const_size 3
-  // CHECK-DAG: return %[[RESULT]] : !shape.size
+func @canonicalize_rank(%arg : tensor<1x2x?xf32>) -> index {
+  // CHECK: %[[RESULT:.*]] = constant 3 : index
+  // CHECK: return %[[RESULT]] : index
   %shape = shape.shape_of %arg : tensor<1x2x?xf32> -> tensor<?xindex>
-  %rank = shape.rank %shape : tensor<?xindex>
-  return %rank : !shape.size
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
 }
 
 // -----
 
 // Do not canonicalize `rank` when shape is derived from unranked tensor.
 // CHECK-LABEL: @dont_canonicalize_rank
-// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> !shape.size
-func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> !shape.size {
-  // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor<?xindex>
-  // CHECK-DAG: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
-  // CHECK-DAG: return %[[SIZE]] : !shape.size
+// CHECK-SAME: (%[[ARG:.*]]: tensor<*xf32>) -> index
+func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> index {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32> -> tensor<?xindex>
+  // CHECK: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
+  // CHECK: return %[[SIZE]] : index
   %shape = shape.shape_of %arg : tensor<*xf32> -> tensor<?xindex>
-  %rank = shape.rank %shape : tensor<?xindex>
-  return %rank : !shape.size
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
 }
 
+// -----
+
 // Canonicalize redundant conversion from `index` to `size` and back.
 // CHECK-LABEL: @index_to_size_to_index
 // CHECK-SAME: (%[[IDX:.*]]: index) -> index
 func @index_to_size_to_index(%index : index) -> index {
   // CHECK: return %[[IDX]] : index
   %size = shape.index_to_size %index
   %result = shape.size_to_index %size
   return %result : index

mlir/test/Dialect/Shape/invalid.mlir

 // -----
 
 func @shape_of(%value_arg : !shape.value_shape,
                %shaped_arg : tensor<?x3x4xf32>) {
   // expected-error@+1 {{if operand is a shaped type then the result must be an extent tensor}}
   %1 = shape.shape_of %shaped_arg : tensor<?x3x4xf32> -> !shape.shape
 }
 
+// -----
+
+func @rank(%arg : !shape.shape) {
+  // expected-error@+1 {{if operand is of type `shape` then the result must be of type `size` to propagate potential errors}}
+  %0 = shape.rank %arg : !shape.shape -> index
+}
+

mlir/test/Dialect/Shape/ops.mlir

 }
 
 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 : !shape.shape
+  %rank = shape.rank %shape : !shape.shape -> !shape.size
   return %rank : !shape.size
 }
 
-func @rank_on_extent_tensor(%shape : tensor<?xindex>) -> !shape.size {
-  %rank = shape.rank %shape : tensor<?xindex>
-  return %rank : !shape.size
+func @rank_on_extent_tensor(%shape : tensor<?xindex>) -> index {
+  %rank = shape.rank %shape : tensor<?xindex> -> index
+  return %rank : index
 }
 
 
 func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
   %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
   return %result : i1
 }