Python on Aurora¶

AI/ML Framework Modules¶

For most Python users on Aurora, a good starting point are the AI/ML framework modules available on the system. For PyTorch users, we recommend

module load frameworks

and for TensorFlow users we provide a separate module accessible with

module load tensorflow
module load mpich/opt/develop-git.6037a7a

Please note that:

• Both modules automatically load a conda environment with GPU-supported builds of the respective framework libraries along with other popular Python and ML packages.
• The frameworks and tensorflow modules may load a different oneAPI compiler SDK than the default module.
• The frameworks and tensorflow modules are updated approximately every quarter.

For more information on PyTorch and TensorFlow on Aurora, please see the respective pages:

• PyTorch
• TensorFlow

Virtual environments via venv¶

While the Anaconda environment automatically loaded with the frameworks and tensorflow modules contains many of the most commonly used Python packages for our users, you may still encounter a scenario in which you need to extend the functionality of the environment (i.e. install additional packages). In this case, we suggest the use of Python virtual environments.

Creating and activating a new virtual environment (venv) is straightforward. After loading the above module, execute:

python3 -m venv /path/to/new/venv --system-site-packages
source /path/to/new/venv/bin/activate

The --system-site-packages flag will make sure that all the packages included in the frameworks or tensorflow modules are still available after sourcing the venv. If, however, you would like to create an empty venv, simply remove this flag. You can always retroactively change the --system-site-packages flag state for this virtual environment by editing venv/pyvenv.cfg and changing the value of include-system-site-packages to true.

To install a different version of a package that is already installed in the base environment, you can use:

pip install --ignore-installed ... # or -I

Any time you wish to use this virtual environment in future shell sessions, be sure to first load the base AI/ML modules (e.g., module load frameworks) before source /path/to/new/venv/bin/activate.

An alternative, although not recommended, approach to creating a venv is to install packages with

pip install --user ...

Intel's Data Parallel Extensions for Python (DPEP)¶

On Aurora, users can access Intel's Python stack comprising of compilers and libraries for programming heterogenous devices, namely the Data Parallel Extensions for Python (DPEP). DPEP is composed of three main packages for programming on CPUs and GPUs:

• dpnp - Data Parallel Extensions for Numpy is a library that implements a subset of Numpy that can be executed on any data parallel device. The subset is a drop-in replacement of core Numpy functions and numerical data types, similar to CuPy for CUDA devices.
• dpctl - Data Parallel Control library provides utilities for device selection, allocation of data on devices, tensor data structure along with Python Array API Standard implementation (deprecated - see details below), and support for creation of user-defined data-parallel extensions.
• numba_dpex - Data Parallel Extensions for Numba is an extension to Numba compiler for programming data-parallel devices similar to developing programs with Numba for CPU or CUDA devices.

The DPEP packages follow the compute-follows-data programming model, meaning that the offload target for a Python library call, or a hand-written kernel using numba-dpex, does not need to be specified directly when making the call. Instead, the offload target is inferred from the input arguments to the library call. With this programming model, the user only needs to specify the offload target when creating the tensor/ndarray objects. For example,

1
2
3
4
5

import dpnp

x_gpu = dpnp.arange(100, device="gpu")
sqx_gpu = dpnp.square(x_gpu) # (1)!
print(sqx_gpu.device) # (2)!

1. dpnp.square() offloads to the "gpu" device
2. sqx_gpu is created on the "gpu" device

Device(level_zero:gpu:0)

However, note that operating on arrays created on different devices will raise an exception.

Accessing the DPEP Packages¶

Users can access the dpnp and dpctl packages by simply loading the latest AI/ML frameworks module with

module load frameworks

However, numba-dpex is not available in the frameworks module, thus using this package requires additional installation steps. The easiest way to get all three DPEP packages is to create a new conda environment and install the DPEP packages with the following recipe:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

module load frameworks
module load cmake
conda create -y --prefix /path/to/dpep_env python=3.12 pip
conda activate /path/to/dpep_env
conda install -y scikit-build
pip install versioneer numpy Cython ninja

# dpctl
git clone https://github.com/IntelPython/dpctl.git
cd dpctl
git checkout 0.21.1
CXX=$(which dpcpp) python setup.py install
cd ..

# dpnp
git clone https://github.com/IntelPython/dpnp.git
cd dpnp
git checkout 0.19.1
CXX=$(which dpcpp) python setup.py install -- -G Ninja -DCMAKE_C_COMPILER:PATH=`which icx` -DCMAKE_CXX_COMPILER:PATH=`which icpx`
cd ..

# numba-dpex
conda install -y numba==0.59* -c conda-forge
git clone https://github.com/argonne-lcf/numba-dpex.git
cd numba-dpex
CXX=$(which dpcpp) python setup.py develop
cd ..

dpnp¶

The dpnp library implements the NumPy API using DPC++ and is meant to serve as a drop-in replacement for NumPy, similar to CuPy for CUDA devices. Therefore, dpnp should be used to port NumPy and CuPy code to Intel GPU, however, please refer to this comparison table to check the current coverage of dpnp API relative to NumPy and CuPy.

Below is a minimal example using dpnp to create and operate on an array allocated on the PVC.

1
2
3
4
5
6

import dpnp as np

x = np.asarray([1, 2, 3]) # (1)!
print("Array x allocated on the device:", x.device)
y = np.sum(x) # (2)!
print("Result y is located on the device:", y.device)

1. np.asarray() creates an array on the default SYCL device, which is the Intal Max 1550 GPU on Aurora. The queue associated with this array is now carried with x, and the pre-compiled kernel for np.sum(x) is submitted to that queue.
2. The result y is allocated on the device and is associated with the queue of x.

Array x allocated on the device: Device(level_zero:gpu:0)
Result y is located on the device: Device(level_zero:gpu:0)

All dpnp array creation routines and random number generators have additional optional keyword arguments (device, queue, and usm_type) which users can leverage to explicitly specify on which device or queue they want the data to be created along with the USM memory type to be used.

1
2
3
4
5
6
7
8

import dpnp as np
from time import perf_counter

x = np.random.randn(1000,1000)
tic = perf_counter()
y = np.matmul(x,x)
y.sycl_queue.wait()
print(f"Execution time: {perf_counter() - tic} sec")

dpctl¶

The dpctl package lets users access devices supported by the DPC++ SYCL runtime. The package exposes features such as device instrospection, execution queue creation, memory allocation, and kernel submission. Below are some of the basic device management functions, but more functionality is available on the dpctl documentation.

1
2
3
4
5
6
7
8
9

import dpctl

dpctl.lsplatform()
num_devices = dpctl.get_num_devices(device_type="gpu") # (1)!
device_list =  dpctl.get_devices(device_type="gpu") # (2)!
print(f"Found {num_devices} GPU devices")
for device in device_list:
    print(f"\t{device}")
print("\nFound CPU devices: ", dpctl.has_cpu_devices()) # (3)!

1. Get the number of GPU devices on the node
2. Get the list of GPU devices on the node
3. Check if CPU devices are available on the node

Intel(R) Level-Zero 1.5
Found 12 GPU devices
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da03430>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da034f0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da035b0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da03530>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da03f70>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8d005770>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8b5a3ab0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8d2467b0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da24ef0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da03fb0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da240b0>
    <dpctl.SyclDevice [backend_type.level_zero, device_type.gpu,  Intel(R) Data Center GPU Max 1550] at 0x154e8da24130>

Found CPU devices:  False

numba-dpex¶

Numba-dpex is Intel's Data Parallel Extension for Numba which allows users to apply Numba's JIT compiler and generate performant, parallel code on Intel's GPU. Its LLVM-based code generator implements a new kernel programming API (kapi) in pure Python that is modeled after the SYCL API. The example below implements and launches simple vector addition as a range kernel. Range kernels implement a basic parallel-for calculation that is ideally suited for embarrassingly parallel operations, such as element-wise computations over n-dimensional arrays.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

import dpnp
import numba_dpex as dpex
from numba_dpex import kernel_api as kapi

# Data parallel kernel implementation of vector sum
@dpex.kernel # (1)!
def vecadd(item: kapi.Item, a, b, c):
    i = item.get_id(0) # (2)!
    c[i] = a[i] + b[i]

N = 1024 # (3)!
a = dpnp.ones(N)
b = dpnp.ones_like(a)
c = dpnp.zeros_like(a)
dpex.call_kernel(vecadd, dpex.Range(N), a, b, c)
assert dpnp.allclose(c,a+b)
print("Sum completed successfully")

1. Decorate the vecadd function as a dpex kernel
2. Get the work item
3. Define the number of work items

The vecadd function, when decorated as a dpex kernel, is compiled with numba-dpex into a data-parallel function to be executed individually by a set of work items (item.get_id(0)). Numba-dpex follows the SPMD programming model, wherein each work item runs the function for a subset of the elements of the input arrays.
The set of work items is defined by the dpex.Range() object and the dpex.call_kernel() call instructs every work item in the range to execute the vecadd kernel for a specific subset of the data. Numba-dpex also follows the compute-follows-data programming model, meaning that the kernel is run on the same device as the dpnp and dpctl arrays passed as inputs.

Note that the numba-dpex kapi allows for more complex data parallel kernels (e.g., nd-range kernels) and the ability to create device callable functions. For these and more features, we refer the users to the numba-dpex documentation.

DLPack¶

Thanks to dpctl supporting the Python Array API standard, both dpnp and dpctl provide interoperability with other Python libraries that follow the same standards, such as Numpy and PyTorch, through DLPack. This allows for zero-copy data access across the Python ecosystem.

An example of using DLPack to pass arrays between dpnp and PyTorch is shown below

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

import dpnp as dp
import torch
import intel_extension_for_pytorch as ipex

t_ary = torch.arange(4).to('xpu') # array [0, 1, 2, 3] on GPU
dp_ary = dp.from_dlpack(t_ary)
t_ary[0] = -2.0 # modify the PyTorch array
print(f'Original PyTorch array: {t_ary}')
print(f'dpnp view of PyTorch array: {dp_ary} on device {dp_ary.device}\n')
del t_ary, dp_ary

dp_ary = dp.arange(4) # array [0, 1, 2, 3] on GPU
t_ary = torch.from_dlpack(dp_ary)
dp_ary[0] = -3.0 # modify the dpnp array
print(f'Original dpnp array: {dp_ary} on device {dp_ary.device}')
print(f'PyTorch view of dpnp array: {t_ary}')

Original PyTorch array: tensor([-2,  1,  2,  3], device='xpu:0')
dpnp view of PyTorch array: [-2  1  2  3] on device Device(level_zero:gpu:0)

Original dpnp array: [-3  1  2  3] on device Device(level_zero:gpu:0)
PyTorch view of dpnp array: tensor([-3,  1,  2,  3], device='xpu:0')