Skip to content

[Codegen] Use overflow-checked math in the CPU and GPU alloc size checks #24745

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.

Sign up for GitHub

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

Open
dimp-pl wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[Codegen] Use overflow-checked math in the CPU and GPU alloc size checks #24745

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 21 additions & 28 deletions compiler/src/iree/compiler/Codegen/Common/GPU/GPUCheckResourceUsage.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@

#include "iree/compiler/Codegen/Common/GPU/Passes.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MathExtras.h"
#include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
Expand All @@ -24,32 +26,6 @@ static unsigned getDatalayoutIndexBitwidth(mlir::FunctionOpInterface func) {
return options.getIndexBitwidth();
}

static int shapedTypeStaticSize(
memref::AllocOp allocOp, ShapedType shapedType,
std::function<unsigned(mlir::FunctionOpInterface)> getIndexBitwidth) {
int allocSize = 1;
for (auto dimSize : shapedType.getShape()) {
if (ShapedType::isDynamic(dimSize)) {
continue;
}
allocSize *= dimSize;
}
if (auto elementType = dyn_cast<ShapedType>(shapedType.getElementType())) {
allocSize *= shapedTypeStaticSize(allocOp, elementType, getIndexBitwidth);
} else {
auto eltTy = shapedType.getElementType();
if (eltTy.isIndex()) {
auto func = allocOp->getParentOfType<mlir::FunctionOpInterface>();
assert(getIndexBitwidth &&
"getIndexBitwidth should have been set earlier");
allocSize *= getIndexBitwidth(func);
} else {
allocSize *= IREE::Util::getTypeBitWidth(shapedType.getElementType());
}
}
return allocSize;
}

/// Returns success if the total shared memory allocation size is less than the
/// limit.
static LogicalResult checkGPUAllocationSize(
Expand All @@ -65,7 +41,7 @@ static LogicalResult checkGPUAllocationSize(
return success();
}

int cumSize = 0;
int64_t cumSize = 0;
for (auto allocOp : allocOps) {
auto allocType = cast<MemRefType>(allocOp.getType());
if (!hasSharedMemoryAddressSpace(allocType)) {
Expand All @@ -77,7 +53,24 @@ static LogicalResult checkGPUAllocationSize(
"has unsupported dynamic shared memory allocations");
}

int allocSize = shapedTypeStaticSize(allocOp, allocType, getIndexBitwidth);
auto func = allocOp->getParentOfType<mlir::FunctionOpInterface>();

@devtbi devtbi Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this can reuse funcOp

FailureOr<int64_t> allocSizeBits =
getStaticShapeSizeInBits(allocType, [&](Type elementType) -> int64_t {

@devtbi devtbi Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

catch-all ref captures should be avoided because of the invisible dependencies they create

if (elementType.isIndex()) {
assert(getIndexBitwidth &&
"getIndexBitwidth should have been set earlier");

@devtbi devtbi Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I'd make this more terse and ... getIndexBitwidth must not be null, or expected....

return getIndexBitwidth(func);
}
return IREE::Util::getTypeBitWidth(elementType);
});
if (failed(allocSizeBits)) {
return emitError(funcOp->getLoc())

@devtbi devtbi Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in line with CPU, this should be allocOp.emitOpError to emit better diagnostics

<< "function '" << funcOp.getName()
<< "' shared memory allocation size overflows the size "
"computation; exceeded the limit of "
<< limit << " bytes";
}
int64_t allocSize = *allocSizeBits;
if (allocOp.getAlignment()) {
int64_t alignmentInBits = *allocOp.getAlignment() * 8;
allocSize =
Expand Down
10 changes: 10 additions & 0 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_check_resource_usage.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,3 +29,13 @@ module {
return
}
}

@devtbi devtbi Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would be nice to have more coverage here e.g. index-element bitwidth path or the nested-shaped-element recursio mentioned in the PR description

// -----

module {
// expected-error @+1 {{shared memory allocation size overflows the size computation}}
func.func @shared_mem_alloc_size_overflow() {
memref.alloc() : memref<9007199254740991x1024x14x14xf32, #gpu.address_space<workgroup>>
return
}
}
44 changes: 11 additions & 33 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUAssignWorkgroupLocalMemory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,32 +49,6 @@ static bool hasWorkgroupLocalMemorySpace(Type type) {
memRefType.getMemorySpace());
}

static LogicalResult checkedAdd(int64_t lhs, int64_t rhs, int64_t &result) {
if (llvm::AddOverflow(lhs, rhs, result)) {
return failure();
}
return success();
}

static LogicalResult checkedMul(int64_t lhs, int64_t rhs, int64_t &result) {
if (llvm::MulOverflow(lhs, rhs, result)) {
return failure();
}
return success();
}

static LogicalResult checkedAlignTo(int64_t value, int64_t alignment,
int64_t &result) {
assert(value >= 0);
assert(alignment > 0);
int64_t remainder = value % alignment;
if (remainder == 0) {
result = value;
return success();
}
return checkedAdd(value, alignment - remainder, result);
}

static bool hasZeroElementShape(MemRefType type) {
return llvm::is_contained(type.getShape(), 0);
}
Expand Down Expand Up @@ -115,15 +89,15 @@ computeStaticElementFootprint(memref::AllocOp allocOp) {
int64_t maxElementOffset = offset;
for (auto [dim, stride] : llvm::zip_equal(type.getShape(), strides)) {
int64_t contribution = 0;
if (failed(checkedMul(dim - 1, stride, contribution)) ||
failed(checkedAdd(maxElementOffset, contribution, maxElementOffset))) {
if (llvm::MulOverflow(dim - 1, stride, contribution) ||
llvm::AddOverflow(maxElementOffset, contribution, maxElementOffset)) {
allocOp.emitOpError("workgroup local memory allocation size overflow");
return failure();
}
}

int64_t footprint = 0;
if (failed(checkedAdd(maxElementOffset, 1, footprint))) {
if (llvm::AddOverflow(maxElementOffset, int64_t(1), footprint)) {
allocOp.emitOpError("workgroup local memory allocation size overflow");
return failure();
}
Expand Down Expand Up @@ -154,7 +128,7 @@ computeAllocationSize(memref::AllocOp allocOp) {
int64_t elementBytes = static_cast<int64_t>(elementByteSize.getFixedValue());

int64_t totalBytes = 0;
if (failed(checkedMul(*elementFootprint, elementBytes, totalBytes))) {
if (llvm::MulOverflow(*elementFootprint, elementBytes, totalBytes)) {
allocOp.emitOpError("workgroup local memory allocation size overflow");
return failure();
}
Expand Down Expand Up @@ -281,7 +255,11 @@ void LLVMCPUAssignWorkgroupLocalMemoryPass::runOnOperation() {
if (failed(alignment)) {
return signalPassFailure();
}
if (failed(checkedAlignTo(currentOffset, *alignment, currentOffset))) {
assert(currentOffset >= 0 && *alignment > 0);
int64_t misalignment = currentOffset % *alignment;
if (misalignment != 0 &&
llvm::AddOverflow(currentOffset, *alignment - misalignment,
currentOffset)) {
allocOp.emitOpError("workgroup local memory allocation size overflow");
return signalPassFailure();
}
Expand All @@ -291,8 +269,8 @@ void LLVMCPUAssignWorkgroupLocalMemoryPass::runOnOperation() {
/*elementFootprint=*/allocationSize->elementFootprint,
});

if (failed(checkedAdd(currentOffset, allocationSize->byteSize,
currentOffset))) {
if (llvm::AddOverflow(currentOffset, allocationSize->byteSize,
currentOffset)) {
allocOp.emitOpError("workgroup local memory allocation size overflow");
return signalPassFailure();
}
Expand Down
37 changes: 26 additions & 11 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUCheckIRBeforeLLVMConversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
#include "iree/compiler/Codegen/LLVMCPU/Utils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/MathExtras.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
Expand Down Expand Up @@ -71,14 +72,20 @@ checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
"all stack allocations need to be hoisted to the entry block of the "
"function");
}
int64_t allocaSize = 1;
auto emitOverflowError = [&]() {
Comment thread
devtbi marked this conversation as resolved.
return allocaOp->emitOpError(
"stack allocation size overflows 64 bits; the allocation is "
"unbounded in practice and cannot live on the stack");

@devtbi devtbi Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • i don't think this unboundedness holds for all cases where this lambda is used? it might be > 64.bit but still bounded
  • Probably makes more sense not to use a lambda

};
auto allocaType = cast<ShapedType>(allocaOp.getType());
for (auto dimSize : allocaType.getShape()) {
if (ShapedType::isDynamic(dimSize)) {
continue;
}
allocaSize *= dimSize;
FailureOr<int64_t> staticSizeBits =
getStaticShapeSizeInBits(allocaType, [](Type elementType) -> int64_t {
return IREE::Util::getTypeBitWidth(elementType);
});
if (failed(staticSizeBits)) {
return emitOverflowError();
}
int64_t allocaSize = *staticSizeBits;
for (auto operand : allocaOp.getDynamicSizes()) {
// Assume vscale is `clAssumedVscaleValue` for determining if the alloca
// is within the stack limit. This should always resolve to a constant
Expand All @@ -90,18 +97,26 @@ checkStackAllocationSize(mlir::FunctionOpInterface funcOp) {
/*vscaleMin=*/assumedVscale,
/*vscaleMax=*/assumedVscale, presburger::BoundType::UB);
if (succeeded(ub)) {
allocaSize *= ub->getSize()->baseSize;
if (llvm::MulOverflow(allocaSize,
static_cast<int64_t>(ub->getSize()->baseSize),
allocaSize)) {
return emitOverflowError();
}
continue;
}
return allocaOp.emitOpError("expected no unbounded stack allocations");
}
allocaSize *= IREE::Util::getTypeBitWidth(allocaType.getElementType());
if (allocaOp.getAlignment()) {
int64_t alignmentInBits = *allocaOp.getAlignment() * 8;
allocaSize =
(llvm::divideCeil(allocaSize, alignmentInBits) * alignmentInBits);
int64_t alignedUnits = llvm::divideCeil(allocaSize, alignmentInBits);
if (llvm::MulOverflow(alignedUnits, alignmentInBits, allocaSize)) {
return emitOverflowError();
}
}
if (llvm::AddOverflow(cumSize, allocaSize / 8, cumSize)) {
return allocaOp->emitOpError(
"cumulative stack allocation size overflows 64 bits");
}
cumSize += allocaSize / 8;
}
if (cumSize > maxAllocationSizeInBytes) {
return funcOp.emitOpError("exceeded stack allocation limit of ")
Expand Down
8 changes: 8 additions & 0 deletions compiler/src/iree/compiler/Codegen/LLVMCPU/test/check_ir_before_llvm_conversion.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,14 @@ func.func @mix_static_and_dynamic_allocas(%arg0: index) {

// -----

func.func @overflowing_static_alloca() {

@devtbi devtbi Aug 24, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would be nice to have more coverage here e.g. for the unbounded dynamic dimension mentioned in the PR description

// expected-error @+1 {{stack allocation size overflows 64 bits}}
%0 = memref.alloca() {alignment = 64 : i64} : memref<9007199254740991x1024x14x14xf32>
return
}

// -----

func.func @non_entry_bb_allocas() {
cf.br ^bb1
^bb1() :
Expand Down
26 changes: 26 additions & 0 deletions compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/MathExtras.h"
#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Analysis/SliceAnalysis.h"
Expand Down Expand Up @@ -94,6 +95,31 @@ bool isEntryPoint(mlir::FunctionOpInterface func) {
return func.isPublic() && getEntryPoint(func);
}

FailureOr<int64_t>
getStaticShapeSizeInBits(ShapedType shapedType,
llvm::function_ref<int64_t(Type)> getElementBitWidth) {
int64_t size = 1;
for (int64_t dimSize : shapedType.getShape()) {
if (ShapedType::isDynamic(dimSize)) {
continue;
}
if (llvm::MulOverflow(size, dimSize, size)) {
return failure();
}
}
Type elementType = shapedType.getElementType();
if (auto shapedElementType = dyn_cast<ShapedType>(elementType)) {
FailureOr<int64_t> elementBits =
getStaticShapeSizeInBits(shapedElementType, getElementBitWidth);
if (failed(elementBits) || llvm::MulOverflow(size, *elementBits, size)) {
return failure();
}
} else if (llvm::MulOverflow(size, getElementBitWidth(elementType), size)) {
return failure();
}
return size;
}

std::optional<StringRef> getConfigCpuFeatures(DictionaryAttr targetConfig) {
auto attr = targetConfig.getAs<StringAttr>(kCpuFeaturesAttrName);
if (attr) {
Expand Down
12 changes: 12 additions & 0 deletions compiler/src/iree/compiler/Codegen/Utils/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,18 @@ bool isEntryPoint(mlir::FunctionOpInterface func);
std::optional<IREE::HAL::ExecutableExportOp>
getEntryPoint(mlir::FunctionOpInterface funcOp);

/// Returns the static size in bits of `shapedType` -- the product of its
/// static extents and the element size -- accumulated with 64-bit overflow
/// checking. Dynamic dimensions are skipped; callers apply their own policy
/// for those. `getElementBitWidth` returns the bit width of a leaf
/// (non-shaped) element type, letting callers control e.g. the target's
/// index-type width. Returns failure on integer overflow: such an allocation
/// necessarily exceeds any real resource limit, so callers should treat
/// failure as "over the limit".
FailureOr<int64_t>
getStaticShapeSizeInBits(ShapedType shapedType,
llvm::function_ref<int64_t(Type)> getElementBitWidth);

/// Returns the dispatch_config op for the `funcOp` by looking up the parent
/// module for a matching function_ref. Returns nullptr if not found.
IREE::Codegen::DispatchConfigOp
Expand Down
Loading