experimental/hip/memory_pools.c - 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/hip/memory_pools.h"

 #include "experimental/hip/dynamic_symbols.h"
 #include "experimental/hip/hip_buffer.h"
 #include "experimental/hip/status_util.h"

 // NOTE: these are currently global for all devices; we could make
 // device-specific ones by malloc() and leaking (with LSAN note) unique string
 // values instead.
 #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_ALLOCATION_TRACKING
 static const char* IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID =
     "HIP pool: device-local reserved";
 static const char* IREE_HAL_HIP_OTHER_POOL_RESERVED_ID =
     "HIP pool: other reserved";
 #endif  // IREE_TRACING_FEATURE_ALLOCATION_TRACKING

 static iree_status_t iree_hal_hip_create_memory_pool(
     const iree_hal_hip_dynamic_symbols_t* hip_symbols, hipDevice_t hip_device,
     iree_hal_hip_memory_pool_params_t params,
     hipMemPool_t* IREE_RESTRICT out_pool) {
   *out_pool = NULL;

   hipMemPoolProps pool_props = {
       .allocType = hipMemAllocationTypePinned,
       // TODO: allow sharing of certain pool memory types by fd/HANDLE.
       .handleTypes = hipMemHandleTypeNone,
       .location =
           {
               .type = hipMemLocationTypeDevice,
               .id = hip_device,
           },
       .reserved = {0},
       .win32SecurityAttributes = NULL,
   };

   hipMemPool_t pool = NULL;

   IREE_HIP_RETURN_IF_ERROR(hip_symbols, hipMemPoolCreate(&pool, &pool_props),
                            "hipMemPoolCreate");
   iree_status_t status = IREE_HIP_RESULT_TO_STATUS(
       hip_symbols,
       hipMemPoolSetAttribute(pool, hipMemPoolAttrReleaseThreshold,
                              &params.release_threshold),
       "hipMemPoolSetAttribute");

   if (iree_status_is_ok(status)) {
     *out_pool = pool;
   } else {
     IREE_HIP_IGNORE_ERROR(hip_symbols, hipMemPoolDestroy(pool));
   }
   return status;
 }

 iree_status_t iree_hal_hip_memory_pools_initialize(
     const iree_hal_hip_dynamic_symbols_t* hip_symbols, hipDevice_t hip_device,
     const iree_hal_hip_memory_pooling_params_t* pooling_params,
     iree_allocator_t host_allocator,
     iree_hal_hip_memory_pools_t* IREE_RESTRICT out_pools) {
   IREE_ASSERT_ARGUMENT(hip_symbols);
   IREE_ASSERT_ARGUMENT(pooling_params);
   IREE_ASSERT_ARGUMENT(out_pools);
   IREE_TRACE_ZONE_BEGIN(z0);

   memset(out_pools, 0, sizeof(*out_pools));
   out_pools->hip_symbols = hip_symbols;
   out_pools->host_allocator = host_allocator;

   iree_status_t status = iree_ok_status();

   if (iree_status_is_ok(status)) {
     status = iree_hal_hip_create_memory_pool(hip_symbols, hip_device,
                                              pooling_params->device_local,
                                              &out_pools->device_local);
   }

   if (iree_status_is_ok(status)) {
     status = iree_hal_hip_create_memory_pool(
         hip_symbols, hip_device, pooling_params->other, &out_pools->other);
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 void iree_hal_hip_memory_pools_deinitialize(
     iree_hal_hip_memory_pools_t* pools) {
   IREE_TRACE_ZONE_BEGIN(z0);

   if (pools->device_local) {
     IREE_HIP_IGNORE_ERROR(pools->hip_symbols,
                           hipMemPoolDestroy(pools->device_local));
     pools->device_local = NULL;
   }

   if (pools->other) {
     IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipMemPoolDestroy(pools->other));
     pools->other = NULL;
   }

   IREE_TRACE_ZONE_END(z0);
 }

 static void iree_hal_hip_memory_pool_track_alloc(
     iree_hal_hip_memory_pools_t* pools, iree_hal_buffer_t* buffer) {
   bool is_device_local = iree_all_bits_set(iree_hal_buffer_memory_type(buffer),
                                            IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL);
   (void)is_device_local;
   iree_device_size_t allocation_size = iree_hal_buffer_allocation_size(buffer);
   (void)allocation_size;
   IREE_TRACE_ALLOC_NAMED(
       is_device_local ? IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID
                       : IREE_HAL_HIP_OTHER_POOL_RESERVED_ID,
       (void*)iree_hal_hip_buffer_device_pointer(buffer), allocation_size);
   IREE_STATISTICS({
     iree_atomic_int64_t* bytes_allocated =
         is_device_local ? &pools->statistics.device_bytes_allocated
                         : &pools->statistics.host_bytes_allocated;
     iree_atomic_fetch_add_int64(bytes_allocated, allocation_size,
                                 iree_memory_order_relaxed);
   });
 }

 static void iree_hal_hip_memory_pool_track_free(
     iree_hal_hip_memory_pools_t* pools, iree_hal_buffer_t* buffer) {
   bool is_device_local = iree_all_bits_set(iree_hal_buffer_memory_type(buffer),
                                            IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL);
   (void)is_device_local;
   IREE_TRACE_FREE_NAMED(is_device_local
                             ? IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID
                             : IREE_HAL_HIP_OTHER_POOL_RESERVED_ID,
                         (void*)iree_hal_hip_buffer_device_pointer(buffer));
   IREE_STATISTICS({
     iree_atomic_int64_t* bytes_freed =
         is_device_local ? &pools->statistics.device_bytes_freed
                         : &pools->statistics.host_bytes_freed;
     iree_device_size_t allocation_size =
         iree_hal_buffer_allocation_size(buffer);
     iree_atomic_fetch_add_int64(bytes_freed, allocation_size,
                                 iree_memory_order_relaxed);
   });
 }

 void iree_hal_hip_memory_pools_merge_statistics(
     iree_hal_hip_memory_pools_t* pools,
     iree_hal_allocator_statistics_t* statistics) {
   IREE_STATISTICS({
     statistics->device_bytes_allocated = iree_atomic_load_int64(
         &pools->statistics.device_bytes_allocated, iree_memory_order_relaxed);
     statistics->host_bytes_allocated = iree_atomic_load_int64(
         &pools->statistics.host_bytes_allocated, iree_memory_order_relaxed);
     statistics->device_bytes_freed = iree_atomic_load_int64(
         &pools->statistics.device_bytes_freed, iree_memory_order_relaxed);
     statistics->host_bytes_freed = iree_atomic_load_int64(
         &pools->statistics.host_bytes_freed, iree_memory_order_relaxed);

     if (pools->device_local) {
       uint64_t pool_peak = 0;
       IREE_HIP_IGNORE_ERROR(
           pools->hip_symbols,
           hipMemPoolGetAttribute(pools->device_local, hipMemPoolAttrUsedMemHigh,
                                  &pool_peak));
       statistics->device_bytes_peak += (iree_device_size_t)pool_peak;
     }
     if (pools->other) {
       uint64_t pool_peak = 0;
       IREE_HIP_IGNORE_ERROR(
           pools->hip_symbols,
           hipMemPoolGetAttribute(pools->other, hipMemPoolAttrUsedMemHigh,
                                  &pool_peak));
       statistics->host_bytes_peak += (iree_device_size_t)pool_peak;
     }
   });
 }

 iree_status_t iree_hal_hip_memory_pools_trim(
     iree_hal_hip_memory_pools_t* pools,
     const iree_hal_hip_memory_pooling_params_t* pooling_params) {
   IREE_HIP_RETURN_IF_ERROR(
       pools->hip_symbols,
       hipMemPoolTrimTo(pools->device_local,
                        pooling_params->device_local.minimum_capacity),
       "hipMemPoolTrimTo");
   IREE_HIP_RETURN_IF_ERROR(
       pools->hip_symbols,
       hipMemPoolTrimTo(pools->other, pooling_params->other.minimum_capacity),
       "hipMemPoolTrimTo");
   return iree_ok_status();
 }

 // NOTE: this is only issued if the buffer is destroyed without having had been
 // scheduled for deallocation asynchronously. When a buffer is scheduled we drop
 // the release callback so that this isn't called and we don't double-free.
 static void iree_hal_hip_async_buffer_release_callback(
     void* user_data, iree_hal_buffer_t* buffer) {
   iree_hal_hip_memory_pools_t* pools = (iree_hal_hip_memory_pools_t*)user_data;
   IREE_TRACE_ZONE_BEGIN(z0);

   hipDeviceptr_t device_ptr = iree_hal_hip_buffer_device_pointer(buffer);
   IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipFree(device_ptr));
   iree_hal_hip_memory_pool_track_free(pools, buffer);

   IREE_TRACE_ZONE_END(z0);
 }

 iree_status_t iree_hal_hip_memory_pools_allocate(
     iree_hal_hip_memory_pools_t* pools, hipStream_t stream,
     iree_hal_allocator_pool_t pool, iree_hal_buffer_params_t params,
     iree_device_size_t allocation_size,
     iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
   IREE_TRACE_ZONE_BEGIN(z0);
   IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)allocation_size);

   iree_hal_buffer_params_canonicalize(&params);

   // TODO: more pools and better selection; this is coarsely deciding between
   // only device local (variables, constants, transients) and other (staging,
   // external) but could use more buffer properties (including usage/export
   // flags) to better isolate the different usage patterns and keep the pools
   // operating with reasonable limits. We should be using the |pool| arg.
   hipMemPool_t memory_pool =
       iree_all_bits_set(params.type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)
           ? pools->device_local
           : pools->other;

   hipDeviceptr_t device_ptr = NULL;
   iree_status_t status = IREE_HIP_RESULT_TO_STATUS(
       pools->hip_symbols,
       hipMallocFromPoolAsync(&device_ptr, (size_t)allocation_size, memory_pool,
                              stream),
       "hipMallocFromPoolAsync");

   // Wrap the allocated HIP buffer in a HAL buffer.
   // NOTE: we don't provide a device allocator because we didn't allocate from
   // one and instead we use a release callback to perform the free if the user
   // doesn't dealloca the buffer.
   iree_hal_buffer_t* buffer = NULL;
   if (iree_status_is_ok(status)) {
     iree_hal_buffer_release_callback_t release_callback = {
         .fn = iree_hal_hip_async_buffer_release_callback,
         .user_data = pools,
     };
     status = iree_hal_hip_buffer_wrap(
         /*device_allocator=*/NULL, params.type, params.access, params.usage,
         allocation_size, /*byte_offset=*/0,
         /*byte_length=*/allocation_size, IREE_HAL_HIP_BUFFER_TYPE_ASYNC,
         device_ptr, /*host_ptr=*/NULL, release_callback, pools->host_allocator,
         &buffer);
   }

   if (iree_status_is_ok(status)) {
     // Update statistics (note that it may not yet be accurate).
     iree_hal_hip_memory_pool_track_alloc(pools, buffer);
     *out_buffer = buffer;
   } else if (buffer) {
     iree_hal_buffer_release(buffer);
   } else {
     IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipFreeAsync(device_ptr, stream));
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 iree_status_t iree_hal_hip_memory_pools_deallocate(
     iree_hal_hip_memory_pools_t* pools, hipStream_t stream,
     iree_hal_buffer_t* buffer) {
   IREE_TRACE_ZONE_BEGIN(z0);
   IREE_TRACE_ZONE_APPEND_VALUE_I64(
       z0, (int64_t)iree_hal_buffer_allocation_size(buffer));

   // Only process the request if the buffer came from an async pool.
   // We may get requests for deallocations on ones that didn't if one part of
   // the application allocated the buffer synchronously and another deallocated
   // it asynchronously.
   iree_status_t status = iree_ok_status();
   if (iree_hal_hip_buffer_type(buffer) == IREE_HAL_HIP_BUFFER_TYPE_ASYNC) {
     // Try to schedule the buffer for freeing.
     hipDeviceptr_t device_ptr = iree_hal_hip_buffer_device_pointer(buffer);
     status = IREE_HIP_RESULT_TO_STATUS(
         pools->hip_symbols, hipFreeAsync(device_ptr, stream), "hipFreeAsync");
     if (iree_status_is_ok(status)) {
       // Drop the release callback so that we don't try to double-free the
       // buffer. Note that we only do this if the HIP free succeeded as
       // otherwise we still need to synchronously deallocate the buffer when it
       // is destroyed.
       iree_hal_hip_buffer_drop_release_callback(buffer);

       // Update statistics (note that it may not yet be accurate).
       iree_hal_hip_memory_pool_track_free(pools, buffer);
     }
   } else {
     // Not allocated via alloca, ignore.
     IREE_TRACE_ZONE_APPEND_TEXT(z0, "ignored sync allocation");
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }
	// 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/hip/memory_pools.h"

	#include "experimental/hip/dynamic_symbols.h"
	#include "experimental/hip/hip_buffer.h"
	#include "experimental/hip/status_util.h"

	// NOTE: these are currently global for all devices; we could make
	// device-specific ones by malloc() and leaking (with LSAN note) unique string
	// values instead.
	#if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_ALLOCATION_TRACKING
	static const char* IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID =
	"HIP pool: device-local reserved";
	static const char* IREE_HAL_HIP_OTHER_POOL_RESERVED_ID =
	"HIP pool: other reserved";
	#endif // IREE_TRACING_FEATURE_ALLOCATION_TRACKING

	static iree_status_t iree_hal_hip_create_memory_pool(
	const iree_hal_hip_dynamic_symbols_t* hip_symbols, hipDevice_t hip_device,
	iree_hal_hip_memory_pool_params_t params,
	hipMemPool_t* IREE_RESTRICT out_pool) {
	*out_pool = NULL;

	hipMemPoolProps pool_props = {
	.allocType = hipMemAllocationTypePinned,
	// TODO: allow sharing of certain pool memory types by fd/HANDLE.
	.handleTypes = hipMemHandleTypeNone,
	.location =
	{
	.type = hipMemLocationTypeDevice,
	.id = hip_device,
	},
	.reserved = {0},
	.win32SecurityAttributes = NULL,
	};

	hipMemPool_t pool = NULL;

	IREE_HIP_RETURN_IF_ERROR(hip_symbols, hipMemPoolCreate(&pool, &pool_props),
	"hipMemPoolCreate");
	iree_status_t status = IREE_HIP_RESULT_TO_STATUS(
	hip_symbols,
	hipMemPoolSetAttribute(pool, hipMemPoolAttrReleaseThreshold,
	&params.release_threshold),
	"hipMemPoolSetAttribute");

	if (iree_status_is_ok(status)) {
	*out_pool = pool;
	} else {
	IREE_HIP_IGNORE_ERROR(hip_symbols, hipMemPoolDestroy(pool));
	}
	return status;
	}

	iree_status_t iree_hal_hip_memory_pools_initialize(
	const iree_hal_hip_dynamic_symbols_t* hip_symbols, hipDevice_t hip_device,
	const iree_hal_hip_memory_pooling_params_t* pooling_params,
	iree_allocator_t host_allocator,
	iree_hal_hip_memory_pools_t* IREE_RESTRICT out_pools) {
	IREE_ASSERT_ARGUMENT(hip_symbols);
	IREE_ASSERT_ARGUMENT(pooling_params);
	IREE_ASSERT_ARGUMENT(out_pools);
	IREE_TRACE_ZONE_BEGIN(z0);

	memset(out_pools, 0, sizeof(*out_pools));
	out_pools->hip_symbols = hip_symbols;
	out_pools->host_allocator = host_allocator;

	iree_status_t status = iree_ok_status();

	if (iree_status_is_ok(status)) {
	status = iree_hal_hip_create_memory_pool(hip_symbols, hip_device,
	pooling_params->device_local,
	&out_pools->device_local);
	}

	if (iree_status_is_ok(status)) {
	status = iree_hal_hip_create_memory_pool(
	hip_symbols, hip_device, pooling_params->other, &out_pools->other);
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	void iree_hal_hip_memory_pools_deinitialize(
	iree_hal_hip_memory_pools_t* pools) {
	IREE_TRACE_ZONE_BEGIN(z0);

	if (pools->device_local) {
	IREE_HIP_IGNORE_ERROR(pools->hip_symbols,
	hipMemPoolDestroy(pools->device_local));
	pools->device_local = NULL;
	}

	if (pools->other) {
	IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipMemPoolDestroy(pools->other));
	pools->other = NULL;
	}

	IREE_TRACE_ZONE_END(z0);
	}

	static void iree_hal_hip_memory_pool_track_alloc(
	iree_hal_hip_memory_pools_t* pools, iree_hal_buffer_t* buffer) {
	bool is_device_local = iree_all_bits_set(iree_hal_buffer_memory_type(buffer),
	IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL);
	(void)is_device_local;
	iree_device_size_t allocation_size = iree_hal_buffer_allocation_size(buffer);
	(void)allocation_size;
	IREE_TRACE_ALLOC_NAMED(
	is_device_local ? IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID
	: IREE_HAL_HIP_OTHER_POOL_RESERVED_ID,
	(void*)iree_hal_hip_buffer_device_pointer(buffer), allocation_size);
	IREE_STATISTICS({
	iree_atomic_int64_t* bytes_allocated =
	is_device_local ? &pools->statistics.device_bytes_allocated
	: &pools->statistics.host_bytes_allocated;
	iree_atomic_fetch_add_int64(bytes_allocated, allocation_size,
	iree_memory_order_relaxed);
	});
	}

	static void iree_hal_hip_memory_pool_track_free(
	iree_hal_hip_memory_pools_t* pools, iree_hal_buffer_t* buffer) {
	bool is_device_local = iree_all_bits_set(iree_hal_buffer_memory_type(buffer),
	IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL);
	(void)is_device_local;
	IREE_TRACE_FREE_NAMED(is_device_local
	? IREE_HAL_HIP_DEVICE_LOCAL_POOL_RESERVED_ID
	: IREE_HAL_HIP_OTHER_POOL_RESERVED_ID,
	(void*)iree_hal_hip_buffer_device_pointer(buffer));
	IREE_STATISTICS({
	iree_atomic_int64_t* bytes_freed =
	is_device_local ? &pools->statistics.device_bytes_freed
	: &pools->statistics.host_bytes_freed;
	iree_device_size_t allocation_size =
	iree_hal_buffer_allocation_size(buffer);
	iree_atomic_fetch_add_int64(bytes_freed, allocation_size,
	iree_memory_order_relaxed);
	});
	}

	void iree_hal_hip_memory_pools_merge_statistics(
	iree_hal_hip_memory_pools_t* pools,
	iree_hal_allocator_statistics_t* statistics) {
	IREE_STATISTICS({
	statistics->device_bytes_allocated = iree_atomic_load_int64(
	&pools->statistics.device_bytes_allocated, iree_memory_order_relaxed);
	statistics->host_bytes_allocated = iree_atomic_load_int64(
	&pools->statistics.host_bytes_allocated, iree_memory_order_relaxed);
	statistics->device_bytes_freed = iree_atomic_load_int64(
	&pools->statistics.device_bytes_freed, iree_memory_order_relaxed);
	statistics->host_bytes_freed = iree_atomic_load_int64(
	&pools->statistics.host_bytes_freed, iree_memory_order_relaxed);

	if (pools->device_local) {
	uint64_t pool_peak = 0;
	IREE_HIP_IGNORE_ERROR(
	pools->hip_symbols,
	hipMemPoolGetAttribute(pools->device_local, hipMemPoolAttrUsedMemHigh,
	&pool_peak));
	statistics->device_bytes_peak += (iree_device_size_t)pool_peak;
	}
	if (pools->other) {
	uint64_t pool_peak = 0;
	IREE_HIP_IGNORE_ERROR(
	pools->hip_symbols,
	hipMemPoolGetAttribute(pools->other, hipMemPoolAttrUsedMemHigh,
	&pool_peak));
	statistics->host_bytes_peak += (iree_device_size_t)pool_peak;
	}
	});
	}

	iree_status_t iree_hal_hip_memory_pools_trim(
	iree_hal_hip_memory_pools_t* pools,
	const iree_hal_hip_memory_pooling_params_t* pooling_params) {
	IREE_HIP_RETURN_IF_ERROR(
	pools->hip_symbols,
	hipMemPoolTrimTo(pools->device_local,
	pooling_params->device_local.minimum_capacity),
	"hipMemPoolTrimTo");
	IREE_HIP_RETURN_IF_ERROR(
	pools->hip_symbols,
	hipMemPoolTrimTo(pools->other, pooling_params->other.minimum_capacity),
	"hipMemPoolTrimTo");
	return iree_ok_status();
	}

	// NOTE: this is only issued if the buffer is destroyed without having had been
	// scheduled for deallocation asynchronously. When a buffer is scheduled we drop
	// the release callback so that this isn't called and we don't double-free.
	static void iree_hal_hip_async_buffer_release_callback(
	void* user_data, iree_hal_buffer_t* buffer) {
	iree_hal_hip_memory_pools_t* pools = (iree_hal_hip_memory_pools_t*)user_data;
	IREE_TRACE_ZONE_BEGIN(z0);

	hipDeviceptr_t device_ptr = iree_hal_hip_buffer_device_pointer(buffer);
	IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipFree(device_ptr));
	iree_hal_hip_memory_pool_track_free(pools, buffer);

	IREE_TRACE_ZONE_END(z0);
	}

	iree_status_t iree_hal_hip_memory_pools_allocate(
	iree_hal_hip_memory_pools_t* pools, hipStream_t stream,
	iree_hal_allocator_pool_t pool, iree_hal_buffer_params_t params,
	iree_device_size_t allocation_size,
	iree_hal_buffer_t** IREE_RESTRICT out_buffer) {
	IREE_TRACE_ZONE_BEGIN(z0);
	IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)allocation_size);

	iree_hal_buffer_params_canonicalize(&params);

	// TODO: more pools and better selection; this is coarsely deciding between
	// only device local (variables, constants, transients) and other (staging,
	// external) but could use more buffer properties (including usage/export
	// flags) to better isolate the different usage patterns and keep the pools
	// operating with reasonable limits. We should be using the \|pool\| arg.
	hipMemPool_t memory_pool =
	iree_all_bits_set(params.type, IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL)
	? pools->device_local
	: pools->other;

	hipDeviceptr_t device_ptr = NULL;
	iree_status_t status = IREE_HIP_RESULT_TO_STATUS(
	pools->hip_symbols,
	hipMallocFromPoolAsync(&device_ptr, (size_t)allocation_size, memory_pool,
	stream),
	"hipMallocFromPoolAsync");

	// Wrap the allocated HIP buffer in a HAL buffer.
	// NOTE: we don't provide a device allocator because we didn't allocate from
	// one and instead we use a release callback to perform the free if the user
	// doesn't dealloca the buffer.
	iree_hal_buffer_t* buffer = NULL;
	if (iree_status_is_ok(status)) {
	iree_hal_buffer_release_callback_t release_callback = {
	.fn = iree_hal_hip_async_buffer_release_callback,
	.user_data = pools,
	};
	status = iree_hal_hip_buffer_wrap(
	/device_allocator=/NULL, params.type, params.access, params.usage,
	allocation_size, /byte_offset=/0,
	/byte_length=/allocation_size, IREE_HAL_HIP_BUFFER_TYPE_ASYNC,
	device_ptr, /host_ptr=/NULL, release_callback, pools->host_allocator,
	&buffer);
	}

	if (iree_status_is_ok(status)) {
	// Update statistics (note that it may not yet be accurate).
	iree_hal_hip_memory_pool_track_alloc(pools, buffer);
	*out_buffer = buffer;
	} else if (buffer) {
	iree_hal_buffer_release(buffer);
	} else {
	IREE_HIP_IGNORE_ERROR(pools->hip_symbols, hipFreeAsync(device_ptr, stream));
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	iree_status_t iree_hal_hip_memory_pools_deallocate(
	iree_hal_hip_memory_pools_t* pools, hipStream_t stream,
	iree_hal_buffer_t* buffer) {
	IREE_TRACE_ZONE_BEGIN(z0);
	IREE_TRACE_ZONE_APPEND_VALUE_I64(
	z0, (int64_t)iree_hal_buffer_allocation_size(buffer));

	// Only process the request if the buffer came from an async pool.
	// We may get requests for deallocations on ones that didn't if one part of
	// the application allocated the buffer synchronously and another deallocated
	// it asynchronously.
	iree_status_t status = iree_ok_status();
	if (iree_hal_hip_buffer_type(buffer) == IREE_HAL_HIP_BUFFER_TYPE_ASYNC) {
	// Try to schedule the buffer for freeing.
	hipDeviceptr_t device_ptr = iree_hal_hip_buffer_device_pointer(buffer);
	status = IREE_HIP_RESULT_TO_STATUS(
	pools->hip_symbols, hipFreeAsync(device_ptr, stream), "hipFreeAsync");
	if (iree_status_is_ok(status)) {
	// Drop the release callback so that we don't try to double-free the
	// buffer. Note that we only do this if the HIP free succeeded as
	// otherwise we still need to synchronously deallocate the buffer when it
	// is destroyed.
	iree_hal_hip_buffer_drop_release_callback(buffer);

	// Update statistics (note that it may not yet be accurate).
	iree_hal_hip_memory_pool_track_free(pools, buffer);
	}
	} else {
	// Not allocated via alloca, ignore.
	IREE_TRACE_ZONE_APPEND_TEXT(z0, "ignored sync allocation");
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}