experimental/metal/direct_allocator.m - 3p/openxla/iree - Git at Google

 // 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

 #include "experimental/metal/direct_allocator.h"

 #import <Metal/Metal.h>

 #include "experimental/metal/metal_buffer.h"
 #include "iree/base/api.h"
 #include "iree/base/target_platform.h"
 #include "iree/base/tracing.h"
 #include "iree/hal/api.h"

 #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_ALLOCATION_TRACKING
 static const char* IREE_HAL_METAL_ALLOCATOR_ID = "METAL";
 #endif  // IREE_TRACING_FEATURE_ALLOCATION_TRACKING

 typedef struct iree_hal_metal_allocator_t {
   // Abstract resource used for injecting reference counting and vtable; must be at offset 0.
   iree_hal_resource_t resource;

   // The device that this allocator is attached to.
   id<MTLDevice> device;
   // The command queue that we can use to issue commands to make buffer contents visible to CPU.
 #if defined(IREE_PLATFORM_MACOS)
   id<MTLCommandQueue> queue;
 #endif  // IREE_PLATFORM_MACOS

   bool is_unified_memory;
   iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode;

   iree_allocator_t host_allocator;

   IREE_STATISTICS(iree_hal_allocator_statistics_t statistics;)
 } iree_hal_metal_allocator_t;

 static const iree_hal_allocator_vtable_t iree_hal_metal_allocator_vtable;

 static iree_hal_metal_allocator_t* iree_hal_metal_allocator_cast(iree_hal_allocator_t* base_value) {
   IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_allocator_vtable);
   return (iree_hal_metal_allocator_t*)base_value;
 }

 static const iree_hal_metal_allocator_t* iree_hal_metal_allocator_const_cast(
     const iree_hal_allocator_t* base_value) {
   IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_allocator_vtable);
   return (const iree_hal_metal_allocator_t*)base_value;
 }

 iree_status_t iree_hal_metal_allocator_create(
     id<MTLDevice> device,
 #if defined(IREE_PLATFORM_MACOS)
     id<MTLCommandQueue> queue,
 #endif  // IREE_PLATFORM_MACOS
     iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode,
     iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator) {
   IREE_ASSERT_ARGUMENT(out_allocator);
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_hal_metal_allocator_t* allocator = NULL;
   iree_status_t status =
       iree_allocator_malloc(host_allocator, sizeof(*allocator), (void**)&allocator);

   if (iree_status_is_ok(status)) {
     iree_hal_resource_initialize(&iree_hal_metal_allocator_vtable, &allocator->resource);
     allocator->device = [device retain];  // +1
     allocator->is_unified_memory = [device hasUnifiedMemory];
     allocator->resource_tracking_mode = resource_tracking_mode;
     allocator->host_allocator = host_allocator;

     *out_allocator = (iree_hal_allocator_t*)allocator;
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 static void iree_hal_metal_allocator_destroy(iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
   iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
   iree_allocator_t host_allocator = allocator->host_allocator;
   IREE_TRACE_ZONE_BEGIN(z0);

   [allocator->device release];  // -1
   iree_allocator_free(host_allocator, allocator);

   IREE_TRACE_ZONE_END(z0);
 }

 static iree_allocator_t iree_hal_metal_allocator_host_allocator(
     const iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
   iree_hal_metal_allocator_t* allocator = (iree_hal_metal_allocator_t*)base_allocator;
   return allocator->host_allocator;
 }

 #if defined(IREE_PLATFORM_MACOS)
 id<MTLCommandQueue> iree_hal_metal_allocator_command_queue(
     const iree_hal_allocator_t* base_allocator) {
   const iree_hal_metal_allocator_t* allocator = (const iree_hal_metal_allocator_t*)base_allocator;
   return allocator->queue;
 }
 #endif  // IREE_PLATFORM_MACOS

 static iree_status_t iree_hal_metal_allocator_trim(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
   return iree_ok_status();
 }

 static void iree_hal_metal_allocator_query_statistics(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator,
     iree_hal_allocator_statistics_t* IREE_RESTRICT out_statistics) {
   IREE_STATISTICS({
     iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
     memcpy(out_statistics, &allocator->statistics, sizeof(*out_statistics));
   });
 }

 static iree_hal_buffer_compatibility_t iree_hal_metal_allocator_query_buffer_compatibility(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator,
     iree_hal_buffer_params_t* IREE_RESTRICT params,
     iree_device_size_t* IREE_RESTRICT allocation_size) {
   // All buffers can be allocated on the heap.
   iree_hal_buffer_compatibility_t compatibility = IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE;

   if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) {
     compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER;
   }

   // Buffers can only be used on the queue if they are device visible.
   if (iree_all_bits_set(params->type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
     if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) {
       compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER;
     }
     if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE)) {
       compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH;
     }
   }

   if (iree_all_bits_set(params->type,
                         IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) {
     // On iOS, we don't have device local + host visible memory. But given the unified memory
     // architecture, it's fine to just request host local + device visible memory.
     // On macOS, for unified memory architecture, it's similar to iOS. Otherwise, we can have
     // device local + host visible memory backed by Managed storage mode.
     iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
     if (allocator->is_unified_memory) {
       params->type &= ~(IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL | IREE_HAL_MEMORY_TYPE_HOST_VISIBLE);
       params->type |= IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
       // We have suggested other configurations than the original request. Mark accordingly.
       compatibility |= IREE_HAL_BUFFER_COMPATIBILITY_LOW_PERFORMANCE;
     }
   }

   // We are now optimal.
   params->type &= ~IREE_HAL_MEMORY_TYPE_OPTIMAL;

   // Guard against the corner case where the requested buffer size is 0. The application is unlikely
   // to do anything when requesting a 0-byte buffer; but it can happen in real world use cases. So
   // we should at least not crash.
   if (*allocation_size == 0) *allocation_size = 4;

   // Align allocation sizes to 4 bytes so shaders operating on 32 bit types can act safely even on
   // buffer ranges that are not naturally aligned.
   *allocation_size = iree_host_align(*allocation_size, 4);

   return compatibility;
 }

 // Returns the corresponding Metal resource options controlling storage modes, CPU caching modes,
 // and hazard tracking modes for the given IREE HAL memory |type|.
 static MTLResourceOptions iree_hal_metal_select_resource_options(
     iree_hal_memory_type_t type, bool is_unified_memory, bool has_init_data,
     iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode) {
   MTLResourceOptions options;

   // Select a storage mode. There are four MTLStorageMode:
   // * Shared: The resource is stored in system memory and is accessible to both the CPU and
   //   the GPU.
   // * Managed: The CPU and GPU may maintain separate copies of the resource, and any changes
   //   must be explicitly synchronized.
   // * Private: The resource can be accessed only by the GPU.
   // * Memoryless: The resource's contents can be accessed only by the GPU and only exist
   //   temporarily during a render pass.
   // macOS has all of the above; MTLStorageModeManaged is not available on iOS.
   //
   // The IREE HAL is modeled after explicit APIs like Vulkan. For buffers visible to both the host
   // and the device, we would like to opt in with the explicit version (MTLStorageManaged) when
   // possible because it should be more performant.
   if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) {
     if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) {
       // Device local + host visible.
       // iree_hal_metal_allocator_query_buffer_compatibility guarantees that we only fall into this
       // case for macOS devices with non-uniform memory.
 #if defined(IREE_PLATFORM_MACOS)
       IREE_ASSERT(!is_unified_memory);
       options = MTLResourceStorageModeManaged;
 #else
       options = MTLResourceStorageModeShared;
 #endif  // IREE_PLATFORM_MACOS
     } else {
       // Device local + host invisible.
       // For unified memory, update it to Shared storage mode if we have initial data.
       options = (is_unified_memory && has_init_data) ? MTLResourceStorageModeShared
                                                      : MTLResourceStorageModePrivate;
     }
   } else {
     if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
       // Host local + device visible.
       options = MTLResourceStorageModeShared;
     } else {
       // Host local + device invisible.
       options = MTLResourceStorageModeShared;
     }
   }

   // Select a CPU cache mode.
   if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_HOST_CACHED)) {
     // The default CPU cache mode for the resource, which guarantees that read and write operations
     // are executed in the expected order.
     options |= MTLResourceCPUCacheModeDefaultCache;
   } else {
     // A write-combined CPU cache mode that is optimized for resources that the CPU writes into, but
     // never reads.
     options |= MTLResourceCPUCacheModeWriteCombined;
   }

   options |= resource_tracking_mode == IREE_HAL_METAL_RESOURCE_HAZARD_TRACKING_MODE_TRACKED
                  ? MTLResourceHazardTrackingModeTracked
                  : MTLResourceHazardTrackingModeUntracked;
   return options;
 }

 static iree_status_t iree_hal_metal_allocator_allocate_buffer(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator,
     const iree_hal_buffer_params_t* IREE_RESTRICT params, iree_device_size_t allocation_size,
     iree_const_byte_span_t initial_data, iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
   iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
   IREE_TRACE_ZONE_BEGIN(z0);
   IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, allocation_size);

   // Coerce options into those required by the current device.
   iree_hal_buffer_params_t compat_params = *params;
   iree_hal_buffer_compatibility_t compatibility =
       iree_hal_metal_allocator_query_buffer_compatibility(base_allocator, &compat_params,
                                                           &allocation_size);
   if (!iree_all_bits_set(compatibility, IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE)) {
 #if IREE_STATUS_MODE
     iree_bitfield_string_temp_t temp0, temp1, temp2;
     iree_string_view_t memory_type_str = iree_hal_memory_type_format(params->type, &temp0);
     iree_string_view_t usage_str = iree_hal_buffer_usage_format(params->usage, &temp1);
     iree_string_view_t compatibility_str =
         iree_hal_buffer_compatibility_format(compatibility, &temp2);
     IREE_TRACE_ZONE_END(z0);
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                             "allocator cannot allocate a buffer with the given parameters; "
                             "memory_type=%.*s, usage=%.*s, compatibility=%.*s",
                             (int)memory_type_str.size, memory_type_str.data, (int)usage_str.size,
                             usage_str.data, (int)compatibility_str.size, compatibility_str.data);
 #else
     IREE_TRACE_ZONE_END(z0);
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                             "allocator cannot allocate a buffer with the given parameters");
 #endif  // IREE_STATUS_MODE
   }

   bool has_init_data = !iree_const_byte_span_is_empty(initial_data);
   MTLResourceOptions options =
       iree_hal_metal_select_resource_options(compat_params.type, allocator->is_unified_memory,
                                              has_init_data, allocator->resource_tracking_mode);
   if (has_init_data && iree_all_bits_set(options, MTLResourceStorageModePrivate)) {
     IREE_TRACE_ZONE_END(z0);
     // This should only happen for macOS devices with non-uniform memory; for uniform memory
     // devices, the allocator compatibility would update it to use Shared storage mode.
     return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "private storage with initial data");
   }

   id<MTLBuffer> metal_buffer = nil;
   // If we chose shared storage mode, we can handle initial data with newByfferWithBytes directly.
   // Otherwise we just create the buffer here, and explicitly transfer range later.
   if (has_init_data && iree_all_bits_set(options, MTLResourceStorageModeShared)) {
     IREE_ASSERT_EQ(allocation_size, initial_data.data_length);
     metal_buffer = [allocator->device newBufferWithBytes:(void*)initial_data.data
                                                   length:initial_data.data_length
                                                  options:options];  // +1
   } else {
     metal_buffer = [allocator->device newBufferWithLength:allocation_size options:options];  // +1
   }
   if (!metal_buffer) {
     IREE_TRACE_ZONE_END(z0);
     return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED, "unable to allocate buffer");
   }
   iree_hal_buffer_t* buffer = NULL;
   iree_status_t status = iree_hal_metal_buffer_wrap(
 #if defined(IREE_PLATFORM_MACOS)
       allocator->queue,
 #endif  // IREE_PLATFORM_MACOS
       metal_buffer, base_allocator, compat_params.type, compat_params.access, compat_params.usage,
       allocation_size, /*byte_offset=*/0, /*byte_length=*/allocation_size,
       iree_hal_buffer_release_callback_null(), &buffer);  // +1

 #if defined(IREE_PLATFORM_MACOS)
   if (iree_status_is_ok(status) && has_init_data &&
       !iree_all_bits_set(options, MTLResourceStorageModePrivate)) {
     // Explicitly handle synchronization for Managed storage mode.
     if (iree_all_bits_set(options, MTLResourceStorageModeManaged)) {
       memcpy(metal_buffer.contents, initial_data.data, initial_data.data_length);
       [metal_buffer didModifyRange:NSMakeRange(0, initial_data.data_length)];
     }
   }
 #endif  // IREE_PLATFORM_MACOS

   if (iree_status_is_ok(status)) {
     IREE_TRACE_ALLOC_NAMED(IREE_HAL_METAL_ALLOCATOR_ID, (void*)iree_hal_metal_buffer_handle(buffer),
                            allocation_size);
     IREE_STATISTICS(iree_hal_allocator_statistics_record_alloc(
         &allocator->statistics, compat_params.type, allocation_size));
     *out_buffer = buffer;
   } else {
     if (buffer) iree_hal_buffer_release(buffer);
   }

   [metal_buffer release];  // -1

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 static void iree_hal_metal_allocator_deallocate_buffer(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator,
     iree_hal_buffer_t* IREE_RESTRICT base_buffer) {
   iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);

   IREE_TRACE_FREE_NAMED(IREE_HAL_METAL_ALLOCATOR_ID,
                         (void*)iree_hal_metal_buffer_handle(base_buffer));
   IREE_STATISTICS(iree_hal_allocator_statistics_record_free(
       &allocator->statistics, iree_hal_buffer_memory_type(base_buffer),
       iree_hal_buffer_allocation_size(base_buffer)));

   iree_hal_buffer_destroy(base_buffer);  // -1
 }

 static iree_status_t iree_hal_metal_allocator_import_buffer(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator,
     const iree_hal_buffer_params_t* IREE_RESTRICT params,
     iree_hal_external_buffer_t* IREE_RESTRICT external_buffer,
     iree_hal_buffer_release_callback_t release_callback,
     iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
   return iree_make_status(IREE_STATUS_UNAVAILABLE, "unsupported importing from external buffer");
 }

 static iree_status_t iree_hal_metal_allocator_export_buffer(
     iree_hal_allocator_t* IREE_RESTRICT base_allocator, iree_hal_buffer_t* IREE_RESTRICT 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) {
   return iree_make_status(IREE_STATUS_UNAVAILABLE, "unsupported exporting to external buffer");
 }

 static const iree_hal_allocator_vtable_t iree_hal_metal_allocator_vtable = {
     .destroy = iree_hal_metal_allocator_destroy,
     .host_allocator = iree_hal_metal_allocator_host_allocator,
     .trim = iree_hal_metal_allocator_trim,
     .query_statistics = iree_hal_metal_allocator_query_statistics,
     .query_buffer_compatibility = iree_hal_metal_allocator_query_buffer_compatibility,
     .allocate_buffer = iree_hal_metal_allocator_allocate_buffer,
     .deallocate_buffer = iree_hal_metal_allocator_deallocate_buffer,
     .import_buffer = iree_hal_metal_allocator_import_buffer,
     .export_buffer = iree_hal_metal_allocator_export_buffer,
 };
	// 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

	#include "experimental/metal/direct_allocator.h"

	#import <Metal/Metal.h>

	#include "experimental/metal/metal_buffer.h"
	#include "iree/base/api.h"
	#include "iree/base/target_platform.h"
	#include "iree/base/tracing.h"
	#include "iree/hal/api.h"

	#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_ALLOCATION_TRACKING
	static const char* IREE_HAL_METAL_ALLOCATOR_ID = "METAL";
	#endif // IREE_TRACING_FEATURE_ALLOCATION_TRACKING

	typedef struct iree_hal_metal_allocator_t {
	// Abstract resource used for injecting reference counting and vtable; must be at offset 0.
	iree_hal_resource_t resource;

	// The device that this allocator is attached to.
	id<MTLDevice> device;
	// The command queue that we can use to issue commands to make buffer contents visible to CPU.
	#if defined(IREE_PLATFORM_MACOS)
	id<MTLCommandQueue> queue;
	#endif // IREE_PLATFORM_MACOS

	bool is_unified_memory;
	iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode;

	iree_allocator_t host_allocator;

	IREE_STATISTICS(iree_hal_allocator_statistics_t statistics;)
	} iree_hal_metal_allocator_t;

	static const iree_hal_allocator_vtable_t iree_hal_metal_allocator_vtable;

	static iree_hal_metal_allocator_t* iree_hal_metal_allocator_cast(iree_hal_allocator_t* base_value) {
	IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_allocator_vtable);
	return (iree_hal_metal_allocator_t*)base_value;
	}

	static const iree_hal_metal_allocator_t* iree_hal_metal_allocator_const_cast(
	const iree_hal_allocator_t* base_value) {
	IREE_HAL_ASSERT_TYPE(base_value, &iree_hal_metal_allocator_vtable);
	return (const iree_hal_metal_allocator_t*)base_value;
	}

	iree_status_t iree_hal_metal_allocator_create(
	id<MTLDevice> device,
	#if defined(IREE_PLATFORM_MACOS)
	id<MTLCommandQueue> queue,
	#endif // IREE_PLATFORM_MACOS
	iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode,
	iree_allocator_t host_allocator, iree_hal_allocator_t** out_allocator) {
	IREE_ASSERT_ARGUMENT(out_allocator);
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_hal_metal_allocator_t* allocator = NULL;
	iree_status_t status =
	iree_allocator_malloc(host_allocator, sizeof(allocator), (void*)&allocator);

	if (iree_status_is_ok(status)) {
	iree_hal_resource_initialize(&iree_hal_metal_allocator_vtable, &allocator->resource);
	allocator->device = [device retain]; // +1
	allocator->is_unified_memory = [device hasUnifiedMemory];
	allocator->resource_tracking_mode = resource_tracking_mode;
	allocator->host_allocator = host_allocator;

	out_allocator = (iree_hal_allocator_t)allocator;
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	static void iree_hal_metal_allocator_destroy(iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
	iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
	iree_allocator_t host_allocator = allocator->host_allocator;
	IREE_TRACE_ZONE_BEGIN(z0);

	[allocator->device release]; // -1
	iree_allocator_free(host_allocator, allocator);

	IREE_TRACE_ZONE_END(z0);
	}

	static iree_allocator_t iree_hal_metal_allocator_host_allocator(
	const iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
	iree_hal_metal_allocator_t* allocator = (iree_hal_metal_allocator_t*)base_allocator;
	return allocator->host_allocator;
	}

	#if defined(IREE_PLATFORM_MACOS)
	id<MTLCommandQueue> iree_hal_metal_allocator_command_queue(
	const iree_hal_allocator_t* base_allocator) {
	const iree_hal_metal_allocator_t* allocator = (const iree_hal_metal_allocator_t*)base_allocator;
	return allocator->queue;
	}
	#endif // IREE_PLATFORM_MACOS

	static iree_status_t iree_hal_metal_allocator_trim(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator) {
	return iree_ok_status();
	}

	static void iree_hal_metal_allocator_query_statistics(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator,
	iree_hal_allocator_statistics_t* IREE_RESTRICT out_statistics) {
	IREE_STATISTICS({
	iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
	memcpy(out_statistics, &allocator->statistics, sizeof(*out_statistics));
	});
	}

	static iree_hal_buffer_compatibility_t iree_hal_metal_allocator_query_buffer_compatibility(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator,
	iree_hal_buffer_params_t* IREE_RESTRICT params,
	iree_device_size_t* IREE_RESTRICT allocation_size) {
	// All buffers can be allocated on the heap.
	iree_hal_buffer_compatibility_t compatibility = IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE;

	if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) {
	compatibility \|= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER;
	}

	// Buffers can only be used on the queue if they are device visible.
	if (iree_all_bits_set(params->type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
	if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_TRANSFER)) {
	compatibility \|= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_TRANSFER;
	}
	if (iree_any_bit_set(params->usage, IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE)) {
	compatibility \|= IREE_HAL_BUFFER_COMPATIBILITY_QUEUE_DISPATCH;
	}
	}

	if (iree_all_bits_set(params->type,
	IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL \| IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) {
	// On iOS, we don't have device local + host visible memory. But given the unified memory
	// architecture, it's fine to just request host local + device visible memory.
	// On macOS, for unified memory architecture, it's similar to iOS. Otherwise, we can have
	// device local + host visible memory backed by Managed storage mode.
	iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
	if (allocator->is_unified_memory) {
	params->type &= ~(IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL \| IREE_HAL_MEMORY_TYPE_HOST_VISIBLE);
	params->type \|= IREE_HAL_MEMORY_TYPE_HOST_LOCAL \| IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
	// We have suggested other configurations than the original request. Mark accordingly.
	compatibility \|= IREE_HAL_BUFFER_COMPATIBILITY_LOW_PERFORMANCE;
	}
	}

	// We are now optimal.
	params->type &= ~IREE_HAL_MEMORY_TYPE_OPTIMAL;

	// Guard against the corner case where the requested buffer size is 0. The application is unlikely
	// to do anything when requesting a 0-byte buffer; but it can happen in real world use cases. So
	// we should at least not crash.
	if (allocation_size == 0) allocation_size = 4;

	// Align allocation sizes to 4 bytes so shaders operating on 32 bit types can act safely even on
	// buffer ranges that are not naturally aligned.
	allocation_size = iree_host_align(allocation_size, 4);

	return compatibility;
	}

	// Returns the corresponding Metal resource options controlling storage modes, CPU caching modes,
	// and hazard tracking modes for the given IREE HAL memory \|type\|.
	static MTLResourceOptions iree_hal_metal_select_resource_options(
	iree_hal_memory_type_t type, bool is_unified_memory, bool has_init_data,
	iree_hal_metal_resource_hazard_tracking_mode_t resource_tracking_mode) {
	MTLResourceOptions options;

	// Select a storage mode. There are four MTLStorageMode:
	// * Shared: The resource is stored in system memory and is accessible to both the CPU and
	// the GPU.
	// * Managed: The CPU and GPU may maintain separate copies of the resource, and any changes
	// must be explicitly synchronized.
	// * Private: The resource can be accessed only by the GPU.
	// * Memoryless: The resource's contents can be accessed only by the GPU and only exist
	// temporarily during a render pass.
	// macOS has all of the above; MTLStorageModeManaged is not available on iOS.
	//
	// The IREE HAL is modeled after explicit APIs like Vulkan. For buffers visible to both the host
	// and the device, we would like to opt in with the explicit version (MTLStorageManaged) when
	// possible because it should be more performant.
	if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)) {
	if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_HOST_VISIBLE)) {
	// Device local + host visible.
	// iree_hal_metal_allocator_query_buffer_compatibility guarantees that we only fall into this
	// case for macOS devices with non-uniform memory.
	#if defined(IREE_PLATFORM_MACOS)
	IREE_ASSERT(!is_unified_memory);
	options = MTLResourceStorageModeManaged;
	#else
	options = MTLResourceStorageModeShared;
	#endif // IREE_PLATFORM_MACOS
	} else {
	// Device local + host invisible.
	// For unified memory, update it to Shared storage mode if we have initial data.
	options = (is_unified_memory && has_init_data) ? MTLResourceStorageModeShared
	: MTLResourceStorageModePrivate;
	}
	} else {
	if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
	// Host local + device visible.
	options = MTLResourceStorageModeShared;
	} else {
	// Host local + device invisible.
	options = MTLResourceStorageModeShared;
	}
	}

	// Select a CPU cache mode.
	if (iree_all_bits_set(type, IREE_HAL_MEMORY_TYPE_HOST_CACHED)) {
	// The default CPU cache mode for the resource, which guarantees that read and write operations
	// are executed in the expected order.
	options \|= MTLResourceCPUCacheModeDefaultCache;
	} else {
	// A write-combined CPU cache mode that is optimized for resources that the CPU writes into, but
	// never reads.
	options \|= MTLResourceCPUCacheModeWriteCombined;
	}

	options \|= resource_tracking_mode == IREE_HAL_METAL_RESOURCE_HAZARD_TRACKING_MODE_TRACKED
	? MTLResourceHazardTrackingModeTracked
	: MTLResourceHazardTrackingModeUntracked;
	return options;
	}

	static iree_status_t iree_hal_metal_allocator_allocate_buffer(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator,
	const iree_hal_buffer_params_t* IREE_RESTRICT params, iree_device_size_t allocation_size,
	iree_const_byte_span_t initial_data, iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
	iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);
	IREE_TRACE_ZONE_BEGIN(z0);
	IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, allocation_size);

	// Coerce options into those required by the current device.
	iree_hal_buffer_params_t compat_params = *params;
	iree_hal_buffer_compatibility_t compatibility =
	iree_hal_metal_allocator_query_buffer_compatibility(base_allocator, &compat_params,
	&allocation_size);
	if (!iree_all_bits_set(compatibility, IREE_HAL_BUFFER_COMPATIBILITY_ALLOCATABLE)) {
	#if IREE_STATUS_MODE
	iree_bitfield_string_temp_t temp0, temp1, temp2;
	iree_string_view_t memory_type_str = iree_hal_memory_type_format(params->type, &temp0);
	iree_string_view_t usage_str = iree_hal_buffer_usage_format(params->usage, &temp1);
	iree_string_view_t compatibility_str =
	iree_hal_buffer_compatibility_format(compatibility, &temp2);
	IREE_TRACE_ZONE_END(z0);
	return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
	"allocator cannot allocate a buffer with the given parameters; "
	"memory_type=%.s, usage=%.s, compatibility=%.*s",
	(int)memory_type_str.size, memory_type_str.data, (int)usage_str.size,
	usage_str.data, (int)compatibility_str.size, compatibility_str.data);
	#else
	IREE_TRACE_ZONE_END(z0);
	return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
	"allocator cannot allocate a buffer with the given parameters");
	#endif // IREE_STATUS_MODE
	}

	bool has_init_data = !iree_const_byte_span_is_empty(initial_data);
	MTLResourceOptions options =
	iree_hal_metal_select_resource_options(compat_params.type, allocator->is_unified_memory,
	has_init_data, allocator->resource_tracking_mode);
	if (has_init_data && iree_all_bits_set(options, MTLResourceStorageModePrivate)) {
	IREE_TRACE_ZONE_END(z0);
	// This should only happen for macOS devices with non-uniform memory; for uniform memory
	// devices, the allocator compatibility would update it to use Shared storage mode.
	return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "private storage with initial data");
	}

	id<MTLBuffer> metal_buffer = nil;
	// If we chose shared storage mode, we can handle initial data with newByfferWithBytes directly.
	// Otherwise we just create the buffer here, and explicitly transfer range later.
	if (has_init_data && iree_all_bits_set(options, MTLResourceStorageModeShared)) {
	IREE_ASSERT_EQ(allocation_size, initial_data.data_length);
	metal_buffer = [allocator->device newBufferWithBytes:(void*)initial_data.data
	length:initial_data.data_length
	options:options]; // +1
	} else {
	metal_buffer = [allocator->device newBufferWithLength:allocation_size options:options]; // +1
	}
	if (!metal_buffer) {
	IREE_TRACE_ZONE_END(z0);
	return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED, "unable to allocate buffer");
	}
	iree_hal_buffer_t* buffer = NULL;
	iree_status_t status = iree_hal_metal_buffer_wrap(
	#if defined(IREE_PLATFORM_MACOS)
	allocator->queue,
	#endif // IREE_PLATFORM_MACOS
	metal_buffer, base_allocator, compat_params.type, compat_params.access, compat_params.usage,
	allocation_size, /byte_offset=/0, /byte_length=/allocation_size,
	iree_hal_buffer_release_callback_null(), &buffer); // +1

	#if defined(IREE_PLATFORM_MACOS)
	if (iree_status_is_ok(status) && has_init_data &&
	!iree_all_bits_set(options, MTLResourceStorageModePrivate)) {
	// Explicitly handle synchronization for Managed storage mode.
	if (iree_all_bits_set(options, MTLResourceStorageModeManaged)) {
	memcpy(metal_buffer.contents, initial_data.data, initial_data.data_length);
	[metal_buffer didModifyRange:NSMakeRange(0, initial_data.data_length)];
	}
	}
	#endif // IREE_PLATFORM_MACOS

	if (iree_status_is_ok(status)) {
	IREE_TRACE_ALLOC_NAMED(IREE_HAL_METAL_ALLOCATOR_ID, (void*)iree_hal_metal_buffer_handle(buffer),
	allocation_size);
	IREE_STATISTICS(iree_hal_allocator_statistics_record_alloc(
	&allocator->statistics, compat_params.type, allocation_size));
	*out_buffer = buffer;
	} else {
	if (buffer) iree_hal_buffer_release(buffer);
	}

	[metal_buffer release]; // -1

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	static void iree_hal_metal_allocator_deallocate_buffer(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator,
	iree_hal_buffer_t* IREE_RESTRICT base_buffer) {
	iree_hal_metal_allocator_t* allocator = iree_hal_metal_allocator_cast(base_allocator);

	IREE_TRACE_FREE_NAMED(IREE_HAL_METAL_ALLOCATOR_ID,
	(void*)iree_hal_metal_buffer_handle(base_buffer));
	IREE_STATISTICS(iree_hal_allocator_statistics_record_free(
	&allocator->statistics, iree_hal_buffer_memory_type(base_buffer),
	iree_hal_buffer_allocation_size(base_buffer)));

	iree_hal_buffer_destroy(base_buffer); // -1
	}

	static iree_status_t iree_hal_metal_allocator_import_buffer(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator,
	const iree_hal_buffer_params_t* IREE_RESTRICT params,
	iree_hal_external_buffer_t* IREE_RESTRICT external_buffer,
	iree_hal_buffer_release_callback_t release_callback,
	iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
	return iree_make_status(IREE_STATUS_UNAVAILABLE, "unsupported importing from external buffer");
	}

	static iree_status_t iree_hal_metal_allocator_export_buffer(
	iree_hal_allocator_t* IREE_RESTRICT base_allocator, iree_hal_buffer_t* IREE_RESTRICT 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) {
	return iree_make_status(IREE_STATUS_UNAVAILABLE, "unsupported exporting to external buffer");
	}

	static const iree_hal_allocator_vtable_t iree_hal_metal_allocator_vtable = {
	.destroy = iree_hal_metal_allocator_destroy,
	.host_allocator = iree_hal_metal_allocator_host_allocator,
	.trim = iree_hal_metal_allocator_trim,
	.query_statistics = iree_hal_metal_allocator_query_statistics,
	.query_buffer_compatibility = iree_hal_metal_allocator_query_buffer_compatibility,
	.allocate_buffer = iree_hal_metal_allocator_allocate_buffer,
	.deallocate_buffer = iree_hal_metal_allocator_deallocate_buffer,
	.import_buffer = iree_hal_metal_allocator_import_buffer,
	.export_buffer = iree_hal_metal_allocator_export_buffer,
	};