fix(rocjitsu): DBT host selection on multi-GPU systems

emulation/rocjitsu/docs/rocjitsu_dbt_guest.md

@@ -26,6 +26,16 @@ running on A0 host) is rejected unless both revisions are set, because the
 required A0 workarounds cannot otherwise be selected. For every other target the
 machine ID already identifies the silicon, so these fields stay `unspecified`.
 
+For either execution backend, `host_gpu_id: 0` selects the first `host_isa`
+match in the effective ROCm-visible device order for that backend's topology.
+The launcher applies
+`ROCR_VISIBLE_DEVICES` first, including numeric reordering and UUID prefixes,
+then `HIP_VISIBLE_DEVICES` or its `CUDA_VISIBLE_DEVICES` fallback. It atomically
+publishes the config path and resolved KFD GPU ID in the per-invocation handoff,
+so the KFD interposer and HSA hook consume one choice without a private
+environment variable. Set a nonzero `host_gpu_id` to pin a particular GPU; the
+launcher rejects the pin if standard ROCm visibility settings hide that device.
+
 ## High Level Flow
 
 User:
@@ -61,9 +71,15 @@ Underneath:
    keeps applications that choose the first GPU on the guest path while leaving
    the host agent alive for execution.
 
-   If `ROCR_VISIBLE_DEVICES` is set, the launcher expands it so the selected
-   host ordinal and appended guest ordinal both remain visible to ROCR before
-   the HSA hook applies this public shadowing.
+   If `ROCR_VISIBLE_DEVICES` is set, the launcher preserves its selected order
+   and expands it with the appended guest ordinal so both sides remain visible
+   internally before the HSA hook applies public shadowing. Because CLR applies
+   `HIP_VISIBLE_DEVICES` or `CUDA_VISIBLE_DEVICES` after HSA agent iteration,
+   the launcher normalizes an existing client selector to equivalent numeric
+   ordinals after ROCR filtering. This keeps host UUID selection meaningful
+   after the host's public identity is replaced by the guest. OpenCL's
+   `GPU_DEVICE_ORDINAL` remains owned by the client runtime and is not used to
+   choose the DBT host.
 
 2. Guest-facing discovery stays guest-shaped.
 

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/dbt_guest_config.cpp

@@ -10,6 +10,7 @@
 #include "flatbuffers/idl.h"
 #include "simulation_config_generated.h"
 
+#include <charconv>
 #include <cstdlib>
 #include <filesystem>
 #include <fstream>
@@ -150,6 +151,39 @@ DbtGuestConfig load_dbt_guest_config_from_file(const std::string &path) {
       false);
 }
 
+void apply_resolved_dbt_host_gpu_id(DbtGuestConfig &config, std::string_view value,
+                                    std::string_view source) {
+  if (!config.enabled || config.host.gpu_id != 0)
+    return;
+
+  uint32_t gpu_id = 0;
+  const char *begin = value.data();
+  const char *end = begin + value.size();
+  auto [ptr, error] = std::from_chars(begin, end, gpu_id);
+  if (error != std::errc{} || ptr != end || gpu_id == 0)
+    throw std::runtime_error(std::string(source) + " must contain a nonzero KFD gpu_id");
+  config.host.gpu_id = gpu_id;
+}
+
+std::optional<DbtRuntimeConfigHandoff> parse_dbt_runtime_config_handoff(std::string_view contents) {
+  const size_t first_newline = contents.find('\n');
+  std::string_view config_path = contents.substr(0, first_newline);
+  if (!config_path.empty() && config_path.back() == '\r')
+    config_path.remove_suffix(1);
+  if (config_path.empty())
+    return std::nullopt;
+
+  std::optional<std::string> resolved_gpu_id;
+  if (first_newline != std::string_view::npos) {
+    std::string_view second_line = contents.substr(first_newline + 1);
+    if (!second_line.empty()) {
+      const size_t second_newline = second_line.find_first_of("\r\n");
+      resolved_gpu_id = second_line.substr(0, second_newline);
+    }
+  }
+  return DbtRuntimeConfigHandoff{std::string(config_path), std::move(resolved_gpu_id)};
+}
+
 std::optional<DbtGuestConfig> load_dbt_guest_config_from_runtime_config() {
   // Try the handoff tiers in priority order, opening the first that exists:
   //   1. $ROCJITSU_INVOCATION_DIR/config_path — the launcher exports this dir before
@@ -176,11 +210,15 @@ std::optional<DbtGuestConfig> load_dbt_guest_config_from_runtime_config() {
   if (!file.is_open())
     return std::nullopt;
 
-  std::string path;
-  std::getline(file, path);
-  if (path.empty())
+  const std::string contents((std::istreambuf_iterator<char>(file)),
+                             std::istreambuf_iterator<char>());
+  std::optional<DbtRuntimeConfigHandoff> handoff = parse_dbt_runtime_config_handoff(contents);
+  if (!handoff)
     return std::nullopt;
-  return load_dbt_guest_config_from_file(path);
+  DbtGuestConfig config = load_dbt_guest_config_from_file(handoff->config_path);
+  if (handoff->resolved_gpu_id)
+    apply_resolved_dbt_host_gpu_id(config, *handoff->resolved_gpu_id, "runtime config handoff");
+  return config;
 }
 
 } // namespace config

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/config/dbt_guest_config.h

@@ -12,6 +12,7 @@
 #include <cstdint>
 #include <optional>
 #include <string>
+#include <string_view>
 
 namespace rocjitsu::fb {
 struct DbtGuestConfig;
@@ -39,7 +40,7 @@ enum class DbtSiliconRevision {
 /// @brief Host target selected for DBT translation and execution.
 struct DbtHostConfig {
   std::string isa;     ///< Host ISA used for DBT output and ROCR execution.
-  uint32_t gpu_id = 0; ///< Host KFD topology gpu_id; 0 matches topology to isa.
+  uint32_t gpu_id = 0; ///< Host KFD topology gpu_id; 0 selects the first ISA match.
   DbtExecutionBackend backend = DbtExecutionBackend::Hardware; ///< Hardware or simulator execution.
   std::string simulator_config_path; ///< Optional external simulator host config.
 };
@@ -64,6 +65,12 @@ struct DbtGuestConfig {
   DbtSiliconRevision host_revision = DbtSiliconRevision::Unspecified;
 };
 
+/// @brief Parsed per-invocation config handoff shared by child runtime layers.
+struct DbtRuntimeConfigHandoff {
+  std::string config_path;
+  std::optional<std::string> resolved_gpu_id;
+};
+
 /// @brief Resolve the simulator host config selected by a DBT guest config.
 /// @details An empty host_config_path selects dbt_config_path itself. Relative
 /// external paths are resolved beside the DBT guest config. A non-empty path
@@ -96,6 +103,17 @@ DbtGuestConfig dbt_guest_from_fb(const fb::DbtGuestConfig *guest);
 /// @throws std::runtime_error on file I/O, parse errors, or invalid config.
 DbtGuestConfig load_dbt_guest_config_from_file(const std::string &path);
 
+/// @brief Apply the launcher's resolved GPU to an automatic DBT host config.
+/// @details Child processes read this value from the second line of the
+/// per-invocation config handoff. Explicit nonzero config values remain authoritative.
+/// @throws std::runtime_error when @p value is not a nonzero KFD gpu_id.
+void apply_resolved_dbt_host_gpu_id(DbtGuestConfig &config, std::string_view value,
+                                    std::string_view source);
+
+/// @brief Parse the config path and optional resolved GPU ID from a runtime handoff.
+/// @returns std::nullopt when the first line does not contain a config path.
+std::optional<DbtRuntimeConfigHandoff> parse_dbt_runtime_config_handoff(std::string_view contents);
+
 /// @brief Load only dbt_guest from the rocjitsu child-process runtime config file.
 ///
 /// @details HSA tools run inside ROCR initialization and must not depend on the

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/kmd/linux/interposer.cpp

@@ -160,12 +160,12 @@ int kfd_ioctl_ret(int r) {
   return r;
 }
 
-/// @brief Read the child-process rocjitsu config path from @p cfg_file.
+/// @brief Read the child-process rocjitsu config handoff from @p cfg_file.
 ///
-/// @details The launcher writes the config path to a runtime file (per-PID
-/// invocation directory) that the interposer reads back for both local
-/// simulation and DBT guest mode.
-std::optional<std::string> child_config_path(const std::string &cfg_file) {
+/// @details The first line is the config path. DBT launches include the resolved
+/// host KFD gpu_id on the second line.
+std::optional<rocjitsu::config::DbtRuntimeConfigHandoff>
+child_config_handoff(const std::string &cfg_file) {
   char cfg_buf[4096]{};
   auto &real = rocjitsu::libc_passthrough();
   int cfg_fd = real.openat(AT_FDCWD, cfg_file.c_str(), O_RDONLY, 0);
@@ -177,9 +177,8 @@ std::optional<std::string> child_config_path(const std::string &cfg_file) {
   if (n <= 0)
     return std::nullopt;
 
-  while (n > 0 && (cfg_buf[n - 1] == '\n' || cfg_buf[n - 1] == '\r'))
-    cfg_buf[--n] = '\0';
-  return std::string(cfg_buf);
+  return rocjitsu::config::parse_dbt_runtime_config_handoff(
+      std::string_view(cfg_buf, static_cast<size_t>(n)));
 }
 
 void *raw_mmap_syscall(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
@@ -1161,31 +1160,34 @@ class InterposerContext {
       cfg_candidates.push_back(invocation_runtime_dir() + "/config_path");
       cfg_candidates.push_back(rocjitsu::rpc_invocation_config_file_path(getpid()));
       cfg_candidates.push_back(rocjitsu::rpc_default_config_file_path());
-      std::optional<std::string> cfg_path;
+      std::optional<rocjitsu::config::DbtRuntimeConfigHandoff> handoff;
       std::string tried_last;
       for (const auto &candidate : cfg_candidates) {
         if (candidate == tried_last)
           continue; // Skip a duplicate tier (e.g. env unset collapses 1 and 2).
         tried_last = candidate;
-        cfg_path = child_config_path(candidate);
-        if (cfg_path)
+        handoff = child_config_handoff(candidate);
+        if (handoff)
           break;
       }
-      if (!cfg_path) {
+      if (!handoff) {
         util::Logger::debug_print("rocjitsu: no child config path");
         in_construction = false;
         return nullptr;
       }
 
       try {
-        auto dbt_guest = rocjitsu::config::load_dbt_guest_config_from_file(*cfg_path);
+        auto dbt_guest = rocjitsu::config::load_dbt_guest_config_from_file(handoff->config_path);
+        if (handoff->resolved_gpu_id)
+          rocjitsu::config::apply_resolved_dbt_host_gpu_id(dbt_guest, *handoff->resolved_gpu_id,
+                                                           "runtime config handoff");
         if (dbt_guest.enabled) {
           LinuxKfd *execution_driver = nullptr;
           const bool simulator_backend =
               dbt_guest.host.backend == rocjitsu::config::DbtExecutionBackend::Simulator;
           if (simulator_backend) {
             const std::string host_config_path = rocjitsu::config::resolve_dbt_host_config_path(
-                *cfg_path, dbt_guest.host.simulator_config_path);
+                handoff->config_path, dbt_guest.host.simulator_config_path);
             if (!create_local_vm(host_config_path)) {
               in_construction = false;
               return nullptr;
@@ -1212,16 +1214,17 @@ class InterposerContext {
           in_construction = false;
           return driver;
         }
+
+        if (!create_local_vm(handoff->config_path)) {
+          in_construction = false;
+          return nullptr;
+        }
       } catch (const std::exception &e) {
         util::Logger::debug_print("rocjitsu: failed to load child config: ", e.what());
         destroy_local_vm();
         in_construction = false;
         return nullptr;
       }
-      if (!create_local_vm(*cfg_path)) {
-        in_construction = false;
-        return nullptr;
-      }
 
       // Publish the fully-constructed, not-yet-running driver BEFORE starting the
       // engine thread: the release-store of active_driver_ pairs with acquire

emulation/rocjitsu/tests/config_test.cpp

@@ -11,6 +11,7 @@
 #include "rocjitsu/isa/arch/amdgpu/cdna3/isa.h"
 #include "rocjitsu/isa/arch/amdgpu/shared/accvgpr_layout.h"
 #include "rocjitsu/kmd/linux/amdgpu_properties.h"
+#include "rocjitsu/kmd/linux/rpc.h"
 #include "rocjitsu/vm/rj_vm.h"
 #include "rocjitsu/vm/soc.h"
 #include "rocjitsu/vm/virtual_machine.h"
@@ -26,10 +27,12 @@ RJ_DIAGNOSTIC_POP
 #include <gtest/gtest.h>
 
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
 #include <filesystem>
 #include <fstream>
 #include <iterator>
+#include <optional>
 #include <stdexcept>
 #include <string>
 #include <string_view>
@@ -49,6 +52,26 @@ std::filesystem::path write_temp_config(std::string_view name, std::string_view
   return path;
 }
 
+class ScopedEnvironmentVariable {
+public:
+  ScopedEnvironmentVariable(const char *name, const std::string &value) : name_(name) {
+    if (const char *original = std::getenv(name))
+      original_ = original;
+    setenv(name_.c_str(), value.c_str(), 1);
+  }
+
+  ~ScopedEnvironmentVariable() {
+    if (original_)
+      setenv(name_.c_str(), original_->c_str(), 1);
+    else
+      unsetenv(name_.c_str());
+  }
+
+private:
+  std::string name_;
+  std::optional<std::string> original_;
+};
+
 TEST(ConfigLoaderTest, LoadCdna4Config) {
   std::string json = CONFIG_DIR_PATH + "/gfx950_cdna4.json";
   auto loaded = config::load_config(json, rocjitsu::kEmbeddedSchema);
@@ -562,6 +585,104 @@ TEST(ConfigLoaderTest, LoadsExplicitHardwareDbtBackendConfig) {
   EXPECT_EQ(dbt.host.backend, config::DbtExecutionBackend::Hardware);
 }
 
+TEST(ConfigLoaderTest, AppliesResolvedDbtHostGpuId) {
+  config::DbtGuestConfig automatic;
+  automatic.enabled = true;
+  automatic.host.backend = config::DbtExecutionBackend::Hardware;
+  config::DbtGuestConfig explicit_id = automatic;
+  explicit_id.host.gpu_id = 8716;
+  config::DbtGuestConfig simulator = automatic;
+  simulator.host.backend = config::DbtExecutionBackend::Simulator;
+
+  config::apply_resolved_dbt_host_gpu_id(automatic, "28851", "test handoff");
+  config::apply_resolved_dbt_host_gpu_id(explicit_id, "28851", "test handoff");
+  config::apply_resolved_dbt_host_gpu_id(simulator, "28851", "test handoff");
+
+  EXPECT_EQ(automatic.host.gpu_id, 28851u);
+  EXPECT_EQ(explicit_id.host.gpu_id, 8716u);
+  EXPECT_EQ(simulator.host.gpu_id, 28851u);
+}
+
+TEST(ConfigLoaderTest, RejectsInvalidResolvedDbtHostGpuId) {
+  const std::array<std::string_view, 5> invalid_values = {"", "0", "not-a-number", "28851 trailing",
+                                                          "4294967296"};
+  for (std::string_view value : invalid_values) {
+    config::DbtGuestConfig dbt;
+    dbt.enabled = true;
+    EXPECT_THROW(config::apply_resolved_dbt_host_gpu_id(dbt, value, "test handoff"),
+                 std::runtime_error)
+        << value;
+  }
+}
+
+TEST(ConfigLoaderTest, ExplicitDbtHostGpuIdOverridesResolvedHandoff) {
+  config::DbtGuestConfig dbt;
+  dbt.enabled = true;
+  dbt.host.gpu_id = 8716;
+
+  EXPECT_NO_THROW(
+      config::apply_resolved_dbt_host_gpu_id(dbt, "invalid-but-ignored", "test handoff"));
+  EXPECT_EQ(dbt.host.gpu_id, 8716u);
+}
+
+TEST(ConfigLoaderTest, ParsesRuntimeConfigHandoff) {
+  const auto dbt = config::parse_dbt_runtime_config_handoff("/tmp/config.json\r\n28851\r\n");
+  ASSERT_TRUE(dbt);
+  EXPECT_EQ(dbt->config_path, "/tmp/config.json");
+  ASSERT_TRUE(dbt->resolved_gpu_id);
+  EXPECT_EQ(*dbt->resolved_gpu_id, "28851");
+
+  const auto non_dbt = config::parse_dbt_runtime_config_handoff("/tmp/config.json");
+  ASSERT_TRUE(non_dbt);
+  EXPECT_EQ(non_dbt->config_path, "/tmp/config.json");
+  EXPECT_FALSE(non_dbt->resolved_gpu_id);
+
+  const auto newline_terminated = config::parse_dbt_runtime_config_handoff("/tmp/config.json\n");
+  ASSERT_TRUE(newline_terminated);
+  EXPECT_FALSE(newline_terminated->resolved_gpu_id);
+  EXPECT_FALSE(config::parse_dbt_runtime_config_handoff("\n28851\n"));
+}
+
+TEST(ConfigLoaderTest, RejectsEmptyResolvedGpuIdLineForEnabledDbt) {
+  const auto handoff = config::parse_dbt_runtime_config_handoff("/tmp/config.json\n\n");
+  ASSERT_TRUE(handoff);
+  ASSERT_TRUE(handoff->resolved_gpu_id);
+
+  config::DbtGuestConfig dbt;
+  dbt.enabled = true;
+  EXPECT_THROW(
+      config::apply_resolved_dbt_host_gpu_id(dbt, *handoff->resolved_gpu_id, "test handoff"),
+      std::runtime_error);
+}
+
+TEST(ConfigLoaderTest, LoadsDbtRuntimeConfigHandoffFromInvocationDirectory) {
+  const std::filesystem::path runtime =
+      std::filesystem::temp_directory_path() / "rocjitsu_runtime_config_handoff_test";
+  std::filesystem::create_directories(runtime);
+  const std::filesystem::path config_path =
+      write_temp_config("rocjitsu_runtime_dbt_config_test.json", R"({
+        "dbt_guest": {
+          "enabled": true,
+          "guest_isa": "gfx950",
+          "host_isa": "gfx942"
+        }
+      })");
+  {
+    std::ofstream handoff(runtime / "config_path");
+    handoff << config_path.string() << "\n28851\n";
+  }
+  ScopedEnvironmentVariable invocation_dir(rocjitsu::kRpcInvocationDirEnv, runtime.string());
+
+  const std::optional<config::DbtGuestConfig> loaded =
+      config::load_dbt_guest_config_from_runtime_config();
+
+  ASSERT_TRUE(loaded);
+  EXPECT_TRUE(loaded->enabled);
+  EXPECT_EQ(loaded->host.gpu_id, 28851u);
+  std::filesystem::remove(config_path);
+  std::filesystem::remove_all(runtime);
+}
+
 TEST(ConfigLoaderTest, RejectsEmptyDbtExecutionBackend) {
   const std::filesystem::path path =
       write_temp_config("rocjitsu_empty_dbt_backend_config_test.json", R"({