Merge pull request #6258 from hzhou/2211_yaksa_pack

datatype: fix builtin_element_size at MPIR_Typerep_unflatten Approved-by: Ken Raffenetti
pmodels · Nov 10, 2022 · a6d3064 · a6d3064
2 parents c06659e + b0cef14
commit a6d3064
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/src/mpi/datatype/typerep/src/typerep_flatten.c b/src/mpi/datatype/typerep/src/typerep_flatten.c
@@ -118,6 +118,11 @@ int MPIR_Typerep_unflatten(MPIR_Datatype * datatype_ptr, void *flattened_type)
     datatype_ptr->true_lb = flatten_hdr->true_lb;
     datatype_ptr->contents = NULL;
     datatype_ptr->flattened = NULL;
+    if (datatype_ptr->basic_type != MPI_DATATYPE_NULL) {
+        datatype_ptr->builtin_element_size = MPIR_Datatype_get_basic_size(datatype_ptr->basic_type);
+    } else {
+        datatype_ptr->builtin_element_size = 0;
+    }
 
 #if (MPICH_DATATYPE_ENGINE == MPICH_DATATYPE_ENGINE_YAKSA)
     int rc = yaksa_unflatten(&datatype_ptr->typerep.handle, flattened_typerep);

diff --git a/src/mpi/datatype/typerep/src/typerep_yaksa_pack.c b/src/mpi/datatype/typerep/src/typerep_yaksa_pack.c
@@ -143,10 +143,7 @@ static int typerep_do_pack(const void *inbuf, MPI_Aint incount, MPI_Datatype dat
         MPIR_Datatype *dtp;
         MPIR_Datatype_get_ptr(datatype, dtp);
         is_contig = dtp->is_contig;
-        /* NOTE: dtp->element_size may not be set if dtp is from a flattened type */
-        if (dtp->basic_type != MPI_DATATYPE_NULL) {
-            element_size = MPIR_Datatype_get_basic_size(datatype);
-        }
+        element_size = dtp->builtin_element_size;
         inbuf_ptr = MPIR_get_contig_ptr(inbuf, dtp->true_lb);
         total_size = incount * dtp->size;
     }