Skip to content

Commit 9a98029

Browse files
ro-iro-i
committed
[OpenMP] Implement omp_get_uid_from_device() / omp_get_device_from_uid()
Use the implementation in libomptarget. If libomptarget is not available, always return the UID / device number of the host / the initial device.
1 parent e66f1b0 commit 9a98029

File tree

15 files changed

+317
-2
lines changed

15 files changed

+317
-2
lines changed

‎offload/include/OpenMP/omp.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,13 @@
3030

3131
extern "C" {
3232

33+
/// Definitions
34+
///{
35+
36+
#define omp_invalid_device -2
37+
38+
///}
39+
3340
/// Type declarations
3441
///{
3542

‎offload/include/omptarget.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,8 @@ extern "C" {
274274
void ompx_dump_mapping_tables(void);
275275
int omp_get_num_devices(void);
276276
int omp_get_device_num(void);
277+
int omp_get_device_from_uid(const char *DeviceUid);
278+
const char *omp_get_uid_from_device(int DeviceNum);
277279
int omp_get_initial_device(void);
278280
void *omp_target_alloc(size_t Size, int DeviceNum);
279281
void omp_target_free(void *DevicePtr, int DeviceNum);

‎offload/libomptarget/OpenMP/API.cpp‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ EXTERN void ompx_dump_mapping_tables() {
4040
using namespace llvm::omp::target::ompt;
4141
#endif
4242

43+
using GenericDeviceTy = llvm::omp::target::plugin::GenericDeviceTy;
44+
4345
void *targetAllocExplicit(size_t Size, int DeviceNum, int Kind,
4446
const char *Name);
4547
void targetFreeExplicit(void *DevicePtr, int DeviceNum, int Kind,
@@ -91,6 +93,59 @@ EXTERN int omp_get_device_num(void) {
9193
return HostDevice;
9294
}
9395

96+
EXTERN int omp_get_device_from_uid(const char *DeviceUid) {
97+
TIMESCOPE();
98+
OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
99+
100+
if (!DeviceUid) {
101+
DP("Call to omp_get_device_from_uid returning omp_invalid_device\n");
102+
return omp_invalid_device;
103+
}
104+
if (strcmp(DeviceUid, GenericDeviceTy::getHostDeviceUid()) == 0) {
105+
DP("Call to omp_get_device_from_uid returning host device number %d\n",
106+
omp_get_initial_device());
107+
return omp_get_initial_device();
108+
}
109+
110+
int DeviceNum = omp_invalid_device;
111+
112+
auto ExclusiveDevicesAccessor = PM->getExclusiveDevicesAccessor();
113+
for (const DeviceTy &Device : PM->devices(ExclusiveDevicesAccessor)) {
114+
const char *Uid = Device.RTL->getDevice(Device.RTLDeviceID).getDeviceUid();
115+
if (Uid && strcmp(DeviceUid, Uid) == 0) {
116+
DeviceNum = Device.DeviceID;
117+
break;
118+
}
119+
}
120+
121+
DP("Call to omp_get_device_from_uid returning %d\n", DeviceNum);
122+
return DeviceNum;
123+
}
124+
125+
EXTERN const char *omp_get_uid_from_device(int DeviceNum) {
126+
TIMESCOPE();
127+
OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
128+
129+
if (DeviceNum == omp_invalid_device) {
130+
DP("Call to omp_get_uid_from_device returning nullptr\n");
131+
return nullptr;
132+
}
133+
if (DeviceNum == omp_get_initial_device()) {
134+
DP("Call to omp_get_uid_from_device returning host device UID\n");
135+
return GenericDeviceTy::getHostDeviceUid();
136+
}
137+
138+
llvm::Expected<DeviceTy &> Device = PM->getDevice(DeviceNum);
139+
if (!Device) {
140+
FATAL_MESSAGE(DeviceNum, "%s", toString(Device.takeError()).c_str());
141+
return nullptr;
142+
}
143+
144+
const char *Uid = Device->RTL->getDevice(Device->RTLDeviceID).getDeviceUid();
145+
DP("Call to omp_get_uid_from_device returning %s\n", Uid);
146+
return Uid;
147+
}
148+
94149
EXTERN int omp_get_initial_device(void) {
95150
TIMESCOPE();
96151
OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));

‎offload/libomptarget/exports‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ VERS1.0 {
4040
omp_get_mapped_ptr;
4141
omp_get_num_devices;
4242
omp_get_device_num;
43+
omp_get_device_from_uid;
44+
omp_get_uid_from_device;
4345
omp_get_initial_device;
4446
omp_target_alloc;
4547
omp_target_free;

‎offload/test/api/omp_device_uid.c‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
// RUN: %libomptarget-compile-run-and-check-generic
2+
3+
#include <omp.h>
4+
#include <stdio.h>
5+
#include <string.h>
6+
7+
// Note that the device UIDs for the "fake" host devices used by libomptarget
8+
// will always be the same as the UID for the initial device (since it *is* the
9+
// same device). The other way round, the device number returned for this UID
10+
// will always be the initial device.
11+
12+
int is_host_device_uid(const char *device_uid) {
13+
return strcmp(device_uid,
14+
omp_get_uid_from_device(omp_get_initial_device())) == 0;
15+
}
16+
17+
int test_omp_device_uid(int device_num) {
18+
const char *device_uid = omp_get_uid_from_device(device_num);
19+
if (device_uid == NULL) {
20+
printf("FAIL for device %d: omp_get_uid_from_device returned NULL\n",
21+
device_num);
22+
return 0;
23+
}
24+
25+
int device_num_from_uid = omp_get_device_from_uid(device_uid);
26+
if (device_num_from_uid != (is_host_device_uid(device_uid)
27+
? omp_get_initial_device()
28+
: device_num)) {
29+
printf(
30+
"FAIL for device %d: omp_get_device_from_uid returned %d (UID: %s)\n",
31+
device_num, device_num_from_uid, device_uid);
32+
return 0;
33+
}
34+
35+
if (device_num == omp_get_initial_device())
36+
return 1;
37+
38+
int success = 1;
39+
40+
// Note that the following code may be executed on the host if the host is the
41+
// device
42+
#pragma omp target map(tofrom : success) device(device_num)
43+
{
44+
int device_num = omp_get_device_num();
45+
46+
// omp_get_uid_from_device() in the device runtime is a dummy function
47+
// returning NULL
48+
const char *device_uid_target = omp_get_uid_from_device(device_num);
49+
50+
// omp_get_device_from_uid() in the device runtime is a dummy function
51+
// returning omp_invalid_device.
52+
device_num_from_uid = omp_get_device_from_uid(device_uid_target);
53+
54+
// Depending on whether we're executing on the device or the host, we either
55+
// got NULL as the device UID or the correct device UID. Consequently,
56+
// omp_get_device_from_uid() either returned omp_invalid_device or the
57+
// correct device number (aka omp_get_initial_device()).
58+
if (device_uid_target ? device_num_from_uid != omp_get_initial_device()
59+
: device_num_from_uid != omp_invalid_device) {
60+
printf("FAIL for device %d (target): omp_get_device_from_uid returned %d "
61+
"(UID: %s)\n",
62+
device_num, device_num_from_uid, device_uid_target);
63+
success = 0;
64+
}
65+
}
66+
67+
return success;
68+
}
69+
70+
int main() {
71+
int num_devices = omp_get_num_devices();
72+
int num_failed = 0;
73+
// (also test initial device aka num_devices)
74+
for (int i = 0; i < num_devices + 1; i++) {
75+
if (!test_omp_device_uid(i)) {
76+
printf("FAIL for device %d\n", i);
77+
num_failed++;
78+
}
79+
}
80+
if (num_failed) {
81+
printf("FAIL\n");
82+
return 1;
83+
}
84+
printf("PASS\n");
85+
return 0;
86+
}
87+
88+
// CHECK: PASS

‎openmp/device/include/DeviceTypes.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ template <typename T> using Constant = __gpu_constant T;
2121
template <typename T> using Local = __gpu_local T;
2222
template <typename T> using Global = __gpu_local T;
2323

24+
#define omp_invalid_device -2
25+
2426
enum omp_proc_bind_t {
2527
omp_proc_bind_false = 0,
2628
omp_proc_bind_true = 1,

‎openmp/device/include/Interface.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,10 @@ int omp_get_num_devices(void);
130130

131131
int omp_get_device_num(void);
132132

133+
int omp_get_device_from_uid(const char *DeviceUid);
134+
135+
const char *omp_get_uid_from_device(int DeviceNum);
136+
133137
int omp_get_num_teams(void);
134138

135139
int omp_get_team_num();

‎openmp/device/src/State.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,12 @@ int omp_get_num_devices(void) { return config::getNumDevices(); }
423423

424424
int omp_get_device_num(void) { return config::getDeviceNum(); }
425425

426+
int omp_get_device_from_uid(const char *DeviceUid) {
427+
return omp_invalid_device;
428+
}
429+
430+
const char *omp_get_uid_from_device(int DeviceNum) { return nullptr; }
431+
426432
int omp_get_num_teams(void) { return mapping::getNumberOfBlocksInKernel(); }
427433

428434
int omp_get_team_num() { return mapping::getBlockIdInKernel(); }

‎openmp/runtime/src/dllexports‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,8 @@ kmp_set_disp_num_buffers 890
544544
omp_get_devices_all_allocator 819
545545
omp_get_memspace_num_resources 820
546546
omp_get_submemspace 821
547+
omp_get_device_from_uid 822
548+
omp_get_uid_from_device 823
547549
%ifndef stub
548550
__kmpc_set_default_allocator
549551
__kmpc_get_default_allocator

‎openmp/runtime/src/include/omp.h.var‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,11 @@
536536

537537
/* OpenMP 5.2 */
538538
extern int __KAI_KMPC_CONVENTION omp_in_explicit_task(void);
539+
#define omp_invalid_device -2
540+
541+
/* OpenMP 6.0 */
542+
extern int __KAI_KMPC_CONVENTION omp_get_device_from_uid(const char *DeviceUid);
543+
extern const char * __KAI_KMPC_CONVENTION omp_get_uid_from_device(int DeviceNum);
539544

540545
/* LLVM Extensions */
541546
extern void *llvm_omp_target_dynamic_shared_alloc(void);

0 commit comments

Comments
 (0)