| // Copyright 2023 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 |
| |
| #import <Metal/Metal.h> |
| |
| #include "experimental/metal/api.h" |
| #include "experimental/metal/metal_device.h" |
| #include "iree/base/api.h" |
| #include "iree/base/target_platform.h" |
| #include "iree/base/tracing.h" |
| #include "iree/hal/api.h" |
| |
| // Maximum device path length we support. The path is always a 16 character hex string. |
| #define IREE_HAL_METAL_MAX_DEVICE_PATH_LENGTH 32 |
| // Maximum device name length we support. Example names: "Apple M1 Pro". |
| #define IREE_HAL_METAL_MAX_DEVICE_NAME_LENGTH 64 |
| |
| // Cast utilities between Metal id<MTLDevice> and IREE opaque iree_hal_device_id_t. |
| #define METAL_DEVICE_TO_DEVICE_ID(device) (iree_hal_device_id_t)((__bridge void*)device) |
| #define DEVICE_ID_TO_METAL_DEVICE(device_id) (__bridge id<MTLDevice>)(device_id) |
| |
| typedef struct iree_hal_metal_driver_t { |
| // Abstract resource used for injecting reference counting and vtable; must be at offset 0. |
| iree_hal_resource_t resource; |
| |
| iree_allocator_t host_allocator; |
| |
| // Identifier used for the driver in the IREE driver registry. We allow overriding so that |
| // multiple Metal versions can be exposed in the same process. |
| iree_string_view_t identifier; |
| |
| // Parameters used to control device behavior. |
| iree_hal_metal_device_params_t device_params; |
| |
| // The list of GPUs available when creating the driver. We retain them here to make sure |
| // id<MTLDevice>, which is used for creating devices and such, remains valid. |
| NSArray<id<MTLDevice>>* devices; |
| } iree_hal_metal_driver_t; |
| |
| static const iree_hal_driver_vtable_t iree_hal_metal_driver_vtable; |
| |
| static iree_hal_metal_driver_t* iree_hal_metal_driver_cast(iree_hal_driver_t* base_value) { |
| IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_driver_vtable); |
| return (iree_hal_metal_driver_t*)base_value; |
| } |
| |
| static const iree_hal_metal_driver_t* iree_hal_metal_driver_const_cast( |
| const iree_hal_driver_t* base_value) { |
| IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_driver_vtable); |
| return (const iree_hal_metal_driver_t*)base_value; |
| } |
| |
| // Returns an retained array of available Metal GPU devices; the caller should release later. |
| static NSArray<id<MTLDevice>>* iree_hal_metal_device_copy() { |
| #if defined(IREE_PLATFORM_MACOS) |
| // For macOS, we might have more then one GPU devices. |
| return MTLCopyAllDevices(); // +1 |
| #else |
| // For other Apple platforms, we only have one GPU device. |
| @autoreleasepool { // Use @autorelasepool to trigger the autorelease carried in NSArray literal. |
| return [@[ MTLCreateSystemDefaultDevice() ] retain]; // +1 |
| } |
| #endif // IREE_PLATFORM_MACOS |
| } |
| |
| static iree_status_t iree_hal_metal_device_check_params( |
| const iree_hal_metal_device_params_t* params) { |
| if (params->arena_block_size < 4096) { |
| return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, |
| "arena block size too small (< 4096 bytes)"); |
| } |
| return iree_ok_status(); |
| } |
| |
| static iree_status_t iree_hal_metal_driver_create_internal( |
| iree_string_view_t identifier, const iree_hal_metal_device_params_t* device_params, |
| iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { |
| iree_hal_metal_driver_t* driver = NULL; |
| iree_host_size_t total_size = iree_sizeof_struct(*driver) + identifier.size; |
| IREE_RETURN_IF_ERROR(iree_allocator_malloc(host_allocator, total_size, (void**)&driver)); |
| |
| iree_hal_resource_initialize(&iree_hal_metal_driver_vtable, &driver->resource); |
| driver->host_allocator = host_allocator; |
| iree_string_view_append_to_buffer(identifier, &driver->identifier, |
| (char*)driver + iree_sizeof_struct(*driver)); |
| driver->device_params = *device_params; |
| |
| // Get all available Metal devices. |
| driver->devices = iree_hal_metal_device_copy(); |
| |
| *out_driver = (iree_hal_driver_t*)driver; |
| return iree_ok_status(); |
| } |
| |
| IREE_API_EXPORT iree_status_t iree_hal_metal_driver_create( |
| iree_string_view_t identifier, const iree_hal_metal_device_params_t* device_params, |
| iree_allocator_t host_allocator, iree_hal_driver_t** out_driver) { |
| IREE_ASSERT_ARGUMENT(out_driver); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, iree_hal_metal_device_check_params(device_params)); |
| iree_status_t status = |
| iree_hal_metal_driver_create_internal(identifier, device_params, host_allocator, out_driver); |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static void iree_hal_metal_driver_destroy(iree_hal_driver_t* base_driver) { |
| iree_hal_metal_driver_t* driver = iree_hal_metal_driver_cast(base_driver); |
| iree_allocator_t host_allocator = driver->host_allocator; |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| [driver->devices release]; // -1 |
| iree_allocator_free(host_allocator, driver); |
| |
| IREE_TRACE_ZONE_END(z0); |
| } |
| |
| // Populates device information from the given Metal physical device handle. |out_device_info| must |
| // point to valid memory and additional data will be appended to |buffer_ptr| and the new pointer is |
| // returned. |
| static iree_status_t iree_hal_metal_populate_device_info(id<MTLDevice> device, uint8_t* buffer_ptr, |
| uint8_t** out_buffer_ptr, |
| iree_hal_device_info_t* out_device_info) { |
| *out_buffer_ptr = buffer_ptr; |
| |
| memset(out_device_info, 0, sizeof(*out_device_info)); |
| |
| out_device_info->device_id = METAL_DEVICE_TO_DEVICE_ID(device); |
| |
| // For Metal devices, we don't have a 128-bit UUID; so just use the 64-bit registry ID here. |
| char device_path[16 + 1] = {0}; |
| snprintf(device_path, sizeof(device_path), "%016" PRIx64, device.registryID); |
| buffer_ptr += iree_string_view_append_to_buffer( |
| iree_make_string_view(device_path, IREE_ARRAYSIZE(device_path) - 1), &out_device_info->path, |
| (char*)buffer_ptr); |
| |
| const char* device_name = [device.name UTF8String]; |
| const size_t name_len = strlen(device_name); |
| if (name_len >= IREE_HAL_METAL_MAX_DEVICE_NAME_LENGTH) { |
| return iree_make_status(IREE_STATUS_OUT_OF_RANGE, "device name out of range"); |
| } |
| buffer_ptr += iree_string_view_append_to_buffer(iree_make_string_view(device_name, name_len), |
| &out_device_info->name, (char*)buffer_ptr); |
| |
| *out_buffer_ptr = buffer_ptr; |
| |
| return iree_ok_status(); |
| } |
| |
| static iree_status_t iree_hal_metal_driver_query_available_devices( |
| iree_hal_driver_t* base_driver, iree_allocator_t host_allocator, |
| iree_host_size_t* out_device_info_count, iree_hal_device_info_t** out_device_infos) { |
| iree_hal_metal_driver_t* driver = iree_hal_metal_driver_cast(base_driver); |
| NSArray<id<MTLDevice>>* devices = driver->devices; |
| unsigned device_count = devices.count; |
| |
| // Allocate the return infos and populate with the devices. |
| iree_hal_device_info_t* device_infos = NULL; |
| iree_host_size_t single_info_size = |
| sizeof(iree_hal_device_info_t) + |
| (IREE_HAL_METAL_MAX_DEVICE_PATH_LENGTH + IREE_HAL_METAL_MAX_DEVICE_NAME_LENGTH) * |
| sizeof(char); |
| iree_host_size_t total_size = device_count * single_info_size; |
| iree_status_t status = iree_allocator_malloc(host_allocator, total_size, (void**)&device_infos); |
| |
| if (iree_status_is_ok(status)) { |
| // Append all path and name strings at the end of the struct. |
| uint8_t* buffer_ptr = (uint8_t*)device_infos + device_count * sizeof(iree_hal_device_info_t); |
| for (iree_host_size_t i = 0; i < device_count; ++i) { |
| status = iree_hal_metal_populate_device_info(devices[i], buffer_ptr, &buffer_ptr, |
| &device_infos[i]); |
| if (!iree_status_is_ok(status)) break; |
| } |
| } |
| if (iree_status_is_ok(status)) { |
| *out_device_info_count = device_count; |
| *out_device_infos = device_infos; |
| } else { |
| iree_allocator_free(host_allocator, device_infos); |
| } |
| return status; |
| } |
| |
| // Returns the GPU family the given |device| supports. Returns 0 if the given |device| does not |
| // belong to a GPU family considered by IREE right now. |
| static MTLGPUFamily iree_hal_metal_apple_gpu_family_query(id<MTLDevice> device) { |
| // Inspect whether the given device is a specific Apple GPU. |
| if ([device supportsFamily:MTLGPUFamilyApple8]) return MTLGPUFamilyApple8; |
| if ([device supportsFamily:MTLGPUFamilyApple7]) return MTLGPUFamilyApple7; |
| if ([device supportsFamily:MTLGPUFamilyApple6]) return MTLGPUFamilyApple6; |
| if ([device supportsFamily:MTLGPUFamilyApple5]) return MTLGPUFamilyApple5; |
| if ([device supportsFamily:MTLGPUFamilyApple4]) return MTLGPUFamilyApple4; |
| if ([device supportsFamily:MTLGPUFamilyApple3]) return MTLGPUFamilyApple3; |
| if ([device supportsFamily:MTLGPUFamilyApple2]) return MTLGPUFamilyApple2; |
| if ([device supportsFamily:MTLGPUFamilyApple1]) return MTLGPUFamilyApple1; |
| |
| // Inspect whether whether the given GPU falls into some common family. |
| if ([device supportsFamily:MTLGPUFamilyCommon3]) return MTLGPUFamilyCommon3; |
| if ([device supportsFamily:MTLGPUFamilyCommon2]) return MTLGPUFamilyCommon2; |
| if ([device supportsFamily:MTLGPUFamilyCommon1]) return MTLGPUFamilyCommon1; |
| |
| return 0; |
| } |
| |
| static const char* iree_hal_metal_get_gpu_family_name(MTLGPUFamily family) { |
| switch (family) { |
| case MTLGPUFamilyApple8: |
| return "apple8(a15/m2)"; |
| case MTLGPUFamilyApple7: |
| return "apple7(a14/m1)"; |
| case MTLGPUFamilyApple6: |
| return "apple6(a13)"; |
| case MTLGPUFamilyApple5: |
| return "apple5(a12)"; |
| case MTLGPUFamilyApple4: |
| return "apple4(a11)"; |
| case MTLGPUFamilyApple3: |
| return "apple3(a9/a10)"; |
| case MTLGPUFamilyApple2: |
| return "apple2(a8)"; |
| case MTLGPUFamilyApple1: |
| return "apple1(a7)"; |
| |
| case MTLGPUFamilyCommon3: |
| return "common3"; |
| case MTLGPUFamilyCommon2: |
| return "common2"; |
| case MTLGPUFamilyCommon1: |
| return "common1"; |
| |
| default: |
| return ""; |
| } |
| } |
| |
| static inline const char* iree_hal_metal_get_argument_buffer_tier_str(MTLGPUFamily family) { |
| if (family >= MTLGPUFamilyApple6 && family <= MTLGPUFamilyApple8) return "2"; |
| if (family >= MTLGPUFamilyApple2 && family <= MTLGPUFamilyApple5) return "1"; |
| return "unknown"; |
| } |
| |
| static inline bool iree_hal_metal_support_simd_matrix_multiply(MTLGPUFamily family) { |
| return family >= MTLGPUFamilyApple7 && family <= MTLGPUFamilyApple8; |
| } |
| |
| static inline bool iree_hal_metal_support_simd_reduce(MTLGPUFamily family) { |
| return family >= MTLGPUFamilyApple7 && family <= MTLGPUFamilyApple8; |
| } |
| |
| static inline bool iree_hal_metal_support_simd_permute(MTLGPUFamily family) { |
| return family >= MTLGPUFamilyApple6 && family <= MTLGPUFamilyApple8; |
| } |
| |
| static inline bool iree_hal_metal_support_simd_shift_and_fill(MTLGPUFamily family) { |
| return family >= MTLGPUFamilyApple8 && family <= MTLGPUFamilyApple8; |
| } |
| |
| static iree_status_t iree_hal_metal_driver_dump_device_info(iree_hal_driver_t* base_driver, |
| iree_hal_device_id_t device_id, |
| iree_string_builder_t* builder) { |
| id<MTLDevice> device = DEVICE_ID_TO_METAL_DEVICE(device_id); |
| MTLGPUFamily apple_gpu_family = iree_hal_metal_apple_gpu_family_query(device); |
| |
| // Dump GPU family information. |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "- gpu-family:")); |
| const char* apple_family_str = iree_hal_metal_get_gpu_family_name(apple_gpu_family); |
| if (apple_family_str) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " ")); |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, apple_family_str)); |
| } |
| if ([device supportsFamily:MTLGPUFamilyMetal3]) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " metal3")); |
| } |
| |
| // Dump memory information. |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_format(builder, "\n- unified-memory: %d", |
| device.hasUnifiedMemory)); |
| |
| // Dump argument buffer tier. |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "\n- argument-buffer-tier: ")); |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring( |
| builder, iree_hal_metal_get_argument_buffer_tier_str(apple_gpu_family))); |
| |
| // Dump resource limits. |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "\n- max-buffer-size: ")); |
| { |
| uint32_t max_buffer_mb = device.maxBufferLength / 1024u / 1024u; |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_format(builder, "%uMB", max_buffer_mb)); |
| } |
| IREE_RETURN_IF_ERROR( |
| iree_string_builder_append_cstring(builder, "\n- max-threadgroup-memory-size: ")); |
| { |
| uint32_t max_memory_kb = device.maxThreadgroupMemoryLength / 1024u; |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_format(builder, "%uKB", max_memory_kb)); |
| } |
| IREE_RETURN_IF_ERROR( |
| iree_string_builder_append_cstring(builder, "\n- max-threads-per-threadgroup: ")); |
| { |
| MTLSize threads = device.maxThreadsPerThreadgroup; |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_format( |
| builder, "(%lu, %lu, %lu)", threads.width, threads.height, threads.depth)); |
| } |
| |
| // Dump SIMD-scoped operation features. |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "\n- simd-scoped-operation:")); |
| if (iree_hal_metal_support_simd_matrix_multiply(apple_gpu_family)) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " matmul")); |
| } |
| if (iree_hal_metal_support_simd_reduce(apple_gpu_family)) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " reduce")); |
| } |
| if (iree_hal_metal_support_simd_permute(apple_gpu_family)) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " permute")); |
| } |
| if (iree_hal_metal_support_simd_shift_and_fill(apple_gpu_family)) { |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, " shift-and-fill")); |
| } |
| IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(builder, "\n")); |
| |
| return iree_ok_status(); |
| } |
| |
| static iree_status_t iree_hal_metal_driver_find_device_by_index(iree_hal_driver_t* base_driver, |
| uint32_t device_index, |
| iree_allocator_t host_allocator, |
| id<MTLDevice>* found_device) { |
| iree_hal_metal_driver_t* driver = iree_hal_metal_driver_cast(base_driver); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (uint64_t)device_index); |
| |
| NSArray<id<MTLDevice>>* devices = driver->devices; |
| if (device_index >= devices.count) { |
| IREE_TRACE_ZONE_END(z0); |
| return iree_make_status(IREE_STATUS_NOT_FOUND, "%d devices enumerated; device #%d not found", |
| (int)devices.count, device_index); |
| } |
| *found_device = devices[device_index]; |
| |
| IREE_TRACE_ZONE_END(z0); |
| return iree_ok_status(); |
| } |
| |
| static iree_status_t iree_hal_metal_driver_create_device_by_id(iree_hal_driver_t* base_driver, |
| iree_hal_device_id_t device_id, |
| iree_host_size_t param_count, |
| const iree_string_pair_t* params, |
| iree_allocator_t host_allocator, |
| iree_hal_device_t** out_device) { |
| iree_hal_metal_driver_t* driver = iree_hal_metal_driver_cast(base_driver); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| id<MTLDevice> device = nil; |
| if (device_id == IREE_HAL_DEVICE_ID_DEFAULT) { |
| // Default to the first Metal device in the list. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR(z0, iree_hal_metal_driver_find_device_by_index( |
| base_driver, device_id, host_allocator, &device)); |
| } else { |
| device = DEVICE_ID_TO_METAL_DEVICE(device_id); |
| } |
| |
| iree_string_view_t device_name = iree_make_cstring_view("metal"); |
| |
| iree_status_t status = iree_hal_metal_device_create(device_name, &driver->device_params, device, |
| host_allocator, out_device); |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static iree_status_t iree_hal_metal_driver_create_device_by_registry_id( |
| iree_hal_driver_t* base_driver, iree_string_view_t driver_name, uint64_t device_registry_id, |
| iree_host_size_t param_count, const iree_string_pair_t* params, iree_allocator_t host_allocator, |
| iree_hal_device_t** out_device) { |
| iree_hal_metal_driver_t* driver = iree_hal_metal_driver_cast(base_driver); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Scan the devices and find the one with the matching registry ID. |
| NSArray<id<MTLDevice>>* devices = driver->devices; |
| id<MTLDevice> found_device = nil; |
| for (iree_host_size_t i = 0, e = devices.count; i < e; ++i) { |
| if (device_registry_id == devices[i].registryID) { |
| found_device = devices[i]; |
| break; |
| } |
| } |
| |
| if (!found_device) { |
| IREE_TRACE_ZONE_END(z0); |
| return iree_make_status(IREE_STATUS_NOT_FOUND, |
| "Metal device with device registry ID %016" PRIx64 " not found", |
| device_registry_id); |
| } |
| |
| iree_status_t status = iree_hal_metal_driver_create_device_by_id( |
| base_driver, METAL_DEVICE_TO_DEVICE_ID(found_device), param_count, params, host_allocator, |
| out_device); |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static iree_status_t iree_hal_metal_driver_create_device_by_path( |
| iree_hal_driver_t* base_driver, iree_string_view_t driver_name, iree_string_view_t device_path, |
| iree_host_size_t param_count, const iree_string_pair_t* params, iree_allocator_t host_allocator, |
| iree_hal_device_t** out_device) { |
| if (iree_string_view_is_empty(device_path)) { |
| return iree_hal_metal_driver_create_device_by_id( |
| base_driver, IREE_HAL_DEVICE_ID_DEFAULT, param_count, params, host_allocator, out_device); |
| } |
| |
| // Try parsing as a device ID. |
| uint8_t device_registry_id[8] = {0}; |
| if (iree_string_view_parse_hex_bytes(device_path, IREE_ARRAYSIZE(device_registry_id), |
| device_registry_id)) { |
| return iree_hal_metal_driver_create_device_by_registry_id( |
| base_driver, driver_name, *(uint64_t*)device_registry_id, param_count, params, |
| host_allocator, out_device); |
| } |
| |
| // Fallback and try to parse as a device index. |
| uint32_t device_index = 0; |
| if (iree_string_view_atoi_uint32(device_path, &device_index)) { |
| id<MTLDevice> found_device; |
| IREE_RETURN_IF_ERROR(iree_hal_metal_driver_find_device_by_index(base_driver, device_index, |
| host_allocator, &found_device)); |
| return iree_hal_metal_driver_create_device_by_id( |
| base_driver, METAL_DEVICE_TO_DEVICE_ID(found_device), param_count, params, host_allocator, |
| out_device); |
| } |
| |
| return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "unsupported device path"); |
| } |
| |
| static const iree_hal_driver_vtable_t iree_hal_metal_driver_vtable = { |
| .destroy = iree_hal_metal_driver_destroy, |
| .query_available_devices = iree_hal_metal_driver_query_available_devices, |
| .dump_device_info = iree_hal_metal_driver_dump_device_info, |
| .create_device_by_id = iree_hal_metal_driver_create_device_by_id, |
| .create_device_by_path = iree_hal_metal_driver_create_device_by_path, |
| }; |
