small fixup. CPU tests failing.

Author: leios

Project.toml

@@ -12,7 +12,6 @@ CellListMap = "69e1c6dd-3888-40e6-b3c8-31ac5f578864"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
-Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
@@ -60,7 +59,6 @@ CellListMap = "0.8.11, 0.9"
 ChainRules = "1.44"
 ChainRulesCore = "1"
 Chemfiles = "0.10.3"
-Colors = "0.10, 0.11, 0.12"
 Combinatorics = "1"
 DataStructures = "0.18"
 Distances = "0.10"
@@ -70,6 +68,7 @@ EzXML = "1"
 FLoops = "0.2"
 ForwardDiff = "0.10.35"
 GLMakie = "0.8, 0.9, 0.10"
+GPUArrays = "10"
 Graphs = "1.8"
 KernelAbstractions = "0.9"
 KernelDensity = "0.5, 0.6"

ext/MollyCUDAExt.jl

@@ -3,10 +3,9 @@ module MollyCUDAExt
 using Molly
 using CUDA
 using ChainRulesCore
+using Atomix
 
-# CUDA specific calls for Molly
-@non_differentiable CUDA.zeros(args...)
-CUDA.Const(nl::NoNeighborList) = nl
+CUDA.Const(nl::Molly.NoNeighborList) = nl
 
 # CUDA.jl kernels
 const WARPSIZE = UInt32(32)
@@ -493,4 +492,9 @@ function specific_pe_4_atoms_kernel!(energy, coords_var, boundary, is_var, js_va
     return nothing
 end
 
+# CUDA specific calls for Molly
+@non_differentiable CUDA.zeros(args...)
+@non_differentiable cuda_threads_blocks_pairwise(args...)
+@non_differentiable cuda_threads_blocks_specific(args...)
+
 end

ext/MollyGLMakieExt.jl

@@ -5,7 +5,6 @@ module MollyGLMakieExt
 
 using Molly
 using GLMakie
-using Colors
 using Unitful
 
 using LinearAlgebra
@@ -95,7 +94,7 @@ function Molly.visualize(coord_logger,
         push!(trail_positions, Observable(PointType.(ustrip_vec.(coords_start))))
         col = parse.(Colorant, color)
         alpha = 1 - (trail_i / (trails + 1))
-        alpha_col = RGBA.(red.(col), green.(col), blue.(col), alpha)
+        alpha_col = GLMakie.RGBAf.(red.(col), green.(col), blue.(col), alpha)
         scatter!(ax, trail_positions[end]; color=alpha_col,  markersize=markersize,
                     transparency=transparency, markerspace=:data, kwargs...)
     end

src/chain_rules.jl

@@ -1,7 +1,6 @@
 # Chain rules to allow differentiable simulations
-
 @non_differentiable n_infinite_dims(args...)
-@non_differentiable zeros(args...)
+@non_differentiable KernelAbstractions.zeros(args...)
 @non_differentiable get_backend(args...)
 @non_differentiable get_array_type(args...)
 @non_differentiable random_velocities(args...)

test/basic.jl

@@ -321,7 +321,7 @@ end
             )
             neighbors_gpu = find_neighbors(sys_gpu, nf_gpu)
             @test length(neighbors_gpu) == n_neighbors_ref
-            CUDA.allowscalar() do
+            GPUArrays.allowscalar() do
                 @test neighbors_gpu[10] isa Tuple{Int32, Int32, Bool}
             end
             @test identical_neighbors(neighbors_gpu, neighbors_ref)
@@ -339,7 +339,7 @@ end
             )
             neighbors_gpu = find_neighbors(sys_gpu, nf_gpu)
             @test length(neighbors_gpu) == n_neighbors_ref
-            AMDGPU.allowscalar() do
+            GPUArrays.allowscalar() do
                 @test neighbors_gpu[10] isa Tuple{Int32, Int32, Bool}
             end
             @test identical_neighbors(neighbors_gpu, neighbors_ref)

test/zygote.jl

@@ -431,7 +431,7 @@ end
     for (test_name, test_fn, test_tol) in test_runs
         for (platform, args) in platform_runs
             f = test_fn(args...)
-            grads_zygote = CUDA.allowscalar() do
+            grads_zygote = GPUArrays.allowscalar() do
                 gradient(f, params_dic)[1]
             end
             @test count(!iszero, values(grads_zygote)) == 67