GPU error with custom pairwise interactions about molly.jl

[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?