Backends and Capability Matrix¶

torch-sla dispatches each solve() call to one of several backends. Pick a backend explicitly via backend="..." or let backend="auto" choose based on device, dtype, problem size, and which optional dependencies are installed.

The current backend lineup and what each supports:

**Capability matrix**¶
Backend	CPU	CUDA	ROCm	Direct	Iterative	Complex	Batched	Distributed	Autograd
`scipy`	✔	--	--	LU / UMFPACK	CG, BiCGStab, GMRES	✔	via batch helpers	--	✔
`pytorch`	✔	✔	✔	--	CG, BiCGStab, GMRES, MINRES, LSQR, LSMR (+ PCG / PBiCGStab)	✔	✔	via `DSparseTensor`	✔
`strumpack`	✔	✔	✔	LU (multifrontal)	--	✔	--	--	✔
`cudss`	--	✔	--	LU / Cholesky / LDL^T / LDL^H	--	✔	--	--	✔
`pyamg`	✔	✔ (V-cycle only)	✔ (V-cycle only)	--	Ruge-Stuben AMG, Smoothed Aggregation	--	--	--	✔
`amgx`	--	✔	--	--	AMG, PCG, PBiCGStab, FGMRES (NVIDIA AmgX)	--	--	--	✔

备注

All six backends are verified correct. Each is checked against a reference solution with the relative residual ‖Ax − b‖ / ‖b‖ at or near machine precision. The two direct GPU paths land well inside that envelope — measured strumpack ≈ 3e-13 and amgx ≈ 5.6e-13 on the verification matrices.

备注

cudss and pyamg are PyPI-installable, but the two native compiled backends — strumpack (torch-strumpack) and amgx (torch-amgx) — ship as prebuilt wheels on GitHub Releases (not PyPI), and each wheel is ABI-tied to a specific CUDA and PyTorch version. See 预编译 wheel：torch-amgx 与 torch-strumpack in the installation guide for the wheel-selection rules and a concrete pip install --no-deps example.

The STRUMPACK backend¶

backend="strumpack" is a portable multifrontal sparse direct solver. Unlike cuDSS (which is NVIDIA CUDA only), STRUMPACK runs on CPU, CUDA, and AMD ROCm from the same API, supports both real and complex matrices, and offers a multifrontal LU factorization. It is fully differentiable: gradients flow through the adjoint (A^H) solve, so it drops into autograd pipelines like the other backends.

In practice STRUMPACK is the answer for a GPU direct solve on hardware where cuDSS cannot go — most importantly AMD ROCm GPUs, where cuDSS is unavailable. It requires the optional torch-strumpack package, published as prebuilt wheels on GitHub Releases (not PyPI; see 预编译 wheel：torch-amgx 与 torch-strumpack) for Linux cpu / cuda / rocm and macOS arm64. Windows (CPU) is supported — STRUMPACK builds on Windows with clang-cl (C/C++) + flang (Fortran) from conda-forge, linked against MSVC-built PyTorch (a clean-env solve gives relative residual ~1.7e-16); a prebuilt Windows wheel via CI is being added:

# Grab the matching wheel from
#   https://github.com/sparsexlab/torch-strumpack/releases
pip install --no-deps <release-url>/torch_strumpack-...-linux_x86_64.whl

Platform availability¶

Direct-solver backends bind to vendor libraries; the table below records which OS each one builds on today.

Backend	Linux	Windows	macOS	Notes
`scipy`	✔	✔	✔	Pure SciPy; UMFPACK optional via `scikit-umfpack`.
`pytorch`	✔	✔	✔	PyTorch-native; CUDA / ROCm path active when `torch.cuda.is_available()` (ROCm torch builds report as `cuda`).
`strumpack`	✔	✔ (CPU)	✔ (arm64)	Multifrontal sparse direct solver (multifrontal LU, real + complex). CPU / CUDA / ROCm on Linux + macOS arm64 via `torch-strumpack` (GitHub-Release wheels, 预编译 wheel：torch-amgx 与 torch-strumpack). Windows (CPU) supported — builds with `clang-cl` + `flang`; prebuilt Windows wheel via CI pending.
`cudss`	✔	✔	--	Requires `nvmath-python[cu12]` + NVIDIA CUDA. macOS is not supported by Nvidia.
`pyamg`	✔	✔	✔	Setup runs on CPU via the optional `pyamg` dependency (`pip install pyamg`); the V-cycle dispatches through `torch.sparse` so the cycle itself runs on whatever device the matrix lives on. Cross-platform AMG: macOS gets CPU AMG, CUDA boxes get GPU V-cycles.

When `backend="auto"` picks what¶

NVIDIA CUDA tensors: try cudss (best direct solver) -> pytorch (iterative fallback).
AMD ROCm tensors: cuDSS is NVIDIA-only and never runs here, so the auto path uses pytorch (iterative) and, when a direct solve is needed, strumpack (portable multifrontal direct solver on ROCm).
CPU tensors, small / medium: prefer scipy LU.
CPU tensors, large or repeated: pytorch CG / BiCGStab keeps the memory footprint flat.

Override via backend="..." whenever you need exact control (e.g. backend="cudss" to force a direct GPU solve on NVIDIA, or backend="strumpack" for a direct GPU solve on AMD ROCm where cuDSS is unavailable).

Direct vs iterative: accuracy and complexity¶

The accuracy tables quote ~1e-14 for the direct backends and ~1e-6 for the iterative ones. The gap is structural, not a bug:

Direct solvers factor the matrix (LU / Cholesky / LDL) and back-substitute. The result is exact up to floating-point round-off -- for a well-conditioned float64 system the relative residual sits near machine epsilon (~1e-14..``1e-16``). There is no convergence knob; you pay the factorization cost once and get a fully accurate answer.
Iterative solvers (CG, BiCGStab, GMRES, ...) refine a guess until the residual ‖Ax − b‖ / ‖b‖ drops below a tolerance you set (atol / rtol, default ~1e-6). They stop at the tolerance, so the answer is only as accurate as you ask for. Tighten atol toward 1e-12 and the residual follows -- at the cost of more iterations. The ill-conditioning of A (its condition number) sets how many iterations each digit costs.

So the iterative ~1e-6 is a default stopping point, not a precision ceiling. The trade-off is what makes the iterative path scale: it never forms the dense fill-in that a factorization does.

Direct vs iterative cost (`n` unknowns, `nnz` non-zeros, `m` iterations)¶
Solver	Time	Space	Accuracy
Direct (LU / Cholesky)	\(O(n^{1.5})\) (2-D) to \(O(n^{2})\) (3-D)	\(O(n\log n)\) to \(O(n^{4/3})\) fill-in	Exact to round-off (`~1e-14`)
Iterative (CG / GMRES)	\(O(m\cdot nnz)\)	\(O(n + nnz)\)	Tolerance-limited (`atol`, default `~1e-6`)

For a sparse PDE matrix nnz = O(n), so an iterative sweep costs \(O(m\,n)\) time in \(O(n)\) memory, while the direct factorization's fill-in is what exhausts memory past a few million unknowns (see the 基准测试). Pick direct when you need the last digits or have many right-hand sides to reuse a factorization on; pick iterative when the matrix is large and ~1e-6 is enough.

Putting it together¶

The capability matrix maps directly to the solve() parameters: any combination where the cell is ✔ is supported:

import torch
from torch_sla import solve, PreconditionerConfig

A_csr = ...                          # any accepted matrix format
b = torch.randn(n)

# Direct GPU solve, automatic Cholesky/LDL^H selection
x = solve(A_csr, b, backend="cudss", matrix_type="auto")

# CPU iterative CG with a tuned SSOR preconditioner
x = solve(A_csr, b,
          backend="pytorch", method="cg",
          preconditioner=PreconditionerConfig(kind="ssor", omega=1.2),
          atol=1e-10, maxiter=5_000)

# CPU iterative CG with a real multi-level AMG preconditioner
# (uses PyAMG when installed, falls back to the lightweight
# 2-level stub otherwise). Reduces the iteration count by 10-100x
# on ill-conditioned PDE problems.
x = solve(A_csr, b,
          backend="pytorch", method="cg",
          preconditioner="amg",  # or PreconditionerConfig(kind="amg", ...)
          atol=1e-10, maxiter=200)

# Diagnostic return -- iteration count + residual
x, info = solve(A_csr, b, return_info=True)
print(info.iter_count, info.residual, info.converged)

Future backends (roadmap)¶

The next wave of backends will extend the table with cross-platform AMG preconditioning and high-end GPU AMG:

Backend	Status	Capability	Notes
`pyamg`	available (this release)	CPU AMG setup + cross-device V-cycle	Already shipping. See above. Standalone solver + `PyAMGHierarchy` for preconditioner re-use.
`amgx`	available (this release)	CUDA AMG + Krylov (Nvidia AmgX)	Linux + Windows only. NVIDIA GPU required (incl. Blackwell `sm_120` on cu12.8). Install the prebuilt wheel from torch-amgx Releases (not PyPI; see 预编译 wheel：torch-amgx 与 torch-strumpack).
`petsc`	investigating	CPU/GPU direct + iterative, distributed (PETSc/hypre BoomerAMG)	Linux + macOS easy; Windows via WSL.