Skip to content

Commit e161516

Browse files
kbenziekbenzie
authored
Merge pull request #1521 from AllanZyne/review/yang/dg2
[DeviceSanitizer] Support GPU DG2 & GEN Device
2 parents 4763308 + febb18b commit e161516

8 files changed

+173
-30
lines changed

‎source/loader/layers/sanitizer/asan_interceptor.cpp

+43-18
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,15 @@ uptr MemToShadow_CPU(uptr USM_SHADOW_BASE, uptr UPtr) {
2727
return USM_SHADOW_BASE + (UPtr >> ASAN_SHADOW_SCALE);
2828
}
2929

30+
uptr MemToShadow_DG2(uptr USM_SHADOW_BASE, uptr UPtr) {
31+
if (UPtr & 0xFFFF000000000000ULL) { // Device USM
32+
return USM_SHADOW_BASE + 0x80000000000ULL +
33+
((UPtr & 0x7FFFFFFFFFFFULL) >> ASAN_SHADOW_SCALE);
34+
} else { // Host/Shared USM
35+
return USM_SHADOW_BASE + (UPtr >> ASAN_SHADOW_SCALE);
36+
}
37+
}
38+
3039
uptr MemToShadow_PVC(uptr USM_SHADOW_BASE, uptr UPtr) {
3140
if (UPtr & 0xFF00000000000000ULL) { // Device USM
3241
return USM_SHADOW_BASE + 0x80000000000ULL +
@@ -56,6 +65,9 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
5665
return UR_RESULT_SUCCESS;
5766
}
5867
if (DeviceInfo->Type == DeviceType::CPU) {
68+
///
69+
/// CPU Device: CPU needs to use a special memset function
70+
///
5971
uptr ShadowBegin = MemToShadow_CPU(DeviceInfo->ShadowOffset, Ptr);
6072
uptr ShadowEnd =
6173
MemToShadow_CPU(DeviceInfo->ShadowOffset, Ptr + Size - 1);
@@ -72,10 +84,25 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
7284
(void *)ShadowBegin, ShadowEnd - ShadowBegin + 1,
7385
(void *)(size_t)Value);
7486
MemSet((void *)ShadowBegin, Value, ShadowEnd - ShadowBegin + 1);
75-
} else if (DeviceInfo->Type == DeviceType::GPU_PVC) {
76-
uptr ShadowBegin = MemToShadow_PVC(DeviceInfo->ShadowOffset, Ptr);
77-
uptr ShadowEnd =
78-
MemToShadow_PVC(DeviceInfo->ShadowOffset, Ptr + Size - 1);
87+
} else {
88+
///
89+
/// GPU Device: GPU needs to manually map physical memory before memset
90+
///
91+
uptr ShadowBegin = 0, ShadowEnd = 0;
92+
93+
if (DeviceInfo->Type == DeviceType::GPU_PVC) {
94+
ShadowBegin = MemToShadow_PVC(DeviceInfo->ShadowOffset, Ptr);
95+
ShadowEnd =
96+
MemToShadow_PVC(DeviceInfo->ShadowOffset, Ptr + Size - 1);
97+
} else if (DeviceInfo->Type == DeviceType::GPU_DG2) {
98+
ShadowBegin = MemToShadow_DG2(DeviceInfo->ShadowOffset, Ptr);
99+
ShadowEnd =
100+
MemToShadow_DG2(DeviceInfo->ShadowOffset, Ptr + Size - 1);
101+
} else {
102+
getContext()->logger.error("Unsupport device type");
103+
return UR_RESULT_ERROR_INVALID_ARGUMENT;
104+
}
105+
79106
assert(ShadowBegin <= ShadowEnd);
80107
{
81108
static const size_t PageSize =
@@ -108,7 +135,9 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
108135
Context, (void *)MappedPtr, PageSize, PhysicalMem, 0,
109136
UR_VIRTUAL_MEM_ACCESS_FLAG_READ_WRITE);
110137
if (URes != UR_RESULT_SUCCESS) {
111-
getContext()->logger.debug("urVirtualMemMap(): {}", URes);
138+
getContext()->logger.debug("urVirtualMemMap({}, {}): {}",
139+
(void *)MappedPtr, PageSize,
140+
URes);
112141
}
113142

114143
// Initialize to zero
@@ -137,9 +166,6 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
137166
getContext()->logger.error("urEnqueueUSMFill(): {}", URes);
138167
return URes;
139168
}
140-
} else {
141-
getContext()->logger.error("Unsupport device type");
142-
return UR_RESULT_ERROR_INVALID_ARGUMENT;
143169
}
144170
return UR_RESULT_SUCCESS;
145171
}
@@ -158,6 +184,7 @@ SanitizerInterceptor::SanitizerInterceptor(logger::Logger &logger)
158184
SanitizerInterceptor::~SanitizerInterceptor() {
159185
DestroyShadowMemoryOnCPU();
160186
DestroyShadowMemoryOnPVC();
187+
DestroyShadowMemoryOnDG2();
161188
}
162189

163190
/// The memory chunk allocated from the underlying allocator looks like this:
@@ -390,6 +417,8 @@ ur_result_t DeviceInfo::allocShadowMemory(ur_context_handle_t Context) {
390417
UR_CALL(SetupShadowMemoryOnCPU(ShadowOffset, ShadowOffsetEnd));
391418
} else if (Type == DeviceType::GPU_PVC) {
392419
UR_CALL(SetupShadowMemoryOnPVC(Context, ShadowOffset, ShadowOffsetEnd));
420+
} else if (Type == DeviceType::GPU_DG2) {
421+
UR_CALL(SetupShadowMemoryOnDG2(Context, ShadowOffset, ShadowOffsetEnd));
393422
} else {
394423
getContext()->logger.error("Unsupport device type");
395424
return UR_RESULT_ERROR_INVALID_ARGUMENT;
@@ -586,12 +615,6 @@ SanitizerInterceptor::insertDevice(ur_device_handle_t Device,
586615

587616
DI = std::make_shared<ur_sanitizer_layer::DeviceInfo>(Device);
588617

589-
// Query device type
590-
DI->Type = GetDeviceType(Device);
591-
if (DI->Type == DeviceType::UNKNOWN) {
592-
return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
593-
}
594-
595618
// Query alignment
596619
UR_CALL(getContext()->urDdiTable.Device.pfnGetInfo(
597620
Device, UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN, sizeof(DI->Alignment),
@@ -786,7 +809,8 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
786809
// Write shadow memory offset for local memory
787810
if (Options(logger).DetectLocals) {
788811
// CPU needn't this
789-
if (DeviceInfo->Type == DeviceType::GPU_PVC) {
812+
if (DeviceInfo->Type == DeviceType::GPU_PVC ||
813+
DeviceInfo->Type == DeviceType::GPU_DG2) {
790814
const size_t LocalMemorySize =
791815
GetDeviceLocalMemorySize(DeviceInfo->Handle);
792816
const size_t LocalShadowMemorySize =
@@ -826,7 +850,8 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
826850
if (Options(logger).DetectPrivates) {
827851
if (DeviceInfo->Type == DeviceType::CPU) {
828852
LaunchInfo.Data->PrivateShadowOffset = DeviceInfo->ShadowOffset;
829-
} else if (DeviceInfo->Type == DeviceType::GPU_PVC) {
853+
} else if (DeviceInfo->Type == DeviceType::GPU_PVC ||
854+
DeviceInfo->Type == DeviceType::GPU_DG2) {
830855
const size_t PrivateShadowMemorySize =
831856
(NumWG * ASAN_PRIVATE_SIZE) >> ASAN_SHADOW_SCALE;
832857

@@ -907,8 +932,8 @@ ur_result_t USMLaunchInfo::updateKernelInfo(const KernelInfo &KI) {
907932
USMLaunchInfo::~USMLaunchInfo() {
908933
[[maybe_unused]] ur_result_t Result;
909934
if (Data) {
910-
auto Type = GetDeviceType(Device);
911-
if (Type == DeviceType::GPU_PVC) {
935+
auto Type = GetDeviceType(Context, Device);
936+
if (Type == DeviceType::GPU_PVC || Type == DeviceType::GPU_DG2) {
912937
if (Data->PrivateShadowOffset) {
913938
Result = getContext()->urDdiTable.USM.pfnFree(
914939
Context, (void *)Data->PrivateShadowOffset);

‎source/loader/layers/sanitizer/asan_shadow_setup.cpp

+74-1
Original file line numberDiff line numberDiff line change
@@ -51,9 +51,21 @@ ur_result_t DestroyShadowMemory() {
5151
} // namespace cpu
5252

5353
namespace pvc {
54-
54+
///
55+
/// USM Allocation Range (56 bits)
56+
/// Host USM : 0x0000_0000_0000_0000 ~ 0x00ff_ffff_ffff_ffff
57+
/// Shared USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff
58+
/// Device USM : 0xff00_0000_0000_0000 ~ 0xff00_ffff_ffff_ffff
59+
///
60+
/// USM Allocation Range (AllocateHostAllocationsInHeapExtendedHost=0)
61+
/// Host USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff
62+
/// Shared USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff
63+
/// Device USM : 0xff00_0000_0000_0000 ~ 0xff00_ffff_ffff_ffff
64+
///
65+
/// Shadow Memory Mapping (SHADOW_SCALE=4, AllocateHostAllocationsInHeapExtendedHost=0)
5566
/// Host/Shared USM : 0x0 ~ 0x07ff_ffff_ffff
5667
/// Device USM : 0x0800_0000_0000 ~ 0x17ff_ffff_ffff
68+
///
5769
constexpr size_t SHADOW_SIZE = 0x180000000000ULL;
5870

5971
uptr LOW_SHADOW_BEGIN;
@@ -98,6 +110,60 @@ ur_result_t DestroyShadowMemory() {
98110

99111
} // namespace pvc
100112

113+
namespace dg2 {
114+
///
115+
/// USM Allocation Range (48 bits)
116+
/// Host/Shared USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff
117+
/// Device USM : 0xffff_8000_0000_0000 ~ 0xffff_ffff_ffff_ffff
118+
///
119+
/// Shadow Memory Mapping (SHADOW_SCALE=4)
120+
/// Host/Shared USM : 0x0 ~ 0x07ff_ffff_ffff
121+
/// Device USM : 0x0800_0000_0000 ~ 0x0fff_ffff_ffff
122+
///
123+
constexpr size_t SHADOW_SIZE = 0x100000000000ULL;
124+
125+
uptr LOW_SHADOW_BEGIN;
126+
uptr HIGH_SHADOW_END;
127+
128+
ur_context_handle_t ShadowContext;
129+
130+
ur_result_t SetupShadowMemory(ur_context_handle_t Context, uptr &ShadowBegin,
131+
uptr &ShadowEnd) {
132+
// Currently, Level-Zero doesn't create independent VAs for each contexts, if we reserve
133+
// shadow memory for each contexts, this will cause out-of-resource error when user uses
134+
// multiple contexts. Therefore, we just create one shadow memory here.
135+
static ur_result_t Result = [&Context]() {
136+
// TODO: Protect Bad Zone
137+
auto Result = getContext()->urDdiTable.VirtualMem.pfnReserve(
138+
Context, nullptr, SHADOW_SIZE, (void **)&LOW_SHADOW_BEGIN);
139+
if (Result == UR_RESULT_SUCCESS) {
140+
HIGH_SHADOW_END = LOW_SHADOW_BEGIN + SHADOW_SIZE;
141+
// Retain the context which reserves shadow memory
142+
ShadowContext = Context;
143+
getContext()->urDdiTable.Context.pfnRetain(Context);
144+
}
145+
return Result;
146+
}();
147+
ShadowBegin = LOW_SHADOW_BEGIN;
148+
ShadowEnd = HIGH_SHADOW_END;
149+
return Result;
150+
}
151+
152+
ur_result_t DestroyShadowMemory() {
153+
static ur_result_t Result = []() {
154+
if (!ShadowContext) {
155+
return UR_RESULT_SUCCESS;
156+
}
157+
auto Result = getContext()->urDdiTable.VirtualMem.pfnFree(
158+
ShadowContext, (const void *)LOW_SHADOW_BEGIN, SHADOW_SIZE);
159+
getContext()->urDdiTable.Context.pfnRelease(ShadowContext);
160+
return Result;
161+
}();
162+
return Result;
163+
}
164+
165+
} // namespace dg2
166+
101167
ur_result_t SetupShadowMemoryOnCPU(uptr &ShadowBegin, uptr &ShadowEnd) {
102168
return cpu::SetupShadowMemory(ShadowBegin, ShadowEnd);
103169
}
@@ -111,4 +177,11 @@ ur_result_t SetupShadowMemoryOnPVC(ur_context_handle_t Context,
111177

112178
ur_result_t DestroyShadowMemoryOnPVC() { return pvc::DestroyShadowMemory(); }
113179

180+
ur_result_t SetupShadowMemoryOnDG2(ur_context_handle_t Context,
181+
uptr &ShadowBegin, uptr &ShadowEnd) {
182+
return dg2::SetupShadowMemory(Context, ShadowBegin, ShadowEnd);
183+
}
184+
185+
ur_result_t DestroyShadowMemoryOnDG2() { return dg2::DestroyShadowMemory(); }
186+
114187
} // namespace ur_sanitizer_layer

‎source/loader/layers/sanitizer/asan_shadow_setup.hpp

+4
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,8 @@ ur_result_t SetupShadowMemoryOnPVC(ur_context_handle_t Context,
2323
uptr &ShadowBegin, uptr &ShadowEnd);
2424
ur_result_t DestroyShadowMemoryOnPVC();
2525

26+
ur_result_t SetupShadowMemoryOnDG2(ur_context_handle_t Context,
27+
uptr &ShadowBegin, uptr &ShadowEnd);
28+
ur_result_t DestroyShadowMemoryOnDG2();
29+
2630
} // namespace ur_sanitizer_layer

‎source/loader/layers/sanitizer/common.hpp

+15
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,21 @@ struct SourceInfo {
138138

139139
enum class DeviceType : uint64_t { UNKNOWN = 0, CPU, GPU_PVC, GPU_DG2 };
140140

141+
inline const char *ToString(DeviceType Type) {
142+
switch (Type) {
143+
case DeviceType::UNKNOWN:
144+
return "UNKNOWN";
145+
case DeviceType::CPU:
146+
return "CPU";
147+
case DeviceType::GPU_PVC:
148+
return "PVC";
149+
case DeviceType::GPU_DG2:
150+
return "DG2";
151+
default:
152+
return "UNKNOWN";
153+
}
154+
}
155+
141156
bool IsInASanContext();
142157

143158
uptr MmapNoReserve(uptr Addr, uptr Size);

‎source/loader/layers/sanitizer/ur_sanddi.cpp

+13-5
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,13 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
2626
auto hDevice = phDevices[i];
2727
std::shared_ptr<DeviceInfo> DI;
2828
UR_CALL(getContext()->interceptor->insertDevice(hDevice, DI));
29+
DI->Type = GetDeviceType(Context, hDevice);
30+
if (DI->Type == DeviceType::UNKNOWN) {
31+
getContext()->logger.error("Unsupport device");
32+
return UR_RESULT_ERROR_INVALID_DEVICE;
33+
}
34+
getContext()->logger.info("Add {} into context {}", ToString(DI->Type),
35+
(void *)Context);
2936
if (!DI->ShadowOffset) {
3037
UR_CALL(DI->allocShadowMemory(Context));
3138
}
@@ -1522,19 +1529,20 @@ ur_result_t context_t::init(ur_dditable_t *dditable,
15221529
ur_result_t result = UR_RESULT_SUCCESS;
15231530

15241531
if (enabledLayerNames.count("UR_LAYER_ASAN")) {
1525-
getContext()->enabledType = SanitizerType::AddressSanitizer;
1532+
enabledType = SanitizerType::AddressSanitizer;
1533+
interceptor = std::make_unique<SanitizerInterceptor>(logger);
15261534
} else if (enabledLayerNames.count("UR_LAYER_MSAN")) {
1527-
getContext()->enabledType = SanitizerType::MemorySanitizer;
1535+
enabledType = SanitizerType::MemorySanitizer;
15281536
} else if (enabledLayerNames.count("UR_LAYER_TSAN")) {
1529-
getContext()->enabledType = SanitizerType::ThreadSanitizer;
1537+
enabledType = SanitizerType::ThreadSanitizer;
15301538
}
15311539

15321540
// Only support AddressSanitizer now
1533-
if (getContext()->enabledType != SanitizerType::AddressSanitizer) {
1541+
if (enabledType != SanitizerType::AddressSanitizer) {
15341542
return result;
15351543
}
15361544

1537-
if (getContext()->enabledType == SanitizerType::AddressSanitizer) {
1545+
if (enabledType == SanitizerType::AddressSanitizer) {
15381546
if (!(dditable->VirtualMem.pfnReserve && dditable->VirtualMem.pfnMap &&
15391547
dditable->VirtualMem.pfnGranularityGetInfo)) {
15401548
die("Some VirtualMem APIs are needed to enable UR_LAYER_ASAN");

‎source/loader/layers/sanitizer/ur_sanitizer_layer.cpp

+1-2
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,7 @@ context_t *getContext() { return context_t::get_direct(); }
1919
///////////////////////////////////////////////////////////////////////////////
2020
context_t::context_t()
2121
: logger(logger::create_logger("sanitizer", false, false,
22-
logger::Level::WARN)),
23-
interceptor(std::make_unique<SanitizerInterceptor>(logger)) {}
22+
logger::Level::WARN)) {}
2423

2524
ur_result_t context_t::tearDown() { return UR_RESULT_SUCCESS; }
2625

‎source/loader/layers/sanitizer/ur_sanitizer_utils.cpp

+21-3
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ ur_device_handle_t GetUSMAllocDevice(ur_context_handle_t Context,
115115
return Device;
116116
}
117117

118-
DeviceType GetDeviceType(ur_device_handle_t Device) {
118+
DeviceType GetDeviceType(ur_context_handle_t Context,
119+
ur_device_handle_t Device) {
119120
ur_device_type_t DeviceType = UR_DEVICE_TYPE_DEFAULT;
120121
[[maybe_unused]] auto Result = getContext()->urDdiTable.Device.pfnGetInfo(
121122
Device, UR_DEVICE_INFO_TYPE, sizeof(DeviceType), &DeviceType, nullptr);
@@ -126,8 +127,25 @@ DeviceType GetDeviceType(ur_device_handle_t Device) {
126127
// TODO: Check fpga is fpga emulator
127128
return DeviceType::CPU;
128129
case UR_DEVICE_TYPE_GPU: {
129-
// TODO: Check device name
130-
return DeviceType::GPU_PVC;
130+
uptr Ptr;
131+
[[maybe_unused]] ur_result_t Result =
132+
getContext()->urDdiTable.USM.pfnDeviceAlloc(
133+
Context, Device, nullptr, nullptr, 4, (void **)&Ptr);
134+
getContext()->logger.debug("GetDeviceType: {}", (void *)Ptr);
135+
assert(Result == UR_RESULT_SUCCESS &&
136+
"getDeviceType() failed at allocating device USM");
137+
// FIXME: There's no API querying the address bits of device, so we guess it by the
138+
// value of device USM pointer (see "USM Allocation Range" in asan_shadow_setup.cpp)
139+
auto Type = DeviceType::UNKNOWN;
140+
if (Ptr >> 48 == 0xff00U) {
141+
Type = DeviceType::GPU_PVC;
142+
} else {
143+
Type = DeviceType::GPU_DG2;
144+
}
145+
Result = getContext()->urDdiTable.USM.pfnFree(Context, (void *)Ptr);
146+
assert(Result == UR_RESULT_SUCCESS &&
147+
"getDeviceType() failed at releasing device USM");
148+
return Type;
131149
}
132150
default:
133151
return DeviceType::UNKNOWN;

‎source/loader/layers/sanitizer/ur_sanitizer_utils.hpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,8 @@ ur_context_handle_t GetContext(ur_queue_handle_t Queue);
3434
ur_context_handle_t GetContext(ur_program_handle_t Program);
3535
ur_context_handle_t GetContext(ur_kernel_handle_t Kernel);
3636
ur_device_handle_t GetDevice(ur_queue_handle_t Queue);
37-
DeviceType GetDeviceType(ur_device_handle_t Device);
37+
DeviceType GetDeviceType(ur_context_handle_t Context,
38+
ur_device_handle_t Device);
3839
std::string GetKernelName(ur_kernel_handle_t Kernel);
3940
size_t GetDeviceLocalMemorySize(ur_device_handle_t Device);
4041
ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel);

0 commit comments

Comments
 (0)