I plan to keep this post as a constantly updating webpage of tips and tricks I’ve encountered while writing papers.

Subfigures for IEEE template

Use the Subcaption

\usepackage{subcaption

with these settings:

%%IEEE%% The following block of 4 lines makes subcaption behave the way IEEEtran expects it to behave.
\DeclareCaptionLabelSeparator{periodspace}{.\quad} %%IEEE%%
\captionsetup{font=footnotesize,labelsep=periodspace,singlelinecheck=false} %%IEEE%%
\captionsetup[sub]{font=footnotesize,singlelinecheck=true} %%IEEE%%
\renewcommand\thesubfigure{\alph{subfigure}} %%IEEE%%

Officially, subcaption is not supported by IEEE, but these settings make it appear as it should.

In the current setup, we can use pip install git+https://github.com/vatai/tadashi.git, which, in theory would take care of everything, however, in practice, since it is building PET and ISL in the background, this can go wrong if some dependencies are not installed. CIBuildWheel¹ solves this problem by building the extensions² in isolation and packaging them into a wheel³.

Obviously, the problem with binary packages is portability! And I’m not just talking about Linux vs Windows, but also withing Linux, a binary built on one Linux may not work on another Linux. To solve this, CIBuildWheel uses something called manylinux. As I understand, to the user (package maintainer), manylinux a set of Docker images, with somewhat standard (read older) glibc and other libraries, so that if you compile against those libraries, you cover a lot of systems. Also there are multiple manylinux Docker images, covering different versions of different libraries. CIBuildWheel builds wheels in these Docker containers and by default generates these Linux wheels (look at the filenames, the cpXYY indicates the Python version X-YY, and the rest of the name includes which manylinux was used – it’s a deep rabbit hole and I’ll stop here with the explanation). In theory it could generate wheels for Windows and Mac too :shrug:.

Prototyping

Before going “all in” to apply this to Tadashi, I made a PoC repo (called Complicated project with extensions, or cpwe for short) to test the critical parts (i.e., parts I didn’t quite understand) of this setup. This also included Cython integration, the new approach I plan to implement to interface the PET/ISL libraries with Python (instead of SWIG, which I currently use). Cython is not just a more “pythonic” approach, it also simplifies many things both in Tadashi’s code and in the build process.

The goals

I originally intended to cover (a simplified version of) everything including uploading binary packages to PyPI, and test pip install cpwe, but then I realised I can do things locally: When I saw all the wheels (.whl files) created by CIBuildWheel, I realised PyPI most likely just stores these .whl files and pip install cpwe probably just selects the appropriate .whl file, downloads it, and installs it, which can also simulated locally with pip install the_appropriate_wheel.whl!

One goal was to set up GitHub actions to generate all the wheels automatically. This was easy to achieve following the CIBuildWheel docs and also easy to verify, bacause the example in the docs uploaded all the wheels as artifacts. These wheels would then be uploaded via GH actions (using twine) to PyPI, something I already implemented in another project.

The other, more important, goal was making a wheel which includes binaries from a separate C library, i.e. has Python code which calls functions from a C library (i.e., ISL). I should elaborate.

Getting the right wheels

To summarise, we want to

• build a Python extension,
• using Cython,
• which would be calling a third party C library (namely ISL).

Extensions: The simple case

Most of the docs describe Cython extensions (or SWIG extensions for that matter) as standalone code, that are converted into C/C++ code and then built into an .so file. It looks something like this:

flowchart LR
    Cython["Cython foo.py or foo.pyx files"]
    C["Generated foo.c or foo.cpp files"]
    SO["Compiled foo.so binary file"]
    Cython --> C --> SO
    style C fill:#ff0,color:#000
    style SO fill:#f9f,color:#000

And that final foo.so file has all the code you need! So the dependencies are simple:

---
title: "Cython extensions: the simple case"
---
classDiagram
    direction LR
    PythonPackage <|-- Extension
    class PythonPackage{
      +cpwe.py Python files
      +build_from_pyproject_toml()
    }
    class Extension{
      +foo.pyx Cython files
      +cythonize_in_setup_py()
    }

Extensions with dependencies

However our situation is more complicated: Our foo.so has an extra dependency (ISL), which itself pulls in other dependencies (namely GMP) which are not necessary installed on the end-user’s system.

---
title: "Cython extensions with a dependencies"
---
classDiagram
    direction LR
    PythonPackage <|-- Extension
    Extension <|-- Dependency
    Dependency <|-- Lv2Dependency
    class PythonPackage{
      +cpwe.py Python files
      +build_from_pyproject_toml()
    }
    class Extension{
      +foo.pyx Cython files
      +cythonize_in_setup_py()
    }
    class Dependency{
      +isl.c files
      +build_with_make()
    }
    class Lv2Dependency{
      +GMP system package
      +apt_install()
    }

And this is why CIBuildWheel is DA BOOOOMB! It takes care of all this!

These are the relevant parts of the pyproject.toml file that configures the whole project:

1
2
3
4
5
6
7
8
9
[build-system]
requires = ["setuptools", "cython", "cibuildwheel"]
build-backend = "setuptools.build_meta"

[tool.cibuildwheel]
build-frontend="build[uv]"
before-all = ["./third_party/install.sh"]
# build = "cp313-manylinux_x86_64"
# manylinux-x86_64-image = "manylinux-mod"

In line 2 we list cibuildwheel as a build dependency (it can just be pip installed as any other Python package).

The block from line 5 covers the cibuildwheel specific settings. Setting the frontend to use uv potentially speeds up the build process. And finally the before-all allows us to call a script which builds our dependency, i.e. ISL.

The commented out lines were hacks to speed up debugging.

• The build field specifies which wheels should be build (uncommenting line 8 would produce only the wheel built agains Python 3.13 and manylinux_x86_64).
• The manylinux-x86_64-image field specifies the Docker image to be used for manylinux-x86_64. I created a custom Docker image with ISL prebuilt which sped up testing by skipping both downloading the GMP system packages from the Alma Linux repo, and building ISL from source. For this I obviously disabled the install.sh script in before-all.

Verification

To verify that everything works as Intended I wrote a small test script. The main steps of the script were the following:

1. Call uvx cibuildwheel, which simulates what happens in the GH action and builds the .whl files locally.
2. Create a temporary Python project, with uv.
3. Installing the cpwe package, i.e., newly generated.whl file in the project’s venv.
4. Running code which invokes cpwe so it calls ISL functions.

uv is used out of convenience to automate the process of creating a venv, activating the venv, installing a wheel to the venv, running a python script in the venv and disabling the venv. Both the directories where wheels are created (step 1) and the directory where the uv project is created (step 2) are delete and regenerated from scratch every time the script executes.

Additional verification

I used the process described above as a quick indicator whether things are working or not. To simulate installing and running the code in the package on a random computer, i.e., to double check that CIBuildWheel actually builds a wheel that has all the code required to be executed on any system (without any dependency outside of the Python ecosystem), I spun up an clean Ubuntu container, I copied the appropriate .whl file to it, installed uv and some other Python packages, and executed a similar script as above. This confirmed that it worked on any system, even if it didn’t have GMP or ISL installed.

I looked inside the venv after installing the wheel to see how this is achieved. Inside the site-packages I found the following files:

.
├── cpwe
│   ├── foo.c # C code (generated by Cython)
│   ├── foo.cpython-313-x86_64-linux-gnu.so #!!! .so file (built from the C code)
│   ├── foo.py # source containing Cython code (written by me)
│   └── isl.pxd # Cython "header" file to declaring the ISL functions (written by me)
...
├── cpwe.libs
│   ├── libgmp-d944b113.so.10.3.2  manylinux container #!!! .so copied from the system package within the
│   └── libisl-8ba7b133.so.23.4.0 #!!! .so file built in the container

CIBuildWheel added all the relevant .so files: the foo.cython*.so generated from my Cython code in foo.py (which would be the only file if this were the simple case) and the libisl*.so* which was Cython extension’s dependency, and libgmp*.so* which was libisl*.so*’s dependency. :tada::tada::tada:

Future work

For now, this really was a minimal working example for building portable wheel packages, and only included ISL as a dependency. I need to see if this works also for PET which includes LLVM as a dependency – which is slightly larger than GMP. :cold_sweat:

In a separate attempt (which I didn’t publish in a repo or blog post), I managed to write a Cython function which was used as the callback parameter in the C function which is PET’s entry point. This callback mechanism was one of the main reasons I wrote most of Tadashi in C, I imagined it would be hard to interface Python code with ISL’s callback setup. Figuring out how to create callbacks in Cython I can more freely mix Python and C code, which would enable a sane way to organise code such that it includes both the current PET and future (Fortran enabled) LLVM/Polly backend.

```plotly
{
  "data": [
    {
      "x": [1, 2, 3, 4],
      "y": [10, 15, 13, 17],
      "type": "scatter"
    },
    {
      "x": [1, 2, 3, 4],
      "y": [16, 5, 11, 9],
      "type": "scatter"
    }
  ]
}
```

Which generates:

{
  "data": [
    {
      "x": [1, 2, 3, 4],
      "y": [10, 15, 13, 17],
      "type": "scatter"
    },
    {
      "x": [1, 2, 3, 4],
      "y": [16, 5, 11, 9],
      "type": "scatter"
    }
  ]
}

Also another example chart.

```plotly
{
  "data": [
    {
      "x": [1, 2, 3, 4],
      "y": [10, 15, 13, 17],
      "mode": "markers"
    },
    {
      "x": [2, 3, 4, 5],
      "y": [16, 5, 11, 9],
      "mode": "lines"
    },
    {
      "x": [1, 2, 3, 4],
      "y": [12, 9, 15, 12],
      "mode": "lines+markers"
    }
  ],
  "layout": {
    "title": {
      "text": "Line and Scatter Plot"
    }
  }
}
```

This is how it looks like:

{
  "data": [
    {
      "x": [1, 2, 3, 4],
      "y": [10, 15, 13, 17],
      "mode": "markers"
    },
    {
      "x": [2, 3, 4, 5],
      "y": [16, 5, 11, 9],
      "mode": "lines"
    },
    {
      "x": [1, 2, 3, 4],
      "y": [12, 9, 15, 12],
      "mode": "lines+markers"
    }
  ],
  "layout": {
    "title": {
      "text": "Line and Scatter Plot"
    }
  }
}

Benchmarking harness specifications

The key functionalities required by the harness are:

• it should have a python interface and
• it should distribute the benchmarking across nodes of a cluster/supercomputer.

Candidate solutions

Celery and Jolt came up as possible solutions, however we ended up trying only Ray and MPI4py.

Both Ray and MPI4py had some sort of implementations for Python futures, and this looked like a good way to implement the benchmarking harness. I opted for MPI4py since it is a better fit for the MPI-based HPC clusters we have access to.

The plan™ was to implement everything using Pythons concurrent.futures on my trusty little laptop, and then just swap out concurrent.futures with mpi4py.futures. But as it is often the case, life wasn’t so simple.

The cluster environment

Also, as it is often the case, the required software is often not available on the cluster. So, the first round of crocodile wrestling was compiling a bunch of libraries and getting them all to work together. That was a pain in the neck, but doable.

The wrench in the gears: Figuring out (how to invoke) MPI4py when we want to use futures

After figuring Tadashi’s dependencies, the time came to “just swap out”™ futures (and Executor) from Python’s concurrent with MPI4py’s implementation.

The right MPI, which supports MPI_Comm_spawn

Something which came up earlier in the development of the benchmarking harness was Fugaku’s support for master-worker jobs/workloads, which uses MPI_Comm_spawn to dynamically spawn processes, and incidentally, MPI4py futures, more precisely the MPIPoolExecutor, also uses MPI_Comm_spawn under the hood. So it was a bit disappointing when I realised OpenMPI doesn’t support MPI_Comm_spawn. However, MPICH, which does support MPI_Comm_spawn, was also available on the cluster and I just needed to recompile MPI4py with MPICH loaded to use it.

Testing went well

Initially, it was a bit hard to wrap my head around how MPI_Comm_spawn works, in my head MPI_COMM_WORLD is everything MPI is/can be aware of, but it turns out, if you have an allocation larger then MPI_COMM_WORLD, MPI still knows about it. This means, if you have an allocation of 10 nodes, you don’t lunch your master/parent program with mpirun -n 10 but with mpirun -n 1 and it will spawn processes on the remaining 9 nodes. So I logged in the cluster, got an interactive node, copy-pasted some example code for MPI4py spawn, and tested it – everything looked fine.

Unwanted behaviour & back to the basics

However, when I swapped concurrent.futures with mpi4py.futures, and wrote a submission script (to be launched by sbatch), things didn’t quite work. First, I realised the state of Tadashi which is in the binary .so files did not get pickled and transferred to the workers. After, temporarily disabling .so dependent code, I tried to rerun things, which did not fail!

However, when, for some reason I remembered to check if the workers are actually being executed on different nodes, it turned out this is not the case: when checking gethostname both master and workers were executed on the same node (and the other allocated nodes remained idle).

3 ways to run thing in Slurm, and finding what threw the wrench in the gears

Ultimately, the proverbial wrench in the gears, (aka bug, aka WTF) came down to the different ways you can launch programs with Slurm: using srun, salloc and sbatch.

srun allocates you resources from a cluster, and runs your binary (I like to think about srun as mpirun, but it is “aware” of the resources). salloc doesn’t run the program, it just allocates resources, and if a command is provided it executes that command only once, i.e. not on all nodes. To utilise all nodes within an allocation obtained by salloc, one would call srun. Finally, sbatch is like srun but instead of getting the allocation and running it immediately (dumping stdout to the terminal), sbatch puts the job/command in the queue, and saves the output into a slurm-*.out file.

Getting it right

To get to the bottom of things, I ended up writing a (pair of) simple MPI programs, spawn_main.c and spawn_child.c, each reporting the hostname. And again I made the mistake of running things from an interactive node, which I obtained using srun -N 3 -p genoa --pty bash. From thin interactive instance, running mpirun -N 1 ./spawn_main gave the desired results: the main and child processes were all reporting different hostnames.

The moment of clarity came when I wanted to present the full example, and wrote a spawn_submit.sh submission script, which I launched with sbatch. Lo and behold, I was back to the undesired behaviour of both main and child instances reporting the same hostname! This meant, that calling mpirun -N 1 ./spawn_main didn’t do the same thing when called from a submission script and when called from an interactive session.

I tried emulating the interactive session by running mpirun inside bash inside srun, i.e.

srun bash -c 'mpirun -N 1 ./spawn_main`

The results was a new kinda of undesired behaviour. Now mpirun saw all the allocated nodes and processes were reporting different hostnames, but each hostname was printed 3 times. It seems srun executed mpirun 3x, but oddly enough the main process was always on the same node (i.e. not on all 3 nodes for the 3x execution of srun).

I obtained the first working solution by adding another if to the monstrosity above which checked the “slurm rank”, i.e.

srun bash -c '[ $SLURM_NODEID = 0 ] && mpirun -N 1 ./spawn_main || true'

Calling mpirun inside srun already felt wrong, and the added complications didn’t improve the situation, so after some googling I fond an SO question asking about launching MPI_Comm_spawn from slurm, and copied the batch script parameters from there and it worked.

It turns out, the missing ingredient was a missing -n 3 for the job allocation. The final working solution looks like this:

#!/usr/bin/bash
#SBATCH -p genoa
#SBATCH -N 3
#SBATCH -n 3
#SBATCH -c 1

# ... snip ...

mpirun -N 1 spawn_main

-N is the number of nodes allocated, -n is the number of tasks (i.e. number of MPI ranks). The solution is valid without -c 1 (number of CPUs per task/rank), but I left it in just in case.

Currently, to run Tadashi, you need to compile some .so files first with CMake. It needs to be built into the build directory under the project root. Finally, adding the project root to PYTHONPATH will allow python to find both the python files and the binary .so files.

SWIG

First thing to clean up was implementing SWIG instead the ad-hoc CDLL approach currently used to call the C/C++ functions from Python. CDLL takes the path of the .so file, which is hard-coded in the Python files (hence the ad-hoc character of the implementation). CDLL also requires the arguments and return values of the functions exposed to Python from the .so files need to be copied from the .h files manually.

The SWIG implementation has the following advantages:

• The arguments and return values of the functions are automagically generated based on the .h files, thus letting us have a single source of truth (well, technically two, since the function prototypes in .h and .cc need to be synced manually, but the compiler catches any discrepancies).
• SWIG generates the wrapper Python file in the same directory as the .so file (as part of the build process), eliminating the need to manually specify the path to the .so.

Building a Python package

To the best of my knowledge, the way to do python packages is to write a pyproject.toml file. However, pyproject.toml does not support build extensions and they must be configured in setup.py. The cmake-build-extension build extension invokes CMake as part of the build process of the wheel.

TODO

• double check

• Python package
• cmake-build-extension

The new problems

Lightbox2

PhotoSwipe

Spotlight JS

Venobox

Let’s try to derive those formulas.

It is easy to spot the \(S_{j-1} := \sum_{t=1}^{j-1} t = \frac{j (j - 1)}{2}\) formula for the sum of all positive integers going up to \(j-1\). Which leads us to:

The same \(S_{j-1}\) formula is also present in the expression for \(j\) (we begin by removing the \(\lfloor \cdot \rfloor\)):

According to (Dash et al., 2021; Al Hajri et al., 2020) this flattened \(\lambda\) loop corresponds to the following \(i\) and \(j\) loop:

N = 5
count = 0
for i in range(N):
    for j in range(i+1, N):
        count += 1
        print(count, ":", i,j)

1 : 0 1
2 : 0 2
3 : 0 3
4 : 0 4
5 : 1 2
6 : 1 3
7 : 1 4
8 : 2 3
9 : 2 4
10 : 3 4

If we implement the flattened loop we see that this is only true in the sense that the set of visited combination is the same, however, the order is different.

import math

N = 5
Nc2 = 10
for L in range(Nc2):
    j = math.floor(math.sqrt(0.25 + 2*L) + 0.5)
    i = L - j*(j-1)//2
    print(L, ":", i, j)

0 : 0 1
1 : 0 2
2 : 1 2
3 : 0 3
4 : 1 3
5 : 2 3
6 : 0 4
7 : 1 4
8 : 2 4
9 : 3 4

So to generate the combinations in the same order as the initial \(i\), \(j\) loops, we need to modify the code as follows:

import math

N = 5
Nc2 = 10
for L in reversed(range(Nc2)):
    j = math.floor(math.sqrt(0.25 + 2*L) + 0.5)
    i = L - j*(j-1)//2
    print(L, ":", N - 1 - j, N - 1 - i)

9 : 0 1
8 : 0 2
7 : 0 3
6 : 0 4
5 : 1 2
4 : 1 3
3 : 1 4
2 : 2 3
1 : 2 4
0 : 3 4

Or alternatively, if we want to modify the original \(i\), \(j\) loop to match the \(\lambda\) loop and the mathematical derivation:

N = 5
count = 0
for j in range(N):
    for i in range(0, j):
        count += 1
        print(count, ":", i, j)

1 : 0 1
2 : 0 2
3 : 1 2
4 : 0 3
5 : 1 3
6 : 2 3
7 : 0 4
8 : 1 4
9 : 2 4
10 : 3 4

Graphical representation

Because \(\lambda \mapsto j(\lambda) = \lfloor \sqrt{1/4 + 2 \lambda} + 1/2 \rfloor\) is monotonically increasing (non-decreasing), returning the \(\lfloor \cdot \rfloor\) which we skipped in the calculations above, means that \(j\) is the largest possible integer such that \(\lambda = i + \frac{j(j - 1)}{2}\) for a non-negative integer \(i\). In the figure we can read \(j=4\) since the blue dots represent \(\sum_{t=1}^{j-1} t\) and the red dots show the calculation of \(i = \lambda - \sum_{t=1}^{j-1} t\).

Overview

We will need 3 things:

• The app/benchmark itself, which we will assume to be given.
• We need to set up a “self-hosted runner”, the program running on the supercomputer, which will be executing the actions such as compiling the app and submitting job to the scheduler of the supercomputer.
• The yaml file which describes when and what should be executed by the “runner”.

I’ll be describing the scenario for GitHub. Other sites like GitLab have similar systems.

Self-hosted runner

In the Github repo settings, on the left-hand side, under Actions there is a Runners page. In the top-right corner there is a green New self-hosted runner button. Clicking on this button brings up a page where you can select the OS and architecture. For supercomputers Linux, x64 is a good choice since usually that is something that will run on the login node.

Below the OS and architecture choice, the page lists the commands needed to install the self-hosted runner. This consists of several sections.

The first section described how to download, validate and extract the runner software. Don’t use these instructions, use the ones from the GitHub settings page.

# Create a folder
$ mkdir actions-runner && cd actions-runner
# Download the latest runner package
$ curl -o actions-runner-linux-x64-2.291.1.tar.gz -L https://github.com/actions/runner/releases/download/v2.291.1/actions-runner-linux-x64-2.291.1.tar.gz
# Optional: Validate the hash
$ echo "1bde3f2baf514adda5f8cf2ce531edd2f6be52ed84b9b6733bf43006d36dcd4c  actions-runner-linux-x64-2.291.1.tar.gz" | shasum -a 256 -c
# Extract the installer
$ tar xzf ./actions-runner-linux-x64-2.291.1.tar.gz

The second section describes how to configure and run the self-hosted runner. Again don’t use these instructions, use the ones provided on the settings page, since the --url and the --token are dependent on the repo you want to add the runner to. The ./config.sh asks a few questions, but generally it is very simple and usually the default answers are acceptable. The last command ./run.sh is the runner itself, it connects to GitHub, and needs to be running to be able to accept workflows/jobs. See note about security.

# Create the runner and start the configuration experience
$ ./config.sh --url https://github.com/<user>/<repo> --token <token>
# Last step, run it!
$ ./run.sh

The third section describes how to enable the runner in the Yaml file which is described in the next section.

# Use this YAML in your workflow file for each job
runs-on: self-hosted

This is the procedure to add a self-hosted runner to a repo. To the best of my knowledge, self-hosted runners can be added to GitHub users or GitHub organisations.

Yaml workflow file

To automatically run commands, we need to create a <name>.yml file in the <repo>/.github/workflows/ directory, for example with the following contents:

# .github/workflows/build-and-submit.yml
name: Build and submit
on: push
jobs:
  build:
    name: Build
    runs-on: [self-hosted,login-node]
    steps:
    - uses: actions/checkout@master
    - name: Create build dir
      run: mkdir build
    - name: Run cmake
      working-directory: ./build
      run: CXX=FCCpx cmake ..
    - name: Build
      run: cmake --build build --clean-first
    - name: Submit
      run: pjsub -g $(stat . -c %G) sub.sh

Each <name>.yml file (which can have any name) describes a workflow, with its name: (which can be any string), and the event when it will be executed. The example above will be exectuted on: push.

Each workflow consists of one or more jobs:. Multiple jobs are, by default, executed in parallel. In the example, for simplicity, there is only one job, with the custom identifier build: (this can be a different identifier e.g. job1:). Each job has a name: (similarly to a workflow), and each job needs to specify where it is should run using the runs-on: value. Without self-hosted runners, we can specify here a docker image (something like ubuntu-20.04), but in our case [self-hosted,login-node] specifies that the job should be executed on a self-hosted runner. The login-node is custom label which can be added to the runner on GitHub.

The main part of a job is the steps: field, which describes a list of steps which are executed sequentially. The job in the example has 5 steps. The first step is an “external” step (like importing a library), which checks out the master branch of the repository. The second, third and fourth steps create a build directory, call cmake in that directory (using the working-directory:), and builds the app using cmake --build. Finally, the last step, calls the command of the supercomputer scheduler to submit the sub.sh script.

Observing the actions

The top bar of a GitHub repository has an “Actions” page.

This page lists the workflows which were executed for the given repository. Clicking on a workflow, brings up a list of jobs defined for that workflow, and clicking on a job brings up the steps of that job. Clicking on a step expands it and displays the

Security

This is obviously a security issue. The runner script ./run.sh should be running all the time, connected to GitHub.com, waiting for jobs. As stated on GitHub, this should be enabled only for private repositories.

Future work

Next, I’d like to figure out how to write a workflow or a job which monitors when the submitted script finishes.

This blog posts is the first in a series of posts about polyhedral compilation, a mathematical model used to describe and reason about certain types of loops, with the aim to generate faster code.

This post revisits “Some efficient solutions to the affine scheduling problem. I. One-dimensional time” by Paul Feautrier, the seminal paper of the field, which describes how to formulate the the search for an optimal schedule as an integer linear programming (ILP) problem.

Overview of the process

Formulated as a source-to-source compilation, the following steps give a (very simplified) overview of the entire process:

• The input is source code with “nice” loops (where “nice” means that the loops satisfy such properties, that they are simple enough to be handled by ILPs).
• Problem/step 1: Finding the “nice” loops in the source code. This is handled by Polyhedral Extraction Tool (PET) which extracts affine description of the source code into ISL objects (named integer tuple sets/relations etc.). The loops in the source code can be marked with scop and endscop #pragmas or PET also has an auto-detect feature.
• Problem/step 2: Find or approximate the dependencies in the code.
• Problem/step 3: Formulate an ILP, which describes the statements from in step 1 and the dependencies from step 2. Given an optimisation objective, the ILP can be solved to find an (optimal) schedule.
• Problem/step 4: Based on the schedule obtained in the previous step generate (improved) source code
• The output is a source code with optimised loops.

This post only addresses (the first half) of Problem/step 3.

Example code: matrix vector product

for (i = 0; i <= n; i++) {
S1: a[i] = 0.0;
    for (j = 0; j <= n; j++)
S2:   a[i] += b[j] * M[i][j];
}

The above code has two relevant statements which access the memory: a[i] = 0.0; labelled as \(S_1\) and a[i] += b[j] * M[i][j]; labelled as \(S_2\). Each of the two statements is executed multiple times, it has multiple instances, for example the instances of statement \(S_1\) are:

• a[0] = 0.0; for \(i = 0\),
• a[1] = 0.0; for \(i = 1\) etc.

Since instances may need to be described by multiple loop variables, we adopt the notation \(\vec{i}\) for vectors in the iteration space, vectors with integer entries, such that the first element corresponds to the outermost and the last to the innermost loop variable. For example

• a[0] += b[1] * M[0][1]; for \(\vec{i} = (i, j) = (0, 1)\) and
• a[2] += b[3] * M[2][3]; for \(\vec{i} = (i, j) = (2, 3)\).

Describing dependencies

Generalised Dependency Graph (GDG)

Vertices/domains

For each statement \(S\) the corresponding vertex of the GDG is labelled with the domain (hence the \(\mathscr{D}\) notation below) of the statement \(S\), i.e. the subset of the iteration space containing the instances of \(S\) executed by the loop.

• \(\mathscr{D}_1 = \\{ i : 0 \le i \le n \\}\) for statement \(S_1\)
• \(\mathscr{D}_2 = \\{ (i, j) : 0 \le i, j \le n \\}\) for statement \(S_2\)

Technically, the domains are not sets, but families of sets, depending on parameters (in this example on the single parameter \(n\)), so the domain for statement 1 is the map \(n \mapsto \\{ i : 0 \le i \le n \\}\), but we omit the “\(n \mapsto\)” part, and treat \(n\) as a constant (but this will be included in a more).

Edges/dependencies

The edges of GDG are the dependencies between two statements and are labelled with a subset of the direct product (a relation, hence the \(\mathscr{R}\) notation below) between the two domains of statements of the start and end of the edge, that is, if \(S'\) and \(S\) are two statements and there is a dependency between the instances \(\vec{i'} \in \mathscr{D}_ {S'}\) and \(\vec{i} \in \mathscr{D}_ S\) then there is and edge from vertex \(\mathscr{D}_{S'}\) to \(\mathscr{D} _S\) labelled with a set that contains \((\vec{i'}, \vec{i})\).

A simplified (ergo very conservative) dependency analysis (there are programs which can perform such analysis) could yield two dependencies:

• \(\mathscr{R}_{1, 2} = \{ \bigl( i', (i, j) \bigr) : i' = i \}\) describes the dependency between \(S_1\) and \(S_2\) which requires for the initialisation in \(S_1\) (a[i] = 0.0) to precede (all instances of) statement \(S_2\) when the two statements share the same value for the loop variable \(i\) (hence \(i' = i\)).
• \(\mathscr{R}_{2, 2} = \{ \bigl( (i', j'), (i, j) : i' = i \land j' < j \}\) describes the dependency of \(S_2\) on itself, which requires, for a given \(i\) (\(i' = i\)) the earlier instances of (in \(j\)) are executed before the later instances (hence \(j' < j\)).

This dependency analysis is very coarse and/or conservative (read poor), we’ll discuss a simple data flow dependency later (which is still quite simple, but a slight improvement over the one above).

Detailed Dependency Graph (DDG)

The GDG is structured: the vertices in GDG are statements, and these statements represent multiple instances, but we actually care about the dependencies between the instances. For this reason the Detailed Dependency Graph “flattens” the graph, and every vertex is an instance of a statement, and the edges are the dependencies between these instances.

Vertices

Edges

where the statement \(\sigma(e)\) is the start, statement \(\delta(e)\) is the end of edge \(e\) (of the GDG).

Schedule

The schedule is a map \(\theta: \Omega \to \mathbb{R}_0^+\) from the set of instances to some non-negative value which is the “date” (or timestamp, or time) of the instance.

Generating code

As mentioned above, generating code is a separate, and very much non-trivial problem. But to get a better feeling how to interpret the schedule \(\theta\) a simplified code generations is presented:

Let \(\mathtt{F}(t) = \\{ (S, \vec{i}) \in \Omega: \theta(S, \vec{i}) = t \\}\), i.e. the set of all instances of all statements which should be executed at time step \(t\). Let \(\mathtt{L} = \max_{(S, \vec{i}) \in \Omega} \theta(S, \vec{i})\).

for (t = 0; t <= L; t++) {
  #pragma omp parallel
  for (inst : F(t))
    execute(inst);
  barrier();
}

Of course, actual code generation is a much harder task than this naive pseudo-code, but it can be handled separately, the objective of this now is how to obtain the optimal schedule.

There is (no) optimal schedule

The paper cites Theorems which say that finding a schedule of arbitrary form is an undecidable problem. Because of this, we restrict ourselves to affine schedules, that is schedules of the form: \(\theta(S, \vec{i}) = \tau_S \vec{i} + \sigma_s \vec{n} + \alpha_s\) for each statement \(S\). The vector \(\vec{n}\) is the vector of parameters, for the example above the vector of length 1 containing \(n\). In this case the triplet \((\tau_S, \sigma_S, \alpha_S)\) completely define \(\theta\) (for a given \(S\)), so the goal is finding a \((\tau_S, \sigma_S, \alpha_S)\) triplet for each statement \(S\).

More advanced dependency analysis

Depth

Descriptions such as GDG and DDG can enable some optimisations.

The depth of an edge is the position until which both instances at the start and the end of the edge share values, and after which the end instance has a larger value, that is \(p_ e\) is the depth of edge \(e\) iff \((\vec{i'}, \vec{i}) \in \mathscr{R} _ e\) and \(i'_ k = i_ k\) for \(1 \le k \le p_ e\) and \(i'_ {p_ e} < i_ {p_ e}\) where \(\vec{i'} = (i'_ 1, i'_ 2, \ldots)\) and \(\vec{i} = (i_ 1, i_2, \ldots)\).

In the example, both edges of the GDG have depth 1:

• \(\mathscr{R}_{1, 2} = \{ \bigl( i', (i, j) \bigr) : i' = i \}\),
• \(\mathscr{R}_{2, 2} = \{ \bigl( (i', j'), (i, j) : i' = i \land j' < j \}\).

In both cases the \(i'=i\) part implies depth \(p_e \ge 1\) and the rest ensures \(p _e \le 1\).

This can be used to infer, that we are allowed to execute the outermost loop in parallel.

Dependence direction vectors

A more detailed description of the dependencies can be given using symbols such as \(<, \le, =, *, \ldots\) combined in a dependence direction vector (the asterisk denotes a wildcard, meaning any relation). Depth can be expressed with DDVs as

Uniform dependence

The case where there is a constant difference between the instances of both ends of an edge, that is when \(i' = i + d\) if \((i', i) \in \mathscr{R}_ e\), the edge \(e\) is said to have a uniform dependence. In this case, instead of keeping track of \(\mathscr{D}_ {\sigma(e)}\), \(\mathscr{D}_ {\delta(e)}\) and the set of \((\vec{i'}, \vec{i})\) pairs, we can just keep track of a single set (polyhedron) of instances \(\mathscr{P}_ e\) and a affine map \(h_ e\) such that \(y \in \mathscr{P}_ e \implies y \in \mathscr{D}_ {\delta(e)} \land h_e(y) \in \mathscr{D} _{\sigma(e)}\) and then

A more detailed analysis shows that the second edge of our example has such a uniform dependency.

Dataflow analysis

A little more advanced (but still very much conservative) dataflow analysis can further restrict the polyhedrons \(\mathscr{R} _{1, 2}\) and \(\mathscr{R} _{2, 2}\). The analysis of the memory reads and writes tells us that only the entries of a[i] updated, they are updated independently for each index \(i\), and making no assumptions about the + operation (such as associativity, which could be used for further optimisations), we observe that

• a[i] is initialised in statement \(S_ 1\) and only the first iteration of the \(j\) loop depends on it: \(\bigl(i', (i, j) \bigr) \in \mathscr{R}_ {1,2} \iff i' = i \land j = 0\) (I think there is a typo in the paper saying \(j = 1\)?). This is reduced as:

  \[\mathscr{P}_ {e _1} = \mathscr{D}_2 \cap \{ (i, j) : j \le 0 \}, \quad h _{e _1}(i, j) = i\]

• a[i] is updated with each iteration of \(j\), so every iteration (instance) of \(j\) depends only on the previous iteration (\(j - 1\)), and this only applies starting from the second iteration (\(j \ge 1\)): \(\bigl( (i', j'), (i, j) \bigr) \in \mathscr{R} _{2,2} \iff i' = i \land j' = j - 1 \land j \ge 1\) (Again, this might be a typo \(j \ge 2\) in the paper?) This is reduced as:

  \[\mathscr{P} _{e _2} = \mathscr{D} _2 \cap \{ (i, j) : j \ge 1 \}, \quad h _{e _2}(i, j) = (i, j - 1)\]

We will continue with these reduced forms.

Formulating the integer linear program

Describing vertices/domains

The \(\mathscr{D}_ S\) domains (including the parameters, represented as \(\vec{n}\)) need to be rewritten in the form where given the parameters \(\vec{n}\) the instance \(\vec{i}\) is in domaind \(\mathscr{D} _S\) iff:

This way, the \((a_ {S_ k}, b_ {S_ k})\) pairs completely describe \(\mathscr{D} _S\) (that is, you can use these vectors to represent them in a computer program).

In the example of \(\mathscr{D} _1\) there are two inequalities, implying \(m _1 = 2\):

implies \(a _{S _1} = (1, 0)\) and \(b _{S _1} = 0\) and

implies \(a _{S _2} = (-1, 1)\) and \(b _{S _2} = 0\).

Domain \(\mathscr{D} _2\) can be described with \(m _2 = 4\) such equations.

Describing edges/dependencies

The edges \(\mathscr{R}_ e\) of the GDG is described by \((c_e, d_e)\) such that:

or for a restricted schedule with the affine map \(\vec{i'} = h_e(\vec{i})\) and the rewritten reduced domain \(\mathscr{P} _e\):

The reduced domains \(\mathscr{P} _{e _1}\) and \(\mathscr{P} _{ e _2}\) can be described similarly as the other domains \(\mathscr{D} _1\) and \(\mathscr{D} _2\).

Describing schedules

The schedule \(\theta(S, \vec{i})\) is also going to be rewritten using a set of \(\mu\) Farkas multipliers. For each statement \(S\) we assume that the schedule can be expressed as:

This captures the information provided by the domains \(\mathscr{D} _S\) captured in the (\(m _S\) number of) \((a _{S _k}, b _{S _k})\) pairs. To combine this with the information from the dependencies/edges we will need the delay corresponding to the edges.

The delay

We assume that if the instance \(\vec{i}\) of a statement \(S\) depends on the instance \(\vec{i'}\) of the statement \(S'\), then there is a delay \(\Delta\) associated with that dependency/edge \(e\). This means that the date of \(S, \vec{i}\) assigned by the schedule \(\theta\) is greater (by at least \(1\)) than the date of \(S', \vec{i'}\):

We assume that this delay can be rewritten with a different set of \(\lambda\) Farkas multipliers (these will be just placeholders to express dependencies between inequalities across inequalities resulting from the dependencies/edges).

Putting it all together

The \(\equiv\) in the last equation was alluding to the next step where we combine the “\(\theta\) equations” expressing the domains and the “\(\Delta\) equations” expressing the dependencies.

On the left side of \(\equiv\) in the expression above we use two instances of the “\(\theta\) equations” (with \(a _S{ _k}\), \(b _S{ _k}\) and \(\mu _S{ _k}\)), on the right “\(\Delta\) equations” (with \(c _{e _k}\), \(d _{e _k}\) and \(\lambda _{e _k}\)) and solve the ILP for the \(\mu _{S _k}\) variables (for each statement \(S\)).

Edge \(e_1 : 1 \to 2\)

For the first edge \(e _1\) between statement \(S_1\) to \(S_2\) the equations from above give rise to the following

The first and second line (except the \(-1\) at the end of it) of the ILP come from the rewritten form of \(\mathscr{D}_2\) and \(\mathscr{D}_1\) from the Describing vertices/domains section, plugged in the “\(\Theta\) equation”, while the third line is the result of taking \(\mathscr{P} _{e _1}\) Dataflow analysis, which is \(-j \ge 0\) and the inequalities from the \(\mathscr{D} _2\) (hence the similarity to the first line).

The previous equation is equivalent to the following system of equations by equating the coefficients of \(i\), \(j\), \(n\) and the constant term.

Edge \(e_2 : 2 \to 2\)

The second edge is a uniform dependency, the schedule for the start and end of the edge, \(\theta(S _2, h(\vec{i}))\) and \(\theta(S _2, \vec{i})\) is nearly identical (difference highlighted in the formulae below).

This results to most of the terms cancelling each other out in the expression \(\theta(S _2, \vec{i}) - \theta(S _1, h(\vec{i}))\) (written with the \(\mu _{S _k}\) Farkas multipliers):

As a result, the loop edge on \(S _2\) results in the following equation (not the lack of \(\lambda _{S _k}\) multipliers).

The calculations

Collecting and rearranging the inequalities for \(e _1 : S _1 \to S _2\) and \(e _2 : S _2 \to S _2\).

Simplifying it gives:

All these manipulations can be performed by algorithms automatically.

One possible result

One valid choice for the \(\mu _{S _k}\) values is:

• \[\mu_{1, 0} = \mu_{2, 1} = \mu_{2, 2} = \mu_{2, 4} = \mu_{1, 1} = \mu_{1, 2} = 0\]
• \[\mu_{2, 0} = \mu_{2, 3} = 1\]
• \[\theta(1, i) = 0\]
• \[\theta(2, i, j) = j + 1\]

Generated code

The resulting schedule is:

• \[\theta(S _1, i) = 0\]
• \[\theta(S _2, i, j) = j + 1\]

Generating code from this is a separate task and will be disucussed in the next blog post, but the paper suggests something similar to:

#pragma omp parallel
for (i = 0; i <= n; n++)
  a[i] = 0.0;
for (j = 0; j <= n; j++)
  #pragma omp parallel
  for (i = 0; i <= n; i++)
    a[i] += b[j] * M[i][j];

Citing this blog post

@misc{vatai2022polytutor1,
  title={Polyhedral compilation: part 1},
  url={https://vatai.github.io/math/compsci/polyhedral/polyhedral-compilation-part-1/},
  author={Vatai, Emil},
  year={2022},
  month={Feb}
}

Feedback

Feedback is very much welcome. I don’t have a comment section set up, but you can raise an issue on GitHub.

On ABCI to get this working, you need to load these modules (some of them might be not needed, I just grabbed a modules.sh file):

module load gcc/9.3.0
module load cuda/11.2/11.2.2
module load cudnn/8.1/8.1.1
module load nccl/2.8/2.8.4-1
module load openmpi/4.0.5
module load python/3.8/3.8.7
module load cmake/3.19

After this we just need to clone the PyTorch repo:

git clone git@github.com:pytorch/pytorch.git

and build it:

python3 setup.py develop --user

This overwrites your current PyTorch installation, and you need to use --upgrade --forece-reinstall with pip3 to install the original one.