gradients correct but with type instability

profile/profile_evaluate.jl

@@ -0,0 +1,81 @@
+
+using ACEpotentials, StaticArrays, BenchmarkTools, 
+      LinearAlgebra, UltraFastACE
+using ACEbase: evaluate, evaluate!, evaluate_d, evaluate_d!, 
+               evaluate_ed, evaluate_ed!
+
+##
+
+elements = [:Si,:O]
+
+model = acemodel(; elements = elements, order = 3, totaldegree = 12)
+pot = model.potential
+mbpot = pot.components[2]
+
+# convert to UFACE format 
+uf_ace = UltraFastACE.uface_from_ace1(mbpot; n_spl_points = 100)
+
+## ------------------------------------
+
+
+Nat = 30; 
+r0 = 0.9 * rnn(:Si); r1 = 1.3 * rnn(:Si)
+Rs = [ (r0 + (r1 - r0) * rand()) * ACE1.Random.rand_sphere() for _=1:Nat ]
+z0 = rand(AtomicNumber.(elements)) 
+Zs = [ rand(AtomicNumber.(elements)) for _ = 1:Nat ]
+
+v1 = evaluate(mbpot, Rs, Zs, z0)
+v2 = evaluate(uf_ace, Rs, Zs, z0)
+
+dv1 = evaluate_d(mbpot, Rs, Zs, z0)
+dv2 = deepcopy(dv1) 
+v3, _ = evaluate_ed!(dv2, uf_ace, Rs, Zs, z0)
+
+@show abs(v3 - v2)
+@show norm(dv1-dv2) / norm(dv1)
+
+##
+
+@info("ACE1 - allocating")
+@btime evaluate($mbpot, $Rs, $Zs, $z0)
+
+@info("ACE1 - pre-allocated")
+tmp = ACE1.alloc_temp(mbpot, length(Rs))
+@btime evaluate!($tmp, $mbpot, $Rs, $Zs, $z0)
+
+@info("UF_ACE")
+@btime evaluate($uf_ace, $Rs, $Zs, $z0)
+
+## 
+
+# @profview let uf_ace = uf_ace, Rs = Rs, Zs = Zs, z0 = z0
+#    for iter = 1:400_000
+#       evaluate(uf_ace, Rs, Zs, z0)
+#    end
+# end
+
+##
+# profiling the gradient code 
+
+@info("profiling the gradient code")
+
+@info("ACE1 - allocating")
+@btime evaluate_d($mbpot, $Rs, $Zs, $z0)
+
+@info("ACE1 - preallocated arrays")
+tmpd = ACE1.alloc_temp_d(mbpot, length(Rs))
+dEs1 = zeros(SVector{3, Float64}, length(Rs))
+@btime evaluate_d!($dEs1, $tmpd, $mbpot, $Rs, $Zs, $z0)
+
+@info("UF_ACE")
+dEs2 = zeros(SVector{3, Float64}, length(Rs))
+evaluate_ed!(dEs2, uf_ace, Rs, Zs, z0)
+@btime evaluate_ed!($dEs2, $uf_ace, $Rs, $Zs, $z0)
+
+##
+
+@profview let uf_ace = uf_ace, Rs = Rs, Zs = Zs, z0 = z0, dEs2 = dEs2 
+   for iter = 1:2_000
+      evaluate_ed!(dEs2, uf_ace, Rs, Zs, z0)
+   end
+end

src/uface.jl

@@ -83,23 +83,39 @@ function ACEbase.evaluate_ed!(∇φ, ace::UFACE_inner, Rs, Zs)
 
    # radial embedding 
    Rn, dRn = evaluate_ed(ace, rbasis, Rs, Zs)
+
+   # _dRn = reduce(vcat, dRn)
+   # @show any(isnan, _dRn)
 
    # angular embedding 
    Zlm, dZlm = evaluate_ylm_ed(ace, Rs)
 
+   # _dZlm = reduce(vcat, dZlm)
+   # @show any(isnan, _dZlm)
+
    # pooling 
    A = ace.abasis((Ez, Rn, Zlm))
 
+   # @show any(isnan, A)
+
    # n correlations  # compute with gradient 
-   φ, ∂φ_∂A = evaluate_and_gradient(ace.aadot, A)   
+   φ, ∂φ_∂A = evaluate_and_gradient(ace.aadot, A)  
+
+   # @show any(isnan, ∂φ_∂A)
 
    # backprop through A 
    ∂φ_∂Ez = BlackHole(TF) 
+   # ∂φ_∂Ez = zeros(TF, size(Ez))
    ∂φ_∂Rn = acquire!(ace.pool, :∂Rn, size(Rn), TF)
    ∂φ_∂Zlm = acquire!(ace.pool, :∂Zlm, size(Zlm), TF)
-   P4ML._pullback_evaluate!((∂φ_∂Ez, ∂φ_∂Rn, ∂φ_∂Zlm), ∂φ_∂A, 
+   fill!(∂φ_∂Rn, zero(TF))
+   fill!(∂φ_∂Zlm, zero(TF))
+   P4ML._pullback_evaluate!((∂φ_∂Ez, unwrap(∂φ_∂Rn), unwrap(∂φ_∂Zlm)), ∂φ_∂A, 
                              ace.abasis, (Ez, Rn, Zlm))
 
+   # @show any(isnan, ∂φ_∂Rn)
+   # @show any(isnan, ∂φ_∂Zlm)
+
    # backprop through the embeddings 
    # depending on whether there is a bottleneck here, this can be 
    # potentially implemented more efficiently without needing writing/reading 
@@ -116,13 +132,19 @@ function ACEbase.evaluate_ed!(∇φ, ace::UFACE_inner, Rs, Zs)
       end
    end
 
+   # _g = reduce(vcat, ∇φ)
+   # @show any(isnan, _g)
+
    # ... and Ylm 
    for i_lm = 1:size(Zlm, 2)
       for j = 1:length(Rs)
          ∇φ[j] += ∂φ_∂Zlm[j, i_lm] * dZlm[j, i_lm]
       end
    end
 
+   # _g = reduce(vcat, ∇φ)
+   # @show any(isnan, _g)
+
    # release the borrowed arrays 
    release!(Zlm)
    release!(Rn)