Metadata-Version: 2.4
Name: gpuastar
Version: 0.1.0
Summary: A* search backends for Python, C++, and CUDA with neural-network heuristics
Project-URL: GitHub, https://github.com/misaghsoltani/GPUAStar/
Author-email: Misagh Soltani <msoltani@email.sc.edu>
License-Expression: MIT
License-File: LICENSE
Keywords: A* search,GPU A*,batch weighted A*,combinatorial optimization,deep learning,heuristic search,puzzle solving
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: <3.15,>=3.10
Requires-Dist: hydra-core<2,>=1.3.2
Requires-Dist: numpy<2,>=1.26; implementation_name == 'cpython' and platform_system == 'Darwin' and platform_machine == 'x86_64' and python_version < '3.13'
Requires-Dist: numpy<3,>=2.0.0; implementation_name == 'cpython' and python_version < '3.13' and (platform_system != 'Darwin' or platform_machine != 'x86_64')
Requires-Dist: numpy<3,>=2.1; implementation_name == 'cpython' and python_version >= '3.13' and python_version < '3.14'
Requires-Dist: numpy<3,>=2.3.2; implementation_name == 'cpython' and python_version >= '3.14'
Requires-Dist: orjson<4,>=3.11; implementation_name == 'cpython' and python_version < '3.14'
Requires-Dist: packaging>=22
Requires-Dist: pydantic<3,>=2.12
Requires-Dist: pyyaml<7,>=6.0.1; python_version >= '3.12' and python_version < '3.13'
Requires-Dist: pyyaml<7,>=6.0.2; python_version >= '3.13' and python_version < '3.14'
Requires-Dist: pyyaml<7,>=6.0.3; python_version >= '3.14'
Requires-Dist: pyyaml<7,>=6; python_version < '3.12'
Requires-Dist: torch<2.3,>=2.2.2; implementation_name == 'cpython' and platform_system == 'Darwin' and platform_machine == 'x86_64' and python_version < '3.13'
Requires-Dist: torch<3,>=2.11; implementation_name == 'cpython' and python_version >= '3.14' and python_version < '3.15' and ((platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))
Requires-Dist: torch<3,>=2.5.1; implementation_name == 'cpython' and platform_system == 'Linux' and platform_machine == 'aarch64' and python_version < '3.13'
Requires-Dist: torch<3,>=2.6; implementation_name == 'cpython' and python_version < '3.14' and ((platform_system == 'Linux' and platform_machine == 'x86_64') or (platform_system == 'Linux' and platform_machine == 'aarch64' and python_version >= '3.13') or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))
Requires-Dist: typing-extensions>=4.10
Provides-Extra: build
Requires-Dist: build<2,>=1.5; extra == 'build'
Requires-Dist: cibuildwheel[uv]<3,>=2.23; (python_version < '3.11') and extra == 'build'
Requires-Dist: cibuildwheel[uv]<5,>=4; (python_version >= '3.11') and extra == 'build'
Requires-Dist: hatch<2,>=1.14; extra == 'build'
Requires-Dist: hatchling<2,>=1.27; extra == 'build'
Requires-Dist: twine<7,>=6; extra == 'build'
Provides-Extra: cpp
Requires-Dist: gpuastar-cpp==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and ((platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))) and extra == 'cpp'
Provides-Extra: cpp-cu-experiments
Requires-Dist: gpuastar-cpp==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and ((platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))) and extra == 'cpp-cu-experiments'
Requires-Dist: gpuastar-cu==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) and extra == 'cpp-cu-experiments'
Requires-Dist: gpuastar-experiments[cpp,runtime]==0.1.0; (implementation_name == 'cpython') and extra == 'cpp-cu-experiments'
Provides-Extra: cpp-experiments
Requires-Dist: gpuastar-cpp==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and ((platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))) and extra == 'cpp-experiments'
Requires-Dist: gpuastar-experiments[cpp,runtime]==0.1.0; (implementation_name == 'cpython') and extra == 'cpp-experiments'
Provides-Extra: cu
Requires-Dist: gpuastar-cu==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) and extra == 'cu'
Provides-Extra: cu-experiments
Requires-Dist: gpuastar-cu==0.1.0; (implementation_name == 'cpython' and python_version < '3.15' and platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) and extra == 'cu-experiments'
Requires-Dist: gpuastar-experiments[runtime]==0.1.0; (implementation_name == 'cpython') and extra == 'cu-experiments'
Provides-Extra: dev
Requires-Dist: mypy; extra == 'dev'
Requires-Dist: pyrefly; extra == 'dev'
Requires-Dist: pyright; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: yamllint; extra == 'dev'
Provides-Extra: experiments
Requires-Dist: gpuastar-experiments[runtime]==0.1.0; (implementation_name == 'cpython') and extra == 'experiments'
Provides-Extra: runtime
Requires-Dist: jinja2<4,>=3.1.2; (implementation_name == 'cpython') and extra == 'runtime'
Requires-Dist: numpy<2,>=1.26; (implementation_name == 'cpython' and platform_system == 'Darwin' and platform_machine == 'x86_64' and python_version < '3.13') and extra == 'runtime'
Requires-Dist: numpy<3,>=2.0.0; (implementation_name == 'cpython' and python_version < '3.13' and (platform_system != 'Darwin' or platform_machine != 'x86_64')) and extra == 'runtime'
Requires-Dist: numpy<3,>=2.1; (implementation_name == 'cpython' and python_version >= '3.13' and python_version < '3.14') and extra == 'runtime'
Requires-Dist: numpy<3,>=2.3.2; (implementation_name == 'cpython' and python_version >= '3.14') and extra == 'runtime'
Requires-Dist: setuptools<82,>=69; (implementation_name == 'cpython') and extra == 'runtime'
Requires-Dist: torch<2.3,>=2.2.2; (implementation_name == 'cpython' and platform_system == 'Darwin' and platform_machine == 'x86_64' and python_version < '3.13') and extra == 'runtime'
Requires-Dist: torch<3,>=2.11; (implementation_name == 'cpython' and python_version >= '3.14' and python_version < '3.15' and ((platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')) or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))) and extra == 'runtime'
Requires-Dist: torch<3,>=2.5.1; (implementation_name == 'cpython' and platform_system == 'Linux' and platform_machine == 'aarch64' and python_version < '3.13') and extra == 'runtime'
Requires-Dist: torch<3,>=2.6; (implementation_name == 'cpython' and python_version < '3.14' and ((platform_system == 'Linux' and platform_machine == 'x86_64') or (platform_system == 'Linux' and platform_machine == 'aarch64' and python_version >= '3.13') or (platform_system == 'Darwin' and platform_machine == 'arm64') or (platform_system == 'Windows' and platform_machine == 'AMD64'))) and extra == 'runtime'
Provides-Extra: test
Requires-Dist: pytest; extra == 'test'
Requires-Dist: pytest-cov; extra == 'test'
Requires-Dist: tomli<3,>=2.1; (python_version < '3.11') and extra == 'test'
Description-Content-Type: text/markdown

# GPUAStar (GPU-A\*): GPU-Accelerated A* Search with Deep Neural Network Heuristics

[![Publication](https://img.shields.io/badge/publication-SoCS-%234285F4?logo=googlescholar&logoColor=%23d0d0d0&cacheSeconds=https://socs26.search-conference.org/accepted-papers#:~:text=GPU%2DAccelerated%20A,Soltani;%20Forest%20Agostinelli)](https://img.shields.io/badge/publication-SoCS-%234285F4?logo=googlescholar&logoColor=%23d0d0d0&cacheSeconds=https://socs26.search-conference.org/accepted-papers#:~:text=GPU%2DAccelerated%20A,Soltani;%20Forest%20Agostinelli)
[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
[![image](https://img.shields.io/pypi/v/gpuastar.svg)](https://pypi.python.org/pypi/gpuastar)
[![image](https://img.shields.io/pypi/l/gpuastar.svg)](https://github.com/misaghsoltani/GPUAStar/blob/main/LICENSE)
[![Pixi Badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/prefix-dev/pixi/main/assets/badge/v0.json&label=package%20manager)](https://pixi.sh)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
![Static Badge](https://img.shields.io/badge/statically%20typed-mypy-039dfc)
<!-- [![Checked with Pyright](https://microsoft.github.io/pyright/img/pyright_badge.svg)](https://microsoft.github.io/pyright/) -->

This repository contains the official implementation of the paper **[GPU-Accelerated A* Search with Deep Neural Network Heuristics](https://socs26.search-conference.org/accepted-papers#:~:text=GPU%2DAccelerated%20A,Soltani;%20Forest%20Agostinelli)** accepted to **The 19th International Symposium on Combinatorial Search (SoCS 2026)**. In this paper we introduce accelerator-resident Batch Weighted A\* Search (BWAS), and across 100 test instances of Rubik's Cube, we show that GPU-A\* search gains over **12× speed-up** compared to the host/CPU-resident batched weighted A\* search (with neural network heuristics which runs inference on GPU) when running the Python implementation of our approach on a single NVIDIA H200 GPU. The 12× speed-up is achieved while preserving the exact same search order as the baseline, expanding the same number of nodes, and without relying on JIT compilation or lower-precision arithmetic (either of which are available in the implementation and provide additional speed-ups).

## Contents

- [GPUAStar (GPU-A\*): GPU-Accelerated A\* Search with Deep Neural Network Heuristics](#gpuastar-gpu-a-gpu-accelerated-a-search-with-deep-neural-network-heuristics)
  - [Contents](#contents)
  - [Quick Start](#quick-start)
    - [Installation](#installation)
    - [Reproducing Paper's Results](#reproducing-papers-results)
  - [Running other Backends](#running-other-backends)
  - [Solver backends](#solver-backends)
  - [Search domains](#search-domains)
  - [Building the compiled extensions](#building-the-compiled-extensions)
  - [Documentation](#documentation)
  - [License](#license)

## Quick Start

### Installation

We use [Pixi](https://pixi.sh) for package and environment managemanet. Using `pixi` you can install `GPUAStar` using the following commands.

```bash
git clone https://github.com/misaghsoltani/GPUAStar.git
cd GPUAStar
pixi install  # This only installs the Python backend, which is enough for reproducing the results reported in the paper
pixi run gpuastar --help
```

The packages can also be installed with pip. Optional components are selected through extras.

```bash
pip install "gpuastar"               # installs the core package, the default Python backends, and CLI
pip install "gpuastar[cpp]"          # adds the compiled C++ backend (gpuastar-cpp)
pip install "gpuastar[cu]"           # adds the CUDA backend (gpuastar-cu)
pip install "gpuastar[experiments]"  # adds example domains, datasets, models, and the experiment suite
pip install "gpuastar[cpp-experiments]"     # adds the C++ backend and compiled example domains
pip install "gpuastar[cu-experiments]"      # adds the CUDA backend and example domains
pip install "gpuastar[cpp-cu-experiments]"  # adds every backend and compiled example domain
```

Supported platforms follow the PyTorch wheel availability declared in [pyproject.toml](pyproject.toml). Upstream PyTorch published no macOS x86_64 wheel after 2.2.2, so on macOS Intel with CPython 3.13 or newer the package installs without Torch and the solver backends report themselves unavailable through `BackendUnavailableError`. The C++ backend supports glibc-based Linux on x86_64 and aarch64, macOS arm64, macOS x86_64 through CPython 3.12, and Windows AMD64. The CUDA backend supports glibc-based Linux on x86_64 and aarch64. Official PyPI Torch wheels for Windows AMD64 are CPU-only, so the CUDA extras do not select `gpuastar-cu` on Windows or macOS. CPython 3.10 through 3.14 is supported, and the compiled backends also publish free-threaded CPython 3.14t (`cp314t`) wheels. Build prerequisites and LibTorch compatibility lines are listed in [docs/installation.md](docs/installation.md).

### Reproducing Paper's Results

For reproducing the results reported in the paper, you can use the following command **after installing `GPUAStar`** (either using `pixi` or from `PyPI`):

```bash
git clone https://github.com/misaghsoltani/GPUAStar.git  # if you haven't already cloned it!
cd GPUAStar
. ./reproduce_paper_results/run_all.sh
```

## Running other Backends

To compile a heuristic checkpoint into an AOTInductor package for the C++ CPU solver:

```bash
pixi run -e dev gpuastar-aoti-export export.model_dir=experiments/src/gpuastar_experiments/saved_models/cube3/current export.env=cube3
```

Each search run writes a `search_metrics.json` file with per-state records (solved flag, solution cost, move count, nodes generated, solve time, iteration count, and per-operation timings). The output format is described in [docs/usage.md](docs/usage.md).

All configuration groups and their default values are documented in [docs/cli.md](docs/cli.md).

## Solver backends

The `backend.language` option selects the solver implementation. Backend discovery uses the `gpuastar.backends` entry-point group and is implemented in [backends.py](src/gpuastar/backends.py), so additional solver packages can register themselves the same way.

- **`python`** ([python.py](src/gpuastar/astar/backends/python.py)): the reference implementation used for generating the results presented in the paper. With `backend.gpu_only=false` it runs the CPU solver in [cpu.py](src/gpuastar/astar/solvers/cpu.py), which uses binary heaps for the open list and dictionaries for the closed set, with heuristic evaluation batched on the configured device. With `backend.gpu_only=true` it runs the GPU-resident search implemented in [solvers/torch.py](src/gpuastar/astar/solvers/torch.py) through [backends/torch.py](src/gpuastar/astar/backends/torch.py), which keeps the open list, the closed set, and all state data in PyTorch tensors on the device. The closed set uses `TorchDict` ([core.py](src/gpuastar/utils/torchdict/core.py)), an open-addressing hash table stored in tensors.
- **`cpp`** ([cpp.py](src/gpuastar/astar/backends/cpp.py)): loads the `gpuastar_cpp` nanobind extension module. With `backend.gpu_only=false` it runs [parallel_weighted_astar.cpp](cpp/cxx/parallel_weighted_astar.cpp), a CPU solver parallelized with OpenMP that evaluates the heuristic through an AOTInductor-compiled model package. With `backend.gpu_only=true` it runs [torch_gpu_astar.cpp](cpp/cxx/torch_gpu_astar.cpp), a solver written against the LibTorch C++ API that keeps the search state in device tensors (GPU-resident search).
- **`cu`** ([cuda.py](src/gpuastar/astar/backends/cuda.py)): loads the `gpuastar-cu` kernel, which is custom kernels that implement the same GPU-resident search algorithm using CUDA programming language. The search loop in [search.py](cu/src/gpuastar_cu/search.py) launches the kernels compiled from [bwas_kernel.cu](cu/csrc/bwas_kernel.cu), which expand states cooperatively within warps and resolve hash table collisions with atomic operations.
- **`auto`**: selects the first available backend.

## Search domains

A search domain is defined by subclassing `Environment` and `State` from [environment_abstract.py](src/gpuastar/environments/environment_abstract.py) and registering the class with `register_environment` or `register_lazy` from [registry.py](src/gpuastar/utils/environments/registry.py). For the C++ solver, states provide their payload either through `colors` or `tiles` sequences or through a `cpp_payload()` method (the `CppPayloadState` protocol in [protocols.py](src/gpuastar/protocols.py)). An environment can set `suggested_compile_max_nodes` to size compiled GPU search artifacts.

Compiled C++ environments are provided by separate packages through the `gpuastar_cpp.environments` entry-point group. The registry interface is declared in [environment_registry.h](cpp/cxx/environment_registry.h), and [`experiments/cpp/`](experiments/cpp/) contains an example plugin package.

The example domains shipped with `gpuastar[experiments]` are described in [docs/environments.md](docs/environments.md).

## Building the compiled extensions

```bash
# Build the C++ solver extension from a clean state
pixi run -e dev cpp-build-from-scratch

# Build the CUDA solver extension
pixi run -e dev cu-build

# Verify that the CUDA extension imports and the backend is available
pixi run -e dev cu-smoke-import
```

After editing the shared C++ headers in `cpp/cxx/` or the plugin sources in `experiments/cpp/`, remove `experiments/cpp/build` and re-run `pixi install -e dev` so the compiled environment plugin matches the current headers. Details are in [docs/installation.md](docs/installation.md).

## Documentation

- [docs/installation.md](docs/installation.md): environment setup and extension builds.
- [docs/usage.md](docs/usage.md): running searches, AOTInductor export, CUDA graphs, profiling, output metrics.
- [docs/cli.md](docs/cli.md): commands and all configuration groups.
- [docs/api_reference.md](docs/api_reference.md): public Python interfaces.
- [docs/python_api.md](docs/python_api.md): calling the solvers from Python code.
- [docs/environments.md](docs/environments.md): environment abstraction and example domains.
- [docs/distributed.md](docs/distributed.md): multi-GPU execution with FSDP2, tensor parallelism, and DTensor.
- [docs/benchmarks.md](docs/benchmarks.md): benchmark runners in the experiments package.

<!-- ## Citation

If you use GPUAStar in your research, please cite it using:

```bibtex
Placeholder
``` -->

## License

GPUAStar is released under the [MIT License](LICENSE).
