PML/UCX: reset converter stack on unordered messages

Signed-off-by: Sergey Oblomov <sergeyo@mellanox.com>

ompi/mca/pml/ucx/pml_ucx_datatype.c

@@ -40,6 +40,7 @@ static void* pml_ucx_generic_datatype_start_unpack(void *context, void *buffer,
 
     OMPI_DATATYPE_RETAIN(datatype);
     convertor->datatype = datatype;
+    convertor->offset = 0;
     opal_convertor_copy_and_prepare_for_recv(ompi_proc_local_proc->super.proc_convertor,
                                              &datatype->super, count, buffer, 0,
                                              &convertor->opal_conv);
@@ -85,8 +86,15 @@ static ucs_status_t pml_ucx_generic_datatype_unpack(void *state, size_t offset,
     iov.iov_base = (void*)src;
     iov.iov_len  = length;
 
+    /* if we detected out-of-order message - apply hack: reset datatype stack
+     * of convertor to allow re-build it on set_position call */
+    if ((offset != convertor->offset) &&
+        !(convertor->opal_conv.flags & OPAL_DATATYPE_FLAG_CONTIGUOUS) ) {
+        opal_convertor_reset_stack_at_beginning(&convertor->opal_conv);
+    }
     opal_convertor_set_position(&convertor->opal_conv, &offset);
     opal_convertor_unpack(&convertor->opal_conv, &iov, &iov_count, &length);
+    convertor->offset = offset + length;
     return UCS_OK;
 }
 

ompi/mca/pml/ucx/pml_ucx_datatype.h

@@ -17,6 +17,7 @@ struct pml_ucx_convertor {
     opal_free_list_item_t     super;
     ompi_datatype_t           *datatype;
     opal_convertor_t          opal_conv;
+    size_t                    offset;
 };
 
 

opal/datatype/opal_convertor.c

@@ -409,6 +409,12 @@ int opal_convertor_create_stack_at_begining( opal_convertor_t* convertor,
 }
 
 
+int32_t opal_convertor_reset_stack_at_beginning( opal_convertor_t* convertor )
+{
+    return opal_convertor_create_stack_at_begining( convertor, opal_datatype_local_sizes );
+}
+
+
 int32_t opal_convertor_set_position_nocheck( opal_convertor_t* convertor,
                                              size_t* position )
 {

opal/datatype/opal_convertor.h

@@ -339,6 +339,11 @@ opal_convertor_set_position( opal_convertor_t* convertor,
     return opal_convertor_set_position_nocheck( convertor, position );
 }
 
+/*
+ *
+ */
+OPAL_DECLSPEC int32_t
+opal_convertor_reset_stack_at_beginning( opal_convertor_t* convertor );
 /*
  *
  */