Nacho

Sparse Tensor Algebra compiler for GPUs and multicore CPUs

Setup

Install first: the CUDA toolkit, oneTBB, Python 3.12 or newer, CMake, and LaTeX, which the benchmark scripts typeset their figures with. We use the following conda environment:

conda create -n nacho_env python=3.12
conda activate nacho_env
conda install conda-forge::tbb-devel conda-forge::cmake nvidia::cuda-toolkit=12.6 texlive-core ghostscript

If you already have your own CUDA toolkit, nvcc has to be findable before configuring, either on PATH or named by CUDACXX:

export PATH=/usr/local/cuda-12.6/bin:$PATH   # or: export CUDACXX=/usr/local/cuda-12.6/bin/nvcc

Then:

pip install -r requirements.txt
pip install torch --index-url https://download.pytorch.org/whl/cu126

Replace cu126 to match the CUDA version on your machine.

Tested with CUDA toolkit 12.6, driver 560.35, gcc 11.4, Python 3.12 and torch 2.8.0+cu128 on an RTX 4090. python/CMakeLists.txt builds for sm_89; override with -DCMAKE_CUDA_ARCHITECTURES.

1. Build nacho:

# Option 1: normal
cmake -S . -B build
cmake --build build -j<N PARALLELISM>

# Option 2: debug
cmake -S . -B build-dbg -DCMAKE_BUILD_TYPE=Debug
cmake --build build-dbg --config Debug -j<N PARALLELISM>

2. Generate kernels into generated/:

./build/compiler                          # every kernel declared in compiler.cpp
./build/compiler --kernels csr_mul,csr_add  # just these

Python bindings

Generated kernels, and the tensor classes they take and return, are exposed to Python through nanobind. The nanobind bindings are also generated by the nacho compiler. A kernel is declared in compiler.cpp — consisting of the tensor formats, the operands over them, and the expression:

Format csr = Format::ordered({
    {"i", LevelFormat::Dense},
    {"j", LevelFormat::Compressed_unique},
}).named("CSR");

TensorType csr_f32 = TensorType(csr, dType::Float32);
Expr a_csr_ij = Tensor::make(csr_f32, "a");
Expr b_csr_ij = Tensor::make(csr_f32, "b");

Kernel("csr_mul").expr(a_csr_ij * b_csr_ij).targets({Target::CPU, Target::GPU}).emit();

.named("CSR") is the Python class name of the layout, giving nacho.CSR_cpu and nacho.CSR_gpu; the kernel is exposed as <cpu|gpu>_<name>_f32. Operands are passed in lexicographic order by name unless .operand_ordering({"b", "a"}) says otherwise.

Before the generated code can be used from Python, it has to be compiled into the extension module:

pip install --no-build-isolation -ve .

import nacho
A, B = nacho.to_csr(A_torch, "cpu"), nacho.to_csr(B_torch, "cpu")
C = nacho.cpu_csr_mul_f32(A, B)          # A, B are nacho.CSR_cpu
C = nacho.gpu_csr_mul_f32(A, B)          # nacho.CSR_gpu, grid size can be optionally provided

Testing

tests/smoke.py checks each generated kernel on both devices against scipy on small random inputs; run it after rebuilding and before any timing work.

python tests/smoke.py

Benchmarks

The benchmarks read their inputs from two datasets, neither of which lives in this repo. Set both before running anything:

export NACHO_SUITESPARSE_DIR=/path/to/suitesparse   # .mtx files, every 2-D benchmark
export NACHO_FROSTT_DIR=/path/to/frostt             # .tns files, frostt_tensors_add.py

Fetching the datasets

benchmarks/datasets/ holds the scripts that populate those two directories. They write to the paths above, so export the variables first.

SuiteSparse dataset

cd "$NACHO_SUITESPARSE_DIR"
bash <nacho>/benchmarks/datasets/download_suitesparse_reals.sh
find . -name '*.tar.gz' -print0 | xargs -0 -P 8 -n 1 tar -xzf

Frostt Tensor Dataset

python benchmarks/datasets/download_frostt.py --all
python benchmarks/datasets/sort_frostt.py --all

The second script dedups and sorts the downloaded frostt tensors, which the kernels require.

Running

Each script directly in benchmarks/ is a standalone entry point. The following script can be used to run all the benchmarks.

numactl --physcpubind 0-15 python benchmarks/csr_add.py  --device both --start 0 --end 1600     # vs cuSPARSE/Taco/MKL
numactl --physcpubind 0-15 python benchmarks/csr_mul.py     --device both --start 0 --end 1600  # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/coo_add.py     --device cuda --start 0 --end 1600  # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/coo_mul.py     --device cuda --start 0 --end 1600  # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/coo_csr_add.py --device both --start 0 --end 1600  # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/csr_add_3.py   --device cuda --start 0 --end 1600  # fused vs unfused
numactl --physcpubind 0-15 python benchmarks/heatmap_dcsr_mul.py --device both                  # compare partitioning schemes
numactl --physcpubind 0-15 python benchmarks/heatmap_csr_add.py --device both                   # skew sweep, synthetic

numactl --physcpubind 0-15 python benchmarks/spgemm.py --start 0 --end 1300                     # vs cuSPARSE
numactl --physcpubind 0-15 python benchmarks/sssmm.py  --start 0 --end 1300                     # fused vs unfused vs cuSPARSE
numactl --physcpubind 0-15 python benchmarks/frostt_tensors_add.py --device both                # CSF3 vs COO3D vs torch
numactl --physcpubind 0-15 python benchmarks/inner_prod.py  --device both                       # FROSTT, CSF3 vs COO3 vs torch

Settings in benchmarks/config.py — dataset paths, iteration counts, output directory — can be overridden by an environment variable of the same name prefixed with NACHO_.

If running cpu-only benchmarks, run

numactl --physcpubind 0-15 python benchmarks/csr_add.py  --device cpu --start 0 --end 1600      # vs Taco/MKL
numactl --physcpubind 0-15 python benchmarks/csr_mul.py     --device cpu --start 0 --end 1600   # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/coo_csr_add.py --device cpu --start 0 --end 1600   # vs PyTorch
numactl --physcpubind 0-15 python benchmarks/heatmap_dcsr_mul.py --device cpu                   # compare partitioning schemes
numactl --physcpubind 0-15 python benchmarks/heatmap_csr_add.py --device cpu                    # skew sweep, synthetic

numactl --physcpubind 0-15 python benchmarks/frostt_tensors_add.py --device cpu                 # CSF3 vs COO3D vs torch
numactl --physcpubind 0-15 python benchmarks/inner_prod.py  --device cpu                        # FROSTT, CSF3 vs COO3 vs torch

If running gpu-only benchmarks, run

python benchmarks/csr_add.py  --device cuda --start 0 --end 1600                                # vs cuSPARSE/Taco
python benchmarks/csr_mul.py     --device cuda --start 0 --end 1600                             # vs PyTorch
python benchmarks/coo_add.py     --device cuda --start 0 --end 1600                             # vs PyTorch
python benchmarks/coo_mul.py     --device cuda --start 0 --end 1600                             # vs PyTorch
python benchmarks/coo_csr_add.py --device cuda --start 0 --end 1600                             # vs PyTorch
python benchmarks/csr_add_3.py   --device cuda --start 0 --end 1600                             # fused vs unfused
python benchmarks/heatmap_dcsr_mul.py --device cuda                                             # compare partitioning schemes
python benchmarks/heatmap_csr_add.py --device cuda                                              # skew sweep, synthetic

python benchmarks/spgemm.py --start 0 --end 1300                                                # vs cuSPARSE
python benchmarks/sssmm.py  --start 0 --end 1300                                                # fused vs unfused vs cuSPARSE
python benchmarks/frostt_tensors_add.py --device cuda                                           # CSF3 vs COO3D vs torch
python benchmarks/inner_prod.py  --device cuda                                                  # FROSTT, CSF3 vs COO3 vs torch

Acknowledgements

A significant portion of the code in this repository is modeled after, or directly taken from, the Halide compiler. That is because they both have done incredible work, and because it is the compiler that we are most familiar with navigating and understanding. As a result, this repository benefits heavily from over a decade of hard work from the Halide developers.