littlehamsterxu-bot opened a new issue, #20125:
URL: https://github.com/apache/tvm/issues/20125
## LLVM backend truncates large AllocBuffer size from int64 to int32
### Summary
`CodeGenLLVM::VisitStmt_(const AllocBufferNode*)` reads the buffer's shape
extent (an `int64_t` value from `IntImmNode::value`), casts it to `int32_t` via
`static_cast`, and passes it to `CreateAlloca` through `ConstInt32`. When the
extent exceeds `INT32_MAX`, the value is truncated (implementation-defined
narrowing), producing a severely undersized stack allocation.
NVPTX and AMDGPU have their own `VisitStmt_` overrides with a related but
distinct issue — they compute the element count in `size_t` (safe), but still
narrow to `ConstInt32` in the `CreateAlloca` call for `kLocal` scope. See
Backend comparison below.
The `C`/`CUDA`/`OpenCL` backends (inheriting from `CodeGenC`) are not
affected — they use `size_t`.
Whether the truncation manifests depends on both the pipeline and the
optimization level:
- **`tvm.compile` (default optimization level)**: `LowerTVMBuiltin`
redirects large allocations to `TVMBackendAllocWorkspace(i64)`. Safe.
- **`pipeline="tirx"` + `opt-level=0` + `disable_lower_builtin`**:
Annotation blocks `LowerTVMBuiltin`; raw codegen runs; `alloca float`
(truncated) + `i64` GEP → **OOB access**.
- **`pipeline="s_tir"` + `opt-level=0`**: `LowerTVMBuiltin` does **not**
apply (annotation or not); raw codegen runs; `alloca float` (truncated) + `i32`
GEP → wrong element accessed but no OOB (index `4294967296` wraps to `0` in
`i32`).
In short, the `int32_t` truncation in `codegen_llvm.cc` is always present;
`LowerTVMBuiltin` can mask it at higher optimization levels. At `opt-level=0`,
both pipelines expose the raw truncation — the only difference is whether the
subsequent GEP is `i64` (TIR-X, OOB crash) or `i32` (S-TIR, silent
wrong-result).
### Affected code
**Upstream `main` (ea0950ab, 2026-08-12)** —
`src/target/llvm/codegen_llvm.cc` ~L2119–2136:
```cpp
void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) {
TVM_FFI_ICHECK_EQ(op->buffer->shape.size(), 1)
<< "LLVM codegen only supports flat 1-d buffer allocation…";
const IntImmNode* dim_imm = op->buffer->shape[0].as<IntImmNode>();
TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack
allocation";
int32_t constant_size = static_cast<int32_t>(dim_imm->value); // ←
TRUNCATION
TVM_FFI_ICHECK_GT(constant_size, 0) << "…"; // check
on WRAPPED value!
builder_->CreateAlloca(
DTypeToLLVMType(op->buffer->dtype),
ConstInt32(constant_size)); // ← also
needs ConstInt64
}
```
`dim_imm->value` is `int64_t`. For `EXTENT = 2**32 + 1 = 4294967297`, the
source value exceeds `INT32_MAX` (2147483647). The `static_cast<int32_t>` is
implementation-defined per the C++ standard (the value is reduced modulo 2³² to
fit the target range on two's-complement platforms, yielding `1` in practice).
Then `ICHECK_GT(1, 0)` passes — the check runs on the already-truncated value.
`CreateAlloca(float, ConstInt32(1))` emits `alloca float` — 1 element instead
of 4294967297.
Note: the function does **not** call
`AllocBufferNode::ConstantAllocationSize()` (which returns
`std::optional<int64_t>` element count at `include/tvm/tirx/stmt.h:292`). It
reads `shape[0]` directly and is limited to 1-D buffers. The truncation is on
the extent value itself.
**TVM 0.25.0.post1 (commit b3e249b7d)** — same bug at line 2020:
```cpp
int32_t constant_size = static_cast<int32_t>(dim_imm->value);
```
**Older versions** (`AllocateNode` API) — `codegen_llvm.cc:1568`:
```cpp
int32_t constant_size = op->ConstantAllocationSize();
```
Same bug, different API. The migration from `AllocateNode` to
`AllocBufferNode` (PR #18865) changed the API call but preserved the truncation.
### Backend comparison
**CPU (`CodeGenLLVM`, base class)** — `src/target/llvm/codegen_llvm.cc`:
```cpp
int32_t constant_size = static_cast<int32_t>(dim_imm->value); // BUG:
int64 → int32
builder_->CreateAlloca(..., ConstInt32(constant_size)); // 32-bit
alloca
```
**NVPTX (`CodeGenNVPTX`)** —
`src/backend/cuda/codegen/llvm/codegen_nvptx.cc`, **overrides** `VisitStmt_`:
```cpp
size_t constant_size = static_cast<size_t>(dim_imm->value); // OK:
computed in size_t
// ... (only for kLocal scope):
builder_->CreateAlloca(..., ConstInt32(constant_size)); //
narrowing: size_t → i32
```
**AMDGPU (`CodeGenAMDGPU`)** —
`src/backend/rocm/codegen/llvm/codegen_amdgpu.cc`, **overrides** `VisitStmt_`:
```cpp
size_t constant_size = static_cast<size_t>(dim_imm->value); // OK:
computed in size_t
// ... (only for kLocal scope):
builder_->CreateAlloca(..., ConstInt32(constant_size)); //
narrowing: size_t → i32
```
| Backend | Own override? | Element count type | CreateAlloca arg | Scope
affected | Severity |
|---|---|---|---|---|---|
| **CPU** | — (base) | `int32_t` ❌ | `ConstInt32` | all | Higher — element
count itself truncates |
| **NVPTX** | Yes | `size_t` ✅ | `ConstInt32` | `kLocal` only | Lower —
count correct, alloca capped at 32-bit |
| **AMDGPU** | Yes | `size_t` ✅ | `ConstInt32` | `kLocal` only | Lower —
same as NVPTX |
| C | Yes (CodeGenC) | `size_t` ✅ | N/A (C array decl) | — | Not affected |
| CUDA | Inherits CodeGenC | `size_t` ✅ | N/A (C array decl) | — | Not
affected |
| OpenCL | Inherits CodeGenC | `size_t` ✅ | N/A (C array decl) | — | Not
affected |
Because NVPTX and AMDGPU have their own overrides, fixing the CPU base class
does **not** automatically fix them. Each needs its `ConstInt32` → `ConstInt64`
changed independently. However, in GPU local memory the risk is lower since
`kLocal` is typically used for small per-thread allocations, not large buffers.
`lower_warp_memory.cc` used `int alloc_size` in older TVM. In v0.25.0+ it
was changed to `int64_t alloc_size`.
### Reproduction
At `opt-level=0`, **both** pipelines expose the raw codegen truncation —
only the GEP index type differs:
#### TIR-X + opt-level=0 + disable_lower_builtin → OOB crash
```python
from tvm.script import tirx as T
EXTENT = 2**32 + 1 # 4294967297 elements
LAST = EXTENT - 1 # 4294967296
@T.prim_func(s_tir=True)
def func(A: T.Buffer((1,), "float32")):
B = T.alloc_buffer(
(EXTENT,), "float32",
scope="global",
annotations={"disable_lower_builtin": True},
)
A[0] = B[LAST]
mod = tvm.tirx.build(
func,
target={"kind": "llvm", "opt-level": 0},
pipeline="tirx",
)
print(mod.inspect_source("ll"))
```
**Observed** (TVM 0.25.0.post1, b3e249b7d):
```llvm
define internal i32 @func_compute_(ptr align 64 %A_ptr) {
%B_ptr = alloca float, align 16 ; 1 element — should be
4294967297
%1 = getelementptr inbounds float, ptr %B_ptr, i64 4294967296 ; offset
~17 GB
%2 = load float, ptr %1, align 16 ; out-of-bounds read
store float %2, ptr %A_ptr, align 64
ret i32 0
}
```
#### S-TIR + opt-level=0 → silent wrong-result (same truncation, masked by
i32 index)
```python
# Same func, same annotations — only pipeline and opt-level changed:
mod = tvm.tirx.build(
func,
target={"kind": "llvm", "opt-level": 0},
pipeline="s_tir", # ← S-TIR instead of TIR-X
)
print(mod.inspect_source("ll"))
```
**Observed**:
```llvm
define internal i32 @func_compute_(ptr align 64 %A_ptr) {
%B_ptr = alloca float, align 16 ; 1 element — SAME TRUNCATION
%1 = getelementptr inbounds float, ptr %B_ptr, i32 0 ; ← i32: 4294967296
wraps to 0!
%2 = load float, ptr %1, align 16 ; reads element 0 instead of
4294967296
store float %2, ptr %A_ptr, align 64
ret i32 0
}
```
Note: `tvm.compile` (default optimization level) is safe — `LowerTVMBuiltin`
routes large allocations to `TVMBackendAllocWorkspace(i64 17179869188)`, and
smaller allocations that escape the workspace heuristic are eliminated by the
LLVM optimizer.
**Steps** (TIR-X OOB case):
1. `shape[0] = 4294967297` (int64)
2. `static_cast<int32_t>(4294967297)` → `1` (implementation-defined, reduces
modulo 2³² on two's-complement; the check in step 3 runs on the
already-truncated value)
3. `TVM_FFI_ICHECK_GT(1, 0)` passes
4. `CreateAlloca(float, ConstInt32(1))` → `alloca float` = 4 bytes
5. `getelementptr ... i64 4294967296` → 17 GB beyond the 4-byte allocation
### Annotation stripping in S-TIR vs TIR-X
When using `pipeline="s_tir"` (default), the `disable_lower_builtin`
annotation is dropped by `StorageRewrite`. When it rebuilds `AllocBuffer`
nodes, only `kVolatile` survives:
`src/tirx/transform/storage_rewrite.cc` ~L744–794 — `PrepareNewAlloc()`
creates a fresh annotations map, only populating `kVolatile`.
However, **this alone does not explain the S-TIR behavior**. At
`opt-level=0`, even without any annotation, S-TIR still produces a truncated
stack `alloca` — `LowerTVMBuiltin` does not apply. The protection at default
`tvm.compile` comes from the combination of `LowerTVMBuiltin` **and** a
sufficiently high optimization level (≥2), not from annotation stripping alone.
The `disable_lower_builtin` annotation is only meaningful in the TIR-X
pipeline (which preserves it), where it explicitly blocks `LowerTVMBuiltin`
even at higher opt-levels.
### Impact
- **Severity: correctness bug.** Two trigger scenarios:
1. `pipeline="tirx"` + `opt-level=0` + `disable_lower_builtin` annotation
→ `alloca float` (truncated) + `i64` GEP → **out-of-bounds memory access**
2. `pipeline="s_tir"` + `opt-level=0` → `alloca float` (truncated) + `i32`
GEP → **wrong element read** (index wraps to 0, no OOB)
- Both require:
- Allocation with extent > `INT32_MAX` (≥ ~2.1 billion elements)
- `opt-level=0` (default `tvm.compile` uses higher opt-level and is safe)
- Default `tvm.compile` is **not affected** — `LowerTVMBuiltin` +
optimization routes large allocations through `TVMBackendAllocWorkspace(i64)`.
### Suggested fix
**CPU (`CodeGenLLVM`)** — three coordinated changes across two files:
1. In `src/target/llvm/codegen_llvm.cc`, widen `constant_size` and the
`CreateAlloca` argument:
```cpp
// Before:
int32_t constant_size = static_cast<int32_t>(dim_imm->value);
// ...
builder_->CreateAlloca(DTypeToLLVMType(op->buffer->dtype),
ConstInt32(constant_size));
// After (Option A — use int64_t throughout):
int64_t constant_size = dim_imm->value;
TVM_FFI_ICHECK_GT(constant_size, 0)
<< "Allocation size must be positive";
builder_->CreateAlloca(
DTypeToLLVMType(op->buffer->dtype), ConstInt64(constant_size));
// After (Option B — explicit clamp + diagnostic):
int64_t extent = dim_imm->value;
TVM_FFI_ICHECK_GT(extent, 0)
<< "Allocation size must be positive";
TVM_FFI_ICHECK_LE(extent, INT32_MAX)
<< "Allocation size " << extent
<< " exceeds maximum supported stack allocation of " << INT32_MAX;
builder_->CreateAlloca(
DTypeToLLVMType(op->buffer->dtype), ConstInt32(extent));
```
Option A is preferred — it matches the int64_t source type and eliminates
the silent truncation. `ConstInt32` → `ConstInt64` is required for the fix to
be effective.
2. In `src/tirx/transform/ir_utils.h`, widen `GetTempAllocaAlignment`'s
second parameter from `int32_t` to `int64_t`:
```cpp
// Before (line 174):
inline int GetTempAllocaAlignment(DataType type, int32_t const_size) {
// After:
inline int GetTempAllocaAlignment(DataType type, int64_t const_size) {
```
This function is called at `codegen_llvm.cc:2028` with `constant_size` as
the argument. If `constant_size` is widened to `int64_t` but the parameter
remains `int32_t`, the narrowing implicit conversion would cause a compilation
warning (or error under `-Werror`). The function body already casts to
`int64_t` internally for its byte-size computation (`int64_t const_s =
static_cast<int64_t>(const_size) * ...`), so widening the parameter is a safe,
mechanical change.
3. The three call sites of `GetTempAllocaAlignment` and their local variable
types:
- `codegen_llvm.cc:2028` — `int32_t constant_size` → will be `int64_t`
after step 1, matches widened parameter
- `codegen_nvptx.cc:104` — `size_t constant_size` → implicitly converts
to `int64_t` (safe for allocation sizes)
- `codegen_amdgpu.cc:117` — `size_t constant_size` → same as NVPTX
All three callers are compatible with the widened parameter without
additional casts.
**NVPTX / AMDGPU** require separate, smaller fixes: change
`ConstInt32(constant_size)` → `ConstInt64(constant_size)` in their respective
`VisitStmt_` overrides. Their element count computation already uses `size_t`,
so only the `CreateAlloca` call needs updating. These are independent of the
CPU fix.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
