Julia: Update memory management (#123)

Author: cortner

LICENSE.md

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 sphericart contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

julia/Project.toml

@@ -1,21 +1,23 @@
 name = "SpheriCart"
 uuid = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 authors = ["Christoph Ortner <[email protected]> and contributors"]
-version = "0.0.3"
+version = "0.1.1"
 
 [deps]
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
 
-[compat] 
-ForwardDiff = "0.10" 
-ObjectPools = "0.3.0"
-OffsetArrays = "1.13" 
-StaticArrays = "1.9"
+[compat]
+Bumper = "0.6.0"
+StrideArrays = "0.1.28"
+ForwardDiff = "0.10"
 LinearAlgebra = "1.8"
+OffsetArrays = "1.13"
+StaticArrays = "1.9"
 julia = "1.8.0, 1.9.0, 1.10.0"
 
 [extras]

julia/src/SpheriCart.jl

@@ -1,6 +1,6 @@
 module SpheriCart
 
-using StaticArrays, OffsetArrays, ObjectPools
+using StaticArrays, OffsetArrays, Bumper, StrideArrays
 
 export SolidHarmonics, 
        compute, 

julia/src/api.jl

@@ -1,3 +1,4 @@
+
 """
 `struct SolidHarmonics` : datatype representing a solid harmonics basis. 
 
@@ -36,7 +37,6 @@ See documentation for more details.
 """
 struct SolidHarmonics{L, NORM, STATIC, T1}
    Flm::OffsetMatrix{T1, Matrix{T1}}
-   cache::TSafe{ArrayPool{FlexArrayCache}}
 end
 
 function SolidHarmonics(L::Integer; 
@@ -45,7 +45,7 @@ function SolidHarmonics(L::Integer;
                         T = Float64) 
    Flm = generate_Flms(L; normalisation = normalisation, T = T)
    @assert eltype(Flm) == T   
-   SolidHarmonics{L, normalisation, static, T}(Flm, TSafe(ArrayPool(FlexArrayCache)))
+   SolidHarmonics{L, normalisation, static, T}(Flm)
 end
 
 @inline (basis::SolidHarmonics)(args...) = compute(basis, args...)
@@ -82,28 +82,23 @@ function compute!(Z::AbstractMatrix,
    nX = length(Rs)
    T = promote_type(T1, T2)
 
-   # allocate temporary arrays from an array cache 
-   temps = (x = acquire!(basis.cache, :x,  (nX, ),    T),
-            y = acquire!(basis.cache, :y,  (nX, ),    T),
-            z = acquire!(basis.cache, :z,  (nX, ),    T), 
-           r² = acquire!(basis.cache, :r2, (nX, ),    T),
-            s = acquire!(basis.cache, :s,  (nX, L+1), T),
-            c = acquire!(basis.cache, :c,  (nX, L+1), T),
-            Q = acquire!(basis.cache, :Q,  (nX, sizeY(L)), T), 
-            Flm = basis.Flm )
-
-   # the actual evaluation kernel 
-   solid_harmonics!(Z, Val{L}(), Rs, temps)
-
-   # release the temporary arrays back into the cache
-   # (don't release Flm!!)
-   release!(temps.x)
-   release!(temps.y)
-   release!(temps.z)
-   release!(temps.r²)
-   release!(temps.s)
-   release!(temps.c)
-   release!(temps.Q)
+   @no_escape begin 
+
+      # allocate temporary arrays from an array cache 
+      temps = (x = @alloc(T, nX), 
+               y = @alloc(T, nX),
+               z = @alloc(T, nX), 
+              r² = @alloc(T, nX),
+               s = @alloc(T, nX, L+1), 
+               c = @alloc(T, nX, L+1),
+               Q = @alloc(T, nX, sizeY(L)),
+             Flm = basis.Flm )
+
+      # the actual evaluation kernel 
+      solid_harmonics!(Z, Val{L}(), Rs, temps)
+
+      nothing
+   end # @no_escape 
 
    return Z 
 end 
@@ -152,28 +147,22 @@ function compute_with_gradients!(
    nX = length(Rs)
    T = promote_type(T1, T2)
 
-   # allocate temporary arrays from an array cache 
-   temps = (x = acquire!(basis.cache, :x,  (nX, ),    T),
-            y = acquire!(basis.cache, :y,  (nX, ),    T),
-            z = acquire!(basis.cache, :z,  (nX, ),    T), 
-           r² = acquire!(basis.cache, :r2, (nX, ),    T),
-            s = acquire!(basis.cache, :s,  (nX, L+1), T),
-            c = acquire!(basis.cache, :c,  (nX, L+1), T),
-            Q = acquire!(basis.cache, :Q,  (nX, sizeY(L)), T),
-            Flm = basis.Flm )
-
-   # the actual evaluation kernel 
-   solid_harmonics_with_grad!(Z, dZ, Val{L}(), Rs, temps)
-
-   # release the temporary arrays back into the cache
-   # (don't release Flm!!)
-   release!(temps.x)
-   release!(temps.y)
-   release!(temps.z)
-   release!(temps.r²)
-   release!(temps.s)
-   release!(temps.c)
-   release!(temps.Q)
+   @no_escape begin 
+
+      # allocate temporary arrays from an array cache 
+      temps = (x = @alloc(T, nX),    
+               y = @alloc(T, nX),    
+               z = @alloc(T, nX),    
+              r² = @alloc(T, nX),    
+               s = @alloc(T, nX, L+1), 
+               c = @alloc(T, nX, L+1), 
+               Q = @alloc(T, nX, sizeY(L)), 
+             Flm = basis.Flm )
+
+      # the actual evaluation kernel 
+      solid_harmonics_with_grad!(Z, dZ, Val{L}(), Rs, temps)
+      nothing
+   end
 
    return Z 
 end 

julia/src/normalisations.jl

@@ -2,7 +2,7 @@
 
 
 # Generates the `F[l, m]` values exactly as described in the 
-# `sphericart` publication. 
+# `sphericart` publication.  This gives L2-orthonormality. 
 function _generate_Flms(L::Integer, ::Union{Val{:sphericart}, Val{:L2}}, T=Float64)
    Flm = OffsetMatrix(zeros(L+1, L+1), (-1, -1))
    for l = 0:L
@@ -14,6 +14,15 @@ function _generate_Flms(L::Integer, ::Union{Val{:sphericart}, Val{:L2}}, T=Float
    return Flm
 end
 
+function _generate_Flms(L::Integer, ::Val{:racah}, T=Float64)
+   Flm = _generate_Flms(L, Val(:L2), T)
+   for l = 0:L 
+      for m = 0:l
+         Flm[l, m] = Flm[l, m] * sqrt(4*pi/(2*l+1))
+      end
+   end
+   return Flm
+end
 
 
 """
@@ -23,8 +32,8 @@ generate_Flms(L; normalisation = :L2, T = Float64)
 generate the `F[l,m]` prefactors in the definitions of the solid harmonics; 
 see `sphericart` publication for details. The default normalisation generates 
 a basis that is L2-orthonormal on the unit sphere. Other normalisations: 
-- `:sphericart` the same as `:L2`
-- `:p4ml` same normalisation as used in `Polynomials4ML.jl`
+- `:sphericart` the same as `:L2`, gives L2-orthonormality, i.e. ∫ |Ylm|² = 1
+- `:racah` gives Racah normalization for which ∫ |Ylm|² = 4π/(2l+1).
 """
 generate_Flms(L::Integer; normalisation = :L2, T = Float64) = 
       _generate_Flms(L, Val(normalisation), T)

julia/src/spherical.jl

@@ -42,12 +42,10 @@ See documentation for more details.
 struct SphericalHarmonics{L, NORM, STATIC, T1}
    solids::SolidHarmonics{L, NORM, STATIC, T1}
    # Flm::OffsetMatrix{T1, Matrix{T1}}
-   cache::TSafe{ArrayPool{FlexArrayCache}}
 end
 
 SphericalHarmonics(L::Integer; kwargs...) = 
-      SphericalHarmonics(SolidHarmonics(L; kwargs...), 
-                         TSafe(ArrayPool(FlexArrayCache)))
+      SphericalHarmonics(SolidHarmonics(L; kwargs...))
 
 
 @inline (basis::SphericalHarmonics)(args...) = compute(basis, args...)
@@ -75,16 +73,16 @@ end
 # --------------------- 
 #  batched api 
 
-function _normalise_Rs!(basis::SphericalHarmonics, 
+function _normalise_Rs!(rs, Rs_norm, 
+                        basis::SphericalHarmonics, 
                         Rs::AbstractVector{SVector{3, T1}}) where {T1}
    nX = length(Rs) 
-   rs = acquire!(basis.cache, :rs, (nX, ), T1)
-   Rs_norm = acquire!(basis.cache, :Rs_norm, (nX, ), SVector{3, T1})
+   @assert length(rs) == length(Rs_norm) == nX
    @inbounds @simd ivdep for i = 1:nX
       rs[i] = norm(Rs[i])
       Rs_norm[i] = Rs[i] / rs[i]
    end
-   return rs, Rs_norm 
+   return nothing 
 end
 
 function _rescale_∇Z2∇Y!(∇Z::AbstractMatrix, Rs_norm, rs)
@@ -99,44 +97,59 @@ function _rescale_∇Z2∇Y!(∇Z::AbstractMatrix, Rs_norm, rs)
 end
 
 function compute(basis::SphericalHarmonics, 
-                         Rs::AbstractVector{<: SVector{3}})
-   rs, Rs_norm = _normalise_Rs!(basis, Rs)
-   Y = compute(basis.solids, Rs_norm)
-   release!(Rs_norm)
-   release!(rs)
+                  Rs::AbstractVector{<: SVector{3, T1}}
+                  ) where {T1}  
+   @no_escape begin     
+      nX = length(Rs)              
+      rs = @alloc(T1, nX)
+      Rs_norm = @alloc(SVector{3, T1}, nX)
+      _normalise_Rs!(rs, Rs_norm, basis, Rs)
+      Y = compute(basis.solids, Rs_norm)
+   end
    return Y
 end
 
 function compute!(Y, basis::SphericalHarmonics, 
-                     Rs::AbstractVector{<: SVector{3}})
-   rs, Rs_norm = _normalise_Rs!(basis, Rs)
-   compute!(Y, basis.solids, Rs_norm)
-   release!(Rs_norm)
-   release!(rs)
+                  Rs::AbstractVector{<: SVector{3, T1}}
+                  ) where {T1}
+   @no_escape begin       
+      nX = length(Rs)            
+      rs = @alloc(T1, nX)
+      Rs_norm = @alloc(SVector{3, T1}, nX)
+      _normalise_Rs!(rs, Rs_norm, basis, Rs)
+      compute!(Y, basis.solids, Rs_norm)
+      nothing 
+   end
    return Y
 end
 
 
 function compute_with_gradients(basis::SphericalHarmonics, 
-                           Rs::AbstractVector{<: SVector{3}})
-   rs, Rs_norm = _normalise_Rs!(basis, Rs)
-   Y, ∇Z = compute_with_gradients(basis.solids, Rs_norm)
-   _rescale_∇Z2∇Y!(∇Z, Rs_norm, rs)
-
-   release!(Rs_norm)
-   release!(rs)
-
+                  Rs::AbstractVector{<: SVector{3, T1}}
+                  ) where {T1} 
+   @no_escape begin     
+      nX = length(Rs)              
+      rs = @alloc(T1, nX)
+      Rs_norm = @alloc(SVector{3, T1}, nX)
+      _normalise_Rs!(rs, Rs_norm, basis, Rs)
+      Y, ∇Z = compute_with_gradients(basis.solids, Rs_norm)
+      _rescale_∇Z2∇Y!(∇Z, Rs_norm, rs)
+      nothing 
+   end 
    return Y, ∇Z
 end 
 
 function compute_with_gradients!(Y, ∇Y, basis::SphericalHarmonics, 
-                                 Rs::AbstractVector{<: SVector{3}})
-   rs, Rs_norm = _normalise_Rs!(basis, Rs)
-   compute_with_gradients!(Y, ∇Y, basis.solids, Rs_norm)
-   _rescale_∇Z2∇Y!(∇Y, Rs_norm, rs)
-
-   release!(Rs_norm)
-   release!(rs)
-
+                        Rs::AbstractVector{<: SVector{3, T1}}
+                        ) where {T1}
+   @no_escape begin      
+      nX = length(Rs)             
+      rs = @alloc(T1, nX)
+      Rs_norm = @alloc(SVector{3, T1}, nX)
+      _normalise_Rs!(rs, Rs_norm, basis, Rs)
+      compute_with_gradients!(Y, ∇Y, basis.solids, Rs_norm)
+      _rescale_∇Z2∇Y!(∇Y, Rs_norm, rs)
+      nothing 
+   end 
    return Y, ∇Y
 end 

julia/test/test_solidharmonics.jl

@@ -99,6 +99,25 @@ for ntest = 1:10
    @test cond(G) < 1.5
 end
 
+## 
+
+@info("check Racah scaling")
+
+L = 3
+racah = SolidHarmonics(L; normalisation=:racah)
+
+for ntest = 1:10
+   local Z
+   Rs = [ rand_sphere() for _ = 1:10_000 ] 
+   Z = compute(racah, Rs)
+   G = (Z' * Z) / length(Rs) * 4 * π
+   dd = [ 4*π/(2*l+1) for l = 0:L for m = -l:l ]
+   D = Diagonal(dd)
+   Drtinv = Diagonal(dd.^(-0.5))
+   @test norm(G - D) < 1.0
+   @test cond( Drtinv * G * Drtinv ) < 1.5
+end
+
 
 ##