Where do these NaNs in my KernelAbstraction derivative come from?

General Usage
, , ,
1

Hi,
I (mostly chatgpt though) am working on an MHD code using KernelAbstractions.jl, Reactant.jl and Enzyme.jl. In one step, a kernel derivative produces NaNs whereas the reactant gradient using plain arrays gives correct gradients (zero).

The function I’m calculating is q = ρ * (v \cdot n), with derivatives
dρ = λ v \cdot n, dv = λ ρ n and dn = λ ρ v, and λ the incoming derivative.
So when (\rho, v, n) = (1,0,0) the derivative is (d\rho, dv, dn) = (0,0,0) and for
(\rho, v, n) = (1,0,1) the derivative is (0,1,0).
I get this result for the array path. For the KernelAbstraction path I get
[0.0, NaN, NaN] and [0.0, NaN, 0.0].

Is there anything wrong with the code?

The MWE is below, environment is pinned to

  • Julia 1.12.6

  • KernelAbstractions 0.9.42

  • Reactant 0.2.284

  • Enzyme 0.13.199

  • CUDA 5.11.3

  • Reactant CPU target on aarch64 Darwin

using CUDA
using Enzyme
using KernelAbstractions
using Reactant

KernelAbstractions.@kernel function product_kernel!(output, state)
    index = @index(Global, Linear)
    rho = @inbounds state[1]
    velocity = @inbounds state[2]
    normal = @inbounds state[3]
    @inbounds output[index] = rho * (velocity * normal)
end

struct KernelLoss end

function (::KernelLoss)(state)
    backend = KernelAbstractions.get_backend(state)
    output = KernelAbstractions.zeros(backend, eltype(state), (1,))
    product_kernel!(backend, 1)(output, state; ndrange=1)
    return sum(output)
end

struct TensorLoss end

function (::TensorLoss)(state)
    rho = state[1:1]
    velocity = state[2:2]
    normal = state[3:3]
    return sum(rho .* (velocity .* normal))
end

struct ValueGradient{F}
    loss::F
end

function (value_gradient::ValueGradient)(state)
    result = Enzyme.gradient(
        Enzyme.ReverseWithPrimal, value_gradient.loss, state)
    return result.val, result.derivs[1]
end

zero_normal = Reactant.to_rarray([1.0, 0.0, 0.0])
unit_normal = Reactant.to_rarray([1.0, 0.0, 1.0])

kernel = Reactant.Compiler.compile(
    ValueGradient(KernelLoss()), (zero_normal,);
    compile_options=Reactant.CompileOptions(raise=true, raise_first=true))
tensor = Reactant.Compiler.compile(
    ValueGradient(TensorLoss()), (zero_normal,);
    compile_options=Reactant.CompileOptions(
        disable_reduce_slice_fusion_passes=true,
        shardy_passes=:none))

for (label, state, expected) in (
        ("n=0", zero_normal, [0.0, 0.0, 0.0]),
        ("n=1", unit_normal, [0.0, 1.0, 0.0]))
    kernel_value, kernel_gradient = kernel(state)
    tensor_value, tensor_gradient = tensor(state)
    kernel_gradient = Array(kernel_gradient)
    tensor_gradient = Array(tensor_gradient)

    @assert Reactant.to_number(kernel_value) ==
            Reactant.to_number(tensor_value) == 0.0
    @assert tensor_gradient == expected
    @assert isnan(kernel_gradient[2])
    if label == "n=0"
        @assert isnan(kernel_gradient[3])
    else
        @assert kernel_gradient[3] == 0.0
    end

    println(label)
    println("  raised KA: value=", Reactant.to_number(kernel_value),
        " gradient=", kernel_gradient)
    println("  tensor:    value=", Reactant.to_number(tensor_value),
        " gradient=", tensor_gradient)
end