Is there a way to bifurcate GPU for simulating multiple?

[cadop]

Since running tests on the HPC takes up time away from others, and warp doesn't auto-scale to GPU(s), I was hoping to simulate multiple devices that I can ensure distribution of data and synchronization is done correctly.

Is there a way to do this? So basically end up with ~ for device in wp.get_simulated_devices(4):

[nvlukasz]

Hi @cadop, there is a way to do this. The idea is to create multiple CUDA contexts on one device, which will be treated as independent Warp CUDA devices at the Python level.

Unfortunately I just found a small regression in our CUDA context creation function. I should be able to fix it quickly and will come back with more instructions.

[nvlukasz]

Created a bug to track this: #402

[nvlukasz]

Ok, the fix has been merged to main and should appear in release 1.5.1.

Here's an example of creating and using some "virtual" devices:

import warp as wp


@wp.kernel
def arange(a: wp.array(dtype=int)):
    tid = wp.tid()
    a[tid] = tid


# ======================================================================
# Create virtual devices
# ======================================================================

NUM_DEVICES = 4
virtual_devices = []

# NOTE: must call wp.init() to initialize wp.context.runtime
wp.init()

for i in range(NUM_DEVICES):
    # create a new CUDA context on cuda:0
    ctx = wp.context.runtime.core.cuda_context_create(0)

    # map as new virtual device
    device = wp.map_cuda_device(f"vcuda:{i}", context=ctx)
    virtual_devices.append(device)

print(f"\nVirtual devices: {virtual_devices}")


# ======================================================================
# Use virtual devices
# ======================================================================

for device in virtual_devices:
    with wp.ScopedDevice(device):
        print(f"\nRunning on {device}")
        n = 16
        a = wp.zeros(n, dtype=int)
        wp.launch(arange, dim=n, inputs=[a])
        print(a)

This creates new CUDA contexts on device 0 using the function wp.context.runtime.core.cuda_context_create(0) and maps the contexts as Warp devices with aliases "vcuda:0", "vcuda:1", etc.

This is an "unofficial" way of adding new devices and is mostly untested, but it should work for emulating multiple GPUs. We can consider adding a proper API for this in the future.

Some caveats:

• Create the new contexts/devices at the beginning of the program and don't delete them.
• Each CUDA context can use a considerable amount of memory on the physical device.
• The different CUDA contexts are asynchronous with respect to each other, so synchronization is needed just like with multiple physical devices. This is actually a good thing if your goal is to emulate multiple GPUs.

[cadop]

Thanks!