pathtracer-diff

A differentiable Monte Carlo path tracer as one kernel, loadable through kernels. The forward pass is a megakernel unidirectional path tracer with five materials, area lights, an importance-sampled environment map, an optional homogeneous medium, and multiple importance sampling; the backward pass replays the identical paths with identical counter-based draws and emits analytic gradients. The reference standards are analytic transport results and finite differences along the same paths.

Rendering is a simulation of light, and making it differentiable turns it into an instrument: give it a photograph and unknown scene parameters, and gradient descent recovers the materials, the lights, the fog, and the geometry that produced the image. That normally requires a research framework; here it is one loadable kernel, and inverse rendering is a bare torch loop: render, compare, backward(), step.

THE DRAGON'S HALL (media/make_hero.py): one continuous shot, every gradient, optimized live through the kernel against the fixed TARGET inset. From a flat gray start, Adam through render() recovers the 128x128 marble floor albedo (49,152 unknowns), the gold trim's conductor tint, the pedestal's plastic albedo, the braziers' warm emission (max err 1.75 of 14), and a 16x64 per-texel emissive rune band (3,072 unknowns, mean err 0.15), around a smooth glass slab and the rough-dielectric crystal dragon. Mid-shot, geometry_grad() alone takes over the dragon's position: 5,205 vertices slide 4.60 scene units onto the pedestal by shadow and silhouette alone, final offset 0.004. In the last phase the participating medium rolls in: sigma_a and sigma_s recover from the hazy target to extinction error 0.014. Image loss falls 8.5e-1 to 1.9e-1 at spp 8; the dim-hall floor and trim retain residual texel error where the scene barely observes them. Stanford dragon courtesy of the Stanford Computer Graphics Laboratory.

Usage

import torch
from kernels import get_kernel

ptd = get_kernel("phanerozoic/pathtracer-diff", version=1, trust_remote_code=True)

wall_tex = torch.full((256, 256, 3), 0.5, device="cuda", requires_grad=True)
scene = ptd.Scene(vertices, faces, material_ids,
                  albedo=[torch.tensor([0.73] * 3), wall_tex],
                  emission=[torch.zeros(3), torch.zeros(3)],
                  uvs=uvs, material_types=[ptd.DIFFUSE, ptd.DIFFUSE],
                  roughness=[0.3, 0.3])
cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8),
                 vfov_deg=39.0)

img = ptd.render(scene, cam, 512, 512, spp=64, max_bounces=4)
loss = (img - target).square().mean()
loss.backward()                       # gradients on wall_tex

grad_verts = ptd.geometry_grad(scene, cam, dloss_dimage)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark. A fixed seed renders the same paths every call, so the Monte Carlo objective is a deterministic function of the parameters and gradient descent sees a smooth landscape at any spp.

API

Symbol Purpose
Scene(vertices, faces, material_ids, albedo, emission, uvs, material_types, roughness, ior, env, medium) triangle-mesh scene; BVH, light list, and the detached environment CDF / medium rate built at construction
Camera(position, look_at, up, vfov_deg) pinhole camera
render(scene, camera, height, width, spp, max_bounces, estimator, seed) differentiable render to [H, W, 3] f32 linear; autograd to albedo, emission, environment texels, and medium sigmas
geometry_grad(scene, camera, grad_image, spp, edge_samples, seed) d(loss)/d(vertex positions) for the direct-lighting term: dual-number interior + edge-sampled shadow and camera silhouettes
DIFFUSE, CONDUCTOR, DIELECTRIC, PLASTIC, ROUGH_DIELECTRIC material type constants
ops.pt_forward / ops.pt_backward / ops.pt_geometry_grad raw kernel launches

Method

One thread owns one pixel and accumulates its spp samples serially, so the image is bitwise deterministic with no atomics in the forward pass. Diffuse vertices sample the cosine hemisphere; conductors sample the GGX visible-normal distribution with the height-correlated Smith term; smooth dielectrics branch on Fresnel; rough plastic mixes a diffuse base with a GGX coat under a 50/50 lobe-mixture pdf; rough dielectrics sample GGX transmission. Direct light is estimated at every NEE-capable vertex by area sampling over emissive faces and, with an environment map, by a detached mixture of its luminance CDF and the uniform sphere; BSDF-sampled hits combine by the balance heuristic.

The backward pass is exact path replay: sampling distributions depend only on geometry, frozen parameters, and the counter-based Philox stream, never on the differentiable parameters, so the backward kernel re-traces identical paths with no stored path state. Every radiance term is a product of per-bounce factors, each affine in its vertex's albedo texels, times a linear emission or environment texel; the replay differentiates the product in closed form with prefix/suffix exclusion products and scatters each factor's gradient into the four texels of its bilinear footprint. Geometry gradients run a separate pass over the direct-lighting term: forward-mode duals for the smooth part, redner-style edge sampling for shadow and camera silhouettes.

Measured

RTX 6000 Ada, 512x512, MIS, 4 bounces; terrain scenes are displaced grids with a 256x256 checker albedo, a GGX conductor block, and a constant environment, built by the vectorized torch SAH builder:

scene build forward backward (replay)
Cornell box, 24 tris, 64 spp - 12.7 ms (1.32 Gpaths/s) 70.2 ms (5.5x)
terrain, 522,254 tris, 16 spp 0.2 s 18.3 ms (229 Mpaths/s) 39.0 ms (2.1x)
terrain, 1,045,470 tris, 16 spp 0.3 s 20.3 ms (207 Mpaths/s) 40.8 ms (2.0x)

Correctness

Measured on RTX 6000 Ada from source; the published v1 variants pass the full 22-test suite through get_kernel on a GeForce RTX 3070 Ti (sm86):

  • Analytic furnace, zero variance: a closed uniform diffuse box returns exactly E * sum a^k per path; the analytic 0.498046875 is met with worst pixel deviation 3.4e-3.
  • Fresnel slabs: a smooth glass slab at normal incidence transmits the internal-bounce series (1-R)/(1+R), measured 0.923141 against 0.923077 (0.007%); a rough-dielectric slab at alpha 0.02 within 0.05%.
  • Beer-Lambert: a purely absorbing medium attenuates a uniform emitter to E * exp(-sigma_a d) within 0.035%.
  • Estimator agreement: MIS, NEE-only, and BRDF-only means agree within Monte Carlo tolerance across Cornell, GGX-under-environment, and rough-plastic scenes.
  • Gradients match finite differences along the same paths: constants 2.2e-4 max relative over nine entries; wall texels 2.6e-3; conductor F0 through GGX and MIS 9.9e-4; environment texels 2.1e-4; emission linear-exact; medium sigmas within 2e-2 gates. Texels no path can see get exactly zero from both sides.
  • Geometry gradients match seed-averaged finite differences: light translation within 15%; occluder vertices within 0.23 worst case.
  • Inverse rendering: a wall color recovers to max error 0.004 in 80 Adam steps; an 8x8 texture drives the loss from 1.8e-4 to 7.2e-14.
  • Deterministic: repeated forward renders are bitwise identical; gradient buffers accumulate through float atomics and are deterministic in expectation but not bitwise.

Requirements and limits

  • NVIDIA GPU with compute capability 8.0+; float32 throughout.
  • At most 64 materials and 16 bounces; two-sided surfaces; no Russian roulette (the fixed bounce budget keeps the replay exact); pinhole camera.
  • render() differentiates albedo, emission, environment texels, and medium sigmas; camera, UVs, roughness, and ior are frozen. geometry_grad() covers the direct-lighting term only.
  • The medium is homogeneous, fills the scene, and is mutually exclusive with an environment map; the environment CDF is detached at Scene construction.
  • Published variants are Linux x86_64; on Windows load_local.py JIT-builds the same source.

References

Kajiya (1986); Vicini et al., "Path Replay Backpropagation" (2021); Nimier-David et al., "Radiative Backpropagation" (2020); Li et al., "Differentiable Monte Carlo Ray Tracing through Edge Sampling" (2018); Heitz (2014, 2018) on GGX and VNDF; Walter et al. (2007); Veach and Guibas (1995); Möller and Trumbore (1997); Wald (2007) on SAH construction.

License

Apache-2.0.

Downloads last month
-
apache-2.0
Supported hardwares new
CUDA
8.08.68.99.010.012.0
GPU
B300
288GB
NVIDIA SXM
B200
192GB
NVIDIA SXM
H200
141GB
NVIDIA SXM
H100
80GB
GPU
H800
80GB
GPU
H20
96GB
GPU
L40s
48GB
GPU
L40
48GB
GPU
L20
48GB
GPU
L4
24GB
DGX Spark
GB10
128GB
GPU
RTX PRO 6000 WS
96GB
GPU
RTX PRO 6000 Max-Q
96GB
GPU
RTX PRO 5000
48GB
GPU
RTX PRO 4500 WS
32GB
GPU
RTX PRO 4000
24GB
GPU
RTX PRO 4000 SFF
24GB
GPU
RTX PRO 2000
16GB
GPU
RTX 6000 Ada
48GB
GPU
RTX 5880 Ada
48GB
RTX
RTX 5000 Ada
32GB
GPU
RTX 4500 Ada
24GB
RTX
RTX 4000 Ada
20GB
RTX
RTX 4000 SFF Ada
20GB
GPU
RTX 3500 Ada Mobile
12GB
GPU
RTX 2000 Ada
16GB
GPU
RTX A6000
48GB
GPU
RTX A5000
8GB
GPU
RTX A5000 Max-Q
16GB
GPU
RTX A5000 Mobile
16GB
GPU
RTX A4000
16GB
GPU
RTX A4000 Max-Q
8GB
GPU
RTX A4000 Mobile
8GB
GPU
RTX A3000 Mobile
6GB
GPU
RTX A2000
6GB
GPU
RTX A2000 Embedded
4GB
GPU
RTX A2000 Max-Q
4GB
GPU
RTX A2000 Mobile
4GB
GPU
A800
40GB
GPU
A100
80GB
GPU
A40
48GB
GPU
A30
24GB
GPU
A10
24GB
GPU
A2
16GB
RTX
RTX 5090
32GB
RTX
RTX 5090 D
32GB
RTX
RTX 5090 Mobile
24GB
RTX
RTX 5080
16GB
RTX
RTX 5080 Mobile
16GB
RTX
RTX 5070
12GB
RTX
RTX 5070 Mobile
8GB
RTX
RTX 5070 Ti
16GB
RTX
RTX 5070 Ti Mobile
12GB
RTX
RTX 5060 Ti
16GB
RTX
RTX 5060
8GB
RTX
RTX 5060 Mobile
8GB
RTX
RTX 5050
8GB
RTX
RTX 5050 Mobile
8GB
RTX
RTX 4090
24GB
RTX
RTX 4090D
24GB
RTX
RTX 4090 Mobile
16GB
RTX
RTX 4080 SUPER
16GB
RTX
RTX 4080
16GB
RTX
RTX 4080 Mobile
12GB
RTX
RTX 4070
12GB
RTX
RTX 4070 Mobile
8GB
RTX
RTX 4070 Ti
12GB
RTX
RTX 4070 Super
12GB
RTX
RTX 4070 Ti Super
16GB
RTX
RTX 4060
8GB
RTX
RTX 4060 Ti
8GB
RTX
RTX 4090 Laptop
16GB
RTX
RTX 4080 Laptop
12GB
RTX
RTX 4070 Laptop
8GB
RTX
RTX 4060 Laptop
8GB
RTX
RTX 4050 Laptop
6GB
RTX
RTX 3090
24GB
RTX
RTX 3090 Ti
24GB
RTX
RTX 3080
12GB
RTX
RTX 3080 Ti
12GB
RTX
RTX 3080 Mobile
16GB
RTX
RTX 3070
8GB
RTX
RTX 3070 Ti
8GB
RTX
RTX 3070 Ti Mobile
8GB
RTX
RTX 3060 Ti
8GB
RTX
RTX 3060
12GB
RTX
RTX 3060 Mobile
6GB
RTX
RTX 3050 Mobile
4GB
GPU
RTX 2050 Mobile
4GB
Jetson
Jetson AGX Orin 64GB
64GB
Jetson
Jetson AGX Orin 32GB
32GB
Jetson
Jetson Orin NX 16GB
16GB
Jetson
Jetson Orin NX 8GB
8GB
Jetson
Jetson Orin Nano 8GB
8GB
Jetson
Jetson Orin Nano 4GB
4GB
OS
linux
Arch
x86_64
Kernel Builder
19aaa64