[mlir][memref] Verify out-of-bounds access for memref.subview

Author: matthias-springer

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

@@ -1859,11 +1859,11 @@ def SubViewOp : MemRef_OpWithOffsetSizesAndStrides<"subview", [
   ]> {
   let summary = "memref subview operation";
   let description = [{
-    The "subview" operation converts a memref type to another memref type
-    which represents a reduced-size view of the original memref as specified by
-    the operation's offsets, sizes and strides arguments.
+    The `subview` operation converts a memref type to a memref type which
+    represents a reduced-size view of the original memref as specified by the
+    operation's offsets, sizes and strides arguments.
 
-    The SubView operation supports the following arguments:
+    The `subview` operation supports the following arguments:
 
     * source: the "base" memref on which to create a "view" memref.
     * offsets: memref-rank number of offsets into the "base" memref at which to
@@ -1876,118 +1876,65 @@ def SubViewOp : MemRef_OpWithOffsetSizesAndStrides<"subview", [
     The representation based on offsets, sizes and strides support a
     partially-static specification via attributes specified through the
     `static_offsets`, `static_sizes` and `static_strides` arguments. A special
-    sentinel value ShapedType::kDynamic encodes that the corresponding entry has
-    a dynamic value.
+    sentinel value `ShapedType::kDynamic` encodes that the corresponding entry
+    has a dynamic value.
 
-    A subview operation may additionally reduce the rank of the resulting view
-    by removing dimensions that are statically known to be of size 1.
+    A `subview` operation may additionally reduce the rank of the resulting
+    view by removing dimensions that are statically known to be of size 1.
+
+    The offsets, sizes and strides must be in-bounds with respect to the source
+    memref. When possible, the static operation verifier will detect
+    out-of-bounds subviews. Subviews that cannot be confirmed to be in-bounds
+    or out-of-bounds based on compile-time information are valid. However,
+    performing an out-of-bounds subview at runtime is undefined behavior.
 
     Example 1:
 
     ```mlir
-    %0 = memref.alloc() : memref<64x4xf32, affine_map<(d0, d1) -> (d0 * 4 + d1)>>
-
-    // Create a sub-view of "base" memref '%0' with offset arguments '%c0',
-    // dynamic sizes for each dimension, and stride arguments '%c1'.
-    %1 = memref.subview %0[%c0, %c0][%size0, %size1][%c1, %c1]
-      : memref<64x4xf32, affine_map<(d0, d1) -> (d0 * 4 + d1)>> to
-        memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + d1 + s0)>>
+    // Subview of static memref with identity layout at dynamic offsets, sizes
+    // and strides.
+    %1 = memref.subview %0[%off0, %off1][%sz0, %sz1][%str0, %str1]
+        : memref<64x4xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
     ```
 
     Example 2:
 
     ```mlir
-    %0 = memref.alloc() : memref<8x16x4xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>>
-
-    // Create a sub-view of "base" memref '%0' with dynamic offsets, sizes,
+    // Subview of static memref with strided layout at static offsets, sizes
     // and strides.
-    // Note that dynamic offsets are represented by the linearized dynamic
-    // offset symbol 's0' in the subview memref layout map, and that the
-    // dynamic strides operands, after being applied to the base memref
-    // strides in each dimension, are represented in the view memref layout
-    // map as symbols 's1', 's2' and 's3'.
-    %1 = memref.subview %0[%i, %j, %k][%size0, %size1, %size2][%x, %y, %z]
-      : memref<8x16x4xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>> to
-        memref<?x?x?xf32,
-          affine_map<(d0, d1, d2)[s0, s1, s2, s3] -> (d0 * s1 + d1 * s2 + d2 * s3 + s0)>>
+    %1 = memref.subview %0[4, 2][8, 2][3, 2]
+        : memref<64x4xf32, strided<[7, 9], offset: 91>> to
+          memref<8x2xf32, strided<[21, 18], offset: 137>>
     ```
 
-    Example 3:
+    Example 4:
 
     ```mlir
-    %0 = memref.alloc() : memref<8x16x4xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>>
-
-    // Subview with constant offsets, sizes and strides.
-    %1 = memref.subview %0[0, 2, 0][4, 4, 4][1, 1, 1]
-      : memref<8x16x4xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)>> to
-        memref<4x4x4xf32, affine_map<(d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2 + 8)>>
+    // Subview of dynamic memref with strided layout at dynamic offsets and
+    // strides, but static sizes.
+    %1 = memref.subview %0[%off0, %off1][4, 4][%str0, %str1]
+        : memref<?x?xf32, strided<[?, ?], offset: ?>> to
+          memref<4x4xf32, strided<[?, ?], offset: ?>>
     ```
 
-    Example 4:
+    Example 5:
 
     ```mlir
-    %0 = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
-
-    // Subview with constant size, but dynamic offsets and
-    // strides. The resulting memref has a static shape, but if the
-    // base memref has an affine map to describe the layout, the result
-    // memref also uses an affine map to describe the layout. The
-    // strides of the result memref is computed as follows:
-    //
-    // Let #map1 represents the layout of the base memref, and #map2
-    // represents the layout of the result memref. A #mapsubview can be
-    // constructed to map an index from the result memref to the base
-    // memref (note that the description below uses more convenient
-    // naming for symbols, while in affine maps, symbols are
-    // represented as unsigned numbers that identify that symbol in the
-    // given affine map.
-    //
-    // #mapsubview = (d0, d1)[o0, o1, t0, t1] -> (d0 * t0 + o0, d1 * t1 + o1)
-    //
-    // where, o0, o1, ... are offsets, and t0, t1, ... are strides. Then,
-    //
-    // #map2 = #map1.compose(#mapsubview)
-    //
-    // If the layout map is represented as
-    //
-    // #map1 = (d0, d1)[s0, s1, s2] -> (d0 * s1 + d1 * s2 + s0)
-    //
-    // then,
-    //
-    // #map2 = (d0, d1)[s0, s1, s2, o0, o1, t0, t1] ->
-    //              (d0 * s1 * t0 + d1 * s2 * t1 + o0 * s1 + o1 * s2 + s0)
-    //
-    // Representing this canonically
-    //
-    // #map2 = (d0, d1)[r0, r1, r2] -> (d0 * r1 + d1 * r2 + r0)
-    //
-    // where, r0 = o0 * s1 + o1 * s2 + s0, r1 = s1 * t0, r2 = s2 * t1.
-    %1 = memref.subview %0[%i, %j][4, 4][%x, %y] :
-      : memref<?x?xf32, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + d1 * s2 + s0)>> to
-        memref<4x4xf32, affine_map<(d0, d1)[r0, r1, r2] -> (d0 * r1 + d1 * r2 + r0)>>
-
-    // Note that the subview op does not guarantee that the result
-    // memref is "inbounds" w.r.t to base memref. It is upto the client
-    // to ensure that the subview is accessed in a manner that is
-    // in-bounds.
+    // Rank-reducing subviews.
+    %1 = memref.subview %0[0, 0, 0][1, 16, 4][1, 1, 1]
+        : memref<8x16x4xf32> to memref<16x4xf32>
+    %3 = memref.subview %2[3, 4, 2][1, 6, 3][1, 1, 1]
+        : memref<8x16x4xf32> to memref<6x3xf32, strided<[4, 1], offset: 210>>
     ```
 
-    Example 5:
-
+    Example 6:
+    
     ```mlir
-    // Rank-reducing subview.
-    %1 = memref.subview %0[0, 0, 0][1, 16, 4][1, 1, 1] :
-      memref<8x16x4xf32> to memref<16x4xf32>
-
-    // Original layout:
-    // (d0, d1, d2) -> (64 * d0 + 16 * d1 + d2)
-    // Subviewed layout:
-    // (d0, d1, d2) -> (64 * (d0 + 3) + 4 * (d1 + 4) + d2 + 2) = (64 * d0 + 4 * d1 + d2 + 210)
-    // After rank reducing:
-    // (d0, d1) -> (4 * d0 + d1 + 210)
-    %3 = memref.subview %2[3, 4, 2][1, 6, 3][1, 1, 1] :
-      memref<8x16x4xf32> to memref<6x3xf32, strided<[4, 1], offset: 210>>
+    // Identity subview. The subview is the full source memref.
+    %1 = memref.subview %0[0, 0, 0] [8, 16, 4] [1, 1, 1]
+        : memref<8x16x4xf32> to memref<8x16x4xf32>
     ```
+
   }];
 
   let arguments = (ins AnyMemRef:$source,

mlir/include/mlir/Interfaces/ViewLikeInterface.h

@@ -45,6 +45,28 @@ unsigned getNumDynamicEntriesUpToIdx(ArrayRef<int64_t> staticVals,
 
 namespace mlir {
 
+/// Result for slice bounds verification;
+struct SliceBoundsVerificationResult {
+  /// If set to "true", the slice bounds verification was successful.
+  bool isValid;
+  /// An error message that can be printed during op verification.
+  std::string errorMessage;
+};
+
+/// Verify that the offsets/sizes/strides-style access into the given shape
+/// is in-bounds. Only static values are verified. If `generateErrorMessage`
+/// is set to "true", an error message is produced that can be printed by the
+///  op verifier.
+SliceBoundsVerificationResult
+verifyInBoundsSlice(ArrayRef<int64_t> shape, ArrayRef<int64_t> staticOffsets,
+                    ArrayRef<int64_t> staticSizes,
+                    ArrayRef<int64_t> staticStrides,
+                    bool generateErrorMessage = false);
+SliceBoundsVerificationResult verifyInBoundsSlice(
+    ArrayRef<int64_t> shape, ArrayRef<OpFoldResult> mixedOffsets,
+    ArrayRef<OpFoldResult> mixedSizes, ArrayRef<OpFoldResult> mixedStrides,
+    bool generateErrorMessage = false);
+
 /// Pattern to rewrite dynamic offsets/sizes/strides of view/slice-like ops as
 /// constant arguments. This pattern assumes that the op has a suitable builder
 /// that takes a result type, a "source" operand and mixed offsets, sizes and
@@ -72,11 +94,20 @@ class OpWithOffsetSizesAndStridesConstantArgumentFolder final
         failed(foldDynamicIndexList(mixedStrides)))
       return failure();
 
-    // Create the new op in canonical form.
+    // Pattern does not apply if the produced op would not verify.
+    SliceBoundsVerificationResult sliceResult = verifyInBoundsSlice(
+        cast<ShapedType>(op.getSource().getType()).getShape(), mixedOffsets,
+        mixedSizes, mixedStrides);
+    if (!sliceResult.isValid)
+      return failure();
+
+    // Compute the new result type.
     auto resultType =
         ResultTypeFn()(op, mixedOffsets, mixedSizes, mixedStrides);
     if (!resultType)
       return failure();
+
+    // Create the new op in canonical form.
     auto newOp =
         rewriter.create<OpType>(op.getLoc(), resultType, op.getSource(),
                                 mixedOffsets, mixedSizes, mixedStrides);

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

@@ -2977,6 +2977,9 @@ static LogicalResult produceSubViewErrorMsg(SliceVerificationResult result,
 LogicalResult SubViewOp::verify() {
   MemRefType baseType = getSourceType();
   MemRefType subViewType = getType();
+  ArrayRef<int64_t> staticOffsets = getStaticOffsets();
+  ArrayRef<int64_t> staticSizes = getStaticSizes();
+  ArrayRef<int64_t> staticStrides = getStaticStrides();
 
   // The base memref and the view memref should be in the same memory space.
   if (baseType.getMemorySpace() != subViewType.getMemorySpace())
@@ -2991,7 +2994,7 @@ LogicalResult SubViewOp::verify() {
   // Compute the expected result type, assuming that there are no rank
   // reductions.
   MemRefType expectedType = SubViewOp::inferResultType(
-      baseType, getStaticOffsets(), getStaticSizes(), getStaticStrides());
+      baseType, staticOffsets, staticSizes, staticStrides);
 
   // Verify all properties of a shaped type: rank, element type and dimension
   // sizes. This takes into account potential rank reductions.
@@ -3025,6 +3028,14 @@ LogicalResult SubViewOp::verify() {
     return produceSubViewErrorMsg(SliceVerificationResult::LayoutMismatch,
                                   *this, expectedType);
 
+  // Verify that offsets, sizes, strides do not run out-of-bounds with respect
+  // to the base memref.
+  SliceBoundsVerificationResult boundsResult =
+      verifyInBoundsSlice(baseType.getShape(), staticOffsets, staticSizes,
+                          staticStrides, /*generateErrorMessage=*/true);
+  if (!boundsResult.isValid)
+    return getOperation()->emitError(boundsResult.errorMessage);
+
   return success();
 }
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

@@ -27,6 +27,7 @@
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
@@ -2352,37 +2353,6 @@ static LogicalResult produceSliceErrorMsg(SliceVerificationResult result,
   }
 }
 
-/// Verify that the offsets/sizes/strides-style access into the given tensor
-/// is in-bounds. Only static information is verified.
-static LogicalResult verifyInBoundsSlice(Operation *op,
-                                         RankedTensorType tensorType,
-                                         ArrayRef<int64_t> staticOffsets,
-                                         ArrayRef<int64_t> staticSizes,
-                                         ArrayRef<int64_t> staticStrides) {
-  for (int64_t i = 0, e = tensorType.getRank(); i < e; ++i) {
-    // Nothing to verify for dynamic source dims.
-    if (tensorType.isDynamicDim(i))
-      continue;
-    // Nothing to verify if the offset is dynamic.
-    if (ShapedType::isDynamic(staticOffsets[i]))
-      continue;
-    if (staticOffsets[i] >= tensorType.getDimSize(i))
-      return op->emitOpError("offset ")
-             << i << " is out-of-bounds: " << staticOffsets[i]
-             << " >= " << tensorType.getDimSize(i);
-    if (ShapedType::isDynamic(staticSizes[i]) ||
-        ShapedType::isDynamic(staticStrides[i]))
-      continue;
-    int64_t lastPos =
-        staticOffsets[i] + (staticSizes[i] - 1) * staticStrides[i];
-    if (lastPos >= tensorType.getDimSize(i))
-      return op->emitOpError("slice along dimension ")
-             << i << " runs out-of-bounds: " << lastPos
-             << " >= " << tensorType.getDimSize(i);
-  }
-  return success();
-}
-
 /// Verifier for ExtractSliceOp.
 LogicalResult ExtractSliceOp::verify() {
   RankedTensorType sourceType = getSourceType();
@@ -2396,8 +2366,13 @@ LogicalResult ExtractSliceOp::verify() {
 
   // Verify that offsets, sizes, strides do not run out-of-bounds with respect
   // to the source tensor.
-  return verifyInBoundsSlice(getOperation(), sourceType, getStaticOffsets(),
-                             getStaticSizes(), getStaticStrides());
+  SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice(
+      sourceType.getShape(), getStaticOffsets(), getStaticSizes(),
+      getStaticStrides(), /*generateErrorMessage=*/true);
+  if (!boundsResult.isValid)
+    return getOperation()->emitError(boundsResult.errorMessage);
+
+  return success();
 }
 
 llvm::SmallBitVector ExtractSliceOp::getDroppedDims() {
@@ -2777,9 +2752,14 @@ LogicalResult InsertSliceOp::verify() {
     return produceSliceErrorMsg(result, *this, expectedType);
 
   // Verify that offsets, sizes, strides do not run out-of-bounds with respect
-  // to the source tensor.
-  return verifyInBoundsSlice(getOperation(), getDestType(), getStaticOffsets(),
-                             getStaticSizes(), getStaticStrides());
+  // to the destination tensor.
+  SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice(
+      getDestType().getShape(), getStaticOffsets(), getStaticSizes(),
+      getStaticStrides(), /*generateErrorMessage=*/true);
+  if (!boundsResult.isValid)
+    return getOperation()->emitError(boundsResult.errorMessage);
+
+  return success();
 }
 
 /// If we have two consecutive InsertSliceOp writing to the same slice, we
@@ -2874,6 +2854,13 @@ class InsertSliceOpConstantArgumentFolder final
         failed(foldDynamicStrideList(mixedStrides)))
       return failure();
 
+    // Pattern does not apply if the produced op would not verify.
+    SliceBoundsVerificationResult sliceResult =
+        verifyInBoundsSlice(insertSliceOp.getDest().getType().getShape(),
+                            mixedOffsets, mixedSizes, mixedStrides);
+    if (!sliceResult.isValid)
+      return failure();
+
     // Create the new op in canonical form.
     auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType(
         insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(),
@@ -3802,9 +3789,14 @@ LogicalResult ParallelInsertSliceOp::verify() {
     return produceSliceErrorMsg(result, *this, expectedType);
 
   // Verify that offsets, sizes, strides do not run out-of-bounds with respect
-  // to the source tensor.
-  return verifyInBoundsSlice(getOperation(), getDestType(), getStaticOffsets(),
-                             getStaticSizes(), getStaticStrides());
+  // to the destination tensor.
+  SliceBoundsVerificationResult boundsResult = verifyInBoundsSlice(
+      getDestType().getShape(), getStaticOffsets(), getStaticSizes(),
+      getStaticStrides(), /*generateErrorMessage=*/true);
+  if (!boundsResult.isValid)
+    return getOperation()->emitError(boundsResult.errorMessage);
+
+  return success();
 }
 
 void ParallelInsertSliceOp::getCanonicalizationPatterns(