| // Copyright 2025 The IREE Authors |
| // |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| #include "iree/hal/drivers/amdgpu/system.h" |
| |
| #include "iree/hal/drivers/amdgpu/executable.h" |
| #include "iree/hal/drivers/amdgpu/util/topology.h" |
| |
| //===----------------------------------------------------------------------===// |
| // Memory Pools and Regions |
| //===----------------------------------------------------------------------===// |
| // Unfortunately HSA has regions (the original concept) and then memory pools |
| // (the new hotness) and the API uses each in different places. They roughly |
| // represent the same underlying memory but do so differently enough that we |
| // can't interchange them and have to basically treat them as independent |
| // concepts everywhere. |
| |
| // NOTE: we could do the filtering inline in the iteration callback but that |
| // requires moving a lot more state into the user data struct and means we can't |
| // emit iree_status_t errors. Instead we iterate over all and then do things in |
| // local for-loops. |
| typedef struct iree_hal_amdgpu_hsa_memory_pool_list_t { |
| // Number of valid entries in |values|. |
| iree_host_size_t count; |
| // Fixed-capacity memory-pool list populated by HSA iteration callbacks. |
| hsa_amd_memory_pool_t values[32]; |
| } iree_hal_amdgpu_hsa_memory_pool_list_t; |
| static hsa_status_t iree_hal_amdgpu_iterate_hsa_memory_pool( |
| hsa_amd_memory_pool_t memory_pool, void* user_data) { |
| iree_hal_amdgpu_hsa_memory_pool_list_t* pool_list = |
| (iree_hal_amdgpu_hsa_memory_pool_list_t*)user_data; |
| if (pool_list->count >= IREE_ARRAYSIZE(pool_list->values)) { |
| return HSA_STATUS_ERROR_OUT_OF_RESOURCES; |
| } |
| pool_list->values[pool_list->count++] = memory_pool; |
| return HSA_STATUS_SUCCESS; |
| } |
| |
| static iree_status_t iree_hal_amdgpu_system_populate_host_memory_pools( |
| const iree_hal_amdgpu_libhsa_t* libhsa, hsa_agent_t host_agent, |
| iree_hal_amdgpu_host_memory_pools_t* host_memory_pools) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Iterate all memory pools on the CPU agent. |
| iree_hal_amdgpu_hsa_memory_pool_list_t all_memory_pools = { |
| .count = 0, |
| }; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_amd_agent_iterate_memory_pools( |
| IREE_LIBHSA(libhsa), host_agent, |
| iree_hal_amdgpu_iterate_hsa_memory_pool, &all_memory_pools)); |
| |
| bool fine_pool_is_kernarg = false; |
| for (iree_host_size_t i = 0; i < all_memory_pools.count; ++i) { |
| hsa_amd_memory_pool_t pool = all_memory_pools.values[i]; |
| |
| // Filter to the global segment only. |
| hsa_region_segment_t segment = 0; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_amd_memory_pool_get_info(IREE_LIBHSA(libhsa), pool, |
| HSA_AMD_MEMORY_POOL_INFO_SEGMENT, |
| &segment)); |
| if (segment != HSA_REGION_SEGMENT_GLOBAL) continue; |
| |
| // Only care about accessible-by-all. |
| bool accessible_by_all = false; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, |
| iree_hsa_amd_memory_pool_get_info( |
| IREE_LIBHSA(libhsa), pool, |
| HSA_AMD_MEMORY_POOL_INFO_ACCESSIBLE_BY_ALL, &accessible_by_all)); |
| if (!accessible_by_all) continue; |
| |
| // Must be able to allocate. This should be true for any pool we query that |
| // matches the other flags. Workgroup-private pools won't have this set. |
| bool alloc_allowed = false; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, |
| iree_hsa_amd_memory_pool_get_info( |
| IREE_LIBHSA(libhsa), pool, |
| HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALLOWED, &alloc_allowed)); |
| if (!alloc_allowed) continue; |
| |
| // Coarse-grained pools are used for write-once/read-many data. |
| // Kernarg-init pools are used for dispatch argument storage. Fine-grained |
| // pools are used for host/device shared state that requires atomics. |
| uint32_t global_flag = 0; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_amd_memory_pool_get_info( |
| IREE_LIBHSA(libhsa), pool, |
| HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &global_flag)); |
| const bool is_kernarg = |
| global_flag & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_KERNARG_INIT; |
| const bool is_fine = |
| global_flag & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_FINE_GRAINED; |
| if ((global_flag & HSA_AMD_MEMORY_POOL_GLOBAL_FLAG_COARSE_GRAINED) && |
| !host_memory_pools->coarse_pool.handle) { |
| host_memory_pools->coarse_pool = pool; |
| } |
| if (is_kernarg && !host_memory_pools->kernarg_pool.handle) { |
| host_memory_pools->kernarg_pool = pool; |
| } |
| if (is_fine && (!host_memory_pools->fine_pool.handle || |
| (fine_pool_is_kernarg && !is_kernarg))) { |
| host_memory_pools->fine_pool = pool; |
| fine_pool_is_kernarg = is_kernarg; |
| } |
| } |
| |
| if (!host_memory_pools->coarse_pool.handle) { |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_make_status(IREE_STATUS_NOT_FOUND, |
| "no accessible-by-all + coarse-grained shared " |
| "memory pool is available in the system")); |
| } |
| if (!host_memory_pools->fine_pool.handle) { |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_make_status(IREE_STATUS_NOT_FOUND, |
| "no accessible-by-all + fine-grained shared " |
| "memory pool is available in the system")); |
| } |
| if (!host_memory_pools->kernarg_pool.handle) { |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_make_status(IREE_STATUS_NOT_FOUND, |
| "no accessible-by-all + kernarg-init shared " |
| "memory pool is available in the system")); |
| } |
| |
| IREE_TRACE_ZONE_END(z0); |
| return iree_ok_status(); |
| } |
| |
| // NOTE: we could do the filtering inline in the iteration callback but that |
| // requires moving a lot more state into the user data struct and means we can't |
| // emit iree_status_t errors. Instead we iterate over all and then do things in |
| // local for-loops. |
| typedef struct iree_hal_amdgpu_hsa_region_list_t { |
| // Number of valid entries in |values|. |
| iree_host_size_t count; |
| // Fixed-capacity region list populated by HSA iteration callbacks. |
| hsa_region_t values[32]; |
| } iree_hal_amdgpu_hsa_region_list_t; |
| static hsa_status_t iree_hal_amdgpu_iterate_hsa_region(hsa_region_t region, |
| void* user_data) { |
| iree_hal_amdgpu_hsa_region_list_t* region_list = |
| (iree_hal_amdgpu_hsa_region_list_t*)user_data; |
| if (region_list->count >= IREE_ARRAYSIZE(region_list->values)) { |
| return HSA_STATUS_ERROR_OUT_OF_RESOURCES; |
| } |
| region_list->values[region_list->count++] = region; |
| return HSA_STATUS_SUCCESS; |
| } |
| |
| static iree_status_t iree_hal_amdgpu_system_populate_host_regions( |
| const iree_hal_amdgpu_libhsa_t* libhsa, hsa_agent_t host_agent, |
| iree_hal_amdgpu_host_memory_pools_t* host_memory_pools) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Iterate all regions on the CPU agent. |
| iree_hal_amdgpu_hsa_region_list_t all_regions = { |
| .count = 0, |
| }; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_agent_iterate_regions(IREE_LIBHSA(libhsa), host_agent, |
| iree_hal_amdgpu_iterate_hsa_region, |
| &all_regions)); |
| |
| for (iree_host_size_t i = 0; i < all_regions.count; ++i) { |
| hsa_region_t region = all_regions.values[i]; |
| |
| // Filter to the global segment only. |
| hsa_region_segment_t segment = 0; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_region_get_info(IREE_LIBHSA(libhsa), region, |
| HSA_REGION_INFO_SEGMENT, &segment)); |
| if (segment != HSA_REGION_SEGMENT_GLOBAL) continue; |
| |
| // Must be able to allocate. This should be true for any pool we query that |
| // matches the other flags. Workgroup-private pools won't have this set. |
| bool alloc_allowed = false; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hsa_region_get_info(IREE_LIBHSA(libhsa), region, |
| HSA_REGION_INFO_RUNTIME_ALLOC_ALLOWED, |
| &alloc_allowed)); |
| if (!alloc_allowed) continue; |
| |
| // Only want fine-grained so we can use atomics. |
| hsa_region_global_flag_t global_flag = 0; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, |
| iree_hsa_region_get_info(IREE_LIBHSA(libhsa), region, |
| HSA_REGION_INFO_GLOBAL_FLAGS, &global_flag)); |
| if (global_flag & HSA_REGION_GLOBAL_FLAG_FINE_GRAINED) { |
| if (!host_memory_pools->fine_region.handle) { // first only |
| host_memory_pools->fine_region = region; |
| } |
| } |
| } |
| |
| iree_status_t status = iree_ok_status(); |
| if (!host_memory_pools->fine_region.handle) { |
| status = iree_make_status(IREE_STATUS_NOT_FOUND, |
| "no accessible-by-all + fine-grained shared " |
| "memory region is available in the system"); |
| } |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // iree_hal_amdgpu_system_t |
| //===----------------------------------------------------------------------===// |
| |
| static void iree_hal_amdgpu_system_deinitialize( |
| iree_hal_amdgpu_system_t* system); |
| |
| static iree_status_t iree_hal_amdgpu_system_initialize( |
| const iree_hal_amdgpu_libhsa_t* libhsa, |
| const iree_hal_amdgpu_topology_t* topology, |
| iree_hal_amdgpu_system_options_t options, iree_allocator_t host_allocator, |
| iree_hal_amdgpu_system_t* out_system) { |
| IREE_ASSERT_ARGUMENT(libhsa); |
| IREE_ASSERT_ARGUMENT(out_system); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| out_system->host_allocator = host_allocator; |
| |
| // Ensure all GPU agents in the topology support compatible ISAs. They should |
| // all be the same today but in the future if we start allowing heterogeneous |
| // (even if just lightly) we'll want to catch issues here. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hal_amdgpu_verify_device_isa_commonality(libhsa, topology)); |
| |
| // Query and validate the system information. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hal_amdgpu_system_info_query(libhsa, &out_system->info)); |
| |
| // Initialize the platform source used for profile clock correlation. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hal_amdgpu_device_clock_source_initialize( |
| &out_system->device_clock_source)); |
| |
| // Copy the libhsa symbol table and retain HSA for the lifetime of the system. |
| // The caller may destroy the provided libhsa after this call returns. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_hal_amdgpu_libhsa_copy(libhsa, &out_system->libhsa)); |
| |
| // Copy the topology - today it's a plain-old-data struct and we can just |
| // memcpy it. This is an implementation detail, though, and in the future if |
| // it allocates anything we'll need to make sure this retains the allocations |
| // or does a deep copy. |
| memcpy(&out_system->topology, topology, sizeof(out_system->topology)); |
| |
| // Preserved as the information may be used by components of the system. |
| out_system->options = options; |
| |
| // Initialize the device library, which will load the builtin executable and |
| // fail if we don't have a supported arch. |
| iree_status_t status = iree_hal_amdgpu_device_library_initialize( |
| &out_system->libhsa, topology, host_allocator, |
| &out_system->device_library); |
| |
| // Find common/shared memory pools. |
| if (iree_status_is_ok(status)) { |
| for (iree_host_size_t host_ordinal = 0; |
| host_ordinal < topology->cpu_agent_count; ++host_ordinal) { |
| iree_hal_amdgpu_host_memory_pools_t* host_memory_pools = |
| &out_system->host_memory_pools[host_ordinal]; |
| status = iree_hal_amdgpu_system_populate_host_memory_pools( |
| libhsa, topology->cpu_agents[host_ordinal], host_memory_pools); |
| if (!iree_status_is_ok(status)) break; |
| status = iree_hal_amdgpu_system_populate_host_regions( |
| &out_system->libhsa, topology->cpu_agents[host_ordinal], |
| host_memory_pools); |
| if (!iree_status_is_ok(status)) break; |
| } |
| } |
| |
| if (!iree_status_is_ok(status)) { |
| iree_hal_amdgpu_system_deinitialize(out_system); |
| } |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static void iree_hal_amdgpu_system_deinitialize( |
| iree_hal_amdgpu_system_t* system) { |
| IREE_ASSERT_ARGUMENT(system); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Unload the device library - no references to it should remain. |
| iree_hal_amdgpu_device_library_deinitialize(&system->device_library); |
| |
| // Release platform clock-sampling state before unloading HSA. |
| iree_hal_amdgpu_device_clock_source_deinitialize( |
| &system->device_clock_source); |
| |
| // This may unload HSA if we were the last retainer in the process. |
| iree_hal_amdgpu_libhsa_deinitialize(&system->libhsa); |
| |
| memset(system, 0, sizeof(*system)); |
| |
| IREE_TRACE_ZONE_END(z0); |
| } |
| |
| iree_status_t iree_hal_amdgpu_system_allocate( |
| const iree_hal_amdgpu_libhsa_t* libhsa, |
| const iree_hal_amdgpu_topology_t* topology, |
| iree_hal_amdgpu_system_options_t options, iree_allocator_t host_allocator, |
| iree_hal_amdgpu_system_t** out_system) { |
| IREE_ASSERT_ARGUMENT(libhsa); |
| IREE_ASSERT_ARGUMENT(out_system); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| *out_system = NULL; |
| |
| iree_hal_amdgpu_system_t* system = NULL; |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_allocator_malloc( |
| host_allocator, |
| sizeof(*system) + topology->cpu_agent_count * |
| sizeof(system->host_memory_pools[0]), |
| (void**)&system)); |
| |
| iree_status_t status = iree_hal_amdgpu_system_initialize( |
| libhsa, topology, options, host_allocator, system); |
| |
| if (iree_status_is_ok(status)) { |
| *out_system = system; |
| } else { |
| iree_allocator_free(host_allocator, system); |
| } |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| void iree_hal_amdgpu_system_free(iree_hal_amdgpu_system_t* system) { |
| if (!system) return; |
| iree_allocator_t host_allocator = system->host_allocator; |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| iree_hal_amdgpu_system_deinitialize(system); |
| |
| iree_allocator_free(host_allocator, system); |
| |
| IREE_TRACE_ZONE_END(z0); |
| } |
