perf[next-dace]: Added automatic setting of chiplet number using amdsmi

pyproject.toml

@@ -467,6 +467,11 @@ url = 'https://gridtools.github.io/pypi/'
 # Add the uv source below to pull dace from the gridtools index instead of PyPI:
 [tool.uv.sources]
 atlas4py = {index = "test.pypi"}
+# `compiler.cuda.chiplet_number`, used to distribute thread-blocks over the chiplets of
+# multi-chiplet AMD GPUs, only exists on this branch and in no released DaCe. The branch
+# reports the same version as the PyPI release, so a version specifier cannot express the
+# requirement. Drop this source once the feature lands in a DaCe release on PyPI.
+dace = {git = "https://github.com/GridTools/dace", branch = "chiplet_old_codegen_updateddace"}
 
 # -- versioningit --
 [tool.versioningit]

src/gt4py/next/program_processors/runners/dace/workflow/common.py

@@ -7,7 +7,9 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import contextlib
+import functools
 import os
+import warnings
 from typing import Any, Final, Generator, Optional, TypeAlias
 
 import dace
@@ -45,6 +47,72 @@
 """
 
 
+_CHIPLET_NUMBER_ENV_VAR: Final[str] = "DACE_compiler_cuda_chiplet_number"
+"""Environment variable used by DaCe to override `compiler.cuda.chiplet_number`."""
+
+
+@functools.cache
+def _query_gpu_chiplet_count() -> Optional[int]:
+    """Number of chiplets (XCDs) of the current AMD GPU, or None if it cannot be determined.
+
+    Cached: the device does not change within a process, and `amdsmi` has to be initialized
+    and shut down around the query, which must happen exactly once. The warning on failure
+    is emitted from here so that it is emitted only once per process, too.
+
+    Note:
+        `amdsmi` is not a GT4Py dependency, it ships with ROCm. Importing it loads
+        `libamd_smi.so` through `ctypes`, which fails on machines without ROCm and not
+        necessarily with an `ImportError`, hence the broad exception handling.
+
+    Note:
+        `amdsmi` ignores `ROCR_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` and enumerates all
+        physical GPUs, so the first processor handle is not necessarily the device that HIP
+        uses. This is only accurate on nodes where all GPUs are of the same model.
+    """
+    try:
+        import amdsmi
+
+        amdsmi.amdsmi_init()
+        try:
+            processor_handles = amdsmi.amdsmi_get_processor_handles()
+            if not processor_handles:
+                raise RuntimeError("`amdsmi` did not report any GPU.")
+            return int(amdsmi.amdsmi_get_gpu_xcd_counter(processor_handles[0]))
+        finally:
+            amdsmi.amdsmi_shut_down()
+    except Exception as e:
+        warnings.warn(
+            UserWarning(
+                f"Could not determine the number of GPU chiplets through `amdsmi`: {e}. "
+                "The distribution of thread-blocks over chiplets is disabled."
+            ),
+            stacklevel=2,
+        )
+        return None
+
+
+def _set_gpu_chiplet_number() -> None:
+    """Configure DaCe to distribute the thread-blocks of a kernel over the GPU chiplets.
+
+    Only meaningful on multi-chiplet AMD GPUs, therefore this must only be called when the
+    DaCe GPU backend is `hip`. Leaves the DaCe configuration untouched when the number of
+    chiplets is unknown, which disables the distribution.
+    """
+    env_value = os.environ.get(_CHIPLET_NUMBER_ENV_VAR, "").strip()
+    if env_value:
+        # `dace.Config.get()` gives precedence to this environment variable on its own, but a
+        #  value that only lives in the environment does not show up in `dace.Config.nondefaults()`
+        #  and would therefore not invalidate the GT4Py build cache, see `compilation.py`.
+        chiplet_count = int(env_value)
+    else:
+        queried_count = _query_gpu_chiplet_count()
+        if queried_count is None:
+            return
+        chiplet_count = queried_count
+
+    dace.Config.set("compiler.cuda.chiplet_number", value=chiplet_count)
+
+
 def set_dace_config(
     device_type: core_defs.DeviceType,
     cmake_build_type: Optional[gtx_config.CMakeBuildType] = None,
@@ -149,6 +217,9 @@ def set_dace_config(
     # This assumes that a process will only use one type of GPU.
     if device_type == core_defs.DeviceType.ROCM:
         dace.Config.set("compiler.cuda.backend", value="hip")
+        # Distributing the grid over chiplets relies on the round-robin thread-block scheduling
+        #  of multi-chiplet AMD GPUs, so it is only configured for the `hip` backend.
+        _set_gpu_chiplet_number()
     elif device_type == core_defs.DeviceType.CUDA:
         dace.Config.set("compiler.cuda.backend", value="cuda")
 

tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_dace_workflow_config.py

@@ -0,0 +1,121 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for the DaCe configuration set up by the dace backend workflow.
+
+Covers the detection of the number of GPU chiplets (XCDs), which no CI machine has the
+hardware for, so `amdsmi` is faked throughout.
+"""
+
+import sys
+import types
+from typing import Any, Optional
+
+import dace
+import pytest
+
+from gt4py._core import definitions as core_defs
+from gt4py.next.program_processors.runners.dace.workflow import common as dace_wf_common
+
+
+_CHIPLET_KEY = "compiler.cuda.chiplet_number"
+
+
+@pytest.fixture(autouse=True)
+def clear_chiplet_query_cache():
+    """The chiplet query is cached per process, which would leak between tests."""
+    dace_wf_common._query_gpu_chiplet_count.cache_clear()
+    yield
+    dace_wf_common._query_gpu_chiplet_count.cache_clear()
+
+
+def make_fake_amdsmi(xcd_count: Optional[int] = None, handles: Any = ("handle",)):
+    """Stand-in for the `amdsmi` module, recording whether it was queried."""
+    calls: list[str] = []
+
+    fake = types.ModuleType("amdsmi")
+    fake.amdsmi_init = lambda *args: calls.append("init")
+    fake.amdsmi_shut_down = lambda *args: calls.append("shut_down")
+    fake.amdsmi_get_processor_handles = lambda: list(handles)
+
+    def get_xcd_counter(handle):
+        calls.append("xcd_counter")
+        assert xcd_count is not None
+        return xcd_count
+
+    fake.amdsmi_get_gpu_xcd_counter = get_xcd_counter
+    fake.calls = calls
+    return fake
+
+
+def chiplet_nondefault() -> Optional[int]:
+    """The chiplet number as it reaches the GT4Py build-cache fingerprint, or None."""
+    nondefaults = dace.Config._data.nondefaults()
+    return nondefaults.get("compiler", {}).get("cuda", {}).get("chiplet_number")
+
+
+def test_chiplet_number_config_key_exists():
+    """Guards the DaCe source pin: setting a key DaCe does not know is a silent no-op."""
+    try:
+        metadata = dace.Config.get_metadata("compiler", "cuda", "chiplet_number")
+    except KeyError as e:
+        raise AssertionError(
+            "The installed DaCe has no `compiler.cuda.chiplet_number`, so GT4Py would be "
+            "setting a configuration entry that nothing reads. GT4Py pins DaCe to a branch "
+            "that provides it, see `[tool.uv.sources]` in `pyproject.toml`."
+        ) from e
+    assert metadata["type"] == "int"
+
+
+def test_chiplet_number_queried_from_amdsmi(monkeypatch):
+    monkeypatch.delenv(dace_wf_common._CHIPLET_NUMBER_ENV_VAR, raising=False)
+    fake = make_fake_amdsmi(xcd_count=6)
+    monkeypatch.setitem(sys.modules, "amdsmi", fake)
+
+    with dace_wf_common.dace_context(device_type=core_defs.DeviceType.ROCM):
+        assert dace.Config.get(_CHIPLET_KEY) == 6
+        assert chiplet_nondefault() == 6
+
+    assert fake.calls == ["init", "xcd_counter", "shut_down"]
+
+
+def test_chiplet_number_env_override_reaches_fingerprint(monkeypatch):
+    """`Config.get` honours the env var on its own, but `nondefaults` would not see it."""
+    monkeypatch.setenv(dace_wf_common._CHIPLET_NUMBER_ENV_VAR, "4")
+    fake = make_fake_amdsmi(xcd_count=6)
+    monkeypatch.setitem(sys.modules, "amdsmi", fake)
+
+    with dace_wf_common.dace_context(device_type=core_defs.DeviceType.ROCM):
+        assert chiplet_nondefault() == 4
+
+    # The environment takes precedence, so the device is never queried.
+    assert fake.calls == []
+
+
+def test_chiplet_number_unset_without_amdsmi(monkeypatch):
+    monkeypatch.delenv(dace_wf_common._CHIPLET_NUMBER_ENV_VAR, raising=False)
+    # A `None` entry in `sys.modules` makes the import fail.
+    monkeypatch.setitem(sys.modules, "amdsmi", None)
+
+    with pytest.warns(UserWarning, match="chiplets"):
+        with dace_wf_common.dace_context(device_type=core_defs.DeviceType.ROCM):
+            # Left at the DaCe default, which disables the distribution.
+            assert dace.Config.get(_CHIPLET_KEY) == 1
+            assert chiplet_nondefault() is None
+
+
+def test_chiplet_number_not_configured_on_cuda(monkeypatch):
+    monkeypatch.delenv(dace_wf_common._CHIPLET_NUMBER_ENV_VAR, raising=False)
+    fake = make_fake_amdsmi(xcd_count=6)
+    monkeypatch.setitem(sys.modules, "amdsmi", fake)
+
+    with dace_wf_common.dace_context(device_type=core_defs.DeviceType.CUDA):
+        assert dace.Config.get(_CHIPLET_KEY) == 1
+        assert chiplet_nondefault() is None
+
+    assert fake.calls == []