Reshuffle block sections of fields

ecmwf-ifs · Apr 2, 2024 · 6975ba2 · 6975ba2
1 parent 9f30946
commit 6975ba2
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/field_shuffle_type_module.fypp b/field_shuffle_type_module.fypp
@@ -46,16 +46,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 = 0
+IF (PRESENT (KBLKOFF)) IBLKOFF = KBLKOFF
 
 SELF%KGPTOT_S  = KGPTOT
 SELF%KLON_S    = KLON_S
@@ -79,7 +84,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
     JLONG = JLONG + 1
     IF (JLONG > SELF%KLON_G) THEN
       JLONG = 1

diff --git a/tests/reshuffle_lastdim.F90 b/tests/reshuffle_lastdim.F90
@@ -24,7 +24,8 @@ PROGRAM RESHUFFLE
 INTEGER, PARAMETER :: NPROMA2 =  6
 INTEGER, PARAMETER :: NFLEVG = 3
 
-INTEGER :: JLON, JBLK, JLEV, JPASS
+INTEGER :: JLON, JBLK, JLEV, JPASS, JBOFF
+INTEGER :: JBLK1L, JBLK1U
 
 INTEGER (KIND=JPIM), ALLOCATABLE :: D1 (:,:,:)
 
@@ -36,30 +37,43 @@ PROGRAM RESHUFFLE
 
 FUNC (JLON, JBLK) = 1000 * JBLK + JLON
 
+DO JBOFF = 1, 2
 ! First pass with modification on CPU, second pass on GPU
 DO JPASS = 1, 2
 
-  ALLOCATE (D1 (NPROMA1, 0:NFLEVG, JBLK1A:JBLK1B))
+  PRINT *, "=========> JPASS = ", JPASS, ", JBOFF = ", JBOFF, " <========="
+
+  IF (JBOFF == 1) THEN
+    JBLK1L = JBLK1A
+    JBLK1U = JBLK1B
+  ELSEIF (JBOFF == 2) THEN
+    JBLK1L = 1 
+    JBLK1U = 7 
+  ELSE
+    CALL FIELD_ABORT ('UNEXPECTED JBOFF')
+  ENDIF
+
+  ALLOCATE (D1 (NPROMA1, 0:NFLEVG, JBLK1L:JBLK1U))
 
   PRINT *, " NPROMA1 = ", NPROMA2
   PRINT *, " LBOUND (D1) = ", LBOUND (D1)
   PRINT *, " UBOUND (D1) = ", UBOUND (D1)
 
-  DO JBLK = JBLK1A, JBLK1B
+  DO JBLK = JBLK1L, JBLK1U
     DO JLON = 1, NPROMA1
       D1 (JLON, :, JBLK) = FUNC (JLON, JBLK)
     ENDDO
   ENDDO
 
-  DO JBLK = JBLK1A, JBLK1B
+  DO JBLK = JBLK1L, JBLK1U
     WRITE (*, '(20I12)') D1 (:, 1, JBLK)
   ENDDO
 
-  CALL FIELD_NEW (FD, DATA=D1, LBOUNDS=[1, 0, JBLK1A])
+  CALL FIELD_NEW (FD, DATA=D1, LBOUNDS=[1, 0, JBLK1L])
 
   ! Reshuffle on NPROMA2 arrays
 
-  CALL FGS%INIT (KGPTOT=NPROMA1*NGPBLKS1, KLON_S=NPROMA1, KLON_G=NPROMA2)
+  CALL FGS%INIT (KGPTOT=NPROMA1*NGPBLKS1, KLON_S=NPROMA1, KLON_G=NPROMA2, KBLKOFF=JBLK1A-JBLK1L)
 
   IF (JPASS == 1) THEN
     Z2 => GATHER_HOST_DATA_RDWR (FGS, FD)
@@ -108,7 +122,7 @@ PROGRAM RESHUFFLE
 
   PRINT *, '------------'
 
-  DO JBLK = JBLK1A, JBLK1B
+  DO JBLK = JBLK1L, JBLK1U
     WRITE (*, '(20I12)') D1 (:, 1, JBLK)
   ENDDO
 
@@ -128,6 +142,7 @@ PROGRAM RESHUFFLE
   DEALLOCATE (D1)
 
 
+ENDDO
 ENDDO
 
 END PROGRAM