IREE PR #24745: How an Overflowing Size Check Let Poison Reach Runtime

IREE PR #24745 fixes a compiler guard that could overflow while calculating whether an allocation was too large.

That failure inverted the purpose of the check.

An allocation derived from a loosely bounded dynamic dimension could be far beyond any real stack or shared-memory limit. The checker multiplied its dimensions in a signed 64-bit integer, overflowed, and continued with a corrupted size. If that corrupted value appeared smaller than the configured limit, the allocation passed validation.

Lowering then carried the impossible allocation into LLVM IR, where the element count could become llvm.mlir.poison. The generated program could load and store through an invalid pointer and return completely wrong results without a useful compiler diagnostic.

The patch changes the resource check from ordinary arithmetic to overflow-checked arithmetic and makes the CPU and GPU paths share the same static-shape calculation.

The failure began with a legal-looking upper bound

Dynamic dimensions cannot always remain dynamic throughout CPU code generation.

One IREE path tries to hoist a dynamically sized stack allocation by computing an upper bound for each dynamic extent. It then creates a statically shaped allocation large enough for that bound and exposes the runtime-sized portion through a view.

That transformation is reasonable when the bound is meaningful, such as a tile size capped at a small constant.

A dynamic value with no useful problem-specific bound can still carry the default index upper bound:

2^53 - 1
=
9,007,199,254,740,991

Treating that number as an allocation dimension does not make the allocation practical. It only makes the dimension statically representable.

The related IREE issue demonstrates a minimal case in which a bound derived from that maximum becomes a 2^52-element i32 stack allocation, approximately 16 pebibytes. The PR itself was reduced from a quantized ResNet downsample path with a dynamic batch and a DequantizeLinear → 1×1 stride-2 Conv sequence.

The reported model compiled without a diagnostic and returned incorrect results.

The checker overflowed before it could reject the allocation

The CPU validation pass was intended to reject stack allocations above a configured byte limit.

Its old calculation followed this shape:

start with 1
        ↓
multiply by every static dimension
        ↓
multiply by upper bounds for dynamic dimensions
        ↓
multiply by the element bit width
        ↓
apply alignment
        ↓
add to the cumulative allocation size
        ↓
compare with the limit

Each step can increase the value.

The old implementation used plain int64_t multiplication and addition for those increases. Once an intermediate result exceeded the signed 64-bit range, the calculation no longer represented the allocation.

The dynamic ResNet-shaped test added by the patch uses:

memref<?x1024x14x14xf32>

with the dynamic dimension bounded above by 2^53 - 1.

Even before converting bits to bytes, the mathematical product is more than six thousand times larger than the maximum positive signed 64-bit integer. A 64-bit accumulator cannot represent it.

The correct result of the resource check is therefore immediate rejection. Continuing to compare a corrupted intermediate value with a 32 KiB stack limit has no useful meaning.

Passing the check did not make the allocation valid

After the broken check accepted the allocation, later lowering still had to express its element count.

The related issue traces the invalid shape into LLVM conversion, where size and stride calculations can fold to llvm.mlir.poison.

Poison is not a large but valid allocation size. It represents a value whose use has undefined semantics in the lowered program.

If the resulting pointer arithmetic reaches loads and stores, the compiler has already lost the memory-safety contract that the allocation check was supposed to protect.

This explains why the visible symptom could be silent numerical corruption rather than a clean out-of-memory error.

The allocation was never genuinely feasible. The compiler only failed to stop it before invalid arithmetic contaminated the lower-level representation.

The shared helper now fails on the first unrepresentable product

The patch introduces getStaticShapeSizeInBits in IREE’s shared Codegen utilities.

Its contract is narrow:

input:
  a shaped type
  a callback that returns the bit width of a leaf element type

output:
  the product of all static extents and the leaf bit width

failure:
  any 64-bit multiplication overflow

The helper recursively handles shaped element types such as vectors nested inside memrefs.

For each static extent, it uses llvm::MulOverflow. Dynamic dimensions are skipped because the CPU and GPU callers apply different policies to them.

The CPU caller subsequently obtains upper bounds for dynamic operands and multiplies each one into the size with another checked operation.

The GPU caller uses the same static-shape helper while retaining its target-specific handling for index bit width and nested shaped elements.

This unification matters because “how many bits does this allocation require?” should not have separate arithmetic definitions for each backend.

Backend policy may differ. Overflow behavior should not.

Alignment and cumulative size are checked too

Catching the dimension product alone would leave two later overflow sites.

An allocation with explicit alignment may be rounded upward. The patch calculates the number of aligned units and uses checked multiplication to reconstruct the aligned bit size.

The pass also adds each allocation to a cumulative total. That addition now uses llvm::AddOverflow.

The resulting decision order is:

calculate static size with checked multiplication
        ↓
apply dynamic upper bounds with checked multiplication
        ↓
round alignment with checked multiplication
        ↓
round bits up to bytes
        ↓
accumulate with checked addition
        ↓
compare the valid byte count with the resource limit

If any intermediate result cannot be represented in 64 bits, the pass emits an error at the allocation operation.

Examples include:

stack allocation size overflows 64 bits

cumulative stack allocation size overflows 64 bits

shared memory allocation size overflows 64 bits

The compiler now fails before LLVM conversion instead of allowing an impossible size to masquerade as a small one.

Bit-sized elements exposed a second boundary

Review also identified a separate problem in the bit-to-byte conversion.

The old GPU code accumulated a size in bits and converted it with integer division:

allocSize / 8

That rounds down.

For an allocation containing one more than a whole number of bytes of i1 elements, rounding down understates the storage requirement.

The merged patch uses ceiling division:

llvm::divideCeil(allocSize, 8)

The CPU test includes 262145xi1, which requires 32,769 bytes rather than 32,768. The GPU test includes 524289xi1, which requires 65,537 bytes rather than 65,536.

These cases are small enough to fit in 64 bits. They verify that a representable size is still measured accurately after overflow handling is added.

CPU and GPU now share arithmetic, not resource policy

The final design does not force the CPU and GPU checks to use identical limits or dynamic-allocation rules.

The CPU path validates stack allocations and resolves upper bounds for dynamic extents before LLVM conversion.

The GPU path validates workgroup shared-memory allocations, uses the target data layout for index width, and rejects unsupported dynamic shared-memory cases in that pass.

What they share is the recursive static-size calculation and its failure semantics.

That distinction is important:

shared arithmetic contract:
  static extents and element widths must be multiplied without overflow

backend policy:
  which dimensions are allowed, which limit applies, and how the error is reported

The patch centralizes the first while preserving the second.

Review broadened the patch beyond the original CPU fix

The initial proposal focused on the CPU stack checker and introduced local wrappers around checked addition and multiplication.

Review requested two changes.

First, the size calculation should live in a common Codegen utility rather than becoming another CPU-only abstraction.

Second, the GPU resource checker should use the same logic so nested element shapes, target index widths, alignment, cumulative totals, and overflow diagnostics stay consistent.

The final patch reflects that review.

It adds CPU tests for:

static-size overflow
dynamic-bound overflow
alignment overflow
cumulative overflow
sub-byte rounding

It adds GPU tests for:

ordinary limit overflow
nested shaped elements
sub-byte rounding
index element width
static-size overflow
alignment overflow
cumulative overflow

The changed review was ultimately approved after those cases were added.

This patch is a guardrail, not the complete upstream fix

The overflow check now rejects an impossible allocation safely.

It does not remove the compiler transformation that can derive enormous static allocations from weak dynamic bounds.

The related IREE issue about hoistOneStaticallyBoundAllocation remained open when this article was prepared. Reviewers explicitly separated the two concerns: the check must fail safely, and the pipeline should also avoid producing such allocations when the bound is merely the type-level maximum.

That leaves two distinct contracts:

allocation formation:
  do not turn a meaningless maximum bound into a practical stack plan

allocation validation:
  reject every size that exceeds the limit or cannot be represented

PR #24745 repairs the second contract.

A later change can improve the first without depending on undefined arithmetic in the safety check.

What the evidence establishes

The merged tests demonstrate that overflow cases now produce compiler diagnostics across the CPU and GPU resource checks. They also protect alignment, cumulative totals, nested shaped elements, target index widths, and sub-byte rounding.

The PR reports an end-to-end model case that previously compiled and returned wrong results.

This review did not independently rebuild IREE or execute that ONNX reproduction. The committed tests operate at the compiler-pass level rather than replaying the complete quantized ResNet pipeline.

The supported conclusion is therefore precise:

IREE’s allocation-size checks could overflow while measuring an oversized allocation, allowing invalid IR to proceed. The merged patch converts the size path to checked arithmetic and rejects unrepresentable allocations before LLVM lowering.

It does not establish that every source of oversized allocation has been removed.

The contract restored by the patch

A resource limit is meaningful only when the measured size is valid.

The corrected compiler boundary is:

derive allocation dimensions
        ↓
calculate size with checked arithmetic
        ↓
reject overflow as over-limit
        ↓
apply backend-specific resource limits
        ↓
lower only representable allocations

An overflow is not evidence that an allocation fits.

It is evidence that the checker no longer has a valid size to compare.

That is the failure signal restored by IREE PR #24745.