AcoSail/llvm-project
8
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix SimplifyAllocConst pattern when we have alloc of negative sizes

Browse Source 
This is UB, but we shouldn't crash the compiler either.

Fixes #61056

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D144978
This commit is contained in:
Mehdi Amini 2023-02-28 15:00:45 -05:00
parent 9a9fce1fed
commit f35ac8a4ff
2 changed files with 20 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
View File

@ -284,7 +284,10 @@ struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
// Check to see if any dimensions operands are constants. If so, we can
// substitute and drop them.
if (llvm::none_of(alloc.getDynamicSizes(), [](Value operand) {
return matchPattern(operand, matchConstantIndex());
APInt constSizeArg;
if (!matchPattern(operand, m_ConstantInt(&constSizeArg)))
return false;
return constSizeArg.isNonNegative();
}))
return failure();
@ -305,11 +308,11 @@ struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
continue;
}
auto dynamicSize = alloc.getDynamicSizes()[dynamicDimPos];
auto *defOp = dynamicSize.getDefiningOp();
if (auto constantIndexOp =
dyn_cast_or_null<arith::ConstantIndexOp>(defOp)) {
APInt constSizeArg;
if (matchPattern(dynamicSize, m_ConstantInt(&constSizeArg)) &&
constSizeArg.isNonNegative()) {
// Dynamic shape dimension will be folded.
newShapeConstants.push_back(constantIndexOp.value());
newShapeConstants.push_back(constSizeArg.getZExtValue());
} else {
// Dynamic shape dimension not folded; copy dynamicSize from old memref.
newShapeConstants.push_back(ShapedType::kDynamic);

12
mlir/test/Dialect/MemRef/canonicalize.mlir
View File

@ -928,3 +928,15 @@ func.func @fold_multiple_memory_space_cast(%arg : memref<?xf32>) -> memref<?xf32
%1 = memref.memory_space_cast %0 : memref<?xf32, 1> to memref<?xf32, 2>
return %1 : memref<?xf32, 2>
}
// -----
// CHECK-lABEL: func @ub_negative_alloc_size
func.func private @ub_negative_alloc_size() -> memref<?x?x?xi1> {
%idx1 = index.constant 1
%c-2 = arith.constant -2 : index
%c15 = arith.constant 15 : index
// CHECK: %[[ALLOC:.*]] = memref.alloc(%c-2) : memref<15x?x1xi1>
%alloc = memref.alloc(%c15, %c-2, %idx1) : memref<?x?x?xi1>
return %alloc : memref<?x?x?xi1>
}