Skip to content

Commit 65c4a53

Browse files
ro-iro-i
authored
[OpenMP] Implement omp_get_uid_from_device() / omp_get_device_from_uid() (#164392)
Use the implementation in libomptarget. If libomptarget is not available, always return the UID / device number of the host / the initial device.
1 parent 38891ba commit 65c4a53

File tree

15 files changed

+310
-2
lines changed

15 files changed

+310
-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
@@ -270,6 +270,8 @@ extern "C" {
270270
void ompx_dump_mapping_tables(void);
271271
int omp_get_num_devices(void);
272272
int omp_get_device_num(void);
273+
int omp_get_device_from_uid(const char *DeviceUid);
274+
const char *omp_get_uid_from_device(int DeviceNum);
273275
int omp_get_initial_device(void);
274276
void *omp_target_alloc(size_t Size, int DeviceNum);
275277
void omp_target_free(void *DevicePtr, int DeviceNum);

‎offload/libomptarget/OpenMP/API.cpp‎

Lines changed: 58 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,
@@ -68,6 +70,62 @@ EXTERN int omp_get_device_num(void) {
6870
return HostDevice;
6971
}
7072

73+
static inline bool is_initial_device_uid(const char *DeviceUid) {
74+
return strcmp(DeviceUid, GenericPluginTy::getHostDeviceUid()) == 0;
75+
}
76+
77+
EXTERN int omp_get_device_from_uid(const char *DeviceUid) {
78+
TIMESCOPE();
79+
OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
80+
81+
if (!DeviceUid) {
82+
DP("Call to omp_get_device_from_uid returning omp_invalid_device\n");
83+
return omp_invalid_device;
84+
}
85+
if (is_initial_device_uid(DeviceUid)) {
86+
DP("Call to omp_get_device_from_uid returning initial device number %d\n",
87+
omp_get_initial_device());
88+
return omp_get_initial_device();
89+
}
90+
91+
int DeviceNum = omp_invalid_device;
92+
93+
auto ExclusiveDevicesAccessor = PM->getExclusiveDevicesAccessor();
94+
for (const DeviceTy &Device : PM->devices(ExclusiveDevicesAccessor)) {
95+
const char *Uid = Device.RTL->getDevice(Device.RTLDeviceID).getDeviceUid();
96+
if (Uid && strcmp(DeviceUid, Uid) == 0) {
97+
DeviceNum = Device.DeviceID;
98+
break;
99+
}
100+
}
101+
102+
DP("Call to omp_get_device_from_uid returning %d\n", DeviceNum);
103+
return DeviceNum;
104+
}
105+
106+
EXTERN const char *omp_get_uid_from_device(int DeviceNum) {
107+
TIMESCOPE();
108+
OMPT_IF_BUILT(ReturnAddressSetterRAII RA(__builtin_return_address(0)));
109+
110+
if (DeviceNum == omp_invalid_device) {
111+
DP("Call to omp_get_uid_from_device returning nullptr\n");
112+
return nullptr;
113+
}
114+
if (DeviceNum == omp_get_initial_device()) {
115+
DP("Call to omp_get_uid_from_device returning initial device UID\n");
116+
return GenericPluginTy::getHostDeviceUid();
117+
}
118+
119+
auto DeviceOrErr = PM->getDevice(DeviceNum);
120+
if (!DeviceOrErr)
121+
FATAL_MESSAGE(DeviceNum, "%s", toString(DeviceOrErr.takeError()).c_str());
122+
123+
const char *Uid =
124+
DeviceOrErr->RTL->getDevice(DeviceOrErr->RTLDeviceID).getDeviceUid();
125+
DP("Call to omp_get_uid_from_device returning %s\n", Uid);
126+
return Uid;
127+
}
128+
71129
EXTERN int omp_get_initial_device(void) {
72130
TIMESCOPE();
73131
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: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
// RUN: %libomptarget-compile-run-and-check-generic
2+
3+
#include <omp.h>
4+
#include <stdio.h>
5+
#include <string.h>
6+
7+
int test_omp_device_uid(int device_num) {
8+
const char *device_uid = omp_get_uid_from_device(device_num);
9+
if (device_uid == NULL) {
10+
printf("FAIL for device %d: omp_get_uid_from_device returned NULL\n",
11+
device_num);
12+
return 0;
13+
}
14+
15+
int device_num_from_uid = omp_get_device_from_uid(device_uid);
16+
if (device_num_from_uid != device_num) {
17+
printf(
18+
"FAIL for device %d: omp_get_device_from_uid returned %d (UID: %s)\n",
19+
device_num, device_num_from_uid, device_uid);
20+
return 0;
21+
}
22+
23+
if (device_num == omp_get_initial_device())
24+
return 1;
25+
26+
int success = 1;
27+
28+
// Note that the following code may be executed on the host if the host is the
29+
// device
30+
#pragma omp target map(tofrom : success) device(device_num)
31+
{
32+
int device_num = omp_get_device_num();
33+
34+
// omp_get_uid_from_device() in the device runtime is a dummy function
35+
// returning NULL
36+
const char *device_uid = omp_get_uid_from_device(device_num);
37+
38+
// omp_get_device_from_uid() in the device runtime is a dummy function
39+
// returning omp_invalid_device.
40+
int device_num_from_uid = omp_get_device_from_uid(device_uid);
41+
42+
// Depending on whether we're executing on the device or the host, we either
43+
// got NULL as the device UID or the correct device UID. Consequently,
44+
// omp_get_device_from_uid() either returned omp_invalid_device or the
45+
// correct device number (aka omp_get_initial_device()).
46+
if (device_uid ? device_num_from_uid != device_num
47+
: device_num_from_uid != omp_invalid_device) {
48+
printf("FAIL for device %d (target): omp_get_device_from_uid returned %d "
49+
"(UID: %s)\n",
50+
device_num, device_num_from_uid, device_uid);
51+
success = 0;
52+
}
53+
}
54+
55+
return success;
56+
}
57+
58+
int main() {
59+
int num_devices = omp_get_num_devices();
60+
int num_failed = 0;
61+
// (also test initial device aka num_devices)
62+
for (int i = 0; i < num_devices + 1; i++) {
63+
if (!test_omp_device_uid(i)) {
64+
printf("FAIL for device %d\n", i);
65+
num_failed++;
66+
}
67+
}
68+
if (num_failed) {
69+
printf("FAIL\n");
70+
return 1;
71+
}
72+
printf("PASS\n");
73+
return 0;
74+
}
75+
76+
// CHECK: PASS

‎openmp/device/include/DeviceTypes.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ 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+
// See definition in OpenMP (omp.h.var/omp_lib.(F90|h).var)
25+
#define omp_invalid_device -2
26+
2427
enum omp_proc_bind_t {
2528
omp_proc_bind_false = 0,
2629
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
@@ -403,6 +403,12 @@ int omp_get_num_devices(void) { return config::getNumDevices(); }
403403

404404
int omp_get_device_num(void) { return config::getDeviceNum(); }
405405

406+
int omp_get_device_from_uid(const char *DeviceUid) {
407+
return omp_invalid_device;
408+
}
409+
410+
const char *omp_get_uid_from_device(int DeviceNum) { return nullptr; }
411+
406412
int omp_get_num_teams(void) { return mapping::getNumberOfBlocksInKernel(); }
407413

408414
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)