ecmwf-ifs · awnawab · Apr 12, 2024 · Mar 28, 2024 · Mar 28, 2024 · Apr 2, 2024
diff --git a/field_RANKSUFF_gather_module.fypp b/field_RANKSUFF_gather_module.fypp
@@ -53,7 +53,8 @@ CLASS (${ft.name}$),   POINTER, INTENT (IN) :: YLF
 
 ${ft.type}$, POINTER :: PTR(${ft.shape}$)
 
-PTR => PAIR_GATHER_${what}$_DATA_${mode}$ (SELF%${ft.name}$_LIST, SELF%LNULL, SELF%LFULL, SELF%KLON_G, SELF%KGPBLKS_G, SELF%YLFINDS, YLF)
+PTR => PAIR_GATHER_${what}$_DATA_${mode}$ (SELF%${ft.name}$_LIST, SELF%LNULL, SELF%LFULL, &
+& SELF%KLON_G, SELF%KGPBLKS_G, SELF%YLFINDS, YLF, SELF%KBLKMIN, SELF%KBLKMAX)
 
 END FUNCTION ${ft.name}$_GATHER_${what}$_DATA_${mode}$
 

diff --git a/field_RANKSUFF_module.fypp b/field_RANKSUFF_module.fypp
@@ -156,6 +156,7 @@ CONTAINS
     LOGICAL, INTENT(IN), OPTIONAL :: SYNC_ON_FINAL
 
     LOGICAL :: LLPERSISTENT
+
 #include "abor1.intfb.h"
 
     LLPERSISTENT = .TRUE.
@@ -166,7 +167,9 @@ CONTAINS
     ELSE
       SELF%PTR => DATA
     ENDIF
+
     SELF%THREAD_BUFFER = .NOT. LLPERSISTENT
+
     CALL SELF%SET_STATUS (NHSTFRESH)
 
     SELF%MAP_DEVPTR = INIT_MAP_DEVPTR
@@ -207,13 +210,18 @@ CONTAINS
     LOGICAL, OPTIONAL,  INTENT(IN) :: MAP_DEVPTR
     ${ft.type}$, OPTIONAL, INTENT(IN) :: INIT_VALUE
     LOGICAL :: IS_DELAYED
+    LOGICAL :: LLPERSISTENT
+
 #include "abor1.intfb.h"
 
     IS_DELAYED = .FALSE.
     IF(PRESENT(DELAYED))THEN
       IS_DELAYED = DELAYED
     ENDIF
 
+    LLPERSISTENT = .FALSE.
+    IF (PRESENT (PERSISTENT)) LLPERSISTENT = PERSISTENT
+
 #:if defined('CUDA')
     ! In the presence of CUDA we pin by default
     SELF%PINNED = INIT_PINNED_VALUE
@@ -237,19 +245,15 @@ CONTAINS
     ELSE
       SELF%LBOUNDS=1
     ENDIF
+
     SELF%UBOUNDS=UBOUNDS
-    SELF%UBOUNDS(${ft.rank}$) = OML_MAX_THREADS ()
 
-    ! By default we allocate thread-local temporaries
-    SELF%THREAD_BUFFER = .TRUE.
+    SELF%THREAD_BUFFER = .NOT. LLPERSISTENT
 
-    IF (PRESENT(PERSISTENT)) THEN
-      IF (PERSISTENT) THEN
-        SELF%THREAD_BUFFER = .FALSE.
-        SELF%LBOUNDS(${ft.rank}$) = 1
-        SELF%UBOUNDS(${ft.rank}$) = UBOUNDS(${ft.rank}$)
-      END IF
-    END IF
+    IF (.NOT. LLPERSISTENT) THEN
+      SELF%LBOUNDS(${ft.rank}$) = 1
+      SELF%UBOUNDS(${ft.rank}$) = OML_MAX_THREADS ()
+    ENDIF
 
     CALL SELF%SET_STATUS (UNALLOCATED)
     IF (PRESENT(INIT_VALUE)) THEN

diff --git a/field_RANKSUFF_shuffle_module.fypp b/field_RANKSUFF_shuffle_module.fypp
@@ -66,12 +66,14 @@ CONTAINS
 #:for what in ['DEVICE', 'HOST']
 #:for mode in ['RDONLY', 'RDWR']
 
-FUNCTION PAIR_${ft.name}$_GATHER_${what}$_DATA_${mode}$ (YDPAIR, LDNULL, LDFULL, KLON, KGPBLKS, YDFINDS, YLF) RESULT (PTR)
+FUNCTION PAIR_${ft.name}$_GATHER_${what}$_DATA_${mode}$ (YDPAIR, LDNULL, LDFULL, KLON, KGPBLKS, YDFINDS, YLF, KBLKMIN, KBLKMAX) RESULT (PTR)
 TYPE (${ft.name}$_SHUFFLE_PAIR), POINTER              :: YDPAIR
 LOGICAL,                                   INTENT (IN) :: LDNULL, LDFULL
 INTEGER (KIND=JPIM),                       INTENT (IN) :: KLON, KGPBLKS
 CLASS (FIELD_3IM),                POINTER              :: YDFINDS 
 CLASS (${ft.name}$),              POINTER, INTENT (IN) :: YLF
+INTEGER (KIND=JPIM),                       INTENT (IN) :: KBLKMIN
+INTEGER (KIND=JPIM),                       INTENT (IN) :: KBLKMAX
 
 ${ft.type}$, POINTER :: PTR(${ft.shape}$), ZTRG(${ft.shape}$), ZTRS(${ft.shape}$)
 ${ft.type}$, POINTER :: PTR1(${ft.shape}$)
@@ -91,6 +93,14 @@ ELSEIF (LDFULL) THEN
   ! Return pointer on packed array
   PTR => GET_${what}$_DATA_${mode}$ (YLF)
 
+  ! Set last lbound to 1
+
+  ILBOUNDS = LBOUND (PTR)
+
+  PTR1 => PTR (${','.join ([':'] * (ft.rank-1))}$, KBLKMIN:KBLKMAX)
+
+  PTR (${','.join (list (map (lambda i: "ILBOUNDS(" + str (i) + "):", range (1, ft.rank))))}$, 1:) => PTR1
+
 ELSE
 
   PTR => GET_${what}$_DATA_RDONLY (YLF)
@@ -104,8 +114,13 @@ ELSE
 
   ILBOUNDS = LBOUND (PTR)
   IUBOUNDS = UBOUND (PTR)
+
+  ILBOUNDS (1) = 1
   IUBOUNDS (1) = KLON
+
+  ILBOUNDS (${ft.rank}$) = 1
   IUBOUNDS (${ft.rank}$) = KGPBLKS
+
   CALL FIELD_NEW (YLPAIR%YLGATH, LBOUNDS=ILBOUNDS, UBOUNDS=IUBOUNDS, PERSISTENT=.TRUE.)
 
   INDS => GET_${what}$_DATA_RDONLY (YDFINDS)
@@ -124,27 +139,37 @@ CONTAINS
 
 SUBROUTINE PAIR_${ft.name}$_${what}$_${mode}$ (KNDS, PTRG, PTRS)
 
-${ft.type}$, INTENT (OUT) :: PTRG(${ft.shape}$)
-${ft.type}$, INTENT (IN)  :: PTRS(${ft.shape}$)
+${ft.type}$, POINTER :: PTRG(${ft.shape}$)
+${ft.type}$, POINTER :: PTRS(${ft.shape}$)
 INTEGER (KIND=JPIM), INTENT (IN)  :: KNDS (:,:,:)
 INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG
 
-#:set dc = ', ' + (':, ' * (ft.rank-2))
+#:set ind = ', '.join (map (lambda i: "J" + str (i), range (2, ft.rank)))
+#:if ft.rank > 2
+INTEGER (KIND=JPIM) :: ${ind}$
+#:set ind = ind + ', '
+#:endif
 
 #:if what == 'DEVICE'
 !$acc parallel loop gang present (PTRG, PTRS, KNDS)
 #:elif what == 'HOST'
-!$OMP PARALLEL DO PRIVATE (JBLKG, JLONG, JBLKS, JLONS)
+!$OMP PARALLEL DO PRIVATE (${ind}$JBLKG, JLONG, JBLKS, JLONS)
 #:endif
 DO JBLKG = 1, SIZE (KNDS, 3)
 #:if what == 'DEVICE'
-!$acc loop vector private (JLONG, JBLKS, JLONS) 
+!$acc loop vector private (${ind}$JLONG, JBLKS, JLONS) 
 #:endif
   DO JLONG = 1, SIZE (KNDS, 2)
     JLONS = KNDS (NLONDIM, JLONG, JBLKG)
     JBLKS = KNDS (NBLKDIM, JLONG, JBLKG)
     IF (JLONS > 0) THEN
-      PTRG (JLONG${dc}$JBLKG) = PTRS (JLONS${dc}$JBLKS)
+#:for i in reversed (range (2, ft.rank))
+      DO J${i}$ = LBOUND (PTRG, ${i}$), UBOUND (PTRG, ${i}$)
+#:endfor
+        PTRG (JLONG, ${ind}$JBLKG) = PTRS (JLONS, ${ind}$JBLKS)
+#:for i in (range (2, ft.rank))
+      ENDDO
+#:endfor
     ENDIF
   ENDDO
 ENDDO
@@ -161,7 +186,7 @@ END FUNCTION PAIR_${ft.name}$_GATHER_${what}$_DATA_${mode}$
 
 SUBROUTINE PAIR_${ft.name}$_SCATTER_DATA (YDPAIR, YDFINDS)
 TYPE (${ft.name}$_SHUFFLE_PAIR), POINTER              :: YDPAIR
-CLASS (FIELD_3IM),                POINTER              :: YDFINDS 
+CLASS (FIELD_3IM),               POINTER              :: YDFINDS 
 
 TYPE (${ft.name}$_SHUFFLE_PAIR), POINTER :: YLPAIR
 ${ft.type}$, POINTER :: PTRG(${ft.shape}$), PTRS(${ft.shape}$)
@@ -178,13 +203,15 @@ DO WHILE (ASSOCIATED (YDPAIR))
       INDS => GET_DEVICE_DATA_RDONLY (YDFINDS)
       PTRG => GET_DEVICE_DATA_RDONLY (YDPAIR%YLGATH)
       PTRS => GET_DEVICE_DATA_RDWR   (YDPAIR%YLSCAT)
+
       CALL ${ft.name}$_SCATTER_DEVICE_KERNEL (INDS, PTRG, PTRS)
 
     ELSEIF (YDPAIR%IWHAT == NHOST) THEN
 
       INDS => GET_HOST_DATA_RDONLY (YDFINDS)
       PTRG => GET_HOST_DATA_RDONLY (YDPAIR%YLGATH)
       PTRS => GET_HOST_DATA_RDWR   (YDPAIR%YLSCAT)
+
       CALL ${ft.name}$_SCATTER_HOST_KERNEL (INDS, PTRG, PTRS)
 
     ENDIF
@@ -202,27 +229,37 @@ CONTAINS
 #:for what in ['DEVICE', 'HOST']
 SUBROUTINE ${ft.name}$_SCATTER_${what}$_KERNEL (KNDS, PTRG, PTRS)
 
-${ft.type}$, INTENT (IN)  :: PTRG(${ft.shape}$)
-${ft.type}$, INTENT (OUT) :: PTRS(${ft.shape}$)
+${ft.type}$, POINTER :: PTRG(${ft.shape}$)
+${ft.type}$, POINTER :: PTRS(${ft.shape}$)
 INTEGER (KIND=JPIM), INTENT (IN) :: KNDS (:,:,:)
 INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG
 
-#:set dc = ', ' + (':, ' * (ft.rank-2))
+#:set ind = ', '.join (map (lambda i: "J" + str (i), range (2, ft.rank)))
+#:if ft.rank > 2
+INTEGER (KIND=JPIM) :: ${ind}$
+#:set ind = ind + ', '
+#:endif
 
 #:if what == 'DEVICE'
 !$acc parallel loop gang present (PTRG, PTRS, KNDS)
 #:elif what == 'HOST'
-!$OMP PARALLEL DO PRIVATE (JBLKG, JLONG, JBLKS, JLONS)
+!$OMP PARALLEL DO PRIVATE (${ind}$JBLKG, JLONG, JBLKS, JLONS)
 #:endif
 DO JBLKG = 1, SIZE (KNDS, 3)
 #:if what == 'DEVICE'
-!$acc loop vector private (JLONG, JBLKS, JLONS) 
+!$acc loop vector private (${ind}$JLONG, JBLKS, JLONS) 
 #:endif
   DO JLONG = 1, SIZE (KNDS, 2)
     JLONS = KNDS (NLONDIM, JLONG, JBLKG)
     JBLKS = KNDS (NBLKDIM, JLONG, JBLKG)
     IF (JLONS > 0) THEN
-      PTRS (JLONS${dc}$JBLKS) = PTRG (JLONG${dc}$JBLKG) 
+#:for i in reversed (range (2, ft.rank))
+      DO J${i}$ = LBOUND (PTRG, ${i}$), UBOUND (PTRG, ${i}$)
+#:endfor
+        PTRS (JLONS, ${ind}$JBLKS) = PTRG (JLONG, ${ind}$JBLKG) 
+#:for i in (range (2, ft.rank))
+      ENDDO
+#:endfor
     ENDIF
   ENDDO
 ENDDO

diff --git a/field_gathscat_type_module.fypp b/field_gathscat_type_module.fypp
@@ -33,37 +33,41 @@ INTEGER (KIND=JPIM), PARAMETER :: NLONDIM = 1, NBLKDIM = 2
 
 CONTAINS
 
-SUBROUTINE INIT_FIELD_GATHSCAT (SELF, YDFCOND, KGPTOT, KLON_S, KLON_G)
+SUBROUTINE INIT_FIELD_GATHSCAT (SELF, YDFCOND, KGPTOT, KLON_S, KLON_G, KBLKMIN, KBLKMAX)
 
 USE FIELD_ABORT_MODULE
 
 CLASS (FIELD_GATHSCAT) :: SELF
 CLASS (FIELD_2LM), POINTER :: YDFCOND
 INTEGER (KIND=JPIM), INTENT (IN) :: KGPTOT
 INTEGER (KIND=JPIM), INTENT (IN), OPTIONAL :: KLON_S, KLON_G
+INTEGER (KIND=JPIM), INTENT (IN), OPTIONAL :: KBLKMIN, KBLKMAX
 
 LOGICAL, POINTER :: LLF (:,:)
 INTEGER (KIND=JPIM), POINTER :: INDS (:,:,:)
 INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG, I1S, I2S
 
-
 LLF => GET_HOST_DATA_RDONLY (YDFCOND)
 
+SELF%KBLKMIN = LBOUND (LLF, 2)
+IF (PRESENT (KBLKMIN)) SELF%KBLKMIN = KBLKMIN
+SELF%KBLKMAX = UBOUND (LLF, 2)
+IF (PRESENT (KBLKMAX)) SELF%KBLKMAX = KBLKMAX
+
 SELF%KLON_S    = SIZE (LLF, 1)
-SELF%KGPBLKS_S = SIZE (LLF, 2)
+SELF%KGPBLKS_S = SELF%KBLKMAX - SELF%KBLKMIN + 1
 SELF%KGPTOT_S  = KGPTOT
 
 ! Reduction
 
 SELF%KGPTOT_G = 0
 
-DO JBLKS = 1, SELF%KGPBLKS_S
+DO JBLKS = SELF%KBLKMIN, SELF%KBLKMAX
   I1S = 1
-  I2S = MIN (SELF%KLON_S, SELF%KGPTOT_S - (JBLKS - 1) * SELF%KLON_S)
+  I2S = MIN (SELF%KLON_S, SELF%KGPTOT_S - (JBLKS - SELF%KBLKMIN) * SELF%KLON_S)
   SELF%KGPTOT_G = SELF%KGPTOT_G + COUNT (LLF (I1S:I2S,JBLKS))
 ENDDO
 
-
 IF (PRESENT (KLON_G)) THEN
   SELF%KLON_G = KLON_G
 ELSE
@@ -76,7 +80,9 @@ SELF%KGPTOT  =  SELF%KGPTOT_G
 SELF%KLON    =  SELF%KLON_G   
 SELF%KGPBLKS =  SELF%KGPBLKS_G
 
+! Does not work anymore with partial block ranges
 SELF%LFULL = SELF%KGPTOT_G == SELF%KGPTOT_S
+!SELF%LFULL = .FALSE.
 SELF%LNULL = SELF%KGPTOT_G == 0
 
 IF (SELF%LNULL) THEN
@@ -92,8 +98,8 @@ ELSE
 
   JBLKG = 1
   JLONG = 1
-  DO JBLKS = 1, SELF%KGPBLKS_S
-    DO JLONS = 1, MIN (SELF%KLON_S, SELF%KGPTOT_S - (JBLKS - 1) * SELF%KLON_S)
+  DO JBLKS = SELF%KBLKMIN, SELF%KBLKMAX
+    DO JLONS = 1, MIN (SELF%KLON_S, SELF%KGPTOT_S - (JBLKS - SELF%KBLKMIN) * SELF%KLON_S)
       IF (LLF (JLONS, JBLKS)) THEN
         IF ((JLONG > SIZE (INDS, 2)) .OR. (JBLKG > SIZE (INDS, 3))) THEN
           CALL FIELD_ABORT ('INIT_FIELD_SHUFFLE: OUT OF BOUNDS')

diff --git a/field_shuffle_type_module.fypp b/field_shuffle_type_module.fypp
@@ -28,6 +28,7 @@ TYPE FIELD_SHUFFLE
   INTEGER (KIND=JPIM) :: KGPBLKS_G = -1, KLON_G = -1, KGPTOT_G = -1
   INTEGER (KIND=JPIM) :: KGPBLKS_S = -1, KLON_S = -1, KGPTOT_S = -1
   INTEGER (KIND=JPIM) :: KGPBLKS   = -1, KLON   = -1, KGPTOT   = -1 ! alias for KGPBLKS, KLON, KGPTOT
+  INTEGER (KIND=JPIM) :: KBLKMIN, KBLKMAX
   LOGICAL :: LFULL = .FALSE.  ! No need to gather/scatter, all columns are OK, return pointers based on original fields
   LOGICAL :: LNULL = .FALSE.  ! No need to gather/scatter, all columns are KO, return pointers on empty arrays
   CLASS (FIELD_3IM), POINTER :: YLFINDS => NULL ()
@@ -46,16 +47,21 @@ INTEGER (KIND=JPIM), PARAMETER :: NLONDIM = 1, NBLKDIM = 2
 
 CONTAINS
 
-SUBROUTINE INIT_FIELD_SHUFFLE (SELF, KGPTOT, KLON_S, KLON_G)
+SUBROUTINE INIT_FIELD_SHUFFLE (SELF, KGPTOT, KLON_S, KLON_G, KBLKOFF)
 
 USE FIELD_ABORT_MODULE
 
 CLASS (FIELD_SHUFFLE) :: SELF
 INTEGER (KIND=JPIM), INTENT (IN) :: KGPTOT
 INTEGER (KIND=JPIM), INTENT (IN) :: KLON_S, KLON_G
+INTEGER (KIND=JPIM), INTENT (IN), OPTIONAL :: KBLKOFF
 
 INTEGER (KIND=JPIM), POINTER :: INDS (:,:,:)
 INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG, I1S, I2S
+INTEGER (KIND=JPIM) :: IBLKOFF
+
+IBLKOFF = 1
+IF (PRESENT (KBLKOFF)) IBLKOFF = KBLKOFF
 
 SELF%KGPTOT_S  = KGPTOT
 SELF%KLON_S    = KLON_S
@@ -69,6 +75,9 @@ SELF%KGPTOT  =  SELF%KGPTOT_G
 SELF%KLON    =  SELF%KLON_G   
 SELF%KGPBLKS =  SELF%KGPBLKS_G
 
+SELF%KBLKMIN = IBLKOFF
+SELF%KBLKMAX = SELF%KGPBLKS+IBLKOFF-1
+
 CALL FIELD_NEW (SELF%YLFINDS, UBOUNDS=[2, SELF%KLON_G, SELF%KGPBLKS_G], PERSISTENT=.TRUE.)
 INDS => GET_HOST_DATA_RDWR (SELF%YLFINDS)
 
@@ -79,7 +88,7 @@ JLONG = 1
 DO JBLKS = 1, SELF%KGPBLKS_S
   DO JLONS = 1, MIN (SELF%KLON_S, SELF%KGPTOT_S - (JBLKS - 1) * SELF%KLON_S)
     INDS (NLONDIM, JLONG, JBLKG) = JLONS
-    INDS (NBLKDIM, JLONG, JBLKG) = JBLKS
+    INDS (NBLKDIM, JLONG, JBLKG) = JBLKS + IBLKOFF-1
     JLONG = JLONG + 1
     IF (JLONG > SELF%KLON_G) THEN
       JLONG = 1

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
@@ -25,6 +25,9 @@ target_compile_definitions( main.x PRIVATE $<${HAVE_CUDA}:_CUDA> )
 
 ## Unit tests
 list(APPEND TEST_FILES
+        gather_scatter_lastdim.F90
+        reshuffle_lastdim.F90
+        test_lastdim.F90
         test_statistics.F90
         test_sizeof.F90
         test_bc.F90