Provide an Implementation of ONEAPI_DEVICE_SELECTOR

[alycm]

Currently SYCL and OpenMP have their own implementations of ONEAPI_DEVICE_SELECTOR, see compiler docs. Those could be eliminated if Unified Runtime provides a single solution for this.

The expectation would be that Unified Runtime implements ONEAPI_DEVICE_SELECTOR directly while matching current behavior as end-users could already be making use of this feature.

Level Zero also has ZE_AFFINITY_MASK, see Level Zero programming guide, which does a similar thing. I believe that ONEAPI_DEVICE_SELECTOR replaces ZE_AFFINITY_MASK but their interaction is something to consider.

[jandres742]

@alycm thanks.

Idea is that oneAPI applications (or in this case, applications using UR) set either ONEAPI_DEVICE_SELECTOR or the ZE_AFFINITY_MASK.

Main idea is that with ONEAPI_DEVICE_SELECTOR, oneAPI libraries expose only a sub-set of devices to the application, while letting all middleware libraries (and UR) having full visibility of all the devices in the system.

This is very useful for instance for MPI+SYCL applications running on multi-device systems, where each rank would have visibility to just 1 or a sub-set of those devices, while MPI would have full visibility to all devices, allowing it to implement any optimizations that rely on the knowledge of the full topology of the system.

Now, if an application choses to use both ONEAPI_DEVICE_SELECTOR and ZE_AFFINITY_MASK, then ZE_AFFINITY_MASK defines what the ONEAPI_DEVICE_SELECTOR may expose, given that L0 sits lower in the stack and closer to the devices.

In that sense, I think UR would just read the ONEAPI_DEVICE_SELECTOR, and mask the devices that L0 already exposes to the layers above:

Example 0 (typical usage):
8 GPU system, d0 to d7
ZE_AFFINITY_MASK not set: L0 exposes zeDevice0 (d0) to zeDevice7(d7)
ONEAPI_DEVICE_SELECTOR set to only expose 2 devices: Application sees only 2 devices, zeDevice0=d0 and zeDevice1=d1

Example 1 (more complex, but still legal):
8 GPU system, d0 to d7
ZE_AFFINITY_MASK set to expose only d3 to d7: L0 exposes zeDevice0 (d3) to zeDevice3(d7)
ONEAPI_DEVICE_SELECTOR set to only expose 2 devices: Application sees only 2 devices, zeDevice0=d3 and zeDevice1=d4

so from the UR point of view, ONEAPI_DEVICE_SELECTOR would expose the devices L0 decides to expose through the L0 device query interfaces.

[alycm]

I can't assign this ticket to him as he isn't a member of the project but @Wee-Free-Scot is working on this.

[Wee-Free-Scot]

Thanks @alycm -- I am indeed working on this. Is there any way I can be added to the project?

[kbenzie]

I've made the request to add you @Wee-Free-Scot

[Wee-Free-Scot]

Goal: there is a single env var (ODS) that consistently affects how devices are exposed by many oneAPI components.

Relevance to UR: "consistently" implies a single code-base for parts of the functionality that are common,

Code/API design concerns:

Some targets need to see all the "real" devices (i.e. ignore the ODS env var): MPI, oneCCL, and oneSHMEM.
Some targets need to see only "selected" devices (i.e. obey the ODS env var): SYCL, OpenMP, CHIP-SPV, Julia
(Other oneAPI components: TBC.)

Of those targets that need "selected" devices,
some permit at most one backend: OpenMP and CHIP-SPV
some permit multiple backends: SYCL
(Others: TBC.)

Common functionality:

Figuring out whether the ODS env var is set.
Reading the ODS env var string from the environment.
Parsing the ODS string into terms.
If all terms are exclusions, prefix an additional "accept any device" term.

Diverging functionality:

Return list of devices, ignoring ODS entirely.
Return list of devices, obeying only ODS terms for backend X (supplied by caller).
Return list of devices, obeying all ODS terms.

Suggestion -- UR needs 3 APIs:

1. urGetRealDevices
2. urGetSelectedDevices
3. urGetSelectedDevicesForPlatform

[pbalcer]

Instead of adding new APIs for enumerating devices, it might be easier to decide whether to use "OneAPI Device Selector" (ODS) based on some flag to urInit, either as an opt-in or an opt-out.

[Wee-Free-Scot]

The flag might need to be more than just yes/no -- some callers permit only one platform/backend and some might need to pass in the name of that one platform/backend.

CHIP-SPV assumes "there is one backend and it is Level Zero", unless an env var is set that switches to "there is only backend and it is OpenCL". CHIP-SPV could read its env var, decide between only-L0 or only-OCL, and then initialize UR passing in backend name, e.g. "level_zero" (string or an enum?).

OpenMP assumes "there is one backend" but does not (AFAIK) offer a way to control which is the chosen backend. There is an open question of how OpenMP should make the choice. One option is to choose the first backend mentioned in ODS.

So the flag would at least need values meaning "ignore ODS" (for MPI), "single platform: <this one>" (for CHIP-SPV), and "multi-platform" (for SYCL). Perhaps there is another value meaning "single platform: you choose for me" (for OpenMP).

[pbalcer]

Sounds like there's some overlap with what we wanted to do in #355.

[kbenzie]

It seems to me that the default path should respect ONEAPI_DEVICE_SELECTOR and that the real devices query should be the opt in path for the environments where that is needed.

Suggestion -- UR needs 3 APIs:

1. urGetRealDevices

2. urGetSelectedDevices

3. urGetSelectedDevicesForPlatform

Since we already have urDeviceGet I think it would make sense for it to remain the default path, i.e. the path that respects ONEAPI_DEVICE_SELECTOR. It already takes a platform handle as the first argument, so I think fulfils API 3.

What differentiates 2. and 3. here? Is urGetSelectedDevices intended to return all devices for all platforms? If so that's a mismatch with how urDeviceGet currently works as we must first enumerate all platforms before enumerating all devices.

The flag might need to be more than just yes/no -- some callers permit only one platform/backend and some might need to pass in the name of that one platform/backend.

Given that ONEAPI_DEVICE_SELECTOR is also filtering backends (which are a 1:1 mapping to platforms in L0, CUDA, and HIP, but not OpenCL) do we also need a way to get the list of "real" platforms/backends?

Instead of adding new APIs for enumerating devices, it might be easier to decide whether to use "OneAPI Device Selector" (ODS) based on some flag to urInit, either as an opt-in or an opt-out.

I've wondered about this also but my mental (which may be incorrect) model for urInit is it would result in a platform not even being initilaized and thus the real devices query would not function as expected in environments.