[rocjitsu][DBI][hook - 1/3] Basic Block Site Selection Layer

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/amdgpu_code_object.cpp

@@ -13,7 +13,6 @@
 
 #include <algorithm>
 #include <cstring>
-#include <iterator>
 #include <optional>
 #include <span>
 #include <utility>
@@ -390,19 +389,7 @@ AmdGpuCodeObject::min_kernel_sgpr_count(rj_code_arch_t arch,
 }
 
 std::optional<uint32_t> AmdGpuCodeObject::min_kernel_sgpr_count(rj_code_arch_t arch) const {
-  const std::span<const uint8_t> image{reinterpret_cast<const uint8_t *>(image_data()),
-                                       image_size()};
-  // Discover descriptors through the shared scanner rather than re-walking sections
-  // here. scan_kernel_descriptors() filters to a single .text extent, so visit each
-  // text section to cover every kernel, then reduce via the span overload.
-  std::vector<KernelDescriptorInfo> kernels;
-  for (const Section *text : text_sections()) {
-    std::vector<KernelDescriptorInfo> found =
-        scan_kernel_descriptors(image, text->sectionOffset(), text->size());
-    kernels.insert(kernels.end(), std::make_move_iterator(found.begin()),
-                   std::make_move_iterator(found.end()));
-  }
-  return min_kernel_sgpr_count(arch, kernels);
+  return min_kernel_sgpr_count(arch, scan_object_kernel_descriptors(*this));
 }
 
 } // namespace rocjitsu

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/analysis/liveness.cpp

@@ -212,6 +212,23 @@ bool accesses_sgprs_indirectly_via_movrel(KernelBlockScope blocks) {
   return false;
 }
 
+bool is_unmodeled_image_access(rj_code_arch_t arch, std::string_view mnemonic) {
+  return arch == ROCJITSU_CODE_ARCH_RDNA4 && mnemonic.starts_with("image_");
+}
+
+bool accesses_vgprs_indirectly(KernelBlockScope blocks, std::span<const uint8_t> text,
+                               rj_code_arch_t arch) {
+  for (BasicBlock *block : blocks) {
+    if (block == nullptr)
+      continue;
+    for (const Instruction &inst : block->instructions()) {
+      if (may_access_vgprs_indirectly(inst, text, arch))
+        return true;
+    }
+  }
+  return false;
+}
+
 LivenessAnalysis::LivenessAnalysis(KernelBlockScope blocks, uint8_t wave_size,
                                    std::unique_ptr<ExecMaskAnalysis> exec,
                                    LivenessAnalysisOptions options,

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/analysis/liveness.h

@@ -23,6 +23,7 @@
 #include <memory>
 #include <optional>
 #include <span>
+#include <string_view>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
@@ -190,6 +191,23 @@ struct LivenessAnalysisOptions {
 /// its own scan.
 [[nodiscard]] bool accesses_sgprs_indirectly_via_movrel(KernelBlockScope blocks);
 
+/// @brief Whether any instruction in @p blocks may access a VGPR its operands do
+///        not name: `v_movrel*` / `v_swaprel*`, or GPR index mode, which
+///        `s_set_gpr_idx_on` or a MODE write enables on CDNA.
+///
+/// @details LivenessAnalysis does not model that displacement either, so a
+/// caller relying on it for VGPR liveness must reject a scope where this holds.
+/// @p text and @p arch let the scan read the MODE bits a write sets.
+[[nodiscard]] bool accesses_vgprs_indirectly(KernelBlockScope blocks, std::span<const uint8_t> text,
+                                             rj_code_arch_t arch);
+
+/// @brief Whether @p mnemonic, decoded for @p arch, is an RDNA4 image
+///        instruction (VIMAGE / VSAMPLE).
+///
+/// @details Their decoders do not report the address VGPRs, so InstDefUse
+/// misses those reads and liveness can report the registers dead.
+[[nodiscard]] bool is_unmodeled_image_access(rj_code_arch_t arch, std::string_view mnemonic);
+
 /// @brief Kernel-global register usage plus deferred backward CFG liveness.
 ///
 /// @details Construction eagerly resolves gfx1250 VGPR banks and scans global

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/dbt/kernel_descriptor_translator.cpp

@@ -42,7 +42,6 @@ static_assert(sizeof(KD) == 64, "AMDHSA kernel descriptor size changed");
 
 constexpr uint32_t kMaxVgprGranulatedField = 63;
 constexpr uint32_t kMaxSgprGranulatedField = 15;
-constexpr uint64_t kKernargPreloadSkipBytes = 256;
 constexpr uint16_t kScalarOperandTtmpBase = 108;
 constexpr uint16_t kTtmpRdna4GridYz = 7;
 constexpr uint16_t kTtmpRdna4GridX = 9;
@@ -197,14 +196,6 @@ constexpr uint16_t kTtmpRdna4GridX = 9;
   return static_cast<uint16_t>(sgpr);
 }
 
-[[nodiscard]] bool uses_kernarg_preload_firmware_skip(rj_code_arch_t arch) {
-  // CDNA3/CDNA4 implement kernarg preloading through the legacy firmware
-  // compatibility window: old firmware enters at the descriptor entry and
-  // compatible firmware enters 256 bytes later. GFX1250 supports preloading in
-  // hardware but has one entry, exactly as specified by the descriptor.
-  return arch == ROCJITSU_CODE_ARCH_CDNA3 || arch == ROCJITSU_CODE_ARCH_CDNA4;
-}
-
 [[nodiscard]] std::optional<uint32_t> kernarg_bytes_to_preserve(const KD &desc) {
   const uint32_t preload_length = kernarg_preload_length(desc);
   if (preload_length == 0)

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/kernel_descriptor_scan.cpp

@@ -4,11 +4,13 @@
 #include "rocjitsu/code/kernel_descriptor_scan.h"
 
 #include "rocjitsu/code/amdgpu_elf.h"
+#include "rocjitsu/code/code_object.h"
 #include "rocjitsu/isa/isa_traits.h"
 #include "rocjitsu/isa/register_set.h"
 
 #include <algorithm>
 #include <cstring>
+#include <iterator>
 #include <limits>
 #include <optional>
 #include <unordered_set>
@@ -190,6 +192,19 @@ scan_kernel_descriptors(std::span<const uint8_t> image, uint64_t text_offset, ui
   return out;
 }
 
+std::vector<KernelDescriptorInfo> scan_object_kernel_descriptors(const CodeObject &object) {
+  const std::span<const uint8_t> image(reinterpret_cast<const uint8_t *>(object.image_data()),
+                                       object.image_size());
+  std::vector<KernelDescriptorInfo> kernels;
+  for (const Section *text : object.text_sections()) {
+    std::vector<KernelDescriptorInfo> found =
+        scan_kernel_descriptors(image, text->sectionOffset(), text->size());
+    kernels.insert(kernels.end(), std::make_move_iterator(found.begin()),
+                   std::make_move_iterator(found.end()));
+  }
+  return kernels;
+}
+
 uint8_t kernel_wavefront_size(rj_code_arch_t arch, const KD &desc) {
   // CDNA kernels are Wave64 in the code objects currently handled here.
   if (arch_is_cdna_4_or_lower(arch))
@@ -318,4 +333,13 @@ uint32_t kernarg_preload_offset(const KD &desc) {
   return RJ_AMDHSA_BITS_GET(desc.kernarg_preload, rocjitsu::amdhsa::KERNARG_PRELOAD_SPEC_OFFSET);
 }
 
+bool uses_kernarg_preload_firmware_skip(rj_code_arch_t arch) {
+  // CDNA3/CDNA4 implement kernarg preloading through the legacy firmware
+  // compatibility window: old firmware enters at the descriptor entry and
+  // compatible firmware enters 256 bytes later. CDNA2 does too, but is left out
+  // (see the header). GFX1250 supports preloading in hardware but has one
+  // entry, exactly as specified by the descriptor.
+  return arch == ROCJITSU_CODE_ARCH_CDNA3 || arch == ROCJITSU_CODE_ARCH_CDNA4;
+}
+
 } // namespace rocjitsu

emulation/rocjitsu/lib/rocjitsu/src/rocjitsu/code/kernel_descriptor_scan.h

@@ -18,6 +18,8 @@
 
 namespace rocjitsu {
 
+class CodeObject;
+
 /// @brief One AMDHSA kernel descriptor located in an ELF image.
 struct KernelDescriptorInfo {
   uint64_t descriptor_file_offset = 0; ///< File offset of the 64-byte descriptor (pre-growth).
@@ -36,6 +38,14 @@ struct KernelDescriptorInfo {
 scan_kernel_descriptors(std::span<const uint8_t> image, uint64_t text_offset, uint64_t text_size,
                         std::optional<size_t> text_section_index = std::nullopt);
 
+/// @brief Every descriptor of @p object: scan_kernel_descriptors() over each
+///        .text section, in section order.
+///
+/// @details Each entry_text_offset is relative to the .text section its entry
+/// lands in, so with more than one .text the offsets do not share an origin.
+[[nodiscard]] std::vector<KernelDescriptorInfo>
+scan_object_kernel_descriptors(const CodeObject &object);
+
 /// @brief SGPRs at the top of a CDNA kernel's allocation that hold special
 ///        registers rather than ordinary ones.
 ///
@@ -121,4 +131,17 @@ kernarg_segment_ptr_sgpr(const rocjitsu::amdhsa::kernel_descriptor_t &desc);
 /// @brief KERNARG_PRELOAD_SPEC_OFFSET: first preloaded kernarg dword index.
 [[nodiscard]] uint32_t kernarg_preload_offset(const rocjitsu::amdhsa::kernel_descriptor_t &desc);
 
+/// @brief Distance past the descriptor entry at which kernarg-preload-compatible
+///        firmware enters a preloading kernel on CDNA2/CDNA3/CDNA4.
+inline constexpr uint64_t kKernargPreloadSkipBytes = 256;
+
+/// @brief Whether DBT treats a kernel of @p arch that preloads kernel arguments
+///        as also entered kKernargPreloadSkipBytes past its descriptor entry.
+///
+/// @details Old firmware enters at the descriptor entry, where the compiler
+/// places a compatibility prologue that loads the preloaded arguments itself.
+/// Covers CDNA3/CDNA4. CDNA2 firmware enters the same way, but DBT does not
+/// handle that entry, so this excludes it.
+[[nodiscard]] bool uses_kernarg_preload_firmware_skip(rj_code_arch_t arch);
+
 } // namespace rocjitsu