Skip to content

Commit 64ff65b

Browse files
AllanZyneAllanZyne
committed
wip
1 parent afdfb91 commit 64ff65b

File tree

5 files changed

+79
-170
lines changed

5 files changed

+79
-170
lines changed

‎source/loader/layers/sanitizer/msan/msan_ddi.cpp‎

Lines changed: 71 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
4545
UR_CALL(DI->allocShadowMemory(Context));
4646
}
4747
CI->DeviceList.emplace_back(hDevice);
48-
CI->AllocInfosMap[hDevice];
4948
}
5049
return UR_RESULT_SUCCESS;
5150
}
@@ -1306,23 +1305,25 @@ ur_result_t UR_APICALL urEnqueueUSMFill(
13061305
///< must not refer to an element of the phEventWaitList array.
13071306
) {
13081307
auto pfnUSMFill = getContext()->urDdiTable.Enqueue.pfnUSMFill;
1309-
13101308
getContext()->logger.debug("==== urEnqueueUSMFill");
13111309

1312-
auto Mem = (uptr)pMem;
1310+
UR_CALL(pfnUSMFill(hQueue, pMem, patternSize, pPattern, size,
1311+
numEventsInWaitList, phEventWaitList, phEvent));
1312+
1313+
const auto Mem = (uptr)pMem;
13131314
auto MemInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Mem);
13141315
if (MemInfoItOp) {
13151316
auto MemInfo = (*MemInfoItOp)->second;
13161317

13171318
const auto &DeviceInfo =
13181319
getMsanInterceptor()->getDeviceInfo(MemInfo->Device);
1319-
UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(
1320-
hQueue, Mem, size, 0, numEventsInWaitList, phEventWaitList,
1321-
phEvent));
1320+
const auto MemShadow = DeviceInfo->Shadow->MemToShadow(Mem);
1321+
1322+
UR_CALL(EnqueueUSMBlockingSet(hQueue, (void *)MemShadow, 0, size, 1,
1323+
phEvent, phEvent));
13221324
}
13231325

1324-
return pfnUSMFill(hQueue, pMem, patternSize, pPattern, size,
1325-
numEventsInWaitList, phEventWaitList, phEvent);
1326+
return UR_RESULT_SUCCESS;
13261327
}
13271328

13281329
///////////////////////////////////////////////////////////////////////////////
@@ -1349,7 +1350,10 @@ ur_result_t UR_APICALL urEnqueueUSMMemcpy(
13491350
auto pfnUSMMemcpy = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy;
13501351
getContext()->logger.debug("==== pfnUSMMemcpy");
13511352

1352-
auto Src = (uptr)pSrc, Dst = (uptr)pDst;
1353+
UR_CALL(pfnUSMMemcpy(hQueue, blocking, pDst, pSrc, size,
1354+
numEventsInWaitList, phEventWaitList, phEvent));
1355+
1356+
const auto Src = (uptr)pSrc, Dst = (uptr)pDst;
13531357
auto SrcInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Src);
13541358
auto DstInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Dst);
13551359

@@ -1359,21 +1363,23 @@ ur_result_t UR_APICALL urEnqueueUSMMemcpy(
13591363

13601364
const auto &DeviceInfo =
13611365
getMsanInterceptor()->getDeviceInfo(SrcInfo->Device);
1362-
UR_CALL(DeviceInfo->Shadow->EnqueueCopyShadow(
1363-
hQueue, blocking, Dst, Src, size, numEventsInWaitList,
1364-
phEventWaitList, phEvent));
1366+
const auto SrcShadow = DeviceInfo->Shadow->MemToShadow(Src);
1367+
const auto DstShadow = DeviceInfo->Shadow->MemToShadow(Dst);
1368+
1369+
UR_CALL(pfnUSMMemcpy(hQueue, blocking, (void *)DstShadow,
1370+
(void *)SrcShadow, size, 1, phEvent, phEvent));
13651371
} else if (DstInfoItOp) {
13661372
auto DstInfo = (*DstInfoItOp)->second;
13671373

13681374
const auto &DeviceInfo =
13691375
getMsanInterceptor()->getDeviceInfo(DstInfo->Device);
1370-
UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(
1371-
hQueue, Dst, size, 0, numEventsInWaitList, phEventWaitList,
1372-
phEvent));
1376+
auto DstShadow = DeviceInfo->Shadow->MemToShadow(Dst);
1377+
1378+
UR_CALL(EnqueueUSMBlockingSet(hQueue, (void *)DstShadow, 0, size, 1,
1379+
phEvent, phEvent));
13731380
}
13741381

1375-
return pfnUSMMemcpy(hQueue, blocking, pDst, pSrc, size, numEventsInWaitList,
1376-
phEventWaitList, phEvent);
1382+
return UR_RESULT_SUCCESS;
13771383
}
13781384

13791385
///////////////////////////////////////////////////////////////////////////////
@@ -1406,20 +1412,25 @@ ur_result_t UR_APICALL urEnqueueUSMFill2D(
14061412
auto pfnUSMFill2D = getContext()->urDdiTable.Enqueue.pfnUSMFill2D;
14071413
getContext()->logger.debug("==== urEnqueueUSMFill2D");
14081414

1409-
auto Mem = (uptr)pMem;
1415+
UR_CALL(pfnUSMFill2D(hQueue, pMem, pitch, patternSize, pPattern, width,
1416+
height, numEventsInWaitList, phEventWaitList,
1417+
phEvent));
1418+
1419+
const auto Mem = (uptr)pMem;
14101420
auto MemInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Mem);
14111421
if (MemInfoItOp) {
14121422
auto MemInfo = (*MemInfoItOp)->second;
14131423

14141424
const auto &DeviceInfo =
14151425
getMsanInterceptor()->getDeviceInfo(MemInfo->Device);
1416-
UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(
1417-
hQueue, Mem, width * height, 0, numEventsInWaitList,
1418-
phEventWaitList, phEvent));
1426+
const auto MemShadow = DeviceInfo->Shadow->MemToShadow(Mem);
1427+
1428+
const char Pattern = 0;
1429+
UR_CALL(pfnUSMFill2D(hQueue, (void *)MemShadow, pitch, 1, &Pattern,
1430+
width, height, 1, phEvent, phEvent));
14191431
}
14201432

1421-
return pfnUSMFill2D(hQueue, pMem, pitch, patternSize, pPattern, width,
1422-
height, numEventsInWaitList, phEventWaitList, phEvent);
1433+
return UR_RESULT_SUCCESS;
14231434
}
14241435

14251436
///////////////////////////////////////////////////////////////////////////////
@@ -1448,8 +1459,43 @@ ur_result_t UR_APICALL urEnqueueUSMMemcpy2D(
14481459
///< kernel execution instance. If phEventWaitList and phEvent are not
14491460
///< NULL, phEvent must not refer to an element of the phEventWaitList array.
14501461
) {
1451-
ur_result_t result = UR_RESULT_SUCCESS;
1452-
return result;
1462+
auto pfnUSMMemcpy2D = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy2D;
1463+
getContext()->logger.debug("==== pfnUSMMemcpy2D");
1464+
1465+
UR_CALL(pfnUSMMemcpy2D(hQueue, blocking, pDst, dstPitch, pSrc, srcPitch,
1466+
width, height, numEventsInWaitList, phEventWaitList,
1467+
phEvent));
1468+
1469+
auto Src = (uptr)pSrc, Dst = (uptr)pDst;
1470+
auto SrcInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Src);
1471+
auto DstInfoItOp = getMsanInterceptor()->findAllocInfoByAddress(Dst);
1472+
1473+
if (SrcInfoItOp && DstInfoItOp) {
1474+
auto SrcInfo = (*SrcInfoItOp)->second;
1475+
auto DstInfo = (*DstInfoItOp)->second;
1476+
1477+
const auto &DeviceInfo =
1478+
getMsanInterceptor()->getDeviceInfo(SrcInfo->Device);
1479+
auto SrcShadow = DeviceInfo->Shadow->MemToShadow(Src);
1480+
auto DstShadow = DeviceInfo->Shadow->MemToShadow(Dst);
1481+
1482+
UR_CALL(pfnUSMMemcpy2D(hQueue, blocking, (void *)DstShadow, dstPitch,
1483+
(void *)SrcShadow, srcPitch, width, height,
1484+
numEventsInWaitList, phEventWaitList, phEvent));
1485+
} else if (DstInfoItOp) {
1486+
auto DstInfo = (*DstInfoItOp)->second;
1487+
1488+
const auto &DeviceInfo =
1489+
getMsanInterceptor()->getDeviceInfo(DstInfo->Device);
1490+
auto DstShadow = DeviceInfo->Shadow->MemToShadow(Dst);
1491+
1492+
const char Pattern = 0;
1493+
UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMFill2D(
1494+
hQueue, (void *)DstShadow, dstPitch, 1, &Pattern, width, height, 1,
1495+
phEvent, phEvent));
1496+
}
1497+
1498+
return UR_RESULT_SUCCESS;
14531499
}
14541500

14551501
///////////////////////////////////////////////////////////////////////////////

‎source/loader/layers/sanitizer/msan/msan_interceptor.cpp‎

Lines changed: 5 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -49,8 +49,7 @@ ur_result_t MsanInterceptor::allocateMemory(ur_context_handle_t Context,
4949
size_t Size, void **ResultPtr) {
5050

5151
auto ContextInfo = getContextInfo(Context);
52-
std::shared_ptr<DeviceInfo> DeviceInfo =
53-
Device ? getDeviceInfo(Device) : nullptr;
52+
std::shared_ptr<DeviceInfo> DeviceInfo = getDeviceInfo(Device);
5453

5554
void *Allocated = nullptr;
5655

@@ -70,15 +69,16 @@ ur_result_t MsanInterceptor::allocateMemory(ur_context_handle_t Context,
7069

7170
AI->print();
7271

73-
// For updating shadow memory
74-
ContextInfo->insertAllocInfo({Device}, AI);
75-
7672
// For memory release
7773
{
7874
std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
7975
m_AllocationMap.emplace(AI->AllocBegin, std::move(AI));
8076
}
8177

78+
ManagedQueue Queue(Context, Device);
79+
DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin,
80+
AI->AllocSize, 0xff);
81+
8282
return UR_RESULT_SUCCESS;
8383
}
8484

@@ -111,8 +111,6 @@ ur_result_t MsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
111111

112112
UR_CALL(prepareLaunch(DeviceInfo, InternalQueue, Kernel, LaunchInfo));
113113

114-
UR_CALL(updateShadowMemory(ContextInfo, DeviceInfo, InternalQueue));
115-
116114
return UR_RESULT_SUCCESS;
117115
}
118116

@@ -137,29 +135,6 @@ ur_result_t MsanInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel,
137135
return Result;
138136
}
139137

140-
ur_result_t
141-
MsanInterceptor::enqueueAllocInfo(std::shared_ptr<DeviceInfo> &DeviceInfo,
142-
ur_queue_handle_t Queue,
143-
std::shared_ptr<MsanAllocInfo> &AI) {
144-
return DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin,
145-
AI->AllocSize, 0xff);
146-
}
147-
148-
ur_result_t
149-
MsanInterceptor::updateShadowMemory(std::shared_ptr<ContextInfo> &ContextInfo,
150-
std::shared_ptr<DeviceInfo> &DeviceInfo,
151-
ur_queue_handle_t Queue) {
152-
auto &AllocInfos = ContextInfo->AllocInfosMap[DeviceInfo->Handle];
153-
std::scoped_lock<ur_shared_mutex> Guard(AllocInfos.Mutex);
154-
155-
for (auto &AI : AllocInfos.List) {
156-
UR_CALL(enqueueAllocInfo(DeviceInfo, Queue, AI));
157-
}
158-
AllocInfos.List.clear();
159-
160-
return UR_RESULT_SUCCESS;
161-
}
162-
163138
ur_result_t MsanInterceptor::registerProgram(ur_program_handle_t Program) {
164139
ur_result_t Result = UR_RESULT_SUCCESS;
165140

‎source/loader/layers/sanitizer/msan/msan_interceptor.hpp‎

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,6 @@ struct ContextInfo {
120120
std::atomic<int32_t> RefCount = 1;
121121

122122
std::vector<ur_device_handle_t> DeviceList;
123-
std::unordered_map<ur_device_handle_t, AllocInfoList> AllocInfosMap;
124123

125124
explicit ContextInfo(ur_context_handle_t Context) : Handle(Context) {
126125
[[maybe_unused]] auto Result =
@@ -129,15 +128,6 @@ struct ContextInfo {
129128
}
130129

131130
~ContextInfo();
132-
133-
void insertAllocInfo(const std::vector<ur_device_handle_t> &Devices,
134-
std::shared_ptr<MsanAllocInfo> &AI) {
135-
for (auto Device : Devices) {
136-
auto &AllocInfos = AllocInfosMap[Device];
137-
std::scoped_lock<ur_shared_mutex> Guard(AllocInfos.Mutex);
138-
AllocInfos.List.emplace_back(AI);
139-
}
140-
}
141131
};
142132

143133
struct USMLaunchInfo {
@@ -264,15 +254,6 @@ class MsanInterceptor {
264254
bool isNormalExit() { return m_NormalExit; }
265255

266256
private:
267-
ur_result_t
268-
updateShadowMemory(std::shared_ptr<msan::ContextInfo> &ContextInfo,
269-
std::shared_ptr<msan::DeviceInfo> &DeviceInfo,
270-
ur_queue_handle_t Queue);
271-
272-
ur_result_t enqueueAllocInfo(std::shared_ptr<msan::DeviceInfo> &DeviceInfo,
273-
ur_queue_handle_t Queue,
274-
std::shared_ptr<MsanAllocInfo> &AI);
275-
276257
/// Initialize Global Variables & Kernel Name at first Launch
277258
ur_result_t prepareLaunch(std::shared_ptr<msan::DeviceInfo> &DeviceInfo,
278259
ur_queue_handle_t Queue,

‎source/loader/layers/sanitizer/msan/msan_shadow.cpp‎

Lines changed: 2 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -133,41 +133,6 @@ ur_result_t MsanShadowMemoryCPU::EnqueuePoisonShadow(
133133
return UR_RESULT_SUCCESS;
134134
}
135135

136-
ur_result_t MsanShadowMemoryCPU::EnqueueCopyShadow(
137-
ur_queue_handle_t Queue, bool Blocking, uptr Dst, uptr Src, uptr Size,
138-
uint32_t NumEvents, const ur_event_handle_t *EventWaitList,
139-
ur_event_handle_t *OutEvent) {
140-
if (Size == 0) {
141-
if (OutEvent) {
142-
UR_CALL(getContext()->urDdiTable.Enqueue.pfnEventsWait(
143-
Queue, NumEvents, EventWaitList, OutEvent));
144-
}
145-
return UR_RESULT_SUCCESS;
146-
}
147-
148-
const uptr SrcShadowBegin = MemToShadow(Src);
149-
const uptr SrcShadowEnd = MemToShadow(Src + Size - 1);
150-
assert(SrcShadowBegin <= SrcShadowEnd);
151-
152-
const uptr DstShadowBegin = MemToShadow(Dst);
153-
const uptr DstShadowEnd = MemToShadow(Dst + Size - 1);
154-
assert(DstShadowBegin <= DstShadowEnd);
155-
156-
assert(SrcShadowEnd - SrcShadowBegin == DstShadowEnd - DstShadowBegin);
157-
158-
// FIXME: host asan will not support to use this function
159-
auto Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
160-
Queue, Blocking, (void *)DstShadowBegin, (void *)SrcShadowBegin,
161-
SrcShadowEnd - SrcShadowBegin + 1, NumEvents, EventWaitList, OutEvent);
162-
163-
getContext()->logger.debug("EnqueueCopyShadow(dst={}, src={}, size={}): {}",
164-
(void *)DstShadowBegin, (void *)SrcShadowBegin,
165-
(void *)(SrcShadowEnd - SrcShadowBegin + 1),
166-
Result);
167-
168-
return Result;
169-
}
170-
171136
ur_result_t MsanShadowMemoryGPU::Setup() {
172137
// Currently, Level-Zero doesn't create independent VAs for each contexts, if we reserve
173138
// shadow memory for each contexts, this will cause out-of-resource error when user uses
@@ -208,7 +173,7 @@ ur_result_t MsanShadowMemoryGPU::Destory() {
208173
return Result;
209174
}
210175

211-
ur_result_t MsanShadowMemoryGPU::EnqueueMappingShadow(
176+
ur_result_t MsanShadowMemoryGPU::EnqueueMapShadow(
212177
ur_queue_handle_t Queue, uptr Ptr, uptr Size,
213178
std::vector<ur_event_handle_t> &EventWaitList,
214179
ur_event_handle_t *OutEvent) {
@@ -290,7 +255,7 @@ ur_result_t MsanShadowMemoryGPU::EnqueuePoisonShadow(
290255

291256
std::vector<ur_event_handle_t> Events(EventWaitList,
292257
EventWaitList + NumEvents);
293-
UR_CALL(EnqueueMappingShadow(Queue, Ptr, Size, Events, OutEvent));
258+
UR_CALL(EnqueueMapShadow(Queue, Ptr, Size, Events, OutEvent));
294259

295260
const uptr ShadowBegin = MemToShadow(Ptr);
296261
const uptr ShadowEnd = MemToShadow(Ptr + Size - 1);
@@ -308,46 +273,6 @@ ur_result_t MsanShadowMemoryGPU::EnqueuePoisonShadow(
308273
return Result;
309274
}
310275

311-
ur_result_t MsanShadowMemoryGPU::EnqueueCopyShadow(
312-
ur_queue_handle_t Queue, bool Blocking, uptr Dst, uptr Src, uptr Size,
313-
uint32_t NumEvents, const ur_event_handle_t *EventWaitList,
314-
ur_event_handle_t *OutEvent) {
315-
if (Size == 0) {
316-
if (OutEvent) {
317-
UR_CALL(getContext()->urDdiTable.Enqueue.pfnEventsWait(
318-
Queue, NumEvents, EventWaitList, OutEvent));
319-
}
320-
return UR_RESULT_SUCCESS;
321-
}
322-
323-
std::vector<ur_event_handle_t> Events(EventWaitList,
324-
EventWaitList + NumEvents);
325-
UR_CALL(EnqueueMappingShadow(Queue, Src, Size, Events, OutEvent));
326-
UR_CALL(EnqueueMappingShadow(Queue, Dst, Size, Events, OutEvent));
327-
328-
const uptr SrcShadowBegin = MemToShadow(Src);
329-
const uptr SrcShadowEnd = MemToShadow(Src + Size - 1);
330-
assert(SrcShadowBegin <= SrcShadowEnd);
331-
332-
const uptr DstShadowBegin = MemToShadow(Dst);
333-
const uptr DstShadowEnd = MemToShadow(Dst + Size - 1);
334-
assert(DstShadowBegin <= DstShadowEnd);
335-
336-
assert(DstShadowEnd - DstShadowBegin == SrcShadowEnd - SrcShadowBegin);
337-
338-
auto Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
339-
Queue, Blocking, (void *)DstShadowBegin, (void *)SrcShadowBegin,
340-
SrcShadowEnd - SrcShadowBegin + 1, Events.size(), Events.data(),
341-
OutEvent);
342-
343-
getContext()->logger.debug("EnqueueCopyShadow(dst={}, src={}, size={}): {}",
344-
(void *)DstShadowBegin, (void *)SrcShadowBegin,
345-
(void *)(SrcShadowEnd - SrcShadowBegin + 1),
346-
Result);
347-
348-
return Result;
349-
}
350-
351276
ur_result_t
352277
MsanShadowMemoryGPU::ReleaseShadow(std::shared_ptr<MsanAllocInfo> AI) {
353278
uptr ShadowBegin = MemToShadow(AI->AllocBegin);

0 commit comments

Comments
 (0)