GPU error with custom pairwise interactions

[bondrewd]

Hi! I have implemented some custom pairwise interactions that use the DistanceNeighborFinder. The implementations work on the CPU. I then tried to use the interactions on the GPU. I followed the examples on the Molly documentation, casting the relevant arrays as CuArray and ensuring all the interactions evaluated to true when using the isbitstype function. However, I get an error that is related to the CUDA.jl package:

Error: LoadError: InvalidIRError: compiling MethodInstance for Molly.pairwise_force_kernel!(... ... ...)
Reason: unsupported call to an unknown function (call to ijl_get_nth_field_checked) 

The error only appears when I use three (or more) pairwise interactions together:

pairwise_interactions = (
    Interaction1(
        use_neighbors=true,
    ),
    Interaction2(
        use_neighbors=true,
    ),
    Interaction3(
        use_neighbors=true,
    ),
)

The CUDA kernel is compiled successfully if I use any pair of interactions (interaction1 + interaction2, interaction1 + interaction3, etc.). By looking at the stack trace, I tracked down the error back to these lines 36-39 of the src/cuda.jl file:

f = force_gpu(inters[1], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
for inter in inters[2:end]
    f += force_gpu(inter, dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
end

I still need to familiarize myself more with the CUDA.jl package, but maybe the instructions generated when using tuples of pairwise interactions are not supported yet for generating CUDA kernels. I tried different versions of the same lines of code, hoping to find an alternative representation that is compatible with CUDA.jl, and found that the following way makes the error disappear:

f = sum(force_gpu(inter, dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special) for inter in inters)

By introducing this change, the CI passes all the tests successfully. I have not yet checked the performance implications, but I can try if there are input files for benchmarking the GPU code.

Finally, I can submit a PR if this change looks ok.

UPDATE: I found that using other combinations of pairwise interactions makes the error reappear. I will further investigate the cause...

[jgreener64]

Thanks for reporting this, I'm not sure I've ever used 3 interactions on the GPU and it looks like that's the issue. One workaround might be to keep the loop but loop 2:N where N is the length of the tuple in the type domain. Using sum in some arrangement would work but may have performance implications.

I'll have a more thorough look next week.

[bondrewd]

Thank you for looking into it! I would like to let you know that the reason for the reappearance of the error after modifying the force evaluation was because of the energy calculation (probably because of the same reason). On src/cuda.jl lines 211-216, I found the same structure as in the force calculation:

pe = potential_energy(inters[1], dr, coord_i, coord_j, atoms[i], atoms[j],
                      boundary, special)
for inter in inters[2:end]
    pe += potential_energy(inter, dr, coord_i, coord_j, atoms[i], atoms[j],
                           boundary, special)
end

Making the changes on both the force and energy calculations now makes the errors completely disappear for any combination of three or more pairwise potentials.

Regarding the loop over 2:N, I tried but found the error reappeared. I will also try to figure out the reason.

[jgreener64]

In particular this occurs when different interactions are used, using three of the same interactions is okay.

This works:

f  = force_gpu(inters[1], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
f += force_gpu(inters[2], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
f += force_gpu(inters[3], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)

But this errors:

f = force_gpu(inters[1], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
for inter_i in 2:3
    f += force_gpu(inters[inter_i], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
end

Extracting it out to a function also errors.

A solution would be to auto-generate the code at the top. I tried metaprogramming for this but couldn't get it to work, I'm not the strongest at that though. I got Base.Cartesian.@nexprs working with a fixed n but couldn't make it a variable in the GPU kernel.

One workaround is to define functions like:

function addforces(inters::Tuple{<:Any, <:Any, <:Any}, dr, coord_i, coord_j, atom_i, atom_j, boundary, special)
    return force_gpu(inters[1], dr, coord_i, coord_j, atom_i, atom_j, boundary, special) +
           force_gpu(inters[2], dr, coord_i, coord_j, atom_i, atom_j, boundary, special) +
           force_gpu(inters[3], dr, coord_i, coord_j, atom_i, atom_j, boundary, special)
end

This works but would require defining multiple functions. Maybe the function definitions could be written with metaprogramming.

I wonder if @vchuravy has any ideas as to why the first example in this comment works in a CUDA kernel and the second doesn't, or if there is an easy workaround.

MWE, for reference:

using Molly, CUDA
boundary = CubicBoundary(1.0u"nm")
coords = CuArray(place_atoms(100, boundary; min_dist=0.1u"nm"))
atoms = CuArray([Atom(σ=0.02u"nm", ϵ=0.1u"kJ * mol^-1") for _ in 1:100])
nf = DistanceNeighborFinder(eligible=CuArray(trues(100, 100)), dist_cutoff=0.2u"nm")
lj = LennardJones(use_neighbors=true)
coul = Coulomb(use_neighbors=true)
ss = SoftSphere(use_neighbors=true)
sys2 = System(coords=coords, atoms=atoms, boundary=boundary, neighbor_finder=nf, pairwise_inters=(lj, coul,))
sys3 = System(coords=coords, atoms=atoms, boundary=boundary, neighbor_finder=nf, pairwise_inters=(lj, coul, ss))
neighbors = find_neighbors(sys2)
forces(sys2, neighbors) # Works
forces(sys3, neighbors) # Errors

[vchuravy]

What is typeof(inters).

But the likely solution is something like:

f += sum(ntuple(Val(N)) do inter_i
    @inline force_gpu(inters[inter_i], dr, coord_i, coord_j, atoms[i], atoms[j], boundary, special)
end)

Important is that N is a compile time constant. Since inters is a tupleN=length(inters) may work. Essentially we are forcing the compiler to unroll this code statically.

[jgreener64]

Thanks Valentin. Indeed the force case seems to work with permutations of the above and length(inters) seems to be available at compile time. This works without slowdown, and with Enzyme:

f_tuple = ntuple(length(inters)) do inter_type_i
    force_gpu(inters[inter_type_i], dr, coord_i, coord_j, atom_i, atom_j, boundary, special)
end
f = sum(f_tuple)

typeof(inters) is a Tuple of 3 different concrete structs:

Tuple{
    LennardJones{false, NoCutoff, Int64, Int64, Unitful.FreeUnits{(kJ, nm^-1, mol^-1), 𝐋 𝐌 𝐍^-1 𝐓^-2, nothing}, Unitful.FreeUnits{(kJ, mol^-1), 𝐋^2 𝐌 𝐍^-1 𝐓^-2, nothing}},
    Coulomb{NoCutoff, Int64, Quantity{Float64, 𝐋^3 𝐌 𝐍^-1 𝐓^-2, Unitful.FreeUnits{(kJ, nm, mol^-1), 𝐋^3 𝐌 𝐍^-1 𝐓^-2, nothing}}, Unitful.FreeUnits{(kJ, nm^-1, mol^-1), 𝐋 𝐌 𝐍^-1 𝐓^-2, nothing}, Unitful.FreeUnits{(kJ, mol^-1), 𝐋^2 𝐌 𝐍^-1 𝐓^-2, nothing}},
    SoftSphere{false, NoCutoff, Unitful.FreeUnits{(kJ, nm^-1, mol^-1), 𝐋 𝐌 𝐍^-1 𝐓^-2, nothing}, Unitful.FreeUnits{(kJ, mol^-1), 𝐋^2 𝐌 𝐍^-1 𝐓^-2, nothing}}
}

Strangely, though, the potential energy case

pe_tuple = ntuple(length(inters)) do inter_type_i
    potential_energy(inters[inter_type_i], dr, coord_i, coord_j, atom_i, atom_j, boundary, special)
end
pe = sum(pe_tuple)

works for tuples of length 2 but fails on tuples of length 3 with a similar error:

Reason: unsupported call to an unknown function (call to ijl_get_nth_field_checked)

potential_energy returns a "simpler" value than force_gpu, a single value (both unitful and unitless cases fail) compared to a SVector, so I was surprised that it didn't work but the force did. Perhaps it is because potential_energy can be applied to more argument types, so it is struggling to infer the return type? I did try annotating the return type, adding an additional function boundary, using @inline and using Val but that did not seem to work.

Any idea what is going on there?