Skip to content

Commit d56ca08

Browse files
committed
[SYCL] Remove OwnZeMemHandle from USMAllocator
OwnZeMemHandle is used to track whether an allocation should be freed by the runtime or not. There is no point in passing this flag to USMAllocator and then to USMFreeHelper since the allocation cannot come from the USMAllocator (it can only be from zeMemAlloc*).
1 parent edd9002 commit d56ca08

File tree

4 files changed

+32
-32
lines changed

4 files changed

+32
-32
lines changed

sycl/plugins/level_zero/pi_level_zero.cpp

Lines changed: 22 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -4189,8 +4189,7 @@ pi_result piMemRetain(pi_mem Mem) {
41894189
// If indirect access tracking is not enabled then this functions just performs
41904190
// zeMemFree. If indirect access tracking is enabled then reference counting is
41914191
// performed.
4192-
static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
4193-
bool OwnZeMemHandle = true) {
4192+
static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr) {
41944193
pi_platform Plt = Context->getPlatform();
41954194
std::unique_lock<pi_shared_mutex> ContextsLock(Plt->ContextsMutex,
41964195
std::defer_lock);
@@ -4210,8 +4209,7 @@ static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
42104209
Context->MemAllocs.erase(It);
42114210
}
42124211

4213-
if (OwnZeMemHandle)
4214-
ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
4212+
ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
42154213

42164214
if (IndirectAccessTrackingEnabled)
42174215
PI_CALL(ContextReleaseHelper(Context));
@@ -4220,7 +4218,7 @@ static pi_result ZeMemFreeHelper(pi_context Context, void *Ptr,
42204218
}
42214219

42224220
static pi_result USMFreeHelper(pi_context Context, void *Ptr,
4223-
bool OwnZeMemHandle);
4221+
bool OwnZeMemHandle = true);
42244222

42254223
pi_result piMemRelease(pi_mem Mem) {
42264224
PI_ASSERT(Mem, PI_ERROR_INVALID_MEM_OBJECT);
@@ -8085,10 +8083,8 @@ static pi_result USMHostAllocImpl(void **ResultPtr, pi_context Context,
80858083
return PI_SUCCESS;
80868084
}
80878085

8088-
static pi_result USMFreeImpl(pi_context Context, void *Ptr,
8089-
bool OwnZeMemHandle) {
8090-
if (OwnZeMemHandle)
8091-
ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
8086+
static pi_result USMFreeImpl(pi_context Context, void *Ptr) {
8087+
ZE_CALL(zeMemFree, (Context->ZeContext, Ptr));
80928088
return PI_SUCCESS;
80938089
}
80948090

@@ -8147,8 +8143,8 @@ void *USMMemoryAllocBase::allocate(size_t Size, size_t Alignment) {
81478143
return Ptr;
81488144
}
81498145

8150-
void USMMemoryAllocBase::deallocate(void *Ptr, bool OwnZeMemHandle) {
8151-
auto Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
8146+
void USMMemoryAllocBase::deallocate(void *Ptr) {
8147+
auto Res = USMFreeImpl(Context, Ptr);
81528148
if (Res != PI_SUCCESS) {
81538149
throw UsmAllocationException(Res);
81548150
}
@@ -8396,8 +8392,13 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
83968392
Context->MemAllocs.erase(It);
83978393
}
83988394

8395+
if (!OwnZeMemHandle) {
8396+
// Memory should not be freed
8397+
return PI_SUCCESS;
8398+
}
8399+
83998400
if (!UseUSMAllocator) {
8400-
pi_result Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
8401+
pi_result Res = USMFreeImpl(Context, Ptr);
84018402
if (IndirectAccessTrackingEnabled)
84028403
PI_CALL(ContextReleaseHelper(Context));
84038404
return Res;
@@ -8416,7 +8417,7 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
84168417
// If memory type is host release from host pool
84178418
if (ZeMemoryAllocationProperties.type == ZE_MEMORY_TYPE_HOST) {
84188419
try {
8419-
Context->HostMemAllocContext->deallocate(Ptr, OwnZeMemHandle);
8420+
Context->HostMemAllocContext->deallocate(Ptr);
84208421
} catch (const UsmAllocationException &Ex) {
84218422
return Ex.getError();
84228423
} catch (...) {
@@ -8444,16 +8445,16 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
84448445
PI_ASSERT(Device, PI_ERROR_INVALID_DEVICE);
84458446

84468447
auto DeallocationHelper =
8447-
[Context, Device, Ptr,
8448-
OwnZeMemHandle](std::unordered_map<ze_device_handle_t, USMAllocContext>
8449-
&AllocContextMap) {
8448+
[Context, Device,
8449+
Ptr](std::unordered_map<ze_device_handle_t, USMAllocContext>
8450+
&AllocContextMap) {
84508451
try {
84518452
auto It = AllocContextMap.find(Device->ZeDevice);
84528453
if (It == AllocContextMap.end())
84538454
return PI_ERROR_INVALID_VALUE;
84548455

84558456
// The right context is found, deallocate the pointer
8456-
It->second.deallocate(Ptr, OwnZeMemHandle);
8457+
It->second.deallocate(Ptr);
84578458
} catch (const UsmAllocationException &Ex) {
84588459
return Ex.getError();
84598460
}
@@ -8479,7 +8480,7 @@ static pi_result USMFreeHelper(pi_context Context, void *Ptr,
84798480
}
84808481
}
84818482

8482-
pi_result Res = USMFreeImpl(Context, Ptr, OwnZeMemHandle);
8483+
pi_result Res = USMFreeImpl(Context, Ptr);
84838484
if (SharedReadOnlyAllocsIterator != Context->SharedReadOnlyAllocs.end()) {
84848485
Context->SharedReadOnlyAllocs.erase(SharedReadOnlyAllocsIterator);
84858486
}
@@ -8494,7 +8495,7 @@ pi_result piextUSMFree(pi_context Context, void *Ptr) {
84948495
std::scoped_lock<pi_shared_mutex> Lock(
84958496
IndirectAccessTrackingEnabled ? Plt->ContextsMutex : Context->Mutex);
84968497

8497-
return USMFreeHelper(Context, Ptr, true /* OwnZeMemHandle */);
8498+
return USMFreeHelper(Context, Ptr);
84988499
}
84998500

85008501
pi_result piextKernelSetArgPointer(pi_kernel Kernel, pi_uint32 ArgIndex,
@@ -9410,11 +9411,11 @@ pi_result _pi_buffer::free() {
94109411
std::scoped_lock<pi_shared_mutex> Lock(
94119412
IndirectAccessTrackingEnabled ? Plt->ContextsMutex : Context->Mutex);
94129413

9413-
PI_CALL(USMFreeHelper(Context, ZeHandle, true));
9414+
PI_CALL(USMFreeHelper(Context, ZeHandle));
94149415
break;
94159416
}
94169417
case allocation_t::free_native:
9417-
PI_CALL(ZeMemFreeHelper(Context, ZeHandle, true));
9418+
PI_CALL(ZeMemFreeHelper(Context, ZeHandle));
94189419
break;
94199420
case allocation_t::unimport:
94209421
ZeUSMImport.doZeUSMRelease(Context->getPlatform()->ZeDriver, ZeHandle);

sycl/plugins/level_zero/pi_level_zero.hpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ class USMMemoryAllocBase : public SystemMemory {
124124
: Context{Ctx}, Device{Dev} {}
125125
void *allocate(size_t Size) override final;
126126
void *allocate(size_t Size, size_t Alignment) override final;
127-
void deallocate(void *Ptr, bool OwnZeMemHandle) override final;
127+
void deallocate(void *Ptr) override final;
128128
};
129129

130130
// Allocation routines for shared memory type

sycl/plugins/unified_runtime/ur/usm_allocator.cpp

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,7 @@ class USMAllocContext::USMAllocImpl {
290290

291291
void *allocate(size_t Size, size_t Alignment, bool &FromPool);
292292
void *allocate(size_t Size, bool &FromPool);
293-
void deallocate(void *Ptr, bool &ToPool, bool OwnZeMemHandle);
293+
void deallocate(void *Ptr, bool &ToPool);
294294

295295
SystemMemory &getMemHandle() { return *MemHandle; }
296296

@@ -332,7 +332,7 @@ Slab::Slab(Bucket &Bkt)
332332

333333
Slab::~Slab() {
334334
unregSlab(*this);
335-
bucket.getMemHandle().deallocate(MemPtr, true /* OwnZeMemHandle */);
335+
bucket.getMemHandle().deallocate(MemPtr);
336336
}
337337

338338
// Return the index of the first available chunk, -1 otherwize
@@ -737,8 +737,7 @@ Bucket &USMAllocContext::USMAllocImpl::findBucket(size_t Size) {
737737
return *(*It);
738738
}
739739

740-
void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
741-
bool OwnZeMemHandle) {
740+
void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool) {
742741
auto *SlabPtr = AlignPtrDown(Ptr, SlabMinSize());
743742

744743
// Lock the map on read
@@ -748,7 +747,7 @@ void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
748747
auto Slabs = getKnownSlabs().equal_range(SlabPtr);
749748
if (Slabs.first == Slabs.second) {
750749
Lk.unlock();
751-
getMemHandle().deallocate(Ptr, OwnZeMemHandle);
750+
getMemHandle().deallocate(Ptr);
752751
return;
753752
}
754753

@@ -779,7 +778,7 @@ void USMAllocContext::USMAllocImpl::deallocate(void *Ptr, bool &ToPool,
779778
// There is a rare case when we have a pointer from system allocation next
780779
// to some slab with an entry in the map. So we find a slab
781780
// but the range checks fail.
782-
getMemHandle().deallocate(Ptr, OwnZeMemHandle);
781+
getMemHandle().deallocate(Ptr);
783782
}
784783

785784
USMAllocContext::USMAllocContext(std::unique_ptr<SystemMemory> MemHandle,
@@ -813,9 +812,9 @@ void *USMAllocContext::allocate(size_t size, size_t alignment) {
813812
return Ptr;
814813
}
815814

816-
void USMAllocContext::deallocate(void *ptr, bool OwnZeMemHandle) {
815+
void USMAllocContext::deallocate(void *ptr) {
817816
bool ToPool;
818-
pImpl->deallocate(ptr, ToPool, OwnZeMemHandle);
817+
pImpl->deallocate(ptr, ToPool);
819818

820819
if (pImpl->getParams().PoolTrace > 2) {
821820
auto MT = pImpl->getParams().memoryTypeName;

sycl/plugins/unified_runtime/ur/usm_allocator.hpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ class SystemMemory {
1717
public:
1818
virtual void *allocate(size_t size) = 0;
1919
virtual void *allocate(size_t size, size_t aligned) = 0;
20-
virtual void deallocate(void *ptr, bool OwnZeMemHandle) = 0;
20+
virtual void deallocate(void *ptr) = 0;
2121
virtual ~SystemMemory() = default;
2222
};
2323

@@ -68,7 +68,7 @@ class USMAllocContext {
6868

6969
void *allocate(size_t size);
7070
void *allocate(size_t size, size_t alignment);
71-
void deallocate(void *ptr, bool OwnZeMemHandle);
71+
void deallocate(void *ptr);
7272

7373
private:
7474
std::unique_ptr<USMAllocImpl> pImpl;

0 commit comments

Comments
 (0)