Skip to content

Commit 2e212e0

Browse files
author
Hugh Delaney
authored
[CUDA][HIP] Fix host task mem migration and add pi entry point for urEnqueueNativeCommandExp (#14353)
The SYCL RT assumes that for devices in the same context, no mem migration needs to occur across devices for a kernel launch or host task. However, a CUdeviceptr is relevant to a specific device, so mem migration must occur between devices in a ctx. If this assumption that the SYCL RT makes about native mems being accessible to all devices in a context, it must hand off the HT lambda to the plugin, so that the plugin can handle the necessary mem migration. This patch uses the new urEnqueueCustomCommandExp to execute the HT lambda, which takes care of mem migration implicitly in the plugin.
1 parent 43286ab commit 2e212e0

File tree

18 files changed

+179
-17
lines changed

18 files changed

+179
-17
lines changed

‎sycl/include/sycl/detail/host_task_impl.hpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ class HostTask {
4747
if (HPI)
4848
HPI->end();
4949
}
50+
51+
friend class DispatchHostTask;
5052
};
5153

5254
class CGHostTask : public CG {

‎sycl/include/sycl/detail/pi.def‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,4 +227,7 @@ _PI_API(piextVirtualMemUnmap)
227227
_PI_API(piextVirtualMemSetAccess)
228228
_PI_API(piextVirtualMemGetInfo)
229229

230+
// Enqueue native command
231+
_PI_API(piextEnqueueNativeCommand)
232+
230233
#undef _PI_API

‎sycl/include/sycl/detail/pi.h‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -195,9 +195,10 @@
195195
// PI_EXT_ONEAPI_DEVICE_INFO_SUPPORTS_VIRTUAL_MEM device info descriptor,
196196
// _pi_virtual_mem_granularity_info enum, _pi_virtual_mem_info enum and
197197
// pi_virtual_access_flags bit flags.
198+
// 15.55 Added piextEnqueueNativeCommand as well as associated types and enums
198199

199200
#define _PI_H_VERSION_MAJOR 15
200-
#define _PI_H_VERSION_MINOR 54
201+
#define _PI_H_VERSION_MINOR 55
201202

202203
#define _PI_STRING_HELPER(a) #a
203204
#define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -512,6 +513,8 @@ typedef enum {
512513

513514
// Virtual memory support
514515
PI_EXT_ONEAPI_DEVICE_INFO_SUPPORTS_VIRTUAL_MEM = 0x2011E,
516+
// Native enqueue
517+
PI_EXT_ONEAPI_DEVICE_INFO_ENQUEUE_NATIVE_COMMAND_SUPPORT = 0x2011F,
515518
} _pi_device_info;
516519

517520
typedef enum {
@@ -1279,6 +1282,7 @@ using pi_image_mem_handle = void *;
12791282
using pi_interop_mem_handle = pi_uint64;
12801283
using pi_interop_semaphore_handle = pi_uint64;
12811284
using pi_physical_mem = _pi_physical_mem *;
1285+
using pi_enqueue_native_command_function = void (*)(pi_queue, void *);
12821286

12831287
typedef struct {
12841288
pi_image_channel_order image_channel_order;
@@ -3201,6 +3205,25 @@ __SYCL_EXPORT pi_result piextSignalExternalSemaphore(
32013205
pi_uint32 num_events_in_wait_list, const pi_event *event_wait_list,
32023206
pi_event *event);
32033207

3208+
/// API to enqueue work through a backend API such that the plugin can schedule
3209+
/// the backend API calls within its own DAG.
3210+
///
3211+
/// \param command_queue is the queue instructed to signal
3212+
/// \param fn is the user submitted native function enqueueing work to a
3213+
/// backend API
3214+
/// \param data is the data that will be used in fn
3215+
/// \param num_mems is the number of mems in mem_list
3216+
/// \param mem_list is the list of mems that are used in fn
3217+
/// \param num_events_in_wait_list is the number of events in the wait list
3218+
/// \param event_wait_list is the list of events to wait on before this
3219+
/// operation
3220+
/// \param event is the returned event representing this operation
3221+
__SYCL_EXPORT pi_result piextEnqueueNativeCommand(
3222+
pi_queue command_queue, pi_enqueue_native_command_function fn, void *data,
3223+
pi_uint32 num_mems, const pi_mem *mem_list,
3224+
pi_uint32 num_events_in_wait_list, const pi_event *event_wait_list,
3225+
pi_event *event);
3226+
32043227
typedef enum {
32053228
_PI_SANITIZE_TYPE_NONE = 0x0,
32063229
_PI_SANITIZE_TYPE_ADDRESS = 0x1,

‎sycl/plugins/cuda/pi_cuda.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1361,6 +1361,16 @@ pi_result piextVirtualMemGetInfo(pi_context context, const void *ptr,
13611361
param_value_size_ret);
13621362
}
13631363

1364+
pi_result
1365+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1366+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1367+
pi_uint32 NumEventsInWaitList,
1368+
const pi_event *EventWaitList, pi_event *Event) {
1369+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1370+
NumEventsInWaitList, EventWaitList,
1371+
Event);
1372+
}
1373+
13641374
const char SupportedVersion[] = _PI_CUDA_PLUGIN_VERSION_STRING;
13651375

13661376
pi_result piPluginInit(pi_plugin *PluginInit) {

‎sycl/plugins/hip/pi_hip.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1364,6 +1364,16 @@ pi_result piextVirtualMemGetInfo(pi_context context, const void *ptr,
13641364
param_value_size_ret);
13651365
}
13661366

1367+
pi_result
1368+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1369+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1370+
pi_uint32 NumEventsInWaitList,
1371+
const pi_event *EventWaitList, pi_event *Event) {
1372+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1373+
NumEventsInWaitList, EventWaitList,
1374+
Event);
1375+
}
1376+
13671377
const char SupportedVersion[] = _PI_HIP_PLUGIN_VERSION_STRING;
13681378

13691379
pi_result piPluginInit(pi_plugin *PluginInit) {

‎sycl/plugins/level_zero/pi_level_zero.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1562,6 +1562,16 @@ pi_result piextVirtualMemGetInfo(pi_context Context, const void *Ptr,
15621562
ParamValueSizeRet);
15631563
}
15641564

1565+
pi_result
1566+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1567+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1568+
pi_uint32 NumEventsInWaitList,
1569+
const pi_event *EventWaitList, pi_event *Event) {
1570+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1571+
NumEventsInWaitList, EventWaitList,
1572+
Event);
1573+
}
1574+
15651575
const char SupportedVersion[] = _PI_LEVEL_ZERO_PLUGIN_VERSION_STRING;
15661576

15671577
pi_result piPluginInit(pi_plugin *PluginInit) { // missing

‎sycl/plugins/native_cpu/pi_native_cpu.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1384,6 +1384,16 @@ pi_result piextVirtualMemGetInfo(pi_context context, const void *ptr,
13841384
param_value_size_ret);
13851385
}
13861386

1387+
pi_result
1388+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1389+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1390+
pi_uint32 NumEventsInWaitList,
1391+
const pi_event *EventWaitList, pi_event *Event) {
1392+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1393+
NumEventsInWaitList, EventWaitList,
1394+
Event);
1395+
}
1396+
13871397
// Initialize function table with stubs.
13881398
#define _PI_API(api) \
13891399
(PluginInit->PiFunctionTable).api = (decltype(&::api))(&api);

‎sycl/plugins/opencl/pi_opencl.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1291,6 +1291,16 @@ pi_result piextVirtualMemGetInfo(pi_context Context, const void *Ptr,
12911291
ParamValueSizeRet);
12921292
}
12931293

1294+
pi_result
1295+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1296+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1297+
pi_uint32 NumEventsInWaitList,
1298+
const pi_event *EventWaitList, pi_event *Event) {
1299+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1300+
NumEventsInWaitList, EventWaitList,
1301+
Event);
1302+
}
1303+
12941304
pi_result piTearDown(void *PluginParameter) {
12951305
return pi2ur::piTearDown(PluginParameter);
12961306
}

‎sycl/plugins/unified_runtime/pi2ur.hpp‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1325,6 +1325,9 @@ inline pi_result piDeviceGetInfo(pi_device Device, pi_device_info ParamName,
13251325
PI_TO_UR_MAP_DEVICE_INFO(
13261326
PI_EXT_ONEAPI_DEVICE_INFO_TIMESTAMP_RECORDING_SUPPORT,
13271327
UR_DEVICE_INFO_TIMESTAMP_RECORDING_SUPPORT_EXP)
1328+
PI_TO_UR_MAP_DEVICE_INFO(
1329+
PI_EXT_ONEAPI_DEVICE_INFO_ENQUEUE_NATIVE_COMMAND_SUPPORT,
1330+
UR_DEVICE_INFO_ENQUEUE_NATIVE_COMMAND_SUPPORT_EXP)
13281331
PI_TO_UR_MAP_DEVICE_INFO(PI_EXT_INTEL_DEVICE_INFO_ESIMD_SUPPORT,
13291332
UR_DEVICE_INFO_ESIMD_SUPPORT)
13301333
PI_TO_UR_MAP_DEVICE_INFO(PI_EXT_ONEAPI_DEVICE_INFO_COMPONENT_DEVICES,
@@ -5722,7 +5725,6 @@ piextVirtualMemGranularityGetInfo(pi_context Context, pi_device Device,
57225725
HANDLE_ERRORS(urVirtualMemGranularityGetInfo(UrContext, UrDevice, InfoType,
57235726
ParamValueSize, ParamValue,
57245727
ParamValueSizeRet));
5725-
57265728
return PI_SUCCESS;
57275729
}
57285730

@@ -5882,4 +5884,30 @@ inline pi_result piextVirtualMemGetInfo(pi_context Context, const void *Ptr,
58825884
// Virtual Memory
58835885
///////////////////////////////////////////////////////////////////////////////
58845886

5887+
///////////////////////////////////////////////////////////////////////////////
5888+
// Enqueue Native Command Extension
5889+
inline pi_result
5890+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
5891+
void *Data, pi_uint32 NumMems, const pi_mem *MemList,
5892+
pi_uint32 NumEventsInWaitList,
5893+
const pi_event *EventWaitList, pi_event *Event) {
5894+
PI_ASSERT(Queue, PI_ERROR_INVALID_QUEUE);
5895+
5896+
auto UrQueue = reinterpret_cast<ur_queue_handle_t>(Queue);
5897+
auto UrFn = reinterpret_cast<void (*)(ur_queue_handle_t, void *)>(Fn);
5898+
const ur_mem_handle_t *UrMemList =
5899+
reinterpret_cast<const ur_mem_handle_t *>(MemList);
5900+
const ur_event_handle_t *UrEventWaitList =
5901+
reinterpret_cast<const ur_event_handle_t *>(EventWaitList);
5902+
ur_event_handle_t *UREvent = reinterpret_cast<ur_event_handle_t *>(Event);
5903+
5904+
HANDLE_ERRORS(urEnqueueNativeCommandExp(
5905+
UrQueue, UrFn, Data, NumMems, UrMemList, nullptr /*pProperties*/,
5906+
NumEventsInWaitList, UrEventWaitList, UREvent));
5907+
5908+
return PI_SUCCESS;
5909+
}
5910+
// Enqueue Native Command Extension
5911+
///////////////////////////////////////////////////////////////////////////////
5912+
58855913
} // namespace pi2ur

‎sycl/plugins/unified_runtime/pi_unified_runtime.cpp‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1447,6 +1447,16 @@ __SYCL_EXPORT pi_result piextSignalExternalSemaphore(
14471447
EventWaitList, Event);
14481448
}
14491449

1450+
pi_result
1451+
piextEnqueueNativeCommand(pi_queue Queue, pi_enqueue_native_command_function Fn,
1452+
void *Data, pi_uint32 NumMems, const pi_mem *Mems,
1453+
pi_uint32 NumEventsInWaitList,
1454+
const pi_event *EventWaitList, pi_event *Event) {
1455+
return pi2ur::piextEnqueueNativeCommand(Queue, Fn, Data, NumMems, Mems,
1456+
NumEventsInWaitList, EventWaitList,
1457+
Event);
1458+
}
1459+
14501460
// This interface is not in Unified Runtime currently
14511461
__SYCL_EXPORT pi_result piPluginInit(pi_plugin *PluginInit) {
14521462
PI_ASSERT(PluginInit, PI_ERROR_INVALID_VALUE);

0 commit comments

Comments
 (0)