26
26
27
27
#include < zet_api.h>
28
28
29
+ #include " ur/usm_allocator_config.hpp"
29
30
#include " ur_bindings.hpp"
30
31
31
32
extern " C" {
@@ -96,6 +97,8 @@ static const bool IndirectAccessTrackingEnabled = [] {
96
97
nullptr ;
97
98
}();
98
99
100
+ static usm_settings::USMAllocatorConfig USMAllocatorConfigInstance;
101
+
99
102
// Map from L0 to PI result.
100
103
static inline pi_result mapError (ze_result_t Result) {
101
104
return ur2piResult (ze2urResult (Result));
@@ -775,18 +778,24 @@ pi_result _pi_context::initialize() {
775
778
auto createUSMAllocators = [this ](pi_device Device) {
776
779
SharedMemAllocContexts.emplace (
777
780
std::piecewise_construct, std::make_tuple (Device->ZeDevice ),
778
- std::make_tuple (std::unique_ptr<SystemMemory>(
779
- new USMSharedMemoryAlloc (this , Device))));
781
+ std::make_tuple (
782
+ std::unique_ptr<SystemMemory>(
783
+ new USMSharedMemoryAlloc (this , Device)),
784
+ USMAllocatorConfigInstance.Configs [usm_settings::MemType::Shared]));
780
785
781
786
SharedReadOnlyMemAllocContexts.emplace (
782
787
std::piecewise_construct, std::make_tuple (Device->ZeDevice ),
783
788
std::make_tuple (std::unique_ptr<SystemMemory>(
784
- new USMSharedReadOnlyMemoryAlloc (this , Device))));
789
+ new USMSharedReadOnlyMemoryAlloc (this , Device)),
790
+ USMAllocatorConfigInstance
791
+ .Configs [usm_settings::MemType::SharedReadOnly]));
785
792
786
793
DeviceMemAllocContexts.emplace (
787
794
std::piecewise_construct, std::make_tuple (Device->ZeDevice ),
788
- std::make_tuple (std::unique_ptr<SystemMemory>(
789
- new USMDeviceMemoryAlloc (this , Device))));
795
+ std::make_tuple (
796
+ std::unique_ptr<SystemMemory>(
797
+ new USMDeviceMemoryAlloc (this , Device)),
798
+ USMAllocatorConfigInstance.Configs [usm_settings::MemType::Device]));
790
799
};
791
800
792
801
// Recursive helper to call createUSMAllocators for all sub-devices
@@ -808,7 +817,8 @@ pi_result _pi_context::initialize() {
808
817
// are device-specific. Host allocations are not device-dependent therefore
809
818
// we don't need a map with device as key.
810
819
HostMemAllocContext = std::make_unique<USMAllocContext>(
811
- std::unique_ptr<SystemMemory>(new USMHostMemoryAlloc (this )));
820
+ std::unique_ptr<SystemMemory>(new USMHostMemoryAlloc (this )),
821
+ USMAllocatorConfigInstance.Configs [usm_settings::MemType::Host]);
812
822
813
823
// We may allocate memory to this root device so create allocators.
814
824
if (SingleRootDevice &&
@@ -8116,12 +8126,6 @@ pi_result USMHostMemoryAlloc::allocateImpl(void **ResultPtr, size_t Size,
8116
8126
return USMHostAllocImpl (ResultPtr, Context, nullptr , Size, Alignment);
8117
8127
}
8118
8128
8119
- MemType USMSharedMemoryAlloc::getMemTypeImpl () { return MemType::Shared; }
8120
-
8121
- MemType USMDeviceMemoryAlloc::getMemTypeImpl () { return MemType::Device; }
8122
-
8123
- MemType USMHostMemoryAlloc::getMemTypeImpl () { return MemType::Host; }
8124
-
8125
8129
void *USMMemoryAllocBase::allocate (size_t Size) {
8126
8130
void *Ptr = nullptr ;
8127
8131
@@ -8150,8 +8154,6 @@ void USMMemoryAllocBase::deallocate(void *Ptr, bool OwnZeMemHandle) {
8150
8154
}
8151
8155
}
8152
8156
8153
- MemType USMMemoryAllocBase::getMemType () { return getMemTypeImpl (); }
8154
-
8155
8157
pi_result piextUSMDeviceAlloc (void **ResultPtr, pi_context Context,
8156
8158
pi_device Device,
8157
8159
pi_usm_mem_properties *Properties, size_t Size,
@@ -9227,7 +9229,7 @@ pi_result _pi_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode,
9227
9229
// The host allocation may already exists, e.g. with imported
9228
9230
// host ptr, or in case of interop buffer.
9229
9231
if (!HostAllocation.ZeHandle ) {
9230
- if (enableBufferPooling () ) {
9232
+ if (USMAllocatorConfigInstance. EnableBuffers ) {
9231
9233
HostAllocation.ReleaseAction = allocation_t ::free;
9232
9234
PI_CALL (piextUSMHostAlloc (pi_cast<void **>(&ZeHandle), Context, nullptr ,
9233
9235
Size, getAlignment ()));
@@ -9280,7 +9282,7 @@ pi_result _pi_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode,
9280
9282
Allocation.Valid = true ;
9281
9283
return PI_SUCCESS;
9282
9284
} else { // Create device allocation
9283
- if (enableBufferPooling () ) {
9285
+ if (USMAllocatorConfigInstance. EnableBuffers ) {
9284
9286
Allocation.ReleaseAction = allocation_t ::free;
9285
9287
PI_CALL (piextUSMDeviceAlloc (pi_cast<void **>(&ZeHandle), Context,
9286
9288
Device, nullptr , Size, getAlignment ()));
@@ -9341,7 +9343,7 @@ pi_result _pi_buffer::getZeHandle(char *&ZeHandle, access_mode_t AccessMode,
9341
9343
// host ptr, or in case of interop buffer.
9342
9344
if (!HostAllocation.ZeHandle ) {
9343
9345
void *ZeHandleHost;
9344
- if (enableBufferPooling () ) {
9346
+ if (USMAllocatorConfigInstance. EnableBuffers ) {
9345
9347
HostAllocation.ReleaseAction = allocation_t ::free;
9346
9348
PI_CALL (piextUSMHostAlloc (&ZeHandleHost, Context, nullptr , Size,
9347
9349
getAlignment ()));
0 commit comments