fix(rccl): fix ASAN GPU target selection, and propagate ASAN to rocSHMEM

projects/rccl/CMakeLists.txt

@@ -199,8 +199,10 @@ endif()
 set(GPU_TARGETS "${DEFAULT_GPUS}" CACHE STRING "Target default GPUs if GPU_TARGETS is not defined.")
 
 
-# Modify GPU architectures for Address Sanitizer builds by appending "xnack+"
+# Adjust GPU architectures for Address Sanitizer builds; per-family rules are
+# in cmake/AsanGpuTargets.cmake.
 if (BUILD_ADDRESS_SANITIZER)
+  include(cmake/AsanGpuTargets.cmake)
   set(_suffix "asan")
   set(_clang_rt_names "libclang_rt.${_suffix}.so" "libclang_rt.${_suffix}-x86_64.so")
   foreach(_clang_rt_name ${_clang_rt_names})
@@ -220,18 +222,13 @@ if (BUILD_ADDRESS_SANITIZER)
   endif()
   message(STATUS "ASAN runtime directory: ${ASAN_RUNTIME_DIR}")
 
-  SET(amdgpu_targets "")
-  foreach(amdgpu_target IN LISTS GPU_TARGETS)
-    if(NOT amdgpu_target STREQUAL "")
-      string(FIND "${amdgpu_target}" ":xnack+" HAS_XNACK_SUFFIX)
-      if(HAS_XNACK_SUFFIX EQUAL -1)
-        list(APPEND amdgpu_targets "${amdgpu_target}:xnack+")
-      else()
-        list(APPEND amdgpu_targets "${amdgpu_target}")
-      endif()
-    endif()
-  endforeach()
-  SET(GPU_TARGETS "${amdgpu_targets}")
+  # Guarded so that "no targets at all" still reaches the curated diagnostic
+  # below. rocm_local_targets() reports failure as the literal "NOTFOUND",
+  # which if() treats as false along with the empty case.
+  if(GPU_TARGETS)
+    rccl_asan_adjust_gpu_targets(GPU_TARGETS ${GPU_TARGETS})
+    message(STATUS "ASAN: GPU targets = ${GPU_TARGETS}")
+  endif()
 endif()
 
 # Check if clang compiler can offload to GPU_TARGETS

projects/rccl/cmake/AsanGpuTargets.cmake

@@ -0,0 +1,63 @@
+# GPU target selection for device AddressSanitizer builds. Device ASAN needs
+# xnack, which gfx9 must request explicitly, gfx1250 has always on (and rejects
+# the suffix), and no other family has. Cases pinned in
+# tests/test_asan_gpu_targets.py; projects/rocshmem repeats the policy inline.
+
+# Parts where xnack is always on, so naming it is an error.
+set(_RCCL_ASAN_XNACK_ALWAYS_ON gfx1250 gfx1250-strict)
+
+# rccl_asan_adjust_gpu_targets(OUT_VAR gfx942 gfx1100 gfx1250)
+#   -> OUT_VAR = gfx942:xnack+;gfx1250
+# Uninstrumentable targets are dropped rather than left to produce an
+# uninstrumented image inside an ASAN build.
+function(rccl_asan_adjust_gpu_targets OUT_VAR)
+  set(_result "")
+  set(_dropped "")
+
+  foreach(_target IN LISTS ARGN)
+    if(_target STREQUAL "")
+      continue()
+    endif()
+
+    # A requested suffix says what the caller wants, not what the part
+    # supports, so classify on the bare name.
+    string(REGEX REPLACE ":.*" "" _base "${_target}")
+    # list(FIND) not IN_LIST: the test loads this via cmake -P, where CMP0057
+    # is unset.
+    list(FIND _RCCL_ASAN_XNACK_ALWAYS_ON "${_base}" _always_on)
+
+    # Drop any xnack term the caller wrote before deciding what to do with it:
+    # it may be ":xnack-", which contradicts the build. Other feature terms
+    # (e.g. ":sramecc+") are part of the requested target and are kept; xnack
+    # sorts last, so re-appending it keeps the target ID canonical.
+    string(REGEX REPLACE ":xnack[+-]" "" _without_xnack "${_target}")
+
+    if(_base MATCHES "^gfx9")
+      list(APPEND _result "${_without_xnack}:xnack+")
+    elseif(NOT _always_on EQUAL -1)
+      # Always on here, so naming xnack at all is an invalid target ID.
+      list(APPEND _result "${_without_xnack}")
+    else()
+      list(APPEND _dropped "${_target}")
+    endif()
+  endforeach()
+
+  # Both xnack spellings of one arch normalise together, and a duplicate
+  # becomes a second add_library() in DeviceLinker.cmake.
+  if(_result)
+    list(REMOVE_DUPLICATES _result)
+  endif()
+
+  if(_dropped)
+    message(STATUS
+      "ASAN: dropping GPU target(s) with no xnack support: ${_dropped}")
+  endif()
+  if(NOT _result)
+    message(FATAL_ERROR
+      "ASAN requires a GPU target that supports xnack (gfx9* with :xnack+, "
+      "or gfx1250 where it is always on). None of the requested GPU targets "
+      "do: ${ARGN}")
+  endif()
+
+  set(${OUT_VAR} "${_result}" PARENT_SCOPE)
+endfunction()

projects/rccl/cmake/ROCSHMEM.cmake

@@ -97,6 +97,16 @@ function(add_rocshmem_targets)
         endif()
         find_package(rocshmem_static)
         if(rocshmem_static_FOUND)
+            # This path returns before the -DASAN=ON propagation below, so an
+            # ASAN rccl links whatever the pre-built tree happens to contain.
+            if(BUILD_ADDRESS_SANITIZER)
+                message(WARNING
+                    "ASAN build using pre-built rocSHMEM at ${ROCSHMEM_INSTALL_DIR}. "
+                    "It is linked into librccl.so as-is; if it was not itself built "
+                    "with -DASAN=ON, rccl gets an uninstrumented rocSHMEM inside an "
+                    "instrumented library. Unset ROCSHMEM_INSTALL_DIR to build it "
+                    "from source with ASAN propagated.")
+            endif()
             set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_INCLUDE_DIR}" PARENT_SCOPE)
             set(ROCSHMEM_LIBRARY     "${ROCSHMEM_LIBRARY}"      PARENT_SCOPE)
             set(ROCSHMEM_SOURCE_DIR  "${ROCSHMEM_SOURCE_DIR}"   CACHE INTERNAL "rocSHMEM source directory")
@@ -124,6 +134,12 @@ function(add_rocshmem_targets)
         if(ENABLE_ROCSHMEM_GIN)
             set(_rocshmem_sdma_opt "-DUSE_SDMA=ON")
         endif()
+        # librocshmem.a is linked into librccl.so, so without this an ASAN
+        # build gets an uninstrumented rocSHMEM inside an instrumented library.
+        set(_rocshmem_asan_opt "")
+        if(BUILD_ADDRESS_SANITIZER)
+            set(_rocshmem_asan_opt "-DASAN=ON")
+        endif()
         message(STATUS "rocSHMEM: building from ${ROCSHMEM_SOURCE_DIR}")
 
         ExternalProject_Add(rocshmem_ext
@@ -141,10 +157,18 @@ function(add_rocshmem_targets)
             CONFIGURE_COMMAND   ""
             BUILD_COMMAND
                 ${CMAKE_COMMAND} -E make_directory build
-                && ${CMAKE_COMMAND} -E chdir build bash -lc "INSTALL_PREFIX=${ROCSHMEM_INSTALL_DIR} ../scripts/build_configs/gda -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} -DUSE_EXTERNAL_MPI=OFF -DGDA_MLX5=ON -DGDA_BNXT=ON -DGDA_IONIC=ON -DBUILD_EXAMPLES=OFF -DBUILD_FUNCTIONAL_TESTS=OFF -DBUILD_UNIT_TESTS=OFF -DBUILD_CTESTS=OFF -DBUILD_TOOLS=OFF -DGPU_TARGETS=${_rocshmem_gpu_targets} ${_rocshmem_sdma_opt} ${_rocshmem_cmake_opts} "
+                && ${CMAKE_COMMAND} -E chdir build bash -lc "INSTALL_PREFIX=${ROCSHMEM_INSTALL_DIR} ../scripts/build_configs/gda -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} -DUSE_EXTERNAL_MPI=OFF -DGDA_MLX5=ON -DGDA_BNXT=ON -DGDA_IONIC=ON -DBUILD_EXAMPLES=OFF -DBUILD_FUNCTIONAL_TESTS=OFF -DBUILD_UNIT_TESTS=OFF -DBUILD_CTESTS=OFF -DBUILD_TOOLS=OFF -DGPU_TARGETS=${_rocshmem_gpu_targets} ${_rocshmem_sdma_opt} ${_rocshmem_asan_opt} ${_rocshmem_cmake_opts} "
             INSTALL_COMMAND ""
         )
 
+        # Rebuild when the sub-build's options change (e.g. ASAN): its build
+        # stamp ignores BUILD_COMMAND. The stamp sits beside the shared install,
+        # so the other config tree overwriting librocshmem.a invalidates it too.
+        set(_rocshmem_opts_stamp "${ROCSHMEM_INSTALL_DIR}/rccl_build_options.txt")
+        file(GENERATE OUTPUT "${_rocshmem_opts_stamp}"
+             CONTENT "${CMAKE_BUILD_TYPE} ${_rocshmem_gpu_targets} ${_rocshmem_sdma_opt} ${_rocshmem_asan_opt} ${_rocshmem_cmake_opts}\n")
+        ExternalProject_Add_StepDependencies(rocshmem_ext build "${_rocshmem_opts_stamp}")
+
         set(ROCSHMEM_INSTALL_DIR "${ROCSHMEM_INSTALL_DIR}"          PARENT_SCOPE)
         set(ROCSHMEM_INCLUDE_DIR "${ROCSHMEM_INSTALL_DIR}/include"  PARENT_SCOPE)
         set(ROCSHMEM_LIBRARY     "${ROCSHMEM_INSTALL_DIR}/lib/librocshmem.a" PARENT_SCOPE)

projects/rccl/tools/scripts/test_runner/tests/test_asan_gpu_targets.py

@@ -0,0 +1,159 @@
+"""Guards the ASAN GPU-target policy in cmake/AsanGpuTargets.cmake.
+
+Each case pins one family's observed clang behaviour, so reclassifying a
+family fails here rather than producing an ASAN build with a silently
+uninstrumented device image. Driven through `cmake -P`, like
+test_device_coverage_cmake.py.
+"""
+
+import shutil
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+
+
+RCCL_ROOT = Path(__file__).resolve().parents[4]
+MODULE = RCCL_ROOT / "cmake" / "AsanGpuTargets.cmake"
+# The same policy, duplicated because rccl builds rocSHMEM as a separate cmake
+# process. Drift between the two shows up only as librocshmem.a built for a
+# different arch set than the librccl.so it links into, so run the cases
+# against both.
+ROCSHMEM_MODULE = (RCCL_ROOT.parent / "rocshmem" / "cmake" / "AsanGpuTargets.cmake")
+
+
+@unittest.skipUnless(shutil.which("cmake"), "cmake not available on PATH")
+class AsanGpuTargetsTest(unittest.TestCase):
+    module = MODULE
+    function = "rccl_asan_adjust_gpu_targets"
+
+    def adjust(self, targets):
+        """Run rccl_asan_adjust_gpu_targets over the given targets.
+
+        Returns (result_list, status_messages). Raises on FATAL_ERROR.
+        """
+        quoted = " ".join(f'"{t}"' for t in targets)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            out_file = Path(temp_dir) / "out.txt"
+            script = Path(temp_dir) / "probe.cmake"
+            script.write_text(
+                f'include("{self.module}")\n'
+                f"{self.function}(RESULT {quoted})\n"
+                f'file(WRITE "{out_file}" "${{RESULT}}")\n'
+            )
+            completed = subprocess.run(
+                ["cmake", "-P", str(script)],
+                capture_output=True, text=True)
+            if completed.returncode != 0:
+                raise AssertionError(completed.stderr.strip())
+            result = out_file.read_text()
+        # cmake -P sends STATUS to stdout and WARNING/FATAL_ERROR to stderr.
+        # Both, so the tests below do not pin the policy to a message level.
+        return ([t for t in result.split(";") if t],
+                completed.stdout + completed.stderr)
+
+    def test_gfx9_gets_xnack_plus_requested(self):
+        """A bare gfx9 name compiles with the sanitizer ignored."""
+        result, _ = self.adjust(["gfx90a", "gfx942", "gfx950"])
+
+        self.assertEqual(result,
+                         ["gfx90a:xnack+", "gfx942:xnack+", "gfx950:xnack+"])
+
+    def test_gfx9_already_requesting_xnack_plus_is_unchanged(self):
+        result, _ = self.adjust(["gfx942:xnack+"])
+
+        self.assertEqual(result, ["gfx942:xnack+"])
+
+    def test_gfx9_requesting_xnack_minus_is_corrected(self):
+        """The old substring test for ":xnack+" produced gfx942:xnack-:xnack+."""
+        result, _ = self.adjust(["gfx942:xnack-"])
+
+        self.assertEqual(result, ["gfx942:xnack+"])
+
+    def test_gfx9_keeps_its_other_target_features(self):
+        """Only the xnack term is ours to rewrite; sramecc+ is what was asked for."""
+        result, _ = self.adjust(["gfx942:sramecc+:xnack-"])
+
+        self.assertEqual(result, ["gfx942:sramecc+:xnack+"])
+
+    def test_both_xnack_variants_of_one_arch_collapse_to_one_entry(self):
+        """A duplicate becomes a second add_library() in DeviceLinker.cmake."""
+        result, _ = self.adjust(["gfx942:xnack+", "gfx942:xnack-"])
+
+        self.assertEqual(result, ["gfx942:xnack+"])
+
+    def test_gfx1250_is_kept_untouched(self):
+        result, _ = self.adjust(["gfx1250"])
+
+        self.assertEqual(result, ["gfx1250"])
+
+    def test_gfx1250_strict_is_kept_untouched(self):
+        """The other entry in _RCCL_ASAN_XNACK_ALWAYS_ON, and the one no other
+        case reaches: drop it from that list and every test here stays green
+        while the target silently moves to the drop arm."""
+        result, _ = self.adjust(["gfx1250-strict", "gfx1250-strict:xnack+"])
+
+        self.assertEqual(result, ["gfx1250-strict"])
+
+    def test_gfx1250_loses_an_explicitly_requested_xnack(self):
+        """clang rejects 'gfx1250:xnack+' outright: xnack is not selectable there."""
+        result, _ = self.adjust(["gfx1250:xnack+"])
+
+        self.assertEqual(result, ["gfx1250"])
+
+    def test_gfx1250_with_and_without_xnack_collapse_to_one_entry(self):
+        result, _ = self.adjust(["gfx1250", "gfx1250:xnack+"])
+
+        self.assertEqual(result, ["gfx1250"])
+
+    def test_archs_without_xnack_are_dropped(self):
+        result, messages = self.adjust(["gfx942", "gfx1100", "gfx1201"])
+
+        self.assertEqual(result, ["gfx942:xnack+"])
+        self.assertIn("gfx1100", messages)
+        self.assertIn("gfx1201", messages)
+
+    def test_full_default_target_list(self):
+        """Matches the old effective set, except gfx1250 is no longer lost.
+
+        The input is a copy of DEFAULT_GPUS in the root CMakeLists.txt and has
+        to be updated by hand when a family is added there.
+        """
+        result, _ = self.adjust([
+            "gfx906", "gfx908", "gfx90a", "gfx942", "gfx950",
+            "gfx1030", "gfx1100", "gfx1101", "gfx1102", "gfx1151",
+            "gfx1200", "gfx1201", "gfx1250"])
+
+        self.assertEqual(result, [
+            "gfx906:xnack+", "gfx908:xnack+", "gfx90a:xnack+",
+            "gfx942:xnack+", "gfx950:xnack+", "gfx1250"])
+
+    def test_empty_result_is_fatal(self):
+        with self.assertRaises(AssertionError) as cm:
+            self.adjust(["gfx1100"])
+
+        self.assertIn("ASAN requires a GPU target that supports xnack",
+                      str(cm.exception))
+
+
+@unittest.skipUnless(shutil.which("cmake"), "cmake not available on PATH")
+@unittest.skipUnless(ROCSHMEM_MODULE.exists(), "projects/rocshmem not in this checkout")
+class RocshmemAsanGpuTargetsTest(AsanGpuTargetsTest):
+    """Every case above, run against the rocSHMEM copy of the policy."""
+
+    module = ROCSHMEM_MODULE
+    function = "rocshmem_asan_adjust_gpu_targets"
+
+    def test_full_default_target_list(self):
+        """rocSHMEM's DEFAULT_GPUS differs from rccl's; see its CMakeLists.txt."""
+        result, _ = self.adjust([
+            "gfx90a", "gfx1100", "gfx1201", "gfx942", "gfx950",
+            "gfx1250", "gfx1250-strict"])
+
+        self.assertEqual(result, [
+            "gfx90a:xnack+", "gfx942:xnack+", "gfx950:xnack+",
+            "gfx1250", "gfx1250-strict"])
+
+
+if __name__ == "__main__":
+    unittest.main()