Skip to content

Commit

Permalink
some changes to the CUDA ext
Browse files Browse the repository at this point in the history
  • Loading branch information
@leios
leios committed Dec 19, 2024
1 parent 89aae31 commit 14cf074
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 42 deletions.
51 changes: 17 additions & 34 deletions ext/MollyCUDAExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ function cuda_threads_blocks_specific(n_inters)
return n_threads_gpu, n_blocks
end

function pairwise_force_gpu!(fs_mat, coords::AbstractArray{SVector{D, C}}, velocities, atoms,
function Molly.pairwise_force_gpu!(fs_mat, coords::CuArray{SVector{D, C}}, velocities, atoms,

Check warning on line 40 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L40 

Added line #L40 was not covered by tests
boundary, pairwise_inters, nbs, step_n, force_units, ::Val{T}) where {D, C, T}
if typeof(nbs) == NoNeighborList
kernel = @cuda launch=false pairwise_force_kernel_nonl!(
Expand Down Expand Up @@ -122,7 +122,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, velocities_var,
function pairwise_force_kernel_nonl!(forces::CuArray{T}, coords_var, velocities_var,

Check warning on line 125 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L125 

Added line #L125 was not covered by tests
atoms_var, boundary, inters, step_n, ::Val{D}, ::Val{F}) where {T, D, F}
coords = CUDA.Const(coords_var)
velocities = CUDA.Const(velocities_var)
Expand Down Expand Up @@ -190,24 +190,7 @@ function pairwise_force_kernel_nonl!(forces::AbstractArray{T}, coords_var, veloc
return nothing
end

@inline function sum_pairwise_forces(inters, atom_i, atom_j, ::Val{F}, special, coord_i, coord_j,
boundary, vel_i, vel_j, step_n) where F
dr = vector(coord_i, coord_j, boundary)
f_tuple = ntuple(length(inters)) do inter_type_i
force_gpu(inters[inter_type_i], dr, atom_i, atom_j, F, special, coord_i, coord_j, boundary,
vel_i, vel_j, step_n)
end
f = sum(f_tuple)
if unit(f[1]) != F
# This triggers an error but it isn't printed
# See https://discourse.julialang.org/t/error-handling-in-cuda-kernels/79692
# for how to throw a more meaningful error
error("wrong force unit returned, was expecting $F but got $(unit(f[1]))")
end
return f
end

function specific_force_gpu!(fs_mat, inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_force_gpu!(fs_mat, inter_list::InteractionList1Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 193 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L193 

Added line #L193 was not covered by tests
velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_1_atoms_kernel!(fs_mat,
Expand All @@ -216,7 +199,7 @@ function specific_force_gpu!(fs_mat, inter_list::InteractionList1Atoms, coords::
return fs_mat
end

function specific_force_gpu!(fs_mat, inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_force_gpu!(fs_mat, inter_list::InteractionList2Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 202 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L202 

Added line #L202 was not covered by tests
velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_2_atoms_kernel!(fs_mat,
Expand All @@ -225,7 +208,7 @@ function specific_force_gpu!(fs_mat, inter_list::InteractionList2Atoms, coords::
return fs_mat
end

function specific_force_gpu!(fs_mat, inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_force_gpu!(fs_mat, inter_list::InteractionList3Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 211 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L211 

Added line #L211 was not covered by tests
velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_3_atoms_kernel!(fs_mat,
Expand All @@ -234,7 +217,7 @@ function specific_force_gpu!(fs_mat, inter_list::InteractionList3Atoms, coords::
return fs_mat
end

function specific_force_gpu!(fs_mat, inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_force_gpu!(fs_mat, inter_list::InteractionList4Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 220 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L220 

Added line #L220 was not covered by tests
velocities, atoms, boundary, step_n, force_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_force_4_atoms_kernel!(fs_mat,
Expand All @@ -243,7 +226,7 @@ function specific_force_gpu!(fs_mat, inter_list::InteractionList4Atoms, coords::
return fs_mat
end

function specific_force_1_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,
function Molly.specific_force_1_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,

Check warning on line 229 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L229 

Added line #L229 was not covered by tests
step_n, is_var, inters_var, ::Val{D}, ::Val{F}) where {D, F}
coords = CUDA.Const(coords_var)
velocities = CUDA.Const(velocities_var)
Expand All @@ -266,7 +249,7 @@ function specific_force_1_atoms_kernel!(forces, coords_var, velocities_var, atom
return nothing
end

function specific_force_2_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,
function Molly.specific_force_2_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,

Check warning on line 252 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L252 

Added line #L252 was not covered by tests
step_n, is_var, js_var, inters_var, ::Val{D}, ::Val{F}) where {D, F}
coords = CUDA.Const(coords_var)
velocities = CUDA.Const(velocities_var)
Expand All @@ -292,7 +275,7 @@ function specific_force_2_atoms_kernel!(forces, coords_var, velocities_var, atom
return nothing
end

function specific_force_3_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,
function Molly.specific_force_3_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,

Check warning on line 278 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L278 

Added line #L278 was not covered by tests
step_n, is_var, js_var, ks_var, inters_var, ::Val{D}, ::Val{F}) where {D, F}
coords = CUDA.Const(coords_var)
velocities = CUDA.Const(velocities_var)
Expand Down Expand Up @@ -320,7 +303,7 @@ function specific_force_3_atoms_kernel!(forces, coords_var, velocities_var, atom
return nothing
end

function specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,
function Molly.specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atoms_var, boundary,

Check warning on line 306 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L306 

Added line #L306 was not covered by tests
step_n, is_var, js_var, ks_var, ls_var, inters_var,
::Val{D}, ::Val{F}) where {D, F}
coords = CUDA.Const(coords_var)
Expand Down Expand Up @@ -352,7 +335,7 @@ function specific_force_4_atoms_kernel!(forces, coords_var, velocities_var, atom
return nothing
end

function pairwise_pe_gpu!(pe_vec_nounits, coords::AbstractArray{SVector{D, C}}, velocities, atoms,
function Molly.pairwise_pe_gpu!(pe_vec_nounits, coords::CuArray{SVector{D, C}}, velocities, atoms,

Check warning on line 338 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L338 

Added line #L338 was not covered by tests
boundary, pairwise_inters, nbs, step_n, energy_units,
::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_pairwise(length(nbs))
Expand All @@ -375,10 +358,10 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun
i, j, special = neighbors[inter_i]
coord_i, coord_j, vel_i, vel_j = coords[i], coords[j], velocities[i], velocities[j]
dr = vector(coord_i, coord_j, boundary)
pe = potential_energy_gpu(inters[1], dr, atoms[i], atoms[j], E, special, coord_i, coord_j,
pe = Molly.potential_energy_gpu(inters[1], dr, atoms[i], atoms[j], E, special, coord_i, coord_j,

Check warning on line 361 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L361 

Added line #L361 was not covered by tests
boundary, vel_i, vel_j, step_n)
for inter in inters[2:end]
pe += potential_energy_gpu(inter, dr, atoms[i], atoms[j], E, special, coord_i, coord_j,
pe += Molly.potential_energy_gpu(inter, dr, atoms[i], atoms[j], E, special, coord_i, coord_j,

Check warning on line 364 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L364 

Added line #L364 was not covered by tests
boundary, vel_i, vel_j, step_n)
end
if unit(pe) != E
Expand All @@ -389,7 +372,7 @@ function pairwise_pe_kernel!(energy, coords_var, velocities_var, atoms_var, boun
return nothing
end

function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList1Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList1Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 375 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L375 

Added line #L375 was not covered by tests
velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_1_atoms_kernel!(
Expand All @@ -398,7 +381,7 @@ function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList1Atoms, coo
return pe_vec_nounits
end

function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList2Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList2Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 384 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L384 

Added line #L384 was not covered by tests
velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_2_atoms_kernel!(
Expand All @@ -407,7 +390,7 @@ function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList2Atoms, coo
return pe_vec_nounits
end

function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList3Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList3Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 393 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L393 

Added line #L393 was not covered by tests
velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_3_atoms_kernel!(
Expand All @@ -416,7 +399,7 @@ function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList3Atoms, coo
return pe_vec_nounits
end

function specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList4Atoms, coords::AbstractArray{SVector{D, C}},
function Molly.specific_pe_gpu!(pe_vec_nounits, inter_list::InteractionList4Atoms, coords::CuArray{SVector{D, C}},

Check warning on line 402 in ext/MollyCUDAExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/MollyCUDAExt.jl#L402 

Added line #L402 was not covered by tests
velocities, atoms, boundary, step_n, energy_units, ::Val{T}) where {D, C, T}
n_threads_gpu, n_blocks = cuda_threads_blocks_specific(length(inter_list))
CUDA.@sync @cuda threads=n_threads_gpu blocks=n_blocks specific_pe_4_atoms_kernel!(
Expand Down
14 changes: 7 additions & 7 deletions src/interactions/implicit_solvent.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -850,7 +850,7 @@ end
numer = 2 * r_d0_strip + 9 * r_d0_strip^5 / 5
I_grad -= 10 * neck_scale * m0 * numer / (denom^2 * unit(dist_cutoff))
end
Atomix.@atomic :monotonic Is[i] += ustrip(unit(dist_cutoff)^-1, I)
Atomix.@atomic Is[i] += ustrip(unit(dist_cutoff)^-1, I)

Check warning on line 853 in src/interactions/implicit_solvent.jl

View check run for this annotation

Codecov / codecov/patch

src/interactions/implicit_solvent.jl#L853 

Added line #L853 was not covered by tests
I_grads[i, j] += ustrip(unit(dist_cutoff)^-2, I_grad)
end
end
Expand Down Expand Up @@ -1036,18 +1036,18 @@ end
dGpol_dalpha2_ij = -Gpol * exp_term * (1 + D_term) / (2 * denominator2)

change_born_force_i = dGpol_dalpha2_ij * Bj
Atomix.@atomic :monotonic born_forces_mod_ustrip[i] += ustrip(change_born_force_i)
Atomix.@atomic born_forces_mod_ustrip[i] += ustrip(change_born_force_i)

Check warning on line 1039 in src/interactions/implicit_solvent.jl

View check run for this annotation

Codecov / codecov/patch

src/interactions/implicit_solvent.jl#L1039 

Added line #L1039 was not covered by tests
if i != j
change_born_force_j = dGpol_dalpha2_ij * Bi
Atomix.@atomic :monotonic born_forces_mod_ustrip[j] += ustrip(change_born_force_j)
Atomix.@atomic born_forces_mod_ustrip[j] += ustrip(change_born_force_j)

Check warning on line 1042 in src/interactions/implicit_solvent.jl

View check run for this annotation

Codecov / codecov/patch

src/interactions/implicit_solvent.jl#L1042 

Added line #L1042 was not covered by tests
fdr = dr * dGpol_dr
if unit(fdr[1]) != F
error("wrong force unit returned, was expecting $F but got $(unit(fdr[1]))")
end
for dim in 1:D
fval = ustrip(fdr[dim])
Atomix.@atomic :monotonic forces[dim, i] += fval
Atomix.@atomic :monotonic forces[dim, j] += -fval
Atomix.@atomic forces[dim, i] += fval
Atomix.@atomic forces[dim, j] += -fval

Check warning on line 1050 in src/interactions/implicit_solvent.jl

View check run for this annotation

Codecov / codecov/patch

src/interactions/implicit_solvent.jl#L1049-L1050 

Added lines #L1049 - L1050 were not covered by tests
end
end
end
Expand Down Expand Up @@ -1088,8 +1088,8 @@ end
end
for dim in 1:D
fval = ustrip(fdr[dim])
Atomix.@atomic :monotonic forces[dim, i] += fval
Atomix.@atomic :monotonic forces[dim, j] += -fval
Atomix.@atomic forces[dim, i] += fval
Atomix.@atomic forces[dim, j] += -fval

Check warning on line 1092 in src/interactions/implicit_solvent.jl

View check run for this annotation

Codecov / codecov/patch

src/interactions/implicit_solvent.jl#L1091-L1092 

Added lines #L1091 - L1092 were not covered by tests
end
end
end
Expand Down
3 changes: 2 additions & 1 deletion src/types.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ export
masses,
charges,
MollyCalculator,
ASECalculator
ASECalculator,
NoNeighborList

const DefaultFloat = Float64

Expand Down

0 comments on commit 14cf074

Please sign in to comment.