cleanup datatype tmps for ialltoallw and friends

leverage nbc infrastructure

Signed-off-by: Howard Pritchard <[email protected]>

Author: hppritcha

ompi/mpi/bindings/ompi_bindings/c.py

@@ -881,6 +881,10 @@ def generate_function(prototype, fn_name, internal_fn, out, enable_count=False):
         return_type = prototype.return_type.construct(abi_type='standard')
         lines.append(f'{return_type.tmp_type_text()} ret_value;')
         for param in params:
+            if param.need_async_cleanup == True:
+                lines.append('int idx  = 0;')
+                break
+        for param in params:
 #           print("param = " + str(param) + " " + str(param.argument))
             if param.init_code:
                 lines.extend(param.init_code)

ompi/mpi/bindings/ompi_bindings/c_type.py

@@ -90,16 +90,24 @@ def callback_wrapper_code(self):
         """Return True if this parameter has callback wrapper code to generate."""
         return False
 
+    @property
+    def need_async_cleanup(self):
+        """Return True if this parameter generates async memory cleanup code."""
+        return False
+
 class StandardABIType(Type):
 
     @property
     def tmpname(self):
-        return f'{self.name}_tmp'
+        return util.abi_tmp_name(self.name)
 
     @property
     def argument(self):
         return self.tmpname
 
+    @staticmethod
+    def async_callback_index(self):
+        return "idx"
 
 @Type.add_type('ERROR_CLASS', abi_type=['ompi'])
 class TypeErrorClass(Type):
@@ -660,13 +668,23 @@ def argument(self):
 @Type.add_type('DATATYPE_ARRAY_ASYNC', abi_type=['standard'])
 class TypeDatatypeArrayAsyncStandard(TypeDatatypeArrayStandard):
 
+    @property
+    def need_async_cleanup(self):
+        return True
+
     @property
     def final_code(self):
-        code = ['{']
+        request_tmp_name = util.abi_tmp_name('request')
+        code = []
+        code.append('if((MPI_SUCCESS == ret_value) && (MPI_REQUEST_NULL != request_tmp)){')
+        code.append(f'ompi_coll_base_nbc_request_t* nb_request = (ompi_coll_base_nbc_request_t*)&{request_tmp_name};')
+        code.append(f'nb_request->data.release_arrays[idx++] = (void *){self.tmpname};')
+        code.append('nb_request->data.release_arrays[idx] = NULL;')
+        code.append('} else {')
+        code.append(f'free({self.tmpname});')
         code.append('}')
         return code
 
-
 @Type.add_type('DATATYPE_ARRAY_OUT', abi_type=['standard'])
 class TypeDatatypeArrayOutStandard(StandardABIType):
 
@@ -722,9 +740,20 @@ def init_code(self):
 @Type.add_type('NEIGHBOR_DATATYPE_ARRAY_ASYNC', abi_type=['standard'])
 class NeighborDatatypeArrayAsyncStandard(NeighborDatatypeArrayStandard):
 
+    @property
+    def need_async_cleanup(self):
+        return True
+
     @property
     def final_code(self):
-        code = ['{']
+        request_tmp_name = util.abi_tmp_name('request')
+        code = []
+        code.append('if((MPI_SUCCESS == ret_value) && (MPI_REQUEST_NULL != request_tmp)){')
+        code.append(f'ompi_coll_base_nbc_request_t* nb_request = (ompi_coll_base_nbc_request_t*)&{request_tmp_name};')
+        code.append(f'nb_request->data.release_arrays[idx++] = (void *){self.tmpname};')
+        code.append('nb_request->data.release_arrays[idx] = NULL;')
+        code.append('} else {')
+        code.append(f'free({self.tmpname});')
         code.append('}')
         return code
 

ompi/mpi/bindings/ompi_bindings/util.py

@@ -197,3 +197,6 @@ def prototype_needs_callback_wrappers(prototype):
     """Should this prototype need a callback wrappers"""
     return any(param.type_name in USER_CALLBACK_NAMES for param in prototype.params)
 
+def abi_tmp_name(name):
+    """Generate standardized tmp name for a supplied name"""
+    return f'{name}_tmp'