Add header caching to shim_test code to match xrt::run

Signed-off-by: Hayden Laccabue <[email protected]>

Author: hlaccabu

test/shim_test/bo_import_export.cpp

@@ -40,6 +40,7 @@ class test_2proc_export_import_bo : public test_2proc
     auto sdev = get_userpf_device(get_dev_id());
     auto dev = sdev.get();
 
+    io_test_bo_set_base::clear_cmd_header_cache();
     io_test_bo_set boset{dev};
     boset.get_bos()[IO_TEST_BO_INPUT].tbo = std::make_shared<bo>(dev, idata.pid, idata.hdl);
     boset.run();
@@ -77,6 +78,7 @@ TEST_export_import_bo(device::id_type id, std::shared_ptr<device>& sdev, const s
 void
 TEST_export_import_bo_single_proc(device::id_type id, std::shared_ptr<device>& sdev, const std::vector<uint64_t>& arg)
 {
+  io_test_bo_set_base::clear_cmd_header_cache();
   auto dev = sdev.get();
 
   // Create IO test BO set and share input BO with same process

test/shim_test/cmd_fence.cpp

@@ -167,6 +167,7 @@ TEST_cmd_fence_device(device::id_type id, std::shared_ptr<device>& sdev, arg_typ
   // Can't fork with opened device.
   sdev.reset();
 
+  io_test_bo_set_base::clear_cmd_header_cache();
   test_2proc_cmd_fence_device t2p(id);
   t2p.run_test();
 }

test/shim_test/dev_info.cpp

@@ -74,7 +74,7 @@ xclbin_info xclbin_infos[] = {
     .type = KERNEL_TYPE_TXN_FULL_ELF,
   },
   {
-    .name = "nop.xclbin",
+    .name = "nop.elf",
     .device = npu3_device_id,
     .revision_id = npu_any_revision_id,
     .ip_name2idx = {
@@ -96,7 +96,7 @@ xclbin_info xclbin_infos[] = {
     .type = KERNEL_TYPE_TXN_FULL_ELF,
   },
   {
-    .name = "nop.xclbin",
+    .name = "nop.elf",
     .device = npu3_device_id1,
     .revision_id = npu_any_revision_id,
     .ip_name2idx = {

test/shim_test/io.cpp

@@ -224,6 +224,33 @@ elf_init_no_arg_cmd(xrt::elf& elf, cuidx_type idx, bool dump, bo& cmd, bo& inst)
 
 } // namespace
 
+std::map<uint32_t, uint32_t> io_test_bo_set_base::s_cmd_header_cache;
+
+void
+io_test_bo_set_base::
+cache_cmd_header(const buffer_handle *cmd_hdl, const ert_start_kernel_cmd *cmd)
+{
+  auto handle = cmd_hdl->get_properties().kmhdl;
+  s_cmd_header_cache.try_emplace(handle, cmd->header);
+}
+
+void
+io_test_bo_set_base::
+restore_cmd_header(const buffer_handle *cmd_hdl, ert_start_kernel_cmd *cmd)
+{
+  auto handle = cmd_hdl->get_properties().kmhdl;
+  auto it = s_cmd_header_cache.find(handle);
+  if (it != s_cmd_header_cache.end())
+    cmd->header = it->second;
+}
+
+void
+io_test_bo_set_base::
+clear_cmd_header_cache()
+{
+  s_cmd_header_cache.clear();
+}
+
 io_test_bo_set_base::
 io_test_bo_set_base(device* dev, const std::string& xclbin_name) :
   m_bo_array{}
@@ -904,6 +931,11 @@ run(const std::vector<fence_handle*>& wait_fences,
 
   auto cbo = m_bo_array[IO_TEST_BO_CMD].tbo.get();
   auto chdl = cbo->get();
+  auto cpkt = reinterpret_cast<ert_start_kernel_cmd *>(cbo->map());
+  
+  restore_cmd_header(chdl, cpkt);
+  cache_cmd_header(chdl, cpkt);
+
   for (const auto& fence : wait_fences)
     hwq->submit_wait(fence);
   hwq->submit_command(chdl);
@@ -948,32 +980,6 @@ get_preemption_checkpoints()
   return 0;
 }
 
-void
-io_test_bo_set_base::
-cache_cmd_header()
-{
-  auto cbo = m_bo_array[IO_TEST_BO_CMD].tbo.get();
-  if (!cbo)
-    return;
-  auto pkt = reinterpret_cast<ert_packet *>(cbo->map());
-  if (!m_cached_cmd_header)
-    m_cached_cmd_header = pkt->header;
-}
-
-void
-io_test_bo_set_base::
-restore_cmd_header()
-{
-  auto cbo = m_bo_array[IO_TEST_BO_CMD].tbo.get();
-  if (!cbo || !m_cached_cmd_header)
-    return;
-
-  auto pkt = reinterpret_cast<ert_packet *>(cbo->map());
-  pkt->header = m_cached_cmd_header;
-  pkt->state = ERT_CMD_STATE_NEW;
-  std::atomic_thread_fence(std::memory_order_seq_cst);
-}
-
 unsigned long
 elf_preempt_io_test_bo_set::
 get_preemption_checkpoints()

test/shim_test/io.h

@@ -70,15 +70,18 @@ class io_test_bo_set_base
   virtual void
   verify_result();
 
-  void
-  cache_cmd_header();
-
-  void
-  restore_cmd_header();
-
   static const char *
   bo_type2name(int type);
 
+  static void
+  cache_cmd_header(const xrt_core::buffer_handle *cmd_hdl, const ert_start_kernel_cmd *cmd);
+
+  static void
+  restore_cmd_header(const xrt_core::buffer_handle *cmd_hdl, ert_start_kernel_cmd *cmd);
+
+  static void
+  clear_cmd_header_cache();
+
   std::array<io_test_bo, IO_TEST_BO_MAX_TYPES>&
   get_bos();
 
@@ -96,7 +99,7 @@ class io_test_bo_set_base
   const int m_FLAG_OPT =      1 << 1;
   const int m_FLAG_NO_FILL =  1 << 2;
   const int m_FLAG_DEV_BUF =  1 << 3;
-  uint32_t m_cached_cmd_header = 0;
+  static std::map<uint32_t, uint32_t> s_cmd_header_cache;
 
   void
   create_data_bo_from_file(io_test_bo& ibo, const std::string filename, int flags);
@@ -206,4 +209,4 @@ class elf_io_gemm_test_bo_set : public io_test_bo_set_base
   std::unique_ptr<xrt_core::buffer_handle> m_dbo;
 };
 
-#endif // _SHIMTEST_IO_H_
+#endif // _SHIMTEST_IO_H_