diff --git a/dace/codegen/targets/cpp.py b/dace/codegen/targets/cpp.py index 046de415f2..eae90d882c 100644 --- a/dace/codegen/targets/cpp.py +++ b/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 diff --git a/dace/codegen/targets/cuda.py b/dace/codegen/targets/cuda.py index cce465ce63..9a48d4d03e 100644 --- a/dace/codegen/targets/cuda.py +++ b/dace/codegen/targets/cuda.py @@ -1053,10 +1053,14 @@ 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): @@ -1064,10 +1068,13 @@ def _emit_copy(self, state_id: int, src_node: nodes.Node, src_storage: dtypes.St 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(): diff --git a/dace/codegen/targets/framecode.py b/dace/codegen/targets/framecode.py index f7d1c73a22..7f1a43a299 100644 --- a/dace/codegen/targets/framecode.py +++ b/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: diff --git a/dace/libraries/blas/environments/cublas.py b/dace/libraries/blas/environments/cublas.py index ef73b511c0..89933bbb5e 100644 --- a/dace/libraries/blas/environments/cublas.py +++ b/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) diff --git a/dace/libraries/blas/environments/rocblas.py b/dace/libraries/blas/environments/rocblas.py index 47e16531ff..3e5bcd326c 100644 --- a/dace/libraries/blas/environments/rocblas.py +++ b/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) diff --git a/dace/libraries/blas/include/dace_cublas.h b/dace/libraries/blas/include/dace_cublas.h index 3547a009d2..62da94c08f 100644 --- a/dace/libraries/blas/include/dace_cublas.h +++ b/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(cublasDestroy(h.second)); } } diff --git a/dace/libraries/blas/include/dace_rocblas.h b/dace/libraries/blas/include/dace_rocblas.h index 00469136a3..b0e7bd6453 100644 --- a/dace/libraries/blas/include/dace_rocblas.h +++ b/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(rocblas_destroy_handle(h.second)); } } diff --git a/dace/libraries/blas/nodes/batched_matmul.py b/dace/libraries/blas/nodes/batched_matmul.py index 0abc94c1fc..739e7095be 100644 --- a/dace/libraries/blas/nodes/batched_matmul.py +++ b/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,7 +328,7 @@ 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}, @@ -335,7 +336,7 @@ def expansion(node, state, sdfg): ({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 diff --git a/dace/libraries/blas/nodes/dot.py b/dace/libraries/blas/nodes/dot.py index 42ce0c0fa8..d21e0136f2 100644 --- a/dace/libraries/blas/nodes/dot.py +++ b/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, diff --git a/dace/libraries/blas/nodes/gemm.py b/dace/libraries/blas/nodes/gemm.py index c78bcdf03b..d31cdc16a7 100644 --- a/dace/libraries/blas/nodes/gemm.py +++ b/dace/libraries/blas/nodes/gemm.py @@ -288,11 +288,13 @@ def expansion(cls, node, state, sdfg): beta = f'{dtype.ctype}({node.beta})' # Set pointer mode to host - call_prefix += f'''{cls.set_pointer_mode}(__dace_{cls.backend}blas_handle, {cls.pointer_host}); + call_prefix += f'''{cls.check_error}( + {cls.set_pointer_mode}(__dace_{cls.backend}blas_handle, {cls.pointer_host})); {dtype.ctype} __alpha = {alpha}; {dtype.ctype} __beta = {beta}; ''' - call_suffix += f'''{cls.set_pointer_mode}(__dace_{cls.backend}blas_handle, {cls.pointer_device});''' + call_suffix += f'''{cls.check_error}( + {cls.set_pointer_mode}(__dace_{cls.backend}blas_handle, {cls.pointer_device}));''' alpha = f'({cdtype} *)&__alpha' beta = f'({cdtype} *)&__beta' else: @@ -310,15 +312,16 @@ def expansion(cls, node, state, sdfg): opt['backend'] = cls.backend opt['backend_op_ta'] = cls.backend_op(opt['ta']) opt['backend_op_tb'] = cls.backend_op(opt['tb']) + opt['check_error'] = cls.check_error - call = '''{backend}blas{func}(__dace_{backend}blas_handle, + call = '''{check_error}({backend}blas{func}(__dace_{backend}blas_handle, {backend_op_ta}, {backend_op_tb}, {M}, {N}, {K}, {alpha}, ({dtype}*){arr_prefix}{x}, {lda}, ({dtype}*){arr_prefix}{y}, {ldb}, {beta}, - ({dtype}*){arr_prefix}_c, {ldc});'''.format_map(opt) + ({dtype}*){arr_prefix}_c, {ldc}));'''.format_map(opt) else: if node.compute_type is not None: acctype = node.compute_type @@ -333,7 +336,7 @@ def expansion(cls, node, state, sdfg): algorithm = node.algorithm call = f''' - {cls.backend}blas{cls.ex_suffix}(__dace_{cls.backend}blas_handle, + {cls.check_error}({cls.backend}blas{cls.ex_suffix}(__dace_{cls.backend}blas_handle, {cls.backend_op(opt['ta'])}, {cls.backend_op(opt['tb'])}, {opt['M']}, {opt['N']}, {opt['K']}, @@ -349,7 +352,7 @@ def expansion(cls, node, state, sdfg): {dtype_to_cudadatatype(opt['cdtype'])}, {opt['ldc']}, {acctype}, - {algorithm}); + {algorithm})); ''' code = (call_prefix + call + call_suffix) @@ -423,6 +426,7 @@ class ExpandGemmCuBLAS(ExpandGemmGPUBLAS): pointer_host = 'CUBLAS_POINTER_MODE_HOST' pointer_device = 'CUBLAS_POINTER_MODE_DEVICE' ex_suffix = 'GemmEx' + check_error = 'dace::blas::CheckCublasError' @classmethod def backend_op(cls, mode: str) -> str: @@ -442,6 +446,7 @@ class ExpandGemmRocBLAS(ExpandGemmGPUBLAS): pointer_host = 'rocblas_pointer_mode_host' pointer_device = 'rocblas_pointer_mode_device' ex_suffix = '_gemm_ex' + check_error = 'dace::blas::CheckRocblasError' @classmethod def backend_op(cls, mode: str) -> str: diff --git a/dace/libraries/blas/nodes/gemv.py b/dace/libraries/blas/nodes/gemv.py index 9ca6368b45..e437f5b095 100644 --- a/dace/libraries/blas/nodes/gemv.py +++ b/dace/libraries/blas/nodes/gemv.py @@ -183,12 +183,13 @@ def expansion(node: 'Gemv', state, sdfg, m=None, n=None, **kwargs): beta = f'{dtype.ctype}({node.beta})' # 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 = {beta}; ''' call_suffix += ''' -cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_DEVICE); +dace::blas::CheckCublasError(cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_DEVICE)); ''' alpha = f'({ctype} *)&alpha' beta = f'({ctype} *)&beta' @@ -197,8 +198,8 @@ def expansion(node: 'Gemv', state, sdfg, m=None, n=None, **kwargs): beta = constants[node.beta] code = (call_prefix + f""" -cublas{func}(__dace_cublas_handle, {trans}, {m}, {n}, {alpha}, _A, {lda}, - _x, {strides_x[0]}, {beta}, _y, {strides_y[0]}); +dace::blas::CheckCublasError(cublas{func}(__dace_cublas_handle, {trans}, {m}, {n}, {alpha}, _A, {lda}, + _x, {strides_x[0]}, {beta}, _y, {strides_y[0]})); """ + call_suffix) tasklet = dace.sdfg.nodes.Tasklet(node.name, diff --git a/dace/libraries/linalg/nodes/transpose.py b/dace/libraries/linalg/nodes/transpose.py index 7676f09a4a..05a0c2dfc7 100644 --- a/dace/libraries/linalg/nodes/transpose.py +++ b/dace/libraries/linalg/nodes/transpose.py @@ -210,9 +210,9 @@ def expansion(node, state, sdfg, **kwargs): _, _, (m, n), (istride, _) = _get_transpose_input(node, state, sdfg) _, _, _, (ostride, _) = _get_transpose_output(node, state, sdfg) - code = (blas_environments.cublas.cuBLAS.handle_setup_code(node) + f"""cublas{func}( + code = (blas_environments.cublas.cuBLAS.handle_setup_code(node) + f"""dace::blas::CheckCublasError(cublas{func}( __dace_cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, - {m}, {n}, {alpha}, ({cdtype}*)_inp, {n}, {beta}, ({cdtype}*)_inp, {m}, ({cdtype}*)_out, {m}); + {m}, {n}, {alpha}, ({cdtype}*)_inp, {n}, {beta}, ({cdtype}*)_inp, {m}, ({cdtype}*)_out, {m})); """) tasklet = dace.sdfg.nodes.Tasklet(node.name, diff --git a/dace/runtime/include/dace/cuda/cudacommon.cuh b/dace/runtime/include/dace/cuda/cudacommon.cuh index 7b39f3c4ba..ec773366d2 100644 --- a/dace/runtime/include/dace/cuda/cudacommon.cuh +++ b/dace/runtime/include/dace/cuda/cudacommon.cuh @@ -58,6 +58,7 @@ struct Context { } ~Context() { delete[] streams; + delete[] internal_streams; delete[] events; } // Keep the first error. One failure tends to produce more, and only the first names the call that diff --git a/tests/codegen/gpu_codegen_error_checks_test.py b/tests/codegen/gpu_codegen_error_checks_test.py new file mode 100644 index 0000000000..caf72a6f10 --- /dev/null +++ b/tests/codegen/gpu_codegen_error_checks_test.py @@ -0,0 +1,166 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Generated GPU code has to check what it calls, and has to stop before it uses what failed. + +These assert on emitted code, so they need a GPU for neither compilation nor a run. +""" +import re + +import dace +from dace import dtypes +from dace.libraries import blas + +BAILOUT = r'if \(__result\)' +EVENT_CALL = r'(?:cuda|hip)(?:EventRecord|StreamWaitEvent)\(' +CUBLAS_CALL = r'cublas[A-Z]\w*\(' + + +def generated_code(sdfg: dace.SDFG) -> str: + """Every generated file for ``sdfg``, joined.""" + return '\n'.join(code.clean_code for code in sdfg.generate_code()) + + +def init_function(code: str, name: str) -> str: + """The body of ``__dace_init_``.""" + match = re.search(r'__dace_init_' + re.escape(name) + r'\(.*?\n\}', code, re.S) + assert match, f'no __dace_init_{name} was emitted, so this test is anchored on nothing' + return match.group(0) + + +def persistent_gpu_transient() -> dace.SDFG: + """A persistent GPU transient, whose allocation is hoisted into the init function.""" + sdfg = dace.SDFG('persistent_gpu_transient') + sdfg.add_array('A', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_array('B', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_transient('T', [20], + dace.float64, + storage=dtypes.StorageType.GPU_Global, + lifetime=dtypes.AllocationLifetime.Persistent) + + state = sdfg.add_state('main') + a = state.add_access('A') + t = state.add_access('T') + b = state.add_access('B') + state.add_nedge(a, t, dace.Memlet('A[0:20]')) + state.add_nedge(t, b, dace.Memlet('T[0:20]')) + sdfg.validate() + return sdfg + + +def cross_stream_consumer() -> dace.SDFG: + """Two independent kernels feeding a third, so one producer is ordered against another stream.""" + sdfg = dace.SDFG('cross_stream_consumer') + for name in ('A', 'B', 'C'): + sdfg.add_array(name, [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_transient('T1', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_transient('T2', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + + state = sdfg.add_state('main') + accesses = {name: state.add_access(name) for name in ('A', 'B', 'C', 'T1', 'T2')} + + def producer(name: str, src: str, dst: str, code: str) -> None: + entry, exit_ = state.add_map(name, {'i': '0:20'}, schedule=dtypes.ScheduleType.GPU_Device) + tasklet = state.add_tasklet(name + '_t', {'inp'}, {'out'}, code) + entry.add_in_connector('IN_x') + entry.add_out_connector('OUT_x') + exit_.add_in_connector('IN_y') + exit_.add_out_connector('OUT_y') + state.add_edge(accesses[src], None, entry, 'IN_x', dace.Memlet(f'{src}[0:20]')) + state.add_edge(entry, 'OUT_x', tasklet, 'inp', dace.Memlet(f'{src}[i]')) + state.add_edge(tasklet, 'out', exit_, 'IN_y', dace.Memlet(f'{dst}[i]')) + state.add_edge(exit_, 'OUT_y', accesses[dst], None, dace.Memlet(f'{dst}[0:20]')) + + producer('k1', 'A', 'T1', 'out = inp * 2') + producer('k2', 'B', 'T2', 'out = inp * 3') + + entry, exit_ = state.add_map('k3', {'i': '0:20'}, schedule=dtypes.ScheduleType.GPU_Device) + tasklet = state.add_tasklet('k3_t', {'p', 'q'}, {'out'}, 'out = p + q') + for conn in ('IN_1', 'IN_2'): + entry.add_in_connector(conn) + for conn in ('OUT_1', 'OUT_2'): + entry.add_out_connector(conn) + exit_.add_in_connector('IN_c') + exit_.add_out_connector('OUT_c') + state.add_edge(accesses['T1'], None, entry, 'IN_1', dace.Memlet('T1[0:20]')) + state.add_edge(accesses['T2'], None, entry, 'IN_2', dace.Memlet('T2[0:20]')) + state.add_edge(entry, 'OUT_1', tasklet, 'p', dace.Memlet('T1[i]')) + state.add_edge(entry, 'OUT_2', tasklet, 'q', dace.Memlet('T2[i]')) + state.add_edge(tasklet, 'out', exit_, 'IN_c', dace.Memlet('C[i]')) + state.add_edge(exit_, 'OUT_c', accesses['C'], None, dace.Memlet('C[0:20]')) + sdfg.validate() + return sdfg + + +def cublas_gemm(alpha: float = 1.0) -> dace.SDFG: + """A GEMM library node expanded onto cuBLAS. ``alpha`` off 1.0 also brings out the pointer mode.""" + sdfg = dace.SDFG('cublas_gemm') + for name in ('A', 'B', 'C'): + sdfg.add_array(name, [20, 20], dace.float64, storage=dtypes.StorageType.GPU_Global) + + state = sdfg.add_state('main') + node = blas.Gemm('gemm', alpha=alpha) + node.implementation = 'cuBLAS' + state.add_node(node) + state.add_edge(state.add_read('A'), None, node, '_a', dace.Memlet('A[0:20, 0:20]')) + state.add_edge(state.add_read('B'), None, node, '_b', dace.Memlet('B[0:20, 0:20]')) + state.add_edge(node, '_c', state.add_write('C'), None, dace.Memlet('C[0:20, 0:20]')) + sdfg.expand_library_nodes() + sdfg.validate() + return sdfg + + +def test_a_failed_target_initializer_stops_before_the_state_it_left_unset(): + """``__dace_init_cuda`` returns early without a gpu_context when no device is present.""" + init = init_function(generated_code(persistent_gpu_transient()), 'persistent_gpu_transient') + initializer = re.search(r'__result \|= __dace_init_cuda\(', init) + assert initializer, 'the CUDA target initializer is not called, so this test is anchored on nothing' + allocation = re.search(r'DACE_GPU_CHECK\(', init) + assert allocation, 'the persistent GPU allocation was not hoisted into the init function' + bailout = re.search(BAILOUT, init[initializer.end():allocation.start()]) + assert bailout, ('the persistent GPU allocation runs even when __dace_init_cuda failed, and every DACE_GPU_CHECK ' + 'in it dereferences the gpu_context that the failed initializer never constructed') + + +def test_the_init_function_still_checks_what_runs_after_the_allocations(): + """Environment and SDFG-level init code can fail too, so the later guard has to stay.""" + init = init_function(generated_code(persistent_gpu_transient()), 'persistent_gpu_transient') + allocation = re.search(r'DACE_GPU_CHECK\(', init) + assert allocation, 'the persistent GPU allocation was not hoisted into the init function' + assert re.search(BAILOUT, init[allocation.end():]), ( + 'nothing checks __result after the allocation and init code, so a failure there returns a live state') + + +def test_cross_stream_event_synchronization_is_checked(): + """A silent EventRecord failure loses the ordering it was supposed to establish.""" + code = generated_code(cross_stream_consumer()) + calls = list(re.finditer(EVENT_CALL, code)) + assert calls, 'no cross-stream event synchronization was emitted, so this test is anchored on nothing' + unchecked = [call.group(0) for call in calls if not code[:call.start()].endswith('DACE_GPU_CHECK(')] + assert not unchecked, f'event synchronization emitted without an error check: {unchecked}' + + +def test_cublas_calls_are_checked(): + """cuBLAS reports through its return value only, so a dropped status is a silently wrong result.""" + for alpha in (1.0, 2.0): + code = generated_code(cublas_gemm(alpha)) + calls = list(re.finditer(CUBLAS_CALL, code)) + assert calls, f'no cuBLAS call was emitted for alpha={alpha}, so this test is anchored on nothing' + unchecked = [ + call.group(0) for call in calls + if not code[:call.start()].rstrip().endswith('dace::blas::CheckCublasError(') + ] + assert not unchecked, f'cuBLAS called without checking its status: {unchecked}' + + +def test_the_cublas_gemm_expansion_still_emits_the_pointer_mode_switch(): + """The check has to wrap the pointer mode switch, not replace it.""" + code = generated_code(cublas_gemm(2.0)) + assert 'cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_HOST)' in code + assert 'cublasSetPointerMode(__dace_cublas_handle, CUBLAS_POINTER_MODE_DEVICE)' in code + + +if __name__ == '__main__': + test_a_failed_target_initializer_stops_before_the_state_it_left_unset() + test_the_init_function_still_checks_what_runs_after_the_allocations() + test_cross_stream_event_synchronization_is_checked() + test_cublas_calls_are_checked() + test_the_cublas_gemm_expansion_still_emits_the_pointer_mode_switch() diff --git a/tests/codegen/gpu_d2h_host_sync_test.py b/tests/codegen/gpu_d2h_host_sync_test.py new file mode 100644 index 0000000000..f3d71bb88d --- /dev/null +++ b/tests/codegen/gpu_d2h_host_sync_test.py @@ -0,0 +1,137 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""A GPU-to-host copy has to be waited for before the host reads its destination. + +The copy is issued as ``MemcpyAsync`` on a stream, so nothing orders it against the host. The +trisolv shape is the one that bites: a device scalar is copied into host memory and the very next +kernel launch takes that scalar *by value*, which means the host reads it while packing the launch +arguments. Pageable host memory hides this today because the driver blocks on it anyway; pinned +memory, HIP and graph capture do not. + +These assert on emitted code, so they need a GPU for neither compilation nor a run. +""" +import re + +import dace +from dace import dtypes + +COPY = r'Memcpy(?:2D)?Async\(' +SYNC = r'StreamSynchronize\(' + + +def generated_code(sdfg: dace.SDFG) -> str: + """Every generated file for ``sdfg``, joined.""" + return '\n'.join(code.clean_code for code in sdfg.generate_code()) + + +def scalar_through_host(scalar_storage: dtypes.StorageType) -> dace.SDFG: + """Device scalar copied into ``scalar_storage``, then consumed by value in a device kernel.""" + sdfg = dace.SDFG('d2h_then_kernel_' + scalar_storage.name) + sdfg.add_array('A', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_array('B', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_scalar('s', dace.float64, storage=scalar_storage, transient=True) + + state = sdfg.add_state('main') + a = state.add_access('A') + s = state.add_access('s') + b = state.add_access('B') + state.add_edge(a, None, s, None, dace.Memlet('A[0]')) + + entry, exit_ = state.add_map('kern', {'i': '0:20'}, schedule=dtypes.ScheduleType.GPU_Device) + tasklet = state.add_tasklet('scale', {'inp', 'sc'}, {'out'}, 'out = inp * sc') + for conn in ('IN_s', 'IN_A'): + entry.add_in_connector(conn) + for conn in ('OUT_s', 'OUT_A'): + entry.add_out_connector(conn) + exit_.add_in_connector('IN_B') + exit_.add_out_connector('OUT_B') + state.add_edge(s, None, entry, 'IN_s', dace.Memlet('s[0]')) + state.add_edge(a, None, entry, 'IN_A', dace.Memlet('A[0:20]')) + state.add_edge(entry, 'OUT_s', tasklet, 'sc', dace.Memlet('s[0]')) + state.add_edge(entry, 'OUT_A', tasklet, 'inp', dace.Memlet('A[i]')) + state.add_edge(tasklet, 'out', exit_, 'IN_B', dace.Memlet('B[i]')) + state.add_edge(exit_, 'OUT_B', b, None, dace.Memlet('B[0:20]')) + sdfg.validate() + return sdfg + + +def device_to_device() -> dace.SDFG: + """The same shape with the intermediate left on the device, so no host ever reads it.""" + sdfg = dace.SDFG('d2d_then_kernel') + sdfg.add_array('A', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_array('B', [20], dace.float64, storage=dtypes.StorageType.GPU_Global) + sdfg.add_array('C', [20], dace.float64, storage=dtypes.StorageType.GPU_Global, transient=True) + + state = sdfg.add_state('main') + a = state.add_access('A') + c = state.add_access('C') + b = state.add_access('B') + state.add_edge(a, None, c, None, dace.Memlet('A[0:20]')) + + entry, exit_ = state.add_map('kern', {'i': '0:20'}, schedule=dtypes.ScheduleType.GPU_Device) + tasklet = state.add_tasklet('scale', {'inp'}, {'out'}, 'out = inp * 2') + entry.add_in_connector('IN_C') + entry.add_out_connector('OUT_C') + exit_.add_in_connector('IN_B') + exit_.add_out_connector('OUT_B') + state.add_edge(c, None, entry, 'IN_C', dace.Memlet('C[0:20]')) + state.add_edge(entry, 'OUT_C', tasklet, 'inp', dace.Memlet('C[i]')) + state.add_edge(tasklet, 'out', exit_, 'IN_B', dace.Memlet('B[i]')) + state.add_edge(exit_, 'OUT_B', b, None, dace.Memlet('B[0:20]')) + sdfg.validate() + return sdfg + + +def host_copy_and_launch(code: str): + """The device-to-host copy and the kernel launch that reads its destination.""" + copy = re.search(COPY + r'[^;]*DeviceToHost[^;]*\)', code) + assert copy, 'no device-to-host copy was emitted, so this test is anchored on nothing' + launch = re.search(r'__dace_runkernel_\w+\(', code[copy.end():]) + assert launch, 'no kernel launch follows the device-to-host copy, so this test is anchored on nothing' + return copy, copy.end() + launch.start() + + +def test_a_device_to_host_copy_is_synchronized_before_its_host_reader(): + code = generated_code(scalar_through_host(dtypes.StorageType.CPU_Heap)) + copy, launch_at = host_copy_and_launch(code) + between = code[copy.end():launch_at] + assert re.search( + SYNC, + between), ('the kernel launch reads the copied scalar by value on the host, but no stream synchronization ' + 'separates it from the asynchronous device-to-host copy') + + +def test_the_synchronization_names_the_stream_the_copy_was_issued_on(): + """Waiting on some other stream orders nothing.""" + code = generated_code(scalar_through_host(dtypes.StorageType.CPU_Heap)) + copy, launch_at = host_copy_and_launch(code) + stream = copy.group(0).rsplit(',', 1)[1].strip().rstrip(')') + sync = re.search(SYNC + re.escape(stream) + r'\)', code[copy.end():launch_at]) + assert sync, f'the synchronization before the kernel launch does not wait on {stream}' + + +def test_pinned_destinations_are_synchronized_too(): + """Pinned memory is where the accidental blocking of pageable copies stops covering for this.""" + code = generated_code(scalar_through_host(dtypes.StorageType.CPU_Pinned)) + copy, launch_at = host_copy_and_launch(code) + assert re.search(SYNC, code[copy.end():launch_at]), ( + 'a copy into pinned host memory is not waited for before the host reads the destination') + + +def test_a_device_to_device_copy_does_not_wait_on_the_host(): + """No host reads the destination, so a host wait would only serialize the stream.""" + code = generated_code(device_to_device()) + copy = re.search(COPY + r'[^;]*DeviceToDevice[^;]*\)', code) + assert copy, 'no device-to-device copy was emitted, so this test is anchored on nothing' + launch = re.search(r'__dace_runkernel_\w+\(', code[copy.end():]) + assert launch, 'no kernel launch follows the device-to-device copy, so this test is anchored on nothing' + between = code[copy.end():copy.end() + launch.start()] + assert not re.search(SYNC, between), ('a device-to-device copy is followed by a host wait, which serializes the ' + 'stream for a destination no host reads') + assert not re.search(r'DeviceToHost', code), 'the device-to-device fixture emitted a host copy after all' + + +if __name__ == '__main__': + test_a_device_to_host_copy_is_synchronized_before_its_host_reader() + test_the_synchronization_names_the_stream_the_copy_was_issued_on() + test_pinned_destinations_are_synchronized_too() + test_a_device_to_device_copy_does_not_wait_on_the_host()