Metadata-Version: 2.4
Name: torchmatch
Version: 1.0.0
Summary: Linear assignment problem solvers for PyTorch.
Author-email: Kurt Stolle <kurt@computer.org>
License-Expression: MIT
Keywords: pytorch,linear assignment problem,lap,lapjv,hungarian algorithm,bipartite matching,cuda
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.13
Classifier: Operating System :: POSIX :: Linux
Requires-Python: <3.14,>=3.13
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.11
Provides-Extra: cpu
Requires-Dist: torch>=2.11; extra == "cpu"
Provides-Extra: cu126
Requires-Dist: torch>=2.11; extra == "cu126"
Provides-Extra: cu128
Requires-Dist: torch>=2.11; extra == "cu128"
Provides-Extra: cu130
Requires-Dist: torch>=2.12; extra == "cu130"
Dynamic: license-file

# torchmatch

Linear assignment problem (LAP) solvers for PyTorch, registered under
`torch.ops.matching`. Inputs are `torch.Tensor`; outputs are integer
row-to-column assignments.

📖 **Documentation**: <https://khwstolle.github.io/torchmatch/>. Tutorials, API reference, and the full benchmark report.

Two algorithmic families share the op surface:

- **Munkres** (Kuhn-Munkres 1957): the primed/starred-zeros 6-step
  augmenting-path procedure. CUDA only; three parallelization
  strategies.
- **Jonker-Volgenant / successive-shortest-path (SSP)**: the
  Dijkstra-like family.

Both families return optimal assignments. They differ in the
augmentation routine, the parallelization story, and constant
factors.

## Install

```bash
pip install torchmatch
```

The Munkres CUDA ops and the CUDA backend of `jv_dense_batch` need a
CUDA-capable PyTorch build. CPU ops run on any platform; the
SIMD-optimized variants need AVX2/FMA at build time (sdist falls back
to the AVX2/FMA-aware JIT path, and `jv_scalar` works without SIMD).

## Usage

```python
import torch
import torchmatch                          # extensions load eagerly at import

cost = torch.randn(8, 8, device="cuda")
row_to_col = torchmatch.munkres_hybrid(cost)

cost_cpu = torch.randn(8, 8)
row_to_col_cpu = torchmatch.jv_scalar(cost_cpu)
```

The same ops live at `torch.ops.matching.<op>`, which is useful inside
`torch.compile` regions or when dispatching by name.

For control over extension loading (preforked workers, latency-sensitive
serving), the loaders stay exported and idempotent:

```python
torchmatch.load_cpu()                      # already ran at import; no-op
torchmatch.load_cuda()
```

Every op returns an `int64` row→col mapping of length `N` (the number
of rows). A `-1` entry marks an unmatched row.

## Ops

### Munkres family (CUDA only)

Primed/starred-zeros Hungarian/Munkres algorithm; three GPU
parallelization strategies. All carry the `cudagraph_unsafe` tag
(host-side syncs read managed-memory flags).

| Op                  | dtype (internal)      | Variant                                                                                                                    |
| ------------------- | --------------------- | -------------------------------------------------------------------------------------------------------------------------- |
| `munkres_classical` | float32, column-major | Single-path augmentation per outer iteration; CUB segmented reductions for column-min                                      |
| `munkres_hybrid`    | float32, row-major    | Adaptive: starts single-path, switches to tree augmentation when single-path stalls                                        |
| `munkres_tree`      | float64, row-major    | Always-tree: parallel BFS finds all vertex-disjoint augmenting paths per outer iteration; cooperative-groups + Thrust scan |

### Jonker-Volgenant family, CPU (single problem)

Successive-shortest-path (Dijkstra-like over reduced costs). All accept
rectangular `(N, M)` cost matrices in float32 or float64.

| Op           | Variant                                                                                                               |
| ------------ | --------------------------------------------------------------------------------------------------------------------- |
| `jv_scalar`  | Sequential reference; no SIMD. Implements Crouse 2016. Rejects NaN / -inf; treats +inf as an infeasible edge.         |
| `jv_dense`   | AVX2 flat-pointer inner loop; rectangular-capable                                                                     |
| `jv_compact` | AVX2-gather inner loop; square-only internal kernel (rectangular inputs padded by the wrapper)                        |

### Jonker-Volgenant family, batched

| Op                          | Devices    | Constraints                                                                                                                                                                                |
| --------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
| `jv_dense_batch`            | CPU + CUDA | CPU: `(B, N, M)`, any size, `at::parallel_for` over per-problem `jv_dense`. CUDA: shared-memory tiled kernel, one block per problem; requires `(B, K, K)` square with `K ≤ MAX_TILE = 64`. |
| `jv_compact_batch`          | CPU        | `at::parallel_for` over per-problem `jv_compact`                                                                                                                                           |
| `jv_dense_batch_unpacked`   | CPU        | Returns `(matches, unmatched_rows, unmatched_cols, n_matched)`, which saves a per-problem Python unpack                                                                                     |
| `jv_compact_batch_unpacked` | CPU        | Same shape, compact variant                                                                                                                                                                |

## torch.compile / torch.export

Every op carries a FakeTensor kernel. The Munkres CUDA ops carry the
`cudagraph_unsafe` tag (host-side syncs); the CUDA backend of
`jv_dense_batch` is fully capturable.

## Build modes

The per-device loaders prefer a prebuilt `.so` shipped in the wheel and
fall back to JIT-compiling the C++/CUDA sources via
`torch.utils.cpp_extension.load`. Both paths register the same
`torch.ops.matching.*` ops; the choice is transparent to callers.

### Building wheels

```bash
# default: builds CPU extension; builds CUDA extension if a CUDA
# toolchain is detected (torch.utils.cpp_extension.CUDA_HOME)
pip wheel . -w dist/

# CPU-only wheel
TORCHMATCH_SKIP_CUDA=1 pip wheel . -w dist/

# CUDA SM targets (default: PyTorch's current-device default)
TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0" pip wheel . -w dist/
```

The build system pairs `setuptools` with `torch.utils.cpp_extension.BuildExtension`.
Sdists ship the full `` tree so the JIT path works on any
CUDA-capable machine, even without a matching wheel.

### Runtime overrides

- `TORCHMATCH_FORCE_JIT=1`: skip the prebuilt `.so` and recompile from
  source via `cpp_extension.load`. Useful for development and for
  diagnosing ABI mismatches.

## Development environment

A self-contained `flake.nix` provides Python 3.13, uv, and a chosen
CUDA toolkit (default 12.8) without touching your global system. With
direnv:

```bash
direnv allow                      # picks devShells.default = cu128
NIX_DEVSHELL_NAME=cpu  direnv reload
NIX_DEVSHELL_NAME=cu130 direnv reload
```

Without direnv:

```bash
nix develop                       # default = cu128
nix develop .#cpu                 # CPU-only (sets TORCHMATCH_SKIP_CUDA=1)
nix develop .#cu126
nix develop .#cu130
```

Once inside the shell:

```bash
uv sync --all-groups              # populate .venv from pyproject.toml
uv run python -m pytest tests/    # run the test suite
```
