extend kahan sum for rank N arrays

Author: jalvesz

doc/specs/stdlib_intrinsics.md

@@ -48,7 +48,7 @@ If `dim` is absent, the output is a scalar of the same `type` and `kind` as to t
 
 #### Description
 
-The `stdlib_sum_kahan` function can replace the intrinsic `sum` for 1D `real` or `complex` arrays. It follows a chunked implementation which maximizes vectorization potential, complemented by an `elemental` kernel based on the [kahan summation](https://en.wikipedia.org/wiki/Kahan_summation_algorithm) strategy to reduce the round-off error:
+The `stdlib_sum_kahan` function can replace the intrinsic `sum` for `real` or `complex` arrays. It follows a chunked implementation which maximizes vectorization potential complemented by an `elemental` kernel based on the [kahan summation](https://en.wikipedia.org/wiki/Kahan_summation_algorithm) strategy to reduce the round-off error:
 
 ```fortran
 elemental subroutine kahan_kernel_<kind>(a,s,c)
@@ -67,6 +67,8 @@ end subroutine
 
 `res = ` [[stdlib_intrinsics(module):stdlib_sum_kahan(interface)]] ` (x [,mask] )`
 
+`res = ` [[stdlib_intrinsics(module):stdlib_sum_kahan(interface)]] ` (x, dim [,mask] )`
+
 #### Status
 
 Experimental
@@ -79,11 +81,13 @@ Pure function.
 
 `x`: 1D array of either `real` or `complex` type. This argument is `intent(in)`.
 
-`mask` (optional): 1D array of `logical` values. This argument is `intent(in)`.
+`dim` (optional): scalar of type `integer` with a value in the range from 1 to n, where n equals the rank of `x`.
+
+`mask` (optional): N-D array of `logical` values, with the same shape as `x`. This argument is `intent(in)`.
 
 #### Output value or Result value
 
-The output is a scalar of `type` and `kind` same as to that of `x`.
+If `dim` is absent, the output is a scalar of the same `type` and `kind` as to that of `x`. Otherwise, an array of rank n-1, where n equals the rank of `x`, and a shape similar to that of `x` with dimension `dim` dropped is returned.
 
 #### Example
 

src/stdlib_intrinsics.fypp

@@ -56,13 +56,14 @@ module stdlib_intrinsics
         !! version: experimental 
         !!
         !!### Summary 
-        !! Sum elements of rank 1 arrays. 
+        !! Sum elements of rank N arrays. 
         !! ([Specification](../page/specs/stdlib_intrinsics.html#stdlib_sum_kahan))
         !!
         !!### Description
         !! 
-        !! This interface provides standard conforming call for sum of elements of rank 1.
+        !! This interface provides standard conforming call for sum of elements of any rank.
         !! The 1-D base implementation follows a chunked approach combined with a kahan kernel for optimizing performance and increasing accuracy.
+        !! The `N-D` interfaces calls upon the `(N-1)-D` implementation. 
         !! Supported data types include `real` and `complex`.
         !!
         #:for rk, rt, rs in RC_KINDS_TYPES
@@ -75,6 +76,19 @@ module stdlib_intrinsics
             logical, intent(in) :: mask(:)
             ${rt}$ :: s
         end function
+        #:for rank in RANKS
+        pure module function stdlib_sum_kahan_${rank}$d_${rs}$( x, mask ) result( s )
+            ${rt}$, intent(in) :: x${ranksuffix(rank)}$
+            logical, intent(in), optional :: mask${ranksuffix(rank)}$
+            ${rt}$ :: s
+        end function
+        pure module function stdlib_sum_kahan_${rank}$d_dim_${rs}$( x , dim, mask ) result( s )
+            ${rt}$, intent(in) :: x${ranksuffix(rank)}$
+            integer, intent(in):: dim
+            logical, intent(in), optional :: mask${ranksuffix(rank)}$
+            ${rt}$ :: s${reduced_shape('x', rank, 'dim')}$
+        end function
+        #:endfor
         #:endfor
     end interface
     public :: stdlib_sum_kahan

src/stdlib_intrinsics_sum.fypp

@@ -9,15 +9,14 @@ submodule(stdlib_intrinsics) stdlib_intrinsics_sum
     use stdlib_kinds
     use stdlib_constants
     implicit none
-
-    integer, parameter :: chunk = 64
 
 contains
 
 !================= 1D Base implementations ============
 ! This implementation is based on https://github.com/jalvesz/fast_math
 #:for rk, rt, rs in RC_KINDS_TYPES
 pure module function stdlib_sum_1d_${rs}$(a) result(s)
+    integer, parameter :: chunk = 64
     ${rt}$, intent(in) :: a(:)
     ${rt}$ :: s
     ${rt}$ :: abatch(chunk)
@@ -39,6 +38,7 @@ pure module function stdlib_sum_1d_${rs}$(a) result(s)
 end function
 
 pure module function stdlib_sum_1d_${rs}$_mask(a,mask) result(s)
+    integer, parameter :: chunk = 64
     ${rt}$, intent(in) :: a(:)
     logical, intent(in) :: mask(:)
     ${rt}$ :: s
@@ -61,13 +61,14 @@ pure module function stdlib_sum_1d_${rs}$_mask(a,mask) result(s)
 end function
 
 pure module function stdlib_sum_kahan_1d_${rs}$(a) result(s)
+    integer, parameter :: chunk = 64
     ${rt}$, intent(in) :: a(:)
     ${rt}$ :: s
     ${rt}$ :: sbatch(chunk)
     ${rt}$ :: cbatch(chunk)
     integer :: i, dr, rr
     ! -----------------------------
-    dr = size(a)/(chunk)
+    dr = size(a)/chunk
     rr = size(a) - dr*chunk     
     sbatch = zero_${rs}$
     cbatch = zero_${rs}$
@@ -83,14 +84,15 @@ pure module function stdlib_sum_kahan_1d_${rs}$(a) result(s)
 end function
 
 pure module function stdlib_sum_kahan_1d_${rs}$_mask(a,mask) result(s)
+    integer, parameter :: chunk = 64
     ${rt}$, intent(in) :: a(:)
     logical, intent(in) :: mask(:)
     ${rt}$ :: s
     ${rt}$ :: sbatch(chunk)
     ${rt}$ :: cbatch(chunk)
     integer :: i, dr, rr
     ! -----------------------------
-    dr = size(a)/(chunk)
+    dr = size(a)/chunk
     rr = size(a) - dr*chunk     
     sbatch = zero_${rs}$
     cbatch = zero_${rs}$
@@ -181,4 +183,78 @@ end function
 #:endfor
 #:endfor
 
+#:for rk, rt, rs in RC_KINDS_TYPES
+#:for rank in RANKS
+pure module function stdlib_sum_kahan_${rank}$d_${rs}$( x , mask ) result( s )
+    ${rt}$, intent(in) :: x${ranksuffix(rank)}$
+    logical, intent(in), optional :: mask${ranksuffix(rank)}$
+    ${rt}$ :: s
+    if(.not.present(mask)) then
+        s = sum_recast(x,size(x))
+    else
+        s = sum_recast_mask(x,mask,size(x))
+    end if
+contains
+    pure ${rt}$ function sum_recast(b,n)
+      integer, intent(in) :: n
+      ${rt}$, intent(in) :: b(n)
+      sum_recast = stdlib_sum_kahan(b)
+    end function
+    pure ${rt}$ function sum_recast_mask(b,m,n)
+      integer, intent(in) :: n
+      ${rt}$, intent(in) :: b(n)
+      logical, intent(in) :: m(n)
+      sum_recast_mask = stdlib_sum_kahan(b,m)
+    end function
+end function
+
+pure module function stdlib_sum_kahan_${rank}$d_dim_${rs}$( x , dim, mask ) result( s )
+    ${rt}$, intent(in) :: x${ranksuffix(rank)}$
+    integer, intent(in):: dim
+    logical, intent(in), optional :: mask${ranksuffix(rank)}$
+    ${rt}$ :: s${reduced_shape('x', rank, 'dim')}$
+    integer :: j 
+
+    if(.not.present(mask)) then
+        if(dim<${rank}$)then
+            do j = 1, size(x,dim=${rank}$)
+                #:if rank == 2
+                s${select_subarray(rank-1, [(rank-1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(rank, 'j')])}$ )
+                #:else
+                s${select_subarray(rank-1, [(rank-1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(rank, 'j')])}$, dim=dim )
+                #:endif
+            end do
+        else
+            do j = 1, size(x,dim=1)
+                #:if rank == 2
+                s${select_subarray(rank-1, [(1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(1, 'j')])}$ )
+                #:else
+                s${select_subarray(rank-1, [(1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(1, 'j')])}$, dim=${rank-1}$ )
+                #:endif
+            end do
+        end if
+    else 
+        if(dim<${rank}$)then
+            do j = 1, size(x,dim=${rank}$)
+                #:if rank == 2
+                s${select_subarray(rank-1, [(rank-1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(rank, 'j')])}$, mask=mask${select_subarray(rank, [(rank, 'j')])}$ )
+                #:else
+                s${select_subarray(rank-1, [(rank-1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(rank, 'j')])}$, dim=dim, mask=mask${select_subarray(rank, [(rank, 'j')])}$ )
+                #:endif
+            end do
+        else
+            do j = 1, size(x,dim=1)
+                #:if rank == 2
+                s${select_subarray(rank-1, [(1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(1, 'j')])}$, mask=mask${select_subarray(rank, [(1, 'j')])}$ )
+                #:else
+                s${select_subarray(rank-1, [(1, 'j')])}$ = stdlib_sum_kahan( x${select_subarray(rank, [(1, 'j')])}$, dim=${rank-1}$, mask=mask${select_subarray(rank, [(1, 'j')])}$ )
+                #:endif
+            end do
+        end if
+    end if
+
+end function
+#:endfor
+#:endfor
+
 end submodule stdlib_intrinsics_sum

test/intrinsics/test_intrinsics.fypp

@@ -127,11 +127,18 @@ subroutine test_sum(error)
         call check(error, abs( sum(x) - stdlib_sum(x) )<tolerance*size(x) , "KO: full ndarray stdlib_sum" )
         if (allocated(error)) return
 
+        call check(error, abs( sum(x) - stdlib_sum_kahan(x) )<tolerance*size(x) , "KO: full ndarray stdlib_sum_kahan" )
+        if (allocated(error)) return
+
         !> sum over specific rank dim
         do i = 1, rank(x)
             call check(error, norm2( sum(x,dim=i) - stdlib_sum(x,dim=i) )<tolerance*size(x) ,&
                         "KO: ndarray stdlib_sum over dim "//to_string(i) )
             if (allocated(error)) return
+
+            call check(error, norm2( sum(x,dim=i) - stdlib_sum_kahan(x,dim=i) )<tolerance*size(x) ,&
+                        "KO: ndarray stdlib_sum_kahan over dim "//to_string(i) )
+            if (allocated(error)) return
         end do
     end block ndarray