Skip to content

Commit

Permalink
WIP - Adjoint LTINV - Adjoint the subroutine.
Browse files Browse the repository at this point in the history
  • Loading branch information
@l90lpa
l90lpa committed Jan 31, 2025
1 parent b2cef3c commit 3b9cfa5
Showing 1 changed file with 72 additions and 71 deletions.
143 changes: 72 additions & 71 deletions src/trans/gpu/internal/ltinvad_mod.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,21 +10,21 @@
! nor does it submit to any jurisdiction.
!

MODULE LTINV_MOD
MODULE LTINVAD_MOD
USE BUFFERED_ALLOCATOR_MOD, ONLY: ALLOCATION_RESERVATION_HANDLE

IMPLICIT NONE

PRIVATE
PUBLIC :: LTINV, LTINV_HANDLE, PREPARE_LTINV
PUBLIC :: LTINVAD, LTINVAD_HANDLE, PREPARE_LTINVAD

TYPE LTINV_HANDLE
TYPE LTINVAD_HANDLE
TYPE(ALLOCATION_RESERVATION_HANDLE) :: HPIA_AND_IN
TYPE(ALLOCATION_RESERVATION_HANDLE) :: HOUTS_AND_OUTA
! TYPE(ALLOCATION_RESERVATION_HANDLE) :: HOUTS_AND_OUTA
END TYPE

CONTAINS
FUNCTION PREPARE_LTINV(ALLOCATOR,KF_UV,KF_SCALARS,LVORGP,LDIVGP,LSCDERS) RESULT(HLTINV)
FUNCTION PREPARE_LTINVAD(ALLOCATOR,KF_UV,KF_SCALARS,LVORGP,LDIVGP,LSCDERS) RESULT(HLTINV)
USE PARKIND_ECTRANS, ONLY: JPIM, JPRBT, JPRD, JPIB
USE TPM_DISTR, ONLY: D
USE TPM_DIM, ONLY: R
Expand All @@ -38,7 +38,7 @@ FUNCTION PREPARE_LTINV(ALLOCATOR,KF_UV,KF_SCALARS,LVORGP,LDIVGP,LSCDERS) RESULT(
INTEGER(KIND=JPIM), INTENT(IN) :: KF_UV,KF_SCALARS
LOGICAL, INTENT(IN) :: LVORGP,LDIVGP,LSCDERS

TYPE(LTINV_HANDLE) :: HLTINV
TYPE(LTINVAD_HANDLE) :: HLTINVAD

INTEGER(KIND=JPIB) :: IALLOC_SZ, IPIA_SZ
INTEGER(KIND=JPIM) :: IOUT_STRIDES0
Expand Down Expand Up @@ -83,23 +83,23 @@ FUNCTION PREPARE_LTINV(ALLOCATOR,KF_UV,KF_SCALARS,LVORGP,LDIVGP,LSCDERS) RESULT(
! ZINP0
IALLOC_SZ = IALLOC_SZ + ALIGN(IIN0_SIZE*C_SIZEOF(ZPRD_DUMMY),128)

HLTINV%HPIA_AND_IN = RESERVE(ALLOCATOR, IALLOC_SZ, "HLTINV_HPIA_AND_IN")
HLTINVAD%HPIA_AND_IN = RESERVE(ALLOCATOR, IALLOC_SZ, "HLTINVAD_HPIA_AND_IN")

IALLOC_SZ = 0
! ZOUTA
IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT_SIZE*C_SIZEOF(ZPRBT_DUMMY),128)
! ZOUTS
IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT_SIZE*C_SIZEOF(ZPRBT_DUMMY),128)
! ZOUTA0
IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT0_SIZE*C_SIZEOF(ZPRD_DUMMY),128)
! ZOUTS0
IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT0_SIZE*C_SIZEOF(ZPRD_DUMMY),128)
! IALLOC_SZ = 0
! ! ZOUTA
! IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT_SIZE*C_SIZEOF(ZPRBT_DUMMY),128)
! ! ZOUTS
! IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT_SIZE*C_SIZEOF(ZPRBT_DUMMY),128)
! ! ZOUTA0
! IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT0_SIZE*C_SIZEOF(ZPRD_DUMMY),128)
! ! ZOUTS0
! IALLOC_SZ = IALLOC_SZ + ALIGN(IOUT0_SIZE*C_SIZEOF(ZPRD_DUMMY),128)

HLTINV%HOUTS_AND_OUTA = RESERVE(ALLOCATOR, IALLOC_SZ, "HLTINV_HOUTS_AND_OUTA")
! HLTINVAD%HOUTS_AND_OUTA = RESERVE(ALLOCATOR, IALLOC_SZ, "HLTINVAD_HOUTS_AND_OUTA")

END FUNCTION PREPARE_LTINV
END FUNCTION PREPARE_LTINVAD

SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
SUBROUTINE LTINVAD(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
& PSPVOR,PSPDIV,PSPSCALAR,PSPSC3A,PSPSC3B,PSPSC2, &
& ZOUTS,ZOUTA,ZOUTS0,ZOUTA0)

Expand All @@ -110,27 +110,28 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
USE TPM_GEOMETRY, ONLY: G
USE BUFFERED_ALLOCATOR_MOD, ONLY: BUFFERED_ALLOCATOR, ASSIGN_PTR, GET_ALLOCATION
USE TPM_DISTR, ONLY: D
USE PRFI1B_MOD, ONLY: PRFI1B
USE VDTUV_MOD, ONLY: VDTUV
USE SPNSDE_MOD, ONLY: SPNSDE
USE LEINV_MOD, ONLY: LEINV_STRIDES, LEINV
USE PRFI1BAD_MOD, ONLY: PRFI1BAD
USE VDTUVAD_MOD, ONLY: VDTUVAD
USE SPNSDEAD_MOD, ONLY: SPNSDEAD
USE LEINV_MOD, ONLY: LEINV_STRIDES
USE LEINVAD_MOD, ONLY: LEINVAD
USE ABORT_TRANS_MOD, ONLY: ABORT_TRANS
USE TPM_FIELDS_GPU, ONLY: FG
USE MPL_MODULE, ONLY: MPL_BARRIER,MPL_ALL_MS_COMM
USE TPM_GEN, ONLY: LSYNC_TRANS
USE TPM_STATS, ONLY: GSTATS => GSTATS_NVTX
USE ISO_C_BINDING, ONLY: C_SIZE_T, C_LOC, C_SIZEOF

!**** *LTINV* - Inverse Legendre transform
!**** *LTINVAD* - adjoint of inverse Legendre transform
!
! Purpose.
! --------
! Tranform from Laplace space to Fourier space, compute U and V
! and north/south derivatives of state variables.
! Adjoint of the "tranform from Laplace space to Fourier space, compute U and V
! and north/south derivatives of state variables".

!** Interface.
! ----------
! *CALL* *LTINV(...)
! *CALL* *LTINVAD(...)

! Explicit arguments :
! --------------------
Expand Down Expand Up @@ -175,14 +176,14 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
INTEGER(KIND=JPIM), INTENT(IN) :: KF_UV
INTEGER(KIND=JPIM), INTENT(IN) :: KF_SCALARS

REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPVOR(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPDIV(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPSCALAR(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPSC2(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPSC3A(:,:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(IN) :: PSPSC3B(:,:,:)
REAL(KIND=JPRBT), POINTER, INTENT(OUT) :: ZOUTS(:), ZOUTA(:)
REAL(KIND=JPRD), POINTER, INTENT(OUT) :: ZOUTS0(:), ZOUTA0(:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPVOR(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPDIV(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSCALAR(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC2(:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC3A(:,:,:)
REAL(KIND=JPRB) ,OPTIONAL,INTENT(INOUT) :: PSPSC3B(:,:,:)
REAL(KIND=JPRBT) , INTENT(IN) :: ZOUTS(:), ZOUTA(:)
REAL(KIND=JPRD) , INTENT(IN) :: ZOUTS0(:), ZOUTA0(:)

INTEGER(KIND=JPIM) :: IFIRST, J3

Expand All @@ -192,7 +193,7 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&

REAL(KIND=JPHOOK) :: ZHOOK_HANDLE
TYPE(BUFFERED_ALLOCATOR), INTENT(IN) :: ALLOCATOR
TYPE(LTINV_HANDLE), INTENT(IN) :: HLTINV
TYPE(LTINVAD_HANDLE), INTENT(IN) :: HLTINV

INTEGER(KIND=JPIM) :: IOUT_STRIDES0
INTEGER(KIND=JPIB) :: IOUT_SIZE
Expand All @@ -214,7 +215,7 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
!* 1. PERFORM LEGENDRE TRANFORM.
! --------------------------

IF (LHOOK) CALL DR_HOOK('LTINV_MOD',0,ZHOOK_HANDLE)
IF (LHOOK) CALL DR_HOOK('LTINVAD_MOD',0,ZHOOK_HANDLE)

! Get all pointers
IF_READIN = 0
Expand Down Expand Up @@ -256,31 +257,31 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
& IALLOC_POS, IALLOC_SZ)
IALLOC_POS = IALLOC_POS + IALLOC_SZ

IALLOC_POS = 1
! IALLOC_POS = 1

! ZOUTA
IALLOC_SZ = ALIGN(IOUT_SIZE*C_SIZEOF(ZOUTA(1)),128)
CALL ASSIGN_PTR(ZOUTA, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
& IALLOC_POS, IALLOC_SZ)
IALLOC_POS = IALLOC_POS + IALLOC_SZ
! ! ZOUTA
! IALLOC_SZ = ALIGN(IOUT_SIZE*C_SIZEOF(ZOUTA(1)),128)
! CALL ASSIGN_PTR(ZOUTA, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
! & IALLOC_POS, IALLOC_SZ)
! IALLOC_POS = IALLOC_POS + IALLOC_SZ

! ZOUTS
IALLOC_SZ = ALIGN(IOUT_SIZE*C_SIZEOF(ZOUTS(1)),128)
CALL ASSIGN_PTR(ZOUTS, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
& IALLOC_POS, IALLOC_SZ)
IALLOC_POS = IALLOC_POS + IALLOC_SZ
! ! ZOUTS
! IALLOC_SZ = ALIGN(IOUT_SIZE*C_SIZEOF(ZOUTS(1)),128)
! CALL ASSIGN_PTR(ZOUTS, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
! & IALLOC_POS, IALLOC_SZ)
! IALLOC_POS = IALLOC_POS + IALLOC_SZ

! ZOUTA0
IALLOC_SZ = ALIGN(IOUT0_SIZE*C_SIZEOF(ZOUTA0(1)),128)
CALL ASSIGN_PTR(ZOUTA0, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
& IALLOC_POS, IALLOC_SZ)
IALLOC_POS = IALLOC_POS + IALLOC_SZ
! ! ZOUTA0
! IALLOC_SZ = ALIGN(IOUT0_SIZE*C_SIZEOF(ZOUTA0(1)),128)
! CALL ASSIGN_PTR(ZOUTA0, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
! & IALLOC_POS, IALLOC_SZ)
! IALLOC_POS = IALLOC_POS + IALLOC_SZ

! ZOUTS0
IALLOC_SZ = ALIGN(IOUT0_SIZE*C_SIZEOF(ZOUTS0(1)),128)
CALL ASSIGN_PTR(ZOUTS0, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
& IALLOC_POS, IALLOC_SZ)
IALLOC_POS = IALLOC_POS + IALLOC_SZ
! ! ZOUTS0
! IALLOC_SZ = ALIGN(IOUT0_SIZE*C_SIZEOF(ZOUTS0(1)),128)
! CALL ASSIGN_PTR(ZOUTS0, GET_ALLOCATION(ALLOCATOR, HLTINV%HOUTS_AND_OUTA),&
! & IALLOC_POS, IALLOC_SZ)
! IALLOC_POS = IALLOC_POS + IALLOC_SZ

! Assign pointers do the different components of PIA
IFIRST = 0
Expand Down Expand Up @@ -345,31 +346,31 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
CALL GSTATS(422,1)

IF (KF_UV > 0) THEN
CALL PRFI1B(PVOR,PSPVOR,KF_UV,UBOUND(PSPVOR,2))
CALL PRFI1B(PDIV,PSPDIV,KF_UV,UBOUND(PSPDIV,2))
CALL PRFI1BAD(PVOR,PSPVOR,KF_UV,UBOUND(PSPVOR,2))
CALL PRFI1BAD(PDIV,PSPDIV,KF_UV,UBOUND(PSPDIV,2))

! Compute U and V for VOR and DIV
CALL VDTUV(KF_UV,ZEPSNM,PVOR,PDIV,PU,PV)
CALL VDTUVAD(KF_UV,ZEPSNM,PVOR,PDIV,PU,PV)
ENDIF

IF (KF_SCALARS > 0) THEN
IF(PRESENT(PSPSCALAR)) THEN
CALL PRFI1B(PSCALARS,PSPSCALAR,KF_SCALARS,UBOUND(PSPSCALAR,2))
CALL PRFI1BAD(PSCALARS,PSPSCALAR,KF_SCALARS,UBOUND(PSPSCALAR,2))
ELSE
IFIRST = 1
IF(PRESENT(PSPSC2) .AND. NF_SC2 > 0) THEN
CALL PRFI1B(PSCALARS(IFIRST:IFIRST+2*NF_SC2-1,:,:),PSPSC2(:,:),NF_SC2,UBOUND(PSPSC2,2))
CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC2-1,:,:),PSPSC2(:,:),NF_SC2,UBOUND(PSPSC2,2))
IFIRST = IFIRST+2*NF_SC2
ENDIF
IF(PRESENT(PSPSC3A) .AND. NF_SC3A > 0) THEN
DO J3=1,UBOUND(PSPSC3A,3)
CALL PRFI1B(PSCALARS(IFIRST:IFIRST+2*NF_SC3A-1,:,:),PSPSC3A(:,:,J3),NF_SC3A,UBOUND(PSPSC3A,2))
CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC3A-1,:,:),PSPSC3A(:,:,J3),NF_SC3A,UBOUND(PSPSC3A,2))
IFIRST = IFIRST+2*NF_SC3A
ENDDO
ENDIF
IF(PRESENT(PSPSC3B) .AND. NF_SC3B > 0) THEN
DO J3=1,UBOUND(PSPSC3B,3)
CALL PRFI1B(PSCALARS(IFIRST:IFIRST+2*NF_SC3B-1,:,:),PSPSC3B(:,:,J3),NF_SC3B,UBOUND(PSPSC3B,2))
CALL PRFI1BAD(PSCALARS(IFIRST:IFIRST+2*NF_SC3B-1,:,:),PSPSC3B(:,:,J3),NF_SC3B,UBOUND(PSPSC3B,2))
IFIRST = IFIRST+2*NF_SC3B
ENDDO
ENDIF
Expand All @@ -382,7 +383,7 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&

! Compute NS derivatives if needed
IF (LSCDERS) THEN
CALL SPNSDE(KF_SCALARS,ZEPSNM,PSCALARS,PSCALARS_NSDER)
CALL SPNSDEAD(KF_SCALARS,ZEPSNM,PSCALARS,PSCALARS_NSDER)
ENDIF

#ifdef OMPGPU
Expand All @@ -404,16 +405,16 @@ SUBROUTINE LTINV(ALLOCATOR,HLTINV,KF_UV,KF_SCALARS,&
! ------------------------------------------------------------------


!* 4. INVERSE LEGENDRE TRANSFORM.
!* 4. Adjoint of INVERSE LEGENDRE TRANSFORM.
! ---------------------------

! Legendre transforms. When converting PIA into ZOUT, we ignore the first entries of LEINV.
! This is because vorticity and divergence is not necessarily converted to GP space.
CALL LEINV(ALLOCATOR,PIA(2*(IF_READIN-IF_LEG)+1:IF_READIN,:,:),ZINP,ZINP0,ZOUTS,ZOUTA,ZOUTS0,ZOUTA0,IF_LEG)
CALL LEINVAD(ALLOCATOR,PIA(2*(IF_READIN-IF_LEG)+1:IF_READIN,:,:),ZINP,ZINP0,ZOUTS,ZOUTA,ZOUTS0,ZOUTA0,IF_LEG)

IF (LHOOK) CALL DR_HOOK('LTINV_MOD',1,ZHOOK_HANDLE)
IF (LHOOK) CALL DR_HOOK('LTINVAD_MOD',1,ZHOOK_HANDLE)
END ASSOCIATE
! ------------------------------------------------------------------
END SUBROUTINE LTINV
END MODULE LTINV_MOD
END SUBROUTINE LTINVAD
END MODULE LTINVAD_MOD

0 comments on commit 3b9cfa5

Please sign in to comment.