Building
Wheel vs JIT runtime paths, build-time and runtime environment variables, CPU SIMD flags, and source layout.
Build vs runtime
torchmatch supports two runtime paths:
- Prebuilt: a
.socompiled for the right Python version and PyTorch build (ABI) ships in the wheel and loads viatorch.ops.load_libraryat first call. No compiler needed at install time. - JIT: when no matching prebuilt
.sois present (source-distribution install (sdist), ABI mismatch, or explicit override viaTORCHMATCH_FORCE_JIT=1), the loader compiles the C++/CUDA sources viatorch.utils.cpp_extension.load. The result caches in$TORCH_EXTENSIONS_DIR.
Both paths register the same torch.ops.assignment.* and torch.ops.transport.* ops; the choice is
transparent to callers.
Building wheels
The build system pairs setuptools with
torch.utils.cpp_extension.BuildExtension.
# default: CPU extension always; CUDA extension when a toolchain is found
# (torch.utils.cpp_extension.CUDA_HOME is not None)
pip wheel . -w dist/
# CPU-only wheel
TORCHMATCH_SKIP_CUDA=1 pip wheel . -w dist/
# CUDA GPU architecture targets (e.g., `8.6` for Ampere) passed to the `nvcc` compiler
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" pip wheel . -w dist/
Build-time environment variables
| Variable | Effect |
|---|---|
TORCHMATCH_SKIP_CUDA=1 | Skip the CUDA extensions; produce a CPU-only wheel |
TORCHMATCH_SKIP_CPU=1 | Skip the CPU extensions (rarely useful) |
TORCHMATCH_SKIP_TRANSPORT=1 | Skip both transport extensions (CPU and CUDA); produces an assignment-only wheel. The torchmatch.transport namespace still imports cleanly because the Sinkhorn-family ops are registered in pure Python; only EXACT_EMD becomes unavailable. |
TORCH_CUDA_ARCH_LIST | Comma-separated CUDA GPU architecture targets (e.g., 8.6 for Ampere) passed to the nvcc compiler |
When the NVIDIA CUDA compiler (nvcc) is absent — meaning torch.utils.cpp_extension.CUDA_HOME is None — the CUDA extensions are skipped silently, so building from an sdist on a
CPU-only host still produces a usable CPU wheel.
Runtime overrides
| Variable | Effect |
|---|---|
TORCHMATCH_FORCE_JIT=1 | Skip the prebuilt .so and recompile from source. Useful for diagnosing ABI mismatches or developing C++ changes. |
TORCH_EXTENSIONS_DIR | (PyTorch-builtin) Where the JIT cache lives. Defaults to ~/.cache/torch_extensions. |
CPU SIMD
The CPU extension builds with -O3 -std=c++17 -mavx2 -mfma by
default, enabling AVX2 and FMA vectorization for best performance on modern x86-64 CPUs.
On older x86 CPUs without AVX2, pip falls back to a source-distribution (sdist) install; the JIT compiler then queries torch.cpu._is_avx2_supported() at build time and omits the AVX2/FMA flags automatically.
The jonker_scalar op works without SIMD.
Source layout
sources/torchmatch/
├── __init__.py # eager-loads the active sub-packages
├── _loader.py # shared prebuilt + JIT plumbing
├── assignment/ # active sub-package
│ ├── __init__.py # public API: solve, Backend, ops, load_cpu/load_cuda
│ ├── _solve.py # dispatcher and Backend enum
│ ├── _greedy.py # pure-PyTorch Kurtzberg 1962 heuristic
│ ├── ops.py # direct handles to torch.ops.assignment.*
│ ├── _cpu.py / _cuda.py # extension loaders
│ ├── cpu/ # CPU sources (jonker_*.{h,cpp}, ops.cpp)
│ └── cuda/ # CUDA sources (munkres.cu, hybrid.cu, lawler.cu, jonker_tiled.cuh, ...)
└── transport/ # optimal transport (OT) sub-package
├── __init__.py # eager-loads matrix; samples is lazy
├── matrix/ # cost-matrix face
│ ├── __init__.py # public API: solve, Backend, ops, load_cpu/load_cuda
│ ├── _solve.py # dispatcher and Backend enum
│ ├── _log_sinkhorn.py / _sinkhorn_divergence.py / _unbalanced_sinkhorn.py
│ ├── _exact_emd.py # Python wrapper around the EMD C++ extension
│ ├── _validate.py / _schedule.py # input validation, epsilon annealing schedule
│ ├── ops.py # direct handles to torch.ops.transport.*
│ ├── _cpu.py / _cuda.py # extension loaders
│ ├── cpu/ # CPU sources (exact_emd_op.cpp, ops.cpp, exact/ network simplex sources)
│ └── cuda/ # CUDA sources (ops.cpp)
└── samples/ # point-cloud face (CUDA-only, Triton kernels)
├── __init__.py # public API: loss
├── _loss.py # entry point
├── _autograd.py # custom_op + register_autograd
├── _solvers.py / _c_transform.py / _cg.py / _hvp.py / _implicit_grad.py
└── kernels/ # Triton kernels (flashstyle_sqeuclid, ...)
setup.py declares up to four setuptools extensions:
torchmatch._assignment_cpu_impl and torchmatch._assignment_cuda_impl
(via CppExtension / CUDAExtension), plus the matching
_transport_cpu_impl and _transport_cuda_impl translation units.
CUDA extensions are included only when a CUDA toolchain is found and
TORCHMATCH_SKIP_CUDA is unset; transport extensions are skipped when
TORCHMATCH_SKIP_TRANSPORT=1.