rocr/aie: Using PDI address instead of cu_mask for dispatch. Automatic hw ctx reconfiguration upon new PDI addition.

runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -64,8 +64,37 @@ static_assert((sizeof(core::ShareableHandle::handle) >= sizeof(uint32_t)) &&
                   (alignof(core::ShareableHandle::handle) >= alignof(uint32_t)),
               "ShareableHandle cannot store a XDNA handle");
 
+/// @brief The number of arguments in the packet payload before we start passing operands
+constexpr uint32_t non_operand_count = 6;
+
 // Index where the operand addresses start in a command.
-constexpr uint32_t operand_starting_index = 5;
+constexpr uint32_t operand_starting_index = non_operand_count - 1;
+
+/// @brief Default amdxdna_cu_config::cu_func when configuring a CU.
+constexpr uint32_t default_cu_func = 0;
+
+/// @brief Calculates the number of operands in a packet given the number of arguments in the
+/// packet.
+/// @param arg_count number of arguments in the packet
+/// @return number of operands in the packet
+static uint32_t GetOperandCount(uint32_t arg_count) {
+  return ((arg_count - non_operand_count) / 2);
+}
+
+/// @brief Flushes operands.
+static void FlushOperands(uint32_t count, hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload) {
+  // Going through all of the operands in the command and flushing them.
+  const uint32_t num_operands = GetOperandCount(count);
+  for (uint32_t operand_iter = 0; operand_iter < num_operands; operand_iter++) {
+    const uint32_t operand_index = operand_starting_index + 2 * operand_iter;
+    const uint64_t operand_addr = Concat<uint64_t>(cmd_pkt_payload->data[operand_index + 1],
+                                                   cmd_pkt_payload->data[operand_index]);
+    const uint32_t operand_size_starting_index = operand_starting_index + 2 * num_operands;
+    const uint32_t operand_bo_size =
+        cmd_pkt_payload->data[operand_size_starting_index + operand_iter];
+    FlushCpuCache(reinterpret_cast<void*>(operand_addr), 0, operand_bo_size);
+  }
+}
 
 XdnaDriver::XdnaDriver(std::string devnode_name)
     : core::Driver(core::DriverType::XDNA, std::move(devnode_name)) {}
@@ -217,12 +246,7 @@ XdnaDriver::AllocateMemory(const core::MemoryRegion &mem_region,
   bo_handle.size = size;
 
   // Close the BO in case of error.
-  MAKE_NAMED_SCOPE_GUARD(bo_guard, [&] {
-    munmap(bo_handle.vaddr, bo_handle.size);
-    drm_gem_close close_bo_args = {};
-    close_bo_args.handle = bo_handle.handle;
-    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
-  });
+  MAKE_NAMED_SCOPE_GUARD(bo_guard, [&] { DestroyBOHandle(bo_handle); });
 
   amdxdna_drm_get_bo_info get_bo_info_args = {};
   get_bo_info_args.handle = create_bo_args.handle;
@@ -280,21 +304,9 @@ hsa_status_t XdnaDriver::CreateQueue(core::Queue &queue) const {
     return HSA_STATUS_ERROR_INVALID_QUEUE;
   }
 
-  auto &aie_queue(static_cast<AieAqlQueue &>(queue));
-  auto &aie_agent(aie_queue.GetAgent());
-
-  // Currently we do not leverage QoS information.
-  amdxdna_qos_info qos_info{0};
-  amdxdna_drm_create_hwctx create_hwctx_args = {};
-  create_hwctx_args.qos_p = reinterpret_cast<uintptr_t>(&qos_info);
-  create_hwctx_args.max_opc = 0x800;
-  create_hwctx_args.num_tiles = static_cast<uint32_t>(aie_agent.GetNumCores());
-
-  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_HWCTX, &create_hwctx_args) < 0) {
-    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
-  }
-
-  aie_queue.SetHwCtxHandle(create_hwctx_args.handle);
+  // Set the hw ctx handle of the queue to invalid to avoid incorrect destruction.
+  auto& aie_queue = static_cast<AieAqlQueue&>(queue);
+  aie_queue.SetHwCtxHandle(AMDXDNA_INVALID_BO_HANDLE);
 
   return HSA_STATUS_SUCCESS;
 }
@@ -398,13 +410,7 @@ hsa_status_t XdnaDriver::InitDeviceHeap() {
   dev_heap_handle.handle = create_bo_args.handle;
 
   // Unmap memory and close the BO in case of error.
-  MAKE_NAMED_SCOPE_GUARD(dev_heap_handle_guard, [&] {
-    munmap(dev_heap_handle.vaddr, dev_heap_handle.size);
-    drm_gem_close close_bo_args = {};
-    close_bo_args.handle = dev_heap_handle.handle;
-    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
-    dev_heap_handle = BOHandle{};
-  });
+  MAKE_NAMED_SCOPE_GUARD(dev_heap_handle_guard, [&] { DestroyBOHandle(dev_heap_handle); });
 
   amdxdna_drm_get_bo_info get_bo_info_args = {};
   get_bo_info_args.handle = dev_heap_handle.handle;
@@ -458,36 +464,34 @@ hsa_status_t XdnaDriver::FreeDeviceHeap() {
   return status;
 }
 
-hsa_status_t XdnaDriver::SyncBos(const std::vector<uint64_t>& bo_addrs,
-                                 const std::vector<uint32_t>& bo_sizes) {
-  if (bo_addrs.size() != bo_sizes.size()) return HSA_STATUS_ERROR;
+hsa_status_t XdnaDriver::ExecCmdAndWait(const BOHandle& cmd_chain_bo_handle,
+                                        const std::vector<uint32_t>& bo_handles,
+                                        AieAqlQueue& aie_queue) {
+  // Submit command chain.
+  amdxdna_drm_exec_cmd exec_cmd = {};
+  exec_cmd.hwctx = aie_queue.GetHwCtxHandle();
+  exec_cmd.type = AMDXDNA_CMD_SUBMIT_EXEC_BUF;
+  exec_cmd.cmd_handles = cmd_chain_bo_handle.handle;
+  exec_cmd.args = reinterpret_cast<uint64_t>(bo_handles.data());
+  exec_cmd.cmd_count = 1;
+  exec_cmd.arg_count = bo_handles.size();
 
-  for (int i = 0; i < bo_addrs.size(); i++) {
-    FlushCpuCache(reinterpret_cast<void*>(bo_addrs[i]), 0, bo_sizes[i]);
-  }
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_EXEC_CMD, &exec_cmd) < 0) return HSA_STATUS_ERROR;
 
-  return HSA_STATUS_SUCCESS;
-}
-
-hsa_status_t XdnaDriver::ExecCmdAndWait(amdxdna_drm_exec_cmd* exec_cmd, uint32_t hw_ctx_handle) {
-  // Submit the cmd
-  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_EXEC_CMD, exec_cmd)) return HSA_STATUS_ERROR;
-
-  // Waiting for command to finish
+  // Waiting for command chain to finish.
   amdxdna_drm_wait_cmd wait_cmd = {};
-  wait_cmd.hwctx = hw_ctx_handle;
+  wait_cmd.hwctx = aie_queue.GetHwCtxHandle();
   wait_cmd.timeout = DEFAULT_TIMEOUT_VAL;
-  wait_cmd.seq = exec_cmd->seq;
+  wait_cmd.seq = exec_cmd.seq;
 
-  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_WAIT_CMD, &wait_cmd)) return HSA_STATUS_ERROR;
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_WAIT_CMD, &wait_cmd) < 0) return HSA_STATUS_ERROR;
 
   return HSA_STATUS_SUCCESS;
 }
 
-hsa_status_t XdnaDriver::RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& bo_args,
-                                        std::vector<uint32_t>& bo_sizes,
-                                        std::vector<uint64_t>& bo_addrs,
-                                        hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload) {
+hsa_status_t XdnaDriver::PrepareBOs(uint32_t count,
+                                    hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload,
+                                    std::vector<uint32_t>& bo_handles) {
   const uint64_t instr_addr =
       Concat<uint64_t>(cmd_pkt_payload->data[CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_IDX + 1],
                        cmd_pkt_payload->data[CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_IDX]);
@@ -496,21 +500,20 @@ hsa_status_t XdnaDriver::RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& b
     return HSA_STATUS_ERROR;
   }
 
-  // Keep track of the handles and addresses before we submit the packet
-  bo_args.push_back(instr_bo_handle.handle);
-  bo_addrs.push_back(reinterpret_cast<uint64_t>(instr_bo_handle.vaddr));
+  // Keep track of the instruction sequence BO.
+  bo_handles.push_back(instr_bo_handle.handle);
 
-  // Adding the instruction sequence size. The packet contains the number of
-  // instructions.
-  uint32_t instr_bo_size =
+  // Flush the instruction sequence. The packet contains the number of instructions.
+  const uint32_t instr_bo_size =
       cmd_pkt_payload->data[CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_SIZE_IDX] * INSTR_SIZE_BYTES;
-  bo_sizes.push_back(instr_bo_size);
+  FlushCpuCache(reinterpret_cast<void*>(instr_addr), 0, instr_bo_size);
 
   // Going through all of the operands in the command, keeping track of the
   // addresses and turning the addresses into handles. The starting index of
   // the operands in a command is `operand_starting_index` and the fields
   // are 32-bits we need to iterate over every two
   const uint32_t num_operands = GetOperandCount(count);
+  bo_handles.reserve(num_operands);
   for (uint32_t operand_iter = 0; operand_iter < num_operands; operand_iter++) {
     const uint32_t operand_index = operand_starting_index + 2 * operand_iter;
     const uint64_t operand_addr = Concat<uint64_t>(cmd_pkt_payload->data[operand_index + 1],
@@ -519,15 +522,15 @@ hsa_status_t XdnaDriver::RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& b
     if (!operand_bo_handle.IsValid()) {
       return HSA_STATUS_ERROR;
     }
-    bo_args.push_back(operand_bo_handle.handle);
-    bo_addrs.push_back(reinterpret_cast<uint64_t>(operand_bo_handle.vaddr));
-  }
 
-  // Going through all of the operands in the command, keeping track of
-  // the sizes of each operand. The size is used to sync the buffer
-  uint32_t operand_size_starting_index = operand_starting_index + 2 * num_operands;
-  for (int operand_iter = 0; operand_iter < num_operands; operand_iter++) {
-    bo_sizes.push_back(cmd_pkt_payload->data[operand_size_starting_index + operand_iter]);
+    // Keep track of the operand BO.
+    bo_handles.push_back(operand_bo_handle.handle);
+
+    // Flush the operand.
+    const uint32_t operand_size_starting_index = operand_starting_index + 2 * num_operands;
+    const uint32_t operand_bo_size =
+        cmd_pkt_payload->data[operand_size_starting_index + operand_iter];
+    FlushCpuCache(reinterpret_cast<void*>(operand_addr), 0, operand_bo_size);
   }
 
   // Transform the instruction sequence address into device address
@@ -537,45 +540,59 @@ hsa_status_t XdnaDriver::RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& b
   return HSA_STATUS_SUCCESS;
 }
 
-hsa_status_t XdnaDriver::CreateCmd(uint32_t size, uint32_t* handle, amdxdna_cmd** cmd) {
-  // Creating the command
+hsa_status_t XdnaDriver::CreateCmdBO(uint32_t size, BOHandle& cmd_bo_handle) {
   amdxdna_drm_create_bo create_cmd_bo = {};
-  create_cmd_bo.type = AMDXDNA_BO_CMD, create_cmd_bo.size = size;
-  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_BO, &create_cmd_bo)) return HSA_STATUS_ERROR;
+  create_cmd_bo.type = AMDXDNA_BO_CMD;
+  create_cmd_bo.size = size;
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_BO, &create_cmd_bo) < 0) {
+    return HSA_STATUS_ERROR;
+  }
+
+  // Close the BO in case of error.
+  MAKE_NAMED_SCOPE_GUARD(cmd_bo_handle_guard, [&] {
+    drm_gem_close close_bo_args = {};
+    close_bo_args.handle = create_cmd_bo.handle;
+    ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+  });
 
   amdxdna_drm_get_bo_info cmd_bo_get_bo_info = {};
   cmd_bo_get_bo_info.handle = create_cmd_bo.handle;
-  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_GET_BO_INFO, &cmd_bo_get_bo_info)) return HSA_STATUS_ERROR;
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_GET_BO_INFO, &cmd_bo_get_bo_info) < 0) {
+    return HSA_STATUS_ERROR;
+  }
 
-  *cmd = static_cast<amdxdna_cmd*>(mmap(nullptr, create_cmd_bo.size, PROT_READ | PROT_WRITE,
-                                        MAP_SHARED, fd_, cmd_bo_get_bo_info.map_offset));
+  void* mem = static_cast<amdxdna_cmd*>(mmap(nullptr, create_cmd_bo.size, PROT_READ | PROT_WRITE,
+                                             MAP_SHARED, fd_, cmd_bo_get_bo_info.map_offset));
+  if (mem == MAP_FAILED) {
+    return HSA_STATUS_ERROR;
+  }
 
-  if (cmd == MAP_FAILED) return HSA_STATUS_ERROR;
+  cmd_bo_handle = BOHandle{mem, create_cmd_bo.handle, size};
 
-  *handle = create_cmd_bo.handle;
+  cmd_bo_handle_guard.Dismiss();
 
   return HSA_STATUS_SUCCESS;
 }
 
 hsa_status_t XdnaDriver::SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uint32_t num_pkts,
-                                        uint32_t num_operands, uint32_t hw_ctx_handle) {
-  // Storing the metadata of the BOs that store the operands and metadata
-  // of the commands we are going to submit
-  std::vector<uint32_t> bo_args;
-  std::vector<uint32_t> bo_sizes;
-  std::vector<uint64_t> bo_addrs;
-  bo_args.reserve(num_operands);
-  bo_sizes.reserve(num_operands);
-  bo_addrs.reserve(num_operands);
+                                        AieAqlQueue& aie_queue) {
+  // Stores instruction and operand BOs.
+  std::vector<uint32_t> bo_handles;
 
-  // Storing the commands that we are going to submit and the
-  // corresponding metadata
-  std::vector<uint32_t> cmd_handles;
-  std::vector<uint32_t> cmd_sizes;
-  std::vector<amdxdna_cmd*> cmds;
-  cmd_handles.reserve(num_pkts);
-  cmd_sizes.reserve(num_pkts);
-  cmds.reserve(num_pkts);
+  // Stores commands that we are going to submit and the corresponding metadata.
+  std::vector<BOHandle> cmd_bo_handles;
+  cmd_bo_handles.reserve(num_pkts);
+  // Unmap and close the command BOs in case of an error.
+  MAKE_NAMED_SCOPE_GUARD(cmd_bo_handles_guard, [&] {
+    for (auto& bo_handle : cmd_bo_handles) {
+      DestroyBOHandle(bo_handle);
+    }
+  });
+
+  // PDI cache. If the cache is updated, a new hardware context will be created for the queue.
+  auto pdi_cache_it = hw_ctx_pdi_cache_map.find(aie_queue.GetHwCtxHandle());
+  auto pdi_cache = (pdi_cache_it != hw_ctx_pdi_cache_map.end()) ? pdi_cache_it->second : PDICache{};
+  bool reconfigure_queue = false;
 
   // Iterating over all the contiguous HSA_AMD_AIE_ERT_CMD_CHAIN packets
   for (uint32_t pkt_iter = 0; pkt_iter < num_pkts; pkt_iter++) {
@@ -584,17 +601,24 @@ hsa_status_t XdnaDriver::SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uin
     hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload =
         reinterpret_cast<hsa_amd_aie_ert_start_kernel_data_t*>(pkt->payload_data);
 
-    // Add the handles for all of the BOs to bo_args as well as rewrite
-    // the command payload handles to contain the actual virtual addresses
-    if (RegisterCmdBOs(pkt->count, bo_args, bo_sizes, bo_addrs, cmd_pkt_payload) !=
-        HSA_STATUS_SUCCESS)
-      return HSA_STATUS_ERROR;
+    // Add the handles for all of the BOs to bo_handles as well as rewrite
+    // the instruction handle to contain the device address
+    hsa_status_t status = PrepareBOs(pkt->count, cmd_pkt_payload, bo_handles);
+    if (status != HSA_STATUS_SUCCESS) {
+      return status;
+    }
 
     // Creating a packet that contains the command to execute the kernel
-    uint32_t cmd_bo_handle = 0;
-    amdxdna_cmd* cmd = nullptr;
-    uint32_t cmd_size = sizeof(amdxdna_cmd) + pkt->count * sizeof(uint32_t);
-    if (CreateCmd(cmd_size, &cmd_bo_handle, &cmd)) return HSA_STATUS_ERROR;
+    const uint32_t cmd_size = sizeof(amdxdna_cmd) + pkt->count * sizeof(uint32_t);
+    BOHandle cmd_bo_handle;
+    status = CreateCmdBO(cmd_size, cmd_bo_handle);
+    if (status != HSA_STATUS_SUCCESS) {
+      return status;
+    }
+    // Unmap and close the command BO in case of an error.
+    MAKE_NAMED_SCOPE_GUARD(cmd_bo_handle_guard, [&] { DestroyBOHandle(cmd_bo_handle); });
+
+    auto* cmd = static_cast<amdxdna_cmd*>(cmd_bo_handle.vaddr);
 
     // Filling in the fields of the command
     cmd->state = pkt->state;
@@ -604,20 +628,58 @@ hsa_status_t XdnaDriver::SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uin
     // Need to increase the size of the command by the size of this structure.
     cmd->count = pkt->count + CMD_COUNT_SIZE_INCREASE;
     cmd->opcode = pkt->opcode;
-    cmd->data[0] = cmd_pkt_payload->cu_mask;
+
+    // Find if the PDI is cached in the queues PDI cache. If even one PDI is not found, the hardware
+    // context will need to be reconfigured and the cache updated.
+    auto pdi_bo_handle = FindBOHandle(cmd_pkt_payload->pdi_addr);
+    if (!pdi_bo_handle.IsValid()) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+    // Determining if the PDI is cached
+    auto cached_pdi_index = pdi_cache.GetIndex(pdi_bo_handle.handle);
+    if (cached_pdi_index == PDICache::NotFound) {
+      // PDI does not exist in the cache.
+      status = pdi_cache.SetNext(pdi_bo_handle, cached_pdi_index);
+      if (status != HSA_STATUS_SUCCESS) {
+        return status;
+      }
+      reconfigure_queue = true;
+    }
+
+    cmd->data[0] = 0x1 << static_cast<uint32_t>(cached_pdi_index);
     memcpy((cmd->data + 1), cmd_pkt_payload->data, 4 * pkt->count);
 
-    // Keeping track of the handle
-    cmd_handles.push_back(cmd_bo_handle);
-    cmds.push_back(cmd);
-    cmd_sizes.push_back(cmd_size);
+    // Keeping track of the command
+    cmd_bo_handles.push_back(cmd_bo_handle);
+    cmd_bo_handle_guard.Dismiss();
+  }
+
+  // If there were PDIs that were not cached, the hardware context needs to be reconfigured.
+  // The cache map will be update with the new hardware context.
+  if (reconfigure_queue) {
+    if (pdi_cache_it != hw_ctx_pdi_cache_map.end()) {
+      hw_ctx_pdi_cache_map.erase(pdi_cache_it);
+    }
+
+    hsa_status_t status = ConfigHwCtx(pdi_cache, aie_queue);
+    if (status != HSA_STATUS_SUCCESS) {
+      return status;
+    }
+
+    // Update cache mapping.
+    hw_ctx_pdi_cache_map.emplace(aie_queue.GetHwCtxHandle(), pdi_cache);
   }
 
   // Creating a packet that contains the command chain
-  uint32_t cmd_chain_bo_handle = 0;
-  amdxdna_cmd* cmd_chain = nullptr;
-  int cmd_chain_size = (cmd_handles.size() + 1) * sizeof(uint32_t);
-  if (CreateCmd(cmd_chain_size, &cmd_chain_bo_handle, &cmd_chain)) return HSA_STATUS_ERROR;
+  const uint32_t cmd_chain_size = (cmd_bo_handles.size() + 1) * sizeof(uint32_t);
+  BOHandle cmd_chain_bo_handle;
+  hsa_status_t status = CreateCmdBO(cmd_chain_size, cmd_chain_bo_handle);
+  if (status != HSA_STATUS_SUCCESS) {
+    return status;
+  }
+  // Unmap and close the command chain BO in case of an error.
+  MAKE_NAMED_SCOPE_GUARD(cmd_chain_bo_handle_guard, [&] { DestroyBOHandle(cmd_chain_bo_handle); });
+
+  auto* cmd_chain = static_cast<amdxdna_cmd*>(cmd_chain_bo_handle.vaddr);
 
   // Writing information to the command buffer
   amdxdna_cmd_chain* cmd_chain_payload = reinterpret_cast<amdxdna_cmd_chain*>(cmd_chain->data);
@@ -625,53 +687,57 @@ hsa_status_t XdnaDriver::SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uin
   // Creating a command chain
   cmd_chain->state = HSA_AMD_AIE_ERT_STATE_NEW;
   cmd_chain->extra_cu_masks = 0;
-  cmd_chain->count = sizeof(amdxdna_cmd_chain) + cmd_handles.size() * sizeof(uint64_t);
+  cmd_chain->count = sizeof(amdxdna_cmd_chain) + cmd_bo_handles.size() * sizeof(uint64_t);
   cmd_chain->opcode = HSA_AMD_AIE_ERT_CMD_CHAIN;
-  cmd_chain_payload->command_count = cmd_handles.size();
+  cmd_chain_payload->command_count = cmd_bo_handles.size();
   cmd_chain_payload->submit_index = 0;
   cmd_chain_payload->error_index = 0;
-  for (int i = 0; i < cmd_handles.size(); i++) {
-    cmd_chain_payload->data[i] = cmd_handles[i];
+  for (size_t i = 0; i < cmd_bo_handles.size(); i++) {
+    cmd_chain_payload->data[i] = cmd_bo_handles[i].handle;
   }
 
-  // Syncing BOs before we execute the command
-  if (SyncBos(bo_addrs, bo_sizes)) return HSA_STATUS_ERROR;
-
   // Removing duplicates in the bo container. The driver will report
   // an error if we provide the same BO handle multiple times.
   // This can happen if any of the BOs are the same across jobs
-  std::sort(bo_args.begin(), bo_args.end());
-  bo_args.erase(std::unique(bo_args.begin(), bo_args.end()), bo_args.end());
-
-  // Filling in the fields to execute the command chain
-  amdxdna_drm_exec_cmd exec_cmd_0 = {};
-  exec_cmd_0.hwctx = hw_ctx_handle;
-  exec_cmd_0.type = AMDXDNA_CMD_SUBMIT_EXEC_BUF;
-  exec_cmd_0.cmd_handles = cmd_chain_bo_handle;
-  exec_cmd_0.args = reinterpret_cast<uint64_t>(bo_args.data());
-  exec_cmd_0.cmd_count = 1;
-  exec_cmd_0.arg_count = bo_args.size();
+  std::sort(bo_handles.begin(), bo_handles.end());
+  bo_handles.erase(std::unique(bo_handles.begin(), bo_handles.end()), bo_handles.end());
 
   // Executing all commands in the command chain
-  ExecCmdAndWait(&exec_cmd_0, hw_ctx_handle);
+  status = ExecCmdAndWait(cmd_chain_bo_handle, bo_handles, aie_queue);
+  if (status != HSA_STATUS_SUCCESS) {
+    return status;
+  }
+
+  for (uint32_t pkt_iter = 0; pkt_iter < num_pkts; pkt_iter++) {
+    hsa_amd_aie_ert_packet_t* pkt = first_pkt + pkt_iter;
+    auto* cmd_pkt_payload =
+        reinterpret_cast<hsa_amd_aie_ert_start_kernel_data_t*>(pkt->payload_data);
+    FlushOperands(pkt->count, cmd_pkt_payload);
+  }
+
+  status = HSA_STATUS_SUCCESS;
 
   // Unmapping and closing the cmd BOs
-  drm_gem_close close_bo_args{0};
-  for (int i = 0; i < cmd_handles.size(); i++) {
-    if (munmap(cmds[i], cmd_sizes[i]) != 0) return HSA_STATUS_ERROR;
-    close_bo_args.handle = cmd_handles[i];
+  cmd_bo_handles_guard.Dismiss();
+  for (auto& command_bo_handle : cmd_bo_handles) {
+    if (munmap(command_bo_handle.vaddr, command_bo_handle.size) != 0) {
+      status = HSA_STATUS_ERROR;
+    }
+    drm_gem_close close_bo_args = {};
+    close_bo_args.handle = command_bo_handle.handle;
     ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
   }
 
   // Unmapping and closing the cmd_chain BO
-  if (munmap(cmd_chain, cmd_chain_size) != 0) return HSA_STATUS_ERROR;
-  close_bo_args.handle = cmd_chain_bo_handle;
+  cmd_chain_bo_handle_guard.Dismiss();
+  if (munmap(cmd_chain, cmd_chain_size) != 0) {
+    status = HSA_STATUS_ERROR;
+  }
+  drm_gem_close close_bo_args = {};
+  close_bo_args.handle = cmd_chain_bo_handle.handle;
   ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
 
-  // Syncing BOs after we execute the command
-  if (SyncBos(bo_addrs, bo_sizes)) return HSA_STATUS_ERROR;
-
-  return HSA_STATUS_SUCCESS;
+  return status;
 }
 
 hsa_status_t XdnaDriver::SPMAcquire(uint32_t preferred_node_id) const {
@@ -697,6 +763,14 @@ hsa_status_t XdnaDriver::IsModelEnabled(bool* enable) const {
   return HSA_STATUS_SUCCESS;
 }
 
+void XdnaDriver::DestroyBOHandle(BOHandle& handle) {
+  munmap(handle.vaddr, handle.size);
+  drm_gem_close close_bo_args = {};
+  close_bo_args.handle = handle.handle;
+  ioctl(fd_, DRM_IOCTL_GEM_CLOSE, &close_bo_args);
+  handle = {};
+}
+
 XdnaDriver::BOHandle XdnaDriver::FindBOHandle(void* mem) const {
   auto it = vmem_addr_mappings.lower_bound(mem);
   if (it == vmem_addr_mappings.cend()) {
@@ -727,6 +801,67 @@ XdnaDriver::BOHandle XdnaDriver::FindBOHandle(void* mem) const {
   return it->second;
 }
 
+hsa_status_t XdnaDriver::ConfigHwCtx(const PDICache& pdi_bo_handles, AieAqlQueue& aie_queue) {
+  const size_t config_cu_param_size =
+      sizeof(amdxdna_hwctx_param_config_cu) + pdi_bo_handles.size() * sizeof(amdxdna_cu_config);
+
+  auto* xdna_config_cu_param =
+      static_cast<amdxdna_hwctx_param_config_cu*>(malloc(config_cu_param_size));
+  if (xdna_config_cu_param == nullptr) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+  MAKE_SCOPE_GUARD([xdna_config_cu_param] { free(xdna_config_cu_param); });
+
+  xdna_config_cu_param->num_cus = pdi_bo_handles.size();
+
+  for (size_t i = 0; i < pdi_bo_handles.size(); i++) {
+    xdna_config_cu_param->cu_configs[i].cu_bo = pdi_bo_handles[i].handle;
+    xdna_config_cu_param->cu_configs[i].cu_func = default_cu_func;
+
+    // Flush the PDI out of the cache
+    FlushCpuCache(pdi_bo_handles[i].vaddr, 0, pdi_bo_handles[i].size);
+  }
+
+  if (aie_queue.GetHwCtxHandle() != AMDXDNA_INVALID_BO_HANDLE) {
+    // Destroy the hardware context
+    // Note: we can do this because we have forced synchronization between
+    // command chains. If we move to a more asynchronous model, we will need to
+    // figure out how hardware context destruction works while applications
+    // are running
+    amdxdna_drm_destroy_hwctx destroy_hwctx_args = {};
+    destroy_hwctx_args.handle = aie_queue.GetHwCtxHandle();
+    if (ioctl(fd_, DRM_IOCTL_AMDXDNA_DESTROY_HWCTX, &destroy_hwctx_args) < 0) {
+      return HSA_STATUS_ERROR;
+    }
+  }
+
+  // Create the new hardware context
+  // Currently we do not leverage QoS information.
+  amdxdna_qos_info qos_info = {};
+  amdxdna_drm_create_hwctx create_hwctx_args = {};
+  create_hwctx_args.qos_p = reinterpret_cast<uintptr_t>(&qos_info);
+  create_hwctx_args.max_opc = 0x800;
+  create_hwctx_args.num_tiles = aie_queue.GetAgent().GetNumCores();
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CREATE_HWCTX, &create_hwctx_args) < 0) {
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
+  // Configure the new hardware context
+  amdxdna_drm_config_hwctx config_hw_ctx_args = {};
+  config_hw_ctx_args.handle = create_hwctx_args.handle;
+  config_hw_ctx_args.param_type = DRM_AMDXDNA_HWCTX_CONFIG_CU;
+  config_hw_ctx_args.param_val = reinterpret_cast<uint64_t>(xdna_config_cu_param);
+  config_hw_ctx_args.param_val_size = static_cast<uint32_t>(config_cu_param_size);
+
+  if (ioctl(fd_, DRM_IOCTL_AMDXDNA_CONFIG_HWCTX, &config_hw_ctx_args) < 0) {
+    return HSA_STATUS_ERROR;
+  }
+
+  aie_queue.SetHwCtxHandle(create_hwctx_args.handle);
+
+  return HSA_STATUS_SUCCESS;
+}
 
 } // namespace AMD
 } // namespace rocr

runtime/hsa-runtime/core/inc/amd_xdna_driver.h

@@ -42,6 +42,8 @@
 #ifndef HSA_RUNTIME_CORE_INC_AMD_XDNA_DRIVER_H_
 #define HSA_RUNTIME_CORE_INC_AMD_XDNA_DRIVER_H_
 
+#include <array>
+#include <climits>
 #include <map>
 #include <memory>
 #include <unordered_map>
@@ -93,8 +95,7 @@ class Queue;
 
 namespace AMD {
 
-/// @brief: The number of arguments in the packet payload before we start passing operands
-constexpr uint32_t NON_OPERAND_COUNT = 6;
+class AieAqlQueue;
 
 // @brief: Used to transform an address into a device address
 constexpr uint32_t DEV_ADDR_BASE = 0x04000000;
@@ -119,14 +120,6 @@ constexpr uint32_t CMD_PKT_PAYLOAD_INSTRUCTION_SEQUENCE_SIZE_IDX = 4;
 /// @brief Environment variable to define job submission timeout
 constexpr uint32_t DEFAULT_TIMEOUT_VAL = 50;
 
-/// @brief: Calculates the number of operands in a packet
-/// given the number of arguments in the packet
-/// @param: arg_count(Input), Number of arguments in the packet
-/// @return: uint32_t, The number of operands in the packet
-inline uint32_t GetOperandCount(uint32_t arg_count) {
-  return ((arg_count - NON_OPERAND_COUNT) / 2);
-}
-
 class XdnaDriver final : public core::Driver {
   /// @brief BO handle information.
   struct BOHandle {
@@ -143,6 +136,48 @@ class XdnaDriver final : public core::Driver {
     constexpr bool IsValid() const { return handle != AMDXDNA_INVALID_BO_HANDLE; }
   };
 
+  /// @brief CU mask size.
+  static constexpr size_t cu_mask_size = sizeof(uint32_t) * CHAR_BIT;
+
+  /// @brief Per hardware context PDI cache.
+  class PDICache {
+    std::array<BOHandle, cu_mask_size> entries = {};
+    size_t entry_count = 0;
+
+   public:
+    /// @brief Sentinel value for entries not found.
+    constexpr static size_t NotFound = cu_mask_size;
+
+    /// @brief Returns the size of the cache.
+    constexpr size_t size() const { return entry_count; }
+
+    /// @brief Returns the index of the BO handle if it is the cache, otherwise @ref NotFound.
+    ///
+    /// This function does a linear search because the mask is small (32 elements).
+    size_t GetIndex(uint32_t pdi_handle) const {
+      for (size_t i = 0; i < entry_count; ++i) {
+        if (entries[i].handle == pdi_handle) {
+          return i;
+        }
+      }
+      return NotFound;
+    }
+
+    /// @brief Sets the next cache entry.
+    hsa_status_t SetNext(const BOHandle& pdi_bo_handle, size_t& index) {
+      if (entry_count == entries.size()) {
+        // cache is full
+        return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+      }
+
+      index = entry_count++;
+      entries[index] = pdi_bo_handle;
+      return HSA_STATUS_SUCCESS;
+    }
+
+    constexpr const BOHandle& operator[](size_t index) const { return entries[index]; }
+  };
+
 public:
   XdnaDriver(std::string devnode_name);
 
@@ -188,7 +223,7 @@ public:
 
   /// @brief Submits @p num_pkts packets in a command chain.
   hsa_status_t SubmitCmdChain(hsa_amd_aie_ert_packet_t* first_pkt, uint32_t num_pkts,
-                              uint32_t num_operands, uint32_t hw_ctx_handle);
+                              AieAqlQueue& aie_queue);
 
   hsa_status_t SPMAcquire(uint32_t preferred_node_id) const override;
   hsa_status_t SPMRelease(uint32_t preferred_node_id) const override;
@@ -199,9 +234,17 @@ public:
   hsa_status_t IsModelEnabled(bool* enable) const override;
 
  private:
+  /// @brief Destroys @p bo_handle.
+  ///
+  /// This function will unmap the virtual address and close the BO, but will not return any status.
+  void DestroyBOHandle(BOHandle& bo_handle);
+
   /// @brief Finds the BO associated with the address.
   BOHandle FindBOHandle(void* mem) const;
 
+  /// @brief Creates a new hardware context with the given PDI BO handles.
+  hsa_status_t ConfigHwCtx(const PDICache& pdi_bo_handles, AieAqlQueue& aie_queue);
+
   hsa_status_t QueryDriverVersion();
 
   /// @brief Allocate device accesible heap space.
@@ -210,35 +253,28 @@ public:
   hsa_status_t InitDeviceHeap();
   hsa_status_t FreeDeviceHeap();
 
-  /// @brief Creates a command BO and returns a pointer to the memory and
-  //          the corresponding handle
+  /// @brief Creates a command BO and returns it to @p bo_info.
   ///
   /// @param size size of memory to allocate
-  /// @param handle A pointer to the BO handle
-  /// @param cmd A pointer to the buffer
-  hsa_status_t CreateCmd(uint32_t size, uint32_t* handle, amdxdna_cmd** cmd);
+  /// @param bo_info allocated BO
+  hsa_status_t CreateCmdBO(uint32_t size, BOHandle& bo_info);
 
-  /// @brief Adds all BOs in a command packet payload to a vector
-  ///         and replaces the handles with a virtual address
+  /// @brief Gets all BOs from a command packet payload, flushes the caches associated with them and
+  /// replaces the instruction virtual address with the device address.
   ///
   /// @param count Number of entries in the command
-  /// @param bo_args A pointer to a vector that contains all bo handles
   /// @param cmd_pkt_payload A pointer to the payload of the command
-  hsa_status_t RegisterCmdBOs(uint32_t count, std::vector<uint32_t>& bo_args,
-                              std::vector<uint32_t>& bo_sizes, std::vector<uint64_t>& bo_addrs,
-                              hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload);
-
-  /// @brief Syncs all BOs referenced in bo_args
-  ///
-  /// @param bo_args vector containing handles of BOs to sync
-  hsa_status_t SyncBos(const std::vector<uint64_t>& bo_args, const std::vector<uint32_t>& bo_sizes);
+  /// @param bo_handles vector that contains all BO handles
+  hsa_status_t PrepareBOs(uint32_t count, hsa_amd_aie_ert_start_kernel_data_t* cmd_pkt_payload,
+                          std::vector<uint32_t>& bo_handles);
 
   /// @brief Executes a command and waits for its completion
   ///
-  /// @param exec_cmd Structure containing the details of the command to execute
-  /// @param hw_ctx_handle the handle of the hardware context to run this
-  /// command
-  hsa_status_t ExecCmdAndWait(amdxdna_drm_exec_cmd* exec_cmd, uint32_t hw_ctx_handle);
+  /// @param cmd_chain_bo_handle command to execute
+  /// @param bo_handles handles associated with the command
+  /// @param aie_queue queue to submit to
+  hsa_status_t ExecCmdAndWait(const BOHandle& cmd_chain_bo_handle,
+                              const std::vector<uint32_t>& bo_handles, AieAqlQueue& aie_queue);
 
   /// TODO: Remove this in the future and rely on the core Runtime
   /// object to track handle allocations. Using the VMEM API for mapping XDNA
@@ -247,6 +283,9 @@ public:
   std::unordered_map<uint32_t, void *> vmem_handle_mappings;
   std::map<void*, BOHandle> vmem_addr_mappings;
 
+  /// @brief Hardware context to PDI cache mapping.
+  std::unordered_map<uint32_t, PDICache> hw_ctx_pdi_cache_map;
+
   /// @brief Virtual address range allocated for the device heap.
   ///
   /// Allocate a large enough space so we can carve out the device heap in

runtime/hsa-runtime/core/runtime/amd_aie_aql_queue.cpp

@@ -218,20 +218,18 @@ hsa_status_t AieAqlQueue::SubmitCmd(void* queue_base, uint64_t read_dispatch_id,
         // Iterating over future packets and seeing how many contiguous HSA_AMD_AIE_ERT_START_CU
         // packets there are. All can be combined into a single chain.
         int num_cont_start_cu_pkts = 1;
-        int num_operands = 0;
         for (int peak_pkt_id = cur_id + 1; peak_pkt_id < write_dispatch_id; peak_pkt_id++) {
           hsa_amd_aie_ert_packet_t* peak_pkt =
               static_cast<hsa_amd_aie_ert_packet_t*>(queue_base) + peak_pkt_id;
           if (peak_pkt->opcode != HSA_AMD_AIE_ERT_START_CU) {
             break;
           }
-          num_operands += GetOperandCount(peak_pkt->count);
           num_cont_start_cu_pkts++;
         }
 
-        // Call into the driver to submit from cur_id to write_dispatch_id
-        hsa_status_t status =
-            driver.SubmitCmdChain(pkt, num_cont_start_cu_pkts, num_operands, hw_ctx_handle_);
+        // Call into the driver to submit from cur_id to write_dispatch_id.
+        // Submitting the command chain might create a new hardware context.
+        hsa_status_t status = driver.SubmitCmdChain(pkt, num_cont_start_cu_pkts, *this);
         if (status != HSA_STATUS_SUCCESS) {
           return status;
         }
@@ -239,13 +237,12 @@ hsa_status_t AieAqlQueue::SubmitCmd(void* queue_base, uint64_t read_dispatch_id,
         cur_id += num_cont_start_cu_pkts;
         break;
       }
-      default: {
+      default:
         return HSA_STATUS_ERROR;
-      }
     }
   }
 
-  return HSA_STATUS_SUCCESS;
+  return HSA_STATUS_ERROR;
 }
 
 void AieAqlQueue::StoreRelease(hsa_signal_value_t value) {

runtime/hsa-runtime/inc/hsa_ext_amd.h

@@ -321,41 +321,15 @@ typedef enum {
  */
 typedef struct hsa_amd_aie_ert_start_kernel_data_s {
   /**
-   * Mandatory CU mask.
+   * Address to the PDI.
    */
-  uint32_t cu_mask;
+  void* pdi_addr;
   /**
-   * Since the CU mask takes up one DWORD this is count - 1 number of DWORDs
-   * (i.e., the remainder of the start kernel payload data).
+   * Opcode, instructions and kernel arguments.
    */
   uint32_t data[];
 } hsa_amd_aie_ert_start_kernel_data_t;
 
-/**
- * Payload data for AIE ERT command chain packets (i.e., when the opcode is
- * HSA_AMD_AIE_ERT_CMD_CHAIN). A command chain is a buffer of commands parsed
- * by the ERT.
- */
-typedef struct hsa_amd_aie_ert_command_chain_data_s {
-  /**
-   * Number of commands in the chain.
-   */
-  uint32_t command_count;
-  /**
-   * Index of last successfully submitted command in the chain.
-   */
-  uint32_t submit_index;
-  /**
-   * Index of failing command if command status is not completed.
-   */
-  uint32_t error_index;
-  uint32_t reserved[3];
-  /**
-   * Address of each command in the chain.
-   */
-  uint64_t data[];
-} hsa_amd_aie_ert_command_chain_data_t;
-
 /**
  * AMD AIE ERT packet. Used for sending a command to an AIE agent.
  */