Skip to content

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

Description

@masonlee277

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

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

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 and from #3019, which involves @localmem and an address-space-3 pointer rather than a default-address-space Complex aggregate allocation.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions