Metadata-Version: 2.4
Name: erfi-pytorch
Version: 1.0.0
Summary: GPU-accelerated imaginary error function for real PyTorch tensors
Author: erfi-pytorch contributors
License-Expression: MIT
Project-URL: Homepage, https://github.com/ZhichaoZhu/erfi_pytorch
Project-URL: Repository, https://github.com/ZhichaoZhu/erfi_pytorch.git
Project-URL: Issues, https://github.com/ZhichaoZhu/erfi_pytorch/issues
Project-URL: Documentation, https://github.com/ZhichaoZhu/erfi_pytorch/blob/main/docs/faddeeva.md
Keywords: pytorch,cuda,triton,special-functions,erfi
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: third_party/faddeeva/LICENSE
Requires-Dist: torch>=2.7
Provides-Extra: test
Requires-Dist: mpmath>=1.3; extra == "test"
Requires-Dist: pytest>=8; extra == "test"
Requires-Dist: scipy>=1.11; extra == "test"
Provides-Extra: benchmark
Requires-Dist: scipy>=1.11; extra == "benchmark"
Provides-Extra: triton
Requires-Dist: triton>=3.3; platform_system == "Linux" and extra == "triton"
Dynamic: license-file

# erfi-pytorch

`erfi-pytorch` provides a forward-only imaginary error function for real
PyTorch tensors:

```python
import torch
from erfi_pytorch import erfi, erfi_

x = torch.linspace(-4, 4, 1_000_000, device="cuda")
y = erfi(x)
```

Repeated latency-sensitive calls can reuse a contiguous output tensor and
avoid allocator overhead:

```python
output = torch.empty_like(x)
erfi(x, out=output)
```

The returned object is `output`. Exact in-place use (`erfi(x, out=x)`) is
supported; partially overlapping input and output views are rejected.

For the lowest fixed overhead, `erfi_(x)` overwrites a contiguous input and
skips allocation plus general output-alias validation. Native CPU wheels also
use a dedicated same-pointer dispatch path rather than passing through the
general `out=` wrapper; tensor size and data pointer extraction are folded
into the same CPython extension call. The coefficient tables are bound once
at import, and dtype-specific one- and two-tensor extension entries keep
invariant metadata out of small in-place and `out=` calls. Native CPU `out=`
also fuses shape, dtype, device, contiguity, gradient, and overlap validation
with execution in one extension boundary. For the common contiguous case,
disjoint address ranges bypass PyTorch's general overlap analyzer and the
already extracted pointers are reused by the evaluator; input dtype and
gradient checks are included before dtype dispatch. The ordinary allocating
`erfi(x)` path likewise
fuses input validation, contiguous conversion, output allocation, dtype
dispatch, and evaluation into one extension call:

```python
erfi_(x)
```

The in-place path is CUDA Graph capture-safe. A graph can copy fresh values
into a fixed buffer and evaluate them with one replay call:

```python
source = torch.linspace(-4, 4, 16, device="cuda")
buffer = torch.empty_like(source)
erfi(source)  # compile and cache the kernel before capture

graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
    buffer.copy_(source)
    erfi_(buffer)

graph.replay()
```

The package supports `torch.float32` and `torch.float64` and preserves tensor
shape, dtype, and device. Its pure-PyTorch graph is compatible with
`torch.compile(fullgraph=True, backend="eager")`, including preallocated
`out=` calls. Inductor compilation depends on a working platform compiler or
Triton installation and is validated separately on supported Linux CUDA
environments.

## Backends

- **Native CPU:** fused C++ path built into platform wheels. It uses OpenMP for
  sufficiently large tensors when the platform compiler supports it. The
  crossover adapts to the active OpenMP thread count. Parallel workers receive
  contiguous, SIMD-aligned chunks and reuse the same SSE2/mixed evaluator as
  the serial path instead of falling back to one scalar iteration per element.
  A small
  input sample selects between the original branch-based Faddeeva evaluator
  and a table-driven evaluator, avoiding branch-misprediction costs for
  shuffled or irregular values without slowing ordered inputs. For
  `|x| <= 4` for float32 and `|x| <= 2.5` for scalar float64,
  dtype-specific
  segmented direct `erfi` polynomials avoid the interval lookup, division,
  and exponential; exact signed zero is returned directly. On x86-64,
  locally ordered inputs evaluate matching polynomial segments with SSE2,
  two elements at a time for float32 and up to four for float64. Float64 keeps
  the direct SIMD range at `|x| <= 2.5`; ordered wider inputs use a mixed loop
  whose outer values retain the high-accuracy scalar Faddeeva evaluator. The
  two binary64 SSE2 Horner chains
  are explicitly interleaved so out-of-order cores can overlap their dependent
  multiply-add sequences; float32
  lanes are widened to binary64 for the existing coefficient accuracy before
  rounding once on store. A small locality sample keeps shuffled inputs on the
  scalar/table path, while zero-filled buffers retain the scalar signed-zero
  fast path.
- **Pure PyTorch:** portable CPU fallback and CUDA implementation.
- **NVRTC CUDA:** fused single-kernel path in CUDA-enabled PyTorch builds,
  including Windows wheels that bundle NVRTC but do not have Triton installed.
  Registered runtimes expose a checked in-place entry that validates dtype,
  gradient state, and contiguity and submits the kernel in one native boundary.
- **Triton:** fused path for contiguous NVIDIA CUDA tensors with at least
  65,536 elements, when Triton is available.

The NVRTC kernel is compiled lazily once per device and dtype, then cached for
the life of the process. Its cubin is also cached across processes under the
user's PyTorch extensions cache. Set `ERFI_PYTORCH_CACHE_DIR` to override that
location. Cache keys are independent of the PyTorch package version when the
CUDA version, GPU architecture, and kernel source remain compatible; existing
version-1 cache entries are migrated without recompilation. The header-free
kernel is compiled through PyTorch's bundled NVRTC
library directly, without importing the C++ extension build toolchain or
requiring a system CUDA toolkit. Latency-sensitive applications should still
make one warmup call during initialization. Platform wheels also use a native
CUDA Driver launcher to minimize Python/ctypes overhead; source installations
without the optional extension retain the Python launcher. The compiler,
module loader, and fallback launcher call NVRTC and the CUDA Driver directly,
so they do not depend on version-specific private PyTorch kernel APIs. Corrupt
or incompatible cache entries are discarded and recompiled. CUDA-free builds
or systems where runtime compilation fails automatically use the pure PyTorch
path. Launch geometry is cached per dtype: float64 uses one-warp blocks to
reduce small-vector latency, while float32 retains 256-thread blocks for
throughput. The native in-place launcher extracts tensor size and pointer and
computes its grid inside the same extension call. Native CUDA `out=` calls
also fuse output validation, input contiguous conversion, pointer extraction,
grid calculation, and kernel launch into one extension boundary. The ordinary
allocating CUDA path similarly folds contiguous conversion, `empty_like`,
pointer/grid extraction, and launch into its native boundary. Compiled runtimes
are registered once per device and dtype with the native extension, which keeps
the kernel and coefficient owners alive. Function/coefficient addresses and
block sizes are parsed to native values at registration; hot in-place calls
therefore launch directly without rebuilding a second Python argument array or
reconverting invariant integers. Hot public calls select that runtime and query
the current raw stream inside the native boundary, bypassing Python runtime
dictionaries and argument unpacking without caching the stream itself.
When exactly one CUDA device is visible, its registered device index object is
also reused instead of querying `Tensor.get_device()` on every launch;
multi-GPU processes retain per-call device selection.
Input dtype/gradient validation is folded into the same boundary; cold or
unavailable-native paths retain Python-side validation before compilation.
The fused kernel
returns
signed zero directly and keeps the common
finite path to one table evaluation plus one exponential; logarithmic scaling
is entered only near overflow.
The native CPU extension is optional when installing from source: if a suitable
C++ compiler is unavailable, installation succeeds and CPU execution falls
back to pure PyTorch. Use `cuda_nvrtc_available()`,
`cuda_native_launcher_available()`, `native_cpu_available()`, and
`triton_available()` to inspect backend capabilities.

## Installation

Install the released package from PyPI:

```bash
pip install erfi-pytorch
```

To install from a local checkout for development and reference tests:

```bash
pip install -e ".[test]"
```

On Linux, install the optional Triton dependency if it is not already
provided by your PyTorch installation:

```bash
pip install -e ".[test,triton]"
```

## Numerical method

For real `x`, the implementation uses

```text
erfi(x) = exp(x^2) Im(w(x)),
```

where `w` is the Faddeeva function. `Im(w(x))` is evaluated with a Taylor
polynomial near zero and a 100-interval table of low-degree polynomial
approximations elsewhere. The native CPU backend evaluates `erfi` directly
with segmented near-minimax polynomials for `|x| <= 4` in float32 and
`|x| <= 2.5` in float64. The x86-64 SSE2 path vectorizes this accurate direct
range while retaining the Faddeeva evaluator outside it. The float32 tail fits
are generated and dense-grid checked by
`tools/generate_direct_float32_tail_coefficients.py`. Near floating-point
overflow, the final magnitude is reconstructed in the log domain so
representable results are not lost to premature overflow in `exp(x^2)`.

The polynomial coefficients originate from Steven G. Johnson's
MIT-licensed Faddeeva implementation. The original license notice is retained
in
[`third_party/faddeeva`](https://github.com/ZhichaoZhu/erfi_pytorch/tree/main/third_party/faddeeva).

The detailed implementation notes are in
[`docs/faddeeva.md`](https://github.com/ZhichaoZhu/erfi_pytorch/blob/main/docs/faddeeva.md).

## License

This project is released under the MIT License. The vendored Faddeeva sources
and material derived from them retain the original Copyright (c) 2012
Massachusetts Institute of Technology attribution and MIT license notice.

## Limitations

- Inputs must be real `torch.float32` or `torch.float64` tensors.
- This release is forward-only. `requires_grad=True` raises an error.
- Triton acceleration currently targets NVIDIA CUDA.
- The fused NVRTC path requires a CUDA-enabled PyTorch installation containing
  NVRTC and a compatible NVIDIA driver; other builds fall back to Triton or
  pure PyTorch as available.

## Benchmark

```bash
python benchmarks/benchmark_erfi.py --dtype float32 --device auto
```

The benchmark covers powers of two from `2^0` through `2^24` by default and
reports CPU timings for `erfi_pytorch.erfi` and `scipy.special.erfi`. When CUDA
is selected and available, it also reports eager PyTorch, compiled PyTorch,
eager dispatch, and compiled dispatch timings. CUDA results include both
per-call synchronization and steady-state CUDA Event batches for allocated,
`out=`, in-place, and fixed-I/O CUDA Graph replay paths. Separate host-wall
batches report submission plus GPU backpressure for allocated, `out=`, and
in-place calls, catching queue-pressure regressions that Event timing alone
cannot expose. Use `--device cpu` on
a CPU-only machine and `--min-exponent`, `--max-exponent`, `--repeats`,
`--cuda-batches`, and `--cuda-batch-calls` to control the run.
Use `--skip-compile` for eager-only latency experiments that should not probe
the optional Triton/Inductor toolchain.
Before timing, it compares the operator against `scipy.special.erfi` and
reports maximum absolute error, maximum and mean relative error, and infinity
mismatches. Infinity mismatches are split into extended finite range and
premature overflow counts so improvements over SciPy are not mistaken for
regressions. Use `--precision-elements` to change the comparison sample count
or `--skip-precision` to run timing only.
