Skip to content

Commit

Permalink
Merge pull request #110 from esmf-org/multitile_undist_redist
Browse files Browse the repository at this point in the history 
Allow I/O of multi-tile arrays and fields with undistributed dimensions

This PR extends the IO routine undist_arraycreate_alldist to work with multi-tile arrays. This allows I/O of multi-tile arrays and fields with undistributed dimensions.

The implementation follows suggestions from @theurich and @oehmke .

I have tested this using billsacks/esmfprojects-multi_tile_io@eea516b. This writes a 6-tile field with 2 ungridded dimensions (dimensions 1 & 3 are gridded, 2 & 4 ungridded). I have manually inspected the output to verify that this field is being written correctly.
  • Loading branch information
@billsacks
billsacks authored Dec 16, 2022
2 parents 6bf5b37 + 115d35e commit 3b11fab
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 34 deletions.
64 changes: 37 additions & 27 deletions src/Infrastructure/IO/src/ESMCI_IO.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1325,46 +1325,40 @@ void IO::undist_arraycreate_alldist(Array *src_array_p, Array **dest_array_p, in

int dimCount= dg->getDimCount();
int deCount = dg->getDELayout()->getDeCount();

int tilecount = dg->getTileCount ();
if (tilecount != 1) {
localrc = ESMF_RC_NOT_IMPL;
std::stringstream errmsg;
errmsg << "tile count of " << tilecount << " != 1 - not supported yet";
if (ESMC_LogDefault.MsgFoundError(localrc, errmsg.str(), ESMC_CONTEXT,
rc)) return;
}
int tileCount = dg->getTileCount ();

const int *arrayToDistGridMap = src_array_p->getArrayToDistGridMap();
const int *undistLBound = src_array_p->getUndistLBound();
const int *undistUBound = src_array_p->getUndistUBound();
const int *minIndexPTile = dg->getMinIndexPDimPTile();
const int *maxIndexPTile = dg->getMaxIndexPDimPTile();
const int *minIndexPDimPTile = dg->getMinIndexPDimPTile();
const int *maxIndexPDimPTile = dg->getMaxIndexPDimPTile();
const int *minIndexPDimPDe = dg->getMinIndexPDimPDe();
const int *maxIndexPDimPDe = dg->getMaxIndexPDimPDe();

// construct the new minIndex and maxIndex and regDecomp taking into account
// how Array dimensions are mapped against DistGrid dimensions
std::vector<int> minIndexNew(rank);
std::vector<int> maxIndexNew(rank);
std::vector<int> minIndexNew(rank*tileCount);
std::vector<int> maxIndexNew(rank*tileCount);
std::vector<int> deBlockList(rank*2*deCount);
std::vector<int> deBlockListLen(3);
deBlockListLen[0]=rank; deBlockListLen[1]=2; deBlockListLen[2]=deCount;
int jj = 0;
for (int i=0; i<rank; i++) {
int j = arrayToDistGridMap[i];
if (j > 0) {
// valid DistGrid dimension
minIndexNew[i] = minIndexPTile[j-1];
maxIndexNew[i] = maxIndexPTile[j-1];
for (int tile=0; tile<tileCount; tile++) {
minIndexNew[tile*rank + i] = minIndexPDimPTile[tile*dimCount + (j-1)];
maxIndexNew[tile*rank + i] = maxIndexPDimPTile[tile*dimCount + (j-1)];
}
for (int k=0; k<deCount; k++){
deBlockList[k*2*rank + i] = minIndexPDimPDe[k*dimCount + (j-1)];
deBlockList[k*2*rank + rank + i] = maxIndexPDimPDe[k*dimCount + (j-1)];
}
} else {
// undistributed dimension
minIndexNew[i] = undistLBound[jj];
maxIndexNew[i] = undistUBound[jj];
for (int tile=0; tile<tileCount; tile++) {
minIndexNew[tile*rank + i] = undistLBound[jj];
maxIndexNew[tile*rank + i] = undistUBound[jj];
}
for (int k=0; k<deCount; k++){
deBlockList[k*2*rank + i] = undistLBound[jj];
deBlockList[k*2*rank + rank + i] = undistUBound[jj];
Expand All @@ -1374,15 +1368,31 @@ void IO::undist_arraycreate_alldist(Array *src_array_p, Array **dest_array_p, in
}

// create the fixed up DistGrid, making sure to use original DELayout
ESMCI::InterArray<int> minIndexInterface(minIndexNew);
ESMCI::InterArray<int> maxIndexInterface(maxIndexNew);
ESMCI::InterArray<int> deBlockListInterface(&deBlockList[0], 3, &deBlockListLen[0]);
int deBlockListLen[] = {rank, 2, deCount};
ESMCI::InterArray<int> deBlockListInterface(&deBlockList[0], 3, deBlockListLen);
DELayout *delayout = dg->getDELayout();
DistGrid *dg_temp = DistGrid::create(&minIndexInterface,
&maxIndexInterface, &deBlockListInterface,
NULL, NULL, NULL, delayout, NULL, &localrc);
if (ESMC_LogDefault.MsgFoundError(localrc, ESMCI_ERR_PASSTHRU, ESMC_CONTEXT, rc)) {
return;
DistGrid *dg_temp;
if (tileCount == 1) {
ESMCI::InterArray<int> minIndexInterface(minIndexNew);
ESMCI::InterArray<int> maxIndexInterface(maxIndexNew);
dg_temp = DistGrid::create(&minIndexInterface,
&maxIndexInterface, &deBlockListInterface,
NULL, NULL, NULL, delayout, NULL, &localrc);
if (ESMC_LogDefault.MsgFoundError(localrc, ESMCI_ERR_PASSTHRU, ESMC_CONTEXT, rc)) {
return;
}
} else {
int dummyLen[] = {rank, tileCount};
ESMCI::InterArray<int> minIndexInterface(&minIndexNew[0], 2, dummyLen);
ESMCI::InterArray<int> maxIndexInterface(&maxIndexNew[0], 2, dummyLen);
ESMCI::InterArray<int> deToTileMapInterface((int*)(dg->getTileListPDe()), 1, &deCount);
dg_temp = DistGrid::create(&minIndexInterface,
&maxIndexInterface, &deBlockListInterface,
&deToTileMapInterface,
NULL, NULL, NULL, delayout, NULL, &localrc);
if (ESMC_LogDefault.MsgFoundError(localrc, ESMCI_ERR_PASSTHRU, ESMC_CONTEXT, rc)) {
return;
}
}

// finally, create the fixed up Array using pointer to original data.
Expand Down
64 changes: 57 additions & 7 deletions src/Infrastructure/IO/tests/ESMF_IO_MultitileUTest.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,8 @@ program ESMF_IO_MultitileUTest
! Fields used for writing:
!
! The following fields make up the field bundle:
type(ESMF_Field) :: field1, field2, field1Copy
real(ESMF_KIND_R8), pointer :: field1Data(:,:), field2Data(:,:), field1CopyData(:,:)
type(ESMF_Field) :: field1, field2, field1Copy, field4d
real(ESMF_KIND_R8), pointer :: field1Data(:,:), field2Data(:,:), field1CopyData(:,:), field4dData(:,:,:,:)
type(ESMF_FieldBundle) :: fieldBundle
! This field is not in the field bundle:
type(ESMF_Field) :: field3
Expand All @@ -62,8 +62,8 @@ program ESMF_IO_MultitileUTest
! Fields used for reading:
!
! The following fields make up the field bundle:
type(ESMF_Field) :: field1Read, field2Read, field1CopyRead
real(ESMF_KIND_R8), pointer :: field1ReadData(:,:), field2ReadData(:,:), field1CopyReadData(:,:)
type(ESMF_Field) :: field1Read, field2Read, field1CopyRead, field4dRead
real(ESMF_KIND_R8), pointer :: field1ReadData(:,:), field2ReadData(:,:), field1CopyReadData(:,:), field4dReadData(:,:,:,:)
type(ESMF_FieldBundle) :: fieldBundleRead
! This field is not in the field bundle:
type(ESMF_Field) :: field3Read
Expand Down Expand Up @@ -145,7 +145,7 @@ program ESMF_IO_MultitileUTest

!------------------------------------------------------------------------
!NEX_UTest_Multi_Proc_Only
write(name, *) "Confirm that FieldBundle-read fields match originals"
write(name, *) "Confirm that simple FieldBundle-read fields match originals"
write(failMsg, *) "Some read-in fields differ from originals"
allEqual = ( &
all(field1ReadData == field1Data) .and. &
Expand All @@ -159,6 +159,19 @@ program ESMF_IO_MultitileUTest
#endif
!------------------------------------------------------------------------

!------------------------------------------------------------------------
!NEX_UTest_Multi_Proc_Only
write(name, *) "Confirm that FieldBundle-read field with ungridded dim matches original"
write(failMsg, *) "Read-in field differs from original"
allEqual = all(field4dReadData == field4dData)
#if (defined ESMF_PIO && (defined ESMF_NETCDF || defined ESMF_PNETCDF))
call ESMF_Test(allEqual, name, failMsg, result, ESMF_SRCLINE)
#else
write(failMsg, *) "Comparison did not fail as expected"
call ESMF_Test(.not. allEqual, name, failMsg, result, ESMF_SRCLINE)
#endif
!------------------------------------------------------------------------

#ifdef ESMF_TESTEXHAUSTIVE
! The following tests don't add much code coverage, so are only done when
! ESMF_TESTEXHAUSTIVE is set
Expand Down Expand Up @@ -341,7 +354,11 @@ subroutine createFields(rc)

integer :: decompPTile(2,6)
type(ESMF_ArraySpec) :: arraySpec
type(ESMF_ArraySpec) :: arraySpec_w_ungridded
type(ESMF_Array) :: array1
real(ESMF_KIND_R8), pointer :: coordPtrX(:,:), coordPtrY(:,:)
integer :: u1, u2, i, j
real :: multiplier

!------------------------------------------------------------------------
! Set up 6-tile grid
Expand All @@ -364,6 +381,8 @@ subroutine createFields(rc)

call ESMF_ArraySpecSet(arraySpec, typekind=ESMF_TYPEKIND_R8, rank=2, rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_ArraySpecSet(arraySpec_w_ungridded, typekind=ESMF_TYPEKIND_R8, rank=4, rc=rc)
if (rc /= ESMF_SUCCESS) return

field1 = ESMF_FieldCreate(grid6tile, arraySpec, name="field1", rc=rc)
if (rc /= ESMF_SUCCESS) return
Expand Down Expand Up @@ -401,6 +420,37 @@ subroutine createFields(rc)
call ESMF_FieldGet(field3Read, farrayPtr=field3ReadData, rc=rc)
if (rc /= ESMF_SUCCESS) return

field4d = ESMF_FieldCreate(grid6tile, arraySpec_w_ungridded, name="field4d", &
ungriddedLBound=[2,15], ungriddedUBound=[4,18], &
! 2nd and 4th dimensions are ungridded dimensions
gridToFieldMap=[1,3], &
rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_FieldGet(field4d, farrayPtr=field4dData, rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_GridGetCoord(grid6tile, coordDim=1, farrayPtr=coordPtrX, rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_GridGetCoord(grid6tile, coordDim=2, farrayPtr=coordPtrY, rc=rc)
do u1 = 2,4
do u2 = 15,18
do i = lbound(field4dData, 1), ubound(field4dData, 1)
do j = lbound(field4dData, 3), ubound(field4dData, 3)
multiplier = 5.**(u2-15)
field4dData(i,u1,j,u2) = u1*multiplier*(coordPtrX(i,j) - coordPtrY(i,j))
end do
end do
end do
end do

field4dRead = ESMF_FieldCreate(grid6tile, arraySpec_w_ungridded, name="field4d", &
ungriddedLBound=[2,15], ungriddedUBound=[4,18], &
! 2nd and 4th dimensions are ungridded dimensions
gridToFieldMap=[1,3], &
rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_FieldGet(field4dRead, farrayPtr=field4dReadData, rc=rc)
if (rc /= ESMF_SUCCESS) return

! Create a copy of field1 that uses the same array, so we can test writing
! the same array twice from a single call.
call ESMF_FieldGet(field1, array=array1, rc=rc)
Expand All @@ -417,12 +467,12 @@ subroutine createFields(rc)

fieldBundle = ESMF_FieldBundleCreate(name="fb", rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_FieldBundleAdd(fieldBundle, [field1, field2, field1Copy], rc=rc)
call ESMF_FieldBundleAdd(fieldBundle, [field1, field2, field1Copy, field4d], rc=rc)
if (rc /= ESMF_SUCCESS) return

fieldBundleRead = ESMF_FieldBundleCreate(name="fbRead", rc=rc)
if (rc /= ESMF_SUCCESS) return
call ESMF_FieldBundleAdd(fieldBundleRead, [field1Read, field2Read, field1CopyRead], rc=rc)
call ESMF_FieldBundleAdd(fieldBundleRead, [field1Read, field2Read, field1CopyRead, field4dRead], rc=rc)
if (rc /= ESMF_SUCCESS) return

end subroutine createFields
Expand Down

0 comments on commit 3b11fab

Please sign in to comment.