[rocr-runtime] Make the DXG topology snapshot safe for multiple consumers

.github/workflows/rocr-runtime-wsl.yml

@@ -121,10 +121,19 @@ jobs:
             exit 1
           fi
 
-          # Build the library
+          # Build the library, and with it the DXG unit tests: they are off by
+          # default so that a packaging build is unaffected, which leaves this
+          # job as the only place they ever get built and run.
           mkdir -p build
           cd build
-          cmake .. -G Ninja -DWIN_SDK="${win_sdk}/shared"
+          cmake .. -G Ninja -DWIN_SDK="${win_sdk}/shared" -DBUILD_ROCDXG_TESTS=ON
           ninja
 
           echo "libhsakmt build completed successfully!"
+
+      - name: Run DXG unit tests
+        shell: wsl-bash {0}
+        run: |
+          set -euo pipefail
+          cd ~/rocm-systems/projects/rocr-runtime/libhsakmt/build
+          ctest --output-on-failure

projects/rocr-runtime/libhsakmt/CMakeLists_wsl.txt

@@ -195,10 +195,43 @@ install ( TARGETS ${ROCDXG_TARGET} EXPORT ${ROCDXG_TARGET}Targets
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT binary
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT binary )
 
-# Public headers (hsakmt.h, hsakmttypes.h, rocdxg_smi.h) are already installed
-# by libhsakmt/core-runtime. librocdxg.pc reuses libhsakmt's includedir so no
-# separate header installation is needed here (and duplicating it breaks
-# TheRock artifact uniqueness checks — see ROCm/TheRock#7005).
+# Install public headers.
+#
+# This is the only configuration that builds librocdxg, and it is mutually
+# exclusive with the libhsakmt branch of the parent CMakeLists.txt - which is
+# what installs include/hsakmt wholesale - so without a rule here the WSL dev
+# package would ship a library, a pkg-config file naming an include directory,
+# and nothing in it. There is no second install rule for these files.
+#
+# hsakmt.h declares the entry points this library exports; hsakmttypes.h is
+# the transitive closure behind it, and is also the header rocprofiler-sdk
+# includes for HsaNodeProperties and HsaStructureSizes.
+# Neither needs anything but <stdbool.h> and <stdint.h>, and tests/dxg checks
+# that claim by compiling against a staged copy of exactly this list. The rest
+# of include/hsakmt describes the native thunk and is deliberately not shipped.
+#
+# TheRock is the one place a core-runtime built by the native branch and this
+# library sit side by side, and there its wsl-rocdxg_dev descriptor excludes
+# include/hsakmt/**. That exclusion, not the absence of a rule here, is what
+# keeps its cross-artifact overlap check quiet.
+set ( ROCDXG_PUBLIC_HEADERS
+  ${CMAKE_CURRENT_SOURCE_DIR}/include/hsakmt/hsakmt.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/include/hsakmt/hsakmttypes.h )
+
+install ( FILES ${ROCDXG_PUBLIC_HEADERS}
+  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/hsakmt COMPONENT dev )
+
+# Pure-logic unit tests. They compile a header out of src/dxg on its own, or a
+# staged copy of ROCDXG_PUBLIC_HEADERS, and link nothing; one more reads the
+# built library's dynamic symbol table. None needs a GPU or a DXCore adapter,
+# so `ctest` here works on a machine with neither. Off by default so a
+# packaging build is unaffected. Kept after the install rule so the header list
+# is already set.
+option ( BUILD_ROCDXG_TESTS "Build the librocdxg unit tests" OFF )
+if ( BUILD_ROCDXG_TESTS )
+  enable_testing()
+  add_subdirectory ( tests/dxg )
+endif()
 
 # Record our usage data for clients find_package calls.
 install ( EXPORT ${ROCDXG_TARGET}Targets
@@ -250,10 +283,9 @@ if(${EXPORT_TO_USER_PACKAGE_REGISTRY})
   export(PACKAGE ${ROCDXG_TARGET})
 endif()
 
-# Since librocdxg.pc and libhsakmt.pc are installed to the same pkgconfig directory,
-# we can directly use libhsakmt's header file path in the includedir.
-# This allows librocdxg to reference the same header files as libhsakmt without
-# duplicating header installation.
+# The includedir librocdxg.pc emits in its Cflags is CMAKE_INSTALL_INCLUDEDIR,
+# which is where the ROCDXG_PUBLIC_HEADERS rule above puts hsakmt/. libhsakmt
+# is not built in this configuration, so nothing else populates it.
 configure_file ( librocdxg.pc.in librocdxg.pc @ONLY )
 
 install ( FILES ${CMAKE_CURRENT_BINARY_DIR}/librocdxg.pc DESTINATION ${CMAKE_INSTALL_LIBDIR}/pkgconfig COMPONENT dev )

projects/rocr-runtime/libhsakmt/include/hsakmt/hsakmt.h

@@ -32,8 +32,14 @@
 extern "C" {
 #endif
 
-/* Forward declaration for debug trap ioctl arguments */
+/* Forward declarations for the KFD ioctl arguments the entry points below take
+ * by pointer. linux/kfd_ioctl.h defines them and neither package installs it,
+ * so without these a consumer's C compiler invents a type scoped to the
+ * prototype and the pointer it passes is a different one.
+ */
 struct kfd_ioctl_dbg_trap_args;
+struct kfd_runtime_info;
+struct kfd_dbg_device_info_entry;
 
 /**
   "Opens" the HSA kernel driver for user-kernel mode communication.

projects/rocr-runtime/libhsakmt/librocdxg.pc.in

@@ -1,4 +1,4 @@
-prefix=${pcfiledir}/../../..
+prefix=${pcfiledir}/../..
 exec_prefix=${prefix}
 libdir=${prefix}/@CMAKE_INSTALL_LIBDIR@
 includedir=${prefix}/@CMAKE_INSTALL_INCLUDEDIR@

projects/rocr-runtime/libhsakmt/src/dxg/debug.cpp

@@ -61,10 +61,15 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtDbgAddressWatch(
 HSAKMT_STATUS HSAKMTAPI hsaKmtCheckRuntimeDebugSupport(void) {
   CHECK_DXG_OPEN();
 
-  // Make sure the WDDMDevices are created.
-  HsaSystemProperties props;
-  hsaKmtAcquireSystemProperties(&props);
-  [](auto&&...) {}(props);
+  // The caller owns the topology snapshot the WDDMDevices belong to; this probe
+  // deliberately takes no reference of its own. Acquiring here would enumerate
+  // the adapters only to destroy them again on return, leaving the caller to
+  // build the very same list a second time.
+  //
+  // An empty device list therefore means no snapshot is held (or the machine
+  // has no adapter). Say so, rather than letting the loop below fall through
+  // and report support for devices nobody has looked at.
+  if (dxg_topology->wdevices_.empty()) return HSAKMT_STATUS_NOT_SUPPORTED;
 
   for (auto&& device : dxg_topology->wdevices_) {
     if (Wkmi::KmdDbgVersion version;
@@ -77,6 +82,12 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtCheckRuntimeDebugSupport(void) {
 }
 
 HSAKMT_STATUS HSAKMTAPI hsaKmtRuntimeEnable(void *rDebug, bool setupTtmp) {
+  /* No acquire here either: the caller holds the snapshot across this call.
+   * ROCr's KfdDriver takes it in Init() before enabling, precisely so the
+   * device list this walks is the one BuildTopology() goes on to use. The
+   * probe reports NOT_SUPPORTED when that list is empty, so a caller that has
+   * not acquired never reaches the loop below.
+   */
   HSAKMT_STATUS result = hsaKmtCheckRuntimeDebugSupport();
 
   if (result)

projects/rocr-runtime/libhsakmt/src/dxg/librocdxg.h

@@ -255,6 +255,10 @@ uint32_t get_direct_link_cpu(uint32_t gpu_node);
 HSAKMT_STATUS topology_sysfs_get_system_props(HsaSystemProperties& props);
 HSAKMT_STATUS topology_get_node_props(HSAuint32 NodeId,
 				      HsaNodeProperties *NodeProperties);
+void topology_clear_snapshot_refs(void);
+void topology_abandon_after_fork(void);
+bool topology_snapshot_is_live(void);
+void topology_drop_snapshot_at_last_close(void);
 HSAKMT_STATUS topology_get_iolink_props(HSAuint32 NodeId,
 					HSAuint32 NumIoLinks,
 					HsaIoLinkProperties *IoLinkProperties);

projects/rocr-runtime/libhsakmt/src/dxg/openclose.cpp

@@ -501,6 +501,16 @@ static void parent_fork_handler(void) { dxg_runtime->hsakmt_mutex.unlock(); }
 static void child_fork_handler(void) {
   dxg_runtime->is_forked = true;
 
+  /* Sever the references the child inherited but never took, before any public
+   * entry point can act on them. CHECK_DXG_OPEN() already rejects calls while
+   * is_forked is set, but zeroing the counts means that even a path that
+   * bypasses it cannot decrement the parent's bookkeeping. Both are plain
+   * stores, which is all an atfork child handler may safely do - the heavier
+   * teardown waits for clear_after_fork().
+   */
+  dxg_runtime->dxg_open_count = 0;
+  topology_clear_snapshot_refs();
+
   /* prepare_fork_handler() locked hsakmt_mutex right before fork() so that
    * no other thread would be mid-operation during the fork snapshot. In the
    * child only this one thread survives, and its TID differs from whichever
@@ -524,6 +534,11 @@ static void child_fork_handler(void) {
 static void clear_after_fork(void) {
   reset_suballocator();
   clear_allocation_map();
+  /* The snapshot's WDDMDevices belong to the parent's DXCore session, which
+   * the runtime reset below discards. Let go of them rather than reusing or
+   * destroying them.
+   */
+  topology_abandon_after_fork();
 
   if (dxg_runtime->dxg_fd >= 0) {
 #if defined(__linux__)
@@ -638,15 +653,22 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtOpenKFD(void) {
   HsaSystemProperties sys_props;
   char *error;
 
-  std::lock_guard<std::recursive_mutex> lck(dxg_runtime->hsakmt_mutex);
-
   /* If the process has forked, the child process must re-initialize
    * it's connection to DXG. Any references tracked by dxg_open_count
-   * belong to the parent
+   * belong to the parent.
+   *
+   * This runs before the lock is taken, and must: clear_after_fork() replaces
+   * *dxg_runtime, so a lock_guard constructed on the old object would unlock
+   * freed memory when it went out of scope. Only the forking thread survives
+   * into the child and child_fork_handler() has already reinitialized the
+   * mutex, so there is nothing here to exclude. is_forked is only ever written
+   * in a single-threaded child, so reading it unlocked races with nothing.
    */
   if (is_forked_child())
     clear_after_fork();
 
+  std::lock_guard<std::recursive_mutex> lck(dxg_runtime->hsakmt_mutex);
+
   if (dxg_runtime->dxg_open_count == 0) {
     static bool atfork_installed = false;
 
@@ -680,7 +702,16 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtOpenKFD(void) {
 
 
 
-    dxg_runtime->dxg_open_count = 1;
+    dxg_runtime->dxg_open_count++;
+
+    /* Only the 0->1 transition owns the suballocator. A later open comes from
+     * a second in-process consumer (on WSL, rocprofiler-sdk reading the KMT
+     * topology alongside the HSA runtime) and the first one's allocations are
+     * still live - resetting the fragment allocator here would throw away the
+     * bookkeeping that tracks them. clear_after_fork() above resets it on the
+     * one path where the inherited bookkeeping really is meaningless.
+     */
+    reset_suballocator();
 
     if (!atfork_installed) {
       /* Atfork handlers cannot be uninstalled and
@@ -699,12 +730,19 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtOpenKFD(void) {
     result = HSAKMT_STATUS_KERNEL_ALREADY_OPENED;
   }
 
-  reset_suballocator();
   return result;
 dxcore_loader_failed:
+  /* Drop the descriptor we just took and forget it. dxg_open_count is still 0,
+   * so nobody else owns it, and leaving dxg_fd pointing at a closed descriptor
+   * makes the next hsaKmtOpenKFD() skip the reopen above and hand out a stale
+   * fd. Note the DxcoreLoader itself latches its result behind a std::once_flag
+   * and so cannot be retried in this process regardless.
+   */
 #if defined(__linux__)
-  close(fd);
+  if (dxg_runtime->dxg_fd >= 0)
+    close(dxg_runtime->dxg_fd);
 #endif
+  dxg_runtime->dxg_fd = -1;
 open_failed:
 
   return result;
@@ -716,6 +754,11 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtCloseKFD(void) {
 
   if (dxg_runtime->dxg_open_count > 0) {
     if (--dxg_runtime->dxg_open_count == 0) {
+      /* Before DXCore goes, and so before the WDDMDevice objects a snapshot
+       * names become pointers into a dead session.
+       */
+      topology_drop_snapshot_at_last_close();
+
       dxg_runtime->HeapFini();
 #if defined(__linux__)
       close(dxg_runtime->dxg_fd);

projects/rocr-runtime/libhsakmt/src/dxg/snapshot_refcount.hpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace wsl {
+namespace thunk {
+
+/* Reference count for the shared KMT topology snapshot.
+ *
+ * More than one component in a process consumes the topology - ROCr, and on
+ * WSL rocprofiler-sdk as well - so hsaKmtAcquireSystemProperties() hands out
+ * references to one snapshot and only the last hsaKmtReleaseSystemProperties()
+ * tears it down.
+ *
+ * The transitions live here, apart from the WDDM objects they guard, so they
+ * can be exercised without a GPU, a DXCore adapter or an open thunk. The
+ * caller still owns the actual teardown; every method below only reports what
+ * the caller must do next. Serialization is the caller's job too: in the thunk
+ * every one of these runs under hsakmtRuntime::hsakmt_mutex.
+ *
+ * librocdxg and libhsa-runtime64 are both built from rocr-runtime and are
+ * expected to ship together, but nothing refuses a runtime out of some other
+ * package - ROCm 7.2.x, where sizeof(HsaNodeProperties) is 376 against this
+ * build's 396, loads against this thunk with no diagnostic. What keeps such a
+ * runtime from tearing down a snapshot another consumer still holds is this
+ * count, which rejects a release it never handed out.
+ */
+class SnapshotRefcount {
+ public:
+  enum class ReleaseAction {
+    /* No reference to give back. An unmatched release must be rejected
+     * rather than allowed to tear down a snapshot somebody else still owns.
+     */
+    kUnbalanced,
+    /* A reference went away but others remain; the snapshot stays. */
+    kKeepSnapshot,
+    /* That was the last reference; the caller drops the snapshot. */
+    kDropSnapshot,
+  };
+
+  uint32_t count() const { return refs_; }
+
+  /* A caller found a snapshot that is already built and took a reference on
+   * it.
+   */
+  void AddReference() { ++refs_; }
+
+  /* A caller built the snapshot and published it. The count is assigned, not
+   * incremented: nothing can hold a reference to a snapshot that did not
+   * exist a moment ago, and a build that failed part way never gets here.
+   */
+  void OnSnapshotPublished() { refs_ = 1; }
+
+  ReleaseAction Release() {
+    if (refs_ == 0) return ReleaseAction::kUnbalanced;
+    if (--refs_ == 0) return ReleaseAction::kDropSnapshot;
+    return ReleaseAction::kKeepSnapshot;
+  }
+
+  /* The snapshot is gone or was never ours: the last release dropped it, a
+   * failed build unwound it, or a fork child disowned what it inherited.
+   * Whatever the reason, no caller holds a reference to it any more.
+   */
+  void Clear() { refs_ = 0; }
+
+ private:
+  uint32_t refs_ = 0;
+};
+
+}  // namespace thunk
+}  // namespace wsl