Add changes to support scratchpad buffer patching for new ELFs (Xilinx#9741)

* Add changes to support scratchpad buffer patching for new ELFs

Signed-off-by: Rahul Bramandlapalli <rbramand@amd.com>
Signed-off-by: rahul <rbramand@amd.com>

* Address comments on PR

Signed-off-by: Rahul Bramandlapalli <rbramand@amd.com>

* Remove code duplicacy

Signed-off-by: rahul <rbramand@amd.com>

---------

Signed-off-by: Rahul Bramandlapalli <rbramand@amd.com>
Signed-off-by: rahul <rbramand@amd.com>

Author: rbramand-xilinx

src/runtime_src/core/common/api/elf_int.h

@@ -175,7 +175,7 @@ struct module_config_aie_gen2
   const buf& preempt_save_data;
   const buf& preempt_restore_data;
 
-  // Size of scratch pad memory
+  // Size of scratch pad memory per column
   size_t scratch_pad_mem_size;
 
   // Control scratch pad memory size (0 if not present)
@@ -214,6 +214,9 @@ struct module_config_aie_gen2_plus
 
   // Reference to dump buffer for debug/trace
   const buf& dump_buf;
+
+  // Size of scratch pad memory per column
+  size_t scratch_pad_mem_size;
   // NOLINTEND
 
   // Parent elf_impl pointer for any mutable operations

src/runtime_src/core/common/api/elf_patcher.cpp

@@ -157,7 +157,7 @@ patch_ctrl57_aie4(uint32_t* bd_data_ptr, uint64_t patch)
 
 void
 symbol_patcher::
-patch_symbol(xrt::bo bo, uint64_t value, bool first)
+patch_symbol(xrt::bo bo, uint64_t value, bool first, bool is_arg)
 {
   if (!m_config)
     throw std::runtime_error("symbol_patcher: config not set");
@@ -176,14 +176,20 @@ patch_symbol(xrt::bo bo, uint64_t value, bool first)
     auto offset = config.offset_to_patch_buffer;
     auto bd_data_ptr = reinterpret_cast<uint32_t*>(base + offset);
 
-    if (!state.dirty) {
-      // first time patching cache bd ptr values using bd ptrs array in patch state
-      std::copy(bd_data_ptr, bd_data_ptr + max_bd_words, state.bd_data_ptrs.begin());
-      state.dirty = true;
-    }
-    else {
-      // not the first time patching, restore bd ptr values from patch state bd ptrs array
-      std::copy(state.bd_data_ptrs.begin(), state.bd_data_ptrs.end(), bd_data_ptr);
+    // If the symbol patched is an argument to kernel we have to cache
+    // original bd data ptrs as args can be changed in subsequent runs.
+    // For non-arg symbols we only patch and sync without caching
+    // as they are patched once and never changed.
+    if (is_arg) {
+      if (!state.dirty) {
+        // first time patching cache bd ptr values using bd ptrs array in patch state
+        std::copy(bd_data_ptr, bd_data_ptr + max_bd_words, state.bd_data_ptrs.begin());
+        state.dirty = true;
+      }
+      else {
+        // not the first time patching, restore bd ptr values from patch state bd ptrs array
+        std::copy(state.bd_data_ptrs.begin(), state.bd_data_ptrs.end(), bd_data_ptr);
+      }
     }
 
     // lambda for calling sync bo

src/runtime_src/core/common/api/elf_patcher.h

@@ -59,6 +59,18 @@ generate_key_string(const std::string& argument_name, buf_type type)
   return argument_name + std::to_string(static_cast<int>(type));
 }
 
+// Get symbol name from section name
+inline std::string
+get_symbol_name_from_section_name(const std::string& section_name)
+{
+  // Symbol name is section name without the grp idx
+  // if sec name is .ctrlpkt-57.grp_idx then sym name is .ctrlpkt-57
+  auto dot_pos = section_name.rfind('.');
+  return (dot_pos != std::string::npos && dot_pos > 0)
+      ? section_name.substr(0, dot_pos)
+      : section_name;
+}
+
 // Symbol type enum for patching schemes
 enum class symbol_type {
   uc_dma_remote_ptr_symbol_kind = 1,
@@ -132,7 +144,7 @@ struct symbol_patcher
 
   // Function to patch a symbol in the buffer.
   void
-  patch_symbol(xrt::bo bo, uint64_t value, bool first);
+  patch_symbol(xrt::bo bo, uint64_t value, bool first, bool is_arg = true);
 
   // static method for patching raw buffers passed by shim tests
   // where the caller handles sync themselves

src/runtime_src/core/common/api/xrt_elf.cpp

@@ -35,6 +35,9 @@ using xrt_core::elf_int::no_ctrl_code_id;
 static constexpr size_t
 operator"" _kb(unsigned long long v)  { return 1024u * v; } // NOLINT
 
+static constexpr size_t
+operator"" _mb(unsigned long long v)  { return 1024u * 1024u * v; } // NOLINT
+
 ///////////////////////////////////////////////////////////////
 // Helper functions for kernel signature demangling and parsing
 ///////////////////////////////////////////////////////////////
@@ -1269,6 +1272,11 @@ class elf_aie_gen2_plus : public elf_impl
         // section to patch is ctrlpkt
         abs_offset += rela->r_offset;
         buf_type = patcher_buf_type::ctrlpkt;
+        // we have multiple ctrlpkt sections and to uniquely identify symbol
+        // for patching we add ctrlpkt section symbol name to key string
+        auto sym_name =
+            xrt_core::elf_patcher::get_symbol_name_from_section_name(patch_sec_name);
+        argnm += sym_name;
       }
       else {
         // section to patch is ctrlcode
@@ -1362,10 +1370,17 @@ class elf_aie_gen2_plus : public elf_impl
 
     static const std::map<std::string, buf> empty_ctrlpkt_map;
 
+    auto scratch_pad_size = (m_platform == elf::platform::aie4  ||
+                             m_platform == elf::platform::aie4a ||
+                             m_platform == elf::platform::aie4z)
+      ? 3_mb  // NOLINT
+      : 0;
+
     return module_config_aie_gen2_plus{
       ctrlcode_it->second,                                                       // ctrlcodes
       ctrlpkt_it != m_ctrlpkt_buf_map.end() ? ctrlpkt_it->second : empty_ctrlpkt_map, // ctrlpkt_bufs
       dump_it != m_dump_buf_map.end() ? dump_it->second : buf::get_empty_buf(),  // dump_buf
+      scratch_pad_size,                                                          // scratch_pad_mem_size
       this                                                                       // elf_parent
     };
   }

src/runtime_src/core/common/api/xrt_module.cpp

@@ -200,6 +200,9 @@ class module_run : public module_impl
   using patcher_config = xrt_core::elf_patcher::patcher_config;
   using symbol_patcher = xrt_core::elf_patcher::symbol_patcher;
 
+  // scratchpad memory symbol name
+  static constexpr const char* Scratch_Pad_Mem_Symbol = "scratch-pad-mem";
+
   // Pointer to shared patcher configs
   // This is created during ELF parsing and shared across module_run instances
   const std::map<std::string, patcher_config>* m_patcher_configs = nullptr;
@@ -262,22 +265,22 @@ class module_run : public module_impl
 
   // Helper function for patching buffer with argument name or index
   bool
-  patch_helper(xrt::bo& bo, const std::string& argnm, size_t index, uint64_t patch,
-               xrt_core::elf_patcher::buf_type type)
+  patch_helper(xrt::bo& bo, uint64_t patch, xrt_core::elf_patcher::buf_type type,
+               const std::string& arg_string, const std::string& index_string,
+               bool is_arg = true)
   {
     // Check if patcher configs exist
     if (!m_patcher_configs || m_patcher_configs->empty())
         return false;
 
-    const auto key_string = xrt_core::elf_patcher::generate_key_string(argnm, type);
+    const auto key_string = xrt_core::elf_patcher::generate_key_string(arg_string, type);
 
     auto config_it = m_patcher_configs->find(key_string);
     auto not_found_use_argument_name = (config_it == m_patcher_configs->end());
     std::string used_key = key_string;
 
     if (not_found_use_argument_name) {
       // Search using index
-      auto index_string = std::to_string(index);
       const auto key_index_string = xrt_core::elf_patcher::generate_key_string(index_string, type);
       config_it = m_patcher_configs->find(key_index_string);
       if (config_it == m_patcher_configs->end())
@@ -293,20 +296,20 @@ class module_run : public module_impl
     }
 
     // Call patch - symbol_patcher owns its state internally
-    patcher_it->second.patch_symbol(bo, patch, m_first_patch);
+    patcher_it->second.patch_symbol(bo, patch, m_first_patch, is_arg);
 
     if (xrt_core::config::get_xrt_debug()) {
       if (not_found_use_argument_name) {
         std::stringstream ss;
         ss << "Patched " << xrt_core::elf_patcher::get_section_name(type)
-           << " using argument index " << index
+           << " using argument index " << index_string
            << " with value " << std::hex << patch;
         xrt_core::message::send( xrt_core::message::severity_level::debug, "xrt_module", ss.str());
       }
       else {
         std::stringstream ss;
         ss << "Patched " << xrt_core::elf_patcher::get_section_name(type)
-           << " using argument name " << argnm
+           << " using argument name " << arg_string
            << " with value " << std::hex << patch;
         xrt_core::message::send( xrt_core::message::severity_level::debug, "xrt_module", ss.str());
       }
@@ -363,8 +366,7 @@ class module_run_aie_gen2 : public module_run
   // map storing xrt::bo that stores pdi data corresponding to each pdi symbol
   std::map<std::string, xrt::bo> m_pdi_bo_map;
 
-  // Symbol names for patching
-  static constexpr const char* Scratch_Pad_Mem_Symbol = "scratch-pad-mem";
+  // Symbol names for patching specific to aie_gen2 platform
   static constexpr const char* Control_Packet_Symbol = "control-packet";
   static constexpr const char* Control_ScratchPad_Symbol = "scratch-pad-ctrl";
 
@@ -457,10 +459,10 @@ class module_run_aie_gen2 : public module_run
       if (!scratchpad_mem)
         throw std::runtime_error("Failed to get scratchpad buffer from context\n");
 
-      patch_helper(m_preempt_save_bo, Scratch_Pad_Mem_Symbol, 0, scratchpad_mem.address(),
-                   xrt_core::elf_patcher::buf_type::preempt_save);
-      patch_helper(m_preempt_restore_bo, Scratch_Pad_Mem_Symbol, 0, scratchpad_mem.address(),
-                   xrt_core::elf_patcher::buf_type::preempt_restore);
+      patch_helper(m_preempt_save_bo, scratchpad_mem.address(),
+                   xrt_core::elf_patcher::buf_type::preempt_save, Scratch_Pad_Mem_Symbol, {}, false);
+      patch_helper(m_preempt_restore_bo, scratchpad_mem.address(),
+                   xrt_core::elf_patcher::buf_type::preempt_restore, Scratch_Pad_Mem_Symbol, {}, false);
 
       if (is_dump_preemption_codes()) {
         std::stringstream ss;
@@ -475,8 +477,8 @@ class module_run_aie_gen2 : public module_run
     // Create control scratchpad buffer and patch if symbol is present
     if (m_config.ctrl_scratch_pad_mem_size > 0) {
       m_ctrl_scratch_pad_mem = xbi::create_bo(m_hwctx, m_config.ctrl_scratch_pad_mem_size, xbi::use_type::ctrl_scratch_pad);
-      patch_helper(m_instr_bo, Control_ScratchPad_Symbol, 0, m_ctrl_scratch_pad_mem.address(),
-                   xrt_core::elf_patcher::buf_type::ctrltext);
+      patch_helper(m_instr_bo, m_ctrl_scratch_pad_mem.address(),
+                   xrt_core::elf_patcher::buf_type::ctrltext, Control_ScratchPad_Symbol, {}, false);
     }
 
     // Patch all PDI addresses using config's pdi symbols
@@ -488,13 +490,13 @@ class module_run_aie_gen2 : public module_run
       auto [it, inserted] = m_pdi_bo_map.emplace(symbol, std::move(pdi_bo));
 
       // Patch instr buffer with PDI address
-      patch_helper(m_instr_bo, symbol, 0, it->second.address(), xrt_core::elf_patcher::buf_type::ctrltext);
+      patch_helper(m_instr_bo, it->second.address(), xrt_core::elf_patcher::buf_type::ctrltext, symbol, {}, false);
     }
 
     // Patch control packet address if present
     if (m_ctrlpkt_bo) {
-      patch_helper(m_instr_bo, Control_Packet_Symbol, 0, m_ctrlpkt_bo.address(),
-                   xrt_core::elf_patcher::buf_type::ctrltext);
+      patch_helper(m_instr_bo, m_ctrlpkt_bo.address(),
+                   xrt_core::elf_patcher::buf_type::ctrltext, Control_Packet_Symbol, {}, false);
     }
 
     // Patch ctrlpkt preemption buffers using config's dynsyms
@@ -507,8 +509,8 @@ class module_run_aie_gen2 : public module_run
       if (bo_itr == m_ctrlpkt_pm_bos.end())
         throw std::runtime_error("Unable to find ctrlpkt pm buffer for symbol " + dynsym);
 
-      patch_helper(m_instr_bo, dynsym, 0, bo_itr->second.address(),
-                   xrt_core::elf_patcher::buf_type::ctrltext);
+      patch_helper(m_instr_bo, bo_itr->second.address(),
+                   xrt_core::elf_patcher::buf_type::ctrltext, dynsym, {}, false);
     }
 
     XRT_DEBUGF("<- module_run_aie_gen2::create_instruction_buf()\n");
@@ -573,22 +575,21 @@ class module_run_aie_gen2 : public module_run
   void
   patch(const std::string& argnm, size_t index, uint64_t value) override
   {
-    bool patched = false;
-
     // patch control-packet buffer
     if (m_ctrlpkt_bo) {
-      if (patch_helper(m_ctrlpkt_bo, argnm, index, value, xrt_core::elf_patcher::buf_type::ctrldata))
-        patched = true;
+      auto type = xrt_core::elf_patcher::buf_type::ctrldata;
+      if (patch_helper(m_ctrlpkt_bo, value, type, argnm, std::to_string(index)))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(argnm, type));
     }
 
     // patch instruction buffer
     if (m_instr_bo) {
-      if (patch_helper(m_instr_bo, argnm, index, value, xrt_core::elf_patcher::buf_type::ctrltext))
-        patched = true;
+      auto type = xrt_core::elf_patcher::buf_type::ctrltext;
+      if (patch_helper(m_instr_bo, value, type, argnm, std::to_string(index)))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(argnm, type));
     }
-
-    if (patched) // if patched, add to patched args set
-      m_patched_args.insert(argnm);
   }
 
   // Sync buffers to device if patching was done
@@ -881,26 +882,49 @@ class module_run_aie_gen2_plus : public module_run
     for (size_t i = 0; i < col_data.size(); ++i) {
       // Find the control-code-* sym-name and patch it in instruction buffer
       auto sym_name = std::string(Control_Code_Symbol) + "-" + std::to_string(i);
-      if (patch_helper(m_buffer, sym_name, std::numeric_limits<size_t>::max(),
-                       m_buffer.address() + offset,
-                       xrt_core::elf_patcher::buf_type::ctrltext))
-        m_patched_args.insert(sym_name);
+      auto type = xrt_core::elf_patcher::buf_type::ctrltext;
+      if (patch_helper(m_buffer, m_buffer.address() + offset, type, sym_name, {}, false))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(sym_name, type));
       offset += col_data[i].size();
     }
 
+    // create scratchpad memory buffer and patch it in ctrlpkt buffers and
+    // instruction buffer if symbol is present in ELF
+    xrt::bo scratchpad_mem;
+    if (m_config.scratch_pad_mem_size > 0) {
+      scratchpad_mem = xrt_core::hw_context_int::get_scratchpad_mem_buf(m_hwctx, m_config.scratch_pad_mem_size);
+      if (!scratchpad_mem)
+        throw std::runtime_error("Failed to get scratchpad buffer from context\n");
+    }
+
     // Patch control packet addresses in instruction buffer
-    for (const auto& [name, ctrlpktbo] : m_ctrlpkt_bos) {
-      // Symbol name is section name without the grp idx
-      // if sec name is .ctrlpkt-57.grp_idx then sym name is .ctrlpkt-57
-      auto dot_pos = name.rfind('.');
-      auto sym_name = (dot_pos != std::string::npos && dot_pos > 0)
-                    ? name.substr(0, dot_pos)
-                    : name;
-
-      if (patch_helper(m_buffer, sym_name, std::numeric_limits<size_t>::max(),
-                       ctrlpktbo.address(),
-                       xrt_core::elf_patcher::buf_type::ctrltext))
-        m_patched_args.insert(sym_name);
+    auto type = xrt_core::elf_patcher::buf_type::buf_type_count;
+    for (auto& [name, ctrlpktbo] : m_ctrlpkt_bos) {
+      auto sym_name =
+          xrt_core::elf_patcher::get_symbol_name_from_section_name(name);
+
+      // Patch scratchpad memory address in ctrlpkt buffer if present
+      if (scratchpad_mem) {
+        type = xrt_core::elf_patcher::buf_type::ctrlpkt;
+        auto symbol = std::string{Scratch_Pad_Mem_Symbol} + sym_name;
+        if (patch_helper(ctrlpktbo, scratchpad_mem.address(), type, symbol, {}, false))
+          m_patched_args.insert(
+              xrt_core::elf_patcher::generate_key_string(symbol, type));
+      }
+
+      type = xrt_core::elf_patcher::buf_type::ctrltext;
+      if (patch_helper(m_buffer, ctrlpktbo.address(), type, sym_name, {}, false))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(sym_name, type));
+    }
+
+    // Patch scratchpad memory address in instruction buffer if present
+    if (scratchpad_mem) {
+      type = xrt_core::elf_patcher::buf_type::ctrltext;
+      if (patch_helper(m_buffer, scratchpad_mem.address(), type, Scratch_Pad_Mem_Symbol, {}, false))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(Scratch_Pad_Mem_Symbol, type));
     }
   }
 
@@ -1015,25 +1039,30 @@ class module_run_aie_gen2_plus : public module_run
   void
   patch(const std::string& argnm, size_t index, uint64_t value) override
   {
-    bool patched = false;
-
+    auto type = xrt_core::elf_patcher::buf_type::ctrltext;
     // patch instruction buffer
-    if (patch_helper(m_buffer, argnm, index, value, xrt_core::elf_patcher::buf_type::ctrltext))
-      patched = true;
+    if (patch_helper(m_buffer, value, type, argnm, std::to_string(index)))
+      m_patched_args.insert(
+          xrt_core::elf_patcher::generate_key_string(argnm, type));
 
     // patch argument in pad section
-    if (patch_helper(m_buffer, argnm, index, value, xrt_core::elf_patcher::buf_type::pad))
-      patched = true;
+    type = xrt_core::elf_patcher::buf_type::pad;
+    if (patch_helper(m_buffer, value, type, argnm, std::to_string(index)))
+      m_patched_args.insert(
+          xrt_core::elf_patcher::generate_key_string(argnm, type));
 
     // New Elfs have multiple ctrlpkt sections
     // Iterate over all ctrlpkt buffers and patch them
-    for (auto& ctrlpkt : m_ctrlpkt_bos) {
-      if (patch_helper(ctrlpkt.second, argnm, index, value, xrt_core::elf_patcher::buf_type::ctrlpkt))
-        patched = true;
+    type = xrt_core::elf_patcher::buf_type::ctrlpkt;
+    for (auto& [name, ctrlpktbo] : m_ctrlpkt_bos) {
+      auto sym_name =
+          xrt_core::elf_patcher::get_symbol_name_from_section_name(name);
+
+      if (patch_helper(ctrlpktbo, value, type, argnm + sym_name,
+                       std::to_string(index) + sym_name))
+        m_patched_args.insert(
+            xrt_core::elf_patcher::generate_key_string(argnm, type));
     }
-
-    if (patched) // if patched, add to patched args set
-      m_patched_args.insert(argnm);
   }
 
   // Sync buffers to device if patching was done