PopTorch Experimental Addons

A collection of addons to PopTorch, with general utility.

Usage

See documentation.

# Tested on Poplar SDK 3.3.0+7857, Ubuntu 20.04, Python 3.8
pip install git+https://github.com/graphcore-research/poptorch-experimental-addons

# Run an example
wget https://raw.githubusercontent.com/graphcore-research/poptorch-experimental-addons/main/examples/sparse_benchmark_spmm.py
python sparse_benchmark_spmm.py

# Import & use
import poptorch_experimental_addons as pea

Libraries

API	Description	Note
pea.autograd_proxy	Custom gradients using a proxy tensor	Example: straight_through_estimator.py
pea.distance_matrix	Pairwise broadcasted distance between two collections of vectors	Supports only L1/L2 distances on IPU
pea.collectives	Collectives across IPU program replicas	All-gather, All-reduce
pea.sharded	1-D Sharded matrix multiplications across IPU program replicas
pea.sparse	Static sparse-dense matmul (forward pass only)	Example: sparse_benchmark_spmm.py

Development

First-time setup:

python3 -m venv .venv
# Add to .venv/bin/activate
# source /PATH_TO_POPLAR_SDK/enable
source .venv/bin/activate
pip install wheel
pip install $POPLAR_SDK_ENABLED/../poptorch-*.whl
pip install -r requirements-dev.txt

We recommend symlinking some system libraries to third_party, to get the IDE helping:

mkdir third_party
ln -s $POPLAR_SDK_ENABLED third_party/poplar
ln -s $(cd $POPLAR_SDK_ENABLED/../popart-* && pwd) third_party/popart
# Add the following to .vscode/settings.json:
# "C_Cpp.default.includePath": [
#     "${workspaceFolder}/third_party/poplar/include",
#     "${workspaceFolder}/third_party/popart/include"
# ],

Note:

• ./dev tests -k PATTERN runs selected tests matching PATTERN
• ./dev python ... runs under Python, after building the native lib & setting PYTHONPATH
• When adding .cpp files, they should also be added to OBJECTS in Makefile
• Custom codelets should have names of form *_codelet.cpp. They don't need to be added to OBJECTS; see here for an example of how to include them in custom ops.
• When extending impl modules, note that we use __all__ to control the public API

License

Copyright (c) 2023 Graphcore Ltd. Licensed under the MIT License.

The included code is released under an MIT license, (see LICENSE).

Our dependencies are (see requirements.txt):

Component	About	License
numpy	Array processing library	BSD 3-Clause
einops	Tensor processing utilities	MIT

We also use additional Python dependencies for development/testing (see requirements-dev.txt).