EnzymeAD / EnzymeAD/Enzyme-JAX

`raise=true`: cannot raise `llvm.alloca` of `Complex{Float64}` (`array<2 x i64>`) inside a KA kernel

Open
#2,608 2 comments 0 reactions 0 assignees View on GitHub
raising
Dominant language
MLIR
Stars
133
Forks
53
Avg merge
1d 10h
Merged PRs (30d)
193

Description

## Summary

Native `Complex{Float64}` arithmetic in a KernelAbstractions kernel fails raise on `llvm.alloca` of `!llvm.array<2 x i64>`. Reproduced on Reactant 0.2.271 (2026-07-09).

## Environment

Julia 1.12.5 · Reactant 0.2.271 · Enzyme 0.13.173 · CUDA 6.2.0 · KernelAbstractions 0.9.42 · GPU backend (NVIDIA L4)

## Reproducer

```julia
using Reactant
using CUDA
using KernelAbstractions
import KernelAbstractions as KA

Reactant.set_default_backend("gpu")
@assert Reactant.is_extension_loaded(Val(:CUDA))

@kernel function complex_ops_kernel!(out, a, b, c)
col = @index(Global)
FT = eltype(out)
x = @inbounds a[col]
y = @inbounds b[col]
t = @inbounds c[col]
z1 = Complex{FT}(x, y)
z2 = Complex{FT}(t, one(FT))
s = sqrt(z1)
p = z1^convert(FT, 0.65)
e = exp(-convert(FT, 2) * Complex{FT}(zero(FT), one(FT)) * z2 * t)
q = (s * z2 - p) / (s * z2 + p)
r = abs((convert(FT, 2) * s) / (s + z2))^2
zsum = p + e + q
zfix = imag(zsum) >= zero(FT) ? Complex{FT}(real(zsum), -convert(FT, 1.0e-4)) : zsum
@inbounds out[col] = real(zfix) + imag(s) + r + real(q) / (one(FT) + abs(e))
end

function complex_forward!(out, a, b, c)
complex_ops_kernel!(KA.get_backend(out))(out, a, b, c; ndrange = length(out))
return nothing
end

a = Reactant.to_rarray([0.5, 1.2, 3.0])
b = Reactant.to_rarray([-0.2, 0.4, 0.1])
c = Reactant.to_rarray([0.3, 0.9, 1.5])
out = Reactant.to_rarray(zeros(3))
compiled = Reactant.@compile raise=true raise_first=true sync=true complex_forward!(out, a, b, c)
compiled(out, a, b, c)
```

## Expected behavior

Native `Complex{Float64}` arithmetic inside the KA kernel should lower and compile without an unsupported aggregate `llvm.alloca`.

## Actual behavior

```text
ERROR: LoadError: CompilationError: MLIR pass pipeline "all" failed
error: cannot raise op to stablehlo%40 = "llvm.alloca"(%8) <{alignment = 8 : i64, elem_type = !llvm.array<2 x i64>}> : (i32) -> !llvm.ptr
note: see current operation: %40 = "llvm.alloca"(%8) <{alignment = 8 : i64, elem_type = !llvm.array<2 x i64>}> : (i32) -> !llvm.ptr
failed to raise func: func.func private @"##call__Z23gpu_complex_ops_kernel_..."(...)
```

## Related issues

This is distinct from the scalar-global `llvm.mlir.addressof @_j_const_*` failure in Reactant.jl [#3043](https://github.com/EnzymeAD/Reactant.jl/issues/3043) and from [#3019](https://github.com/EnzymeAD/Reactant.jl/issues/3019), which involves `@localmem` and an address-space-3 pointer rather than a default-address-space Complex aggregate allocation.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the Reactant.@compile raise=true reproducer for complex_ops_kernel! and trace the reported llvm.alloca of !llvm.array<2 x i64> through the raising pipeline. Compare the distinctions from related issues #3043 and #3019; done means native Complex{Float64} arithmetic in the KA kernel compiles without the unsupported aggregate allocation error.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.