diff --git a/runtime/src/iree/hal/drivers/amdgpu/allocator.c b/runtime/src/iree/hal/drivers/amdgpu/allocator.c index dd4f3b416030..72a5d9129b54 100644 --- a/runtime/src/iree/hal/drivers/amdgpu/allocator.c +++ b/runtime/src/iree/hal/drivers/amdgpu/allocator.c @@ -1445,6 +1445,9 @@ static iree_status_t iree_hal_amdgpu_allocator_export_buffer( iree_hal_external_buffer_type_t requested_type, iree_hal_external_buffer_flags_t requested_flags, iree_hal_external_buffer_t* IREE_RESTRICT out_external_buffer) { + (void)base_allocator; + IREE_ASSERT_ARGUMENT(buffer); + IREE_ASSERT_ARGUMENT(out_external_buffer); if (IREE_UNLIKELY(requested_flags != IREE_HAL_EXTERNAL_BUFFER_FLAG_NONE)) { return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "unsupported AMDGPU external buffer export flags: " @@ -1452,14 +1455,24 @@ static iree_status_t iree_hal_amdgpu_allocator_export_buffer( requested_flags); } + iree_hal_buffer_t* allocated_buffer = iree_hal_buffer_allocated_buffer(buffer); + void* base_ptr = iree_hal_amdgpu_buffer_device_pointer(allocated_buffer); + if (IREE_UNLIKELY(!base_ptr)) { + return iree_make_status( + IREE_STATUS_UNAVAILABLE, + "AMDGPU buffer has no HSA allocation pointer to export"); + } + const iree_device_size_t byte_offset = iree_hal_buffer_byte_offset(buffer); + const iree_device_size_t byte_length = iree_hal_buffer_byte_length(buffer); + void* view_ptr = (uint8_t*)base_ptr + byte_offset; + + memset(out_external_buffer, 0, sizeof(*out_external_buffer)); + out_external_buffer->flags = requested_flags; + out_external_buffer->type = requested_type; + out_external_buffer->size = byte_length; + switch (requested_type) { case IREE_HAL_EXTERNAL_BUFFER_TYPE_DEVICE_ALLOCATION: { - if (IREE_UNLIKELY(buffer != iree_hal_buffer_allocated_buffer(buffer))) { - return iree_make_status( - IREE_STATUS_UNAVAILABLE, - "AMDGPU subspan buffers cannot yet be exported as external device " - "allocations"); - } if (IREE_UNLIKELY( !iree_all_bits_set(iree_hal_buffer_memory_type(buffer), IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL))) { @@ -1468,22 +1481,24 @@ static iree_status_t iree_hal_amdgpu_allocator_export_buffer( "AMDGPU buffer memory type is not supported for export as an " "external device allocation"); } + out_external_buffer->handle.device_allocation.ptr = + (uint64_t)(uintptr_t)view_ptr; + return iree_ok_status(); + } - void* device_ptr = iree_hal_amdgpu_buffer_device_pointer(buffer); - if (IREE_UNLIKELY(!device_ptr)) { + case IREE_HAL_EXTERNAL_BUFFER_TYPE_HOST_ALLOCATION: { + if (IREE_UNLIKELY( + !iree_all_bits_set(iree_hal_buffer_memory_type(buffer), + IREE_HAL_MEMORY_TYPE_HOST_LOCAL))) { return iree_make_status( IREE_STATUS_UNAVAILABLE, - "AMDGPU buffer has no HSA device allocation pointer to export"); + "AMDGPU buffer memory type is not supported for export as an " + "external host allocation"); } - out_external_buffer->flags = requested_flags; - out_external_buffer->type = requested_type; - out_external_buffer->size = iree_hal_buffer_allocation_size(buffer); - out_external_buffer->handle.device_allocation.ptr = - (uint64_t)(uintptr_t)device_ptr; + out_external_buffer->handle.host_allocation.ptr = view_ptr; return iree_ok_status(); } - case IREE_HAL_EXTERNAL_BUFFER_TYPE_HOST_ALLOCATION: case IREE_HAL_EXTERNAL_BUFFER_TYPE_OPAQUE_FD: case IREE_HAL_EXTERNAL_BUFFER_TYPE_OPAQUE_WIN32: return iree_make_status( diff --git a/runtime/src/iree/hal/drivers/amdgpu/logical_device.c b/runtime/src/iree/hal/drivers/amdgpu/logical_device.c index 376fad8bdc7b..296fe454e380 100644 --- a/runtime/src/iree/hal/drivers/amdgpu/logical_device.c +++ b/runtime/src/iree/hal/drivers/amdgpu/logical_device.c @@ -1887,6 +1887,47 @@ static iree_status_t iree_hal_amdgpu_logical_device_query_i64( system->topology.gpu_agent_queue_count; return iree_ok_status(); } + if (iree_string_view_equal(key, IREE_SV("warp_size")) || + iree_string_view_equal(key, IREE_SV("wavefront_size"))) { + if (logical_device->physical_device_count == 0) { + return iree_make_status(IREE_STATUS_UNAVAILABLE, + "logical device has no physical devices"); + } + *out_value = + (int64_t)logical_device->physical_devices[0]->wavefront_size; + return iree_ok_status(); + } + if (iree_string_view_equal(key, IREE_SV("gfxip"))) { + if (logical_device->physical_device_count == 0) { + return iree_make_status( + IREE_STATUS_INTERNAL, + "logical device has no physical devices (initialization incomplete)"); + } + const iree_hal_amdgpu_gfxip_version_t version = + logical_device->physical_devices[0]->isa.target_id.version; + *out_value = ((int64_t)version.major << 16) | + ((int64_t)version.minor << 8) | (int64_t)version.stepping; + return iree_ok_status(); + } + if (iree_string_view_equal(key, IREE_SV("memory.total")) || + iree_string_view_equal(key, IREE_SV("memory.free"))) { + uint64_t total = 0; + for (iree_host_size_t i = 0; + i < logical_device->physical_device_count; ++i) { + iree_hal_amdgpu_physical_device_t* physical_device = + logical_device->physical_devices[i]; + hsa_amd_memory_pool_t pool = + physical_device->coarse_block_pools.large.memory_pool; + if (!pool.handle) continue; + size_t pool_size = 0; + IREE_RETURN_IF_ERROR(iree_hsa_amd_memory_pool_get_info( + IREE_LIBHSA(&system->libhsa), pool, + HSA_AMD_MEMORY_POOL_INFO_SIZE, &pool_size)); + total += (uint64_t)pool_size; + } + *out_value = (int64_t)total; + return iree_ok_status(); + } } else if (iree_string_view_equal(category, IREE_SV("amdgpu.device"))) { if (iree_string_view_equal(key, IREE_SV("physical_device.count"))) { *out_value = (int64_t)logical_device->physical_device_count;