Add cmd bo restore header in umq shim

Signed-off-by: Hayden Laccabue <hlaccabu@amd.com>

Author: hlaccabu

src/shim/umq/hwq.cpp

@@ -185,6 +185,14 @@ hwq_umq::
 issue_single_exec_buf(const cmd_buffer *cmd_bo, bool last_of_chain)
 {
   auto cmd = reinterpret_cast<ert_start_kernel_cmd *>(cmd_bo->vaddr());
+  auto handle = cmd_bo->id().handle;
+
+  // Restore header if this is a resubmission
+  auto it = m_cmd_header_cache.find(handle);
+  if (it != m_cmd_header_cache.end()) {
+    cmd->header = it->second;
+  }
+
   auto dpu = get_ert_dpu_data(cmd);
 
   // Sanity check
@@ -198,6 +206,10 @@ issue_single_exec_buf(const cmd_buffer *cmd_bo, bool last_of_chain)
     shim_err(EINVAL, "No dpu data, invalid exec buf");
   }
 
+  if (it == m_cmd_header_cache.end()) {
+    m_cmd_header_cache[handle] = cmd->header;
+  }
+
   auto slot_idx = get_next_avail_slot();
 
   if (get_ert_dpu_data_next(dpu))
@@ -316,6 +328,14 @@ issue_command(const cmd_buffer *cmd_bo)
     return issue_single_exec_buf(cmd_bo, true);
 
   // Runlist command submission.
+  auto handle = cmd_bo->id().handle;
+
+  // Restore chain BO header is this is a resubmission
+  auto it = m_cmd_header_cache.find(handle);
+  if (it != m_cmd_header_cache.end()) {
+    cmd->header = it->second;
+  }
+
   auto payload = get_ert_cmd_chain_data(cmd);
   if (payload->command_count == 0 || payload->command_count > 100)
     shim_err(EINVAL, "Runlist exec buf with bad num of subcmds: %zx", payload->command_count);
@@ -329,6 +349,12 @@ issue_command(const cmd_buffer *cmd_bo)
     seq = issue_single_exec_buf(subcmd, i == payload->command_count - 1);
     subcmds.push_back(subcmd);
   }
+
+  // Cache chain BO header on first submission
+  if (it == m_cmd_header_cache.end()) {
+    m_cmd_header_cache[handle] = cmd->header;
+  }
+
   return seq;
 }
 

src/shim/umq/hwq.h

@@ -56,6 +56,10 @@ class hwq_umq : public hwq
 
   uint64_t
   issue_single_exec_buf(const cmd_buffer *cmd_bo, bool last_of_chain);
+
+  // Cache for command buffer headers (firmware corrupts them)
+  // Key: BO handle, Value: cached header word
+  std::map<uint32_t, uint32_t> m_cmd_header_cache;
 };
 
 }

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

@@ -948,32 +948,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,12 +70,6 @@ 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);
 
@@ -96,7 +90,6 @@ 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;
 
   void
   create_data_bo_from_file(io_test_bo& ibo, const std::string filename, int flags);