Device Agnostic Code - sickkids-mri/mrrecon GitHub Wiki

It is necessary to perform computationally expensive image reconstructions using the GPU. For completeness, the same functions should also be able to run on the CPU. Basically, NumPy arrays are used when computation on the CPU is desired, and CuPy arrays are used when computation on the GPU is desired. Device agnosticism is possible because NumPy and CuPy have nearly identical interfaces.

The 'xp' Module

When writing device agnostic code, instead of using np to refer to the NumPy module and cp to refer to the CuPy module, the name xp is used to refer to either module, where the value for xp is determined at runtime.

This technique originated from CuPy, and SigPy builds on it by making the technique possible to use in the case where CuPy is not installed (for users who are not using GPU).

The following examples will demonstrate how the device can be determined, which would then determine the value of xp.

Examples

Specifying the Device Explicitly

A function that creates new arrays can have the device provided explicitly as an argument. The variable device is an integer that indicates which device should be used. The integer -1 refers to the CPU, 0 refers to the first GPU, 1 refers to the second GPU, 2 refers to the third GPU, and so on. Once the device is known, the value for xp is straightforward to determine.

import sigpy as sp


# Note: Importing numpy and cupy was not necessary

def create_array(device=-1):
    # It is good practice to use the CPU as the default device

    # First convert the device integer to the sp.Device object
    device = sp.Device(device)

    # Get the array module
    xp = device.xp

    # Create array
    array = xp.zeros((8, 512, 512), dtype=xp.complex64)
    return array


a = create_array(-1)  # This array is on regular RAM

a = create_array(0)  # This array is on GPU RAM

Inferring the Device From the Input Array

Functions that perform operations with/on arrays can infer the device from the input array. This makes sense to do when computation should be performed on the same device as the input array (which should be the vast majority of cases).

import sigpy as sp

def sum(array):
    # Get the array module from the input array
    xp = sp.get_array_module(array)

    # Can also get the device from the input array
    device = sp.get_device(array)
    xp = device.xp

    out = xp.sum(array)
    return out

Multi-GPU Management

When multiple GPUs are available, which GPU new arrays should be created on must be specified. Otherwise new arrays are created on the default GPU. The GPU to be used is specified using a context manager, as shown below.

import sigpy as sp

device0 = sp.Device(0)
device1 = sp.Device(1)

xp = device0.xp  # CuPy module

# Creates array on the default GPU (usually the first GPU)
a = xp.ones((2, 3))

with device0:
    a = xp.ones((2, 3))  # On GPU 0

with device1:
    b = xp.ones((2, 3))  # On GPU 1

# c = a + b
# This will raise an error because the input arrays are on different GPUs

# Move the array on GPU 0 to GPU 1
a = sp.to_device(a, device1)

# Both arrays are now on GPU 1

# c = a + b
# The output array is on the default GPU, which would raise an error if
# the default GPU is not GPU 1

with device1:
    c = a + b  # The output array is on GPU 1