[mlir][Interfaces] Split getArrayAttrMaxRanks into 3 functions

[springerm]

Summary

getArrayAttrMaxRanks returns the number of offsets/sizes/strides of ops that implement the OffsetSizeAndStrideOpInterface. This revision splits the this interface method into 3 methods, in preparation of extracting the "offsets" part of the interface into a new interface. Also provide a default implementation for ops that implement the ViewLikeOpInterface.

Depends On: D156864

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

def MemRef_ReinterpretCastOp
   let extraClassDeclaration = extraBaseClassDeclaration # [{
     // The result of the op is always a ranked memref.
     MemRefType getType() { return getResult().getType(); }
     Value getViewSource() { return getSource(); }
 
     /// Return the rank of the result type.
     unsigned getResultRank() { return getType().getRank(); }
 
-    /// Return the expected rank of each of the`static_offsets`, `static_sizes`
-    /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrMaxRanks() {
-      unsigned resultRank = getType().getRank();
-      return {1, resultRank, resultRank};
-    }
+    // Return the expected size of each of the `static_offsets`, `static_sizes`
+    // and `static_strides` attributes.
+    int64_t getOffsetsRank() { return 1; }
+    int64_t getSizesRank() { return getType().getRank(); }
+    int64_t getStridesRank() { return getType().getRank(); }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return a vector of all the static or dynamic sizes of the op, while
     /// statically inferring the sizes of the dynamic sizes, when possible.
     /// This is best effort.
static Type inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            ArrayRef<int64_t> staticSizes,
                                            ArrayRef<int64_t> staticStrides);
     static Type inferRankReducedResultType(ArrayRef<int64_t> resultShape,
                                            MemRefType sourceMemRefType,
                                            ArrayRef<OpFoldResult> staticOffsets,
                                            ArrayRef<OpFoldResult> staticSizes,
                                            ArrayRef<OpFoldResult> staticStrides);
 
-    /// Return the expected rank of each of the`static_offsets`, `static_sizes`
-    /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrMaxRanks() {
-      unsigned rank = getSourceType().getRank();
-      return {rank, rank, rank};
-    }
-
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return the dimensions of the source type that are dropped when
     /// the result is rank-reduced.
     llvm::SmallBitVector getDroppedDims();
 

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

static RankedTensorType inferCanonicalRankReducedResultType(
       ArrayRef<int64_t> staticStrides);
     static RankedTensorType inferCanonicalRankReducedResultType(
       unsigned resultRank,
       RankedTensorType sourceRankedTensorType,
       ArrayRef<OpFoldResult> staticOffsets,
       ArrayRef<OpFoldResult> staticSizes,
       ArrayRef<OpFoldResult> staticStrides);
 
-    /// Return the expected rank of each of the`static_offsets`, `static_sizes`
-    /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrMaxRanks() {
-      unsigned rank = getSourceType().getRank();
-      return {rank, rank, rank};
-    }
+    // Return the expected size of each of the `static_offsets`, `static_sizes`
+    // and `static_strides` attributes.
+    int64_t getOffsetsRank() { return getSourceType().getRank(); }
+    int64_t getSizesRank() { return getSourceType().getRank(); }
+    int64_t getStridesRank() { return getSourceType().getRank(); }
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return the dimensions of the source that are dropped in the
     /// result when the result is rank-reduced.
     llvm::SmallBitVector getDroppedDims();
RankedTensorType getType() {
       return getResultType();
     }
 
     /// The `dest` type is the same as the result type.
     RankedTensorType getDestType() {
       return getResultType();
     }
 
-    /// Return the expected rank of each of the`static_offsets`, `static_sizes`
-    /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrMaxRanks() {
-      unsigned rank = getResultType().getRank();
-      return {rank, rank, rank};
-    }
+    // Return the expected size of each of the `static_offsets`, `static_sizes`
+    // and `static_strides` attributes.
+    int64_t getOffsetsRank() { return getDestType().getRank(); }
+    int64_t getSizesRank() { return getDestType().getRank(); }
+    int64_t getStridesRank() { return getDestType().getRank(); }
 
     /// Return the dimensions of the dest that are omitted to insert a source
     /// when the result is rank-extended.
     llvm::SmallBitVector getDroppedDims();
 
     /// Return the number of leading operands before the `offsets`, `sizes` and
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; }
RankedTensorType getDestType() {
       return ::llvm::cast<RankedTensorType>(getDest().getType());
     }
 
     ParallelCombiningOpInterface getParallelCombiningParent() {
       return dyn_cast<ParallelCombiningOpInterface>(
         getOperation()->getParentOp());
     }
 
-    /// Return the expected rank of each of the `static_offsets`, `static_sizes`
-    /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrMaxRanks() {
-      unsigned rank = getDestType().getRank();
-      return {rank, rank, rank};
-    }
+    // Return the expected size of each of the `static_offsets`, `static_sizes`
+    // and `static_strides` attributes.
+    int64_t getOffsetsRank() { return getDestType().getRank(); }
+    int64_t getSizesRank() { return getDestType().getRank(); }
+    int64_t getStridesRank() { return getDestType().getRank(); }
 
     /// Return the number of leading operands before `offsets`, `sizes` and
     /// `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 1; }
 
     /// Return the OpResult of the enclosing ForallOp that is
     /// corresponding to this ParallelInsertSliceOp.
     OpResult getTiedOpResult();

mlir/include/mlir/Interfaces/ViewLikeInterface.h

bool sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
     llvm::function_ref<bool(OpFoldResult, OpFoldResult)> cmp);
 
 /// Helper method to compute the number of dynamic entries of `staticVals`,
 /// up to `idx`.
 unsigned getNumDynamicEntriesUpToIdx(ArrayRef<int64_t> staticVals,
                                      unsigned idx);
 
+/// Default implementation for `getOffsetsRank`, `getSizesRank` and
+/// `getStridesRank` of `OffsetSizeAndStrideOpInterface` that tries to infer the
+/// sizes of the corresponding arrays with the `ViewLikeOpInterface`.
+FailureOr<int64_t> defaultArrayRank(Operation *op);
+
 } // namespace detail
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/ViewLikeInterface.h.inc"
 
 namespace mlir {
 

mlir/include/mlir/Interfaces/ViewLikeInterface.td

def ViewLikeOpInterface : OpInterface<"ViewLikeOpInterface"> {
   ];
 }
 
 def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface"> {
   let description = [{
     Common interface for ops that allow specifying mixed dynamic and static
     offsets, sizes and strides variadic operands.
     Ops that implement this interface need to expose the following methods:
-      1. `getArrayAttrMaxRanks` to specify the length of static integer
+      1. `getOffsetSizeAndStrideRank` to specify the length of static integer
           attributes.
       2. `offsets`, `sizes` and `strides` variadic operands.
       3. `static_offsets`, resp. `static_sizes` and `static_strides` integer
           array attributes.
       4. `getOffsetSizeAndStrideStartOperandIndex` method that specifies the
          starting index of the OffsetSizeAndStrideOpInterface operands
 
     The invariants of this interface are:
-      1. `static_offsets`, `static_sizes` and `static_strides` have length
-          exactly `getArrayAttrMaxRanks()`[0] (resp. [1], [2]).
-      2. `offsets`, `sizes` and `strides` have each length at most
-         `getArrayAttrMaxRanks()`[0] (resp. [1], [2]).
+      1. `static_offsets`, `static_sizes` and `static_strides` have a length of
+          exactly `getOffsetsRank`, `getSizesRank` and `getStridesRank`,
+          respectively.
+      2. `offsets`, `sizes` and `strides` have each length of at most
+         `getOffsetsRank`, `getSizesRank` and `getStridesRank`, respectively.
       3. if an entry of `static_offsets` (resp. `static_sizes`,
          `static_strides`) is equal to a special sentinel value, namely
          `ShapedType::kDynamic`, then the corresponding entry is a dynamic
          offset (resp. size, stride).
       4. a variadic `offset` (resp. `sizes`, `strides`) operand  must be present
          for each dynamic offset (resp. size, stride).
       5. `offsets`, `sizes` and `strides` operands are specified in this order
          at operand index starting at `getOffsetSizeAndStrideStartOperandIndex`.
StaticInterfaceMethod<
         and `strides` operands.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getOffsetSizeAndStrideStartOperandIndex",
       /*args=*/(ins)
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Return the expected rank of each of the`static_offsets`, `static_sizes`
-        and `static_strides` attributes.
+        Return the expected size of the `static_offsets` attribute.
       }],
-      /*retTy=*/"std::array<unsigned, 3>",
-      /*methodName=*/"getArrayAttrMaxRanks",
-      /*args=*/(ins)
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getOffsetsRank",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        FailureOr<int64_t> rank = ::mlir::detail::defaultArrayRank($_op);
+        if (succeeded(rank))
+          return *rank;
+        llvm_unreachable("getOffsetsRank not implemented");
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the expected size of the `static_sizes` attribute.
+      }],
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getSizesRank",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        FailureOr<int64_t> rank = ::mlir::detail::defaultArrayRank($_op);
+        if (succeeded(rank))
+          return *rank;
+        llvm_unreachable("getOffsetsRank not implemented");
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return the expected size of the `static_strides` attribute.
+      }],
+      /*retTy=*/"int64_t",
+      /*methodName=*/"getStridesRank",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        FailureOr<int64_t> rank = ::mlir::detail::defaultArrayRank($_op);
+        if (succeeded(rank))
+          return *rank;
+        llvm_unreachable("getOffsetsRank not implemented");
+      }]
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the dynamic offset operands.
       }],
       /*retTy=*/"::mlir::OperandRange",
       /*methodName=*/"getOffsets",
       /*args=*/(ins),

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

return os << "range " << range.offset << ":" << range.size << ":"
             << range.stride;
 }
 
 /// Return the list of Range (i.e. offset, size, stride). Each Range
 /// entry contains either the dynamic value or a ConstantIndexOp constructed
 /// with `b` at location `loc`.
 SmallVector<Range, 8> mlir::getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
                                               OpBuilder &b, Location loc) {
-  std::array<unsigned, 3> ranks = op.getArrayAttrMaxRanks();
-  assert(ranks[0] == ranks[1] && "expected offset and sizes of equal ranks");
-  assert(ranks[1] == ranks[2] && "expected sizes and strides of equal ranks");
+  int64_t offsetsRank = op.getOffsetsRank();
+  assert(offsetsRank == op.getSizesRank() &&
+         "expected offset and sizes of equal ranks");
+  assert(op.getSizesRank() == op.getStridesRank() &&
+         "expected sizes and strides of equal ranks");
   SmallVector<Range, 8> res;
-  unsigned rank = ranks[0];
-  res.reserve(rank);
-  for (unsigned idx = 0; idx < rank; ++idx) {
+  res.reserve(offsetsRank);
+  for (unsigned idx = 0; idx < offsetsRank; ++idx) {
     Value offset =
         op.isDynamicOffset(idx)
             ? op.getDynamicOffset(idx)
             : b.create<arith::ConstantIndexOp>(loc, op.getStaticOffset(idx));
     Value size =
         op.isDynamicSize(idx)
             ? op.getDynamicSize(idx)
             : b.create<arith::ConstantIndexOp>(loc, op.getStaticSize(idx));

mlir/lib/Interfaces/ViewLikeInterface.cpp

LogicalResult mlir::verifyListOfOperandsOrIntegers(Operation *op,
   if (values.size() != expectedNumDynamicEntries)
     return op->emitError("expected ")
            << expectedNumDynamicEntries << " dynamic " << name << " values";
   return success();
 }
 
 LogicalResult
 mlir::detail::verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op) {
-  std::array<unsigned, 3> maxRanks = op.getArrayAttrMaxRanks();
   // Offsets can come in 2 flavors:
-  //   1. Either single entry (when maxRanks == 1).
+  //   1. Either single entry (when getOffsetsRank() == 1).
   //   2. Or as an array whose rank must match that of the mixed sizes.
   // So that the result type is well-formed.
-  if (!(op.getMixedOffsets().size() == 1 && maxRanks[0] == 1) && // NOLINT
+  if (!(op.getMixedOffsets().size() == 1 &&
+        op.getOffsetsRank() == 1) && // NOLINT
       op.getMixedOffsets().size() != op.getMixedSizes().size())
     return op->emitError(
                "expected mixed offsets rank to match mixed sizes rank (")
            << op.getMixedOffsets().size() << " vs " << op.getMixedSizes().size()
            << ") so the rank of the result type is well-formed.";
   // Ranks of mixed sizes and strides must always match so the result type is
   // well-formed.
   if (op.getMixedSizes().size() != op.getMixedStrides().size())
     return op->emitError(
                "expected mixed sizes rank to match mixed strides rank (")
            << op.getMixedSizes().size() << " vs " << op.getMixedStrides().size()
            << ") so the rank of the result type is well-formed.";
 
-  if (failed(verifyListOfOperandsOrIntegers(
-          op, "offset", maxRanks[0], op.getStaticOffsets(), op.getOffsets())))
+  if (failed(verifyListOfOperandsOrIntegers(op, "offset", op.getOffsetsRank(),
+                                            op.getStaticOffsets(),
+                                            op.getOffsets())))
     return failure();
   if (failed(verifyListOfOperandsOrIntegers(
-          op, "size", maxRanks[1], op.getStaticSizes(), op.getSizes())))
+          op, "size", op.getSizesRank(), op.getStaticSizes(), op.getSizes())))
     return failure();
-  if (failed(verifyListOfOperandsOrIntegers(
-          op, "stride", maxRanks[2], op.getStaticStrides(), op.getStrides())))
+  if (failed(verifyListOfOperandsOrIntegers(op, "stride", op.getStridesRank(),
+                                            op.getStaticStrides(),
+                                            op.getStrides())))
     return failure();
   return success();
 }
 
 static char getLeftDelimiter(AsmParser::Delimiter delimiter) {
   switch (delimiter) {
   case AsmParser::Delimiter::Paren:
     return '(';
bool mlir::detail::sameOffsetsSizesAndStrides(
   return true;
 }
 
 unsigned mlir::detail::getNumDynamicEntriesUpToIdx(ArrayRef<int64_t> staticVals,
                                                    unsigned idx) {
   return std::count_if(staticVals.begin(), staticVals.begin() + idx,
                        [&](int64_t val) { return ShapedType::isDynamic(val); });
 }
+
+FailureOr<int64_t> mlir::detail::defaultArrayRank(Operation *op) {
+  if (auto viewLikeOp = dyn_cast<ViewLikeOpInterface>(op))
+    if (auto shapedType =
+            dyn_cast<ShapedType>(viewLikeOp.getViewSource().getType()))
+      if (shapedType.hasRank())
+        return shapedType.getRank();
+  return failure();
+}