pathtracer-diff / CARD.md
phanerozoic's picture
Card: standardized form
5b3a263 verified
|
Raw
History Blame
9.25 kB
metadata
library_name: kernels
license: apache-2.0

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.