1import json, time
 2import numpy as np
 3import torch
 4
 5SEED = 2025
 6torch.manual_seed(SEED); np.random.seed(SEED)
 7device = "cuda" if torch.cuda.is_available() else "cpu"
 8try:
 9    if device == "cuda": torch.cuda.set_device(0); torch.cuda.empty_cache()
10except Exception:
11    device = "cpu"
12
13
14def _work_dtype(x):
15    return torch.float64 if x.dtype == torch.float64 else torch.float32
16
17def truncated(x, c, xmin=-2.0, h=0.25):
18    """Vectorized fixed truncated-power cubic cardinal spline contraction."""
19    n = c.shape[-1]; dt = _work_dtype(x)
20    xx, cc = x.to(dt), c.to(dt)
21    u = ((xx - xmin) / h).clamp(0., float(n))
22    i = torch.arange(n, device=x.device, dtype=dt)
23    r = torch.arange(5, device=x.device, dtype=dt)
24    a = torch.tensor([1., -4., 6., -4., 1.], device=x.device, dtype=dt)
25    z = u[..., None, None] - i[None, :, None] - r[None, None, :]
26    basis = (torch.relu(z).pow(3) * a).sum(-1) / 6.
27    return (basis * cc).sum(-1).to(x.dtype)
28
29def cox_de_boor(x, c, xmin=-2.0, h=0.25):
30    """Reference Cox-de Boor evaluation of B_3(u-i), i=0,...,N-1."""
31    n = c.shape[-1]; dt = _work_dtype(x)
32    xx, cc = x.to(dt), c.to(dt)
33    u = ((xx - xmin) / h).clamp(0., float(n))
34    m = n + 3
35    knots = torch.arange(m, device=x.device, dtype=dt)
36    b = ((u[..., None] >= knots) & (u[..., None] < knots + 1.)).to(dt)
37    for d in range(1, 4):
38        ii = torch.arange(m-d, device=x.device, dtype=dt)
39        b = ((u[..., None]-ii)/d)*b[..., :m-d] + ((ii+d+1.-u[..., None])/d)*b[..., 1:m-d+1]
40    return (b[..., :n] * cc).sum(-1).to(x.dtype)
41
42def sync():
43    if device == "cuda": torch.cuda.synchronize()
44
45def timed(fn, x, c, warm=30, reps=150):
46    for _ in range(warm): fn(x,c)
47    sync(); t0=time.perf_counter()
48    for _ in range(reps): fn(x,c)
49    sync(); return (time.perf_counter()-t0)*1000/reps
50
51def main():
52    n=16
53    x=torch.randn(4096,device=device)*1.8-.3; c=torch.randn(4096,n,device=device)
54    x1=x.clone().requires_grad_(True); x2=x.clone().requires_grad_(True)
55    y1,y2=truncated(x1,c),cox_de_boor(x2,c)
56    g1=torch.autograd.grad(y1.sum(),x1)[0]; g2=torch.autograd.grad(y2.sum(),x2)[0]
57    xo=torch.tensor([-100.,-10.,0.,10.,100.],device=device); yo=truncated(xo,torch.ones(5,n,device=device))
58    xb=torch.randn(128,32,32,device=device)*1.5-.2; cb=torch.randn(128,32,32,n,device=device)
59    ft=timed(truncated,xb,cb); fc=timed(cox_de_boor,xb,cb)
60    def step(fn):
61        xx=xb.detach().requires_grad_(True); cc=cb.detach().requires_grad_(True)
62        (fn(xx,cc).square().mean()).backward()
63    st=timed(lambda a,b:step(truncated),xb,cb,warm=8,reps=30)
64    sc=timed(lambda a,b:step(cox_de_boor),xb,cb,warm=8,reps=30)
65    # Valid precision comparison; unlike the original version, float64 is preserved.
66    xd=(torch.rand(20000,device=device,dtype=torch.float64)*16-0)*.25-2
67    cd=torch.randn(20000,n,device=device,dtype=torch.float64)
68    ed=(truncated(xd,cd)-cox_de_boor(xd,cd)).abs()
69    result={"device":device,"n":n,"forward_max_abs_error":(y1-y2).abs().max().item(),
70      "forward_relative_mean_error":(y1-y2).abs().mean().item()/(y2.abs().mean().item()+1e-12),
71      "gradient_max_abs_error":(g1-g2).abs().max().item(),"out_of_range_finite":bool(torch.isfinite(yo).all()),
72      "out_of_range_outputs":yo.detach().cpu().tolist(),"forward_ms_truncated":ft,"forward_ms_cox_de_boor":fc,
73      "forward_speedup":fc/ft,"step_ms_truncated":st,"step_ms_cox_de_boor":sc,"step_speedup":sc/st,
74      "float64_max_abs_error":ed.max().item(),"float64_mean_abs_error":ed.mean().item()}
75    print(json.dumps(result,indent=2))
76
77if __name__=='__main__': main()