From c63e01724cc9b1863705e9f047bf58529ff78e83 Mon Sep 17 00:00:00 2001 From: Yiannis Papadopoulos Date: Fri, 28 Mar 2025 15:29:34 -0400 Subject: [PATCH] rocr/aie: Using PDI address instead of cu_mask for dispatch. Automatic hw ctx reconfiguration upon new PDI addition. --- .../core/driver/xdna/amd_xdna_driver.cpp | 411 ++++++++++++------ .../hsa-runtime/core/inc/amd_xdna_driver.h | 101 +++-- .../core/runtime/amd_aie_aql_queue.cpp | 13 +- runtime/hsa-runtime/inc/hsa_ext_amd.h | 32 +- 4 files changed, 351 insertions(+), 206 deletions(-) diff --git a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp b/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp index 83cb0d1f90..4db998a67b 100644 --- a/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp +++ b/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(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(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(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(&qos_info); - create_hwctx_args.max_opc = 0x800; - create_hwctx_args.num_tiles = static_cast(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(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& bo_addrs, - const std::vector& 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& 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(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(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& bo_args, - std::vector& bo_sizes, - std::vector& 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& bo_handles) { const uint64_t instr_addr = Concat(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& 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(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(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(cmd_pkt_payload->data[operand_index + 1], @@ -519,15 +522,15 @@ hsa_status_t XdnaDriver::RegisterCmdBOs(uint32_t count, std::vector& b if (!operand_bo_handle.IsValid()) { return HSA_STATUS_ERROR; } - bo_args.push_back(operand_bo_handle.handle); - bo_addrs.push_back(reinterpret_cast(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(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& 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(mmap(nullptr, create_cmd_bo.size, PROT_READ | PROT_WRITE, - MAP_SHARED, fd_, cmd_bo_get_bo_info.map_offset)); + void* mem = static_cast(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 bo_args; - std::vector bo_sizes; - std::vector 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 bo_handles; - // Storing the commands that we are going to submit and the - // corresponding metadata - std::vector cmd_handles; - std::vector cmd_sizes; - std::vector 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 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(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(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(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(cmd_chain_bo_handle.vaddr); // Writing information to the command buffer amdxdna_cmd_chain* cmd_chain_payload = reinterpret_cast(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(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(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(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(&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(xdna_config_cu_param); + config_hw_ctx_args.param_val_size = static_cast(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 diff --git a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h b/runtime/hsa-runtime/core/inc/amd_xdna_driver.h index 55058faa08..41eeed203c 100644 --- a/runtime/hsa-runtime/core/inc/amd_xdna_driver.h +++ b/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 +#include #include #include #include @@ -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 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& bo_args, - std::vector& bo_sizes, std::vector& 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& bo_args, const std::vector& 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& 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& 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 vmem_handle_mappings; std::map vmem_addr_mappings; + /// @brief Hardware context to PDI cache mapping. + std::unordered_map 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 diff --git a/runtime/hsa-runtime/core/runtime/amd_aie_aql_queue.cpp b/runtime/hsa-runtime/core/runtime/amd_aie_aql_queue.cpp index 2d1c51c725..a82baa2b35 100644 --- a/runtime/hsa-runtime/core/runtime/amd_aie_aql_queue.cpp +++ b/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(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) { diff --git a/runtime/hsa-runtime/inc/hsa_ext_amd.h b/runtime/hsa-runtime/inc/hsa_ext_amd.h index 2593b25ab7..68670ef1c8 100644 --- a/runtime/hsa-runtime/inc/hsa_ext_amd.h +++ b/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. */