Skip to content

Commit 7c5597e

Browse files
leiosleios
committed
bits and bobs
1 parent e8364a3 commit 7c5597e

File tree

9 files changed

+90
-91
lines changed

9 files changed

+90
-91
lines changed

‎src/chain_rules.jl‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,7 @@ function ChainRulesCore.rrule(::typeof(pairwise_force_gpu), coords::AbstractArra
354354
ArrayType = find_arrat_type(coords)
355355
n_atoms = length(atoms)
356356
z = zero(T)
357-
fs_mat = zeros(backend, T, D, n_atoms)
357+
fs_mat = KernelAbstractions.zeros(backend, T, D, n_atoms)
358358
d_coords = zero(coords)
359359
d_atoms = ArrayType([Atom(charge=z, mass=z, σ=z, ϵ=z) for _ in 1:n_atoms])
360360
n_threads_gpu = gpu_threads_blocks_pairwise(length(nbs))
@@ -422,7 +422,7 @@ function ChainRulesCore.rrule(::typeof(pairwise_pe_gpu), coords::AbstractArray{S
422422
ArrayType = get_array_type(coords)
423423
n_atoms = length(atoms)
424424
z = zero(T)
425-
pe_vec = zeros(backend, T, 1)
425+
pe_vec = KernelAbstractions.zeros(backend, T, 1)
426426
d_pe_vec = ArrayType([d_pe_vec_arg[1]])
427427
d_coords = zero(coords)
428428
d_atoms = ArrayType([Atom(charge=z, mass=z, σ=z, ϵ=z) for _ in 1:n_atoms])
@@ -535,7 +535,7 @@ function ChainRulesCore.rrule(::typeof(specific_force_gpu), inter_list,
535535

536536
function specific_force_gpu_pullback(d_fs_mat)
537537
backend = get_backend(coords)
538-
fs_mat = zeros(backend, T, D, length(coords))
538+
fs_mat = KernelAbstractions.zeros(backend, T, D, length(coords))
539539
d_inter_list = zero(inter_list)
540540
d_coords = zero(coords)
541541
n_threads_gpu = get_threads_blocks_specific(length(inter_list))
@@ -658,7 +658,7 @@ function ChainRulesCore.rrule(::typeof(specific_pe_gpu), inter_list,
658658
function specific_pe_gpu_pullback(d_pe_vec_arg)
659659
backend = get_backend(coords)
660660
array_type = get_array_type(coords)
661-
pe_vec = zeros(backend, T, 1)
661+
pe_vec = KernelAbstractions.zeros(backend, T, 1)
662662
d_pe_vec = ArrayType([d_pe_vec_arg[1]])
663663
d_inter_list = zero(inter_list)
664664
d_coords = zero(coords)
@@ -754,14 +754,14 @@ function ChainRulesCore.rrule(::typeof(gbsa_born_gpu), coords::AbstractArray{SVe
754754
function gbsa_born_gpu_pullback(d_args)
755755
backend = get_backend(coords)
756756
n_atoms = length(coords)
757-
d_Is = d_args[1] == ZeroTangent() ? zeros(backend, T, n_atoms) : d_args[1]
758-
d_I_grads = d_args[2] == ZeroTangent() ? zeros(backend, T, n_atoms, n_atoms) : d_args[2]
759-
Is = zeros(backend, T, n_atoms)
760-
I_grads = zeros(backend, T, n_atoms, n_atoms)
757+
d_Is = d_args[1] == ZeroTangent() ? KernelAbstractions.zeros(backend, T, n_atoms) : d_args[1]
758+
d_I_grads = d_args[2] == ZeroTangent() ? KernelAbstractions.zeros(backend, T, n_atoms, n_atoms) : d_args[2]
759+
Is = KernelAbstractions.zeros(backend, T, n_atoms)
760+
I_grads = KernelAbstractions.zeros(backend, T, n_atoms, n_atoms)
761761
d_coords = zero(coords)
762762
d_offset_radii = zero(offset_radii)
763763
d_scaled_offset_radii = zero(scaled_offset_radii)
764-
grad_neck_scale = zeros(backend, T, 1)
764+
grad_neck_scale = KernelAbstractions.zeros(backend, T, 1)
765765
d_d0s = zero(d0s)
766766
d_m0s = zero(m0s)
767767
n_inters = n_atoms ^ 2
@@ -862,14 +862,14 @@ function ChainRulesCore.rrule(::typeof(gbsa_force_1_gpu), coords::AbstractArray{
862862
function gbsa_force_1_gpu_pullback(d_args)
863863
backend = get_backend(coords)
864864
n_atoms = length(coords)
865-
d_fs_mat = d_args[1] == ZeroTangent() ? zeros(backend, T, D, n_atoms) : d_args[1]
866-
d_born_forces_mod_ustrip = d_args[2] == ZeroTangent() ? zeros(backend, T, n_atoms) : d_args[2]
867-
fs_mat = zeros(backend, T, D, n_atoms)
868-
born_forces_mod_ustrip = zeros(backend, T, n_atoms)
865+
d_fs_mat = d_args[1] == ZeroTangent() ? KernelAbstractions.zeros(backend, T, D, n_atoms) : d_args[1]
866+
d_born_forces_mod_ustrip = d_args[2] == ZeroTangent() ? KernelAbstractions.zeros(backend, T, n_atoms) : d_args[2]
867+
fs_mat = KernelAbstractions.zeros(backend, T, D, n_atoms)
868+
born_forces_mod_ustrip = KernelAbstractions.zeros(backend, T, n_atoms)
869869
d_coords = zero(coords)
870-
grad_factor_solute = zeros(backend, T, 1)
871-
grad_factor_solvent = zeros(backend, T, 1)
872-
grad_kappa = zeros(backend, T, 1)
870+
grad_factor_solute = KernelAbstractions.zeros(backend, T, 1)
871+
grad_factor_solvent = KernelAbstractions.zeros(backend, T, 1)
872+
grad_kappa = KernelAbstractions.zeros(backend, T, 1)
873873
d_Bs = zero(Bs)
874874
d_chs = zero(chs)
875875
n_inters = n_atoms_to_n_pairs(n_atoms) + n_atoms
@@ -931,7 +931,7 @@ function ChainRulesCore.rrule(::typeof(gbsa_force_2_gpu), coords::AbstractArray{
931931
function gbsa_force_2_gpu_pullback(d_fs_mat)
932932
backend = get_backend(coords)
933933
n_atoms = length(coords)
934-
fs_mat = zeros(backend, T, D, n_atoms)
934+
fs_mat = KernelAbstractions.zeros(backend, T, D, n_atoms)
935935
d_coords = zero(coords)
936936
d_born_forces = zero(born_forces)
937937
d_offset_radii = zero(offset_radii)

‎src/energy.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ function potential_energy(sys::System{D, AT, T}, neighbors=nothing;
229229
n_threads::Integer=Threads.nthreads()) where {D, AT <: AbstractGPUArray, T}
230230
n_atoms = length(sys)
231231
val_ft = Val(T)
232-
pe_vec = zeros(get_backend(sys.coords), T, 1)
232+
pe_vec = KernelAbstractions.zeros(get_backend(sys.coords), T, 1)
233233

234234
pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters))
235235
if length(pairwise_inters_nonl) > 0

‎src/force.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,7 @@ function forces(sys::System{D, AT, T}, neighbors=nothing;
307307
n_threads::Integer=Threads.nthreads()) where {D, T, AT <: AbstractGPUArray}
308308
n_atoms = length(sys)
309309
val_ft = Val(T)
310-
fs_mat = zeros(get_backend(sys.coords), T, D, n_atoms)
310+
fs_mat = KernelAbstractions.zeros(get_backend(sys.coords), T, D, n_atoms)
311311

312312
pairwise_inters_nonl = filter(!use_neighbors, values(sys.pairwise_inters))
313313
if length(pairwise_inters_nonl) > 0

‎src/interactions/implicit_solvent.jl‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -733,8 +733,8 @@ function gbsa_born_gpu(coords::AbstractArray{SVector{D, C}}, offset_radii, scale
733733
::Val{T}) where {D, C, T}
734734
backend = get_backend(coords)
735735
n_atoms = length(coords)
736-
Is_nounits = zeros(backend, T, n_atoms)
737-
I_grads_nounits = zeros(backend, T, n_atoms, n_atoms)
736+
Is_nounits = KernelAbstractions.zeros(backend, T, n_atoms)
737+
I_grads_nounits = KernelAbstractions.zeros(backend, T, n_atoms, n_atoms)
738738
n_inters = n_atoms ^ 2
739739
n_threads_gpu = gpu_threads_blocks_gbsa(n_inters)
740740

@@ -748,11 +748,11 @@ function gbsa_born_gpu(coords::AbstractArray{SVector{D, C}}, offset_radii, scale
748748
return Is, I_grads
749749
end
750750

751-
@kernel function gbsa_born_kernel!(Is, I_grads, @Const(coords_var),
752-
@Const(offset_radii_var),
753-
@Const(scaled_offset_radii_var),
751+
@kernel function gbsa_born_kernel!(Is, I_grads, @Const(coords),
752+
@Const(offset_radii),
753+
@Const(scaled_offset_radii),
754754
dist_cutoff, offset, neck_scale, neck_cut,
755-
@Const(d0s_var), @Const(m0s_var), boundary,
755+
@Const(d0s), @Const(m0s), boundary,
756756
::Val{C}) where C
757757

758758
n_atoms = length(coords)
@@ -928,8 +928,8 @@ function gbsa_force_1_gpu(coords::AbstractArray{SVector{D, C}}, boundary, dist_c
928928
force_units) where {D, C, T}
929929
backend = get_backend(coords)
930930
n_atoms = length(coords)
931-
fs_mat = zeros(backend, T, D, n_atoms)
932-
born_forces_mod_ustrip = zeros(backend, T, n_atoms)
931+
fs_mat = KernelAbstractions.zeros(backend, T, D, n_atoms)
932+
born_forces_mod_ustrip = KernelAbstractions.zeros(backend, T, n_atoms)
933933
n_inters = n_atoms_to_n_pairs(n_atoms) + n_atoms
934934
n_threads_gpu = gpu_threads_blocks_gbsa(n_inters)
935935

@@ -946,7 +946,7 @@ function gbsa_force_2_gpu(coords::AbstractArray{SVector{D, C}}, boundary, dist_c
946946
force_units, ::Val{T}) where {D, C, T}
947947
backend = get_backend(coords)
948948
n_atoms = length(coords)
949-
fs_mat = zeros(backend, T, D, n_atoms)
949+
fs_mat = KernelAbstractions.zeros(backend, T, D, n_atoms)
950950
n_inters = n_atoms ^ 2
951951
n_threads_gpu = gpu_threads_blocks_gbsa(n_inters)
952952

@@ -959,9 +959,9 @@ function gbsa_force_2_gpu(coords::AbstractArray{SVector{D, C}}, boundary, dist_c
959959
end
960960

961961
@kernel function gbsa_force_1_kernel!(forces, born_forces_mod_ustrip,
962-
@Const(coords_var), boundary, dist_cutoff,
962+
@Const(coords), boundary, dist_cutoff,
963963
factor_solute, factor_solvent, kappa,
964-
@Const(Bs_var), @Const(atom_charges_var),
964+
@Const(Bs), @Const(atom_charges),
965965
::Val{D}, ::Val{F}) where {D, F}
966966

967967
n_atoms = length(coords)
@@ -1015,10 +1015,10 @@ end
10151015
end
10161016
end
10171017

1018-
@kernel function gbsa_force_2_kernel!(forces, born_forces, @Const(coords_var),
1019-
boundary, dist_cutoff, @Const(or_var),
1020-
@Const(sor_var), @Const(Bs_var),
1021-
@Const(B_grads_var), @Const(I_grads_var),
1018+
@kernel function gbsa_force_2_kernel!(forces, born_forces, @Const(coords),
1019+
boundary, dist_cutoff, @Const(or),
1020+
@Const(sor), @Const(Bs),
1021+
@Const(B_grads), @Const(I_grads),
10221022
::Val{D}, ::Val{F}) where {D, F}
10231023

10241024
n_atoms = length(coords)

0 commit comments

Comments
 (0)