torchmatch 1.0.0__tar.gz

torchmatch-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Kurt Stolle
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchmatch-1.0.0/PKG-INFO

@@ -0,0 +1,186 @@
+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
+```

torchmatch-1.0.0/README.md

@@ -0,0 +1,156 @@
+# 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
+```

torchmatch-1.0.0/pyproject.toml

@@ -0,0 +1,174 @@
+[build-system]
+# Pinned to the same range as [project.dependencies]: a wider build-system
+# constraint lets the isolated build env resolve a torch whose libtorch ABI
+# does not match the runtime torch, producing a .so that fails to load with
+# `undefined symbol: torch::Library::_def(...)` style errors.
+requires = ["setuptools>=70", "wheel", "torch>=2.11,<2.13"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torchmatch"
+description = "Linear assignment problem solvers for PyTorch."
+readme = "README.md"
+version = "1.0.0"
+requires-python = ">=3.13, <3.14"
+keywords = [
+  "pytorch",
+  "linear assignment problem",
+  "lap",
+  "lapjv",
+  "hungarian algorithm",
+  "bipartite matching",
+  "cuda",
+]
+authors = [{ name = "Kurt Stolle", email = "kurt@computer.org" }]
+license = "MIT"
+license-files = ["LICENSE"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Scientific/Engineering :: Mathematics",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.13",
+  "Operating System :: POSIX :: Linux",
+]
+dependencies = ["torch>=2.11"]
+
+[project.optional-dependencies]
+cpu = ["torch>=2.11"]
+cu126 = ["torch>=2.11"]
+cu128 = ["torch>=2.11"]
+cu130 = ["torch>=2.12"]
+
+[dependency-groups]
+dev = [
+  "pytest>=8.4.2",
+  "pytest-benchmark>=5.1.0",
+  "pytest-timeout>=2.4.0",
+  "ruff>=0.9.0",
+  "scipy>=1.10",
+]
+docs = [
+  "mkdocs-material>=9.5",
+  "mkdocstrings[python]>=0.27",
+  "pymdown-extensions>=10.0",
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu126"
+url = "https://download.pytorch.org/whl/cu126"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", extra = "cpu" },
+  { index = "pytorch-cu126", extra = "cu126" },
+  { index = "pytorch-cu128", extra = "cu128" },
+  { index = "pytorch-cu130", extra = "cu130" },
+]
+
+[tool.uv]
+# [build-system].requires uses a separate isolated env that does NOT see
+# [tool.uv.sources], so pip would pull torch from PyPI's default index
+# (which bundles CUDA 13 in current releases) and produce a .so whose
+# torch ABI / CUDA version does not match the venv torch. Disable build
+# isolation for torchmatch so the build reuses the venv's pinned torch
+# (cu126 / cu128 / cu130, picked by the user-selected extra).
+no-build-isolation-package = ["torchmatch"]
+
+conflicts = [
+  [
+    { extra = "cpu" },
+    { extra = "cu126" },
+    { extra = "cu128" },
+    { extra = "cu130" },
+  ],
+]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["sources"]
+
+[tool.setuptools.package-dir]
+"" = "sources"
+
+[tool.setuptools.package-data]
+"torchmatch" = [
+  "py.typed",
+  "cuda/*.cpp",
+  "cuda/*.cu",
+  "cuda/*.cuh",
+  "cpu/*.cpp",
+  "cpu/*.h",
+]
+
+[tool.ruff]
+target-version = 'py313'
+line-length = 88
+
+[tool.ruff.lint]
+select = [
+  "ARG",
+  "B",
+  "C",
+  "C4",
+  "D",
+  "E",
+  "EM",
+  "F",
+  "FURB",
+  "G",
+  "I",
+  "ICN",
+  "N",
+  "NPY",
+  "PERF",
+  "PGH",
+  "PIE",
+  "PL",
+  "PT",
+  "PTH",
+  "PYI",
+  "RUF",
+  "RET",
+  "SIM",
+  "T20",
+  "TRY",
+  "UP",
+]
+ignore = ["PLR2004", "PLW2901", "D212", "D203", "D107"]
+
+[tool.ruff.lint.per-file-ignores]
+# pytest.mark.parametrize functions take one parameter per fixture, so
+# PLR0913's 5-argument cap fires on every multi-parametrize test.
+"tests/*" = ["D", "T", "PLC0415", "PLR0913", "PLR1714"]
+# Utility scripts: print is expected (CLI tools), docstrings are optional,
+# RUF002 ambiguous-unicode rules are docstring-only false positives.
+"scripts/*" = ["D", "T", "RUF002", "C901", "PERF401"]
+
+[tool.ruff.format]
+docstring-code-format = true
+docstring-code-line-length = "dynamic"
+
+[tool.pytest.ini_options]
+timeout = 10

torchmatch-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torchmatch-1.0.0/setup.py

@@ -0,0 +1,97 @@
+"""
+Build the C++/CUDA extensions for torchmatch.
+
+Project metadata lives in ``pyproject.toml``; this file defines only
+``ext_modules`` so :class:`torch.utils.cpp_extension.BuildExtension` can
+compile the native sources into precompiled ``.so`` files.
+
+Build-time behaviour
+--------------------
+- ``TORCHMATCH_SKIP_CUDA=1`` skips the CUDA extension entirely.
+- ``TORCHMATCH_SKIP_CPU=1`` skips the CPU extension.
+- When no CUDA toolchain is detected (``torch.utils.cpp_extension.CUDA_HOME``
+  is None), the CUDA extension drops out silently, so an sdist install
+  on a CPU-only box still produces a usable wheel with just the CPU
+  extension; CUDA ops then fall back to JIT at runtime on machines that
+  do have a toolchain.
+- ``TORCH_CUDA_ARCH_LIST`` controls which SM versions ``nvcc`` targets;
+  honor it when set, otherwise PyTorch's default kicks in.
+- CPU compile flags default to ``-O3 -mavx2 -mfma``. Older CPUs without
+  AVX2 install from sdist (runtime JIT detects flags from
+  ``/proc/cpuinfo``).
+
+Runtime behaviour
+-----------------
+``torchmatch.load_cpu()`` / ``torchmatch.load_cuda()`` look for a
+prebuilt ``_cpu_impl*.so`` / ``_cuda_impl*.so`` next to the package and
+``dlopen`` it via ``torch.ops.load_library``. When no prebuilt is
+present, they JIT-compile via ``torch.utils.cpp_extension.load``. Set
+``TORCHMATCH_FORCE_JIT=1`` to bypass the prebuilt path.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from setuptools import Extension, setup
+from torch.utils.cpp_extension import (
+    CUDA_HOME,
+    BuildExtension,
+    CppExtension,
+    CUDAExtension,
+)
+
+CSRC = Path("sources/torchmatch")
+CUDA_CSRC = CSRC / "cuda"
+CPU_CSRC = CSRC / "cpu"
+
+SKIP_CUDA = os.environ.get("TORCHMATCH_SKIP_CUDA") == "1"
+SKIP_CPU = os.environ.get("TORCHMATCH_SKIP_CPU") == "1"
+
+
+def _ext_modules() -> list[Extension]:
+    mods: list[Extension] = []
+
+    if not SKIP_CPU:
+        mods.append(
+            CppExtension(
+                name="torchmatch._cpu_impl",
+                sources=[
+                    str(CPU_CSRC / "jv_scalar.cpp"),
+                    str(CPU_CSRC / "jv_dense_core.cpp"),
+                    str(CPU_CSRC / "ops.cpp"),
+                ],
+                include_dirs=[str(CPU_CSRC)],
+                extra_compile_args={
+                    "cxx": ["-O3", "-std=c++17", "-mavx2", "-mfma"],
+                },
+            )
+        )
+
+    if not SKIP_CUDA and CUDA_HOME is not None:
+        mods.append(
+            CUDAExtension(
+                name="torchmatch._cuda_impl",
+                sources=[
+                    str(CUDA_CSRC / "ops.cpp"),
+                    str(CUDA_CSRC / "dispatch.cu"),
+                    str(CUDA_CSRC / "munkres_classical.cu"),
+                    str(CUDA_CSRC / "munkres_hybrid.cu"),
+                    str(CUDA_CSRC / "munkres_tree.cu"),
+                ],
+                include_dirs=[str(CUDA_CSRC)],
+                extra_compile_args={
+                    "cxx": ["-O3"],
+                    "nvcc": ["-O3", "-std=c++17"],
+                },
+            )
+        )
+
+    return mods
+
+
+setup(
+    ext_modules=_ext_modules(),
+    cmdclass={"build_ext": BuildExtension},
+)