[HAL][AMDGPU] Export host allocations and device properties

runtime/src/iree/hal/drivers/amdgpu/allocator.c

@@ -1445,21 +1445,34 @@ 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: "
                             "0x%x",
                             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(

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;