Building CPU optimized packages for conda-forge

Modern CPUs have powerful features that can make your code run significantly faster. One of these features is called Single Instruction, Multiple Data (SIMD) instruction sets. A recent article released by the authors of the open-source physics library box2d showcases how using modern CPU instruction sets can make a huge difference in performance. The article achieves a 2x speedup specifically, by using something called "Single Instruction, Multiple Data" or SIMD instructions, we'll get to what this actually means in a bit.

Some examples of libraries specifically targeting these instructions and yielding a considerable improvement are¹:

1. OpenCV color conversion functionality, ~25x faster on ARM CPUs with Neon: opencv#19883

2. PyTorch softmax, min and max 3x-4x faster for bfloat16 with AVX2/AVX512 on x86-64: pytorch#55202, and up to 2x-10x with uint8 for +, >>, min: pytorch#89284

3. As mentioned in the introduction, ~2.5x faster speed-up of 2D collision checking.

While certain specialized libraries like NumPy and PyTorch have always made use of the full potential of your hardware by using dynamic dispatching, other libraries need to be compiled with the right flags to enable these optimizations. Conda-forge (like many other software distributions) aims to be as compatible with not-so-recent hardware as possible, so these optimizations are not enabled by default.

But recently, it became possible to target newer CPU instruction sets on conda-forge directly! Let's quickly cover what SIMD and related terms mean and go over the basics on how to build cpu-optimized packages using either conda-build or rattler-build.

What does SIMD/AVX/NEON even mean?

CPUs execute instructions to perform tasks, and modern CPUs support instruction sets allow enabling processing multiple data points in parallel—which is why the instruction set is called "Single Instruction, Multiple Data" or SIMD.

There are a number of different instruction sets that can be categorized as SIMD, the key ones are:

1. SSE (Streaming SIMD Extensions): An older instruction set that allows the CPU to perform the same operation on multiple data points at once.

2. AVX (Advanced Vector Extensions): A more advanced instruction set that extends SSE with more powerful operations for faster data processing.

3. Neon: An ARM-specific SIMD instruction set found in Apple Silicon, mobile and embedded devices.

Using these SIMD instructions can greatly improve code performance by reducing the number of instructions needed for data processing. Compilers can automatically leverage these SIMD instructions, but your CPU must support the specific sets. Some SIMD sets have been available for years, newer ones may not be supported on all hardware, particularly older devices.

Libraries like NumPy make extensive use of these instructions. These instructions can be enabled in the following ways:

• Runtime Selection: Code is compiled for multiple hardware targets, and the best version is chosen at runtime. This approach can boost performance but requires complex engineering and increases the package size.

• Just-In-Time Compilation: Libraries like Numba or Pythran can compile code at runtime in order to optimize the code for the specific hardware. This approach can be very powerful but requires additional (large!) dependencies.

• Installation Time Selection: The best compiled program is selected during installation, reducing complexity and package size while optimizing for the specific hardware. This approach is supported by the conda-forge ecosystem, simplifying the process and still optimizing performance.

How can I use it today?

If a package that you are maintaining or interested in is available on conda-forge, you can enable these optimizations. By adding these sections to the meta.yaml or recipe.yaml in the conda-forge feedstock, you can start making use of the optimizations today

Recipe v1 (recipe.yaml):

context:
  build: 0

build:
  # Prioritize builds with a higher microarch level.
  number: ${{ build|int + (microarch_level|int) * 100 }}

requirements:
  build:
    - if: microarch_level|int > 0
      then: x86_64-microarch-level ${{ microarch_level }}
    - ${{ compiler('c') }}

microarch_level:
  - if: not(x86_64)
    then:
      - 0
    else:
      - 1
      - 3
      - 4

Recipe v0 (meta.yaml):

{% set build = 0 %}

build:
  # Prioritize builds with a higher microarch level.
  number: {{ build }}          # [not x86_64]
  number: {{ build + 100 }}    # [x86_64 and microarch_level == 1]
  number: {{ build + 300 }}    # [x86_64 and microarch_level == 3]
  number: {{ build + 400 }}    # [x86_64 and microarch_level == 4]

requirements:
  build:
    - x86_64-microarch-level {{ microarch_level }}  # [x86_64]
    - {{ compiler('c') }}

microarch_level:  # [x86_64]
  - 1  # [x86_64]
  - 3  # [x86_64]
  - 4  # [x86_64]

What does the above code adds for both recipe.yaml and meta.yaml?

1. The multiple build numbers allows the solver to prioritize these variants if these are available. Newer architectures get a higher build number to prioritize those over older architectures. E.g AVX gets a higher build number than SSE.

2. A requirement is added on the microarch package that makes sure that the required compiler flags are set and the package will only run on hardware that supports it.

Also refer to the conda-forge knowledge base on this. This is all that's needed, to enable the users of the package to make use of the compiler optimizations. For a recently merged example in a real-life recipe see the following PR

How to use with Pixi?

Pixi always locks environments for a specific platform including capabilities. To be compatible with most systems Pixi defaults to lower capability systems. But it is possible to include a more capable platform for which Pixi should create an environment.

[workspace]
channels = ["conda-forge"]
name = "archspec"
platforms = [
  # The archspec definition here targets a more capable CPU architecture
  { platform = "linux-64", archspec = "x86_64_v3" },
  # This platform is used as the base for systems without `x86_64_v3`. 
  { platform = "linux-64" },
]

[dependencies]
box2d = "*"

Running pixi lock on this examples outputs something like:

❯ pixi lock
✔ Updated lock file
Environment: default
  + C _openmp_mutex                            4.5 20_gnu
  + C libgcc                                   16.2.0 ha9f2e26_4
  + C libgomp                                  16.2.0 he0feb66_4
  + C libstdcxx                                16.2.0 h934c35e_4
Platform: default:linux-64
  + C box2d                                    3.1.1 h3e4d06c_0
Platform: default:linux-64-archspec-x86-64-v3
  + C _x86_64-microarch-level                  3 5_level3
  + C box2d                                    3.1.1 hd97bcb3_1

The output shows that we get different variants of box2d for the different platforms. The variant for the linux-64-archspec-x86-64-v3 platform has been build with AVX2 support.

When running a Pixi command, the first platform that matches your machine is automatically selected. This provides you with the most capable packages that match your machine.

Conclusions

As you can see it is fairly straightforward! If any of the packages you maintain benefit from SIMD operations you might want to give this a try!

To recap, to enable SIMD optimizations in your conda-forge package:

1. Add the x86_64-microarch-level package as a build requirement.

2. Set the build number based on the microarch_level in the meta.yaml or recipe.yaml.

3. Add the microarch_level key to the conda_build_config.yaml or variants.yaml file.

As always, feel free to ask us any questions. You can join our Discord and have a chat about building your packages, reach us on X or follow projects on our GitHub.

The Footnotes

1. These numbers are partially taken from: pypackaging-native