diff --git a/pyproject.toml b/pyproject.toml index f7e4ae9d01..4e7aac28eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ dependencies = [ 'click>=8.0.0', 'cmake>=3.22', 'cytoolz>=1.0.1', - 'dace>=2.0.0a4', + 'dace==2.3.36', 'deepdiff>=8.1.0', 'devtools>=0.6', 'factory-boy>=3.3.3', @@ -479,6 +479,9 @@ url = 'https://gridtools.github.io/pypi/' # dace = {index = "gridtools"} [tool.uv.sources] atlas4py = {index = "test.pypi"} +dace = [ + {git = "https://github.com/philip-paul-mueller/dace", branch = "phimuell__new-gpu-codegen-dev"} +] # -- versioningit -- [tool.versioningit] diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py index aa34736c8a..34dc1fd633 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py @@ -1057,6 +1057,16 @@ def gt_gpu_apply_mempool(sdfg: dace.SDFG) -> None: Args: sdfg: The SDFG that should be processed. """ + + # TODO(phimuell): Reverse once the new codegen has caught up. + gpu_backend = dace.Config.get("compiler.cuda.backend") + if gpu_backend != "cuda": + warnings.warn( + f"GPU Memory-Pool is only implemented for `CUDA` and not for `{gpu_backend}`.", + stacklevel=0, + ) + return + for _, _, desc in sdfg.arrays_recursive(): if ( isinstance(desc, dace.data.Array) diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/common.py b/src/gt4py/next/program_processors/runners/dace/workflow/common.py index 738caffc13..62f4166d4f 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/common.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/common.py @@ -117,6 +117,28 @@ def set_dace_config( # This setting allows to throw an exception if any implicit Copy-Map slips thorugh. dace.Config.set("compiler.cuda.allow_implicit_memlet_to_map", value=False) + # Use the new GPU code generator + # NOTE: In the CI file we export the variable to force the experimental code gen to be used. + dace.Config.set("compiler.cuda.implementation", value="experimental") + + # Skip GPU Sync at the end. + # NOTE: That this will most likely break the UNIT tests, but should not be a problem + # for the blueline. + # NOTE: There are also functions that are supposed to achieve this. But they are + # currently useless, we have to work a bit more with them, but it is probably + # possible. + dace.Config.set("compiler", "cuda", "synchronize_on_exit", value=False) + + # NEW CODEGEN: + # In the old codegen we had to use the default stream. The effect was that the "same" + # stream was used across all SDFGs in ICON4Py. After switching to the new codegen, + # a single stream was used (there was a bug that made the default stream not used + # correctly [see](https://github.com/spcl/dace/pull/2408)). The net effect was that + # allocations got slower. We think this is because there are now multiple stream + # involved which makes it harder for the driver to schedule them. Thus, for good + # performance, we still have to use the default stream. But for other reasons as + # before. + # OLD CODEGEN: # In some stencils, for example `apply_diffusion_to_w`, the cuda codegen messes # up with the cuda streams, i.e. it allocates N streams but uses N+1. The first # idea was to use just one stream. However, even in that case the generator diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py index 8beecdd7e3..cc4fd43a17 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -106,6 +106,10 @@ def make_sdfg_call_async(sdfg: dace.SDFG, gpu: bool) -> None: Todo: Revisit this function once DaCe changes its behaviour in this regard. """ + # TODO(phimuell, edopao): Revisit this function after we understand the new + # code generator better. + return + # This is only a problem on GPU. # TODO(phimuell): Figuring out what about OpenMP. if not gpu: @@ -284,6 +288,10 @@ def make_sdfg_call_sync(sdfg: dace.SDFG, gpu: bool) -> None: work that runs on the GPU. Furthermore, all work is scheduled on the default stream. """ + # TODO(phimuell, edopao): Revisit this function after we understand the new + # code generator better. + return + if not gpu: # This is only a problem on GPU. Dace uses OpenMP on CPU and # the OpenMP parallel region creates a synchronization point. diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py index 08d6d4997b..ce70fd9f74 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_translation.py @@ -201,6 +201,7 @@ def _check_cpu_sdfg_call(sdfg: dace.SDFG) -> None: assert not _are_streams_synchronized(sdfg) +@pytest.mark.skip("To revisit after switch to new code gen.") @pytest.mark.parametrize( "make_async_sdfg_call", [False, True], @@ -242,6 +243,7 @@ def test_generate_sdfg_async_call(make_async_sdfg_call: bool, device_type: core_ _check_sdfg_without_async_call(sdfg) +@pytest.mark.skip("To revisit after switch to new code gen.") def test_generate_sdfg_async_call_no_map(device_type: core_defs.DeviceType): """Verify that the flag `async_sdfg_call=True` has no effect on an SDFG that does not contain any GPU map.""" @@ -367,6 +369,7 @@ def _make_multi_state_sdfg_3( return sdfg, first_state, second_state +@pytest.mark.skip("To revisit after switch to new code gen.") @pytest.mark.parametrize( "multi_state_config", [ diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_gpu_utils.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_gpu_utils.py index dd4927bb03..db3f247d22 100644 --- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_gpu_utils.py +++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_gpu_utils.py @@ -158,7 +158,6 @@ def test_set_gpu_properties(method: int): sdfg = dace.SDFG(util.unique_name("gpu_properties_test")) state = sdfg.add_state(is_start_block=True) - map_entries: dict[int, dace_nodes.MapEntry] = {} for dim in [1, 2, 3, 4]: shape = (10,) * dim sdfg.add_array( @@ -167,7 +166,7 @@ def test_set_gpu_properties(method: int): sdfg.add_array( f"B_{dim}", shape=shape, dtype=dace.float64, storage=dace.StorageType.GPU_Global ) - _, me, _ = state.add_mapped_tasklet( + state.add_mapped_tasklet( f"map_{dim}", map_ranges={f"__i{i}": f"0:{s}" for i, s in enumerate(shape)}, inputs={"__in": dace.Memlet(f"A_{dim}[{','.join(f'__i{i}' for i in range(dim))}]")}, @@ -175,7 +174,7 @@ def test_set_gpu_properties(method: int): outputs={"__out": dace.Memlet(f"B_{dim}[{','.join(f'__i{i}' for i in range(dim))}]")}, external_edges=True, ) - map_entries[dim] = me + del state sdfg.validate() if method == 0: @@ -203,6 +202,11 @@ def test_set_gpu_properties(method: int): else: raise ValueError(f"Unknown method {method}") + # Because of the inplace reconstruction all references to graph objects are destroyed. + map_entries: dict[int, dace_nodes.MapEntry] = {} + for node in sdfg.states()[0].nodes(): + if isinstance(node, dace_nodes.MapEntry): + map_entries[int(node.label[4])] = node map1, map2, map3, map4 = (map_entries[d].map for d in [1, 2, 3, 4]) # It takes the normal block size and does not regulate anything. @@ -258,6 +262,7 @@ def test_set_gpu_properties_1D(): map_entries[dim] = me sdfg.validate() + # `get_set_gpu_blocksize()` is non destructive, so `map_entries` are still pointing into the SDFG. sdfg.apply_gpu_transformations() gtx_dace_fieldview_gpu_utils.gt_set_gpu_blocksize( sdfg=sdfg, @@ -322,6 +327,7 @@ def test_set_gpu_properties_2D_3D(): map_entries[dim] = me sdfg.validate() + # `get_set_gpu_blocksize()` is non destructive, so `map_entries` are still pointing into the SDFG. sdfg.apply_gpu_transformations() gtx_dace_fieldview_gpu_utils.gt_set_gpu_blocksize( sdfg=sdfg, diff --git a/uv.lock b/uv.lock index f853aee738..8c6e2f1fa3 100644 --- a/uv.lock +++ b/uv.lock @@ -1206,8 +1206,8 @@ wheels = [ [[package]] name = "dace" -version = "2.0.0a4" -source = { registry = "https://pypi.org/simple" } +version = "2.3.36" +source = { git = "https://github.com/philip-paul-mueller/dace?branch=phimuell__new-gpu-codegen-dev#fe326334ff71cb4b65fec9860d9f95907caf1870" } dependencies = [ { name = "astunparse" }, { name = "dill" }, @@ -1223,7 +1223,6 @@ dependencies = [ { name = "sympy" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3a/07/a0410822c1aaef052c0015e7a8399a7d7da76a11e0a97939027c15d125ee/dace-2.0.0a4.tar.gz", hash = "sha256:f45f84b59a34846e0049347559416c1fac39c739aa6cbc791e27e1973d58b1ad", size = 6026693, upload-time = "2026-06-25T14:09:32.009Z" } [[package]] name = "debugpy" @@ -1859,7 +1858,7 @@ requires-dist = [ { name = "cupy-cuda13x", marker = "extra == 'cuda13'", specifier = ">=14.0" }, { name = "cupy-rocm-7-0", marker = "extra == 'rocm7'", specifier = ">=14.0" }, { name = "cytoolz", specifier = ">=1.0.1" }, - { name = "dace", specifier = ">=2.0.0a4" }, + { name = "dace", git = "https://github.com/philip-paul-mueller/dace?branch=phimuell__new-gpu-codegen-dev" }, { name = "deepdiff", specifier = ">=8.1.0" }, { name = "devtools", specifier = ">=0.6" }, { name = "factory-boy", specifier = ">=3.3.3" },