Adding KernelAbstractions tooling for Molly and tests

Author: leios

.gitignore

@@ -2,7 +2,7 @@
 *.jl.*.cov
 *.jl.mem
 docs/build
-/Manifest.toml
+*Manifest.toml
 benchmark/tune.json
 benchmark/results
 .vscode/settings.json

Project.toml

@@ -8,7 +8,6 @@ Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 AtomsBase = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
 AtomsCalculators = "a3e0e189-c65a-42c1-833c-339540406eb1"
 BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -22,7 +21,9 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EzXML = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
 FLoops = "cc61a311-1640-44b5-9fba-1b764f453329"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KernelDensity = "5ab0869b-81aa-558d-bb23-cbf5423bbe9b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
@@ -42,10 +43,12 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [weakdeps]
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [extensions]
 MollyGLMakieExt = "GLMakie"
 MollyPythonCallExt = "PythonCall"
+MollyCUDAExt = "CUDA"
 
 [compat]
 Atomix = "0.1"
@@ -68,6 +71,7 @@ FLoops = "0.2"
 ForwardDiff = "0.10.35"
 GLMakie = "0.8, 0.9, 0.10"
 Graphs = "1.8"
+KernelAbstractions = "0.9"
 KernelDensity = "0.5, 0.6"
 LinearAlgebra = "1.9"
 NearestNeighbors = "0.4"

benchmark/benchmarks.jl

@@ -62,7 +62,8 @@ const starting_velocities = [random_velocity(atom_mass, 1.0u"K") for i in 1:n_at
 const starting_coords_f32 = [Float32.(c) for c in starting_coords]
 const starting_velocities_f32 = [Float32.(c) for c in starting_velocities]
 
-function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool)
+function test_sim(nl::Bool, parallel::Bool, f32::Bool,
+                  ArrayType::Type{AT}) where AT <: AbstractArray
     n_atoms = 400
     n_steps = 200
     atom_mass = f32 ? 10.0f0u"g/mol" : 10.0u"g/mol"
@@ -72,34 +73,27 @@ function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool)
     r0 = f32 ? 0.2f0u"nm" : 0.2u"nm"
     bonds = [HarmonicBond(k=k, r0=r0) for i in 1:(n_atoms ÷ 2)]
     specific_inter_lists = (InteractionList2Atoms(
-        gpu ? CuArray(Int32.(collect(1:2:n_atoms))) : Int32.(collect(1:2:n_atoms)),
-        gpu ? CuArray(Int32.(collect(2:2:n_atoms))) : Int32.(collect(2:2:n_atoms)),
-        gpu ? CuArray(bonds) : bonds,
+        ArrayType(Int32.(collect(1:2:n_atoms))),
+        ArrayType(Int32.(collect(2:2:n_atoms))),
+        ArrayType(bonds),
     ),)
 
     neighbor_finder = NoNeighborFinder()
     cutoff = DistanceCutoff(f32 ? 1.0f0u"nm" : 1.0u"nm")
     pairwise_inters = (LennardJones(use_neighbors=false, cutoff=cutoff),)
     if nl
         neighbor_finder = DistanceNeighborFinder(
-            eligible=gpu ? CuArray(trues(n_atoms, n_atoms)) : trues(n_atoms, n_atoms),
+            eligible=ArrayType(trues(n_atoms, n_atoms)),
             n_steps=10,
             dist_cutoff=f32 ? 1.5f0u"nm" : 1.5u"nm",
         )
         pairwise_inters = (LennardJones(use_neighbors=true, cutoff=cutoff),)
     end
 
-    if gpu
-        coords = CuArray(deepcopy(f32 ? starting_coords_f32 : starting_coords))
-        velocities = CuArray(deepcopy(f32 ? starting_velocities_f32 : starting_velocities))
-        atoms = CuArray([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
-                              ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms])
-    else
-        coords = deepcopy(f32 ? starting_coords_f32 : starting_coords)
-        velocities = deepcopy(f32 ? starting_velocities_f32 : starting_velocities)
-        atoms = [Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
-                      ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms]
-    end
+    coords = ArrayType(deepcopy(f32 ? starting_coords_f32 : starting_coords))
+    velocities = ArrayType(deepcopy(f32 ? starting_velocities_f32 : starting_velocities))
+    atoms = ArrayType([Atom(charge=f32 ? 0.0f0 : 0.0, mass=atom_mass, σ=f32 ? 0.2f0u"nm" : 0.2u"nm",
+                            ϵ=f32 ? 0.2f0u"kJ * mol^-1" : 0.2u"kJ * mol^-1") for i in 1:n_atoms])
 
     sys = System(
         atoms=atoms,
@@ -117,22 +111,22 @@ function test_sim(nl::Bool, parallel::Bool, f32::Bool, gpu::Bool)
 end
 
 runs = [
-    ("CPU"       , [false, false, false, false]),
-    ("CPU f32"   , [false, false, true , false]),
-    ("CPU NL"    , [true , false, false, false]),
-    ("CPU f32 NL", [true , false, true , false]),
+    ("CPU"       , [false, false, false, Array]),
+    ("CPU f32"   , [false, false, true , Array]),
+    ("CPU NL"    , [true , false, false, Array]),
+    ("CPU f32 NL", [true , false, true , Array]),
 ]
 if run_parallel_tests
-    push!(runs, ("CPU parallel"       , [false, true , false, false]))
-    push!(runs, ("CPU parallel f32"   , [false, true , true , false]))
-    push!(runs, ("CPU parallel NL"    , [true , true , false, false]))
-    push!(runs, ("CPU parallel f32 NL", [true , true , true , false]))
+    push!(runs, ("CPU parallel"       , [false, true , false, Array]))
+    push!(runs, ("CPU parallel f32"   , [false, true , true , Array]))
+    push!(runs, ("CPU parallel NL"    , [true , true , false, Array]))
+    push!(runs, ("CPU parallel f32 NL", [true , true , true , Array]))
 end
-if run_gpu_tests
-    push!(runs, ("GPU"       , [false, false, false, true]))
-    push!(runs, ("GPU f32"   , [false, false, true , true]))
-    push!(runs, ("GPU NL"    , [true , false, false, true]))
-    push!(runs, ("GPU f32 NL", [true , false, true , true]))
+if run_cuda_tests
+    push!(runs, ("GPU"       , [false, false, false, CuArray]))
+    push!(runs, ("GPU f32"   , [false, false, true , CuArray]))
+    push!(runs, ("GPU NL"    , [true , false, false, CuArray]))
+    push!(runs, ("GPU f32 NL", [true , false, true , CuArray]))
 end
 
 for (name, args) in runs

benchmark/protein.jl

@@ -11,7 +11,7 @@ const data_dir = normpath(dirname(pathof(Molly)), "..", "data")
 const ff_dir = joinpath(data_dir, "force_fields")
 const openmm_dir = joinpath(data_dir, "openmm_6mrr")
 
-function setup_system(gpu::Bool, f32::Bool, units::Bool)
+function setup_system(ArrayType::AbstractArray, f32::Bool, units::Bool)
     T = f32 ? Float32 : Float64
     ff = MolecularForceField(
         T,
@@ -27,7 +27,7 @@ function setup_system(gpu::Bool, f32::Bool, units::Bool)
     sys = System(
         joinpath(data_dir, "6mrr_equil.pdb"),
         ff;
-        velocities=gpu ? CuArray(velocities) : velocities,
+        velocities=ArrayType(velocities),
         units=units,
         gpu=gpu,
         dist_cutoff=(units ? dist_cutoff * u"nm" : dist_cutoff),
@@ -42,15 +42,15 @@ end
 
 runs = [
     # run_name                             gpu    parr   f32    units
-    ("CPU 1 thread"                      , false, false, false, true ),
-    ("CPU 1 thread f32"                  , false, false, true , true ),
-    ("CPU 1 thread f32 nounits"          , false, false, true , false),
-    ("CPU $n_threads threads"            , false, true , false, true ),
-    ("CPU $n_threads threads f32"        , false, true , true , true ),
-    ("CPU $n_threads threads f32 nounits", false, true , true , false),
-    ("GPU"                               , true , false, false, true ),
-    ("GPU f32"                           , true , false, true , true ),
-    ("GPU f32 nounits"                   , true , false, true , false),
+    ("CPU 1 thread"                      , Array, false, false, true ),
+    ("CPU 1 thread f32"                  , Array, false, true , true ),
+    ("CPU 1 thread f32 nounits"          , Array, false, true , false),
+    ("CPU $n_threads threads"            , Array, true , false, true ),
+    ("CPU $n_threads threads f32"        , Array, true , true , true ),
+    ("CPU $n_threads threads f32 nounits", Array, true , true , false),
+    ("GPU"                               , CuArray, false, false, true ),
+    ("GPU f32"                           , CuArray, false, true , true ),
+    ("GPU f32 nounits"                   , CuArray, false, true , false),
 ]
 
 for (run_name, gpu, parallel, f32, units) in runs

src/cuda.jl renamed to ext/MollyCUDAExt.jl

@@ -1,3 +1,12 @@
+module MollyCUDAExt
+
+using Molly
+using CUDA
+
+# CUDA specific calls for Molly
+@non_differentiable CUDA.zeros(args...)
+CUDA.Const(nl::NoNeighborList) = nl
+
 # CUDA.jl kernels
 const WARPSIZE = UInt32(32)
 
@@ -29,7 +38,7 @@ function cuda_threads_blocks_specific(n_inters)
     return n_threads_gpu, n_blocks
 end
 
-function pairwise_force_gpu(coords::AbstractArray{SVector{D, C}}, atoms, boundary,
+function pairwise_force_gpu(coords::CuArray{SVector{D, C}}, atoms, boundary,
                             pairwise_inters, nbs, force_units, ::Val{T}) where {D, C, T}
     fs_mat = CUDA.zeros(T, D, length(atoms))
 
@@ -112,7 +121,7 @@ That's why the calculations are done in the following order:
     h | 1 2 3 4 5 6
 ```
 =#
-function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, atoms_var, boundary, inters,
+function pairwise_force_kernel_nonl!(forces::CuArray{T}, coords_var, atoms_var, boundary, inters,
                                      ::Val{D}, ::Val{F}) where {T, D, F}
     coords = CUDA.Const(coords_var)
     atoms = CUDA.Const(atoms_var)
@@ -193,7 +202,7 @@ end
     return f
 end
 
-function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_force_gpu(inter_list::InteractionList1Atoms, coords::CuArray{SVector{D, C}},
                             boundary, force_units, ::Val{T}) where {D, C, T}
     fs_mat = CUDA.zeros(T, D, length(coords))
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -202,7 +211,7 @@ function specific_force_gpu(inter_list::InteractionList1Atoms, coords::AbstractA
     return fs_mat
 end
 
-function specific_force_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_force_gpu(inter_list::InteractionList2Atoms, coords::CuArray{SVector{D, C}},
                             boundary, force_units, ::Val{T}) where {D, C, T}
     fs_mat = CUDA.zeros(T, D, length(coords))
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -212,7 +221,7 @@ function specific_force_gpu(inter_list::InteractionList2Atoms, coords::AbstractA
     return fs_mat
 end
 
-function specific_force_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_force_gpu(inter_list::InteractionList3Atoms, coords::CuArray{SVector{D, C}},
                             boundary, force_units, ::Val{T}) where {D, C, T}
     fs_mat = CUDA.zeros(T, D, length(coords))
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -222,7 +231,7 @@ function specific_force_gpu(inter_list::InteractionList3Atoms, coords::AbstractA
     return fs_mat
 end
 
-function specific_force_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_force_gpu(inter_list::InteractionList4Atoms, coords::CuArray{SVector{D, C}},
                             boundary, force_units, ::Val{T}) where {D, C, T}
     fs_mat = CUDA.zeros(T, D, length(coords))
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -328,7 +337,7 @@ function specific_force_4_atoms_kernel!(forces, coords_var, boundary, is_var, js
     return nothing
 end
 
-function pairwise_pe_gpu(coords::AbstractArray{SVector{D, C}}, atoms, boundary,
+function pairwise_pe_gpu(coords::CuArray{SVector{D, C}}, atoms, boundary,
                          pairwise_inters, nbs, energy_units, ::Val{T}) where {D, C, T}
     pe_vec = CUDA.zeros(T, 1)
     n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs))
@@ -363,7 +372,7 @@ function pairwise_pe_kernel!(energy, coords_var, atoms_var, boundary, inters, ne
     return nothing
 end
 
-function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::CuArray{SVector{D, C}},
                          boundary, energy_units, ::Val{T}) where {D, C, T}
     pe_vec = CUDA.zeros(T, 1)
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -372,7 +381,7 @@ function specific_pe_gpu(inter_list::InteractionList1Atoms, coords::AbstractArra
     return pe_vec
 end
 
-function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::CuArray{SVector{D, C}},
                          boundary, energy_units, ::Val{T}) where {D, C, T}
     pe_vec = CUDA.zeros(T, 1)
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -381,7 +390,7 @@ function specific_pe_gpu(inter_list::InteractionList2Atoms, coords::AbstractArra
     return pe_vec
 end
 
-function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::CuArray{SVector{D, C}},
                          boundary, energy_units, ::Val{T}) where {D, C, T}
     pe_vec = CUDA.zeros(T, 1)
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -391,7 +400,7 @@ function specific_pe_gpu(inter_list::InteractionList3Atoms, coords::AbstractArra
     return pe_vec
 end
 
-function specific_pe_gpu(inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}},
+function specific_pe_gpu(inter_list::InteractionList4Atoms, coords::CuArray{SVector{D, C}},
                          boundary, energy_units, ::Val{T}) where {D, C, T}
     pe_vec = CUDA.zeros(T, 1)
     n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
@@ -482,3 +491,5 @@ function specific_pe_4_atoms_kernel!(energy, coords_var, boundary, is_var, js_va
     end
     return nothing
 end
+
+end

ext/MollyPythonCallExt.jl

@@ -6,7 +6,7 @@ module MollyPythonCallExt
 using Molly
 using PythonCall
 import AtomsCalculators
-using CUDA
+using GPUArrays
 using StaticArrays
 using Unitful
 
@@ -91,7 +91,7 @@ end
 
 uconvert_vec(x...) = uconvert.(x...)
 
-function AtomsCalculators.forces(sys::System{D, G, T},
+function AtomsCalculators.forces(sys::System{D, AT, T},
                                  ase_calc::ASECalculator;
                                  kwargs...) where {D, G, T}
     update_ase_calc!(ase_calc, sys)
@@ -105,7 +105,7 @@ function AtomsCalculators.forces(sys::System{D, G, T},
     else
         fs_unit = uconvert_vec.(sys.force_units, fs * u"eV/Å")
     end
-    return G ? CuArray(fs_unit) : fs_unit
+    return AT <: AbstractGPUArray ? AT(fs_unit) : fs_unit
 end
 
 function AtomsCalculators.potential_energy(sys::System{D, G, T},

src/Molly.jl

@@ -13,7 +13,8 @@ using ChainRules
 using ChainRulesCore
 import Chemfiles
 using Combinatorics
-using CUDA
+using KernelAbstractions
+using GPUArrays
 using DataStructures
 using Distances
 using Distributions
@@ -41,7 +42,7 @@ include("types.jl")
 include("units.jl")
 include("spatial.jl")
 include("cutoffs.jl")
-include("cuda.jl")
+include("kernels.jl")
 include("force.jl")
 include("interactions/lennard_jones.jl")
 include("interactions/soft_sphere.jl")

src/analysis.jl

@@ -162,7 +162,7 @@ Calculate the hydrodynamic radius of a set of coordinates.
 function hydrodynamic_radius(coords::AbstractArray{SVector{D, T}}, boundary) where {D, T}
     n_atoms = length(coords)
     diag_cpu = Diagonal(ones(T, n_atoms))
-    diag = isa(coords, CuArray) ? CuArray(diag_cpu) : diag_cpu
+    diag = isa(coords, AbstractGPUArray) ? get_array_type(coords)((diag_cpu)) : diag_cpu
     # Other approaches to removing the diagonal Inf didn't work with Zygote
     dists = distances(coords, boundary) .+ diag
     sum_inv_dists = sum(inv.(dists)) - sum(inv(diag))