From 80187933465cc65c5ae80e7014058139983ce393 Mon Sep 17 00:00:00 2001 From: Hui Zhou Date: Mon, 3 Jun 2024 15:21:03 -0500 Subject: [PATCH] python/binding_c: Add streamsync directive Extend the functionality of MPIX stream by allowing passing a stream communication with local gpu stream to regular MPI functions. The semantics is to run an implicit gpu stream synchronize before the MPI operation. This is semantically serializing the MPI oepration with the GPU stream, albeit in a heavy handed way. --- maint/local_python/binding_c.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/maint/local_python/binding_c.py b/maint/local_python/binding_c.py index 302e991a5ed..108b61de670 100644 --- a/maint/local_python/binding_c.py +++ b/maint/local_python/binding_c.py @@ -1474,6 +1474,16 @@ def dump_function_normal(func): # ---- def dump_body_of_routine(): do_threadcomm = False + if RE.search(r'streamsync', func['extra'], re.IGNORECASE): + if '_has_comm' in func: + G.out.append("MPIR_Stream *stream = MPIR_stream_comm_get_local_stream(%s_ptr);" % func['_has_comm']) + elif '_has_win' in func: + G.out.append("MPIR_Stream *stream = MPIR_stream_comm_get_local_stream(%s_ptr->comm_ptr);" % func['_has_win']) + else: + raise Exception("streamsync not supported in %s" % func['name']) + dump_if_open("stream && stream->type == MPIR_STREAM_GPU") + G.out.append("MPL_gpu_stream_synchronize(stream->u.gpu_stream);") + dump_if_close() if RE.search(r'threadcomm', func['extra'], re.IGNORECASE): do_threadcomm = True G.out.append("#ifdef ENABLE_THREADCOMM") @@ -1506,6 +1516,7 @@ def dump_body_of_routine(): if do_threadcomm: dump_if_close() + # ---- G.out.append("/* ... body of routine ... */")