fix(kpack): fall through to next matching kpack archive on kernel-not-found

shared/kpack/runtime/src/loader.cpp

@@ -7,6 +7,7 @@
 #include <filesystem>
 #include <msgpack.hpp>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "isa_target_match.h"
@@ -414,8 +415,11 @@ kpack_error_t kpack_load_code_object(kpack_cache_t cache,
     // to find an archive that has a matching architecture.
     KPACK_DEBUG(cache, "trying architecture: %s", arch);
 
-    kpack_archive_t archive = nullptr;
-    std::string matched_arch;
+    // A single architecture's kernels may be split across several archives
+    // (e.g. an xnack-variant kpack holding a handful of kernels plus the bare
+    // kpack holding the rest), so collect every matching archive rather than
+    // committing to the most specific one.
+    std::vector<std::pair<kpack_archive_t, std::string>> candidates;
 
     // Find archive containing a compatible architecture
     // Lock only for cache lookup, release before kernel fetch
@@ -435,34 +439,38 @@ kpack_error_t kpack_load_code_object(kpack_cache_t cache,
 
           auto archive_it = cache->archives.find(archive_path);
           if (archive_it != cache->archives.end()) {
-            archive = archive_it->second;
-            matched_arch = target;
-            return true;  // found — stop searching
+            candidates.emplace_back(archive_it->second, target);
           }
         }
-        return false;  // try next compatible target
+        return false;  // keep collecting from every compatible target
       });
     }  // Release cache->archive_mutex before kernel fetch
 
-    if (!archive) {
+    if (candidates.empty()) {
       continue;
     }
 
     // Fetch kernel using the matched architecture (which may be a subset of
     // the agent's full ISA, e.g., bare "gfx942" for a release build).
-    err = kpack_get_kernel(archive, lookup_key.c_str(), matched_arch.c_str(),
-                           &kernel_data, &kernel_size);
-    if (err == KPACK_SUCCESS) {
-      KPACK_DEBUG(cache, "  found kernel: %zu bytes", kernel_size);
+    for (const auto& [archive, matched_arch] : candidates) {
+      err = kpack_get_kernel(archive, lookup_key.c_str(), matched_arch.c_str(),
+                             &kernel_data, &kernel_size);
+      if (err == KPACK_SUCCESS) {
+        KPACK_DEBUG(cache, "  found kernel: %zu bytes", kernel_size);
+        break;
+      }
+      // Archive was found with matching architecture but the kernel was not
+      // present in it (e.g. an xnack-variant kpack that is missing a kernel
+      // present in the base/generic kpack). Continue to the next candidate
+      // rather than returning immediately — a less-specific but still ISA-
+      // compatible archive (e.g. bare gfx90a.kpack) may contain the kernel.
+      KPACK_DEBUG(cache, "  kernel not found in this archive (error %d), trying next candidate", err);
+      last_err = err;
+    }
+
+    if (kernel_data) {
       break;
     }
-    // Archive was found with matching architecture but the kernel was not
-    // present in it (e.g. an xnack-variant kpack that is missing a kernel
-    // present in the base/generic kpack). Continue to the next candidate
-    // rather than returning immediately — a less-specific but still ISA-
-    // compatible archive (e.g. bare gfx90a.kpack) may contain the kernel.
-    KPACK_DEBUG(cache, "  kernel not found in this archive (error %d), trying next candidate", err);
-    last_err = err;
   }
 
   if (!kernel_data) {

shared/kpack/runtime/tests/generate_test_data.py

@@ -41,7 +41,7 @@ def generate_noop_archive(output_dir: Path) -> None:
     archive.add_kernel(prepared2)
 
     # Kernel 3: bin/testapp#0 @ gfx900
-    kernel3_data = b"KERNEL3_APP_GFX900" + b"\xFF" * 150
+    kernel3_data = b"KERNEL3_APP_GFX900" + b"\xff" * 150
     prepared3 = archive.prepare_kernel("bin/testapp#0", "gfx900", kernel3_data)
     archive.add_kernel(prepared3)
 
@@ -94,6 +94,109 @@ def generate_zstd_archive(output_dir: Path) -> None:
     print(f"  - bin/hiptest#0: gfx1100 ({len(kernel3_data)} bytes)")
 
 
+def generate_fallthrough_archives(output_dir: Path) -> None:
+    """Generate archives for the xnack-split archive-fallthrough regression test.
+
+    Reproduces the ROCm/TheRock#7081 archive layout: a kernel's full set of
+    kpack archives is split across a more-specific (xnack) archive holding
+    unrelated kernels and a less-specific (bare) archive holding the actual
+    kernel. The loader must fall through from the xnack archive to the bare
+    archive on a kernel-not-found miss instead of giving up immediately.
+    """
+    # xnack-specific archive: matches the agent ISA more specifically, but
+    # does not contain the kernel under test.
+    xnack_archive = PackedKernelArchive(
+        group_name="test",
+        gfx_arch_family="gfx90aX",
+        gfx_arches=["gfx90a:xnack-"],
+        compressor=NoOpCompressor(),
+    )
+    other_data = b"OTHER_KERNEL_XNACK_MINUS" + b"\x00" * 50
+    prepared = xnack_archive.prepare_kernel(
+        "lib/libother.so#0", "gfx90a:xnack-", other_data
+    )
+    xnack_archive.add_kernel(prepared)
+    xnack_archive.finalize_archive()
+    xnack_archive.write(output_dir / "test_fallthrough_xnack.kpack")
+    print(
+        f"Generated fallthrough xnack archive: {output_dir / 'test_fallthrough_xnack.kpack'}"
+    )
+    print("  - gfx90a:xnack-, 1 unrelated kernel (lib/libother.so#0)")
+
+    # Bare/generic archive: less specific, holds the kernel under test.
+    generic_archive = PackedKernelArchive(
+        group_name="test",
+        gfx_arch_family="gfx90aX",
+        gfx_arches=["gfx90a"],
+        compressor=NoOpCompressor(),
+    )
+    target_data = b"GENERIC_GFX90A_KERNEL_DATA" + b"\x00" * 50
+    prepared2 = generic_archive.prepare_kernel(
+        "lib/libtest.so#0", "gfx90a", target_data
+    )
+    generic_archive.add_kernel(prepared2)
+    generic_archive.finalize_archive()
+    generic_archive.write(output_dir / "test_fallthrough_generic.kpack")
+    print(
+        f"Generated fallthrough generic archive: {output_dir / 'test_fallthrough_generic.kpack'}"
+    )
+    print("  - gfx90a, 1 kernel (lib/libtest.so#0)")
+
+
+def generate_fallthrough_two_feature_archives(output_dir: Path) -> None:
+    """Generate archives for a two-feature (sramecc+xnack) fallthrough test.
+
+    Reproduces the ROCM-31644 archive layout on gfx950: the agent reports
+    "gfx950:sramecc+:xnack+", but no archive is tagged with "sramecc+" at
+    all (rocFFT's build strips feature flags before compiling — see
+    projects/rocfft/CMakeLists.txt). Only two archives exist for the group:
+    one tagged "gfx950:xnack+" holding an unrelated binary's kernel (e.g.
+    hipFFT-test), and one bare "gfx950" holding the target kernel (e.g.
+    rocFFT-test). The loader must walk the full compatible-target power set
+    ("gfx950:sramecc+:xnack+" -> "gfx950:sramecc+" -> "gfx950:xnack+" ->
+    "gfx950"), skipping the two candidates that have no archive at all,
+    missing the kernel in the "gfx950:xnack+" archive, and finally finding
+    it in the bare "gfx950" archive.
+    """
+    xnack_archive = PackedKernelArchive(
+        group_name="test",
+        gfx_arch_family="gfx950X",
+        gfx_arches=["gfx950:xnack+"],
+        compressor=NoOpCompressor(),
+    )
+    other_data = b"OTHER_KERNEL_HIPFFT_TEST" + b"\x00" * 50
+    prepared = xnack_archive.prepare_kernel(
+        "lib/libother.so#0", "gfx950:xnack+", other_data
+    )
+    xnack_archive.add_kernel(prepared)
+    xnack_archive.finalize_archive()
+    xnack_archive.write(output_dir / "test_fallthrough_2feat_xnack.kpack")
+    print(
+        f"Generated 2-feature fallthrough xnack archive: "
+        f"{output_dir / 'test_fallthrough_2feat_xnack.kpack'}"
+    )
+    print("  - gfx950:xnack+, 1 unrelated kernel (lib/libother.so#0)")
+
+    generic_archive = PackedKernelArchive(
+        group_name="test",
+        gfx_arch_family="gfx950X",
+        gfx_arches=["gfx950"],
+        compressor=NoOpCompressor(),
+    )
+    target_data = b"GENERIC_GFX950_ROCFFT_TEST" + b"\x00" * 50
+    prepared2 = generic_archive.prepare_kernel(
+        "lib/libtest.so#0", "gfx950", target_data
+    )
+    generic_archive.add_kernel(prepared2)
+    generic_archive.finalize_archive()
+    generic_archive.write(output_dir / "test_fallthrough_2feat_generic.kpack")
+    print(
+        f"Generated 2-feature fallthrough generic archive: "
+        f"{output_dir / 'test_fallthrough_2feat_generic.kpack'}"
+    )
+    print("  - gfx950, 1 kernel (lib/libtest.so#0)")
+
+
 def generate_test_manifests(output_dir: Path) -> None:
     """Generate .kpm manifest files for testing manifest resolution."""
     # Manifest for test_noop.kpack (gfx900, gfx906)
@@ -175,6 +278,10 @@ def main() -> None:
     print()
     generate_zstd_archive(output_dir)
     print()
+    generate_fallthrough_archives(output_dir)
+    print()
+    generate_fallthrough_two_feature_archives(output_dir)
+    print()
     generate_test_manifests(output_dir)
     print()
     print("Done!")

shared/kpack/runtime/tests/test_loader_api.cpp

@@ -778,6 +778,98 @@ TEST(LoaderAPITest, LoadCodeObject_KernelNotFound) {
   EXPECT_EQ(code_object, nullptr);
 }
 
+// Regression test for ROCm/TheRock#7081: when a requested architecture's
+// kernels are split across multiple kpack archives (e.g. a more-specific
+// "gfx90a:xnack-" archive plus a less-specific "gfx90a" archive), a
+// kernel-not-found miss in the first matching archive must not abort the
+// search — the loader must fall through to the next compatible archive.
+//
+// test_fallthrough_xnack.kpack declares architecture "gfx90a:xnack-" and
+// contains only an unrelated kernel ("lib/libother.so#0").
+// test_fallthrough_generic.kpack declares architecture "gfx90a" and contains
+// the kernel under test ("lib/libtest.so#0").
+//
+// Both archives are listed as search paths. For agent ISA "gfx90a:xnack-",
+// for_each_compatible_target() yields "gfx90a:xnack-" before "gfx90a", so the
+// xnack archive is tried first and matches the architecture but not the
+// kernel. Before the fix, the loader returned KPACK_ERROR_KERNEL_NOT_FOUND
+// immediately at that point. With the fix, it falls through and finds the
+// kernel in the generic archive.
+TEST(LoaderAPITest, LoadCodeObject_FallsThroughToNextArchiveOnKernelNotFound) {
+  std::string assets_dir = get_test_assets_dir();
+  if (assets_dir.empty()) {
+    GTEST_SKIP() << "ROCM_KPACK_TEST_ASSETS_DIR not set";
+  }
+
+  CacheGuard cache;
+  ASSERT_EQ(kpack_cache_create(cache.ptr()), KPACK_SUCCESS);
+
+  // Search paths list the more-specific (xnack) archive first, matching the
+  // real torch device wheel layout where the xnack archive is also listed
+  // ahead of the bare-arch archive.
+  auto metadata = make_hipk_metadata(
+      "lib/libtest.so",
+      {"test_fallthrough_xnack.kpack", "test_fallthrough_generic.kpack"});
+  std::string binary_path = assets_dir + "/fake_binary.so";
+
+  const char* arch_list[] = {"gfx90a:xnack-"};
+  void* code_object = nullptr;
+  size_t size = 0;
+
+  kpack_error_t err =
+      kpack_load_code_object(cache.get(), metadata.data(), binary_path.c_str(),
+                             0, arch_list, 1, &code_object, &size);
+
+  ASSERT_EQ(err, KPACK_SUCCESS);
+  ASSERT_NE(code_object, nullptr);
+  EXPECT_EQ(std::memcmp(code_object, "GENERIC_GFX90A_KERNEL_DATA", 26), 0);
+
+  kpack_free_code_object(code_object);
+}
+
+// Regression test for ROCM-31644: a two-feature agent ISA
+// ("gfx950:sramecc+:xnack+") where no archive is tagged with "sramecc+" at
+// all — two of the four compatible-target candidates
+// ("gfx950:sramecc+:xnack+", "gfx950:sramecc+") have no archive whatsoever,
+// not merely a kernel-not-found miss. The loader must skip those cleanly,
+// miss the kernel in the "gfx950:xnack+" archive (holding an unrelated
+// binary's kernel), and fall through to the bare "gfx950" archive that
+// actually holds it. This is the exact shape rocFFT-test hits on gfx950
+// ASAN builds: rocFFT's own CMakeLists.txt strips all feature flags from
+// GPU_TARGETS before compiling ("rocFFT builds generic code for each
+// target"), so its kernels always land in the bare, untagged archive even
+// when a sibling binary in the same kpack group (e.g. hipFFT-test) is
+// correctly tagged "xnack+".
+TEST(LoaderAPITest,
+     LoadCodeObject_FallsThroughAcrossTwoFeatureCandidatesOnKernelNotFound) {
+  std::string assets_dir = get_test_assets_dir();
+  if (assets_dir.empty()) {
+    GTEST_SKIP() << "ROCM_KPACK_TEST_ASSETS_DIR not set";
+  }
+
+  CacheGuard cache;
+  ASSERT_EQ(kpack_cache_create(cache.ptr()), KPACK_SUCCESS);
+
+  auto metadata = make_hipk_metadata("lib/libtest.so",
+                                     {"test_fallthrough_2feat_xnack.kpack",
+                                      "test_fallthrough_2feat_generic.kpack"});
+  std::string binary_path = assets_dir + "/fake_binary.so";
+
+  const char* arch_list[] = {"gfx950:sramecc+:xnack+"};
+  void* code_object = nullptr;
+  size_t size = 0;
+
+  kpack_error_t err =
+      kpack_load_code_object(cache.get(), metadata.data(), binary_path.c_str(),
+                             0, arch_list, 1, &code_object, &size);
+
+  ASSERT_EQ(err, KPACK_SUCCESS);
+  ASSERT_NE(code_object, nullptr);
+  EXPECT_EQ(std::memcmp(code_object, "GENERIC_GFX950_ROCFFT_TEST", 26), 0);
+
+  kpack_free_code_object(code_object);
+}
+
 TEST(LoaderAPITest, LoadCodeObject_ArchiveFileNotFound) {
   std::string assets_dir = get_test_assets_dir();
   if (assets_dir.empty()) {