Emit host-side synchronization for GPU-to-host copies

conftest.py

@@ -0,0 +1,74 @@
+# Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved.
+"""
+Root pytest configuration file.
+"""
+import os
+import subprocess
+
+import pytest
+
+
+def parse_worker_index(worker_id: str) -> int:
+    """Parse the digits in a worker id or MPI rank string (e.g. 'gw3' -> 3, '12' -> 12).
+
+    Falls back to 0 for text with no digits (e.g. the non-xdist 'master' worker id).
+    """
+    digits = ''.join(char for char in worker_id if char.isdigit())
+    return int(digits) if digits else 0
+
+
+def list_cuda_devices() -> list:
+    """Enumerate CUDA device indices via `nvidia-smi -L`, without touching CUDA in this process."""
+    try:
+        result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, timeout=10)
+    except (OSError, subprocess.TimeoutExpired):
+        return []
+    if result.returncode != 0:
+        return []
+    device_lines = [line for line in result.stdout.splitlines() if line.startswith('GPU ')]
+    return [str(index) for index in range(len(device_lines))]
+
+
+def pick_gpu_worker_device(worker_id: str, device_pool: list) -> str:
+    """Round-robin a worker id (xdist worker or MPI rank) onto one device of device_pool."""
+    index = parse_worker_index(worker_id)
+    return device_pool[index % len(device_pool)]
+
+
+def resolve_worker_id() -> str:
+    """Id that places this process in a device rotation.
+
+    The pytest-xdist worker id if present, else the first set MPI/job-launcher rank env var
+    (dace.sdfg.sdfg.LAUNCHER_RANK_VARS, the same list dace itself reads for per-rank build
+    folders), else '' if this process is neither an xdist worker nor an MPI rank.
+    """
+    xdist_worker = os.environ.get('PYTEST_XDIST_WORKER')
+    if xdist_worker is not None:
+        return xdist_worker
+    from dace.sdfg.sdfg import LAUNCHER_RANK_VARS
+    for var in LAUNCHER_RANK_VARS:
+        rank = os.environ.get(var)
+        if rank:
+            return rank
+    return ''
+
+
+def pin_worker_to_gpu() -> None:
+    """Pin this worker process (pytest-xdist or an MPI-launched rank) to a single GPU.
+
+    Without this, every worker/rank sees all GPUs and piles its CUDA context onto device 0, which
+    is the root cause of flaky 'invalid device ordinal' failures under high worker/rank counts.
+    """
+    worker_id = resolve_worker_id()
+    if not worker_id:
+        return
+    preset = os.environ.get('CUDA_VISIBLE_DEVICES', '')
+    device_pool = [entry.strip() for entry in preset.split(',') if entry.strip()] if preset.strip() \
+        else list_cuda_devices()
+    if len(device_pool) <= 1:
+        return
+    os.environ['CUDA_VISIBLE_DEVICES'] = pick_gpu_worker_device(worker_id, device_pool)
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    pin_worker_to_gpu()

dace/codegen/targets/cpp.py

@@ -1397,18 +1397,29 @@ def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_str
 
             if (isinstance(edge.dst, nodes.AccessNode) and hasattr(edge.dst, '_cuda_stream')
                     and edge.dst._cuda_stream != node._cuda_stream):
-                callsite_stream.write(
-                    """DACE_GPU_CHECK({backend}EventRecord(__state->gpu_context->events[{ev}], {src_stream}));
+                # Stream assignment gives a cross-stream edge its own event. Event 0 belongs to some
+                # other edge, so recording into it when this edge has none breaks that edge's
+                # ordering instead of establishing this one -- let the host wait instead.
+                if hasattr(edge, "_cuda_event"):
+                    callsite_stream.write(
+                        """DACE_GPU_CHECK({backend}EventRecord(__state->gpu_context->events[{ev}], {src_stream}));
 DACE_GPU_CHECK({backend}StreamWaitEvent({dst_stream}, __state->gpu_context->events[{ev}], 0));""".format(
-                        ev=edge._cuda_event if hasattr(edge, "_cuda_event") else 0,
-                        src_stream=cudastream,
-                        dst_stream=common.gpu_stream_expr(edge.dst._cuda_stream),
-                        backend=backend,
-                    ),
-                    cfg,
-                    state_id,
-                    [edge.src, edge.dst],
-                )
+                            ev=edge._cuda_event,
+                            src_stream=cudastream,
+                            dst_stream=common.gpu_stream_expr(edge.dst._cuda_stream),
+                            backend=backend,
+                        ),
+                        cfg,
+                        state_id,
+                        [edge.src, edge.dst],
+                    )
+                else:
+                    callsite_stream.write(
+                        "DACE_GPU_CHECK(%sStreamSynchronize(%s));" % (backend, cudastream),
+                        cfg,
+                        state_id,
+                        [edge.src, edge.dst],
+                    )
                 continue
 
             # If a view, get the relevant access node
@@ -1421,24 +1432,37 @@ def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_str
             for e in dfg.out_edges(dstnode):
                 if isinstance(e.dst, nodes.AccessNode):
                     continue
-                # If no stream at destination: synchronize stream with host.
+                # If no stream at destination: the consumer runs on the host, so wait for the stream.
                 if not hasattr(e.dst, "_cuda_stream"):
-                    pass
-                    # Done at destination
-
-                # If different stream at destination: record event and wait
-                # for it in target stream.
-                elif e.dst._cuda_stream != node._cuda_stream:
                     callsite_stream.write(
-                        """{backend}EventRecord(__state->gpu_context->events[{ev}], {src_stream});
-    {backend}StreamWaitEvent({dst_stream}, __state->gpu_context->events[{ev}], 0);""".format(
-                            ev=e._cuda_event if hasattr(e, "_cuda_event") else 0,
-                            src_stream=cudastream,
-                            dst_stream=common.gpu_stream_expr(e.dst._cuda_stream),
-                            backend=backend,
-                        ),
+                        "DACE_GPU_CHECK(%sStreamSynchronize(%s));" % (backend, cudastream),
                         cfg,
                         state_id,
                         [e.src, e.dst],
                     )
+
+                # If different stream at destination: record event and wait
+                # for it in target stream.
+                elif e.dst._cuda_stream != node._cuda_stream:
+                    # Same as above: without an event of its own there is nothing to record into.
+                    if hasattr(e, "_cuda_event"):
+                        callsite_stream.write(
+                            """DACE_GPU_CHECK({backend}EventRecord(__state->gpu_context->events[{ev}], {src_stream}));
+DACE_GPU_CHECK({backend}StreamWaitEvent({dst_stream}, __state->gpu_context->events[{ev}], 0));""".format(
+                                ev=e._cuda_event,
+                                src_stream=cudastream,
+                                dst_stream=common.gpu_stream_expr(e.dst._cuda_stream),
+                                backend=backend,
+                            ),
+                            cfg,
+                            state_id,
+                            [e.src, e.dst],
+                        )
+                    else:
+                        callsite_stream.write(
+                            "DACE_GPU_CHECK(%sStreamSynchronize(%s));" % (backend, cudastream),
+                            cfg,
+                            state_id,
+                            [e.src, e.dst],
+                        )
                 # Otherwise, no synchronization necessary

dace/codegen/targets/cuda.py

@@ -1053,21 +1053,28 @@ def _emit_copy(self, state_id: int, src_node: nodes.Node, src_storage: dtypes.St
             else:
                 if max_streams >= 0:
                     print('WARNING: Undefined stream, reverting to default')
-                if dst_location == 'Host':
-                    is_sync = True
                 cudastream = 'nullptr'
 
+            # The host can read a host-located destination as soon as the copy is done, so the copy
+            # has to be waited for. Stream assignment stamps host containers that sit inside a GPU
+            # dataflow chain, so the stamp says nothing about who reads them.
+            if dst_location == 'Host':
+                is_sync = True
+
             # Handle case of impending kernel/tasklet on another stream
             if max_streams >= 0:
                 for e in state_dfg.out_edges(dst_node):
                     if isinstance(e.dst, nodes.AccessNode):
                         continue
                     if not hasattr(e.dst, '_cuda_stream'):
                         is_sync = True
-                    elif not hasattr(e, '_cuda_event'):
-                        is_sync = True
                     elif e.dst._cuda_stream != cudastream:
-                        syncwith[e.dst._cuda_stream] = e._cuda_event
+                        # A consumer on another stream is ordered by an event, or by the host when
+                        # stream assignment did not leave one.
+                        if hasattr(e, '_cuda_event'):
+                            syncwith[e.dst._cuda_stream] = e._cuda_event
+                        else:
+                            is_sync = True
 
                 cudastream = common.gpu_stream_expr(cudastream)
 
@@ -1283,8 +1290,10 @@ def _emit_copy(self, state_id: int, src_node: nodes.Node, src_storage: dtypes.St
 
             # Post-copy synchronization
             if is_sync:
-                # Synchronize with host (done at destination)
-                pass
+                # Every copy emitted above is asynchronous, so the host has to wait for the stream
+                # before it may read the destination.
+                callsite_stream.write('DACE_GPU_CHECK(%sStreamSynchronize(%s));\n' % (self.backend, cudastream), cfg,
+                                      state_id, [src_node, dst_node])
             else:
                 # Synchronize with other streams as necessary
                 for streamid, event in syncwith.items():

dace/codegen/targets/framecode.py

@@ -300,6 +300,15 @@ def generate_footer(self, sdfg: SDFG, global_stream: CodeIOStream, callsite_stre
             if target.has_initializer:
                 callsite_stream.write(
                     '__result |= __dace_init_%s(__state%s);' % (target.target_name, initparamnames_comma), sdfg)
+        # A failed target initializer leaves its part of the state struct unset, and everything below
+        # allocates against it -- persistent GPU arrays dereference __state->gpu_context, which
+        # __dace_init_cuda never constructs when it bails out on a missing device. Leave here first.
+        callsite_stream.write(f"""
+    if (__result) {{
+        delete __state;
+        return nullptr;
+    }}
+""", sdfg)
         for env in self.environments:
             init_code = _get_or_eval_sdfg_first_arg(env.init_code, sdfg)
             if init_code:

dace/libraries/blas/environments/cublas.py

@@ -35,7 +35,7 @@ def handle_setup_code(node):
         code = """\
 const int __dace_cuda_device = {location};
 cublasHandle_t &__dace_cublas_handle = __state->cublas_handle.Get(__dace_cuda_device);
-cublasSetStream(__dace_cublas_handle, __dace_current_stream);\n"""
+dace::blas::CheckCublasError(cublasSetStream(__dace_cublas_handle, __dace_current_stream));\n"""
 
         return code.format(location=location)
 

dace/libraries/blas/environments/rocblas.py

@@ -35,7 +35,7 @@ def handle_setup_code(node):
         code = """\
 const int __dace_cuda_device = {location};
 rocblas_handle &__dace_rocblas_handle = __state->rocblas_handle.Get(__dace_cuda_device);
-rocblas_set_stream(__dace_rocblas_handle, __dace_current_stream);\n"""
+dace::blas::CheckRocblasError(rocblas_set_stream(__dace_rocblas_handle, __dace_current_stream));\n"""
 
         return code.format(location=location)
 

dace/libraries/blas/include/dace_cublas.h

@@ -170,9 +170,11 @@ class CublasHandle {
     return f->second;
   }
 
+  // A destructor that throws terminates the process. Teardown failures have nowhere left to go, so
+  // they are dropped rather than turned into a crash that hides whatever the program computed.
   ~CublasHandle() {
     for (auto& h : handles_) {
-      CheckCublasError(cublasDestroy(h.second));
+      static_cast<void>(cublasDestroy(h.second));
     }
   }
 

dace/libraries/blas/include/dace_rocblas.h

@@ -173,9 +173,11 @@ class RocblasHandle {
     return f->second;
   }
 
+  // A destructor that throws terminates the process. Teardown failures have nowhere left to go, so
+  // they are dropped rather than turned into a crash that hides whatever the program computed.
   ~RocblasHandle() {
     for (auto& h : handles_) {
-      CheckRocblasError(rocblas_destroy_handle(h.second));
+      static_cast<void>(rocblas_destroy_handle(h.second));
     }
   }
 

dace/libraries/blas/nodes/batched_matmul.py

@@ -308,12 +308,13 @@ def expansion(node, state, sdfg):
                 alpha = f'{dtype.ctype}({node.alpha})'
 
             # Set pointer mode to host
-            call_prefix += f'''cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_HOST);
+            call_prefix += f'''dace::blas::CheckCublasError(
+                cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_HOST));
                 {dtype.ctype} alpha = {alpha};
                 {dtype.ctype} beta = 0;
                 '''
             call_suffix += '''
-    cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_DEVICE);
+    dace::blas::CheckCublasError(cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_DEVICE));
                 '''
             beta = f'({cdtype} *)&beta'
             alpha = f'({cdtype} *)&alpha'
@@ -327,15 +328,15 @@ def expansion(node, state, sdfg):
 
         # Matrix multiplication
         if (node.compute_type is None and node.accumulator_type is None and node.algorithm is None):
-            call = '''cublas{func}StridedBatched(__dace_cublas_handle,
+            call = '''dace::blas::CheckCublasError(cublas{func}StridedBatched(__dace_cublas_handle,
                 CUBLAS_OP_{ta}, CUBLAS_OP_{tb},
                 {M}, {N}, {K},
                 {alpha},
                 ({dtype}*){array_prefix}{x}, {lda}, {stride_a},
                 ({dtype}*){array_prefix}{y}, {ldb}, {stride_b},
                 {beta},
                 ({dtype}*){array_prefix}_c, {ldc}, {stride_c},
-                {BATCH});'''.format_map(opt)
+                {BATCH}));'''.format_map(opt)
         else:
             if node.compute_type is not None:
                 acctype = node.compute_type
@@ -350,7 +351,7 @@ def expansion(node, state, sdfg):
                 algorithm = node.algorithm
 
             call = f'''
-            cublasGemmStridedBatchedEx(__dace_cublas_handle,
+            dace::blas::CheckCublasError(cublasGemmStridedBatchedEx(__dace_cublas_handle,
                 CUBLAS_OP_{opt['ta']}, CUBLAS_OP_{opt['tb']},
                 {opt['M']}, {opt['N']}, {opt['K']},
                 {alpha},
@@ -365,7 +366,7 @@ def expansion(node, state, sdfg):
                 {dtype_to_cudadatatype(opt['cdtype'])},
                 {opt['ldc']}, {opt['stride_c']},
                 {opt['BATCH']},
-                {acctype}, {algorithm});
+                {acctype}, {algorithm}));
             '''
 
         code = call_prefix + call + call_suffix

dace/libraries/blas/nodes/dot.py

@@ -125,11 +125,11 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
 
         code = environments.cublas.cuBLAS.handle_setup_code(node)
         if node.accumulator_type is None:
-            code += f"""cublas{func}(__dace_cublas_handle, {n}, _x, {stride_x}, _y,
-                             {stride_y}, _result);"""
+            code += f"""dace::blas::CheckCublasError(cublas{func}(__dace_cublas_handle, {n}, _x, {stride_x}, _y,
+                             {stride_y}, _result));"""
         else:
             code += f"""
-            cublasDotEx(
+            dace::blas::CheckCublasError(cublasDotEx(
                 __dace_cublas_handle,
                 {n},
                 _x,
@@ -140,7 +140,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs):
                 {stride_y},
                 _result,
                 {blas_helpers.dtype_to_cudadatatype(desc_res.dtype)},
-                {blas_helpers.dtype_to_cudadatatype(node.accumulator_type)});
+                {blas_helpers.dtype_to_cudadatatype(node.accumulator_type)}));
             """
 
         tasklet = dace.sdfg.nodes.Tasklet(node.name,