Functor GPU JIT

ESCAPE can evaluate functors on a GPU for large arrays. Set f.device = "gpu" and the library builds a GPU kernel the first time you need it, then reuses that kernel from a cache.

This is useful when you evaluate the same expression on millions of points. For small arrays or single numbers, the CPU path is usually simpler and fast enough.

Requirements

• An NVIDIA GPU with a working driver
• The CUDA Toolkit installed (so ESCAPE can build GPU kernels; typically via CUDA_PATH / CUDA_HOME, or NVIDIA tools on PATH)
• On Windows, Visual Studio build tools (same setup used to build ESCAPE)

The first GPU evaluation of a new expression can take tens of seconds while the kernel is generated. Later calls on the same functor are much faster because the kernel is cached.

If CUDA is not installed or not visible to Jupyter, setting f.device = "gpu" will fail when the kernel is built. Make sure NVIDIA's tools folder is on the PATH that Jupyter uses (it can differ from your terminal PATH).

import numpy as np
import escape as esc

esc.require("0.9.8")

Loading material database from C:\dev\escape-core\python\src\escape\scattering\..\data\mdb\materials.db

Quick start

Build a functor as usual, then switch it to the GPU:

x = esc.var("x")
f_jit = esc.sin(x) + 1.0
f_jit.device = "gpu"

print("Scalar value f_jit(0.5) =", f_jit(0.5))
print("Current device:", f_jit.device)
f_jit.show()

Scalar value f_jit(0.5) = 1.4794254302978516
Current device: gpu

How it works

1. f.device = "gpu" — array evaluation uses a GPU kernel. ESCAPE generates that kernel on demand and caches it.
2. f.device = "cpu" (default) — array evaluation stays on the CPU.
3. Scalars — single-point calls always use the CPU expression path; GPU mode is for arrays.

For fair timings, pass contiguous float64 arrays (for example np.linspace(...).copy()). That avoids extra host-side copies and warning messages.

Reading the benchmarks

GPU timings are not always faster than CPU — that can be normal with the current GPU path.

What each GPU call does today

For every array evaluation, ESCAPE moves data to the GPU, runs the kernel, waits for it to finish, then copies results back. That fixed overhead matters a lot for smaller arrays.

When GPU tends to help

Very large arrays (tens of millions of points) and expressions that are costly per point. Smaller jobs often stay faster on CPU.

How timings are shown below

Print labels name the device and array size clearly. %timeit reports wall time for evaluating the full array once the kernel is already built (run the accuracy cell first so the first-time build cost is not mixed into the comparison).

Example 1 — Gaussian peak (large arrays)

$$f(x) = A \exp\!\left(-\frac{(x - \mu)^2}{2\sigma^2}\right)$$

We compare CPU and GPU on the same functor by switching device. Accuracy is checked first (max absolute difference vs CPU), then wall time for a large array.

x = esc.var("x")
A = esc.par("A", 1.0)
mu = esc.par("mu", 0.0, userlim=[-1, 1])
sigma = esc.par("sigma", 1.0)
f_gauss = A * esc.exp(-(x - mu) ** 2 / (2 * sigma ** 2))


def make_x(n: int) -> np.ndarray:
    return np.linspace(-8.0, 8.0, n, dtype=np.float64).copy()


def check_accuracy(n: int):
    xb = make_x(n)
    f_gauss.device = "cpu"
    y_cpu = np.asarray(f_gauss(xb), dtype=np.float64)
    f_gauss.device = "gpu"
    y_gpu = np.asarray(f_gauss(xb), dtype=np.float64)
    return float(np.max(np.abs(y_gpu - y_cpu)))


print("Accuracy (max |GPU - CPU|):")
for n in (1_000_000, 5_000_000, 10_000_000):
    err = check_accuracy(n)
    print(f"  n = {n:>10,}   max|gpu-cpu| = {err:.3e}")

Accuracy (max |GPU - CPU|):
  n =  1,000,000   max|gpu-cpu| = 1.184e-07
  n =  5,000,000   max|gpu-cpu| = 1.196e-07
  n = 10,000,000   max|gpu-cpu| = 1.206e-07

f_gauss.device = "gpu"
f_gauss.show()

n = 100_000_000
xb = make_x(n)

f_gauss.device = "cpu"
print(f"Timing — Gaussian, n = {n:,}")
print("  CPU:")
%timeit f_gauss(xb)

f_gauss.device = "gpu"
print("  GPU:")
%timeit f_gauss(xb)

Timing — Gaussian, n = 100,000,000
  CPU:
366 ms ± 29.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
  GPU:
195 ms ± 2.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Example 2 — Mixed trig + exp

$$f(x) = a \sin(b x) + e^{-x^2}$$

x = esc.var("x")
a = esc.par("a", 2.5)
b = esc.par("b", -1.3)
f_mix = a * esc.sin(b * x) + esc.exp(-x ** 2)

n = 8_000_000
xb = np.linspace(-4.0, 4.0, n, dtype=np.float64).copy()

f_mix.device = "cpu"
y_cpu = np.asarray(f_mix(xb), dtype=np.float64)
f_mix.device = "gpu"
y_gpu = np.asarray(f_mix(xb), dtype=np.float64)

print(f"Accuracy — mixed, n = {n:,}")
print(f"  max|gpu-cpu| = {float(np.max(np.abs(y_gpu - y_cpu))):.3e}")

Accuracy — mixed, n = 8,000,000
  max|gpu-cpu| = 2.500e-06

print(f"Timing — mixed, n = {n:,}")
f_mix.device = "cpu"
print("  CPU:")
%timeit f_mix(xb)

f_mix.device = "gpu"
print("  GPU:")
%timeit f_mix(xb)

Timing — mixed, n = 8,000,000
  CPU:
28.4 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  GPU:
15.1 ms ± 113 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Example 3 — Decaying exponential (10M points)

x = esc.var("x")
a = esc.par("a", 1.0)
f_exp = a * esc.exp(-x ** 2)

n = 10_000_000
xb = np.linspace(-12.0, 12.0, n, dtype=np.float64).copy()

print(f"Timing — exponential, n = {n:,}")
f_exp.device = "cpu"
print("  CPU:")
%timeit f_exp(xb)

f_exp.device = "gpu"
print("  GPU:")
%timeit f_exp(xb)

Timing — exponential, n = 10,000,000
  CPU:
27.5 ms ± 1.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
  GPU:
19.6 ms ± 433 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Example 4 — Switching back to CPU

The same functor can move between devices. Setting device = "cpu" turns off the GPU array path — useful for debugging or A/B checks.

x = esc.var("x")
f = x * 2.0

n = 2_000_000
xb = np.linspace(-3.0, 3.0, n, dtype=np.float64).copy()

f.device = "cpu"
y_ref = np.asarray(f(xb), dtype=np.float64)

f.device = "gpu"
y_gpu = np.asarray(f(xb), dtype=np.float64)
f.device = "cpu"
y_cpu = np.asarray(f(xb), dtype=np.float64)

print(f"Accuracy — linear, n = {n:,}")
print(f"  max|gpu-ref| = {float(np.max(np.abs(y_gpu - y_ref))):.3e}")
print(f"  max|cpu-ref| = {float(np.max(np.abs(y_cpu - y_ref))):.3e}")

print()
print(f"Timing — linear, n = {n:,}")
f.device = "gpu"
print("  GPU:")
%timeit f(xb)
f.device = "cpu"
print("  CPU:")
%timeit f(xb)

Accuracy — linear, n = 2,000,000
  max|gpu-ref| = 2.384e-07
  max|cpu-ref| = 0.000e+00

Timing — linear, n = 2,000,000
  GPU:
3.96 ms ± 31.1 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
  CPU:
1.95 ms ± 22.7 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Example 5 — Two variables

For two variables, pack points as [x0, y0, x1, y1, ...]. A moderate 2D grid is enough for a quick check (full large tensors use a lot of memory).

x = esc.var("x")
y = esc.var("y")
f2 = esc.exp(-x) * y + 1.0
f2.device = "gpu"

nx = ny = 256
xs = np.linspace(-1.0, 1.0, nx, dtype=np.float64)
ys = np.linspace(-1.0, 1.0, ny, dtype=np.float64)
X, Y = np.meshgrid(xs, ys, indexing="xy")
pts = np.column_stack([X.ravel(), Y.ravel()]).ravel().copy()

z = np.asarray(f2(pts), dtype=np.float64)
z_ref = (np.exp(-X) * Y + 1.0).ravel()

print(f"Accuracy — 2D grid {nx}x{ny} ({pts.size // 2:,} points)")
print(f"  max|gpu-ref| = {float(np.max(np.abs(z - z_ref))):.3e}")

print("Timing — GPU:")
%timeit f2(pts)

Accuracy — 2D grid 256x256 (65,536 points)
  max|gpu-ref| = 3.819e-07
Timing — GPU:
449 μs ± 6.31 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

API summary

Item	Meaning
f.device = "gpu"	Use GPU for array evaluation (kernel built on first use, then cached)
f.device = "cpu"	Use CPU for array evaluation (default)

See also

• GPU JIT unit tests live under python/tests/core/ (GPU device tests)

Distribution functors on GPU

Compare CPU and GPU for the built-in distributions (normal, gamma, schulz, lognorm, uniform, triangular). Small differences (around 1e-7 to 1e-8) are normal floating-point variation between devices.

x = esc.var("x")
mean_p = esc.par("mean", 10.0)
sigma_p = esc.par("sigma", 2.0)
fwhm_p = esc.par("fwhm", 4.0)

distributions = {
    "normal":     esc.normal(x, mean_p, fwhm_p),
    "gamma":      esc.gamma(x, mean_p, sigma_p),
    "schulz":     esc.schulz(x, mean_p, sigma_p),
    "lognorm":    esc.lognorm(x, mean_p, sigma_p),
    "uniform":    esc.uniform(x, mean_p, fwhm_p),
    "triangular": esc.triangular(x, mean_p, fwhm_p),
}

xarr = np.linspace(0.1, 25.0, 500, dtype=np.float64).copy()
tol = 1e-5

print(f"Distributions — n = {xarr.size:,}  (pass if max|cpu-gpu| < {tol:.0e})")
for name, f in distributions.items():
    f.device = "cpu"
    cpu_vals = np.asarray(f(xarr), dtype=np.float64)
    f.device = "gpu"
    gpu_vals = np.asarray(f(xarr), dtype=np.float64)
    f.device = "cpu"
    max_diff = float(np.max(np.abs(cpu_vals - gpu_vals)))
    status = "OK" if max_diff < tol else "CHECK"
    print(f"  {name:12s}  max|cpu-gpu| = {max_diff:.2e}  {status}")

Distributions — n = 500  (pass if max|cpu-gpu| < 1e-05)
  normal        max|cpu-gpu| = 4.87e-08  OK
  gamma         max|cpu-gpu| = 1.06e-07  OK
  schulz        max|cpu-gpu| = 1.44e-07  OK
  lognorm       max|cpu-gpu| = 2.05e-07  OK
  uniform       max|cpu-gpu| = 0.00e+00  OK
  triangular    max|cpu-gpu| = 2.97e-08  OK

Convolution on GPU

$$h(x) = \int f(\tau)\, g(x - \tau)\, d\tau$$

x = esc.var("x")
tau = esc.var("tau")
mean_f = esc.par("mean_f", 10.0)
sigma_f = esc.par("sigma_f", 1.0)
mean_g = esc.par("mean_g", 0.0)
sigma_g = esc.par("sigma_g", 1.5)

f_conv = esc.normal(tau, mean_f, sigma_f)
g_conv = esc.normal(x - tau, mean_g, sigma_g)

h_conv = esc.convolution(
    f_conv, g_conv, tau, 0.0, 30.0, numpoints=61, epsabs=1e-8, epsrel=1e-8, maxiter=200
)

xarr = np.linspace(2.0, 18.0, 200, dtype=np.float64).copy()
tol = 1e-6

h_conv.device = "cpu"
cpu_vals = np.asarray(h_conv(xarr), dtype=np.float64)
h_conv.device = "gpu"
gpu_vals = np.asarray(h_conv(xarr), dtype=np.float64)
h_conv.device = "cpu"

max_diff = float(np.max(np.abs(cpu_vals - gpu_vals)))
status = "OK" if max_diff < tol else "CHECK"
print(f"Convolution — n = {xarr.size}")
print(f"  max|cpu-gpu| = {max_diff:.2e}  {status}")

Convolution — n = 200
  max|cpu-gpu| = 1.17e-38  OK

Weighted average on GPU

$$H(x) = \int f(\tau)\, g(x, \tau)\, d\tau$$

API: esc.average_normal(f, fwhm, x, ...). Below we average a gently oscillating integrand sin(x)^2 over a normal resolution width, sweep several FWHM values, then overlay CPU and GPU curves for one representative width.

x = esc.var("x")
xarr = np.linspace(1.0, 10.0, 200, dtype=np.float64).copy()
tol = 2e-5

print(f"Weighted average of sin(x)^2 — n = {xarr.size}  (pass if max|cpu-gpu| < {tol:.0e})")
print(f"{'fwhm':>8}  {'max|cpu-gpu|':>14}  status")

for fwhm_val in (0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0):
    fwhm = esc.par("fwhm", float(fwhm_val))
    h_avg = esc.average_normal(
        esc.sin(x) ** 2, fwhm, x, epsrel=1e-8, epsabs=1e-8, maxiter=None, numpoints=121
    )
    h_avg.device = "cpu"
    cpu_vals = np.asarray(h_avg(xarr), dtype=np.float64)
    h_avg.device = "gpu"
    gpu_vals = np.asarray(h_avg(xarr), dtype=np.float64)
    h_avg.device = "cpu"
    max_diff = float(np.max(np.abs(cpu_vals - gpu_vals)))
    status = "OK" if max_diff < tol else "CHECK"
    print(f"{fwhm_val:8g}  {max_diff:14.3e}  {status}")

Weighted average of sin(x)^2 — n = 200  (pass if max|cpu-gpu| < 2e-05)
    fwhm    max|cpu-gpu|  status
    0.01       8.542e-06  OK
    0.02       1.237e-05  OK
    0.05       2.847e-06  OK
     0.1       1.615e-06  OK
     0.2       1.065e-06  OK
     0.5       9.395e-07  OK
       1       6.409e-07  OK
       2       3.289e-07  OK

# Overlay CPU vs GPU for a gentle oscillation (fwhm = 0.5)
fwhm_plot = esc.par("fwhm", 0.5, userlim=[0.01, 5.0])
h_gentle = esc.average_normal(
    esc.sin(x) ** 2, fwhm_plot, x, epsrel=1e-8, epsabs=1e-8, maxiter=None, numpoints=121
)
coords = np.linspace(0.5, 8.0, 400, dtype=np.float64).copy()

h_gentle.device = "cpu"
y_cpu = np.asarray(h_gentle(coords), dtype=np.float64)
h_gentle.device = "gpu"
y_gpu = np.asarray(h_gentle(coords), dtype=np.float64)

max_diff = float(np.max(np.abs(y_cpu - y_gpu)))
print(f"Gentle sin(x)^2 — max|cpu-gpu| = {max_diff:.3e}")

esc.overlay(y_cpu, y_gpu, coordinates=coords).config(
    title="Weighted average: sin(x)^2 (gentle)",
    xlabel="x",
    ylabel="H(x)",
    labels=["CPU", "GPU"],
)

Gentle sin(x)^2 — max|cpu-gpu| = 9.004e-07

Rapid oscillations — when GPU can look worse

If the integrand oscillates many times under the resolution window, the GPU path (float32 quadrature nodes/weights and --use_fast_math) accumulates more rounding error than the CPU double-precision integrator. The curves can then separate visibly even though the same expression is evaluated.

Below: sin(1000 x)^2 with the same averaging setup. Expect a much larger max|cpu-gpu| than for sin(x)^2.

# Same averaging API, but a rapidly oscillating integrand
fwhm_fast = esc.par("fwhm_fast", 0.5, userlim=[0.01, 5.0])
h_fast = esc.average_normal(
    esc.sin(1000 * x) ** 2,
    fwhm_fast,
    x,
    epsrel=1e-8,
    epsabs=1e-8,
    maxiter=None,
    numpoints=121,
)
coords_fast = np.linspace(1.0, 10.0, 400, dtype=np.float64).copy()

h_fast.device = "cpu"
y_cpu_fast = np.asarray(h_fast(coords_fast), dtype=np.float64)
h_fast.device = "gpu"
y_gpu_fast = np.asarray(h_fast(coords_fast), dtype=np.float64)

max_diff_fast = float(np.max(np.abs(y_cpu_fast - y_gpu_fast)))
mean_abs = float(np.mean(np.abs(y_cpu_fast)))
rel = max_diff_fast / mean_abs if mean_abs > 0 else float("nan")

print("Rapid sin(1000 x)^2 — CPU vs GPU")
print(f"  max|cpu-gpu| = {max_diff_fast:.3e}")
print(f"  mean|cpu|    = {mean_abs:.3e}")
print(f"  relative     = {rel:.3e}")
print()
print(
    "Why GPU can be worse here: the GPU integration path uses float32 Gauss-Kronrod "
    "tables and fast math. With ~1000 oscillations under the resolution window, "
    "those float errors add up. The CPU path stays in double precision, so the "
    "two results diverge. Prefer a smoother integrand, or accept larger tolerance "
    "when the physical signal truly oscillates this fast."
)

esc.overlay(y_cpu_fast, y_gpu_fast, coordinates=coords_fast).config(
    title="Weighted average: sin(1000 x)^2 (rapid) — CPU vs GPU",
    xlabel="x",
    ylabel="H(x)",
    labels=["CPU", "GPU"],
)

Rapid sin(1000 x)^2 — CPU vs GPU
  max|cpu-gpu| = 1.298e-04
  mean|cpu|    = 5.000e-01
  relative     = 2.596e-04

Why GPU can be worse here: the GPU integration path uses float32 Gauss-Kronrod tables and fast math. With ~1000 oscillations under the resolution window, those float errors add up. The CPU path stays in double precision, so the two results diverge. Prefer a smoother integrand, or accept larger tolerance when the physical signal truly oscillates this fast.