diff --git a/components/omega/src/infra/IOStream.cpp b/components/omega/src/infra/IOStream.cpp index b1bf48edef5e..ebbf0c2d9105 100644 --- a/components/omega/src/infra/IOStream.cpp +++ b/components/omega/src/infra/IOStream.cpp @@ -727,24 +727,28 @@ int IOStream::computeDecomp( } // Get dimension and size information for each dimension - I4 GlobalSize = 1; - LocalSize = 1; - std::vector DimLengthGlobal(NDims); - std::vector DimIDs(NDims); - std::vector DimOffsets(NDims); + // For the offset calculation, we also create dimension and dimension offsets + // in index order up to the max of 5 dimensions. The active indices are + // at the end in array index order. + I4 GlobalSize = 1; + LocalSize = 1; constexpr I4 MaxDims = 5; - std::vector OffsetLoopLim(MaxDims, 1); // loop limits for later offset + std::vector DimLengthsGlob(NDims, 1); + std::vector DimOffsets(MaxDims); + std::vector DimLengthGlob(MaxDims, 1); // lengths padded to MaxDims + std::vector DimLengthLoc(MaxDims, 1); // lengths padded to MaxDims for (int IDim = 0; IDim < NDims; ++IDim) { + I4 StartDim = MaxDims - NDims; std::string DimName = DimNames[IDim]; - DimIDs[IDim] = AllDimIDs[DimName]; std::shared_ptr ThisDim = Dimension::get(DimName); DimLengths[IDim] = ThisDim->getLengthLocal(); - DimLengthGlobal[IDim] = ThisDim->getLengthGlobal(); - DimOffsets[IDim] = ThisDim->getOffset(); + DimLengthsGlob[IDim] = ThisDim->getLengthGlobal(); + DimOffsets[StartDim + IDim] = ThisDim->getOffset(); + DimLengthLoc[StartDim + IDim] = DimLengths[IDim]; + DimLengthGlob[StartDim + IDim] = DimLengthsGlob[IDim]; LocalSize *= DimLengths[IDim]; - GlobalSize *= DimLengthGlobal[IDim]; - OffsetLoopLim[IDim] = DimLengths[IDim]; + GlobalSize *= DimLengthsGlob[IDim]; } // Create the data decomposition based on dimension information @@ -754,45 +758,61 @@ int IOStream::computeDecomp( // The linear offset along each dimension has already been computed // by the dimension class. - std::vector Offset(LocalSize); + std::vector Offset(LocalSize, -1); // Compute strides in linear space for each dimension - std::vector Strides(MaxDims, 1); - for (int IDim = 1; IDim < NDims; ++IDim) { - Strides[IDim] = Strides[IDim - 1] * DimLengthGlobal[IDim - 1]; + // For the proper stride in storage order, we start from the rightmost + // index in index order. + std::vector StrideGlob(MaxDims, 1); + std::vector StrideLoc(MaxDims, 1); + for (int IDim = MaxDims - 2; IDim >= 0; --IDim) { + StrideGlob[IDim] = StrideGlob[IDim + 1] * DimLengthGlob[IDim + 1]; + StrideLoc[IDim] = StrideLoc[IDim + 1] * DimLengthLoc[IDim + 1]; } // Compute full array offsets based on each dimensions linear offset - I4 Add = 0; // linear address of offset vector - for (int N = 0; N < OffsetLoopLim[4]; ++N) { + // Skip -1 entries that are outside owned domain + for (int N = 0; N < DimLengthLoc[0]; ++N) { I4 NGlob = 0; if (NDims >= 5) - NGlob = (DimOffsets[4])(N); - for (int M = 0; M < OffsetLoopLim[3]; ++M) { + NGlob = (DimOffsets[0])(N); + if (NGlob < 0) + continue; // index outside owned domain + for (int M = 0; M < DimLengthLoc[1]; ++M) { I4 MGlob = 0; if (NDims >= 4) - MGlob = (DimOffsets[3])(M); - for (int K = 0; K < OffsetLoopLim[2]; ++K) { + MGlob = (DimOffsets[1])(M); + if (MGlob < 0) + continue; + for (int K = 0; K < DimLengthLoc[2]; ++K) { I4 KGlob = 0; if (NDims >= 3) KGlob = (DimOffsets[2])(K); - for (int J = 0; J < OffsetLoopLim[1]; ++J) { + if (KGlob < 0) + continue; + for (int J = 0; J < DimLengthLoc[3]; ++J) { I4 JGlob = 0; if (NDims >= 2) - JGlob = (DimOffsets[1])(J); - for (int I = 0; I < OffsetLoopLim[0]; ++I) { - I4 IGlob = (DimOffsets[0])(I); - Offset[Add] = IGlob * Strides[0] + JGlob * Strides[1] + - KGlob * Strides[2] + MGlob * Strides[3] + - NGlob * Strides[4]; - ++Add; + JGlob = (DimOffsets[3])(J); + if (JGlob < 0) + continue; + for (int I = 0; I < DimLengthLoc[4]; ++I) { + I4 IGlob = (DimOffsets[4])(I); + if (IGlob < 0) + continue; + int Add = I * StrideLoc[4] + J * StrideLoc[3] + + K * StrideLoc[2] + M * StrideLoc[1] + + N * StrideLoc[0]; + Offset[Add] = IGlob * StrideGlob[4] + JGlob * StrideGlob[3] + + KGlob * StrideGlob[2] + MGlob * StrideGlob[1] + + NGlob * StrideGlob[0]; } } } } } - Err = OMEGA::IO::createDecomp(DecompID, MyIOType, NDims, DimLengthGlobal, + Err = OMEGA::IO::createDecomp(DecompID, MyIOType, NDims, DimLengthsGlob, LocalSize, Offset, OMEGA::IO::DefaultRearr); if (Err != 0) { LOG_ERROR("Error creating decomp for field {} in stream {}", FieldName, diff --git a/components/omega/test/infra/IOStreamTest.cpp b/components/omega/test/infra/IOStreamTest.cpp index d887ae74694d..3864c3ecefec 100644 --- a/components/omega/test/infra/IOStreamTest.cpp +++ b/components/omega/test/infra/IOStreamTest.cpp @@ -253,6 +253,9 @@ int main(int argc, char **argv) { } // Force read the latest restart and check the results + // We use this log message to act as a barrier/delay to make sure the + // restart write has finished before we read. + LOG_INFO("Test reading final restart"); bool ForceRead = true; Err1 = IOStream::read("RestartRead", *ModelClock, ReqMetadata, ForceRead); TestEval("Restart force read", Err1, ErrRef, Err);