erfi-pytorch

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

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:

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:

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:

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:

pip install erfi-pytorch

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

pip install -e ".[test]"

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

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

Numerical method

For real x, the implementation uses

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.

The detailed implementation notes are in 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

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.