Merge pull request #3562 from eisenhauer/LaterShapeCollection

Collect Shape at EndStep (and restore DataWrite test)
ornladios · Mar 28, 2023 · d86ff45 · d86ff45
2 parents 6d62443 + 376e2eb
commit d86ff45
Show file tree

Hide file tree

Showing 6 changed files with 350 additions and 14 deletions.
diff --git a/source/adios2/toolkit/format/bp5/BP5Serializer.cpp b/source/adios2/toolkit/format/bp5/BP5Serializer.cpp
@@ -448,6 +448,7 @@ BP5Serializer::CreateWriterRec(void *Variable, const char *Name, DataType Type,
     if (Type == DataType::String)
         ElemSize = sizeof(char *);
     Rec->Key = Variable;
+    Rec->Shape = VB->m_ShapeID;
     Rec->FieldID = Info.RecCount;
     Rec->DimCount = DimCount;
     Rec->Type = (int)Type;
@@ -801,6 +802,7 @@ void BP5Serializer::Marshal(void *Variable, const char *Name,
         if (!AlreadyWritten)
         {
             if (Shape)
+                // this will be overwritten in EndStep with the final shape
                 MetaEntry->Shape = CopyDims(DimCount, Shape);
             else
                 MetaEntry->Shape = NULL;
@@ -849,15 +851,10 @@ void BP5Serializer::Marshal(void *Variable, const char *Name,
         {
             /* already got some metadata, add blocks */
             size_t PreviousDBCount = MetaEntry->DBCount;
-            //  Assume shape is still valid   (modify this if shape /global
-            //  dimensions can change )
-            // Also assume Dims is always right and consistent, otherwise,
-            // bad things
-            if (Shape && MetaEntry->Shape)
-            {
-                // Shape can change with later writes, so must overwrite
-                memcpy(MetaEntry->Shape, Shape, DimCount * sizeof(Shape[0]));
-            }
+            //
+            // we don't snag shape again here, because we're going to grab it at
+            // EndStep
+            //
             MetaEntry->DBCount += DimCount;
             MetaEntry->BlockCount++;
             MetaEntry->Count =
@@ -1135,9 +1132,34 @@ BufferV *BP5Serializer::ReinitStepData(BufferV *DataBuffer,
     return tmp;
 }
 
+void BP5Serializer::CollectFinalShapeValues()
+{
+    for (int i = 0; i < Info.RecCount; i++)
+    {
+        BP5WriterRec Rec = &Info.RecList[i];
+        if (Rec->Shape == ShapeID::GlobalArray)
+        {
+            core::VariableBase *VB =
+                static_cast<core::VariableBase *>(Rec->Key);
+            struct BP5MetadataInfoStruct *MBase =
+                (struct BP5MetadataInfoStruct *)MetadataBuf;
+            int AlreadyWritten = BP5BitfieldTest(MBase, Rec->FieldID);
+            if (!AlreadyWritten)
+                continue;
+
+            MetaArrayRec *MetaEntry =
+                (MetaArrayRec *)((char *)(MetadataBuf) + Rec->MetaOffset);
+
+            memcpy(MetaEntry->Shape, VB->Shape().data(),
+                   Rec->DimCount * sizeof(size_t));
+        }
+    }
+}
+
 BP5Serializer::TimestepInfo BP5Serializer::CloseTimestep(int timestep,
                                                          bool forceCopyDeferred)
 {
+    // EndStep()
     std::vector<MetaMetaInfoBlock> Formats;
     if (!Info.MetaFormat && Info.MetaFieldCount)
     {
@@ -1245,6 +1267,8 @@ BP5Serializer::TimestepInfo BP5Serializer::CloseTimestep(int timestep,
 
     ProcessDeferredMinMax();
 
+    CollectFinalShapeValues();
+
     void *MetaDataBlock = FFSencode(MetaEncodeBuffer, Info.MetaFormat,
                                     MetadataBuf, &MetaDataSize);
     BufferFFS *Metadata =

diff --git a/source/adios2/toolkit/format/bp5/BP5Serializer.h b/source/adios2/toolkit/format/bp5/BP5Serializer.h
@@ -219,6 +219,7 @@ class BP5Serializer : virtual public BP5Base
                                  DataType Type, size_t ElemSize,
                                  size_t DimCount);
     void ValidateWriterRec(BP5WriterRec Rec, void *Variable);
+    void CollectFinalShapeValues();
     void RecalcMarshalStorageSize();
     void RecalcAttributeStorageSize();
     void AddSimpleField(FMFieldList *FieldP, int *CountP, const char *Name,

diff --git a/testing/adios2/engine/staging-common/CMakeLists.txt b/testing/adios2/engine/staging-common/CMakeLists.txt
@@ -15,7 +15,7 @@ if(ADIOS2_HAVE_SST)
   gtest_add_tests_helper(Threads MPI_NONE "" Engine.Staging. ".FileStream" EXTRA_ARGS "FileStream")
 endif()
 
-set (PROGS TestCommonWrite
+set (PROGS TestCommonWrite TestShapeChangingWrite
     TestCommonWriteModes
     TestCommonWriteAttrs
     TestCommonWriteLocal
@@ -117,7 +117,7 @@ if(ADIOS2_HAVE_MPI AND MPIEXEC_EXECUTABLE)
 endif()
 
 set (SIMPLE_TESTS "1x1;1x1.Attrs;1x1.ModAttrs;1x1DefSync;1x1VarDestruction;1x1.SpanMinMax;1x1.Local2;1x1Struct;1x1Joined")
-set (BP5File_ONLY_TESTS "1x1DataWrite")
+set (BP5File_ONLY_TESTS "1x1DataWrite;1x1.ShapeChange")
 
 set (SIMPLE_FORTRAN_TESTS "")
 if(ADIOS2_HAVE_Fortran)
@@ -234,7 +234,7 @@ endif()
 
 #  BP5 tests
 if(ADIOS2_HAVE_BP5)
-    set (BP5_TESTS ${ALL_SIMPLE_TESTS} ${SPECIAL_TESTS})
+    set (BP5_TESTS ${ALL_SIMPLE_TESTS} ${SPECIAL_TESTS} ${BP5File_ONLY_TESTS})
     # Delayed reader not worth testing on file engines
     list (FILTER BP5_TESTS EXCLUDE REGEX "DelayedReader")
     # Discard not a feature of BP5

diff --git a/testing/adios2/engine/staging-common/TestCommonRead.cpp b/testing/adios2/engine/staging-common/TestCommonRead.cpp
@@ -239,7 +239,8 @@ TEST_F(CommonReadTest, ADIOS2CommonRead1D8)
                 engine.BlocksInfo(var_r64, engine.CurrentStep());
             EXPECT_EQ(i8Info.size(), writerSize);
             EXPECT_EQ(i16Info.size(), writerSize);
-            EXPECT_EQ(i32Info.size(), writerSize);
+            EXPECT_TRUE((i32Info.size() == writerSize) ||
+                        (i32Info.size() == writerSize * 3));
             EXPECT_EQ(i64Info.size(), writerSize);
             EXPECT_EQ(r32Info.size(), writerSize);
             EXPECT_EQ(r64Info.size(), writerSize);