pull out inner function for mutlithreaded nested loops to make logic clearer

Author: maximerischard

src/GPE.jl

@@ -215,6 +215,20 @@ function update_mll!(gp::GPE; noise::Bool=true, domean::Bool=true, kern::Bool=tr
     gp
 end
 
+function _dmll_kern_row!(dmll, buf, k, ααinvcKI, X, data, j, dim, nparams)
+    # diagonal
+    dKij_dθ!(buf, k, X, X, data, j, j, dim, nparams)
+    @inbounds for iparam in 1:nparams
+        dmll[iparam] += buf[iparam] * ααinvcKI[j, j] / 2.0
+    end
+    # off-diagonal
+    @inbounds for i in 1:j-1
+        dKij_dθ!(buf, k, X, X, data, i, j, dim, nparams)
+        @simd for iparam in 1:nparams
+            dmll[iparam] += buf[iparam] * ααinvcKI[i, j]
+        end
+    end
+end
 """
     dmll_kern!((dmll::AbstractVector, k::Kernel, X::AbstractMatrix, data::KernelData, ααinvcKI::AbstractMatrix))
 
@@ -236,18 +250,8 @@ function dmll_kern!(dmll::AbstractVector, k::Kernel, X::AbstractMatrix, data::Ke
         kthread = kcopies[Threads.threadid()]
         bufthread = buffercopies[Threads.threadid()]
         dmllthread = dmllcopies[Threads.threadid()]
-        # diagonal
-        dKij_dθ!(bufthread, kthread, X, X, data, j, j, dim, nparams)
-        for iparam in 1:nparams
-            dmllthread[iparam] += bufthread[iparam] * ααinvcKI[j, j] / 2.0
-        end
-        # off-diagonal
-        for i in j+1:nobs
-            dKij_dθ!(bufthread, kthread, X, X, data, i, j, dim, nparams)
-            @simd for iparam in 1:nparams
-                dmllthread[iparam] += bufthread[iparam] * ααinvcKI[i, j]
-            end
-        end
+        _dmll_kern_row!(dmllthread, bufthread, kthread, 
+                        ααinvcKI, X, data, j, dim, nparams)
     end
 
     dmll[:] = sum(dmllcopies) # sum up the results from all threads

src/kernels/kernels.jl

@@ -36,20 +36,28 @@ function cov(k::Kernel, X1::AbstractMatrix, X2::AbstractMatrix, data::KernelData
     cov!(cK, k, X1, X2, data)
 end
 
+function _cov_row!(cK, k, X::AbstractMatrix, data, j, dim)
+    cK[j,j] = cov_ij(k, X, X, data, j, j, dim)
+    @inbounds for i in 1:j-1
+        cK[i,j] = cov_ij(k, X, X, data, i, j, dim)
+        cK[j,i] = cK[i,j]
+    end
+end
 function cov!(cK::AbstractMatrix, k::Kernel, X::AbstractMatrix, data::KernelData=EmptyData())
     dim, nobs = size(X)
     (nobs,nobs) == size(cK) || throw(ArgumentError("cK has size $(size(cK)) and X has size $(size(X))"))
     kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()] # in case k is not threadsafe (e.g. ADkernel)
     @inbounds Threads.@threads for j in 1:nobs
         kthread = kcopies[Threads.threadid()]
-        cK[j,j] = cov_ij(kthread, X, X, data, j, j, dim)
-        for i in 1:j-1
-            cK[i,j] = cov_ij(kthread, X, X, data, i, j, dim)
-            cK[j,i] = cK[i,j]
-        end
+        _cov_row!(cK, k, X, data, j, dim)
     end
     return cK
 end
+function _cov_row!(cK, k, X1::AbstractMatrix, X2::AbstractMatrix, data, i, dim, nobs2)
+    @inbounds for j in 1:nobs2
+        cK[i,j] = cov_ij(k, X1, X2, data, i, j, dim)
+    end
+end
 """
     cov!(cK::AbstractMatrix, k::Kernel, X1::AbstractMatrix, X2::AbstractMatrix, data::KernelData=EmptyData())
 
@@ -67,9 +75,7 @@ function cov!(cK::AbstractMatrix, k::Kernel, X1::AbstractMatrix, X2::AbstractMat
     kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()]
     @inbounds Threads.@threads for i in 1:nobs1
         kthread = kcopies[Threads.threadid()]
-        for j in 1:nobs2
-            cK[i,j] = cov_ij(kthread, X1, X2, data, i, j, dim)
-        end
+        _cov_row!(cK, kthread, X1, X2, data, i, dim, nobs2)
     end
     return cK
 end
@@ -97,20 +103,28 @@ cov(k::Kernel, X::AbstractMatrix, data::KernelData=EmptyData()) = cov(k, X, X, d
     end
 end
 
+function _grad_slice_row!(dK, k, X::AbstractMatrix, data, j, p, dim)
+    dK[j,j] = dKij_dθp(k,X,X,data,j,j,p,dim)
+    @inbounds @simd for i in 1:(j-1)
+        dK[i,j] = dKij_dθp(k,X,X,data,i,j,p,dim)
+        dK[j,i] = dK[i,j]
+    end
+end
 function grad_slice!(dK::AbstractMatrix, k::Kernel, X::AbstractMatrix, data::KernelData, p::Int)
     dim, nobs = size(X)
     (nobs,nobs) == size(dK) || throw(ArgumentError("dK has size $(size(dK)) and X has size $(size(X))"))
     kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()]
     @inbounds Threads.@threads for j in 1:nobs
         kthread = kcopies[Threads.threadid()]
-        dK[j,j] = dKij_dθp(kthread,X,X,data,j,j,p,dim)
-        @simd for i in 1:(j-1)
-            dK[i,j] = dKij_dθp(kthread,X,X,data,i,j,p,dim)
-            dK[j,i] = dK[i,j]
-        end
+        _grad_slice_row!(dK, kthread, X, data, j, p, dim)
     end
     return dK
 end
+function _grad_slice_row!(dK, k, X1::AbstractMatrix, X2::AbstractMatrix, data, i, p, dim, nobs2)
+    @inbounds @simd for j in 1:nobs2
+        dK[i,j] = dKij_dθp(k,X1,X2,data,i,j,p,dim)
+    end
+end
 function grad_slice!(dK::AbstractMatrix, k::Kernel, X1::AbstractMatrix, X2::AbstractMatrix, data::KernelData, p::Int)
     if X1 === X2
         return grad_slice!(dK, k, X1, data, p)
@@ -123,9 +137,7 @@ function grad_slice!(dK::AbstractMatrix, k::Kernel, X1::AbstractMatrix, X2::Abst
     kcopies = [deepcopy(k) for _ in 1:Threads.nthreads()]
     @inbounds Threads.@threads for i in 1:nobs1
         kthread = kcopies[Threads.threadid()]
-        @simd for j in 1:nobs2
-            dK[i,j] = dKij_dθp(kthread,X1,X2,data,i,j,p,dim)
-        end
+        _grad_slice_row!(dK, kthread, X1, X2, data, i, p, dim, nobs2)
     end
     return dK
 end