draft multithreaded forces

Project.toml

@@ -7,7 +7,9 @@ version = "0.0.1-dev"
 ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e"
 ACEpotentials = "3b96b61c-0fcc-4693-95ed-1ef9f35fcc53"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+ChunkSplitters = "ae650224-84b6-46f8-82ea-d812ca08434e"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
+Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
@@ -24,8 +26,8 @@ StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 
 [compat]
 NamedTupleTools = "0.14"
-julia = "1"
 Polynomials4ML = "0.2.7"
+julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/julip_calculator.jl

@@ -56,15 +56,26 @@ function energy_new(ace::UFACE, at::Atoms)
 end
 
 
-function forces_new!(F, ace::UFACE, at::Atoms)
+function forces_new(ace::UFACE, at::Atoms; 
+                     domain = 1:length(at), 
+                     nlist = neighbourlist(at, cutoff(ace)) )
+   TF = eltype(eltype(at.X)) 
+   F = zeros(SVector{3, Float64}, length(at))
+   forces_new!(F, ace, at; domain = domain, nlist = nlist)
+   return F 
+end
+
+function forces_new!(F, ace::UFACE, at::Atoms; 
+                     domain = 1:length(at), 
+                     nlist = neighbourlist(at, cutoff(ace))
+                  )
    TF = eltype(eltype(at.X))
-   nlist = neighbourlist(at, cutoff(ace))
    maxneigs = ACEpotentials.JuLIP.maxneigs(nlist) 
    Rs = acquire!(ace.pool, :calc_Rs, (maxneigs,), SVector{3, TF})
    Zs = acquire!(ace.pool, :calc_Zs, (maxneigs,), AtomicNumber)
    tmp = (R = unwrap(Rs), Z = unwrap(Zs),)
 
-   for i = 1:length(at) 
+   for i in domain 
       Js, Rs, Zs = ACEpotentials.JuLIP.Potentials.neigsz!(tmp, nlist, at, i)
       z0 = at.Z[i] 
       _, dEs = evaluate_ed(ace, Rs, Zs, z0)
@@ -81,4 +92,18 @@ function forces_new!(F, ace::UFACE, at::Atoms)
    release!(Zs)
 
    return F
-end
+end
+
+
+
+using Folds, ChunkSplitters
+
+function forces_new_mt(ace::UFACE, at::Atoms; 
+                        domain = 1:length(at), 
+                        executor = ThreadedEx(),
+                        ntasks = Threads.nthreads())
+   nlist = neighbourlist(at, cutoff(ace))
+   return Folds.sum( collect(chunks(domain, ntasks)), executor ) do (d, _)
+      forces_new(ace, at; domain=d, nlist = nlist)
+   end
+end

test/test_julip.jl

@@ -83,7 +83,18 @@ F2 = forces(jpot, at)
 @info("New calculator")
 F = zeros(SVector{3, Float64}, length(at))
 @btime UltraFastACE.forces_new!($F, $uf_ace, $at)
+@btime UltraFastACE.forces_new($uf_ace, $at)
 
+## 
+
+JuLIP.usethreads!(true)
+at = rand_struct(; rep = 7)
+@info("Multithreaded - JuLIP")
+@time forces(jpot, at)
+@info("Single threaded - new")
+@time UltraFastACE.forces_new(uf_ace, at)
+@info("Multithreaded - new")
+@time UltraFastACE.forces_new_mt(uf_ace, at)
 
 # @time forces(pot, at)
 # @time forces(jpot, at)
@@ -95,4 +106,5 @@ F = zeros(SVector{3, Float64}, length(at))
       # UltraFastACE.energy_new(uf_ace, at)
       UltraFastACE.forces_new!(F, uf_ace, at)
    end
-end
+end
+