Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir] Verify non-negative offset and size #72059

Merged
merged 13 commits into from
Nov 16, 2023
1 change: 1 addition & 0 deletions mlir/include/mlir/Interfaces/ViewLikeInterface.td
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
for each dynamic offset (resp. size, stride).
5. `offsets`, `sizes` and `strides` operands are specified in this order
at operand index starting at `getOffsetSizeAndStrideStartOperandIndex`.
6. `offsets` and `sizes` operands are non-negative.

This interface is useful to factor out common behavior and provide support
for carrying or injecting static behavior through the use of the static
Expand Down
18 changes: 16 additions & 2 deletions mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2621,6 +2621,17 @@ Type SubViewOp::inferResultType(MemRefType sourceMemRefType,
dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets);
dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes);
dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides);

// If one of the offsets or sizes is invalid, fail the canonicalization.
// These checks also occur in the verifier, but they are needed here
// because some dynamic dimensions may have been constant folded.
for (int64_t offset : staticOffsets)
if (offset < 0 && !ShapedType::isDynamic(offset))
return {};
for (int64_t size : staticSizes)
if (size < 0 && !ShapedType::isDynamic(size))
return {};

return SubViewOp::inferResultType(sourceMemRefType, staticOffsets,
staticSizes, staticStrides);
}
Expand Down Expand Up @@ -3094,8 +3105,11 @@ struct SubViewReturnTypeCanonicalizer {
ArrayRef<OpFoldResult> mixedSizes,
ArrayRef<OpFoldResult> mixedStrides) {
// Infer a memref type without taking into account any rank reductions.
MemRefType nonReducedType = cast<MemRefType>(SubViewOp::inferResultType(
op.getSourceType(), mixedOffsets, mixedSizes, mixedStrides));
auto resTy = SubViewOp::inferResultType(op.getSourceType(), mixedOffsets,
mixedSizes, mixedStrides);
if (!resTy)
return {};
MemRefType nonReducedType = cast<MemRefType>(resTy);

// Directly return the non-rank reduced type if there are no dropped dims.
llvm::SmallBitVector droppedDims = op.getDroppedDims();
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1261,7 +1261,7 @@ struct StaticTensorGenerate : public OpRewritePattern<GenerateOp> {

for (int64_t newdim : newShape) {
// This check also occurs in the verifier, but we need it here too
// since intermediate passes may have some replaced dynamic dimensions
// since intermediate passes may have replaced some dynamic dimensions
// by constants.
if (newdim < 0 && !ShapedType::isDynamic(newdim))
return failure();
Expand Down
11 changes: 11 additions & 0 deletions mlir/lib/Interfaces/ViewLikeInterface.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,17 @@ mlir::detail::verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op) {
if (failed(verifyListOfOperandsOrIntegers(
op, "stride", maxRanks[2], op.getStaticStrides(), op.getStrides())))
return failure();

for (int64_t offset : op.getStaticOffsets()) {
if (offset < 0 && !ShapedType::isDynamic(offset))
return op->emitError("expected offsets to be non-negative, but got ")
<< offset;
}
for (int64_t size : op.getStaticSizes()) {
if (size < 0 && !ShapedType::isDynamic(size))
return op->emitError("expected sizes to be non-negative, but got ")
<< size;
}
return success();
}

Expand Down
11 changes: 11 additions & 0 deletions mlir/test/Dialect/MemRef/canonicalize.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,17 @@ func.func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {

// -----

// CHECK-LABEL: func @no_fold_subview_negative_size
// CHECK: %[[SUBVIEW:.+]] = memref.subview
// CHECK: return %[[SUBVIEW]]
func.func @no_fold_subview_negative_size(%input: memref<4x1024xf32>) -> memref<?x256xf32, strided<[1024, 1], offset: 2304>> {
%cst = arith.constant -13 : index
%0 = memref.subview %input[2, 256] [%cst, 256] [1, 1] : memref<4x1024xf32> to memref<?x256xf32, strided<[1024, 1], offset: 2304>>
return %0 : memref<?x256xf32, strided<[1024, 1], offset: 2304>>
}

// -----

// CHECK-LABEL: func @no_fold_of_store
// CHECK: %[[cst:.+]] = memref.cast %arg
// CHECK: memref.store %[[cst]]
Expand Down
16 changes: 16 additions & 0 deletions mlir/test/Dialect/MemRef/invalid.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -611,6 +611,22 @@ func.func @invalid_view(%arg0 : index, %arg1 : index, %arg2 : index) {

// -----

func.func @invalid_subview(%input: memref<4x1024xf32>) -> memref<2x256xf32, strided<[1024, 1], offset: 2304>> {
// expected-error@+1 {{expected offsets to be non-negative, but got -1}}
%0 = memref.subview %input[-1, 256] [2, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, strided<[1024, 1], offset: 2304>>
return %0 : memref<2x256xf32, strided<[1024, 1], offset: 2304>>
}

// -----

func.func @invalid_subview(%input: memref<4x1024xf32>) -> memref<2x256xf32, strided<[1024, 1], offset: 2304>> {
// expected-error@+1 {{expected sizes to be non-negative, but got -1}}
%0 = memref.subview %input[2, 256] [-1, 256] [1, 1] : memref<4x1024xf32> to memref<2x256xf32, strided<[1024, 1], offset: 2304>>
return %0 : memref<2x256xf32, strided<[1024, 1], offset: 2304>>
}

// -----

func.func @invalid_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
%0 = memref.alloc() : memref<8x16x4xf32>
// expected-error@+1 {{expected mixed offsets rank to match mixed sizes rank (2 vs 3) so the rank of the result type is well-formed}}
Expand Down