[Codegen] Use overflow-checked math in the CPU and GPU alloc size checks by dimp-pl · Pull Request #24745 · iree-org/iree

dimp-pl · 2026-07-22T14:46:00Z

checkStackAllocationSize accumulates the allocation size in a plain int64_t. Allocas sized by the upper bound of an unbounded dynamic dimension (the util.assume.int default, 2^53-1) overflow it, the wrapped value passes the limit check and the alloca reaches ConvertToLLVM, where its element count folds to llvm.mlir.poison and the program loads and stores through a garbage pointer at runtime.

The reproduction was reduced from the quantized ResNet downsample construct: an ONNX model with a dynamic batch doing DequantizeLinear -> 1x1 stride-2 Conv compiles without diagnostics and returns completely wrong results.

Move getStaticShapeSizeInBits to Codegen/Utils (taken from #24684) and use it from both LLVMCPUCheckIRBeforeLLVMConversion and GPUCheckResourceUsage. Callers supply the leaf element bit width,
which lets the GPU check keep its index-bitwidth handling and the recursion into shaped element types. On the CPU side the remaining scalar steps are overflow-checked in place with llvm::MulOverflow / llvm::AddOverflow. The pre-existing checkedAdd/checkedMul/checkedAlignTo wrappers in LLVMCPUAssignWorkgroupLocalMemory.cpp are replaced with llvm::MulOverflow / llvm::AddOverflow.

devtbi

right idea, I'd like to see all size calculations unified

devtbi · 2026-07-22T15:08:57Z

 }

+LogicalResult checkedAdd(int64_t lhs, int64_t rhs, int64_t &result) {
+  if (llvm::AddOverflow(lhs, rhs, result)) {


this is 1) the wrong place for such abstractions and 2) I'm not sure wrapper functions to wrap an if-statement are good design

Moved it to Codegen/Utils.cpp, borrowing getStaticShapeSizeInBits from your PR #24684.

egebeysel · 2026-07-22T16:35:00Z

JFYI at first glance, these seem relevant:

#24483
#24627
#24532

egebeysel · 2026-07-22T16:37:33Z

Allocas sized by the upper bound of an unbounded dynamic dimension (the util.assume.int default, 2^53-1)

This guardrail is good, though I think independent of that, we should not have these allocas. Seems like a tiling/fusion bug. Can you create an issue with a reproducer so that we can also take a look at what's causing those?

`checkStackAllocationSize` accumulates the allocation size in a plain int64_t. Allocas sized by the upper bound of an unbounded dynamic dimension (the util.assume.int default, 2^53-1) overflow it, the wrapped value passes the limit check and the alloca reaches ConvertToLLVM, where its element count folds to llvm.mlir.poison and the program loads and stores through a garbage pointer at runtime. The reproduction was reduced from the quantized ResNet downsample construct: an ONNX model with a dynamic batch doing DequantizeLinear -> 1x1 stride-2 Conv compiles without diagnostics and returns completely wrong results. Move `getStaticShapeSizeInBits` to Codegen/Utils (taken from iree-org#24684) and use it from both LLVMCPUCheckIRBeforeLLVMConversion and GPUCheckResourceUsage. Callers supply the leaf element bit width, which lets the GPU check keep its index-bitwidth handling and the recursion into shaped element types. On the CPU side the remaining scalar steps are overflow-checked in place with `llvm::MulOverflow` / `llvm::AddOverflow`. The pre-existing `checkedAdd`/`checkedMul`/ `checkedAlignTo` wrappers in LLVMCPUAssignWorkgroupLocalMemory.cpp are likewise replaced with `llvm::MulOverflow` / `llvm::AddOverflow`. Signed-off-by: Zmicier Prybysh <zprybysh@baylibre.com>

dimp-pl · 2026-07-23T16:45:38Z

@egebeysel thanks for the context, #24483/#24627 definitely seem relevant. I opened a new issue with an easy reproducer for this fusion issue -- #24752.

dimp-pl · 2026-08-04T12:12:43Z

@devtbi i think i addressed all your points, can i ask you to re-review?

devtbi

sorry for taking so long, I missed to publish the review and forgot about it.
Thanks a lot for unifying LLVMCPU and GPU - that will help a lot :)

review in summary: missing a bit of testing coverage + a few small improvements. Let me know if you want/need help.

devtbi · 2026-08-05T05:44:59Z

+    auto emitOverflowError = [&]() {
+      return allocaOp->emitOpError(
+          "stack allocation size overflows 64 bits; the allocation is "
+          "unbounded in practice and cannot live on the stack");


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

devtbi · 2026-08-05T05:46:54Z

    return
  }
 }
+


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

devtbi · 2026-08-05T05:49:53Z

-    int allocSize = shapedTypeStaticSize(allocOp, allocType, getIndexBitwidth);
+    auto func = allocOp->getParentOfType<mlir::FunctionOpInterface>();
+    FailureOr<int64_t> allocSizeBits =
+        getStaticShapeSizeInBits(allocType, [&](Type elementType) -> int64_t {


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

devtbi · 2026-08-05T05:51:37Z

+        getStaticShapeSizeInBits(allocType, [&](Type elementType) -> int64_t {
+          if (elementType.isIndex()) {
+            assert(getIndexBitwidth &&
+                   "getIndexBitwidth should have been set earlier");


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

devtbi · 2026-08-05T14:18:12Z

      allocSize =
          (llvm::divideCeil(allocSize, alignmentInBits) * alignmentInBits);
    }
    cumSize += allocSize / 8;


might overflow

devtbi · 2026-08-24T13:22:52Z


 // -----

+func.func @overflowing_static_alloca() {


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

devtbi · 2026-08-24T13:27:03Z

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


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

devtbi · 2026-08-24T13:28:37Z

    }

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


I think this can reuse funcOp

devtbi · 2026-08-24T13:38:49Z

      allocSize =
          (llvm::divideCeil(allocSize, alignmentInBits) * alignmentInBits);
    }
    cumSize += allocSize / 8;


not modified, but actually, this looks fishy... @AGindinson, I think this should use divideCeil shouldn't it?

dimp-pl requested a review from hanhanW as a code owner July 22, 2026 14:46

devtbi requested changes Jul 22, 2026

View reviewed changes

dimp-pl force-pushed the llvmcpu-stack-check-overflow branch from 84fc675 to 486baa5 Compare July 23, 2026 16:10

dimp-pl requested review from Groverkss, MaheshRavishankar, kuhar and qedawkins as code owners July 23, 2026 16:10

dimp-pl changed the title ~~[Codegen][LLVMCPU] Use overflow-checked math in the stack allocation size check~~ [Codegen] Use overflow-checked math in the stack allocation size check Jul 23, 2026

dimp-pl changed the title ~~[Codegen] Use overflow-checked math in the stack allocation size check~~ [Codegen] Use overflow-checked math in the CPU and GPU alloc size checks Jul 23, 2026

dimp-pl mentioned this pull request Jul 23, 2026

[CPU] Dequantize feeding a strided conv is not fused, whole-tensor temp lands on the stack #24752

Open

egebeysel requested a review from devtbi July 24, 2026 13:32

devtbi requested changes Aug 24, 2026

View reviewed changes

Uh oh!

Conversation

dimp-pl commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

devtbi left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

egebeysel commented Jul 22, 2026

Uh oh!

egebeysel commented Jul 22, 2026

Uh oh!

dimp-pl commented Jul 23, 2026

Uh oh!

dimp-pl commented Aug 4, 2026

Uh oh!

devtbi left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dimp-pl commented Jul 22, 2026 •

edited

Loading

devtbi left a comment •

edited

Loading