pathtracer-diff

A differentiable Monte Carlo path tracer as one kernel, loadable through kernels. The forward pass is a megakernel unidirectional path tracer: Lambertian diffuse, GGX conductor (VNDF-sampled, height-correlated Smith), smooth dielectric, rough plastic, and rough dielectric materials; area lights with per-texel emission and an importance-sampled equirectangular environment map; an optional homogeneous participating medium; multiple importance sampling (balance heuristic) by default; per-texel albedo with bilinear filtering; a binned-SAH BVH built in vectorized torch in fractions of a second at a million triangles. The backward pass replays the identical paths with the identical counter-based random draws and emits analytic gradients to albedo texels, emission texels, environment texels, and medium coefficients; a companion pass differentiates the direct-lighting image with respect to vertex positions, including shadow and camera silhouettes. Inverse rendering runs as a bare torch loop around the kernel, with no rendering framework in the 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, 30 fps with a progress bar. 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, reported as measured. 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)
panel_emi = torch.full((16, 32, 3), 5.0, device="cuda", requires_grad=True)
scene = ptd.Scene(vertices, faces, material_ids,
                  albedo=[torch.tensor([0.73] * 3),   # constants are 1x1 textures
                          wall_tex,                   # [H, W, 3] per-texel albedo
                          torch.tensor([0.9, 0.6, 0.3])],
                  emission=[torch.zeros(3), torch.zeros(3), panel_emi],
                  uvs=uvs,
                  material_types=[ptd.DIFFUSE, ptd.DIFFUSE, ptd.PLASTIC],
                  roughness=[0.3, 0.3, 0.2])
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)  # [H, W, 3] linear
loss = (img - target).square().mean()
loss.backward()            # gradients on wall_tex and panel_emi

grad_verts = ptd.geometry_grad(scene, cam, dloss_dimage)  # [V, 3] vertex grads

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; that is the configuration an inverse-rendering loop wants. estimator selects "mis" (default), "nee", or "brdf".

API

How it works

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 (the throughput multiplier is then exactly the albedo); conductors sample the GGX visible-normal distribution (Heitz 2018) with the height-correlated Smith term, so the continuation weight is F * G2/G1; smooth dielectrics branch on the Fresnel term; rough plastic mixes a diffuse base with a GGX coat (F0 = 0.04) under a 50/50 lobe-mixture pdf; rough dielectrics sample GGX transmission. With a medium, propagation distances are drawn from a detached exponential rate frozen at Scene construction, and each segment carries the transmittance ratio exp(-sigma_t d) / exp(-sbar d). Direct light is estimated at every NEE-capable vertex (and at medium scattering vertices) by area sampling over the emissive faces and, when an environment map is present, by sampling a detached mixture of its luminance CDF and the uniform sphere; BSDF-sampled hits on emitters and environment misses are combined with the balance heuristic, so all three estimators agree in expectation at any bounce budget.

The backward pass is exact path replay. Sampling distributions depend only on geometry, frozen material parameters (roughness, ior), the construction-time environment CDF and medium rate, and the counter-based Philox stream keyed by (seed, pixel, sample), never on the differentiable parameters, so the backward kernel re-traces identical paths with identical draws and no stored path state. Every radiance term is a product of per-bounce factors, each affine in its vertex's albedo texels (Schlick Fresnel is affine in F0; dielectric factors are constant), times a linear emission or environment texel; the replay differentiates the product in closed form with prefix/suffix exclusion products (no division, so zero albedos are safe) and scatters each factor's gradient into the four texels of its bilinear footprint with the sampling weights, the exact adjoint of the fetch. Medium sigma-gradients are the closed-form log-derivatives -(D_total + d_term) for sigma_a, plus N_scatter/sigma_s for sigma_s. Albedo gradients up to 2048 texels and emission gradients up to 512 texels accumulate in per-block shared memory and flush once; larger textures and environment gradients use direct global atomics.

Geometry gradients are a separate pass over the direct-lighting transport term. The interior (smooth) part re-traces camera rays and light samples under detached sampling and pushes forward-mode dual numbers through the receiver hit, the shading geometry, and the light-area measure, one pass per perturbed vertex coordinate. The boundary (visibility) part is redner-style edge sampling: shadow silhouettes are sampled as (edge point, light point) pairs, projected to the receiver, tested for a clean lit/blocked crossing, and accumulated with the exact projective sweep velocity, the edge-to-curve arc stretch, and the receiver-area-to-pixel measure; camera silhouettes are sampled on the image plane, where the projected edge sweeps front-minus-background direct radiance across pixels.

Correctness

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

• Analytic furnace, zero variance. Closed uniform diffuse box under BRDF sampling: every path returns exactly E * sum a^k; the analytic value 0.498046875 is met with worst pixel deviation 3.4e-3. A diffuse plane under a constant environment reproduces albedo * c the same way.
• Fresnel slabs. A smooth glass slab (ior 1.5) at normal incidence transmits the internal-bounce series (1-R)/(1+R): measured 0.923141 against 0.923077 analytic (0.007%). A rough-dielectric slab at alpha = 0.02 transmits the same series within 0.05%.
• Beer-Lambert. A purely absorbing medium attenuates a uniform emitter to E * exp(-sigma_a d) within 0.035%; with scattering on, MIS and BRDF-only means agree on a foggy Cornell box.
• Estimator agreement. MIS, NEE-only, and BRDF-only means agree within Monte Carlo tolerance on a Cornell box; MIS vs BRDF-only agree for a GGX conductor under an environment map and for rough plastic, which jointly exercises the VNDF weight, the GGX pdf, the mixed-lobe pdf, and the environment pdf.
• Gradients match finite differences along the same paths. Constants 2.2e-4 max relative error over nine entries; wall-texture texels 2.6e-3; conductor F0 through GGX sampling and MIS 9.9e-4; plastic albedo through both lobes 2.8e-4; environment texels 2.1e-4 (detached CDF); emission texels linear-exact; medium sigmas 2e-2 gates. Texels no path can see get exactly zero from both the replay and the differences.
• Geometry gradients match seed-averaged finite differences. Interior: translating an area light matches central differences of the direct image within 15%. Boundary: each occluder vertex's gradient (shadow sweep + its own image silhouette) matches an 8-seed-averaged FD within 0.23 worst-case (6-8% on three of four vertices; a single-seed FD of a visibility discontinuity is flip noise at these scales).
• Constants are 1x1 textures. Both albedo forms render equal images to float rounding and receive matching gradients.
• Inverse rendering. A wall color (3 unknowns) recovers to max error 0.004 in 80 Adam steps; an 8x8 wall texture (192 unknowns) drives the image loss from 1.8e-4 to 7.2e-14; a 4x8 emissive panel texture recovers from the room it lights; all through the kernel alone.
• Deterministic. Repeated forward renders are bitwise identical; gradient buffers accumulate through float atomics and are deterministic in expectation but not bitwise.

Measured

RTX 6000 Ada, 512x512, MIS, 4 bounces. The terrain scenes are displaced grids with a 256x256 checker albedo texture, a GGX conductor block, and a constant environment map; the SAH build is the vectorized torch builder.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+; float32 throughout.
• Materials: Lambertian diffuse, GGX conductor, smooth dielectric, rough plastic (F0 = 0.04 coat), rough dielectric (BSDF-sampled only; GGX roughness clamped to >= 0.01); per-texel emission; pinhole camera; at most 64 materials and 16 bounces; two-sided surfaces; no Russian roulette (fixed bounce budget keeps the replay exact).
• render() differentiates albedo texels/constants (all material types; conductor F0 tint), emission texels, environment texels, and medium sigmas. Camera, UVs, roughness, and ior are frozen.
• geometry_grad() differentiates vertex positions for the direct-lighting term only: diffuse receivers and emitters carry the radiance jumps, indirect bounces are not differentiated, and uv motion across textures is not differentiated. Boundary terms cover shadow silhouettes and camera silhouettes.
• The medium is homogeneous (isotropic phase), fills the scene, and is mutually exclusive with an environment map; its sampling rate is detached at Scene construction (scene.med_sbar).
• The environment sampling CDF is detached at Scene construction (a 0.5 uniform-sphere mixture keeps every direction reachable, so optimizing an env map from an arbitrary initialization stays unbiased); rebuild the Scene to re-importance-sample a changed map.
• Texture and env shapes are fixed at Scene construction; UVs wrap.
• Published variants are Linux x86_64; on Windows load_local.py JIT-builds the same source and exposes the identical API.

References

Kajiya, "The Rendering Equation" (SIGGRAPH 1986); Vicini, Speierer, Jakob, "Path Replay Backpropagation" (SIGGRAPH 2021); Nimier-David, Speierer, Ruiz, Jakob, "Radiative Backpropagation" (SIGGRAPH 2020); Li, Aittala, Durand, Lehtinen, "Differentiable Monte Carlo Ray Tracing through Edge Sampling" (SIGGRAPH Asia 2018); Heitz, "Sampling the GGX Distribution of Visible Normals" (JCGT 2018); Heitz, "Understanding the Masking-Shadowing Function" (JCGT 2014); Walter et al., "Microfacet Models for Refraction through Rough Surfaces" (EGSR 2007); Veach and Guibas, "Optimally Combining Sampling Techniques" (SIGGRAPH 1995); Möller and Trumbore, "Fast, Minimum Storage Ray-Triangle Intersection" (1997); Duff et al., "Building an Orthonormal Basis, Revisited" (JCGT 2017); Wald, "On fast Construction of SAH-based Bounding Volume Hierarchies" (2007).

License

Apache-2.0.