GPU kernel cannot use non-literal global variables
- Dominant language
- Python
- Stars
- 16.8k
- Forks
- 603
- Avg merge
- 4d 23h
- Merged PRs (30d)
- 6
Description
## Summary
GPU kernels currently seem unable to use `non-literal global variables` that are initialized from top-level expressions/conversions.
Literal globals work correctly, but globals derived from other globals or top-level computations appear as `0`.
## Motivation
The original trigger was `math.pi32`.
According to the Math module, `pi64` is stored as a literal global, while `pi32` is derived from it via a cast/conversion. In GPU kernels, `pi64` works, but `pi32` reads as zero.
That led me to reduce the problem into a smaller reproducer involving top-level globals initialized through conversions.
## Problem Reproduction in `Math Module`
```
import math
import gpu
def close(a: float, b: float, epsilon: float = 1e-6):
return math.fabs(a - b) <= epsilon
@gpu.kernel
def kernel(out):
i = gpu.thread.x
out[i] = float(math.pi32)
out = [0.0]
kernel_args = (out,)
kernel(*kernel_args, grid=1, block=1)
print("out: ", out[0])
print("math.pi32: ", float(math.pi32))
assert close(out[0], float(math.pi32), 1e-5)
print("PASS pi32_gpu_f32")
```
## Result of `Math Module` Reproduction
```
$ codon run pi32_gpu_f32.codon
out: 0
math.pi32: 3.14159
AssertionError: Assert failed (/home/swchoi/src/test_code/codon_math_matrix/generated/cases/nvptx/pi32_gpu_f32.codon:18)
Raised from:
/home/swchoi/src/test_code/codon_math_matrix/generated/cases/nvptx/pi32_gpu_f32.codon:18:1
Backtrace:
[0x77492d8d0f13] main.0 at /home/swchoi/src/test_code/codon_math_matrix/generated/cases/nvptx/pi32_gpu_f32.codon:18:1
Aborted (core dumped)
```
As you can see that, reference pi32 in kernel appears to be zero.
### Minimal Reproducer
```
import math
import gpu
# Base literals.
G_F64 = 3.14159265358979323846
G_I64 = 123456789
# Derived globals initialized via top-level conversion.
G_F32_FROM_F64 = float32(G_F64)
G_F64_FROM_I64 = float(G_I64)
G_F32_FROM_I64 = float32(G_I64)
def close(a: float, b: float, epsilon: float = 1e-6):
return math.fabs(a - b) <= epsilon
@gpu.kernel
def kernel(out_f64, out_f32_from_f64, out_f64_from_i64, out_f32_from_i64):
i = gpu.thread.x
out_f64[i] = G_F64
out_f32_from_f64[i] = float(G_F32_FROM_F64)
out_f64_from_i64[i] = G_F64_FROM_I64
out_f32_from_i64[i] = float(G_F32_FROM_I64)
out_f64 = [0.0]
out_f32_from_f64 = [0.0]
out_f64_from_i64 = [0.0]
out_f32_from_i64 = [0.0]
kernel(out_f64, out_f32_from_f64, out_f64_from_i64, out_f32_from_i64, grid=1, block=1)
print("host G_F64 =", G_F64)
print("host G_F32_FROM_F64 =", float(G_F32_FROM_F64))
print("host G_F64_FROM_I64 =", G_F64_FROM_I64)
print("host G_F32_FROM_I64 =", float(G_F32_FROM_I64))
print("kernel out_f64 =", out_f64[0])
print("kernel out_f32_from_f64 =", out_f32_from_f64[0])
print("kernel out_f64_from_i64 =", out_f64_from_i64[0])
print("kernel out_f32_from_i64 =", out_f32_from_i64[0])
assert close(out_f64[0], G_F64)
assert close(out_f32_from_f64[0], float(G_F32_FROM_F64))
assert close(out_f64_from_i64[0], G_F64_FROM_I64)
assert close(out_f32_from_i64[0], float(G_F32_FROM_I64))
print("PASS global_cast_init_repro_nvptx")
```
## Result of MRE
```
$ ./global_cast_init_repro_nvptx
host G_F64 = 3.14159
host G_F32_FROM_F64 = 3.14159
host G_F64_FROM_I64 = 1.23457e+08
host G_F32_FROM_I64 = 1.23457e+08
kernel out_f64 = 3.14159
kernel out_f32_from_f64 = 0
kernel out_f64_from_i64 = 0
kernel out_f32_from_i64 = 0
Segmentation fault (core dumped)
```
## Insight
Only literal Global `G_F64` can be read correctly in the kernel.
Globals derived from top-level computation/conversion appear as `zero`.
## Suspected Cause
From Reading the LLVM IR code generated by `--llvm` option, my current understanding is
Codon does `not necessarily materialize globals` as final compile-time constants in LLVM IR.
Instead, their values are established at runtime `in main.0`
Host/device module separation and pruning happen earlier.
The relevant IR shape for the globals looks like this:
```
@.G_F32_FROM_F64.0 = private global float 0.000000e+00, !dbg !462
@.G_F32_FROM_I64.0 = private global float 0.000000e+00, !dbg !464
@.G_F64.0 = private global double 0.000000e+00, !dbg !466
@.G_F64_FROM_I64.0 = private global double 0.000000e+00, !dbg !468
```
and the traversal logic I was looking at is:
```
void exploreGV(llvm::GlobalValue *G, llvm::SmallPtrSetImpl &keep) {
if (keep.contains(G))
return;
keep.insert(G);
if (auto *F = llvm::dyn_cast(G)) {
for (auto I = llvm::inst_begin(F), E = inst_end(F); I != E; ++I) {
for (auto &U : I->operands()) {
if (auto *G2 = llvm::dyn_cast(U.get()))
exploreGV(G2, keep);
}
}
}
}
```
My suspicion is that this fails for derived globals because the global itself is emitted as a `zero-initialized storage object`, while its actual value is established later through runtime initialization.
So there is no static initializer chain in Compile Time IR for `exploreGV` to follow back to the root dependency.
In other words, the device module sees the storage, but not the computation that defines its value.
## Why I think this is a design question, not just a local bug
I think this raises a broader question about what semantics Codon wants for globals used by accelerators.
There seem to be a few possible directions:
- Freeze eligible globals as compile-time constants for device code (This is how Numba Cuda in Python deals with GV!)
- Similar to how some Python GPU DSLs effectively snapshot globals for kernels. (Idk what about taichi btw)
- This would make top-level pure expressions/casts usable in kernels as constants.
- Require explicit host/device separation for globals
- Closer to C/CUDA-style models where device-visible globals are explicitly represented.
- Disallow or restrict runtime-initialized globals in accelerator kernels
- In that case, perhaps only literal/constant-foldable globals should be allowed in kernels, with a diagnostic.
For comparison, tools like Numba appear to treat globals more like captured/snapshotted environment values for compiled kernels rather than dynamically tracking later host-side mutation.
## Question for maintainers
What direction would maintainers prefer here?
More specifically:
- Should top-level pure/derived globals (e.g. casts from other globals) be expected to work in GPU kernels?
- Is the intended model that kernel-visible globals must be compile-time constants only?
- If runtime-initialized globals are not meant to be supported on accelerators, should this become a frontend restriction / diagnostic instead of silently producing zero-initialized behavior?
As far as I know, I couldn't find any docs or issues related to this issue.
If so, would you please kindly let me know?
If you need any reference for other code or insight, I can share them on discord too! Please let me know.
Always Thanks For Your Great Job, Maintainers and Contributors!
Contributor guide
Assessment
This issue has not been assessed yet.