Adapt to partial block ranges

ecmwf-ifs · Apr 2, 2024 · 21b7d70 · 21b7d70
1 parent 90f5591
commit 21b7d70
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 39 deletions.
diff --git a/field_RANKSUFF_shuffle_module.fypp b/field_RANKSUFF_shuffle_module.fypp
@@ -132,8 +132,8 @@ CONTAINS
 
 SUBROUTINE PAIR_${ft.name}$_${what}$_${mode}$ (KNDS, PTRG, PTRS)
 
-${ft.type}$, INTENT (INOUT) :: 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
 
@@ -212,8 +212,8 @@ 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 (INOUT) :: 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
 
@@ -224,11 +224,11 @@ INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG
 #:elif what == 'HOST'
 !$OMP PARALLEL DO PRIVATE (JBLKG, JLONG, JBLKS, JLONS)
 #:endif
-DO JBLKG = LBOUND (PTRG, ${ft.rank}$), UBOUND (PTRG, ${ft.rank}$)
+DO JBLKG = 1, SIZE (KNDS, 3)
 #:if what == 'DEVICE'
 !$acc loop vector private (JLONG, JBLKS, JLONS) 
 #:endif
-  DO JLONG = 1, SIZE (PTRG, 1)
+  DO JLONG = 1, SIZE (KNDS, 2)
     JLONS = KNDS (NLONDIM, JLONG, JBLKG)
     JBLKS = KNDS (NBLKDIM, JLONG, JBLKG)
     IF (JLONS > 0) THEN

diff --git a/field_gathscat_type_module.fypp b/field_gathscat_type_module.fypp
@@ -78,7 +78,9 @@ SELF%KGPTOT  =  SELF%KGPTOT_G
 SELF%KLON    =  SELF%KLON_G   
 SELF%KGPBLKS =  SELF%KGPBLKS_G
 
-SELF%LFULL = SELF%KGPTOT_G == SELF%KGPTOT_S
+! 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
@@ -101,7 +103,7 @@ ELSE
           CALL FIELD_ABORT ('INIT_FIELD_SHUFFLE: OUT OF BOUNDS')
         ENDIF
         INDS (NLONDIM, JLONG, JBLKG) = JLONS
-        INDS (NBLKDIM, JLONG, JBLKG) = JBLKS
+        INDS (NBLKDIM, JLONG, JBLKG) = JBLKS+JBOFF
         JLONG = JLONG + 1
         IF (JLONG > SELF%KLON_G) THEN
           JLONG = 1

diff --git a/field_shuffle_type_module.fypp b/field_shuffle_type_module.fypp
@@ -59,7 +59,7 @@ INTEGER (KIND=JPIM), POINTER :: INDS (:,:,:)
 INTEGER (KIND=JPIM) :: JLONS, JBLKS, JLONG, JBLKG, I1S, I2S
 INTEGER (KIND=JPIM) :: IBLKOFF
 
-IBLKOFF = 0
+IBLKOFF = 1
 IF (PRESENT (KBLKOFF)) IBLKOFF = KBLKOFF
 
 SELF%KGPTOT_S  = KGPTOT
@@ -84,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 + IBLKOFF
+    INDS (NBLKDIM, JLONG, JBLKG) = JBLKS + IBLKOFF-1
     JLONG = JLONG + 1
     IF (JLONG > SELF%KLON_G) THEN
       JLONG = 1

diff --git a/tests/gather_scatter_lastdim.F90 b/tests/gather_scatter_lastdim.F90
@@ -26,16 +26,18 @@ PROGRAM GATHER_SCATTER
 
 LOGICAL :: TRIG (NPROMA,JBLKA:JBLKB)
 
-INTEGER :: I, J, JPASS
+INTEGER :: I, J, JPASS, JBOFF
+INTEGER :: JBLKL, JBLKU, JBLK
 
 INTEGER (KIND=JPIM), ALLOCATABLE :: D (:,:,:)
 CLASS (FIELD_3IM), POINTER :: FD => NULL()
 
 INTEGER (KIND=JPIM), POINTER :: FILTERED_D(:,:,:) => NULL()
 
+DO JBOFF = 1, 2
 DO JPASS = 1, 3
 
-  WRITE (0, '("===========> ",A4," <===========")') CLTYPE (JPASS)
+  WRITE (0, '("===========> PASS = ",A4,", JBOFF ",I0," <===========")') CLTYPE (JPASS), JBOFF
 
   !CREATE A FILTER TO USE WITH THE GATHSCAT
 
@@ -53,15 +55,25 @@ PROGRAM GATHER_SCATTER
     CASE (3)
       TRIG = .TRUE.
   END SELECT
+
+  IF (JBOFF == 1) THEN
+    JBLKL = JBLKA
+    JBLKU = JBLKB
+  ELSEIF (JBOFF == 2) THEN
+    JBLKL = 1
+    JBLKU = 7
+  ELSE
+    CALL FIELD_ABORT ('UNEXPECTED JBOFF')
+  ENDIF
 
   CALL FIELD_NEW (FTRIG, DATA=TRIG, LBOUNDS=[1,JBLKA])
 
   !CREATE THE FIELD TO BE FILTERED BY GATHSCAT
-  ALLOCATE (D (NPROMA, 0:NFLEVG, JBLKA:JBLKB))
+  ALLOCATE (D (NPROMA, 0:NFLEVG, JBLKL:JBLKU))
 
   D = 1
 
-  CALL FIELD_NEW (FD, DATA=D, LBOUNDS=[1,0,JBLKA])
+  CALL FIELD_NEW (FD, DATA=D, LBOUNDS=[1,0,JBLKL])
 
   WRITE (0, *) " LBOUND (D) = ", LBOUND (D)
   WRITE (0, *) " UBOUND (D) = ", UBOUND (D)
@@ -71,7 +83,7 @@ PROGRAM GATHER_SCATTER
   CALL FGS%INIT (FTRIG, NPROMA*NGPBLKS)
 
   FILTERED_D => GATHER_HOST_DATA_RDWR (FGS, FD)
-  FILTERED_D = 2 !NOT ALL OF D WILL BE MODIFIED, ONLY THE FILTERED DATA
+  FILTERED_D (:,:,:) = 2 !NOT ALL OF D WILL BE MODIFIED, ONLY THE FILTERED DATA
 
   WRITE (0, *) " LBOUND (FILTERED_D) = ", LBOUND (FILTERED_D)
   WRITE (0, *) " UBOUND (FILTERED_D) = ", UBOUND (FILTERED_D)
@@ -87,33 +99,50 @@ PROGRAM GATHER_SCATTER
 
   CALL FGS%SCATTER ()
 
-  SELECT CASE (JPASS)
-    CASE (1)
-      DO I = 1, NPROMA
-        IF (MOD (I,2) == 0)THEN
-          IF (.NOT. ALL(D(I,:,:)==2)) THEN
+  DO JBLK = JBLKL, JBLKU
+    SELECT CASE (JPASS)
+      CASE (1)
+        IF ((JBLKA <= JBLK) .AND. (JBLK <= JBLKB)) THEN
+          DO I = 1, NPROMA
+            IF (MOD (I, 2) == 0)THEN
+              IF (ANY (D (I,:,JBLK) /= 2)) THEN
+                CALL FIELD_ABORT ("ERROR")
+              ENDIF
+            ELSE
+              IF (ANY (D (I,:,JBLK) /= 1)) THEN
+                CALL FIELD_ABORT ("ERROR")
+              ENDIF
+            ENDIF
+          ENDDO
+        ELSE
+          IF (ANY (D (:,:,JBLK) /= 1)) THEN
+            CALL FIELD_ABORT ("ERROR")
+          ENDIF
+        ENDIF
+      CASE (2)
+! Nothing has been touched
+        IF (ANY (D (:,:,JBLK) /= 1)) THEN
+          CALL FIELD_ABORT ("ERROR")
+        ENDIF
+      CASE (3)
+! All blocks between JBLKA and JBLKB have changed
+        IF ((JBLKA <= JBLK) .AND. (JBLK <= JBLKB)) THEN
+          IF (ANY (D (:,:,JBLK) /= 2)) THEN
             CALL FIELD_ABORT ("ERROR")
           ENDIF
         ELSE
-          IF (.NOT. ALL(D(I,:,:)==1)) THEN
+          IF (ANY (D (:,:,JBLK) /= 1)) THEN
             CALL FIELD_ABORT ("ERROR")
           ENDIF
         ENDIF
-      ENDDO
-    CASE (2)
-      IF (.NOT. ALL (D == 1)) THEN
-        CALL FIELD_ABORT ("ERROR")
-      ENDIF
-    CASE (3)
-      IF (.NOT. ALL (D == 2)) THEN
-        CALL FIELD_ABORT ("ERROR")
-      ENDIF
-  END SELECT
+    END SELECT
+  ENDDO
 
   DEALLOCATE (D)
 
   CALL FIELD_DELETE (FTRIG)
 
 ENDDO
+ENDDO
 
 END PROGRAM
diff --git a/tests/reshuffle_lastdim.F90 b/tests/reshuffle_lastdim.F90
@@ -55,7 +55,7 @@ PROGRAM RESHUFFLE
 
   ALLOCATE (D1 (NPROMA1, 0:NFLEVG, JBLK1L:JBLK1U))
 
-  PRINT *, " NPROMA1 = ", NPROMA2
+  PRINT *, " NPROMA1 = ", NPROMA1
   PRINT *, " LBOUND (D1) = ", LBOUND (D1)
   PRINT *, " UBOUND (D1) = ", UBOUND (D1)
 
@@ -73,7 +73,7 @@ PROGRAM RESHUFFLE
 
   ! Reshuffle on NPROMA2 arrays
 
-  CALL FGS%INIT (KGPTOT=NPROMA1*NGPBLKS1, KLON_S=NPROMA1, KLON_G=NPROMA2, KBLKOFF=JBLK1A-JBLK1L)
+  CALL FGS%INIT (KGPTOT=NPROMA1*NGPBLKS1, KLON_S=NPROMA1, KLON_G=NPROMA2, KBLKOFF=JBLK1A)
 
   IF (JPASS == 1) THEN
     Z2 => GATHER_HOST_DATA_RDWR (FGS, FD)
@@ -126,14 +126,23 @@ PROGRAM RESHUFFLE
     WRITE (*, '(20I12)') D1 (:, 1, JBLK)
   ENDDO
 
-  DO JBLK = JBLK1A, JBLK1B
+  DO JBLK = JBLK1L, JBLK1U
     DO JLEV = 0, NFLEVG
       DO JLON = 1, NPROMA1
-        IF (D1 (JLON, JLEV, JBLK) /= (JPASS + 1) * FUNC (JLON, JBLK)) THEN
-          PRINT *, " JPASS = ", JPASS, " JLON = ", JLON, " JLEV = ", JLEV, &
-        & " JBLK = ", JBLK, " D1 = ", D1 (JLON, JLEV, JBLK), &
-        & (JPASS + 1) * FUNC (JLON, JBLK)
-          CALL FIELD_ABORT ('VALUE ERROR')
+        IF ((JBLK1A .LE. JBLK) .AND. (JBLK .LE. JBLK1B)) THEN
+          IF (D1 (JLON, JLEV, JBLK) /= (JPASS + 1) * FUNC (JLON, JBLK)) THEN
+            PRINT *, " JPASS = ", JPASS, " JLON = ", JLON, " JLEV = ", JLEV, &
+          & " JBLK = ", JBLK, " D1 = ", D1 (JLON, JLEV, JBLK), &
+          & (JPASS + 1) * FUNC (JLON, JBLK)
+            CALL FIELD_ABORT ('VALUE ERROR')
+          ENDIF
+        ELSE
+          IF (D1 (JLON, JLEV, JBLK) /= FUNC (JLON, JBLK)) THEN
+            PRINT *, " JPASS = ", JPASS, " JLON = ", JLON, " JLEV = ", JLEV, &
+          & " JBLK = ", JBLK, " D1 = ", D1 (JLON, JLEV, JBLK), &
+          & FUNC (JLON, JBLK)
+            CALL FIELD_ABORT ('VALUE ERROR')
+          ENDIF
         ENDIF
       ENDDO
     ENDDO