From 6c0134e454dfa832d081df62065061f77e3cae18 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Thu, 6 Aug 2026 15:56:27 +0200 Subject: [PATCH 01/11] Give each rank its own build cache root, under cache_distaware Ranks of one job derive the same build folder, so ranks that each compile build on top of each other and can load a library another rank is still writing. Eight processes running one GPU test out of one folder failed six times; with a folder each, none. The new cache_distaware config entry names the build cache root after the rank the launcher (MPI, Flux, Slurm) advertises. It is off by default, because sharing one build is also a valid setup: distributed_compile has rank 0 build and every other rank load its folder. That path now pins the broadcast folder on the ranks that hold the SDFG, the others being free to pass None. --- dace/config_schema.yml | 10 +++ dace/sdfg/sdfg.py | 31 ++++++++- dace/sdfg/utils.py | 7 +- tests/custom_build_folder_test.py | 103 ++++++++++++++++++++++++++++++ 4 files changed, 148 insertions(+), 3 deletions(-) diff --git a/dace/config_schema.yml b/dace/config_schema.yml index ee7a11746b..5baf488c08 100644 --- a/dace/config_schema.yml +++ b/dace/config_schema.yml @@ -738,6 +738,16 @@ required: potentially build time, but disallows executing SDFGs in parallel and caching of more than one simultaneous SDFG. + cache_distaware: + type: bool + default: false + title: Distribution-aware build cache + description: > + Give every rank of a job its own build folder, named after the rank its launcher + (MPI, Flux, Slurm) advertises. Without it, ranks that each compile share one folder + and can load a library another rank is still writing. Leave it off when only one + rank compiles, as ``dace.sdfg.utils.distributed_compile`` does. + store_history: type: bool default: true diff --git a/dace/sdfg/sdfg.py b/dace/sdfg/sdfg.py index 1ee07e4a2a..2ac48e0c0b 100644 --- a/dace/sdfg/sdfg.py +++ b/dace/sdfg/sdfg.py @@ -35,6 +35,20 @@ ShapeType = Sequence[Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic]] RankType = Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic] +#: How a launcher tells a task its rank, most specific first. Read instead of importing mpi4py, +#: which is optional and initializes MPI. All are job-unique; node-local counters are not. +LAUNCHER_RANK_VARS = ( + 'OMPI_COMM_WORLD_RANK', # Open MPI and the vendor MPIs built on it + 'MV2_COMM_WORLD_RANK', # MVAPICH2 + 'PMIX_RANK', # Open MPI 4+, Slurm pmix + 'PMI_RANK', # MPICH, Intel MPI, Cray MPICH + 'PMI_ID', # older MPICH + 'FLUX_TASK_RANK', # Flux + 'PALS_RANKID', # HPE/Cray PALS + 'ALPS_APP_PE', # Cray ALPS + 'SLURM_PROCID', # srun with no MPI +) + if TYPE_CHECKING: from dace.codegen.instrumentation.report import InstrumentationReport from dace.codegen.instrumentation.data.data_report import InstrumentedDataReport @@ -42,6 +56,21 @@ from dace.sdfg.analysis.schedule_tree.treenodes import ScheduleTreeRoot +def build_folder_root() -> str: + """The build cache root, one per rank if ``cache_distaware`` is on and a launcher set a rank. + + Ranks that each compile otherwise share a folder and can load each other's half-written library. + """ + base = Config.get('default_build_folder') + if not Config.get_bool('cache_distaware'): + return base + for var in LAUNCHER_RANK_VARS: + rank = os.environ.get(var) + if rank: + return f'{base}_rank{rank}' + return base + + class NestedDict(dict): def __init__(self, mapping=None): @@ -1217,7 +1246,7 @@ def build_folder(self) -> str: if self._build_folder is not None: return self._build_folder cache_config = Config.get('cache') - base_folder = Config.get('default_build_folder') + base_folder = build_folder_root() if cache_config == 'single': # Always use the same directory, overwriting any other program, # preventing parallelism and caching of multiple programs, but diff --git a/dace/sdfg/utils.py b/dace/sdfg/utils.py index fa052cb324..629e6a102b 100644 --- a/dace/sdfg/utils.py +++ b/dace/sdfg/utils.py @@ -1672,15 +1672,16 @@ def load_precompiled_sdfg(*args, **kwargs) -> csdfg.CompiledSDFG: return sdfg_compiler.load_precompiled_sdfg(*args, **kwargs) -def distributed_compile(sdfg: SDFG, comm, *, validate: bool = True) -> csdfg.CompiledSDFG: +def distributed_compile(sdfg: Optional[SDFG], comm, *, validate: bool = True) -> csdfg.CompiledSDFG: """ Compiles an SDFG in rank 0 of MPI communicator ``comm``. Then, the compiled SDFG is loaded in all other ranks. - :param sdfg: SDFG to be compiled. + :param sdfg: SDFG to be compiled. Ranks other than 0 only load, and may pass ``None``. :param comm: MPI communicator. ``Intracomm`` is the base mpi4py communicator class. :param validate: If True, validates the SDFG prior to generating code. :return: Compiled SDFG. :note: This method can be used only if the module mpi4py is installed. + :note: Only rank 0 builds, so a rank holding the SDFG is pinned to rank 0's folder. :todo: Relocate this function to `dace.codegen.compiler`. """ @@ -1695,6 +1696,8 @@ def distributed_compile(sdfg: SDFG, comm, *, validate: bool = True) -> csdfg.Com # Broadcasts build folder. folder = comm.bcast(folder, root=0) + if sdfg is not None: + sdfg.build_folder = folder # Loads compiled SDFG. if rank > 0: diff --git a/tests/custom_build_folder_test.py b/tests/custom_build_folder_test.py index d1d22fb3ac..7015e250ef 100644 --- a/tests/custom_build_folder_test.py +++ b/tests/custom_build_folder_test.py @@ -1,8 +1,12 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import os +import pytest import tempfile +from dace.sdfg import sdfg as sdfg_module +from dace.sdfg import utils as sdfg_utils + @dace.program def customprog(A: dace.float64[20]): @@ -24,5 +28,104 @@ def test_custom_build_folder(): del csdfg +@pytest.fixture +def unlaunched(monkeypatch): + """Drop the rank and cache settings the surrounding environment exports, which override config.""" + for var in sdfg_module.LAUNCHER_RANK_VARS: + monkeypatch.delenv(var, raising=False) + for var in ('DACE_cache', 'DACE_cache_distaware', 'DACE_default_build_folder'): + monkeypatch.delenv(var, raising=False) + return monkeypatch + + +@pytest.mark.parametrize('rank_var', sdfg_module.LAUNCHER_RANK_VARS) +def test_distaware_gives_each_rank_its_own_cache_root(unlaunched, rank_var): + """Ranks that each compile would otherwise build into one folder and load a half-written .so.""" + unlaunched.setenv('DACE_cache_distaware', '1') + + unlaunched.setenv(rank_var, '0') + rank0 = sdfg_module.build_folder_root() + unlaunched.setenv(rank_var, '1') + + assert sdfg_module.build_folder_root() != rank0 + + +@pytest.mark.parametrize('cache_mode', ['name', 'hash', 'unique', 'single']) +def test_every_cache_mode_builds_under_the_rank_root(unlaunched, cache_mode): + """Splitting the root rather than the SDFG name separates the ranks in every mode.""" + sdfg = dace.SDFG('rankprobe') + unlaunched.setenv('SLURM_PROCID', '3') + + with dace.config.set_temporary('cache', value=cache_mode): + unlaunched.setenv('DACE_cache_distaware', '1') + ranked = sdfg.build_folder + unlaunched.delenv('DACE_cache_distaware') + + assert sdfg.build_folder != ranked + + +def test_ranks_share_a_build_folder_unless_asked_otherwise(unlaunched): + """The default has to stay: distributed_compile has rank 0 build where every other rank looks.""" + sdfg = dace.SDFG('rankprobe') + + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '0') + rank0 = sdfg.build_folder + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '1') + + assert sdfg.build_folder == rank0 + + +def test_a_process_no_launcher_started_keeps_its_folder(unlaunched): + """No launcher is not rank 0: a lone process keeps the folder it always had, distaware or not.""" + unlaunched.setenv('DACE_cache_distaware', '1') + + assert sdfg_module.build_folder_root() == dace.Config.get('default_build_folder') + + +class OneRankOfAJob: + """Stands in for an mpi4py communicator, with the ranks taking their turn in this one process.""" + + def __init__(self, rank: int): + self.rank = rank + self.broadcast = None + + def Get_rank(self) -> int: + return self.rank + + def bcast(self, value, root: int = 0): + if self.rank == root: + self.broadcast = value + return self.broadcast + + def Barrier(self): + pass + + +def test_distributed_compile_puts_every_rank_in_rank_0_folder(unlaunched, tmp_path): + """Only rank 0 builds, so the others must look where it built and not where they would have.""" + unlaunched.setenv('DACE_default_build_folder', str(tmp_path)) + unlaunched.setenv('DACE_cache_distaware', '1') + + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '0') + builder = OneRankOfAJob(0) + csdfg = sdfg_utils.distributed_compile(customprog.to_sdfg(), builder) + del csdfg # Close the library, so the loading rank below opens it fresh + + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '1') + loader = OneRankOfAJob(1) + loader.broadcast = builder.broadcast + sdfg = customprog.to_sdfg() + assert sdfg.build_folder != builder.broadcast, "rank 1 was looking in rank 0's folder regardless" + + csdfg = sdfg_utils.distributed_compile(sdfg, loader) + + assert sdfg.build_folder == builder.broadcast + del csdfg + + # A rank that only loads is free to hold no SDFG at all, as tests/library/mpi does. + csdfg = sdfg_utils.distributed_compile(None, loader) + del csdfg + + if __name__ == '__main__': test_custom_build_folder() From 3a903d165b880b48b87dec58e7c03c19e8b27370 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 10 Aug 2026 14:56:28 +0200 Subject: [PATCH 02/11] Update default value for cache_distaware Changed default value of cache_distaware from false to true. --- dace/config_schema.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dace/config_schema.yml b/dace/config_schema.yml index 5baf488c08..83bcb9458d 100644 --- a/dace/config_schema.yml +++ b/dace/config_schema.yml @@ -740,7 +740,7 @@ required: cache_distaware: type: bool - default: false + default: true title: Distribution-aware build cache description: > Give every rank of a job its own build folder, named after the rank its launcher From 20cb2c793435efbe9d04df824f1e8fcdf26c1af8 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 10 Aug 2026 16:39:54 +0200 Subject: [PATCH 03/11] test: cover cache_distaware default flip in build folder tests Ranked-vs-shared assertions assumed distaware defaulted off, so clearing the env override fell through to the new true default and compared a rank-suffixed path against itself. Wrap the old off-path assertions in an explicit distaware=False context and add structural per-rank-root assertions for the new on-by-default behavior, for every cache mode. --- tests/custom_build_folder_test.py | 34 ++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/tests/custom_build_folder_test.py b/tests/custom_build_folder_test.py index 7015e250ef..8c5e5de8a6 100644 --- a/tests/custom_build_folder_test.py +++ b/tests/custom_build_folder_test.py @@ -51,28 +51,52 @@ def test_distaware_gives_each_rank_its_own_cache_root(unlaunched, rank_var): @pytest.mark.parametrize('cache_mode', ['name', 'hash', 'unique', 'single']) -def test_every_cache_mode_builds_under_the_rank_root(unlaunched, cache_mode): +def test_every_cache_mode_builds_under_the_rank_root(unlaunched, cache_mode, tmp_path): """Splitting the root rather than the SDFG name separates the ranks in every mode.""" + unlaunched.setenv('DACE_default_build_folder', str(tmp_path)) sdfg = dace.SDFG('rankprobe') unlaunched.setenv('SLURM_PROCID', '3') with dace.config.set_temporary('cache', value=cache_mode): + # distaware defaults on: every mode's root gets the rank suffix. unlaunched.setenv('DACE_cache_distaware', '1') ranked = sdfg.build_folder + assert os.path.dirname(ranked) == f'{tmp_path}_rank3' unlaunched.delenv('DACE_cache_distaware') - assert sdfg.build_folder != ranked + # Turning distaware off is how a caller opts back into the old shared root. + with dace.config.set_temporary('cache_distaware', value=False): + assert sdfg.build_folder != ranked + assert os.path.dirname(sdfg.build_folder) == str(tmp_path) + assert os.path.basename(sdfg.build_folder) == os.path.basename(ranked) + + +def test_ranks_share_a_build_folder_when_distaware_is_off(unlaunched): + """Turning distaware off restores the old default: distributed_compile has rank 0 build + where every other rank looks.""" + with dace.config.set_temporary('cache_distaware', value=False): + sdfg = dace.SDFG('rankprobe') + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '0') + rank0 = sdfg.build_folder + unlaunched.setenv('OMPI_COMM_WORLD_RANK', '1') -def test_ranks_share_a_build_folder_unless_asked_otherwise(unlaunched): - """The default has to stay: distributed_compile has rank 0 build where every other rank looks.""" + assert sdfg.build_folder == rank0 + + +def test_ranks_do_not_share_a_build_folder_by_default(unlaunched, tmp_path): + """distaware defaults on: ranks that each compile must not land in one folder.""" + unlaunched.setenv('DACE_default_build_folder', str(tmp_path)) sdfg = dace.SDFG('rankprobe') unlaunched.setenv('OMPI_COMM_WORLD_RANK', '0') rank0 = sdfg.build_folder unlaunched.setenv('OMPI_COMM_WORLD_RANK', '1') + rank1 = sdfg.build_folder - assert sdfg.build_folder == rank0 + assert rank0 == os.path.join(f'{tmp_path}_rank0', 'rankprobe') + assert rank1 == os.path.join(f'{tmp_path}_rank1', 'rankprobe') + assert rank0 != rank1 def test_a_process_no_launcher_started_keeps_its_folder(unlaunched): From c205df5115d9d8cabd1c77f06da305844316626f Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Mon, 10 Aug 2026 17:26:18 +0200 Subject: [PATCH 04/11] test: pin cache mode in the rank-default-split test The exact-leaf assertion assumed cache mode 'name'. A workflow whose DACE_cache resolves to anything else (env or a persisted config value the unlaunched fixture does not clear) flipped the leaf to a hash suffix and broke the path match. Pin it explicitly like the sibling tests pin their env, same env-wins-over-config precedence used to root-cause the distaware default flip. --- tests/custom_build_folder_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/custom_build_folder_test.py b/tests/custom_build_folder_test.py index 8c5e5de8a6..e7100f6428 100644 --- a/tests/custom_build_folder_test.py +++ b/tests/custom_build_folder_test.py @@ -87,6 +87,7 @@ def test_ranks_share_a_build_folder_when_distaware_is_off(unlaunched): def test_ranks_do_not_share_a_build_folder_by_default(unlaunched, tmp_path): """distaware defaults on: ranks that each compile must not land in one folder.""" unlaunched.setenv('DACE_default_build_folder', str(tmp_path)) + unlaunched.setenv('DACE_cache', 'name') # pin the leaf naming policy so the exact path holds sdfg = dace.SDFG('rankprobe') unlaunched.setenv('OMPI_COMM_WORLD_RANK', '0') From 90d3c7593a789a93abd043a1e2cba7425f435301 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 12:32:07 +0200 Subject: [PATCH 05/11] ci: pin xdist workers and MPI ranks to GPUs All workers/ranks see every GPU and pile CUDA contexts onto device 0, which flakes as invalid device ordinal (101) under -n 32 on cscs CI. --- conftest.py | 74 ++++++++++++++++++++++++ tests/gpu_worker_pinning_test.py | 96 ++++++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+) create mode 100644 conftest.py create mode 100644 tests/gpu_worker_pinning_test.py diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000000..669e2cc9f8 --- /dev/null +++ b/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() diff --git a/tests/gpu_worker_pinning_test.py b/tests/gpu_worker_pinning_test.py new file mode 100644 index 0000000000..9c59a3daff --- /dev/null +++ b/tests/gpu_worker_pinning_test.py @@ -0,0 +1,96 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Unit tests for the root conftest.py GPU worker/rank pinning logic. +""" +import importlib.util +import os +import pathlib +import types + +from dace.sdfg.sdfg import LAUNCHER_RANK_VARS + +CONFTEST_PATH = pathlib.Path(__file__).resolve().parent.parent / 'conftest.py' + + +def load_root_conftest() -> types.ModuleType: + spec = importlib.util.spec_from_file_location('dace_root_conftest_under_test', CONFTEST_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +root_conftest = load_root_conftest() + + +def clear_worker_env(monkeypatch) -> None: + monkeypatch.delenv('PYTEST_XDIST_WORKER', raising=False) + for var in LAUNCHER_RANK_VARS: + monkeypatch.delenv(var, raising=False) + monkeypatch.delenv('CUDA_VISIBLE_DEVICES', raising=False) + + +def test_pick_gpu_worker_device_preset_pool_selection(): + device = root_conftest.pick_gpu_worker_device('gw2', ['0', '1', '2', '3']) + assert device == '2' + + +def test_parse_worker_index_gw_prefix(): + assert root_conftest.parse_worker_index('gw7') == 7 + + +def test_parse_worker_index_fallback_no_digits(): + assert root_conftest.parse_worker_index('master') == 0 + assert root_conftest.parse_worker_index('gw') == 0 + + +def test_pick_gpu_worker_device_modulo_wrap(): + assert root_conftest.pick_gpu_worker_device('gw4', ['0', '1']) == '0' + assert root_conftest.pick_gpu_worker_device('gw5', ['0', '1']) == '1' + + +def test_resolve_worker_id_prefers_xdist_over_mpi_rank(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('PYTEST_XDIST_WORKER', 'gw1') + monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '5') + assert root_conftest.resolve_worker_id() == 'gw1' + + +def test_resolve_worker_id_falls_back_to_launcher_rank_var(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '3') + assert root_conftest.resolve_worker_id() == '3' + + +def test_resolve_worker_id_uses_first_set_rank_var_in_order(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('PMI_RANK', '9') + monkeypatch.setenv('SLURM_PROCID', '1') + assert root_conftest.resolve_worker_id() == '9' + + +def test_resolve_worker_id_empty_when_neither_xdist_nor_mpi(monkeypatch): + clear_worker_env(monkeypatch) + assert root_conftest.resolve_worker_id() == '' + + +def test_pin_worker_to_gpu_uses_preset_pool_and_rank_var(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '1') + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '2,3') + root_conftest.pin_worker_to_gpu() + assert os.environ['CUDA_VISIBLE_DEVICES'] == '3' + + +def test_pin_worker_to_gpu_noop_without_worker_id(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '0,1') + root_conftest.pin_worker_to_gpu() + assert os.environ['CUDA_VISIBLE_DEVICES'] == '0,1' + + +def test_pin_worker_to_gpu_noop_single_device_pool(monkeypatch): + clear_worker_env(monkeypatch) + monkeypatch.setenv('PYTEST_XDIST_WORKER', 'gw3') + monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '0') + root_conftest.pin_worker_to_gpu() + assert os.environ['CUDA_VISIBLE_DEVICES'] == '0' From a4fa937c338727a1bb0432b81ba931a73739943c Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 12:51:38 +0200 Subject: [PATCH 06/11] test: drop subsumed gpu pinning cases --- tests/gpu_worker_pinning_test.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/gpu_worker_pinning_test.py b/tests/gpu_worker_pinning_test.py index 9c59a3daff..bc99bc5017 100644 --- a/tests/gpu_worker_pinning_test.py +++ b/tests/gpu_worker_pinning_test.py @@ -29,11 +29,6 @@ def clear_worker_env(monkeypatch) -> None: monkeypatch.delenv('CUDA_VISIBLE_DEVICES', raising=False) -def test_pick_gpu_worker_device_preset_pool_selection(): - device = root_conftest.pick_gpu_worker_device('gw2', ['0', '1', '2', '3']) - assert device == '2' - - def test_parse_worker_index_gw_prefix(): assert root_conftest.parse_worker_index('gw7') == 7 @@ -55,12 +50,6 @@ def test_resolve_worker_id_prefers_xdist_over_mpi_rank(monkeypatch): assert root_conftest.resolve_worker_id() == 'gw1' -def test_resolve_worker_id_falls_back_to_launcher_rank_var(monkeypatch): - clear_worker_env(monkeypatch) - monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '3') - assert root_conftest.resolve_worker_id() == '3' - - def test_resolve_worker_id_uses_first_set_rank_var_in_order(monkeypatch): clear_worker_env(monkeypatch) monkeypatch.setenv('PMI_RANK', '9') @@ -68,11 +57,6 @@ def test_resolve_worker_id_uses_first_set_rank_var_in_order(monkeypatch): assert root_conftest.resolve_worker_id() == '9' -def test_resolve_worker_id_empty_when_neither_xdist_nor_mpi(monkeypatch): - clear_worker_env(monkeypatch) - assert root_conftest.resolve_worker_id() == '' - - def test_pin_worker_to_gpu_uses_preset_pool_and_rank_var(monkeypatch): clear_worker_env(monkeypatch) monkeypatch.setenv('OMPI_COMM_WORLD_RANK', '1') @@ -86,11 +70,3 @@ def test_pin_worker_to_gpu_noop_without_worker_id(monkeypatch): monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '0,1') root_conftest.pin_worker_to_gpu() assert os.environ['CUDA_VISIBLE_DEVICES'] == '0,1' - - -def test_pin_worker_to_gpu_noop_single_device_pool(monkeypatch): - clear_worker_env(monkeypatch) - monkeypatch.setenv('PYTEST_XDIST_WORKER', 'gw3') - monkeypatch.setenv('CUDA_VISIBLE_DEVICES', '0') - root_conftest.pin_worker_to_gpu() - assert os.environ['CUDA_VISIBLE_DEVICES'] == '0' From 9f71786be9691450fccda1698eb6c2390b6ac6f2 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 15:22:15 +0200 Subject: [PATCH 07/11] Emit host-side synchronization for GPU-to-host copies The synchronization site after a device-to-host copy was a bare pass claiming the destination emitted it, and the destination-side helper early-returns on any node that stream assignment stamped, so nothing was emitted at all. Pageable host memory hides this because the driver blocks on it anyway; pinned memory, HIP and graph capture do not. Emit the StreamSynchronize where the copy is issued, force it for any host-located destination regardless of the stamp, and only fall back to it for a cross-stream consumer that stream assignment left without an event. --- dace/codegen/targets/cpp.py | 10 +- dace/codegen/targets/cuda.py | 23 ++-- tests/codegen/gpu_d2h_host_sync_test.py | 137 ++++++++++++++++++++++++ 3 files changed, 160 insertions(+), 10 deletions(-) create mode 100644 tests/codegen/gpu_d2h_host_sync_test.py diff --git a/dace/codegen/targets/cpp.py b/dace/codegen/targets/cpp.py index 046de415f2..6b13b2c856 100644 --- a/dace/codegen/targets/cpp.py +++ b/dace/codegen/targets/cpp.py @@ -1421,10 +1421,14 @@ 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 + callsite_stream.write( + "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. 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/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() From d0f88aa72088da6bd9484f91ff960fd75e6c80f6 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 16:50:11 +0200 Subject: [PATCH 08/11] Bail out of SDFG init before allocating against a failed target __dace_init_cuda returns early, without constructing the gpu_context, when no GPU-capable device is present. The init function only checked __result after it had already run the environment init code, the SDFG-level init code and every persistent allocation -- and a persistent GPU array allocates through DACE_GPU_CHECK, which dereferences the gpu_context that was never constructed. Check __result right after the target initializers. The later check stays, because the code it guards can fail on its own. --- dace/codegen/targets/framecode.py | 9 +++ .../codegen/gpu_codegen_error_checks_test.py | 69 +++++++++++++++++++ 2 files changed, 78 insertions(+) create mode 100644 tests/codegen/gpu_codegen_error_checks_test.py 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/tests/codegen/gpu_codegen_error_checks_test.py b/tests/codegen/gpu_codegen_error_checks_test.py new file mode 100644 index 0000000000..8e92560b89 --- /dev/null +++ b/tests/codegen/gpu_codegen_error_checks_test.py @@ -0,0 +1,69 @@ +# 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 + +BAILOUT = r'if \(__result\)' + + +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 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') + + +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() From d23ee8227e9151e1787389306cc58c34db11fa46 Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 17:05:55 +0200 Subject: [PATCH 09/11] Check cross-stream event synchronization, and stop faking a missing event --- dace/codegen/targets/cpp.py | 66 ++++++++++++------- .../codegen/gpu_codegen_error_checks_test.py | 55 ++++++++++++++++ 2 files changed, 98 insertions(+), 23 deletions(-) diff --git a/dace/codegen/targets/cpp.py b/dace/codegen/targets/cpp.py index 6b13b2c856..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 @@ -1433,16 +1444,25 @@ def synchronize_streams(sdfg, cfg, dfg, state_id, node, scope_exit, callsite_str # 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, - ), - cfg, - state_id, - [e.src, e.dst], - ) + # 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/tests/codegen/gpu_codegen_error_checks_test.py b/tests/codegen/gpu_codegen_error_checks_test.py index 8e92560b89..99043e6b3c 100644 --- a/tests/codegen/gpu_codegen_error_checks_test.py +++ b/tests/codegen/gpu_codegen_error_checks_test.py @@ -9,6 +9,7 @@ from dace import dtypes BAILOUT = r'if \(__result\)' +EVENT_CALL = r'(?:cuda|hip)(?:EventRecord|StreamWaitEvent)\(' def generated_code(sdfg: dace.SDFG) -> str: @@ -43,6 +44,50 @@ def persistent_gpu_transient() -> dace.SDFG: 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 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') @@ -64,6 +109,16 @@ def test_the_init_function_still_checks_what_runs_after_the_allocations(): '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}' + + 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() From 3b1949d8bee8da689dfe2b54815672933e22185b Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 17:14:53 +0200 Subject: [PATCH 10/11] Check the status every emitted cuBLAS and rocBLAS call returns --- dace/libraries/blas/environments/cublas.py | 2 +- dace/libraries/blas/environments/rocblas.py | 2 +- dace/libraries/blas/nodes/batched_matmul.py | 13 +++--- dace/libraries/blas/nodes/dot.py | 8 ++-- dace/libraries/blas/nodes/gemm.py | 17 +++++--- dace/libraries/blas/nodes/gemv.py | 9 ++-- dace/libraries/linalg/nodes/transpose.py | 4 +- .../codegen/gpu_codegen_error_checks_test.py | 42 +++++++++++++++++++ 8 files changed, 73 insertions(+), 24 deletions(-) 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/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/tests/codegen/gpu_codegen_error_checks_test.py b/tests/codegen/gpu_codegen_error_checks_test.py index 99043e6b3c..caf72a6f10 100644 --- a/tests/codegen/gpu_codegen_error_checks_test.py +++ b/tests/codegen/gpu_codegen_error_checks_test.py @@ -7,9 +7,11 @@ 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: @@ -88,6 +90,24 @@ def producer(name: str, src: str, dst: str, code: str) -> None: 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') @@ -118,7 +138,29 @@ def test_cross_stream_event_synchronization_is_checked(): 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() From 6f2782855dd873402bd77da461e0bb411084cd5e Mon Sep 17 00:00:00 2001 From: Yakup Koray Budanaz Date: Tue, 11 Aug 2026 17:14:53 +0200 Subject: [PATCH 11/11] Release the internal stream array, and stop throwing out of BLAS handle teardown --- dace/libraries/blas/include/dace_cublas.h | 4 +++- dace/libraries/blas/include/dace_rocblas.h | 4 +++- dace/runtime/include/dace/cuda/cudacommon.cuh | 1 + 3 files changed, 7 insertions(+), 2 deletions(-) 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/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