diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -284,7 +284,10 @@ // 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,12 +308,13 @@ continue; } auto dynamicSize = alloc.getDynamicSizes()[dynamicDimPos]; - auto *defOp = dynamicSize.getDefiningOp(); - if (auto constantIndexOp = - dyn_cast_or_null(defOp)) { + APInt constSizeArg; + if (matchPattern(dynamicSize, m_ConstantInt(&constSizeArg)) + && constSizeArg.isNonNegative()) { // Dynamic shape dimension will be folded. - newShapeConstants.push_back(constantIndexOp.value()); - } else { + newShapeConstants.push_back(constSizeArg.getZExtValue()); + } + else { // Dynamic shape dimension not folded; copy dynamicSize from old memref. newShapeConstants.push_back(ShapedType::kDynamic); dynamicSizes.push_back(dynamicSize); diff --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir --- a/mlir/test/Dialect/MemRef/canonicalize.mlir +++ b/mlir/test/Dialect/MemRef/canonicalize.mlir @@ -928,3 +928,15 @@ %1 = memref.memory_space_cast %0 : memref to memref return %1 : memref } + +// ----- + +// CHECK-lABEL: func @ub_negative_alloc_size +func.func private @ub_negative_alloc_size() -> memref { + %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 + return %alloc : memref +}