ND arrays: use LAPACK internals via loops

perazz · perazz · commit 5f6b5e894745 · 2024-09-27T10:59:16.000+02:00
diff --git a/include/common.fypp b/include/common.fypp
@@ -298,4 +298,105 @@ ${prefix + joinstr.join([line.strip() for line in txt.split("\n")]) + suffix}$
   #:endif
 #:enddef
 
+#!
+#! Generates a list of loop variables
+#!
+#! Args:
+#!   varname(str): Name of the variable to be used as prefix
+#!   n      (int): Number of loop variables to be created
+#!   offset (int): Optional index offset
+#!
+#! Returns:
+#!   Variable definition string 
+#!
+#! E.g.,
+#!  loop_variables('j', 5)
+#!   -> "j1, j2, j3, j4, j5
+#!
+#:def loop_variables(varname, n, offset=0)
+  #:assert n > 0
+    #:call join_lines(joinstr=", ")
+      #:for i in range(1, n + 1)
+        ${varname}$${i+offset}$
+      #:endfor
+    #:endcall
+#:enddef
+
+#! Generates an array shape specifier from an N-D array size
+#!
+#! Args:
+#!   name  (str): Name of the original variable
+#!   rank  (int): Rank of the original variable
+#!   offset(int): optional offset of the dimension loop (default = 0)
+#!
+#! Returns:
+#!   Array rank suffix string enclosed in braces
+#!
+#! E.g.,
+#!   shape_from_array_size('mat', 5)}$
+#!   -> (size(mat,1),size(mat,2),size(mat,3),size(mat,4),size(mat,5))
+#!   shape_from_array_size('mat', 5, 2)}$
+#!   -> (size(mat,3),size(mat,4),size(mat,5),size(mat,6),size(mat,7))
+#!
+#:def shape_from_array_size(name, rank, offset=0)
+  #:assert rank > 0
+  #:call join_lines(joinstr=", ", prefix="(", suffix=")")
+    #:for i in range(1, rank + 1)
+      size(${name}$,${i+offset}$)
+    #:endfor
+  #:endcall
+#:enddef
+
+#! Generates an array shape specifier from an N-D array of sizes
+#!
+#! Args:
+#!   name  (str): Name of the original variable
+#!   rank  (int): Rank of the original variable
+#!   offset(int): optional offset of the dimension loop (default = 0)
+#!
+#! Returns:
+#!   Array rank suffix string enclosed in braces
+#!
+#! E.g.,
+#!   shape_from_array_data('mat', 5)}$
+#!   -> (1:mat(1),1:mat(2),1:mat(3),1:mat(4),1:mat(5))
+#!   shape_from_array_data('mat', 5, 2)}$
+#!   -> (1:mat(3),1:mat(4),1:mat(5),1:mat(6),1:mat(7))
+#!
+#:def shape_from_array_data(name, rank, offset=0)
+  #:assert rank > 0
+  #:call join_lines(joinstr=", ", prefix="(", suffix=")")
+    #:for i in range(1, rank + 1)
+      1:${name}$(${i+offset}$)
+    #:endfor
+  #:endcall
+#:enddef
+
+#!
+#! Start a sequence of loop with indexed variables over an N-D array
+#!
+#! Args:
+#!   varname    (str): Name of the variable to be used as prefix
+#!   matname    (str): Name of the variable to be used as array
+#!   n          (int): Number of nested loops to be created (1=innermost; n=outermost)
+#!   dim_offset (int): Optional dimension offset (1st loop is over dimension 1+dim_offset)
+#!   intent     (str): Optional indentation. Default: 8 spaces
+#!
+#!
+#:def loop_variables_start(varname, matname, n, dim_offset=0, indent=" "*8)
+  #:assert n > 0
+    #:for i in range(1, n + 1)
+${indent}$do ${varname}$${n+1+dim_offset-i}$ = lbound(${matname}$, ${n+1+dim_offset-i}$), ubound(${matname}$, ${n+1+dim_offset-i}$)
+    #:endfor
+#:enddef
+
+#:def loop_variables_end(n, indent=" "*8)
+#:assert n > 0
+  #:call join_lines(joinstr="; ",prefix=indent)
+    #:for i in range(1, n + 1)
+       enddo 
+    #:endfor
+  #:endcall  
+#:enddef
+
 #:endmute
diff --git a/src/stdlib_linalg.fypp b/src/stdlib_linalg.fypp
@@ -1234,7 +1234,7 @@ module stdlib_linalg
         #:for rank in range(2, MAXRANK + 1)
         pure module subroutine norm_${rank}$D_to_${rank-1}$D_${ii}$_${ri}$(a, nrm, order, dim, err)
            !> Input matrix a[..]
-           ${rt}$, intent(in), target :: a${ranksuffix(rank)}$
+           ${rt}$, intent(in) :: a${ranksuffix(rank)}$
            !> Dimension the norm is computed along
            integer(ilp), intent(in) :: dim        
            !> Norm of the matrix. (Same shape as `a`, with `dim` dropped).
diff --git a/src/stdlib_linalg_norms.fypp b/src/stdlib_linalg_norms.fypp
@@ -264,7 +264,7 @@ submodule(stdlib_linalg) stdlib_linalg_norms
     ! Internal implementation 
     pure module subroutine norm_${rank}$D_to_${rank-1}$D_${ii}$_${ri}$(a, nrm, order, dim, err)
         !> Input matrix a[..]
-        ${rt}$, intent(in), target :: a${ranksuffix(rank)}$
+        ${rt}$, intent(in) :: a${ranksuffix(rank)}$
         !> Dimension to collapse by computing the norm w.r.t other dimensions
         !  (dim must be defined before it is used for `nrm`)
         integer(ilp), intent(in) :: dim        
@@ -276,11 +276,17 @@ submodule(stdlib_linalg) stdlib_linalg_norms
         type(linalg_state_type), intent(out), optional :: err           
         
         type(linalg_state_type) :: err_
-        integer(ilp) :: sze,norm_request
+        integer(ilp) :: sze,lda,norm_request,${loop_variables('j',rank-1,1)}$
+        logical :: contiguous_data
         real(${rk}$) :: rorder
+        integer(ilp), dimension(${rank}$) :: spe,spack,perm,iperm
+        integer(ilp), dimension(${rank}$), parameter :: dim_range = [(lda,lda=1_ilp,${rank}$_ilp)]
+        ${rt}$, allocatable :: apack${ranksuffix(rank)}$
         intrinsic :: abs, sum, sqrt, norm2, maxval, minval, conjg
         
-        sze = size(a,kind=ilp)
+        ! Input matrix properties
+        sze = size (a,kind=ilp)
+        spe = shape(a,kind=ilp)
         
         ! Initialize norm to zero
         nrm = 0.0_${rk}$
@@ -304,28 +310,60 @@ submodule(stdlib_linalg) stdlib_linalg_norms
         if (err_%error()) then 
             call linalg_error_handling(err_,err)
             return
-        endif                   
+        endif     
         
-        select case(norm_request)            
-            case(NORM_ONE)
-                nrm = sum( abs(a) , dim = dim )                
-            case(NORM_TWO)                
-                #:if rt.startswith('complex')
-                nrm = sqrt( real( sum( a * conjg(a) , dim = dim ), ${rk}$) )
-                #:else
-                nrm = norm2( a , dim = dim )
-                #:endif
-            case(NORM_INF)
-                nrm = maxval( abs(a) , dim = dim )
-            case(NORM_MINUSINF)
-                nrm = minval( abs(a) , dim = dim )
-            case (NORM_POW_FIRST:NORM_POW_LAST)
-                rorder = 1.0_${rk}$ / norm_request
-                nrm = sum( abs(a) ** norm_request , dim = dim ) ** rorder
-            case default
-                err_ = linalg_state_type(this,LINALG_INTERNAL_ERROR,'invalid norm type after checking')
-                call linalg_error_handling(err_,err)            
-        end select
+        ! The norm's leading dimension
+        lda = spe(dim)    
+        
+        ! Check if input column data is contiguous
+        contiguous_data = dim==1 .or. all(norm_request/=[NORM_ONE,NORM_TWO])   
+        
+        ! Get packed data with the norm dimension as the first dimension
+        if (.not.contiguous_data) then 
+            
+            ! Permute array to map dim to 1
+            perm = [dim,pack(dim_range,dim_range/=dim)]
+            iperm(perm) = dim_range            
+            spack = spe(perm)            
+            apack = reshape(a, shape=spack, order=iperm)                 
+            
+${loop_variables_start('j', 'apack', rank-1, 1," "*12)}$
+            select case(norm_request)            
+                case(NORM_ONE)                
+                    nrm(${loop_variables('j',rank-1,1)}$) = &
+                        asum(lda,apack(:, ${loop_variables('j',rank-1,1)}$),incx=1_ilp)
+                case(NORM_TWO)                
+                    nrm(${loop_variables('j',rank-1,1)}$) = &
+                        nrm2(lda,apack(:, ${loop_variables('j',rank-1,1)}$),incx=1_ilp)                    
+            end select            
+${loop_variables_end(rank-1," "*12)}$            
+            
+        else
+            
+            select case(norm_request)      
+                case(NORM_ONE)             
+${loop_variables_start('j', 'a', rank-1, 1," "*20)}$                    
+                        nrm(${loop_variables('j',rank-1,1)}$) = &
+                            asum(lda,a(:, ${loop_variables('j',rank-1,1)}$),incx=1_ilp)
+${loop_variables_end(rank-1," "*20)}$                                    
+                case(NORM_TWO)                
+${loop_variables_start('j', 'a', rank-1, 1," "*20)}$                    
+                        nrm(${loop_variables('j',rank-1,1)}$) = &
+                            nrm2(lda,a(:, ${loop_variables('j',rank-1,1)}$),incx=1_ilp)                    
+${loop_variables_end(rank-1," "*20)}$                                    
+                case(NORM_INF)
+                    nrm = maxval( abs(a) , dim = dim )
+                case(NORM_MINUSINF)
+                    nrm = minval( abs(a) , dim = dim )
+                case (NORM_POW_FIRST:NORM_POW_LAST)
+                    rorder = 1.0_${rk}$ / norm_request
+                    nrm = sum( abs(a) ** norm_request , dim = dim ) ** rorder
+                case default
+                    err_ = linalg_state_type(this,LINALG_INTERNAL_ERROR,'invalid norm type after checking')
+                    call linalg_error_handling(err_,err)            
+            end select            
+            
+        endif        
         
     end subroutine norm_${rank}$D_to_${rank-1}$D_${ii}$_${ri}$
 
diff --git a/test/linalg/test_linalg_norm.fypp b/test/linalg/test_linalg_norm.fypp
@@ -159,17 +159,17 @@ module test_linalg_norm
         
         ! Test some norms
         call check(error,abs(norm(a,2) - norm2(a))<tol*norm(a,2),&
-                         'Euclidean norm does not match `norm2` intrinsic')
+                         'Euclidean norm does not match ${rt}$ `norm2` intrinsic')
         if (allocated(error)) return               
         
         ! Infinity norms
         call check(error,abs(norm(b,2)-norm2(b))<tol*norm(b,2),&
-                         'Dimmed Euclidean norm does not match `norm2` intrinsic')
+                         'Dimmed Euclidean norm does not match ${rt}$ `norm2` intrinsic')
         if (allocated(error)) return                       
         
         ! Test norm as collapsed around dimension
         do dim = 1, ndim
-            write(msg,"('Not all dim=',i0,' Euclidean norms match `norm2` intrinsic')") dim
+            write(msg,"('Not all dim=',i0,' Euclidean norms match ${rt}$ `norm2` intrinsic')") dim
             call check(error,all(abs(norm(b,2,dim)-norm2(b,dim))<tol*max(1.0_${rk}$,norm(b,2,dim))),&
                        trim(msg)) 
             if (allocated(error)) return             
