[mlir] Extend semantic of OffsetSizeAndStrideOpInterface.

OffsetSizeAndStrideOpInterface now have the ability to specify only a leading subset of
offset, sizes, strides operands/attributes.
The size of that leading subset must be limited by the corresponding entry in `getArrayAttrMaxRanks` to avoid overflows.
Missing trailing dimensions are assumed to span the whole range (i.e. [0 .. dim)).
This brings more natural semantics to slice-like op on top of subview and is a simplifies to removing all uses of SliceOp in dependent projects.

Differential revision: https://reviews.llvm.org/D95441

Author: nicolasvasilache

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

@@ -1986,7 +1986,7 @@ def MemRefReinterpretCastOp:
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrRanks() {
+    std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned resultRank = getResult().getType().cast<ShapedType>().getRank();
       return {1, resultRank, resultRank};
     }
@@ -2983,7 +2983,7 @@ def SubViewOp : BaseOpWithOffsetSizesAndStrides<
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrRanks() {
+    std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
@@ -3097,7 +3097,7 @@ def SubTensorOp : BaseOpWithOffsetSizesAndStrides<
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrRanks() {
+    std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned rank = getSourceType().getRank();
       return {rank, rank, rank};
     }
@@ -3184,7 +3184,7 @@ def SubTensorInsertOp : BaseOpWithOffsetSizesAndStrides<
 
     /// Return the expected rank of each of the`static_offsets`, `static_sizes`
     /// and `static_strides` attributes.
-    std::array<unsigned, 3> getArrayAttrRanks() {
+    std::array<unsigned, 3> getArrayAttrMaxRanks() {
       unsigned rank = getType().getRank();
       return {rank, rank, rank};
     }

mlir/include/mlir/Interfaces/ViewLikeInterface.td

@@ -35,7 +35,7 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
     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. `getArrayAttrRanks` to specify the length of static integer
+      1. `getArrayAttrMaxRanks` 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
@@ -45,9 +45,9 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
 
     The invariants of this interface are:
       1. `static_offsets`, `static_sizes` and `static_strides` have length
-          exactly `getArrayAttrRanks()`[0] (resp. [1], [2]).
+          at most `getArrayAttrMaxRanks()`[0] (resp. [1], [2]).
       2. `offsets`, `sizes` and `strides` have each length at most
-         `getArrayAttrRanks()`[0] (resp. [1], [2]).
+         length `static_offsets` (resp. `static_sizes`, `static_strides`).
       3. if an entry of `static_offsets` (resp. `static_sizes`,
          `static_strides`) is equal to a special sentinel value, namely
          `ShapedType::kDynamicStrideOrOffset` (resp. `ShapedType::kDynamicSize`,
@@ -81,7 +81,7 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
         and `static_strides` attributes.
       }],
       /*retTy=*/"std::array<unsigned, 3>",
-      /*methodName=*/"getArrayAttrRanks",
+      /*methodName=*/"getArrayAttrMaxRanks",
       /*args=*/(ins)
     >,
     InterfaceMethod<
@@ -166,9 +166,8 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         SmallVector<OpFoldResult, 4> res;
-        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();
         unsigned numDynamic = 0;
-        unsigned count = ranks[getOffsetOperandGroupPosition()];
+        unsigned count = $_op.static_offsets().size();
         for (unsigned idx = 0; idx < count; ++idx) {
           if (isDynamicOffset(idx))
             res.push_back($_op.offsets()[numDynamic++]);
@@ -188,9 +187,8 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         SmallVector<OpFoldResult, 4> res;
-        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();
         unsigned numDynamic = 0;
-        unsigned count = ranks[getSizeOperandGroupPosition()];
+        unsigned count = $_op.static_sizes().size();
         for (unsigned idx = 0; idx < count; ++idx) {
           if (isDynamicSize(idx))
             res.push_back($_op.sizes()[numDynamic++]);
@@ -210,9 +208,8 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         SmallVector<OpFoldResult, 4> res;
-        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();
         unsigned numDynamic = 0;
-        unsigned count = ranks[getStrideOperandGroupPosition()];
+        unsigned count = $_op.static_strides().size();
         for (unsigned idx = 0; idx < count; ++idx) {
           if (isDynamicStride(idx))
             res.push_back($_op.strides()[numDynamic++]);

mlir/integration_test/Dialect/Standard/CPU/test_subview.mlir

@@ -0,0 +1,63 @@
+// RUN: mlir-opt %s -convert-std-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext | FileCheck %s
+
+global_memref "private" constant @__constant_5x3xf32 : memref<5x3xf32> =
+dense<[[0.0, 1.0, 2.0],
+       [3.0, 4.0, 5.0],
+       [6.0, 7.0, 8.0],
+       [9.0, 10.0, 11.0],
+       [12.0, 13.0, 14.0]]>
+
+func @main() {
+  %0 = get_global_memref @__constant_5x3xf32 : memref<5x3xf32>
+
+  /// Subview with only leading operands.
+  %1 = subview %0[2][3][1]: memref<5x3xf32> to memref<3x3xf32, offset: 6, strides: [3, 1]>
+  %unranked = memref_cast %1 : memref<3x3xf32, offset: 6, strides: [3, 1]> to memref<*xf32>
+  call @print_memref_f32(%unranked) : (memref<*xf32>) -> ()
+
+  //      CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
+  // CHECK-SAME: rank = 2 offset = 6 sizes = [3, 3] strides = [3, 1] data =
+  // CHECK-NEXT: [
+  // CHECK-SAME:  [6,   7,   8],
+  // CHECK-NEXT:  [9,   10,   11],
+  // CHECK-NEXT:  [12,   13,   14]
+  // CHECK-SAME: ]
+
+  /// Regular subview.
+  %2 = subview %0[0, 2][5, 1][1, 1]: memref<5x3xf32> to memref<5x1xf32, offset: 2, strides: [3, 1]>
+  %unranked2 = memref_cast %2 : memref<5x1xf32, offset: 2, strides: [3, 1]> to memref<*xf32>
+  call @print_memref_f32(%unranked2) : (memref<*xf32>) -> ()
+
+  //      CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
+  // CHECK-SAME: rank = 2 offset = 2 sizes = [5, 1] strides = [3, 1] data =
+  // CHECK-NEXT: [
+  // CHECK-SAME:  [2],
+  // CHECK-NEXT:  [5],
+  // CHECK-NEXT:  [8],
+  // CHECK-NEXT:  [11],
+  // CHECK-NEXT:  [14]
+  // CHECK-SAME: ]
+
+  /// Rank-reducing subview.
+  %3 = subview %0[0, 2][5, 1][1, 1]: memref<5x3xf32> to memref<5xf32, offset: 2, strides: [3]>
+  %unranked3 = memref_cast %3 : memref<5xf32, offset: 2, strides: [3]> to memref<*xf32>
+  call @print_memref_f32(%unranked3) : (memref<*xf32>) -> ()
+
+  //      CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
+  // CHECK-SAME: rank = 1 offset = 2 sizes = [5] strides = [3] data =
+  // CHECK-NEXT: [2,  5,  8,  11,  14]
+
+  /// Rank-reducing subview with only leading operands.
+  %4 = subview %0[1][1][1]: memref<5x3xf32> to memref<3xf32, offset: 3, strides: [1]>
+  %unranked4 = memref_cast %4 : memref<3xf32, offset: 3, strides: [1]> to memref<*xf32>
+  call @print_memref_f32(%unranked4) : (memref<*xf32>) -> ()
+  //      CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
+  // CHECK-SAME: rank = 1 offset = 3 sizes = [3] strides = [1] data =
+  // CHECK-NEXT: [3,  4,  5]
+
+  return
+}
+
+func private @print_memref_f32(%ptr : memref<*xf32>)

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

@@ -3336,8 +3336,14 @@ struct SubViewOpLowering : public ConvertOpToLLVMPattern<SubViewOp> {
       targetMemRef.setConstantOffset(rewriter, loc, offset);
     } else {
       Value baseOffset = sourceMemRef.offset(rewriter, loc);
-      for (unsigned i = 0; i < inferredShapeRank; ++i) {
+      // `inferredShapeRank` may be larger than the number of offset operands
+      // because of trailing semantics. In this case, the offset is guaranteed
+      // to be interpreted as 0 and we can just skip the extra dimensions.
+      for (unsigned i = 0, e = std::min(inferredShapeRank,
+                                        subViewOp.getMixedOffsets().size());
+           i < e; ++i) {
         Value offset =
+            // TODO: need OpFoldResult ODS adaptor to clean this up.
             subViewOp.isDynamicOffset(i)
                 ? operands[subViewOp.getIndexOfDynamicOffset(i)]
                 : rewriter.create<LLVM::ConstantOp>(
@@ -3350,31 +3356,47 @@ struct SubViewOpLowering : public ConvertOpToLLVMPattern<SubViewOp> {
     }
 
     // Update sizes and strides.
+    SmallVector<OpFoldResult> mixedSizes = subViewOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = subViewOp.getMixedStrides();
+    assert(mixedSizes.size() == mixedStrides.size() &&
+           "expected sizes and strides of equal length");
     for (int i = inferredShapeRank - 1, j = resultShapeRank - 1;
          i >= 0 && j >= 0; --i) {
       if (!mask[i])
         continue;
 
-      Value size =
-          subViewOp.isDynamicSize(i)
-              ? operands[subViewOp.getIndexOfDynamicSize(i)]
-              : rewriter.create<LLVM::ConstantOp>(
-                    loc, llvmIndexType,
-                    rewriter.getI64IntegerAttr(subViewOp.getStaticSize(i)));
-      targetMemRef.setSize(rewriter, loc, j, size);
-      Value stride;
-      if (!ShapedType::isDynamicStrideOrOffset(strides[i])) {
+      // `i` may overflow subViewOp.getMixedSizes because of trailing semantics.
+      // In this case, the size is guaranteed to be interpreted as Dim and the
+      // stride as 1.
+      Value size, stride;
+      if (static_cast<unsigned>(i) >= mixedSizes.size()) {
+        size = rewriter.create<LLVM::DialectCastOp>(
+            loc, llvmIndexType,
+            rewriter.create<DimOp>(loc, subViewOp.source(), i));
         stride = rewriter.create<LLVM::ConstantOp>(
-            loc, llvmIndexType, rewriter.getI64IntegerAttr(strides[i]));
+            loc, llvmIndexType, rewriter.getI64IntegerAttr(1));
       } else {
-        stride =
-            subViewOp.isDynamicStride(i)
-                ? operands[subViewOp.getIndexOfDynamicStride(i)]
+        // TODO: need OpFoldResult ODS adaptor to clean this up.
+        size =
+            subViewOp.isDynamicSize(i)
+                ? operands[subViewOp.getIndexOfDynamicSize(i)]
                 : rewriter.create<LLVM::ConstantOp>(
                       loc, llvmIndexType,
-                      rewriter.getI64IntegerAttr(subViewOp.getStaticStride(i)));
-        stride = rewriter.create<LLVM::MulOp>(loc, stride, strideValues[i]);
+                      rewriter.getI64IntegerAttr(subViewOp.getStaticSize(i)));
+        if (!ShapedType::isDynamicStrideOrOffset(strides[i])) {
+          stride = rewriter.create<LLVM::ConstantOp>(
+              loc, llvmIndexType, rewriter.getI64IntegerAttr(strides[i]));
+        } else {
+          stride = subViewOp.isDynamicStride(i)
+                       ? operands[subViewOp.getIndexOfDynamicStride(i)]
+                       : rewriter.create<LLVM::ConstantOp>(
+                             loc, llvmIndexType,
+                             rewriter.getI64IntegerAttr(
+                                 subViewOp.getStaticStride(i)));
+          stride = rewriter.create<LLVM::MulOp>(loc, stride, strideValues[i]);
+        }
       }
+      targetMemRef.setSize(rewriter, loc, j, size);
       targetMemRef.setStride(rewriter, loc, j, stride);
       j--;
     }

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

@@ -712,6 +712,10 @@ static LogicalResult verify(InitTensorOp op) {
                                             ShapedType::isDynamic)))
     return failure();
 
+  if (op.static_sizes().size() != static_cast<unsigned>(resultType.getRank()))
+    return op->emitError("expected ")
+           << resultType.getRank() << " sizes values";
+
   Type expectedType =
       InitTensorOp::inferResultType(staticSizes, resultType.getElementType());
   if (resultType != expectedType) {

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

@@ -2209,6 +2209,8 @@ static ParseResult parseMemRefReinterpretCastOp(OpAsmParser &parser,
   return parser.addTypeToList(dstType, result.types);
 }
 
+// TODO: ponder whether we want to allow missing trailing sizes/strides that are
+// completed automatically, like we have for subview and subtensor.
 static LogicalResult verify(MemRefReinterpretCastOp op) {
   // The source and result memrefs should be in the same memory space.
   auto srcType = op.source().getType().cast<BaseMemRefType>();
@@ -2833,16 +2835,28 @@ Wrapper operator*(Wrapper a, int64_t b) {
 /// static representation of offsets, sizes and strides. Special sentinels
 /// encode the dynamic case.
 Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
-                                ArrayRef<int64_t> staticOffsets,
-                                ArrayRef<int64_t> staticSizes,
-                                ArrayRef<int64_t> staticStrides) {
+                                ArrayRef<int64_t> leadingStaticOffsets,
+                                ArrayRef<int64_t> leadingStaticSizes,
+                                ArrayRef<int64_t> leadingStaticStrides) {
+  // A subview may specify only a leading subset of offset/sizes/strides in
+  // which case we complete with offset=0, sizes from memref type and strides=1.
   unsigned rank = sourceMemRefType.getRank();
-  (void)rank;
-  assert(staticOffsets.size() == rank &&
-         "unexpected staticOffsets size mismatch");
-  assert(staticSizes.size() == rank && "unexpected staticSizes size mismatch");
-  assert(staticStrides.size() == rank &&
-         "unexpected staticStrides size mismatch");
+  assert(leadingStaticOffsets.size() <= rank &&
+         "unexpected leadingStaticOffsets overflow");
+  assert(leadingStaticSizes.size() <= rank &&
+         "unexpected leadingStaticSizes overflow");
+  assert(leadingStaticStrides.size() <= rank &&
+         "unexpected leadingStaticStrides overflow");
+  auto staticOffsets = llvm::to_vector<4>(leadingStaticOffsets);
+  auto staticSizes = llvm::to_vector<4>(leadingStaticSizes);
+  auto staticStrides = llvm::to_vector<4>(leadingStaticStrides);
+  unsigned numTrailingOffsets = rank - staticOffsets.size();
+  unsigned numTrailingSizes = rank - staticSizes.size();
+  unsigned numTrailingStrides = rank - staticStrides.size();
+  staticOffsets.append(numTrailingOffsets, 0);
+  llvm::append_range(staticSizes,
+                     sourceMemRefType.getShape().take_back(numTrailingSizes));
+  staticStrides.append(numTrailingStrides, 1);
 
   // Extract source offset and strides.
   int64_t sourceOffset;
@@ -3197,7 +3211,7 @@ raw_ostream &mlir::operator<<(raw_ostream &os, Range &range) {
 /// with `b` at location `loc`.
 SmallVector<Range, 8> mlir::getOrCreateRanges(OffsetSizeAndStrideOpInterface op,
                                               OpBuilder &b, Location loc) {
-  std::array<unsigned, 3> ranks = op.getArrayAttrRanks();
+  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");
   SmallVector<Range, 8> res;
@@ -3484,16 +3498,18 @@ static ParseResult parseSubTensorOp(OpAsmParser &parser,
 /// static representation of offsets, sizes and strides. Special sentinels
 /// encode the dynamic case.
 Type SubTensorOp::inferResultType(RankedTensorType sourceRankedTensorType,
-                                  ArrayRef<int64_t> staticOffsets,
-                                  ArrayRef<int64_t> staticSizes,
-                                  ArrayRef<int64_t> staticStrides) {
+                                  ArrayRef<int64_t> leadingStaticOffsets,
+                                  ArrayRef<int64_t> leadingStaticSizes,
+                                  ArrayRef<int64_t> leadingStaticStrides) {
+  // A subtensor may specify only a leading subset of offset/sizes/strides in
+  // which case we complete with offset=0, sizes from memref type and strides=1.
   unsigned rank = sourceRankedTensorType.getRank();
-  (void)rank;
-  assert(staticOffsets.size() == rank &&
-         "unexpected staticOffsets size mismatch");
-  assert(staticSizes.size() == rank && "unexpected staticSizes size mismatch");
-  assert(staticStrides.size() == rank &&
-         "unexpected staticStrides size mismatch");
+  assert(leadingStaticSizes.size() <= rank &&
+         "unexpected leadingStaticSizes overflow");
+  auto staticSizes = llvm::to_vector<4>(leadingStaticSizes);
+  unsigned numTrailingSizes = rank - staticSizes.size();
+  llvm::append_range(staticSizes, sourceRankedTensorType.getShape().take_back(
+                                      numTrailingSizes));
   return RankedTensorType::get(staticSizes,
                                sourceRankedTensorType.getElementType());
 }

mlir/lib/Interfaces/ViewLikeInterface.cpp

@@ -18,12 +18,12 @@ using namespace mlir;
 #include "mlir/Interfaces/ViewLikeInterface.cpp.inc"
 
 LogicalResult mlir::verifyListOfOperandsOrIntegers(
-    Operation *op, StringRef name, unsigned expectedNumElements, ArrayAttr attr,
+    Operation *op, StringRef name, unsigned maxNumElements, ArrayAttr attr,
     ValueRange values, llvm::function_ref<bool(int64_t)> isDynamic) {
-  /// Check static and dynamic offsets/sizes/strides breakdown.
-  if (attr.size() != expectedNumElements)
-    return op->emitError("expected ")
-           << expectedNumElements << " " << name << " values";
+  /// Check static and dynamic offsets/sizes/strides does not overflow type.
+  if (attr.size() > maxNumElements)
+    return op->emitError("expected <= ")
+           << maxNumElements << " " << name << " values";
   unsigned expectedNumDynamicEntries =
       llvm::count_if(attr.getValue(), [&](Attribute attr) {
         return isDynamic(attr.cast<IntegerAttr>().getInt());
@@ -35,17 +35,35 @@ LogicalResult mlir::verifyListOfOperandsOrIntegers(
 }
 
 LogicalResult mlir::verify(OffsetSizeAndStrideOpInterface op) {
-  std::array<unsigned, 3> ranks = op.getArrayAttrRanks();
+  std::array<unsigned, 3> maxRanks = op.getArrayAttrMaxRanks();
+  // Offsets can come in 2 flavors:
+  //   1. Either single entry (when maxRanks == 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) &&
+      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", ranks[0], op.static_offsets(), op.offsets(),
+          op, "offset", maxRanks[0], op.static_offsets(), op.offsets(),
           ShapedType::isDynamicStrideOrOffset)))
     return failure();
-  if (failed(verifyListOfOperandsOrIntegers(op, "size", ranks[1],
+  if (failed(verifyListOfOperandsOrIntegers(op, "size", maxRanks[1],
                                             op.static_sizes(), op.sizes(),
                                             ShapedType::isDynamic)))
     return failure();
   if (failed(verifyListOfOperandsOrIntegers(
-          op, "stride", ranks[2], op.static_strides(), op.strides(),
+          op, "stride", maxRanks[2], op.static_strides(), op.strides(),
           ShapedType::isDynamicStrideOrOffset)))
     return failure();
   return success();