Per-texel albedo textures with bilinear-footprint adjoints

Browse files

Files changed (10) hide show

CARD.md +112 -90
README.md +113 -91
dev/report_local.py +199 -29
local_bind.cpp +90 -42
pathtracer_diff_cuda/pathtracer.cu +538 -233
pathtracer_diff_cuda/pathtracer_launch.h +37 -21
tests/test_pathtracer_diff.py +212 -0
torch-ext/pathtracer_diff/__init__.py +380 -81
torch-ext/torch_binding.cpp +118 -57
torch-ext/torch_binding.h +14 -5

CARD.md CHANGED Viewed

@@ -6,18 +6,17 @@ license: apache-2.0
 # pathtracer-diff
 A differentiable Monte Carlo path tracer as one kernel, loadable through
-`kernels`. The forward pass renders a linear-radiance image with a megakernel
-unidirectional path tracer (diffuse + emissive materials, per-texel albedo
-textures with bilinear filtering, cosine-weighted BRDF sampling, optional
-next-event estimation, BVH traversal); the backward pass replays the
-identical paths with the identical counter-based random draws and emits
-analytic gradients to the albedo texels and the per-material emission.
-Inverse rendering then runs as a bare torch loop around the kernel, with no
-rendering framework in the loop: render, compare to a target, `backward()`,
-step. Per-texel albedo is what makes the inverse problem non-trivial: a
-texture is thousands of unknowns estimated through indirect illumination,
-where per-material constants make the pixel values merely multilinear in a
-handful of scalars.
 ## Usage
@@ -28,67 +27,73 @@ 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)
-albedo = [torch.tensor([0.73] * 3),      # constants are 1x1 textures
-          torch.tensor([0.65, 0.05, 0.05]),
-          wall_tex]                      # [H, W, 3] per-texel albedo
-scene = ptd.Scene(vertices, faces, material_ids, albedo, emission, uvs=uvs)
 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 emission)
 ```
 `version` selects the release branch; `trust_remote_code` is required by
 `kernels` for publishers without the trusted-publisher mark.
-`albedo` is either an `[M, 3]` tensor of per-material constants or a list of
-M per-material textures `[Hm, Wm, 3]`; constants are treated as 1x1 textures,
-so both forms share one code path and both receive gradients. A fixed `seed`
-renders with 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.
 ## API
 | Symbol | Purpose |
 |---|---|
-| `Scene(vertices, faces, material_ids, albedo, emission, uvs)` | triangle-mesh scene; the BVH and the emissive-face light list are built at construction; `uvs` is `[V, 2]` per vertex or `[F, 3, 2]` per corner, textures wrap (repeat); texture shapes are fixed at construction |
 | `Camera(position, look_at, up, vfov_deg)` | pinhole camera |
-| `render(scene, camera, height, width, spp, max_bounces, nee, seed)` | differentiable render to `[H, W, 3]` f32 linear radiance; autograd to the scene's albedo (constants or texels) and emission |
 | `ops.pt_forward` / `ops.pt_backward` | raw kernel launches |
 ## 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. Each
-sample is a unidirectional path: the albedo at a hit is a bilinear fetch
-from the material's texture (repeat wrap), cosine-weighted hemisphere
-sampling makes the diffuse throughput multiplier exactly that albedo per
-bounce, and with next-event estimation on, emission is gathered at the
-camera hit while every later vertex estimates direct light by area sampling
-over the emissive-face list. NEE is skipped at the final vertex so the NEE
-and BRDF-only estimators cover camera-to-light paths of the same maximum
-length and agree in expectation at any bounce budget, not just in the limit.
-The backward pass is exact path replay. The random stream is counter-based
-Philox keyed by `(seed, pixel, sample)`, and sampling directions depend only
-on geometry and that stream, never on albedo or emission, so the backward
-kernel re-traces the identical paths with the identical draws and no stored
-path state. Each radiance term is a product of per-bounce albedo factors, a
-geometry scalar, and an emission value; the replay differentiates it in
-closed form with prefix/suffix exclusion products (no division, so zero
-albedos are safe) and scatters each bounce's gradient into the four texels
-of its bilinear footprint with the sampling weights, the exact adjoint of
-the fetch. Texture-gradient buffers up to 2048 texels are accumulated in
-per-block shared memory and flushed once (per-material constants and small
-textures take this path); larger textures use direct global atomics, where
-the contention is spread across the texel buffer. Parameter gradients are
-exactly the derivative of the rendered estimate: interior terms only, no
-visibility/silhouette gradients, which is the correct and complete
-derivative for materials and emission under fixed geometry.
 ## Correctness
@@ -96,57 +101,74 @@ Measured on RTX 6000 Ada (sm89) from a source build; the published `v1`
 variants additionally pass the full suite through `get_kernel` on a GeForce
 RTX 3070 Ti (Ampere sm86, Linux, torch 2.13, cu126):
-- **Analytic furnace, zero variance.** In a closed uniform box with BRDF
-  sampling, every path returns exactly `E * sum a^k` independent of the
-  random walk; at `a = 0.5`, `E = 0.25`, 8 bounces the analytic value is
-  0.498046875 and the worst pixel deviates 3.4e-3 (64x64, 4 spp).
-- **Estimator agreement.** The NEE and BRDF-only estimators agree to 0.10%
-  on a diffuse room with an area light (128x128; 128 vs 512 spp).
-- **Gradients match finite differences.** With a fixed seed the loss is
-  smooth in the parameters along the same paths; analytic replay gradients
-  match central differences to 3.1e-4 max relative error over nine
-  constant albedo/emission entries, and to 1.1e-3 over five texels of a
-  4x4 wall texture through the bilinear-footprint adjoint.
-- **Constants are 1x1 textures.** The two albedo forms render images equal
-  to float rounding and receive matching gradients.
-- **Inverse rendering.** Recovering a wall color (3 unknowns) reaches max
-  abs error 0.003 in 80 Adam steps; recovering an 8x8 wall texture (192
-  unknowns) drives the image loss from 1.4e-4 to 4.8e-14 and the mean texel
-  error below 5e-4 in 200 steps, both through the kernel alone.
-- **Deterministic.** Repeated forward renders are bitwise identical
-  (`torch.equal`); gradient buffers accumulate through float atomics and are
-  deterministic in expectation but not bitwise.
 ## Measured
-RTX 6000 Ada, Cornell box with an interior blocker (24 triangles, 4
-materials), 512x512, 64 spp, 4 bounces, NEE on:
-| pass | time | rate |
-|---|---|---|
-| forward | 9.6 ms | 1.76 Gpaths/s |
-| backward (replay) | 35.7 ms | 3.74x forward |
 ## Requirements and limits
 - NVIDIA GPU with compute capability 8.0+; float32 throughout.
-- Diffuse (Lambertian) + emissive materials; pinhole camera; at most 64
-  materials and 16 bounces; two-sided surfaces; emission is per-material
-  constant (textures carry albedo only).
-- Differentiable parameters are the albedo texels/constants and per-material
-  emission. Geometry, camera, and UVs are not differentiable inputs, and
-  there are no visibility/silhouette gradients.
-- Texture shapes are fixed at Scene construction; UVs wrap (repeat).
-- 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); Möller and
-Trumbore, "Fast, Minimum Storage Ray-Triangle Intersection" (1997); Duff et
-al., "Building an Orthonormal Basis, Revisited" (JCGT 2017).
 ## License

 # 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),
+and smooth dielectric materials; area lights and an importance-sampled
+equirectangular environment map; 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 the
+albedo texels, the per-material emission, and the environment texels.
+Inverse rendering runs as a bare torch loop around the kernel, with no
+rendering framework in the loop: render, compare, `backward()`, step.
 ## Usage
 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)
+env = torch.full((64, 128, 3), 0.8, 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=emission, uvs=uvs,
+                  material_types=[ptd.DIFFUSE, ptd.DIFFUSE, ptd.CONDUCTOR],
+                  roughness=[0.3, 0.3, 0.2], env=env)
 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, emission, and env
 ```
 `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
 | Symbol | Purpose |
 |---|---|
+| `Scene(vertices, faces, material_ids, albedo, emission, uvs, material_types, roughness, ior, env)` | triangle-mesh scene; BVH, light list, and the detached environment CDF are built at construction; albedo is `[M, 3]` constants or a list of `[Hm, Wm, 3]` textures; `env` is an optional `[Eh, Ew, 3]` equirect map; texture/env shapes are fixed 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 radiance; autograd to albedo texels, emission, and env texels |
+| `DIFFUSE`, `CONDUCTOR`, `DIELECTRIC` | material type constants |
 | `ops.pt_forward` / `ops.pt_backward` | raw kernel launches |
 ## 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 between reflection and refraction. Direct light is estimated
+at every non-delta vertex 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 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. Texture-gradient buffers up to 2048 texels accumulate in per-block
+shared memory and flush once; larger textures and environment gradients
+use direct global atomics spread across the buffer. Gradients are the
+exact derivative of the rendered estimate under frozen sampling: interior
+terms only, no visibility/silhouette gradients.
 ## Correctness
 variants additionally pass the full suite through `get_kernel` on a GeForce
 RTX 3070 Ti (Ampere sm86, Linux, torch 2.13, cu126):
+- **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 slab.** A smooth glass slab (ior 1.5) at normal incidence
+  transmits the internal-bounce series `(1-R)/(1+R)`; measured 0.9222
+  against 0.923077 analytic (0.10%).
+- **Estimator agreement.** MIS, NEE-only, and BRDF-only means agree within
+  Monte Carlo tolerance on a Cornell box, and MIS vs BRDF-only agree on a
+  GGX conductor under an environment map, which jointly exercises the VNDF
+  weight, the GGX pdf, and the environment pdf.
+- **Gradients match finite differences along the same paths.** Constants
+  4.6e-4 max relative error over nine entries; wall-texture texels 5.2e-3;
+  conductor F0 through GGX sampling and MIS 8.4e-5; environment texels
+  2.1e-4 (detached CDF). Texels no path can see get exactly zero from both
+  the replay and the differences.
+- **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 6.5e-14 with mean texel error below 5e-4 in
+  200 steps, both 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.
+| scene | build | forward | backward (replay) |
+|---|---|---|---|
+| Cornell box, 24 tris, 64 spp | - | 12.8 ms (1.31 Gpaths/s) | 59.0 ms (4.6x) |
+| terrain, 522,254 tris, 16 spp | 0.2 s | 18.2 ms (230 Mpaths/s) | 35.3 ms (1.9x) |
+| terrain, 1,045,470 tris, 16 spp | 0.3 s | 21.6 ms (194 Mpaths/s) | 37.9 ms (1.8x) |
 ## Requirements and limits
 - NVIDIA GPU with compute capability 8.0+; float32 throughout.
+- Materials: Lambertian diffuse, GGX conductor (roughness clamped to
+  >= 0.01), smooth dielectric; per-material emission; pinhole camera; at
+  most 64 materials and 16 bounces; two-sided surfaces; no Russian
+  roulette (fixed bounce budget keeps the replay exact).
+- Differentiable parameters: albedo texels/constants (all material types;
+  conductor F0 tint), per-material emission, environment texels. Geometry,
+  camera, UVs, roughness, and ior are frozen, and there are no
+  visibility/silhouette gradients.
+- 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); Heitz, "Sampling
+the GGX Distribution of Visible Normals" (JCGT 2018); Heitz,
+"Understanding the Masking-Shadowing Function" (JCGT 2014); 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

README.md CHANGED Viewed

@@ -6,18 +6,17 @@ license: apache-2.0
 # pathtracer-diff
 A differentiable Monte Carlo path tracer as one kernel, loadable through
-`kernels`. The forward pass renders a linear-radiance image with a megakernel
-unidirectional path tracer (diffuse + emissive materials, per-texel albedo
-textures with bilinear filtering, cosine-weighted BRDF sampling, optional
-next-event estimation, BVH traversal); the backward pass replays the
-identical paths with the identical counter-based random draws and emits
-analytic gradients to the albedo texels and the per-material emission.
-Inverse rendering then runs as a bare torch loop around the kernel, with no
-rendering framework in the loop: render, compare to a target, `backward()`,
-step. Per-texel albedo is what makes the inverse problem non-trivial: a
-texture is thousands of unknowns estimated through indirect illumination,
-where per-material constants make the pixel values merely multilinear in a
-handful of scalars.
 ## Usage
@@ -28,67 +27,73 @@ 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)
-albedo = [torch.tensor([0.73] * 3),      # constants are 1x1 textures
-          torch.tensor([0.65, 0.05, 0.05]),
-          wall_tex]                      # [H, W, 3] per-texel albedo
-scene = ptd.Scene(vertices, faces, material_ids, albedo, emission, uvs=uvs)
 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 emission)
 ```
 `version` selects the release branch; `trust_remote_code` is required by
 `kernels` for publishers without the trusted-publisher mark.
-`albedo` is either an `[M, 3]` tensor of per-material constants or a list of
-M per-material textures `[Hm, Wm, 3]`; constants are treated as 1x1 textures,
-so both forms share one code path and both receive gradients. A fixed `seed`
-renders with 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.
 ## API
 | Symbol | Purpose |
 |---|---|
-| `Scene(vertices, faces, material_ids, albedo, emission, uvs)` | triangle-mesh scene; the BVH and the emissive-face light list are built at construction; `uvs` is `[V, 2]` per vertex or `[F, 3, 2]` per corner, textures wrap (repeat); texture shapes are fixed at construction |
 | `Camera(position, look_at, up, vfov_deg)` | pinhole camera |
-| `render(scene, camera, height, width, spp, max_bounces, nee, seed)` | differentiable render to `[H, W, 3]` f32 linear radiance; autograd to the scene's albedo (constants or texels) and emission |
 | `ops.pt_forward` / `ops.pt_backward` | raw kernel launches |
 ## 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. Each
-sample is a unidirectional path: the albedo at a hit is a bilinear fetch
-from the material's texture (repeat wrap), cosine-weighted hemisphere
-sampling makes the diffuse throughput multiplier exactly that albedo per
-bounce, and with next-event estimation on, emission is gathered at the
-camera hit while every later vertex estimates direct light by area sampling
-over the emissive-face list. NEE is skipped at the final vertex so the NEE
-and BRDF-only estimators cover camera-to-light paths of the same maximum
-length and agree in expectation at any bounce budget, not just in the limit.
-The backward pass is exact path replay. The random stream is counter-based
-Philox keyed by `(seed, pixel, sample)`, and sampling directions depend only
-on geometry and that stream, never on albedo or emission, so the backward
-kernel re-traces the identical paths with the identical draws and no stored
-path state. Each radiance term is a product of per-bounce albedo factors, a
-geometry scalar, and an emission value; the replay differentiates it in
-closed form with prefix/suffix exclusion products (no division, so zero
-albedos are safe) and scatters each bounce's gradient into the four texels
-of its bilinear footprint with the sampling weights, the exact adjoint of
-the fetch. Texture-gradient buffers up to 2048 texels are accumulated in
-per-block shared memory and flushed once (per-material constants and small
-textures take this path); larger textures use direct global atomics, where
-the contention is spread across the texel buffer. Parameter gradients are
-exactly the derivative of the rendered estimate: interior terms only, no
-visibility/silhouette gradients, which is the correct and complete
-derivative for materials and emission under fixed geometry.
 ## Correctness
@@ -96,58 +101,75 @@ Measured on RTX 6000 Ada (sm89) from a source build; the published `v1`
 variants additionally pass the full suite through `get_kernel` on a GeForce
 RTX 3070 Ti (Ampere sm86, Linux, torch 2.13, cu126):
-- **Analytic furnace, zero variance.** In a closed uniform box with BRDF
-  sampling, every path returns exactly `E * sum a^k` independent of the
-  random walk; at `a = 0.5`, `E = 0.25`, 8 bounces the analytic value is
-  0.498046875 and the worst pixel deviates 3.4e-3 (64x64, 4 spp).
-- **Estimator agreement.** The NEE and BRDF-only estimators agree to 0.10%
-  on a diffuse room with an area light (128x128; 128 vs 512 spp).
-- **Gradients match finite differences.** With a fixed seed the loss is
-  smooth in the parameters along the same paths; analytic replay gradients
-  match central differences to 3.1e-4 max relative error over nine
-  constant albedo/emission entries, and to 1.1e-3 over five texels of a
-  4x4 wall texture through the bilinear-footprint adjoint.
-- **Constants are 1x1 textures.** The two albedo forms render images equal
-  to float rounding and receive matching gradients.
-- **Inverse rendering.** Recovering a wall color (3 unknowns) reaches max
-  abs error 0.003 in 80 Adam steps; recovering an 8x8 wall texture (192
-  unknowns) drives the image loss from 1.4e-4 to 4.8e-14 and the mean texel
-  error below 5e-4 in 200 steps, both through the kernel alone.
-- **Deterministic.** Repeated forward renders are bitwise identical
-  (`torch.equal`); gradient buffers accumulate through float atomics and are
-  deterministic in expectation but not bitwise.
 ## Measured
-RTX 6000 Ada, Cornell box with an interior blocker (24 triangles, 4
-materials), 512x512, 64 spp, 4 bounces, NEE on:
-| pass | time | rate |
-|---|---|---|
-| forward | 9.6 ms | 1.76 Gpaths/s |
-| backward (replay) | 35.7 ms | 3.74x forward |
 ## Requirements and limits
 - NVIDIA GPU with compute capability 8.0+; float32 throughout.
-- Diffuse (Lambertian) + emissive materials; pinhole camera; at most 64
-  materials and 16 bounces; two-sided surfaces; emission is per-material
-  constant (textures carry albedo only).
-- Differentiable parameters are the albedo texels/constants and per-material
-  emission. Geometry, camera, and UVs are not differentiable inputs, and
-  there are no visibility/silhouette gradients.
-- Texture shapes are fixed at Scene construction; UVs wrap (repeat).
-- 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); Möller and
-Trumbore, "Fast, Minimum Storage Ray-Triangle Intersection" (1997); Duff et
-al., "Building an Orthonormal Basis, Revisited" (JCGT 2017).
 ## 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:
+Lambertian diffuse, GGX conductor (VNDF-sampled, height-correlated Smith),
+and smooth dielectric materials; area lights and an importance-sampled
+equirectangular environment map; 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 the
+albedo texels, the per-material emission, and the environment texels.
+Inverse rendering runs as a bare torch loop around the kernel, with no
+rendering framework in the loop: render, compare, `backward()`, step.
 ## Usage
 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)
+env = torch.full((64, 128, 3), 0.8, 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=emission, uvs=uvs,
+                  material_types=[ptd.DIFFUSE, ptd.DIFFUSE, ptd.CONDUCTOR],
+                  roughness=[0.3, 0.3, 0.2], env=env)
 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, emission, and env
 ```
 `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
 | Symbol | Purpose |
 |---|---|
+| `Scene(vertices, faces, material_ids, albedo, emission, uvs, material_types, roughness, ior, env)` | triangle-mesh scene; BVH, light list, and the detached environment CDF are built at construction; albedo is `[M, 3]` constants or a list of `[Hm, Wm, 3]` textures; `env` is an optional `[Eh, Ew, 3]` equirect map; texture/env shapes are fixed 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 radiance; autograd to albedo texels, emission, and env texels |
+| `DIFFUSE`, `CONDUCTOR`, `DIELECTRIC` | material type constants |
 | `ops.pt_forward` / `ops.pt_backward` | raw kernel launches |
 ## 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 between reflection and refraction. Direct light is estimated
+at every non-delta vertex 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 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. Texture-gradient buffers up to 2048 texels accumulate in per-block
+shared memory and flush once; larger textures and environment gradients
+use direct global atomics spread across the buffer. Gradients are the
+exact derivative of the rendered estimate under frozen sampling: interior
+terms only, no visibility/silhouette gradients.
 ## Correctness
 variants additionally pass the full suite through `get_kernel` on a GeForce
 RTX 3070 Ti (Ampere sm86, Linux, torch 2.13, cu126):
+- **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 slab.** A smooth glass slab (ior 1.5) at normal incidence
+  transmits the internal-bounce series `(1-R)/(1+R)`; measured 0.9222
+  against 0.923077 analytic (0.10%).
+- **Estimator agreement.** MIS, NEE-only, and BRDF-only means agree within
+  Monte Carlo tolerance on a Cornell box, and MIS vs BRDF-only agree on a
+  GGX conductor under an environment map, which jointly exercises the VNDF
+  weight, the GGX pdf, and the environment pdf.
+- **Gradients match finite differences along the same paths.** Constants
+  4.6e-4 max relative error over nine entries; wall-texture texels 5.2e-3;
+  conductor F0 through GGX sampling and MIS 8.4e-5; environment texels
+  2.1e-4 (detached CDF). Texels no path can see get exactly zero from both
+  the replay and the differences.
+- **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 6.5e-14 with mean texel error below 5e-4 in
+  200 steps, both 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.
+| scene | build | forward | backward (replay) |
+|---|---|---|---|
+| Cornell box, 24 tris, 64 spp | - | 12.8 ms (1.31 Gpaths/s) | 59.0 ms (4.6x) |
+| terrain, 522,254 tris, 16 spp | 0.2 s | 18.2 ms (230 Mpaths/s) | 35.3 ms (1.9x) |
+| terrain, 1,045,470 tris, 16 spp | 0.3 s | 21.6 ms (194 Mpaths/s) | 37.9 ms (1.8x) |
 ## Requirements and limits
 - NVIDIA GPU with compute capability 8.0+; float32 throughout.
+- Materials: Lambertian diffuse, GGX conductor (roughness clamped to
+  >= 0.01), smooth dielectric; per-material emission; pinhole camera; at
+  most 64 materials and 16 bounces; two-sided surfaces; no Russian
+  roulette (fixed bounce budget keeps the replay exact).
+- Differentiable parameters: albedo texels/constants (all material types;
+  conductor F0 tint), per-material emission, environment texels. Geometry,
+  camera, UVs, roughness, and ior are frozen, and there are no
+  visibility/silhouette gradients.
+- 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); Heitz, "Sampling
+the GGX Distribution of Visible Normals" (JCGT 2018); Heitz,
+"Understanding the Masking-Shadowing Function" (JCGT 2014); 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.

dev/report_local.py CHANGED Viewed

@@ -222,36 +222,206 @@ print(f"texture inverse: 8x8x3 (192 unknowns) recovered to mean abs err "
       f"{err.mean().item():.3f} (max {err.max().item():.3f}), "
       f"loss {first:.2e} -> {loss.item():.2e} in 200 steps")
-# ---- throughput
-scene = cornell(with_box=True)
-H, W, spp, B = 512, 512, 64, 4
-img = ptd.render(scene, cam, H, W, spp=spp, max_bounces=B)  # warmup + autograd off
-torch.cuda.synchronize()
-ts = []
-for _ in range(3):
-    t0 = time.perf_counter()
-    img = ptd.render(scene, cam, H, W, spp=spp, max_bounces=B)
-    torch.cuda.synchronize()
-    ts.append(time.perf_counter() - t0)
-fwd = sorted(ts)[1]
-paths = H * W * spp
-print(f"forward: {H}x{W} spp {spp} bounces {B} ({scene.n_faces} tris): "
-      f"{fwd*1e3:.1f} ms, {paths/fwd/1e6:.0f} Mpaths/s")
 alb = scene.albedo.clone().requires_grad_(True)
 scene.albedo = alb
-g = torch.ones(H, W, 3, device="cuda")
-img = ptd.render(scene, cam, H, W, spp=spp, max_bounces=B)
-img.backward(g)  # warmup
-torch.cuda.synchronize()
-ts = []
-for _ in range(3):
-    alb.grad = None
-    img = ptd.render(scene, cam, H, W, spp=spp, max_bounces=B)
-    torch.cuda.synchronize()
     t0 = time.perf_counter()
-    img.backward(g)
-    torch.cuda.synchronize()
-    ts.append(time.perf_counter() - t0)
-bwd = sorted(ts)[1]
-print(f"backward (replay): {bwd*1e3:.1f} ms, {bwd/fwd:.2f}x forward")

       f"{err.mean().item():.3f} (max {err.max().item():.3f}), "
       f"loss {first:.2e} -> {loss.item():.2e} in 200 steps")
+# ---- Fresnel slab analytic
+verts, faces, mat = [], [], []
+def quad_w(a, b, c, d, m):
+    base = len(verts)
+    verts.extend([a, b, c, d])
+    faces.extend(quad(base, base + 1, base + 2, base + 3))
+    mat.extend([m, m])
+quad_w((-10, -10, 5), (10, -10, 5), (10, 10, 5), (-10, 10, 5), 0)
+add_box(verts, faces, mat, (-10, -10, 2.0), (10, 10, 2.2), 1)
+slab = ptd.Scene(verts, faces, mat, albedo=[[0, 0, 0], [0, 0, 0]],
+                 emission=[[1.0] * 3, [0, 0, 0]],
+                 material_types=[ptd.DIFFUSE, ptd.DIELECTRIC], ior=[1.5, 1.5])
+scam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
+img = ptd.render(slab, scam, 32, 32, spp=1024, max_bounces=10,
+                 estimator="brdf", seed=5)
+R = (0.5 / 2.5) ** 2
+exp_T = (1 - R) / (1 + R)
+got = img[12:20, 12:20].mean().item()
+print(f"fresnel slab (ior 1.5): analytic T {exp_T:.6f}, measured {got:.6f}, "
+      f"rel err {abs(got-exp_T)/exp_T:.4%}")
+# ---- conductor F0 FD (through GGX + MIS)
+H = W = 48
+torch.manual_seed(0)
+Wr = torch.rand(H, W, 3, device="cuda")
+base_alb = torch.tensor([[0.73] * 3, [0.9, 0.5, 0.2], [0.12, 0.45, 0.15],
+                         [0.78] * 3], device="cuda")
+def cond_loss(alb):
+    verts, faces, mat = [], [], []
+    def wall(a, b, c, d, m):
+        b0 = len(verts)
+        verts.extend([a, b, c, d])
+        faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
+        mat.extend([m, m])
+    X, Y, Z = 5.56, 5.488, 5.592
+    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
+    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
+    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)
+    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)
+    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
+    wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
+         (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
+    sc = ptd.Scene(verts, faces, mat, albedo=alb,
+                   emission=[[0, 0, 0], [0, 0, 0], [0, 0, 0],
+                             [18.4, 15.6, 8.0]],
+                   material_types=[0, 1, 0, 0], roughness=[0.3, 0.2, 0.3, 0.3])
+    return (ptd.render(sc, cam, H, W, spp=8, max_bounces=3, seed=13) * Wr).sum()
+alb = base_alb.clone().requires_grad_(True)
+cond_loss(alb).backward()
+worst = 0.0
+for (i, c) in [(1, 0), (1, 2), (0, 1)]:
+    h = 2e-3
+    ap = base_alb.clone(); ap[i, c] += h
+    am = base_alb.clone(); am[i, c] -= h
+    fd = (cond_loss(ap) - cond_loss(am)).item() / (2 * h)
+    worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
+print(f"conductor F0 gradients vs central differences (GGX+MIS): "
+      f"max rel err {worst:.2e} over 3 entries")
+# ---- env texel FD
+torch.manual_seed(3)
+env0 = torch.rand(4, 8, 3, device="cuda") * 1.5 + 0.2
+verts, faces, mat = [], [], []
+b0 = len(verts)
+verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
+faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
+mat.extend([0, 0])
+ecam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.5), vfov_deg=40.0)
+torch.manual_seed(0)
+Wre = torch.rand(32, 32, 3, device="cuda")
+esc = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3], emission=[[0, 0, 0]],
+                env=env0)
+def env_loss(e):
+    esc.env = e.to("cuda")
+    return (ptd.render(esc, ecam, 32, 32, spp=8, max_bounces=3, seed=11)
+            * Wre).sum()
+e = env0.clone().requires_grad_(True)
+env_loss(e).backward()
+worst = 0.0
+for (y, x, c) in [(0, 0, 0), (1, 4, 1), (1, 7, 2)]:
+    h = 5e-3
+    ep = env0.clone(); ep[y, x, c] += h
+    em2 = env0.clone(); em2[y, x, c] -= h
+    fd = (env_loss(ep) - env_loss(em2)).item() / (2 * h)
+    worst = max(worst, abs(e.grad[y, x, c].item() - fd) / abs(fd))
+print(f"environment texel gradients vs central differences: "
+      f"max rel err {worst:.2e} over 3 texels (detached CDF)")
+# ---- throughput helper
+def bench(scene, bcam, H, W, spp, B, label, grad_leaf=None):
+    img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
+    torch.cuda.synchronize()
+    ts = []
+    for _ in range(3):
+        t0 = time.perf_counter()
+        img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
+        torch.cuda.synchronize()
+        ts.append(time.perf_counter() - t0)
+    fwd = sorted(ts)[1]
+    paths = H * W * spp
+    line = (f"{label}: forward {fwd*1e3:.1f} ms ({paths/fwd/1e6:.0f} Mpaths/s)")
+    if grad_leaf is not None:
+        g = torch.ones(H, W, 3, device="cuda")
+        img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
+        img.backward(g)
+        torch.cuda.synchronize()
+        ts = []
+        for _ in range(3):
+            grad_leaf.grad = None
+            img = ptd.render(scene, bcam, H, W, spp=spp, max_bounces=B)
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            img.backward(g)
+            torch.cuda.synchronize()
+            ts.append(time.perf_counter() - t0)
+        bwd = sorted(ts)[1]
+        line += f"; backward {bwd*1e3:.1f} ms ({bwd/fwd:.2f}x)"
+    print(line)
+# cornell (continuity with v1 numbers)
+scene = cornell(with_box=True)
 alb = scene.albedo.clone().requires_grad_(True)
 scene.albedo = alb
+bench(scene, cam, 512, 512, 64, 4, "cornell 24 tris, 512x512 spp 64 b4", alb)
+# ---- large scenes: displaced terrain + conductor box + checker + env
+def terrain_scene(n):
+    ax = torch.linspace(0, 10, n + 1)
+    X, Z = torch.meshgrid(ax, ax, indexing="ij")
+    Y = (0.35 * torch.sin(X * 1.7) * torch.cos(Z * 1.3) +
+         0.2 * torch.sin(X * 4.1 + 1.0) * torch.sin(Z * 3.7) +
+         0.08 * torch.sin(X * 9.3) * torch.cos(Z * 8.1))
+    verts = torch.stack([X, Y, Z], dim=-1).reshape(-1, 3)
+    ii = torch.arange(n)
+    jj = torch.arange(n)
+    I, J = torch.meshgrid(ii, jj, indexing="ij")
+    v00 = (I * (n + 1) + J).reshape(-1)
+    v10 = ((I + 1) * (n + 1) + J).reshape(-1)
+    v01 = (I * (n + 1) + J + 1).reshape(-1)
+    v11 = ((I + 1) * (n + 1) + J + 1).reshape(-1)
+    f1 = torch.stack([v00, v10, v11], dim=1)
+    f2 = torch.stack([v00, v11, v01], dim=1)
+    faces = torch.cat([f1, f2], dim=0)
+    mat = torch.zeros(faces.shape[0], dtype=torch.int64)
+    uv = torch.stack([X / 10.0, Z / 10.0], dim=-1).reshape(-1, 2)
+    verts_l, faces_l, mat_l = [verts], [faces], [mat]
+    base = verts.shape[0]
+    bx = torch.tensor([(4.2, 0.4, 4.2), (5.8, 0.4, 4.2), (5.8, 0.4, 5.8),
+                       (4.2, 0.4, 5.8), (4.2, 2.2, 4.2), (5.8, 2.2, 4.2),
+                       (5.8, 2.2, 5.8), (4.2, 2.2, 5.8)])
+    verts_l.append(bx)
+    bq = [[0, 1, 2], [0, 2, 3], [4, 7, 6], [4, 6, 5], [0, 4, 5], [0, 5, 1],
+          [3, 2, 6], [3, 6, 7], [0, 3, 7], [0, 7, 4], [1, 5, 6], [1, 6, 2]]
+    faces_l.append(torch.tensor(bq, dtype=torch.int64) + base)
+    mat_l.append(torch.ones(12, dtype=torch.int64))
+    verts = torch.cat(verts_l)
+    faces = torch.cat(faces_l)
+    mat = torch.cat(mat_l)
+    uv = torch.cat([uv, torch.zeros(8, 2)])
+    yy, xx = torch.meshgrid(torch.arange(256), torch.arange(256),
+                            indexing="ij")
+    checker = (((yy // 16 + xx // 16) % 2).float() * 0.45 + 0.3)
+    tex = torch.stack([checker, checker * 0.9, checker * 0.7], dim=-1)
+    env = torch.full((16, 32, 3), 0.9)
     t0 = time.perf_counter()
+    sc = ptd.Scene(verts, faces, mat,
+                   albedo=[tex.cuda(), torch.tensor([0.9, 0.65, 0.35])],
+                   emission=[[0, 0, 0], [0, 0, 0]],
+                   material_types=[ptd.DIFFUSE, ptd.CONDUCTOR],
+                   roughness=[0.3, 0.15], uvs=uv, env=env)
+    build = time.perf_counter() - t0
+    return sc, build
+tcam = ptd.Camera(position=(11.0, 4.5, 11.0), look_at=(5.0, 0.5, 5.0),
+                  vfov_deg=45.0)
+for n in (511, 723):
+    sc, build = terrain_scene(n)
+    ntri = sc.n_faces
+    tex_leaf = sc.albedo_textures[0].requires_grad_(True)
+    sc.albedo_textures[0] = tex_leaf
+    bench(sc, tcam, 512, 512, 16, 4,
+          f"terrain {ntri} tris (SAH build {build:.1f} s), 512x512 spp 16 b4",
+          tex_leaf)

local_bind.cpp CHANGED Viewed

@@ -9,35 +9,83 @@
 namespace {
 void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                 torch::Tensor nodes_f, torch::Tensor nodes_i,
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
-                torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                int64_t seed, torch::Tensor image) {
-  TORCH_CHECK(tris.is_cuda() && tris.is_contiguous(), "tris");
-  TORCH_CHECK(uvs.numel() == tris.size(0) * 6, "uvs [F, 3, 2]");
-  TORCH_CHECK(image.is_cuda() && image.is_contiguous() && image.dim() == 3,
-              "image [H, W, 3]");
-  TORCH_CHECK(tex_hdr.size(0) <= 64, "at most 64 materials");
-  TORCH_CHECK(max_bounces >= 1 && max_bounces <= 16, "bounces in [1,16]");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  int L = (int)light_faces.numel();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
-  ptd_forward_launch(
-      tris.const_data_ptr<float>(), mat_ids.const_data_ptr<int>(),
-      uvs.const_data_ptr<float>(), (int)tris.size(0),
-      nodes_f.const_data_ptr<float>(), nodes_i.const_data_ptr<int>(),
-      (int)nodes_f.size(0), L ? light_faces.const_data_ptr<int>() : nullptr,
-      L ? light_cdf.const_data_ptr<float>() : nullptr, L,
-      (float)total_light_area, tex.const_data_ptr<float>(),
-      tex_hdr.const_data_ptr<int>(), (int)tex.size(0),
-      emission.const_data_ptr<float>(),
-      (int)tex_hdr.size(0), cam_h.const_data_ptr<float>(),
-      (int)image.size(0), (int)image.size(1), (int)spp, (int)max_bounces,
-      nee ? 1 : 0, (long long)seed, image.data_ptr<float>(), stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
@@ -46,30 +94,30 @@ void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
-                 torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                 int64_t seed, torch::Tensor grad_image,
-                 torch::Tensor grad_tex, torch::Tensor grad_emission) {
-  TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous(), "grad_image");
-  TORCH_CHECK(grad_tex.sizes() == tex.sizes(), "grad_tex must match tex");
-  TORCH_CHECK(tex_hdr.size(0) <= 64, "at most 64 materials");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  int L = (int)light_faces.numel();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
-  ptd_backward_launch(
-      tris.const_data_ptr<float>(), mat_ids.const_data_ptr<int>(),
-      uvs.const_data_ptr<float>(), (int)tris.size(0),
-      nodes_f.const_data_ptr<float>(), nodes_i.const_data_ptr<int>(),
-      (int)nodes_f.size(0), L ? light_faces.const_data_ptr<int>() : nullptr,
-      L ? light_cdf.const_data_ptr<float>() : nullptr, L,
-      (float)total_light_area, tex.const_data_ptr<float>(),
-      tex_hdr.const_data_ptr<int>(), (int)tex.size(0),
-      emission.const_data_ptr<float>(),
-      (int)tex_hdr.size(0), cam_h.const_data_ptr<float>(),
-      (int)grad_image.size(0), (int)grad_image.size(1), (int)spp,
-      (int)max_bounces, nee ? 1 : 0, (long long)seed,
-      grad_image.const_data_ptr<float>(), grad_tex.data_ptr<float>(),
-      grad_emission.data_ptr<float>(), stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }

 namespace {
+const float* fptr(const torch::Tensor& t) {
+  return t.numel() ? t.const_data_ptr<float>() : nullptr;
+}
+const int* iptr(const torch::Tensor& t) {
+  return t.numel() ? t.const_data_ptr<int>() : nullptr;
+}
+PtdSceneArgs pack_args(const torch::Tensor& tris, const torch::Tensor& mat_ids,
+                       const torch::Tensor& uvs, const torch::Tensor& nodes_f,
+                       const torch::Tensor& nodes_i,
+                       const torch::Tensor& light_faces,
+                       const torch::Tensor& light_cdf, double total_light_area,
+                       const torch::Tensor& tex, const torch::Tensor& tex_hdr,
+                       const torch::Tensor& emission,
+                       const torch::Tensor& mat_type,
+                       const torch::Tensor& mat_rough,
+                       const torch::Tensor& mat_ior, const torch::Tensor& env,
+                       const torch::Tensor& env_cdf_m,
+                       const torch::Tensor& env_cdf_c,
+                       const torch::Tensor& env_pdf, int64_t env_w,
+                       int64_t env_h, int64_t max_bounces) {
+  TORCH_CHECK(tris.is_cuda() && tris.is_contiguous(), "tris");
+  TORCH_CHECK(uvs.numel() == tris.size(0) * 6, "uvs [F, 3, 2]");
+  TORCH_CHECK(tex_hdr.size(0) <= 64, "at most 64 materials");
+  TORCH_CHECK(max_bounces >= 1 && max_bounces <= 16, "bounces in [1,16]");
+  PtdSceneArgs a;
+  a.tris = tris.const_data_ptr<float>();
+  a.mat_ids = mat_ids.const_data_ptr<int>();
+  a.uvs = uvs.const_data_ptr<float>();
+  a.n_faces = (int)tris.size(0);
+  a.nodes_f = nodes_f.const_data_ptr<float>();
+  a.nodes_i = nodes_i.const_data_ptr<int>();
+  a.n_nodes = (int)nodes_f.size(0);
+  a.light_faces = iptr(light_faces);
+  a.light_cdf = fptr(light_cdf);
+  a.n_lights = (int)light_faces.numel();
+  a.total_light_area = (float)total_light_area;
+  a.tex = tex.const_data_ptr<float>();
+  a.tex_hdr = tex_hdr.const_data_ptr<int>();
+  a.n_texels = (int)tex.size(0);
+  a.emission = emission.const_data_ptr<float>();
+  a.mat_type = mat_type.const_data_ptr<int>();
+  a.mat_rough = mat_rough.const_data_ptr<float>();
+  a.mat_ior = mat_ior.const_data_ptr<float>();
+  a.n_mats = (int)tex_hdr.size(0);
+  a.env = fptr(env);
+  a.env_w = (int)env_w;
+  a.env_h = (int)env_h;
+  a.env_cdf_m = fptr(env_cdf_m);
+  a.env_cdf_c = fptr(env_cdf_c);
+  a.env_pdf = fptr(env_pdf);
+  return a;
+}
 void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                 torch::Tensor nodes_f, torch::Tensor nodes_i,
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
+                torch::Tensor mat_type, torch::Tensor mat_rough,
+                torch::Tensor mat_ior, torch::Tensor env,
+                torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                int64_t mode, int64_t seed, torch::Tensor image) {
+  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
+                             light_cdf, total_light_area, tex, tex_hdr,
+                             emission, mat_type, mat_rough, mat_ior, env,
+                             env_cdf_m, env_cdf_c, env_pdf, env_w, env_h,
+                             max_bounces);
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
+  ptd_forward_launch(&a, cam_h.const_data_ptr<float>(), (int)image.size(0),
+                     (int)image.size(1), (int)spp, (int)max_bounces,
+                     (int)mode, (long long)seed, image.data_ptr<float>(),
+                     stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
+                 torch::Tensor mat_type, torch::Tensor mat_rough,
+                 torch::Tensor mat_ior, torch::Tensor env,
+                 torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                 torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                 torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                 int64_t mode, int64_t seed, torch::Tensor grad_image,
+                 torch::Tensor grad_tex, torch::Tensor grad_emission,
+                 torch::Tensor grad_env) {
+  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
+                             light_cdf, total_light_area, tex, tex_hdr,
+                             emission, mat_type, mat_rough, mat_ior, env,
+                             env_cdf_m, env_cdf_c, env_pdf, env_w, env_h,
+                             max_bounces);
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
+  ptd_backward_launch(&a, cam_h.const_data_ptr<float>(),
+                      (int)grad_image.size(0), (int)grad_image.size(1),
+                      (int)spp, (int)max_bounces, (int)mode, (long long)seed,
+                      grad_image.const_data_ptr<float>(),
+                      grad_tex.data_ptr<float>(),
+                      grad_emission.data_ptr<float>(),
+                      grad_env.numel() ? grad_env.data_ptr<float>() : nullptr,
+                      stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }

pathtracer_diff_cuda/pathtracer.cu CHANGED Viewed

@@ -1,28 +1,29 @@
 // pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
 //
-// Forward: a megakernel unidirectional path tracer (diffuse + emissive
-// materials, per-texel albedo textures with bilinear filtering, cosine-
-// weighted BRDF sampling, optional next-event estimation, BVH traversal).
-// One thread owns one pixel and accumulates its spp samples serially, so
-// the image is bitwise deterministic.
 //
-// Backward: exact path replay. Sampling directions depend only on geometry
-// and the counter-based Philox stream, never on the differentiable
-// parameters (albedo texels and per-material emission), so the backward
-// kernel re-traces the identical paths with the identical draws and
-// accumulates analytic gradients with no stored path state (path-replay
-// backpropagation, interior terms only: no visibility/silhouette
-// gradients). Each bounce's albedo factor is a bilinear combination of four
-// texels; its gradient is scattered to those texels with the bilinear
-// weights, the exact adjoint of the fetch.
-//
-// Radiance estimator with NEE on: emission is gathered at the camera hit
-// only; every later vertex estimates direct light by area sampling, and NEE
-// is skipped at the final vertex so the NEE and BRDF-only estimators cover
-// camera-to-light paths of the same maximum length. With NEE off: emission
-// is gathered at every hit (with cosine sampling the diffuse throughput
-// factor is exactly the albedo, which makes a closed uniform furnace a
-// zero-variance analytic test case).
 #include <cuda_runtime.h>
 #include <curand_kernel.h>
@@ -36,14 +37,16 @@ namespace {
 constexpr int kThreads = 128;
 constexpr int kMaxBounces = 16;
 constexpr int kMaxMats = 64;
-// Texture-gradient buffers up to this many texels are accumulated in
-// dynamic shared memory per block and flushed once; larger buffers take
-// direct global atomics (contention is then spread across many texels).
-constexpr int kSharedTexels = 2048;
 constexpr int kStack = 64;
 constexpr float kInvPi = 0.31830988618379067154f;
 constexpr float kRayEps = 1e-4f;
 constexpr float kShadowEps = 1e-3f;
 // ---------------------------------------------------------------- float3 ops
 __device__ __forceinline__ float3 f3(float x, float y, float z) {
@@ -67,6 +70,7 @@ __device__ __forceinline__ float3 cross(float3 a, float3 b) {
 __device__ __forceinline__ float3 normalize(float3 a) {
   return a * rsqrtf(fmaxf(dot(a, a), 1e-30f));
 }
 // Duff et al. 2017, "Building an Orthonormal Basis, Revisited".
 __device__ __forceinline__ void onb(float3 n, float3& t, float3& b) {
@@ -79,24 +83,33 @@ __device__ __forceinline__ void onb(float3 n, float3& t, float3& b) {
 // ---------------------------------------------------------------- scene view
 struct DevScene {
-  const float* tris;   // [F, 9] v0 v1 v2
   const int* mats;     // [F]
-  const float* uvs;    // [F, 3, 2] per-corner uv
-  const float* nf;     // [N, 6] bbox lo.xyz hi.xyz
-  const int* ni;       // [N, 3] (left,right,0) | (start,count,1)
   int n_nodes;
-  const int* lf;       // [L] emissive face indices
-  const float* lcdf;   // [L] normalized area cdf
   int nl;
-  float larea;         // total emissive area
-  const float* tex;    // [T, 3] flat albedo texels, all materials
-  const int* thdr;     // [M, 3] (texel offset, W, H) per material
-  int nt;              // T, total texel count
   const float* emi;    // [M, 3]
   int nm;
 };
-struct DevCam {  // pos, fwd, right*tan(v/2)*aspect, up*tan(v/2)
   float p[3], f[3], r[3], u[3];
 };
@@ -155,6 +168,7 @@ __device__ int bvh_closest(const DevScene& sc, float3 ro, float3 rd,
                            float tmin, float& tbest, float& bu, float& bv,
                            float3& ngbest) {
   float roa[3] = {ro.x, ro.y, ro.z};
   float inv[3];
   inv_dir(rd, inv);
   int stack[kStack];
@@ -165,7 +179,7 @@ __device__ int bvh_closest(const DevScene& sc, float3 ro, float3 rd,
     int nid = stack[--sp];
     if (!slab(&sc.nf[nid * 6], roa, inv, tbest)) continue;
     const int* n = &sc.ni[nid * 3];
-    if (n[2]) {
       for (int f = n[0]; f < n[0] + n[1]; ++f) {
         float t, u, v;
         float3 ng;
@@ -178,8 +192,11 @@ __device__ int bvh_closest(const DevScene& sc, float3 ro, float3 rd,
         }
       }
     } else if (sp + 2 <= kStack) {
-      stack[sp++] = n[0];
-      stack[sp++] = n[1];
     }
   }
   return best;
@@ -197,7 +214,7 @@ __device__ bool bvh_occluded(const DevScene& sc, float3 ro, float3 rd,
     int nid = stack[--sp];
     if (!slab(&sc.nf[nid * 6], roa, inv, tmax)) continue;
     const int* n = &sc.ni[nid * 3];
-    if (n[2]) {
       for (int f = n[0]; f < n[0] + n[1]; ++f) {
         float t, u, v;
         float3 ng;
@@ -212,8 +229,8 @@ __device__ bool bvh_occluded(const DevScene& sc, float3 ro, float3 rd,
   return false;
 }
-__device__ __forceinline__ int pick_light(const float* cdf, int nl, float r) {
-  int lo = 0, hi = nl - 1;
   while (lo < hi) {
     int mid = (lo + hi) >> 1;
     if (cdf[mid] < r) lo = mid + 1; else hi = mid;
@@ -221,18 +238,10 @@ __device__ __forceinline__ int pick_light(const float* cdf, int nl, float r) {
   return lo;
 }
-// Bilinear albedo fetch with repeat wrap; also emits the four texel indices
-// and weights (the footprint) for the backward scatter.
-__device__ __forceinline__ void sample_albedo(const DevScene& sc, int m,
-                                              int face, float bu, float bv,
-                                              float rgb[3], int idx4[4],
-                                              float w4[4]) {
-  const float* U = &sc.uvs[face * 6];
-  float w0 = 1.0f - bu - bv;
-  float u = w0 * U[0] + bu * U[2] + bv * U[4];
-  float v = w0 * U[1] + bu * U[3] + bv * U[5];
-  const int* h = &sc.thdr[m * 3];
-  int off = h[0], W = h[1], H = h[2];
   float uu = u * (float)W - 0.5f;
   float vv = v * (float)H - 0.5f;
   float fx = uu - floorf(uu);
@@ -249,52 +258,240 @@ __device__ __forceinline__ void sample_albedo(const DevScene& sc, int m,
   rgb[0] = rgb[1] = rgb[2] = 0.0f;
 #pragma unroll
   for (int i = 0; i < 4; ++i) {
-    const float* t = &sc.tex[idx4[i] * 3];
     rgb[0] += w4[i] * t[0];
     rgb[1] += w4[i] * t[1];
     rgb[2] += w4[i] * t[2];
   }
 }
-// Texel-gradient scatter: shared staging when the whole buffer fits,
-// direct global atomics otherwise.
 __device__ __forceinline__ void scatter_tex(float* s_gt, float* g_tex,
                                             int idx, int c, float v) {
-  if (s_gt) {
-    atomicAdd(&s_gt[idx * 3 + c], v);
-  } else {
-    atomicAdd(&g_tex[idx * 3 + c], v);
   }
 }
-// ------------------------------------------------------------- path tracing
-// One camera path: forward accumulates radiance into acc; backward (GRAD)
-// accumulates parameter gradients (texel scatter for albedo, shared buffer
-// for emission). The RNG draw order is identical in both instantiations by
-// construction.
 template <bool GRAD>
 __device__ void trace_path(const DevScene& sc, float3 ro, float3 rd,
-                           curandStatePhilox4_32_10_t& st, int B, int nee,
                            float3& acc, const float gs[3], float* g_tex,
-                           float* s_gt, float* s_ge) {
-  float T[3] = {1.0f, 1.0f, 1.0f};      // prefix albedo product
-  float albs[kMaxBounces][3];           // per-bounce sampled albedo
-  int fp_i[kMaxBounces][4];             // per-bounce texel footprint
-  float fp_w[kMaxBounces][4];
-  int mats[kMaxBounces];
-  float suf[kMaxBounces + 1][3];        // suffix products workspace
   for (int k = 0; k < B; ++k) {
-    float tbest = 3.0e38f;
     float bu = 0.0f, bv = 0.0f;
     float3 ng;
     int face = bvh_closest(sc, ro, rd, kRayEps, tbest, bu, bv, ng);
-    if (face < 0) break;
     int m = sc.mats[face];
     float3 x = ro + rd * tbest;
     float3 n = normalize(ng);
-    if (dot(n, rd) > 0.0f) n = f3(-n.x, -n.y, -n.z);
     float am[3];
     int idx4[4];
@@ -302,118 +499,231 @@ __device__ void trace_path(const DevScene& sc, float3 ro, float3 rd,
     sample_albedo(sc, m, face, bu, bv, am, idx4, w4);
     const float* em = &sc.emi[m * 3];
-    // Emission gathered at the camera hit (NEE on) or at every hit (NEE off).
-    if (k == 0 || !nee) {
-      if (!GRAD) {
-        acc.x += T[0] * em[0];
-        acc.y += T[1] * em[1];
-        acc.z += T[2] * em[2];
-      } else {
-        for (int c = 0; c < 3; ++c)
-          atomicAdd(&s_ge[m * 3 + c], gs[c] * T[c]);
-        if (k > 0) {
-          // dC/d texel: exclusion product over bounces j < k, scattered to
-          // bounce j's bilinear footprint.
-          suf[k][0] = suf[k][1] = suf[k][2] = 1.0f;
-          for (int i = k - 1; i >= 0; --i)
-            for (int c = 0; c < 3; ++c) suf[i][c] = suf[i + 1][c] * albs[i][c];
-          float pref[3] = {1.0f, 1.0f, 1.0f};
-          for (int j = 0; j < k; ++j) {
-            for (int c = 0; c < 3; ++c) {
-              float base = gs[c] * em[c] * pref[c] * suf[j + 1][c];
-#pragma unroll
-              for (int i = 0; i < 4; ++i)
-                scatter_tex(s_gt, g_tex, fp_i[j][i], c, base * fp_w[j][i]);
-              pref[c] *= albs[j][c];
-            }
-          }
-        }
       }
     }
-    if (GRAD) {
-      albs[k][0] = am[0];
-      albs[k][1] = am[1];
-      albs[k][2] = am[2];
 #pragma unroll
-      for (int i = 0; i < 4; ++i) {
-        fp_i[k][i] = idx4[i];
-        fp_w[k][i] = w4[i];
       }
-      mats[k] = m;
     }
-    // Next-event estimation against the area-sampled light list, skipped at
-    // the final vertex to match the BRDF-only estimator's truncation.
-    if (nee && sc.nl > 0 && k + 1 < B) {
       float r1 = curand_uniform(&st);
       float r2 = curand_uniform(&st);
       float r3 = curand_uniform(&st);
-      int li = pick_light(sc.lcdf, sc.nl, r1);
-      int lface = sc.lf[li];
-      const float* lv = &sc.tris[lface * 9];
-      float3 lv0 = f3(lv[0], lv[1], lv[2]);
-      float3 le1 = f3(lv[3], lv[4], lv[5]) - lv0;
-      float3 le2 = f3(lv[6], lv[7], lv[8]) - lv0;
-      float su = sqrtf(r2);
-      float b0 = 1.0f - su;
-      float b1 = r3 * su;
-      float3 y = lv0 + le1 * b0 + le2 * b1;  // uniform on the triangle
-      float3 ln = normalize(cross(le1, le2));
-      float3 dvec = y - x;
-      float d2 = fmaxf(dot(dvec, dvec), 1e-12f);
-      float d = sqrtf(d2);
-      float3 wi = dvec * (1.0f / d);
-      if (dot(ln, dvec) > 0.0f) ln = f3(-ln.x, -ln.y, -ln.z);
-      float cx = dot(n, wi);
-      float cy = -dot(ln, wi);
-      int lm = sc.mats[lface];
-      if (cx > 1e-6f && cy > 1e-6f && d > kShadowEps * 2.0f) {
-        float3 xo = x + n * kRayEps;
-        if (!bvh_occluded(sc, xo, wi, d - kShadowEps)) {
-          float S = cx * cy / d2 * sc.larea * kInvPi;
-          const float* el = &sc.emi[lm * 3];
-          if (!GRAD) {
-            acc.x += T[0] * am[0] * el[0] * S;
-            acc.y += T[1] * am[1] * el[1] * S;
-            acc.z += T[2] * am[2] * el[2] * S;
-          } else {
-            // Term = prod_{j<=k} a_j * E_l * S; albs[k] is already stored.
-            suf[k + 1][0] = suf[k + 1][1] = suf[k + 1][2] = 1.0f;
-            for (int i = k; i >= 0; --i)
-              for (int c = 0; c < 3; ++c)
-                suf[i][c] = suf[i + 1][c] * albs[i][c];
-            for (int c = 0; c < 3; ++c)
-              atomicAdd(&s_ge[lm * 3 + c], gs[c] * suf[0][c] * S);
-            float pref[3] = {1.0f, 1.0f, 1.0f};
-            for (int j = 0; j <= k; ++j) {
-              for (int c = 0; c < 3; ++c) {
-                float base = gs[c] * el[c] * S * pref[c] * suf[j + 1][c];
-#pragma unroll
-                for (int i = 0; i < 4; ++i)
-                  scatter_tex(s_gt, g_tex, fp_i[j][i], c, base * fp_w[j][i]);
-                pref[c] *= albs[j][c];
-              }
             }
           }
         }
       }
     }
-    // Cosine-weighted continuation: BRDF*cos/pdf = albedo exactly.
-    float rd1 = curand_uniform(&st);
-    float rd2 = curand_uniform(&st);
-    float rr = sqrtf(rd1);
-    float phi = 6.2831853071795864769f * rd2;
-    float3 tang, bit;
-    onb(n, tang, bit);
-    float3 nd = tang * (rr * cosf(phi)) + bit * (rr * sinf(phi)) +
-                n * sqrtf(fmaxf(0.0f, 1.0f - rd1));
-    rd = normalize(nd);
     ro = x + n * kRayEps;
-    T[0] *= am[0];
-    T[1] *= am[1];
-    T[2] *= am[2];
   }
 }
@@ -428,7 +738,7 @@ __device__ __forceinline__ float3 camera_ray(const DevCam& cam, int px, int py,
 }
 __global__ void k_forward(DevScene sc, DevCam cam, int H, int W, int spp,
-                          int B, int nee, unsigned long long seed,
                           float* img) {
   int pid = blockIdx.x * blockDim.x + threadIdx.x;
   if (pid >= H * W) return;
@@ -441,8 +751,8 @@ __global__ void k_forward(DevScene sc, DevCam cam, int H, int W, int spp,
     float jx = curand_uniform(&st);
     float jy = curand_uniform(&st);
     float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
-    trace_path<false>(sc, ro, rd, st, B, nee, acc, nullptr, nullptr, nullptr,
-                      nullptr);
   }
   float inv_spp = 1.0f / (float)spp;
   img[pid * 3 + 0] = acc.x * inv_spp;
@@ -451,12 +761,13 @@ __global__ void k_forward(DevScene sc, DevCam cam, int H, int W, int spp,
 }
 __global__ void k_backward(DevScene sc, DevCam cam, int H, int W, int spp,
-                           int B, int nee, unsigned long long seed,
-                           const float* gimg, float* g_tex, float* ge) {
   extern __shared__ float smem[];
-  float* s_ge = smem;                     // [nm * 3]
   bool staged = sc.nt <= kSharedTexels;
-  float* s_gt = staged ? smem + sc.nm * 3 : nullptr;  // [nt * 3] when staged
   int nm3 = sc.nm * 3;
   int nt3 = staged ? sc.nt * 3 : 0;
   for (int i = threadIdx.x; i < nm3 + nt3; i += blockDim.x) smem[i] = 0.0f;
@@ -476,7 +787,8 @@ __global__ void k_backward(DevScene sc, DevCam cam, int H, int W, int spp,
       float jx = curand_uniform(&st);
       float jy = curand_uniform(&st);
       float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
-      trace_path<true>(sc, ro, rd, st, B, nee, acc, gs, g_tex, s_gt, s_ge);
     }
   }
@@ -488,28 +800,32 @@ __global__ void k_backward(DevScene sc, DevCam cam, int H, int W, int spp,
   }
 }
-DevScene make_scene(const float* tris, const int* mat_ids, const float* uvs,
-                    const float* nodes_f, const int* nodes_i, int n_nodes,
-                    const int* light_faces, const float* light_cdf,
-                    int n_lights, float total_light_area, const float* tex,
-                    const int* tex_hdr, int n_texels, const float* emission,
-                    int n_mats) {
   DevScene sc;
-  sc.tris = tris;
-  sc.mats = mat_ids;
-  sc.uvs = uvs;
-  sc.nf = nodes_f;
-  sc.ni = nodes_i;
-  sc.n_nodes = n_nodes;
-  sc.lf = light_faces;
-  sc.lcdf = light_cdf;
-  sc.nl = n_lights;
-  sc.larea = total_light_area;
-  sc.tex = tex;
-  sc.thdr = tex_hdr;
-  sc.nt = n_texels;
-  sc.emi = emission;
-  sc.nm = n_mats;
   return sc;
 }
@@ -526,41 +842,30 @@ DevCam make_cam(const float* cam) {
 }  // namespace
-extern "C" void ptd_forward_launch(
-    const float* tris, const int* mat_ids, const float* uvs, int /*n_faces*/,
-    const float* nodes_f, const int* nodes_i, int n_nodes,
-    const int* light_faces, const float* light_cdf, int n_lights,
-    float total_light_area, const float* tex, const int* tex_hdr,
-    int n_texels, const float* emission, int n_mats, const float* cam, int H,
-    int W, int spp, int max_bounces, int nee, long long seed, float* image,
-    cudaStream_t stream) {
-  DevScene sc = make_scene(tris, mat_ids, uvs, nodes_f, nodes_i, n_nodes,
-                           light_faces, light_cdf, n_lights, total_light_area,
-                           tex, tex_hdr, n_texels, emission, n_mats);
   DevCam dc = make_cam(cam);
   int blocks = (H * W + kThreads - 1) / kThreads;
   k_forward<<<blocks, kThreads, 0, stream>>>(
-      sc, dc, H, W, spp, max_bounces, nee, (unsigned long long)seed, image);
-}
-extern "C" void ptd_backward_launch(
-    const float* tris, const int* mat_ids, const float* uvs, int /*n_faces*/,
-    const float* nodes_f, const int* nodes_i, int n_nodes,
-    const int* light_faces, const float* light_cdf, int n_lights,
-    float total_light_area, const float* tex, const int* tex_hdr,
-    int n_texels, const float* emission, int n_mats, const float* cam, int H,
-    int W, int spp, int max_bounces, int nee, long long seed,
-    const float* grad_image, float* grad_tex, float* grad_emission,
-    cudaStream_t stream) {
-  DevScene sc = make_scene(tris, mat_ids, uvs, nodes_f, nodes_i, n_nodes,
-                           light_faces, light_cdf, n_lights, total_light_area,
-                           tex, tex_hdr, n_texels, emission, n_mats);
   DevCam dc = make_cam(cam);
   int blocks = (H * W + kThreads - 1) / kThreads;
-  bool staged = n_texels <= kSharedTexels;
-  size_t smem = (size_t)(n_mats * 3 + (staged ? n_texels * 3 : 0)) *
-                sizeof(float);
   k_backward<<<blocks, kThreads, smem, stream>>>(
-      sc, dc, H, W, spp, max_bounces, nee, (unsigned long long)seed,
-      grad_image, grad_tex, grad_emission);
 }

 // pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
 //
+// Forward: a megakernel unidirectional path tracer. Materials: Lambertian
+// diffuse, GGX conductor (height-correlated Smith, VNDF sampling), and
+// smooth dielectric (Fresnel-branched reflection/refraction). Lights: any
+// emissive face (area-sampled) and an optional importance-sampled
+// equirectangular environment map. Estimators: multiple importance
+// sampling (balance heuristic) by default, with legacy NEE-only and
+// BRDF-only modes. Per-texel albedo with bilinear filtering; binned-SAH
+// BVH with near-child-first traversal. One thread owns one pixel and
+// accumulates its spp samples serially, so the image is bitwise
+// deterministic.
 //
+// Backward: exact path replay. Sampling distributions depend only on
+// geometry, frozen material parameters (roughness, ior), the detached
+// environment CDF, and the counter-based Philox stream -- never on the
+// differentiable parameters (albedo texels, per-material emission,
+// environment texels). The backward kernel therefore 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 own vertex's
+// albedo texels (diffuse: a; conductor: Schlick F(F0) is affine in F0;
+// dielectric: constant), times an emission or environment texel (linear).
+// 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 -- the exact adjoint of the fetch.
+// Interior terms only: no visibility/silhouette gradients.
 #include <cuda_runtime.h>
 #include <curand_kernel.h>
 constexpr int kThreads = 128;
 constexpr int kMaxBounces = 16;
 constexpr int kMaxMats = 64;
+constexpr int kSharedTexels = 2048;  // shared-staged texture-gradient limit
 constexpr int kStack = 64;
+constexpr float kPi = 3.14159265358979323846f;
 constexpr float kInvPi = 0.31830988618379067154f;
 constexpr float kRayEps = 1e-4f;
 constexpr float kShadowEps = 1e-3f;
+constexpr float kEnvDist = 1e30f;
+enum Mode { kModeBrdf = 0, kModeNee = 1, kModeMis = 2 };
+enum MatType { kDiffuse = 0, kConductor = 1, kDielectric = 2 };
 // ---------------------------------------------------------------- float3 ops
 __device__ __forceinline__ float3 f3(float x, float y, float z) {
 __device__ __forceinline__ float3 normalize(float3 a) {
   return a * rsqrtf(fmaxf(dot(a, a), 1e-30f));
 }
+__device__ __forceinline__ float3 neg(float3 a) { return f3(-a.x, -a.y, -a.z); }
 // Duff et al. 2017, "Building an Orthonormal Basis, Revisited".
 __device__ __forceinline__ void onb(float3 n, float3& t, float3& b) {
 // ---------------------------------------------------------------- scene view
 struct DevScene {
+  const float* tris;   // [F, 9]
   const int* mats;     // [F]
+  const float* uvs;    // [F, 3, 2]
+  const float* nf;     // [N, 6]
+  const int* ni;       // [N, 3] internal (l, r, axis<<1) | leaf (start, count, odd)
   int n_nodes;
+  const int* lf;       // [L] emissive faces
+  const float* lcdf;   // [L]
   int nl;
+  float larea;
+  const float* tex;    // [T, 3] albedo texels
+  const int* thdr;     // [M, 3] (offset, W, H)
+  int nt;
   const float* emi;    // [M, 3]
+  const int* mtype;    // [M] material type
+  const float* mrough; // [M] GGX alpha (conductor)
+  const float* mior;   // [M] ior (dielectric)
   int nm;
+  // environment map (equirect) + detached sampling tables
+  const float* env;      // [Eh*Ew, 3] or null
+  int ew, eh;
+  const float* ecdf_m;   // [Eh] marginal cdf over rows
+  const float* ecdf_c;   // [Eh*Ew] conditional cdf per row
+  const float* epdf;     // [Eh*Ew] table pdf over the image (integrates to 1)
 };
+struct DevCam {
   float p[3], f[3], r[3], u[3];
 };
                            float tmin, float& tbest, float& bu, float& bv,
                            float3& ngbest) {
   float roa[3] = {ro.x, ro.y, ro.z};
+  float dira[3] = {rd.x, rd.y, rd.z};
   float inv[3];
   inv_dir(rd, inv);
   int stack[kStack];
     int nid = stack[--sp];
     if (!slab(&sc.nf[nid * 6], roa, inv, tbest)) continue;
     const int* n = &sc.ni[nid * 3];
+    if (n[2] & 1) {
       for (int f = n[0]; f < n[0] + n[1]; ++f) {
         float t, u, v;
         float3 ng;
         }
       }
     } else if (sp + 2 <= kStack) {
+      int axis = n[2] >> 1;
+      int near = (dira[axis] >= 0.0f) ? n[0] : n[1];
+      int far = (dira[axis] >= 0.0f) ? n[1] : n[0];
+      stack[sp++] = far;
+      stack[sp++] = near;
     }
   }
   return best;
     int nid = stack[--sp];
     if (!slab(&sc.nf[nid * 6], roa, inv, tmax)) continue;
     const int* n = &sc.ni[nid * 3];
+    if (n[2] & 1) {
       for (int f = n[0]; f < n[0] + n[1]; ++f) {
         float t, u, v;
         float3 ng;
   return false;
 }
+__device__ __forceinline__ int cdf_pick(const float* cdf, int n, float r) {
+  int lo = 0, hi = n - 1;
   while (lo < hi) {
     int mid = (lo + hi) >> 1;
     if (cdf[mid] < r) lo = mid + 1; else hi = mid;
   return lo;
 }
+// Bilinear fetch with repeat wrap over a flat texel block.
+__device__ __forceinline__ void bilinear(const float* block, int off, int W,
+                                         int H, float u, float v, float rgb[3],
+                                         int idx4[4], float w4[4]) {
   float uu = u * (float)W - 0.5f;
   float vv = v * (float)H - 0.5f;
   float fx = uu - floorf(uu);
   rgb[0] = rgb[1] = rgb[2] = 0.0f;
 #pragma unroll
   for (int i = 0; i < 4; ++i) {
+    const float* t = &block[idx4[i] * 3];
     rgb[0] += w4[i] * t[0];
     rgb[1] += w4[i] * t[1];
     rgb[2] += w4[i] * t[2];
   }
 }
+__device__ __forceinline__ void sample_albedo(const DevScene& sc, int m,
+                                              int face, float bu, float bv,
+                                              float rgb[3], int idx4[4],
+                                              float w4[4]) {
+  const float* U = &sc.uvs[face * 6];
+  float w0 = 1.0f - bu - bv;
+  float u = w0 * U[0] + bu * U[2] + bv * U[4];
+  float v = w0 * U[1] + bu * U[3] + bv * U[5];
+  const int* h = &sc.thdr[m * 3];
+  bilinear(sc.tex, h[0], h[1], h[2], u, v, rgb, idx4, w4);
+}
+// ----------------------------------------------------------- environment map
+__device__ __forceinline__ void dir_to_equirect(float3 d, float& u, float& v) {
+  float phi = atan2f(d.z, d.x);                 // [-pi, pi]
+  float theta = acosf(fminf(fmaxf(d.y, -1.0f), 1.0f));  // [0, pi], y-up
+  u = (phi + kPi) / (2.0f * kPi);
+  v = theta / kPi;
+}
+__device__ __forceinline__ float3 equirect_to_dir(float u, float v) {
+  float phi = u * 2.0f * kPi - kPi;
+  float theta = v * kPi;
+  float st = sinf(theta);
+  return f3(st * cosf(phi), cosf(theta), st * sinf(phi));
+}
+// Solid-angle pdf of the detached env sampling mixture (0.5 table + 0.5
+// uniform sphere).
+__device__ __forceinline__ float env_pdf(const DevScene& sc, float3 d) {
+  float u, v;
+  dir_to_equirect(d, u, v);
+  int x = min(sc.ew - 1, (int)(u * sc.ew));
+  int y = min(sc.eh - 1, (int)(v * sc.eh));
+  float st = fmaxf(sinf((y + 0.5f) * kPi / sc.eh), 1e-4f);
+  float p_img = sc.epdf[y * sc.ew + x];  // per-texel image pdf (sums to 1)
+  float p_tab = p_img * (float)(sc.ew * sc.eh) / (2.0f * kPi * kPi * st);
+  return 0.5f * p_tab + 0.5f * (1.0f / (4.0f * kPi));
+}
+__device__ __forceinline__ void env_fetch(const DevScene& sc, float3 d,
+                                          float rgb[3], int idx4[4],
+                                          float w4[4]) {
+  float u, v;
+  dir_to_equirect(d, u, v);
+  // clamp v away from the poles to keep the bilinear rows in range
+  v = fminf(fmaxf(v, 0.5f / sc.eh), 1.0f - 0.5f / sc.eh);
+  bilinear(sc.env, 0, sc.ew, sc.eh, u, v, rgb, idx4, w4);
+}
+// ------------------------------------------------------------------- GGX
+__device__ __forceinline__ float ggx_lambda(float a2, float cs) {
+  float c2 = cs * cs;
+  float t2 = fmaxf(0.0f, 1.0f - c2) / fmaxf(c2, 1e-12f);
+  return 0.5f * (-1.0f + sqrtf(1.0f + a2 * t2));
+}
+__device__ __forceinline__ float ggx_d(float a2, float ch) {
+  float d = ch * ch * (a2 - 1.0f) + 1.0f;
+  return a2 / fmaxf(kPi * d * d, 1e-20f);
+}
+// Heitz 2018 VNDF sampling; wi in local frame (z up), wi.z > 0.
+__device__ __forceinline__ float3 ggx_sample_vndf(float3 wi, float alpha,
+                                                  float u1, float u2) {
+  float3 vh = normalize(f3(alpha * wi.x, alpha * wi.y, wi.z));
+  float lensq = vh.x * vh.x + vh.y * vh.y;
+  float3 T1 = lensq > 1e-12f ? f3(-vh.y, vh.x, 0.0f) * rsqrtf(lensq)
+                             : f3(1.0f, 0.0f, 0.0f);
+  float3 T2 = cross(vh, T1);
+  float r = sqrtf(u1);
+  float phi = 2.0f * kPi * u2;
+  float t1 = r * cosf(phi);
+  float t2 = r * sinf(phi);
+  float s = 0.5f * (1.0f + vh.z);
+  t2 = (1.0f - s) * sqrtf(fmaxf(0.0f, 1.0f - t1 * t1)) + s * t2;
+  float3 nh = T1 * t1 + T2 * t2 +
+              vh * sqrtf(fmaxf(0.0f, 1.0f - t1 * t1 - t2 * t2));
+  return normalize(f3(alpha * nh.x, alpha * nh.y, fmaxf(1e-6f, nh.z)));
+}
+// pdf of wo under VNDF sampling (solid angle).
+__device__ __forceinline__ float ggx_pdf(float a2, float3 wi, float3 wo) {
+  float3 h = normalize(wi + wo);
+  float ch = fmaxf(h.z, 1e-6f);
+  float wih = fmaxf(dot(wi, h), 1e-6f);
+  float g1 = 1.0f / (1.0f + ggx_lambda(a2, wi.z));
+  return g1 * ggx_d(a2, ch) * wih / fmaxf(wi.z, 1e-6f) / (4.0f * wih);
+}
+// Fresnel (dielectric, unpolarized), cos_i >= 0; n1 -> n2.
+__device__ __forceinline__ float fresnel_dielectric(float cos_i, float eta) {
+  float s2 = (1.0f - cos_i * cos_i) / (eta * eta);
+  if (s2 >= 1.0f) return 1.0f;  // TIR
+  float cos_t = sqrtf(1.0f - s2);
+  float rs = (cos_i - eta * cos_t) / (cos_i + eta * cos_t);
+  float rp = (eta * cos_i - cos_t) / (eta * cos_i + cos_t);
+  return 0.5f * (rs * rs + rp * rp);
+}
+// Schlick term split so F(F0) = F0 * (1 - s) + s: affine in F0.
+__device__ __forceinline__ float schlick_s(float cos_h) {
+  float m = fminf(fmaxf(1.0f - cos_h, 0.0f), 1.0f);
+  float m2 = m * m;
+  return m2 * m2 * m;
+}
+__device__ __forceinline__ float mis_w(float pa, float pb) {
+  return pa / fmaxf(pa + pb, 1e-20f);
+}
+// ------------------------------------------------------------- path tracing
 __device__ __forceinline__ void scatter_tex(float* s_gt, float* g_tex,
                                             int idx, int c, float v) {
+  if (s_gt) atomicAdd(&s_gt[idx * 3 + c], v);
+  else atomicAdd(&g_tex[idx * 3 + c], v);
+}
+// Per-bounce factor record for the replay: value and d(value)/d(albedo
+// texel channel), both per channel, plus the texel footprint.
+struct Factor {
+  float val[3];
+  float dc[3];
+  int fi[4];
+  float fw[4];
+};
+// Accumulate one radiance term and (GRAD) its gradients:
+//   term = (prod over the nk factors fs, channelwise) * S * E
+// where fs holds the continuation factors and, for NEE-type terms, the
+// vertex's own BSDF factor as the last entry. T is the product of the
+// continuation factors only; tf is the last factor's value (null for
+// emission/env gathers, whose product is T alone). E is emission (index
+// em_mat >= 0) or environment (footprint env_i/env_w).
+template <bool GRAD>
+__device__ void add_term(const Factor* fs, int nk, const float T[3],
+                         const float* tf, float S, const float E[3],
+                         int em_mat, const int env_i[4], const float env_w[4],
+                         float3& acc, const float gs[3], float* g_tex,
+                         float* s_gt, float* s_ge, float* g_env) {
+  float P[3] = {T[0], T[1], T[2]};
+  if (tf) {
+    P[0] *= tf[0];
+    P[1] *= tf[1];
+    P[2] *= tf[2];
+  }
+  if (!GRAD) {
+    acc.x += P[0] * S * E[0];
+    acc.y += P[1] * S * E[1];
+    acc.z += P[2] * S * E[2];
+    return;
+  }
+  // gradient to the emission / environment texels: g * prodA * S
+  if (em_mat >= 0) {
+    for (int c = 0; c < 3; ++c)
+      atomicAdd(&s_ge[em_mat * 3 + c], gs[c] * P[c] * S);
+  } else if (env_i) {
+    for (int c = 0; c < 3; ++c) {
+      float base = gs[c] * P[c] * S;
+#pragma unroll
+      for (int i = 0; i < 4; ++i)
+        atomicAdd(&g_env[env_i[i] * 3 + c], base * env_w[i]);
+    }
+  }
+  // gradient to each factor's albedo texels via exclusion products
+  if (nk <= 0) return;
+  float suf[kMaxBounces + 1][3];
+  suf[nk][0] = suf[nk][1] = suf[nk][2] = 1.0f;
+  for (int i = nk - 1; i >= 0; --i)
+    for (int c = 0; c < 3; ++c) suf[i][c] = suf[i + 1][c] * fs[i].val[c];
+  float pref[3] = {1.0f, 1.0f, 1.0f};
+  for (int j = 0; j < nk; ++j) {
+    for (int c = 0; c < 3; ++c) {
+      if (fs[j].dc[c] != 0.0f) {
+        float base = gs[c] * E[c] * S * pref[c] * suf[j + 1][c] * fs[j].dc[c];
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+          scatter_tex(s_gt, g_tex, fs[j].fi[i], c, base * fs[j].fw[i]);
+      }
+      pref[c] *= fs[j].val[c];
+    }
   }
 }
 template <bool GRAD>
 __device__ void trace_path(const DevScene& sc, float3 ro, float3 rd,
+                           curandStatePhilox4_32_10_t& st, int B, int mode,
                            float3& acc, const float gs[3], float* g_tex,
+                           float* s_gt, float* s_ge, float* g_env) {
+  float T[3] = {1.0f, 1.0f, 1.0f};
+  Factor fs[kMaxBounces];
+  int nk = 0;                 // factors recorded so far (= bounces done)
+  float prev_pdf = 0.0f;      // solid-angle pdf of the previous BSDF sample
+  bool prev_delta = true;     // camera counts as delta (w = 1 on first hit)
+  bool has_env = sc.env != nullptr;
   for (int k = 0; k < B; ++k) {
+    float tbest = kEnvDist;
     float bu = 0.0f, bv = 0.0f;
     float3 ng;
     int face = bvh_closest(sc, ro, rd, kRayEps, tbest, bu, bv, ng);
+    if (face < 0) {
+      // Environment gather with MIS against the env-NEE at the previous
+      // vertex (weight 1 after the camera or a delta vertex, or outside
+      // MIS mode).
+      if (has_env) {
+        float E[3];
+        int ei[4];
+        float ew[4];
+        env_fetch(sc, rd, E, ei, ew);
+        float w = 1.0f;
+        if (mode == kModeMis && !prev_delta)
+          w = mis_w(prev_pdf, env_pdf(sc, rd));
+        if (mode != kModeNee || prev_delta)
+          add_term<GRAD>(fs, nk, T, nullptr, w, E, -1, ei, ew, acc, gs,
+                         g_tex, s_gt, s_ge, g_env);
+      }
+      break;
+    }
     int m = sc.mats[face];
+    int mt = sc.mtype[m];
     float3 x = ro + rd * tbest;
     float3 n = normalize(ng);
+    bool backface = dot(n, rd) > 0.0f;
+    if (backface) n = neg(n);
     float am[3];
     int idx4[4];
     sample_albedo(sc, m, face, bu, bv, am, idx4, w4);
     const float* em = &sc.emi[m * 3];
+    // Emission gather. BRDF mode: always. NEE mode: camera hit only. MIS:
+    // always, weighted against the light-NEE at the previous vertex.
+    bool emissive = em[0] > 0.0f || em[1] > 0.0f || em[2] > 0.0f;
+    if (emissive) {
+      float w = 1.0f;
+      bool add = true;
+      if (mode == kModeNee) {
+        add = (k == 0);
+      } else if (mode == kModeMis && !prev_delta && sc.nl > 0) {
+        float cy = fabsf(dot(normalize(ng), rd));
+        float p_l = (tbest * tbest) / fmaxf(cy * sc.larea, 1e-12f);
+        w = mis_w(prev_pdf, p_l);
+      }
+      if (add) {
+        float E[3] = {em[0], em[1], em[2]};
+        add_term<GRAD>(fs, nk, T, nullptr, w, E, m, nullptr, nullptr, acc, gs,
+                       g_tex, s_gt, s_ge, g_env);
       }
     }
+    // ---------------- dielectric: delta lobes, no NEE, unit factor
+    if (mt == kDielectric) {
+      float r1 = curand_uniform(&st);
+      float r2 = curand_uniform(&st);
+      (void)r2;
+      float eta = backface ? sc.mior[m] : 1.0f / sc.mior[m];  // n1/n2
+      float ci = -dot(rd, n);
+      float F = fresnel_dielectric(ci, 1.0f / eta);
+      float3 nd;
+      if (r1 < F) {
+        nd = rd + n * (2.0f * ci);
+        ro = x + n * kRayEps;
+      } else {
+        float s2 = eta * eta * (1.0f - ci * ci);
+        float ct = sqrtf(fmaxf(0.0f, 1.0f - s2));
+        nd = rd * eta + n * (eta * ci - ct);
+        ro = x - n * kRayEps;
+      }
+      rd = normalize(nd);
+      if (nk < kMaxBounces) {
+        Factor& f = fs[nk++];
+        for (int c = 0; c < 3; ++c) { f.val[c] = 1.0f; f.dc[c] = 0.0f; }
 #pragma unroll
+        for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
       }
+      prev_delta = true;
+      prev_pdf = 0.0f;
+      continue;
     }
+    // Local frame; wi = toward camera.
+    float3 tang, bit;
+    onb(n, tang, bit);
+    float3 wi_w = neg(rd);
+    float3 wi = f3(dot(wi_w, tang), dot(wi_w, bit), dot(wi_w, n));
+    wi.z = fmaxf(wi.z, 1e-6f);
+    float alpha = fmaxf(sc.mrough[m], 0.01f);
+    float a2 = alpha * alpha;
+    // ---------------- light NEE (area lights)
+    bool nee_here = (mode == kModeNee && sc.nl > 0 && k + 1 < B) ||
+                    (mode == kModeMis && sc.nl > 0);
+    if ((mode != kModeBrdf) && sc.nl > 0 &&
+        (mode == kModeMis || k + 1 < B)) {
       float r1 = curand_uniform(&st);
       float r2 = curand_uniform(&st);
       float r3 = curand_uniform(&st);
+      if (nee_here) {
+        int li = cdf_pick(sc.lcdf, sc.nl, r1);
+        int lface = sc.lf[li];
+        const float* lv = &sc.tris[lface * 9];
+        float3 lv0 = f3(lv[0], lv[1], lv[2]);
+        float3 le1 = f3(lv[3], lv[4], lv[5]) - lv0;
+        float3 le2 = f3(lv[6], lv[7], lv[8]) - lv0;
+        float su = sqrtf(r2);
+        float3 y = lv0 + le1 * (1.0f - su) + le2 * (r3 * su);
+        float3 ln = normalize(cross(le1, le2));
+        float3 dvec = y - x;
+        float d2 = fmaxf(dot(dvec, dvec), 1e-12f);
+        float d = sqrtf(d2);
+        float3 wo_w = dvec * (1.0f / d);
+        if (dot(ln, dvec) > 0.0f) ln = neg(ln);
+        float cx = dot(n, wo_w);
+        float cy = -dot(ln, wo_w);
+        int lm = sc.mats[lface];
+        if (cx > 1e-6f && cy > 1e-6f && d > kShadowEps * 2.0f &&
+            !bvh_occluded(sc, x + n * kRayEps, wo_w, d - kShadowEps)) {
+          float3 wo = f3(dot(wo_w, tang), dot(wo_w, bit), cx);
+          // f * cos / pdf_A, with pdf_A -> S_geo = cx*cy/d2 * larea
+          float S_geo = cx * cy / d2 * sc.larea;
+          float tf_val[3], tf_dc[3];
+          float S;
+          if (mt == kDiffuse) {
+            S = S_geo * kInvPi;
+            for (int c = 0; c < 3; ++c) { tf_val[c] = am[c]; tf_dc[c] = 1.0f; }
+          } else {  // conductor
+            float3 h = normalize(wi + wo);
+            float sgl = schlick_s(fmaxf(dot(wi, h), 0.0f));
+            float G2 = 1.0f / (1.0f + ggx_lambda(a2, wi.z) +
+                               ggx_lambda(a2, wo.z));
+            float gg = ggx_d(a2, fmaxf(h.z, 1e-6f)) * G2 /
+                       (4.0f * wi.z * fmaxf(wo.z, 1e-6f));
+            S = S_geo * gg;
+            for (int c = 0; c < 3; ++c) {
+              tf_val[c] = am[c] * (1.0f - sgl) + sgl;
+              tf_dc[c] = 1.0f - sgl;
             }
           }
+          float w = 1.0f;
+          if (mode == kModeMis) {
+            float p_lw = d2 / fmaxf(cy * sc.larea, 1e-12f);
+            float p_bw = (mt == kDiffuse) ? fmaxf(wo.z, 0.0f) * kInvPi
+                                          : ggx_pdf(a2, wi, wo);
+            w = mis_w(p_lw, p_bw);
+          }
+          if (nk < kMaxBounces) {
+            Factor& f = fs[nk];
+            for (int c = 0; c < 3; ++c) { f.val[c] = tf_val[c]; f.dc[c] = tf_dc[c]; }
+#pragma unroll
+            for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
+            float E[3] = {sc.emi[lm * 3], sc.emi[lm * 3 + 1], sc.emi[lm * 3 + 2]};
+            add_term<GRAD>(fs, nk + 1, T, f.val, S * w, E, lm, nullptr, nullptr,
+                           acc, gs, g_tex, s_gt, s_ge, g_env);
+          }
         }
       }
     }
+    // ---------------- environment NEE (MIS mode only)
+    if (mode == kModeMis && has_env) {
+      float e1 = curand_uniform(&st);
+      float e2 = curand_uniform(&st);
+      float e3 = curand_uniform(&st);
+      float3 wo_w;
+      if (e1 < 0.5f) {  // table
+        int y = cdf_pick(sc.ecdf_m, sc.eh, e2);
+        int xcol = cdf_pick(&sc.ecdf_c[y * sc.ew], sc.ew, e3);
+        float u = (xcol + 0.5f) / sc.ew;
+        float v = (y + 0.5f) / sc.eh;
+        wo_w = equirect_to_dir(u, v);
+      } else {  // uniform sphere
+        float z = 1.0f - 2.0f * e2;
+        float rxy = sqrtf(fmaxf(0.0f, 1.0f - z * z));
+        float ph = 2.0f * kPi * e3;
+        wo_w = f3(rxy * cosf(ph), z, rxy * sinf(ph));
+      }
+      float cx = dot(n, wo_w);
+      if (cx > 1e-6f &&
+          !bvh_occluded(sc, x + n * kRayEps, wo_w, kEnvDist)) {
+        float p = env_pdf(sc, wo_w);
+        float3 wo = f3(dot(wo_w, tang), dot(wo_w, bit), cx);
+        float tf_val[3], tf_dc[3];
+        float S;
+        if (mt == kDiffuse) {
+          S = cx * kInvPi / p;
+          for (int c = 0; c < 3; ++c) { tf_val[c] = am[c]; tf_dc[c] = 1.0f; }
+        } else {
+          float3 h = normalize(wi + wo);
+          float sgl = schlick_s(fmaxf(dot(wi, h), 0.0f));
+          float G2 = 1.0f / (1.0f + ggx_lambda(a2, wi.z) +
+                             ggx_lambda(a2, wo.z));
+          float gg = ggx_d(a2, fmaxf(h.z, 1e-6f)) * G2 /
+                     (4.0f * wi.z * fmaxf(wo.z, 1e-6f));
+          S = cx * gg / p;
+          for (int c = 0; c < 3; ++c) {
+            tf_val[c] = am[c] * (1.0f - sgl) + sgl;
+            tf_dc[c] = 1.0f - sgl;
+          }
+        }
+        float p_bw = (mt == kDiffuse) ? fmaxf(wo.z, 0.0f) * kInvPi
+                                      : ggx_pdf(a2, wi, wo);
+        float w = mis_w(p, p_bw);
+        float E[3];
+        int ei[4];
+        float ewt[4];
+        env_fetch(sc, wo_w, E, ei, ewt);
+        if (nk < kMaxBounces) {
+          Factor& f = fs[nk];
+          for (int c = 0; c < 3; ++c) { f.val[c] = tf_val[c]; f.dc[c] = tf_dc[c]; }
+#pragma unroll
+          for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
+          add_term<GRAD>(fs, nk + 1, T, f.val, S * w, E, -1, ei, ewt, acc, gs,
+                         g_tex, s_gt, s_ge, g_env);
+        }
+      }
+    }
+    // ---------------- BSDF continuation
+    float c1 = curand_uniform(&st);
+    float c2 = curand_uniform(&st);
+    float3 wo;
+    float fv[3], fdc[3];
+    if (mt == kDiffuse) {
+      float rr = sqrtf(c1);
+      float phi = 2.0f * kPi * c2;
+      wo = f3(rr * cosf(phi), rr * sinf(phi), sqrtf(fmaxf(0.0f, 1.0f - c1)));
+      prev_pdf = fmaxf(wo.z, 1e-8f) * kInvPi;
+      for (int c = 0; c < 3; ++c) { fv[c] = am[c]; fdc[c] = 1.0f; }
+    } else {  // conductor: VNDF sample, weight = F * G2/G1
+      float3 h = ggx_sample_vndf(wi, alpha, c1, c2);
+      float wih = dot(wi, h);
+      wo = h * (2.0f * wih) - wi;
+      if (wo.z <= 1e-6f) break;  // sampled below the horizon: zero weight
+      float sgl = schlick_s(fmaxf(wih, 0.0f));
+      float li = ggx_lambda(a2, wi.z);
+      float lo = ggx_lambda(a2, wo.z);
+      float gw = (1.0f + li) / (1.0f + li + lo);  // G2/G1
+      prev_pdf = ggx_pdf(a2, wi, wo);
+      for (int c = 0; c < 3; ++c) {
+        fv[c] = (am[c] * (1.0f - sgl) + sgl) * gw;
+        fdc[c] = (1.0f - sgl) * gw;
+      }
+    }
+    prev_delta = false;
+    rd = normalize(tang * wo.x + bit * wo.y + n * wo.z);
     ro = x + n * kRayEps;
+    if (nk < kMaxBounces) {
+      Factor& f = fs[nk++];
+      for (int c = 0; c < 3; ++c) { f.val[c] = fv[c]; f.dc[c] = fdc[c]; }
+#pragma unroll
+      for (int i = 0; i < 4; ++i) { f.fi[i] = idx4[i]; f.fw[i] = w4[i]; }
+    }
+    T[0] *= fv[0];
+    T[1] *= fv[1];
+    T[2] *= fv[2];
   }
 }
 }
 __global__ void k_forward(DevScene sc, DevCam cam, int H, int W, int spp,
+                          int B, int mode, unsigned long long seed,
                           float* img) {
   int pid = blockIdx.x * blockDim.x + threadIdx.x;
   if (pid >= H * W) return;
     float jx = curand_uniform(&st);
     float jy = curand_uniform(&st);
     float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
+    trace_path<false>(sc, ro, rd, st, B, mode, acc, nullptr, nullptr, nullptr,
+                      nullptr, nullptr);
   }
   float inv_spp = 1.0f / (float)spp;
   img[pid * 3 + 0] = acc.x * inv_spp;
 }
 __global__ void k_backward(DevScene sc, DevCam cam, int H, int W, int spp,
+                           int B, int mode, unsigned long long seed,
+                           const float* gimg, float* g_tex, float* ge,
+                           float* g_env) {
   extern __shared__ float smem[];
+  float* s_ge = smem;
   bool staged = sc.nt <= kSharedTexels;
+  float* s_gt = staged ? smem + sc.nm * 3 : nullptr;
   int nm3 = sc.nm * 3;
   int nt3 = staged ? sc.nt * 3 : 0;
   for (int i = threadIdx.x; i < nm3 + nt3; i += blockDim.x) smem[i] = 0.0f;
       float jx = curand_uniform(&st);
       float jy = curand_uniform(&st);
       float3 rd = camera_ray(cam, px, py, W, H, jx, jy);
+      trace_path<true>(sc, ro, rd, st, B, mode, acc, gs, g_tex, s_gt, s_ge,
+                       g_env);
     }
   }
   }
 }
+DevScene make_scene(const PtdSceneArgs& a) {
   DevScene sc;
+  sc.tris = a.tris;
+  sc.mats = a.mat_ids;
+  sc.uvs = a.uvs;
+  sc.nf = a.nodes_f;
+  sc.ni = a.nodes_i;
+  sc.n_nodes = a.n_nodes;
+  sc.lf = a.light_faces;
+  sc.lcdf = a.light_cdf;
+  sc.nl = a.n_lights;
+  sc.larea = a.total_light_area;
+  sc.tex = a.tex;
+  sc.thdr = a.tex_hdr;
+  sc.nt = a.n_texels;
+  sc.emi = a.emission;
+  sc.mtype = a.mat_type;
+  sc.mrough = a.mat_rough;
+  sc.mior = a.mat_ior;
+  sc.nm = a.n_mats;
+  sc.env = a.env;
+  sc.ew = a.env_w;
+  sc.eh = a.env_h;
+  sc.ecdf_m = a.env_cdf_m;
+  sc.ecdf_c = a.env_cdf_c;
+  sc.epdf = a.env_pdf;
   return sc;
 }
 }  // namespace
+extern "C" void ptd_forward_launch(const PtdSceneArgs* args, const float* cam,
+                                   int H, int W, int spp, int max_bounces,
+                                   int mode, long long seed, float* image,
+                                   cudaStream_t stream) {
+  DevScene sc = make_scene(*args);
   DevCam dc = make_cam(cam);
   int blocks = (H * W + kThreads - 1) / kThreads;
   k_forward<<<blocks, kThreads, 0, stream>>>(
+      sc, dc, H, W, spp, max_bounces, mode, (unsigned long long)seed, image);
+}
+extern "C" void ptd_backward_launch(const PtdSceneArgs* args, const float* cam,
+                                    int H, int W, int spp, int max_bounces,
+                                    int mode, long long seed,
+                                    const float* grad_image, float* grad_tex,
+                                    float* grad_emission, float* grad_env,
+                                    cudaStream_t stream) {
+  DevScene sc = make_scene(*args);
   DevCam dc = make_cam(cam);
   int blocks = (H * W + kThreads - 1) / kThreads;
+  bool staged = args->n_texels <= kSharedTexels;
+  size_t smem = (size_t)(args->n_mats * 3 +
+                         (staged ? args->n_texels * 3 : 0)) * sizeof(float);
   k_backward<<<blocks, kThreads, smem, stream>>>(
+      sc, dc, H, W, spp, max_bounces, mode, (unsigned long long)seed,
+      grad_image, grad_tex, grad_emission, grad_env);
 }

pathtracer_diff_cuda/pathtracer_launch.h CHANGED Viewed

@@ -6,27 +6,43 @@
 extern "C" {
-void ptd_forward_launch(
-    const float* tris, const int* mat_ids, const float* uvs, int n_faces,
-    const float* nodes_f, const int* nodes_i, int n_nodes,
-    const int* light_faces, const float* light_cdf, int n_lights,
-    float total_light_area,
-    const float* tex, const int* tex_hdr, int n_texels,
-    const float* emission, int n_mats,
-    const float* cam,
-    int H, int W, int spp, int max_bounces, int nee,
-    long long seed, float* image, cudaStream_t stream);
-void ptd_backward_launch(
-    const float* tris, const int* mat_ids, const float* uvs, int n_faces,
-    const float* nodes_f, const int* nodes_i, int n_nodes,
-    const int* light_faces, const float* light_cdf, int n_lights,
-    float total_light_area,
-    const float* tex, const int* tex_hdr, int n_texels,
-    const float* emission, int n_mats,
-    const float* cam,
-    int H, int W, int spp, int max_bounces, int nee,
-    long long seed, const float* grad_image,
-    float* grad_tex, float* grad_emission, cudaStream_t stream);
 }  // extern "C"

 extern "C" {
+// All device pointers; ints/floats are host scalars. env* may be null/0
+// when the scene has no environment map.
+struct PtdSceneArgs {
+  const float* tris;         // [F, 9]
+  const int* mat_ids;        // [F]
+  const float* uvs;          // [F, 3, 2]
+  int n_faces;
+  const float* nodes_f;      // [N, 6]
+  const int* nodes_i;        // [N, 3]
+  int n_nodes;
+  const int* light_faces;    // [L]
+  const float* light_cdf;    // [L]
+  int n_lights;
+  float total_light_area;
+  const float* tex;          // [T, 3]
+  const int* tex_hdr;        // [M, 3]
+  int n_texels;
+  const float* emission;     // [M, 3]
+  const int* mat_type;       // [M] 0 diffuse | 1 conductor | 2 dielectric
+  const float* mat_rough;    // [M]
+  const float* mat_ior;      // [M]
+  int n_mats;
+  const float* env;          // [Eh*Ew, 3] or null
+  int env_w, env_h;
+  const float* env_cdf_m;    // [Eh]
+  const float* env_cdf_c;    // [Eh*Ew]
+  const float* env_pdf;      // [Eh*Ew], sums to 1
+};
+void ptd_forward_launch(const PtdSceneArgs* args, const float* cam, int H,
+                        int W, int spp, int max_bounces, int mode,
+                        long long seed, float* image, cudaStream_t stream);
+void ptd_backward_launch(const PtdSceneArgs* args, const float* cam, int H,
+                         int W, int spp, int max_bounces, int mode,
+                         long long seed, const float* grad_image,
+                         float* grad_tex, float* grad_emission,
+                         float* grad_env, cudaStream_t stream);
 }  // extern "C"

tests/test_pathtracer_diff.py CHANGED Viewed

@@ -353,6 +353,208 @@ def test_texture_inverse_recovery():
     assert err.mean().item() < 0.08, err.mean().item()
 @requires_cuda
 @pytest.mark.kernels_ci
 def test_validation():
@@ -366,6 +568,16 @@ def test_validation():
     cam = ptd.Camera(position=(0.5, 0.5, 0.5), look_at=(0.5, 0.5, 1.0))
     with pytest.raises(ValueError):
         ptd.render(scene, cam, 8, 8, spp=1, max_bounces=17)
     img = ptd.render(scene, cam, 8, 8, spp=2, max_bounces=2)
     assert img.shape == (8, 8, 3) and img.dtype == torch.float32
     assert torch.isfinite(img).all()

     assert err.mean().item() < 0.08, err.mean().item()
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_estimators_agree_three_ways():
+    """BRDF-only, NEE-only, and MIS target the same integral at the same
+    bounce budget; their image means must agree within Monte Carlo
+    tolerance."""
+    scene = cornell(with_box=True)
+    cam = cornell_camera()
+    m_mis = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
+                       estimator="mis", seed=1).mean().item()
+    m_nee = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
+                       estimator="nee", seed=2).mean().item()
+    m_brdf = ptd.render(scene, cam, 64, 64, spp=1024, max_bounces=4,
+                        estimator="brdf", seed=3).mean().item()
+    assert abs(m_mis - m_nee) / m_nee < 0.03, (m_mis, m_nee)
+    assert abs(m_mis - m_brdf) / m_brdf < 0.05, (m_mis, m_brdf)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_conductor_under_env_estimators_agree():
+    """A GGX conductor under a constant environment: MIS and BRDF-only
+    means must agree, which jointly validates the VNDF weight, the GGX pdf,
+    and the environment pdf."""
+    verts, faces, mat = [], [], []
+    def quad_w(a, b, c, d, m):
+        base = len(verts)
+        verts.extend([a, b, c, d])
+        faces.extend(quad(base, base + 1, base + 2, base + 3))
+        mat.extend([m, m])
+    quad_w((-4, 0, -4), (4, 0, -4), (4, 0, 4), (-4, 0, 4), 0)   # floor
+    add_box(verts, faces, mat, (-1, 0, -1), (1, 2, 1), 1)       # metal box
+    env = torch.full((8, 16, 3), 1.0)
+    scene = ptd.Scene(verts, faces, mat,
+                      albedo=[[0.5, 0.5, 0.5], [0.9, 0.6, 0.3]],
+                      emission=[[0, 0, 0], [0, 0, 0]],
+                      material_types=[ptd.DIFFUSE, ptd.CONDUCTOR],
+                      roughness=[0.3, 0.25], env=env)
+    cam = ptd.Camera(position=(5, 3, 5), look_at=(0, 1, 0), vfov_deg=45.0)
+    m_mis = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4,
+                       estimator="mis", seed=1).mean().item()
+    m_brdf = ptd.render(scene, cam, 64, 64, spp=512, max_bounces=4,
+                        estimator="brdf", seed=2).mean().item()
+    assert abs(m_mis - m_brdf) / m_brdf < 0.04, (m_mis, m_brdf)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_fresnel_slab_analytic():
+    """A smooth glass slab in front of a uniform emitter, viewed at normal
+    incidence: the transmitted fraction sums the internal-bounce series to
+    (1-R)/(1+R) with R = ((n-1)/(n+1))^2."""
+    verts, faces, mat = [], [], []
+    def quad_w(a, b, c, d, m):
+        base = len(verts)
+        verts.extend([a, b, c, d])
+        faces.extend(quad(base, base + 1, base + 2, base + 3))
+        mat.extend([m, m])
+    quad_w((-10, -10, 5), (10, -10, 5), (10, 10, 5), (-10, 10, 5), 0)  # emitter
+    add_box(verts, faces, mat, (-10, -10, 2.0), (10, 10, 2.2), 1)      # slab
+    n_ior = 1.5
+    scene = ptd.Scene(verts, faces, mat,
+                      albedo=[[0, 0, 0], [0, 0, 0]],
+                      emission=[[1.0, 1.0, 1.0], [0, 0, 0]],
+                      material_types=[ptd.DIFFUSE, ptd.DIELECTRIC],
+                      ior=[1.5, n_ior])
+    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=10.0)
+    img = ptd.render(scene, cam, 32, 32, spp=256, max_bounces=10,
+                     estimator="brdf", seed=5)
+    R = ((n_ior - 1) / (n_ior + 1)) ** 2
+    expected = (1 - R) / (1 + R)
+    got = img[12:20, 12:20].mean().item()
+    assert abs(got - expected) / expected < 0.01, (got, expected)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_env_furnace_diffuse():
+    """A diffuse plane under a constant environment c: with BRDF sampling
+    every sample returns exactly albedo * c (cosine pdf cancellation plus a
+    guaranteed env hit), and MIS agrees in expectation."""
+    verts, faces, mat = [], [], []
+    base = len(verts)
+    verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
+    faces.extend(quad(base, base + 1, base + 2, base + 3))
+    mat.extend([0, 0])
+    c = 0.8
+    a = 0.6
+    env = torch.full((8, 16, 3), c)
+    scene = ptd.Scene(verts, faces, mat, albedo=[[a] * 3],
+                      emission=[[0, 0, 0]], env=env)
+    cam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.2),
+                     vfov_deg=30.0)
+    img = ptd.render(scene, cam, 32, 32, spp=4, max_bounces=3,
+                     estimator="brdf", seed=0)
+    assert (img - a * c).abs().max().item() < 2e-3, \
+        (img.min().item(), img.max().item(), a * c)
+    m_mis = ptd.render(scene, cam, 32, 32, spp=256, max_bounces=3,
+                       estimator="mis", seed=1).mean().item()
+    assert abs(m_mis - a * c) / (a * c) < 0.02, (m_mis, a * c)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_env_texel_gradients_match_fd():
+    """Environment texels are linear in the estimate and the sampling CDF is
+    detached, so same-seed central differences are exact."""
+    torch.manual_seed(3)
+    env0 = torch.rand(4, 8, 3, device="cuda") * 1.5 + 0.2
+    verts, faces, mat = [], [], []
+    base = len(verts)
+    verts.extend([(-5, 0, -5), (5, 0, -5), (5, 0, 5), (-5, 0, 5)])
+    faces.extend(quad(base, base + 1, base + 2, base + 3))
+    mat.extend([0, 0])
+    cam = ptd.Camera(position=(0, 2.0, 0.01), look_at=(0, 0, 0.5),
+                     vfov_deg=40.0)
+    torch.manual_seed(0)
+    Wr = torch.rand(32, 32, 3, device="cuda")
+    scene = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3],
+                      emission=[[0, 0, 0]], env=env0)
+    def loss_of(e):
+        # swap texel values under the construction-time (detached) CDF, the
+        # frozen-sampling semantics the kernel defines
+        scene.env = e.to("cuda")
+        img = ptd.render(scene, cam, 32, 32, spp=8, max_bounces=3, seed=11)
+        return (img * Wr).sum()
+    e = env0.clone().requires_grad_(True)
+    loss_of(e).backward()
+    worst = 0.0
+    for (y, x, c) in [(0, 0, 0), (1, 4, 1), (1, 7, 2)]:  # upper hemisphere
+        h = 5e-3
+        ep = env0.clone(); ep[y, x, c] += h
+        em = env0.clone(); em[y, x, c] -= h
+        fd = (loss_of(ep) - loss_of(em)).item() / (2 * h)
+        an = e.grad[y, x, c].item()
+        assert fd != 0.0, (y, x, c)
+        worst = max(worst, abs(an - fd) / abs(fd))
+    assert worst < 2e-2, worst
+    # a texel no path can see gets exactly zero from both sides
+    assert e.grad[3, 0, 0].item() == 0.0
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_conductor_albedo_gradients_match_fd():
+    """Conductor F0 enters every term through Schlick Fresnel, which is
+    affine in F0, so replay gradients must match same-seed central
+    differences through GGX sampling and MIS."""
+    cam = cornell_camera()
+    H = W = 48
+    torch.manual_seed(0)
+    Wr = torch.rand(H, W, 3, device="cuda")
+    base_alb = torch.tensor([[0.73, 0.73, 0.73], [0.9, 0.5, 0.2],
+                             [0.12, 0.45, 0.15], [0.78, 0.78, 0.78]],
+                            device="cuda")
+    def loss_of(alb):
+        verts, faces, mat = [], [], []
+        def wall(a, b, c, d, m):
+            b0 = len(verts)
+            verts.extend([a, b, c, d])
+            faces.extend(quad(b0, b0 + 1, b0 + 2, b0 + 3))
+            mat.extend([m, m])
+        X, Y, Z = 5.56, 5.488, 5.592
+        wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)
+        wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)
+        wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)
+        wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)   # conductor
+        wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)
+        wall((2.13, Y - 0.001, 2.27), (3.43, Y - 0.001, 2.27),
+             (3.43, Y - 0.001, 3.32), (2.13, Y - 0.001, 3.32), 3)
+        scene = ptd.Scene(verts, faces, mat, albedo=alb,
+                          emission=[[0, 0, 0], [0, 0, 0], [0, 0, 0],
+                                    [18.4, 15.6, 8.0]],
+                          material_types=[ptd.DIFFUSE, ptd.CONDUCTOR,
+                                          ptd.DIFFUSE, ptd.DIFFUSE],
+                          roughness=[0.3, 0.2, 0.3, 0.3])
+        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=13)
+        return (img * Wr).sum()
+    alb = base_alb.clone().requires_grad_(True)
+    loss_of(alb).backward()
+    worst = 0.0
+    for (i, c) in [(1, 0), (1, 2), (0, 1)]:
+        h = 2e-3
+        ap = base_alb.clone(); ap[i, c] += h
+        am = base_alb.clone(); am[i, c] -= h
+        fd = (loss_of(ap) - loss_of(am)).item() / (2 * h)
+        an = alb.grad[i, c].item()
+        assert fd != 0.0, (i, c)
+        worst = max(worst, abs(an - fd) / abs(fd))
+    assert worst < 2e-2, worst
 @requires_cuda
 @pytest.mark.kernels_ci
 def test_validation():
     cam = ptd.Camera(position=(0.5, 0.5, 0.5), look_at=(0.5, 0.5, 1.0))
     with pytest.raises(ValueError):
         ptd.render(scene, cam, 8, 8, spp=1, max_bounces=17)
+    with pytest.raises(ValueError):
+        ptd.render(scene, cam, 8, 8, spp=1, estimator="mis", nee=True)
+    with pytest.raises(ValueError):
+        ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
+                  emission=[[1.0] * 3], material_types=[7])
+    env = torch.full((4, 8, 3), 1.0)
+    scene_env = ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
+                          emission=[[1.0] * 3], env=env)
+    with pytest.raises(ValueError):
+        ptd.render(scene_env, cam, 8, 8, spp=1, estimator="nee")
     img = ptd.render(scene, cam, 8, 8, spp=2, max_bounces=2)
     assert img.shape == (8, 8, 3) and img.dtype == torch.float32
     assert torch.isfinite(img).all()

torch-ext/pathtracer_diff/__init__.py CHANGED Viewed

@@ -1,28 +1,32 @@
 """pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
-Forward renders a linear-radiance image with a megakernel path tracer
-(diffuse + emissive materials, per-texel albedo textures with bilinear
-filtering, cosine-weighted BRDF sampling, optional next-event estimation,
-BVH); one thread owns one pixel, so the image is bitwise deterministic.
-Backward replays the identical paths with the identical counter-based
-Philox draws and accumulates analytic gradients to the albedo texels and
-the per-material emission, scattering each bounce's gradient into the four
-texels of its bilinear footprint with the sampling weights (the exact
-adjoint of the fetch). Sampling directions never depend on the
-differentiable parameters, so the replay is exact. No visibility or
-silhouette gradients: geometry is not a differentiable input.
     from kernels import get_kernel
     ptd = get_kernel("phanerozoic/pathtracer-diff", version=1, trust_remote_code=True)
-    scene = ptd.Scene(vertices, faces, material_ids, albedo, emission, uvs=uvs)
     cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8))
-    img = ptd.render(scene, cam, 256, 256, spp=64)  # autograd to albedo/emission
-    img.sum().backward()
-`albedo` is either an [M, 3] tensor of per-material constants or a list of
-M per-material textures [Hm, Wm, 3]; constants are treated as 1x1 textures,
-so the two forms share one code path and both receive gradients.
 """
 import math
@@ -32,15 +36,172 @@ from ._ops import ops
 MAX_MATERIALS = 64
 MAX_BOUNCES = 16
-__all__ = ["Scene", "Camera", "render", "ops", "MAX_MATERIALS", "MAX_BOUNCES"]
-def _build_bvh(tris, leaf_size=4):
-    """Median-split BVH over [F, 9] triangles (pure torch, CPU). Returns
-    (nodes_f [N, 6], nodes_i [N, 3], order [F]) with leaves as contiguous
-    ranges of the reordered face list. nodes_i rows are (left, right, 0) or
-    (start, count, 1)."""
     tris = tris.to(torch.float32)
     F = tris.shape[0]
     v = tris.reshape(F, 3, 3)
@@ -50,35 +211,103 @@ def _build_bvh(tris, leaf_size=4):
     nodes_f, nodes_i, order = [], [], []
-    def rec(idx):
         nid = len(nodes_f)
         nodes_f.append(None)
         nodes_i.append(None)
         blo = lo[idx].amin(dim=0)
         bhi = hi[idx].amax(dim=0)
         pad = 1e-5 * (1.0 + blo.abs() + bhi.abs())
         nodes_f[nid] = torch.cat([blo - pad, bhi + pad])
-        if idx.numel() <= leaf_size:
             start = len(order)
             order.extend(idx.tolist())
-            nodes_i[nid] = (start, idx.numel(), 1)
-            return nid
-        ax = int((bhi - blo).argmax())
-        srt = idx[torch.argsort(cen[idx, ax], stable=True)]
-        mid = idx.numel() // 2
-        left = rec(srt[:mid])
-        right = rec(srt[mid:])
-        nodes_i[nid] = (left, right, 0)
-        return nid
-    import sys
-    old = sys.getrecursionlimit()
-    sys.setrecursionlimit(max(old, 64 + 2 * int(math.log2(max(F, 2))) * 64))
-    try:
-        rec(torch.arange(F))
-    finally:
-        sys.setrecursionlimit(old)
     return (torch.stack(nodes_f),
             torch.tensor(nodes_i, dtype=torch.int32),
             torch.tensor(order, dtype=torch.int64))
@@ -118,18 +347,22 @@ def _as_texture(t, device):
 class Scene:
-    """Triangle-mesh scene with per-texel albedo and per-material emission.
-    vertices [V, 3] float, faces [F, 3] int, material_ids [F] int.
-    albedo: [M, 3] tensor of per-material constants (treated as 1x1
-    textures) or a list of M textures [Hm, Wm, 3]; either form may require
-    grad. emission [M, 3] float32 (may require grad). uvs: None (all zeros),
-    [V, 2] per-vertex, or [F, 3, 2] per corner; textures wrap (repeat).
-    Texture shapes are fixed at construction; faces are reordered by the
-    BVH build, and emissive faces become the area-sampled light list."""
     def __init__(self, vertices, faces, material_ids, albedo, emission,
-                 uvs=None, device="cuda"):
         vertices = torch.as_tensor(vertices, dtype=torch.float32)
         faces = torch.as_tensor(faces, dtype=torch.int64)
         material_ids = torch.as_tensor(material_ids, dtype=torch.int64)
@@ -160,6 +393,20 @@ class Scene:
         if int(material_ids.max()) >= M or int(material_ids.min()) < 0:
             raise ValueError("material_ids out of range")
         shapes = ([(1, 1)] * M if self.albedo is not None else
                   [(int(t.shape[0]), int(t.shape[1]))
                    for t in self.albedo_textures])
@@ -201,6 +448,35 @@ class Scene:
             total = 0.0
             cdf = torch.zeros(0, dtype=torch.float32)
         self.device = device
         self.tris = tris.to(device).contiguous()
         self.mat_ids = mat_ids.to(device).contiguous()
@@ -217,11 +493,8 @@ class Scene:
         return self.tris.shape[0]
     def _flat_albedo(self):
-        """Pack the albedo into the flat [T, 3] texel buffer inside the
-        autograd graph, so texel gradients flow back to the caller's
-        tensors. Texture shapes must match construction."""
         if self.albedo is not None:
-            if self.albedo.shape[0] * 1 != self.tex_hdr.shape[0]:
                 raise ValueError("albedo/material count changed")
             return self.albedo.reshape(-1, 3)
         parts = []
@@ -232,52 +505,78 @@ class Scene:
             parts.append(t.reshape(-1, 3))
         return torch.cat(parts, dim=0)
 class _RenderFn(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, tex_flat, emission, scene, cam_t, H, W, spp, max_bounces,
-                nee, seed):
         image = torch.empty(H, W, 3, device=tex_flat.device,
                             dtype=torch.float32)
         ops.pt_forward(scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f,
                        scene.nodes_i, scene.light_faces, scene.light_cdf,
                        float(scene.total_light_area),
                        tex_flat.detach().contiguous(), scene.tex_hdr,
-                       emission.detach().contiguous(), cam_t, spp,
-                       max_bounces, nee, seed, image)
         ctx.scene = scene
         ctx.cam_t = cam_t
-        ctx.params = (H, W, spp, max_bounces, nee, seed)
-        ctx.save_for_backward(tex_flat, emission)
         return image
     @staticmethod
     def backward(ctx, grad_image):
-        tex_flat, emission = ctx.saved_tensors
         scene, cam_t = ctx.scene, ctx.cam_t
-        H, W, spp, max_bounces, nee, seed = ctx.params
         ga = torch.zeros_like(tex_flat)
         ge = torch.zeros_like(emission)
         ops.pt_backward(scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f,
                         scene.nodes_i, scene.light_faces, scene.light_cdf,
                         float(scene.total_light_area),
                         tex_flat.detach().contiguous(), scene.tex_hdr,
-                        emission.detach().contiguous(), cam_t, spp,
-                        max_bounces, nee, seed, grad_image.contiguous(), ga,
-                        ge)
-        return ga, ge, None, None, None, None, None, None, None, None
-def render(scene, camera, height, width, spp=64, max_bounces=4, nee=True,
-           seed=0):
     """Render a linear-radiance image [H, W, 3] float32, differentiable with
-    respect to the scene's albedo (constants or texels) and emission. A fixed
-    seed makes the estimate a deterministic function of the parameters
-    (correlated samples), which is what inverse-rendering loops want."""
     if not (1 <= max_bounces <= MAX_BOUNCES):
         raise ValueError(f"max_bounces must be in [1, {MAX_BOUNCES}]")
     cam_t = camera.tensor(height, width, device=scene.device)
-    tex_flat = scene._flat_albedo()
-    return _RenderFn.apply(tex_flat, scene.emission, scene, cam_t, height,
-                           width, int(spp), int(max_bounces), bool(nee),
                            int(seed))

 """pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
+Forward renders a linear-radiance image with a megakernel path tracer:
+Lambertian diffuse, GGX conductor (VNDF-sampled, height-correlated Smith),
+and smooth dielectric materials; area lights and an importance-sampled
+equirectangular environment map; multiple importance sampling by default;
+per-texel albedo with bilinear filtering; binned-SAH BVH. One thread owns
+one pixel, so the image is bitwise deterministic.
+Backward is exact path replay: sampling distributions depend only on
+geometry, frozen material parameters (roughness, ior), the environment CDF
+(detached at Scene construction), and the counter-based Philox stream,
+never on the differentiable parameters. Every radiance term is a product
+of per-bounce factors that are affine in their vertex's albedo texels
+(Schlick Fresnel is affine in F0), times a linear emission or environment
+texel, so the replay differentiates in closed form and scatters through
+the exact bilinear-footprint adjoint. Gradients flow to albedo texels,
+per-material emission, and environment texels. No visibility/silhouette
+gradients: geometry, camera, roughness, and ior are frozen.
     from kernels import get_kernel
     ptd = get_kernel("phanerozoic/pathtracer-diff", version=1, trust_remote_code=True)
+    scene = ptd.Scene(vertices, faces, material_ids, albedo, emission,
+                      uvs=uvs, material_types=types, roughness=rough,
+                      ior=ior, env=env_map)
     cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8))
+    img = ptd.render(scene, cam, 512, 512, spp=64)
+    img.sum().backward()   # grads on albedo texels / emission / env texels
 """
 import math
 MAX_MATERIALS = 64
 MAX_BOUNCES = 16
+DIFFUSE, CONDUCTOR, DIELECTRIC = 0, 1, 2
+_MODES = {"brdf": 0, "nee": 1, "mis": 2}
+__all__ = ["Scene", "Camera", "render", "ops", "MAX_MATERIALS", "MAX_BOUNCES",
+           "DIFFUSE", "CONDUCTOR", "DIELECTRIC"]
+def _build_bvh(tris, leaf_size=4, sah_bins=16):
+    """Binned-SAH BVH over [F, 9] triangles, built level-wise with fully
+    vectorized torch ops (segmented reductions + stable-sort partitioning),
+    so million-triangle scenes build in seconds. Returns (nodes_f [N, 6],
+    nodes_i [N, 3], order [F]). nodes_i rows are (left, right, axis << 1)
+    for internal nodes and (start, count, 1) for leaves; internal children
+    are the (lower, upper) coordinate halves along the split axis, so
+    traversal can visit the near child first."""
+    dev = "cuda" if torch.cuda.is_available() else "cpu"
+    tris = tris.to(torch.float32).to(dev)
+    F = tris.shape[0]
+    v = tris.reshape(F, 3, 3)
+    tlo = v.amin(dim=1)
+    thi = v.amax(dim=1)
+    cen = (tlo + thi) * 0.5
+    perm = torch.arange(F, device=dev)
+    # active segments: contiguous [start, end) ranges of perm, one per node
+    starts = torch.zeros(1, dtype=torch.long, device=dev)
+    ends = torch.full((1,), F, dtype=torch.long, device=dev)
+    ids = torch.zeros(1, dtype=torch.long, device=dev)
+    INF = float("inf")
+    nf_parts, ni_parts = [], []  # per-node rows appended in id order
+    n_nodes = 1
+    def half_area(lo, hi):
+        d = (hi - lo).clamp(min=0)
+        return d[..., 0] * d[..., 1] + d[..., 1] * d[..., 2] + \
+            d[..., 2] * d[..., 0]
+    while ids.numel() > 0:
+        A = ids.numel()
+        seg_len = ends - starts
+        total = int(seg_len.sum())
+        # active positions and per-position segment index
+        csum = torch.cumsum(seg_len, 0) - seg_len
+        pos = starts.repeat_interleave(seg_len) + \
+            (torch.arange(total, device=dev) -
+             csum.repeat_interleave(seg_len))
+        s = torch.arange(A, device=dev).repeat_interleave(seg_len)
+        tri = perm[pos]
+        # per-segment node bounds and centroid bounds
+        nlo = torch.full((A, 3), INF, device=dev)
+        nhi = torch.full((A, 3), -INF, device=dev)
+        clo = torch.full((A, 3), INF, device=dev)
+        chi = torch.full((A, 3), -INF, device=dev)
+        s3 = s.unsqueeze(1).expand(-1, 3)
+        nlo.scatter_reduce_(0, s3, tlo[tri], reduce="amin", include_self=True)
+        nhi.scatter_reduce_(0, s3, thi[tri], reduce="amax", include_self=True)
+        clo.scatter_reduce_(0, s3, cen[tri], reduce="amin", include_self=True)
+        chi.scatter_reduce_(0, s3, cen[tri], reduce="amax", include_self=True)
+        pad = 1e-5 * (1.0 + nlo.abs() + nhi.abs())
+        bounds = torch.cat([nlo - pad, nhi + pad], dim=1)  # [A, 6]
+        leaf_mask = seg_len <= leaf_size
+        split_mask = ~leaf_mask
+        # choose axis (largest centroid extent) and bin the centroids
+        ext = (chi - clo).clamp(min=0)
+        axis = ext.argmax(dim=1)                       # [A]
+        ax_pos = axis[s]
+        coord = cen[tri].gather(1, ax_pos.unsqueeze(1)).squeeze(1)
+        seg_lo = clo.gather(1, axis.unsqueeze(1)).squeeze(1)[s]
+        seg_ext = ext.gather(1, axis.unsqueeze(1)).squeeze(1)[s]
+        t = (coord - seg_lo) / seg_ext.clamp(min=1e-30)
+        bin_i = (t * sah_bins).long().clamp(0, sah_bins - 1)
+        degen = seg_ext <= 1e-12  # all centroids equal on the axis
+        # per-(segment, bin) counts and bounds
+        key = s * sah_bins + bin_i
+        counts = torch.bincount(key, minlength=A * sah_bins) \
+            .reshape(A, sah_bins)
+        blo = torch.full((A * sah_bins, 3), INF, device=dev)
+        bhi = torch.full((A * sah_bins, 3), -INF, device=dev)
+        k3 = key.unsqueeze(1).expand(-1, 3)
+        blo.scatter_reduce_(0, k3, tlo[tri], reduce="amin", include_self=True)
+        bhi.scatter_reduce_(0, k3, thi[tri], reduce="amax", include_self=True)
+        blo = blo.reshape(A, sah_bins, 3)
+        bhi = bhi.reshape(A, sah_bins, 3)
+        # prefix (left) and suffix (right) running bounds and counts
+        plo = torch.cummin(blo, dim=1).values
+        phi = torch.cummax(bhi, dim=1).values
+        pn = torch.cumsum(counts, dim=1)
+        slo = torch.flip(torch.cummin(torch.flip(blo, [1]), dim=1).values, [1])
+        shi = torch.flip(torch.cummax(torch.flip(bhi, [1]), dim=1).values, [1])
+        sn = torch.flip(torch.cumsum(torch.flip(counts, [1]), dim=1), [1])
+        # SAH cost of splitting after bin b (b = 0..bins-2)
+        nl = pn[:, :-1].float()
+        nr = sn[:, 1:].float()
+        cost = half_area(plo[:, :-1], phi[:, :-1]) * nl + \
+            half_area(slo[:, 1:], shi[:, 1:]) * nr
+        cost = torch.where((nl > 0) & (nr > 0), cost,
+                           torch.full_like(cost, INF))
+        best_cost, best_bin = cost.min(dim=1)
+        sah_ok = torch.isfinite(best_cost) & ~degen[csum]  # per segment
+        # side: left if bin <= best_bin (SAH) else lower positional half
+        # (median fallback for degenerate segments)
+        side_left = bin_i <= best_bin[s]
+        within = pos - starts[s]
+        med_left = within < (seg_len[s] // 2)
+        use_sah = sah_ok[s]
+        side_left = torch.where(use_sah, side_left, med_left)
+        side_left = side_left & split_mask[s]
+        # per-segment left count; guard SAH splits that put everything on
+        # one side (possible only via numerics) with the median fallback
+        nl_seg = torch.zeros(A, dtype=torch.long, device=dev)
+        nl_seg.scatter_add_(0, s, side_left.long())
+        bad = split_mask & ((nl_seg == 0) | (nl_seg == seg_len))
+        if bool(bad.any()):
+            fix = bad[s]
+            side_left = torch.where(fix, med_left & split_mask[s], side_left)
+            nl_seg = torch.zeros(A, dtype=torch.long, device=dev)
+            nl_seg.scatter_add_(0, s, side_left.long())
+        # partition each segment in place: stable sort by (segment, right?)
+        k2 = s * 2 + (~side_left).long()
+        order2 = torch.argsort(k2, stable=True)
+        perm[pos] = tri[order2]
+        # allocate children for splitting segments, emit node rows
+        n_split = int(split_mask.sum())
+        child_rank = torch.cumsum(split_mask.long(), 0) - 1
+        left_ids = n_nodes + 2 * child_rank
+        right_ids = left_ids + 1
+        ni = torch.empty(A, 3, dtype=torch.int64, device=dev)
+        ni[:, 0] = torch.where(split_mask, left_ids, starts)
+        ni[:, 1] = torch.where(split_mask, right_ids, seg_len)
+        ni[:, 2] = torch.where(split_mask, axis * 2,
+                               torch.ones_like(axis))
+        nf_parts.append((ids, bounds))
+        ni_parts.append((ids, ni))
+        # next level
+        mid = starts + nl_seg
+        new_starts = torch.cat([starts[split_mask], mid[split_mask]])
+        new_ends = torch.cat([mid[split_mask], ends[split_mask]])
+        new_ids = torch.cat([left_ids[split_mask], right_ids[split_mask]])
+        n_nodes += 2 * n_split
+        starts, ends, ids = new_starts, new_ends, new_ids
+    nodes_f = torch.empty(n_nodes, 6, device=dev)
+    nodes_i = torch.empty(n_nodes, 3, dtype=torch.int64, device=dev)
+    for idv, rows in nf_parts:
+        nodes_f[idv] = rows
+    for idv, rows in ni_parts:
+        nodes_i[idv] = rows
+    return (nodes_f.cpu(), nodes_i.to(torch.int32).cpu(), perm.cpu())
+def _build_bvh_reference(tris, leaf_size=4, force_split=8, sah_bins=16):
+    """Recursive reference builder (kept for cross-checking the vectorized
+    build in development)."""
     tris = tris.to(torch.float32)
     F = tris.shape[0]
     v = tris.reshape(F, 3, 3)
     nodes_f, nodes_i, order = [], [], []
+    def half_area(blo, bhi):
+        d = (bhi - blo).clamp(min=0)
+        return d[0] * d[1] + d[1] * d[2] + d[2] * d[0]
+    stack = []
+    def alloc(idx):
         nid = len(nodes_f)
         nodes_f.append(None)
         nodes_i.append(None)
+        stack.append((nid, idx))
+        return nid
+    root = alloc(torch.arange(F))
+    while stack:
+        nid, idx = stack.pop()
+        n = idx.numel()
         blo = lo[idx].amin(dim=0)
         bhi = hi[idx].amax(dim=0)
         pad = 1e-5 * (1.0 + blo.abs() + bhi.abs())
         nodes_f[nid] = torch.cat([blo - pad, bhi + pad])
+        if n <= leaf_size:
             start = len(order)
             order.extend(idx.tolist())
+            nodes_i[nid] = (start, n, 1)
+            continue
+        cb_lo = cen[idx].amin(dim=0)
+        cb_hi = cen[idx].amax(dim=0)
+        ext = cb_hi - cb_lo
+        axis = int(ext.argmax())
+        split_pos = None
+        if float(ext[axis]) > 1e-12:
+            c = cen[idx, axis]
+            edges = torch.linspace(float(cb_lo[axis]), float(cb_hi[axis]),
+                                   sah_bins + 1)
+            b = torch.bucketize(c, edges[1:-1])
+            counts = torch.bincount(b, minlength=sah_bins)
+            # per-bin bounds over the node's triangles
+            binlo = torch.full((sah_bins, 3), float("inf"))
+            binhi = torch.full((sah_bins, 3), float("-inf"))
+            binlo.scatter_reduce_(0, b.unsqueeze(1).expand(-1, 3), lo[idx],
+                                  reduce="amin", include_self=True)
+            binhi.scatter_reduce_(0, b.unsqueeze(1).expand(-1, 3), hi[idx],
+                                  reduce="amax", include_self=True)
+            best_cost, best_bin = None, None
+            nl = 0
+            llo = torch.full((3,), float("inf"))
+            lhi = torch.full((3,), float("-inf"))
+            pre = []
+            for i in range(sah_bins - 1):
+                if counts[i] > 0:
+                    llo = torch.minimum(llo, binlo[i])
+                    lhi = torch.maximum(lhi, binhi[i])
+                nl += int(counts[i])
+                pre.append((nl, llo.clone(), lhi.clone()))
+            nr = 0
+            rlo = torch.full((3,), float("inf"))
+            rhi = torch.full((3,), float("-inf"))
+            for i in range(sah_bins - 1, 0, -1):
+                if counts[i] > 0:
+                    rlo = torch.minimum(rlo, binlo[i])
+                    rhi = torch.maximum(rhi, binhi[i])
+                nr += int(counts[i])
+                nl_i, llo_i, lhi_i = pre[i - 1]
+                if nl_i == 0 or nr == 0:
+                    continue
+                cost = (half_area(llo_i, lhi_i) * nl_i +
+                        half_area(rlo, rhi) * nr)
+                if best_cost is None or float(cost) < best_cost:
+                    best_cost = float(cost)
+                    best_bin = i - 1
+            if best_bin is not None:
+                parent_area = half_area(blo, bhi)
+                leaf_cost = float(n) * float(parent_area)
+                split_cost = 0.125 * float(parent_area) + best_cost
+                if split_cost < leaf_cost or n > force_split:
+                    mask = b <= best_bin
+                    left_idx = idx[mask]
+                    right_idx = idx[~mask]
+                    if left_idx.numel() > 0 and right_idx.numel() > 0:
+                        split_pos = (left_idx, right_idx)
+        if split_pos is None:
+            if n > force_split and float(ext[axis]) > 1e-12:
+                srt = idx[torch.argsort(cen[idx, axis], stable=True)]
+                mid = n // 2
+                split_pos = (srt[:mid], srt[mid:])
+            else:
+                start = len(order)
+                order.extend(idx.tolist())
+                nodes_i[nid] = (start, n, 1)
+                continue
+        left = alloc(split_pos[0])
+        right = alloc(split_pos[1])
+        nodes_i[nid] = (left, right, axis << 1)
+    assert root == 0
     return (torch.stack(nodes_f),
             torch.tensor(nodes_i, dtype=torch.int32),
             torch.tensor(order, dtype=torch.int64))
 class Scene:
+    """Triangle-mesh scene.
+    vertices [V, 3], faces [F, 3], material_ids [F]. albedo: [M, 3]
+    constants or a list of M textures [Hm, Wm, 3] (constants are 1x1
+    textures; both forms receive gradients). emission [M, 3] (may require
+    grad). uvs: None, [V, 2], or [F, 3, 2]; textures wrap. material_types
+    [M] of DIFFUSE|CONDUCTOR|DIELECTRIC (default all DIFFUSE); roughness
+    [M] GGX alpha for conductors (default 0.3, clamped >= 0.01); ior [M]
+    for dielectrics (default 1.5). env: optional [Eh, Ew, 3] equirect
+    radiance map (may require grad); its importance-sampling CDF is built
+    detached at construction. Texture and env shapes are fixed at
+    construction."""
     def __init__(self, vertices, faces, material_ids, albedo, emission,
+                 uvs=None, material_types=None, roughness=None, ior=None,
+                 env=None, device="cuda"):
         vertices = torch.as_tensor(vertices, dtype=torch.float32)
         faces = torch.as_tensor(faces, dtype=torch.int64)
         material_ids = torch.as_tensor(material_ids, dtype=torch.int64)
         if int(material_ids.max()) >= M or int(material_ids.min()) < 0:
             raise ValueError("material_ids out of range")
+        def mvec(x, default, dtype):
+            if x is None:
+                return torch.full((M,), default, dtype=dtype, device=device)
+            t = torch.as_tensor(x, dtype=dtype).reshape(-1).to(device)
+            if t.numel() != M:
+                raise ValueError("per-material array must have M entries")
+            return t
+        self.mat_type = mvec(material_types, DIFFUSE, torch.int32).contiguous()
+        if int(self.mat_type.max()) > 2 or int(self.mat_type.min()) < 0:
+            raise ValueError("material_types must be 0|1|2")
+        self.mat_rough = mvec(roughness, 0.3, torch.float32).contiguous()
+        self.mat_ior = mvec(ior, 1.5, torch.float32).contiguous()
         shapes = ([(1, 1)] * M if self.albedo is not None else
                   [(int(t.shape[0]), int(t.shape[1]))
                    for t in self.albedo_textures])
             total = 0.0
             cdf = torch.zeros(0, dtype=torch.float32)
+        # environment map + detached sampling tables
+        if env is not None:
+            env = torch.as_tensor(env, dtype=torch.float32)
+            if env.dim() != 3 or env.shape[2] != 3 or env.shape[0] < 2 or \
+                    env.shape[1] < 2:
+                raise ValueError("env must be [Eh, Ew, 3]")
+            self.env = env.to(device)
+            eh, ew = int(env.shape[0]), int(env.shape[1])
+            lum = self.env.detach().mean(dim=2).cpu().double() + 1e-8
+            sint = torch.sin((torch.arange(eh, dtype=torch.float64) + 0.5)
+                             * math.pi / eh).clamp(min=1e-4)
+            roww = (lum.sum(dim=1) * sint)
+            row_p = roww / roww.sum()
+            cdf_m = row_p.cumsum(0).float()
+            col_p = lum / lum.sum(dim=1, keepdim=True)
+            cdf_c = col_p.cumsum(1).float()
+            pdf_img = (row_p.unsqueeze(1) * col_p).float()  # sums to 1
+            self.env_w, self.env_h = ew, eh
+            self.env_cdf_m = cdf_m.to(device).contiguous()
+            self.env_cdf_c = cdf_c.reshape(-1).to(device).contiguous()
+            self.env_pdf = pdf_img.reshape(-1).to(device).contiguous()
+        else:
+            self.env = None
+            self.env_w = self.env_h = 0
+            z = torch.zeros(0, dtype=torch.float32, device=device)
+            self.env_cdf_m = z
+            self.env_cdf_c = z
+            self.env_pdf = z
         self.device = device
         self.tris = tris.to(device).contiguous()
         self.mat_ids = mat_ids.to(device).contiguous()
         return self.tris.shape[0]
     def _flat_albedo(self):
         if self.albedo is not None:
+            if self.albedo.shape[0] != self.tex_hdr.shape[0]:
                 raise ValueError("albedo/material count changed")
             return self.albedo.reshape(-1, 3)
         parts = []
             parts.append(t.reshape(-1, 3))
         return torch.cat(parts, dim=0)
+    def _flat_env(self):
+        if self.env is None:
+            return torch.zeros(0, 3, dtype=torch.float32, device=self.device)
+        if (int(self.env.shape[0]), int(self.env.shape[1])) != \
+                (self.env_h, self.env_w):
+            raise ValueError("env shape is fixed at construction")
+        return self.env.reshape(-1, 3)
 class _RenderFn(torch.autograd.Function):
     @staticmethod
+    def forward(ctx, tex_flat, emission, env_flat, scene, cam_t, H, W, spp,
+                max_bounces, mode, seed):
         image = torch.empty(H, W, 3, device=tex_flat.device,
                             dtype=torch.float32)
         ops.pt_forward(scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f,
                        scene.nodes_i, scene.light_faces, scene.light_cdf,
                        float(scene.total_light_area),
                        tex_flat.detach().contiguous(), scene.tex_hdr,
+                       emission.detach().contiguous(), scene.mat_type,
+                       scene.mat_rough, scene.mat_ior,
+                       env_flat.detach().contiguous(), scene.env_cdf_m,
+                       scene.env_cdf_c, scene.env_pdf, scene.env_w,
+                       scene.env_h, cam_t, spp, max_bounces, mode, seed,
+                       image)
         ctx.scene = scene
         ctx.cam_t = cam_t
+        ctx.params = (H, W, spp, max_bounces, mode, seed)
+        ctx.save_for_backward(tex_flat, emission, env_flat)
         return image
     @staticmethod
     def backward(ctx, grad_image):
+        tex_flat, emission, env_flat = ctx.saved_tensors
         scene, cam_t = ctx.scene, ctx.cam_t
+        H, W, spp, max_bounces, mode, seed = ctx.params
         ga = torch.zeros_like(tex_flat)
         ge = torch.zeros_like(emission)
+        genv = torch.zeros_like(env_flat)
         ops.pt_backward(scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f,
                         scene.nodes_i, scene.light_faces, scene.light_cdf,
                         float(scene.total_light_area),
                         tex_flat.detach().contiguous(), scene.tex_hdr,
+                        emission.detach().contiguous(), scene.mat_type,
+                        scene.mat_rough, scene.mat_ior,
+                        env_flat.detach().contiguous(), scene.env_cdf_m,
+                        scene.env_cdf_c, scene.env_pdf, scene.env_w,
+                        scene.env_h, cam_t, spp, max_bounces, mode, seed,
+                        grad_image.contiguous(), ga, ge, genv)
+        return (ga, ge, genv, None, None, None, None, None, None, None, None)
+def render(scene, camera, height, width, spp=64, max_bounces=4,
+           estimator=None, nee=None, seed=0):
     """Render a linear-radiance image [H, W, 3] float32, differentiable with
+    respect to the scene's albedo texels, emission, and environment texels.
+    estimator: "mis" (default), "nee", or "brdf"; the legacy `nee` bool maps
+    True->"nee", False->"brdf". A fixed seed renders the same paths every
+    call, so the Monte Carlo objective is a deterministic function of the
+    parameters."""
     if not (1 <= max_bounces <= MAX_BOUNCES):
         raise ValueError(f"max_bounces must be in [1, {MAX_BOUNCES}]")
+    if estimator is not None and nee is not None:
+        raise ValueError("pass estimator or nee, not both")
+    if estimator is None:
+        estimator = "mis" if nee is None else ("nee" if nee else "brdf")
+    if estimator not in _MODES:
+        raise ValueError('estimator must be "mis", "nee", or "brdf"')
+    if estimator == "nee" and scene.env is not None:
+        raise ValueError('environment maps require the "mis" or "brdf" estimator')
     cam_t = camera.tensor(height, width, device=scene.device)
+    return _RenderFn.apply(scene._flat_albedo(), scene.emission,
+                           scene._flat_env(), scene, cam_t, height, width,
+                           int(spp), int(max_bounces), _MODES[estimator],
                            int(seed))

torch-ext/torch_binding.cpp CHANGED Viewed

@@ -10,13 +10,28 @@
 namespace {
-void check_scene(const torch::Tensor& tris, const torch::Tensor& mat_ids,
-                 const torch::Tensor& uvs, const torch::Tensor& nodes_f,
-                 const torch::Tensor& nodes_i,
-                 const torch::Tensor& light_faces,
-                 const torch::Tensor& light_cdf, const torch::Tensor& tex,
-                 const torch::Tensor& tex_hdr, const torch::Tensor& emission,
-                 const torch::Tensor& cam, int64_t spp, int64_t max_bounces) {
   TORCH_CHECK(tris.is_cuda() && tris.is_contiguous() &&
                   tris.dtype() == torch::kFloat32 && tris.dim() == 2 &&
                   tris.size(1) == 9,
@@ -37,13 +52,9 @@ void check_scene(const torch::Tensor& tris, const torch::Tensor& mat_ids,
                   nodes_i.dtype() == torch::kInt32 && nodes_i.dim() == 2 &&
                   nodes_i.size(1) == 3 && nodes_i.size(0) == nodes_f.size(0),
               "nodes_i must be contiguous CUDA i32 [N, 3]");
-  TORCH_CHECK(light_faces.is_cuda() && light_faces.is_contiguous() &&
-                  light_faces.dtype() == torch::kInt32,
-              "light_faces must be contiguous CUDA i32 [L]");
-  TORCH_CHECK(light_cdf.is_cuda() && light_cdf.is_contiguous() &&
-                  light_cdf.dtype() == torch::kFloat32 &&
                   light_cdf.numel() == light_faces.numel(),
-              "light_cdf must be contiguous CUDA f32 [L]");
   TORCH_CHECK(tex.is_cuda() && tex.is_contiguous() &&
                   tex.dtype() == torch::kFloat32 && tex.dim() == 2 &&
                   tex.size(1) == 3,
@@ -52,16 +63,60 @@ void check_scene(const torch::Tensor& tris, const torch::Tensor& mat_ids,
                   tex_hdr.dtype() == torch::kInt32 && tex_hdr.dim() == 2 &&
                   tex_hdr.size(1) == 3,
               "tex_hdr must be contiguous CUDA i32 [M, 3]");
   TORCH_CHECK(emission.is_cuda() && emission.is_contiguous() &&
                   emission.dtype() == torch::kFloat32 &&
-                  emission.dim() == 2 && emission.size(1) == 3 &&
-                  emission.size(0) == tex_hdr.size(0),
               "emission must be contiguous CUDA f32 [M, 3]");
-  TORCH_CHECK(tex_hdr.size(0) <= 64, "at most 64 materials");
-  TORCH_CHECK(cam.numel() == 12, "cam must have 12 elements");
   TORCH_CHECK(spp >= 1, "spp must be >= 1");
   TORCH_CHECK(max_bounces >= 1 && max_bounces <= 16,
               "max_bounces must be in [1, 16]");
 }
 }  // namespace
@@ -71,30 +126,28 @@ void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
-                torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                int64_t seed, torch::Tensor image) {
-  check_scene(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces, light_cdf,
-              tex, tex_hdr, emission, cam, spp, max_bounces);
   TORCH_CHECK(image.is_cuda() && image.is_contiguous() &&
                   image.dtype() == torch::kFloat32 && image.dim() == 3 &&
                   image.size(2) == 3,
               "image must be contiguous CUDA f32 [H, W, 3]");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  int L = (int)light_faces.numel();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
-  ptd_forward_launch(
-      tris.const_data_ptr<float>(), mat_ids.const_data_ptr<int>(),
-      uvs.const_data_ptr<float>(), (int)tris.size(0),
-      nodes_f.const_data_ptr<float>(), nodes_i.const_data_ptr<int>(),
-      (int)nodes_f.size(0), L ? light_faces.const_data_ptr<int>() : nullptr,
-      L ? light_cdf.const_data_ptr<float>() : nullptr, L,
-      (float)total_light_area, tex.const_data_ptr<float>(),
-      tex_hdr.const_data_ptr<int>(), (int)tex.size(0),
-      emission.const_data_ptr<float>(),
-      (int)tex_hdr.size(0), cam_h.const_data_ptr<float>(),
-      (int)image.size(0), (int)image.size(1), (int)spp, (int)max_bounces,
-      nee ? 1 : 0, (long long)seed, image.data_ptr<float>(), stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
@@ -103,11 +156,19 @@ void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
-                 torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                 int64_t seed, torch::Tensor grad_image,
-                 torch::Tensor grad_tex, torch::Tensor grad_emission) {
-  check_scene(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces, light_cdf,
-              tex, tex_hdr, emission, cam, spp, max_bounces);
   TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
                   grad_image.dtype() == torch::kFloat32 &&
                   grad_image.dim() == 3 && grad_image.size(2) == 3,
@@ -118,24 +179,19 @@ void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
   TORCH_CHECK(grad_emission.is_cuda() && grad_emission.is_contiguous() &&
                   grad_emission.sizes() == emission.sizes(),
               "grad_emission must match emission");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
-  int L = (int)light_faces.numel();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
-  ptd_backward_launch(
-      tris.const_data_ptr<float>(), mat_ids.const_data_ptr<int>(),
-      uvs.const_data_ptr<float>(), (int)tris.size(0),
-      nodes_f.const_data_ptr<float>(), nodes_i.const_data_ptr<int>(),
-      (int)nodes_f.size(0), L ? light_faces.const_data_ptr<int>() : nullptr,
-      L ? light_cdf.const_data_ptr<float>() : nullptr, L,
-      (float)total_light_area, tex.const_data_ptr<float>(),
-      tex_hdr.const_data_ptr<int>(), (int)tex.size(0),
-      emission.const_data_ptr<float>(),
-      (int)tex_hdr.size(0), cam_h.const_data_ptr<float>(),
-      (int)grad_image.size(0), (int)grad_image.size(1), (int)spp,
-      (int)max_bounces, nee ? 1 : 0, (long long)seed,
-      grad_image.const_data_ptr<float>(), grad_tex.data_ptr<float>(),
-      grad_emission.data_ptr<float>(), stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
@@ -144,7 +200,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "pt_forward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
       " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
       " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emission,"
-      " Tensor cam, int spp, int max_bounces, bool nee, int seed,"
       " Tensor! image) -> ()");
   ops.impl("pt_forward", torch::kCUDA, &pt_forward);
@@ -152,8 +210,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "pt_backward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
       " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
       " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emission,"
-      " Tensor cam, int spp, int max_bounces, bool nee, int seed,"
-      " Tensor grad_image, Tensor! grad_tex, Tensor! grad_emission) -> ()");
   ops.impl("pt_backward", torch::kCUDA, &pt_backward);
 }

 namespace {
+const float* fptr(const torch::Tensor& t) {
+  return t.numel() ? t.const_data_ptr<float>() : nullptr;
+}
+const int* iptr(const torch::Tensor& t) {
+  return t.numel() ? t.const_data_ptr<int>() : nullptr;
+}
+PtdSceneArgs pack_args(const torch::Tensor& tris, const torch::Tensor& mat_ids,
+                       const torch::Tensor& uvs, const torch::Tensor& nodes_f,
+                       const torch::Tensor& nodes_i,
+                       const torch::Tensor& light_faces,
+                       const torch::Tensor& light_cdf, double total_light_area,
+                       const torch::Tensor& tex, const torch::Tensor& tex_hdr,
+                       const torch::Tensor& emission,
+                       const torch::Tensor& mat_type,
+                       const torch::Tensor& mat_rough,
+                       const torch::Tensor& mat_ior, const torch::Tensor& env,
+                       const torch::Tensor& env_cdf_m,
+                       const torch::Tensor& env_cdf_c,
+                       const torch::Tensor& env_pdf, int64_t env_w,
+                       int64_t env_h, int64_t spp, int64_t max_bounces,
+                       int64_t mode) {
   TORCH_CHECK(tris.is_cuda() && tris.is_contiguous() &&
                   tris.dtype() == torch::kFloat32 && tris.dim() == 2 &&
                   tris.size(1) == 9,
                   nodes_i.dtype() == torch::kInt32 && nodes_i.dim() == 2 &&
                   nodes_i.size(1) == 3 && nodes_i.size(0) == nodes_f.size(0),
               "nodes_i must be contiguous CUDA i32 [N, 3]");
+  TORCH_CHECK(light_faces.dtype() == torch::kInt32 &&
                   light_cdf.numel() == light_faces.numel(),
+              "light list mismatch");
   TORCH_CHECK(tex.is_cuda() && tex.is_contiguous() &&
                   tex.dtype() == torch::kFloat32 && tex.dim() == 2 &&
                   tex.size(1) == 3,
                   tex_hdr.dtype() == torch::kInt32 && tex_hdr.dim() == 2 &&
                   tex_hdr.size(1) == 3,
               "tex_hdr must be contiguous CUDA i32 [M, 3]");
+  int64_t M = tex_hdr.size(0);
+  TORCH_CHECK(M <= 64, "at most 64 materials");
   TORCH_CHECK(emission.is_cuda() && emission.is_contiguous() &&
                   emission.dtype() == torch::kFloat32 &&
+                  emission.sizes() == torch::IntArrayRef({M, 3}),
               "emission must be contiguous CUDA f32 [M, 3]");
+  TORCH_CHECK(mat_type.dtype() == torch::kInt32 && mat_type.numel() == M,
+              "mat_type must be i32 [M]");
+  TORCH_CHECK(mat_rough.dtype() == torch::kFloat32 && mat_rough.numel() == M,
+              "mat_rough must be f32 [M]");
+  TORCH_CHECK(mat_ior.dtype() == torch::kFloat32 && mat_ior.numel() == M,
+              "mat_ior must be f32 [M]");
+  if (env.numel()) {
+    TORCH_CHECK(env.is_cuda() && env.is_contiguous() &&
+                    env.dtype() == torch::kFloat32 &&
+                    env.numel() == env_w * env_h * 3,
+                "env must be contiguous CUDA f32 [Eh*Ew, 3]");
+    TORCH_CHECK(env_cdf_m.numel() == env_h &&
+                    env_cdf_c.numel() == env_w * env_h &&
+                    env_pdf.numel() == env_w * env_h,
+                "env tables mismatch");
+  }
   TORCH_CHECK(spp >= 1, "spp must be >= 1");
   TORCH_CHECK(max_bounces >= 1 && max_bounces <= 16,
               "max_bounces must be in [1, 16]");
+  TORCH_CHECK(mode >= 0 && mode <= 2, "mode must be 0|1|2");
+  PtdSceneArgs a;
+  a.tris = tris.const_data_ptr<float>();
+  a.mat_ids = mat_ids.const_data_ptr<int>();
+  a.uvs = uvs.const_data_ptr<float>();
+  a.n_faces = (int)tris.size(0);
+  a.nodes_f = nodes_f.const_data_ptr<float>();
+  a.nodes_i = nodes_i.const_data_ptr<int>();
+  a.n_nodes = (int)nodes_f.size(0);
+  a.light_faces = iptr(light_faces);
+  a.light_cdf = fptr(light_cdf);
+  a.n_lights = (int)light_faces.numel();
+  a.total_light_area = (float)total_light_area;
+  a.tex = tex.const_data_ptr<float>();
+  a.tex_hdr = tex_hdr.const_data_ptr<int>();
+  a.n_texels = (int)tex.size(0);
+  a.emission = emission.const_data_ptr<float>();
+  a.mat_type = mat_type.const_data_ptr<int>();
+  a.mat_rough = mat_rough.const_data_ptr<float>();
+  a.mat_ior = mat_ior.const_data_ptr<float>();
+  a.n_mats = (int)M;
+  a.env = fptr(env);
+  a.env_w = (int)env_w;
+  a.env_h = (int)env_h;
+  a.env_cdf_m = fptr(env_cdf_m);
+  a.env_cdf_c = fptr(env_cdf_c);
+  a.env_pdf = fptr(env_pdf);
+  return a;
 }
 }  // namespace
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
+                torch::Tensor mat_type, torch::Tensor mat_rough,
+                torch::Tensor mat_ior, torch::Tensor env,
+                torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                int64_t mode, int64_t seed, torch::Tensor image) {
+  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
+                             light_cdf, total_light_area, tex, tex_hdr,
+                             emission, mat_type, mat_rough, mat_ior, env,
+                             env_cdf_m, env_cdf_c, env_pdf, env_w, env_h, spp,
+                             max_bounces, mode);
   TORCH_CHECK(image.is_cuda() && image.is_contiguous() &&
                   image.dtype() == torch::kFloat32 && image.dim() == 3 &&
                   image.size(2) == 3,
               "image must be contiguous CUDA f32 [H, W, 3]");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
+  ptd_forward_launch(&a, cam_h.const_data_ptr<float>(), (int)image.size(0),
+                     (int)image.size(1), (int)spp, (int)max_bounces,
+                     (int)mode, (long long)seed, image.data_ptr<float>(),
+                     stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
+                 torch::Tensor mat_type, torch::Tensor mat_rough,
+                 torch::Tensor mat_ior, torch::Tensor env,
+                 torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                 torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                 torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                 int64_t mode, int64_t seed, torch::Tensor grad_image,
+                 torch::Tensor grad_tex, torch::Tensor grad_emission,
+                 torch::Tensor grad_env) {
+  PtdSceneArgs a = pack_args(tris, mat_ids, uvs, nodes_f, nodes_i, light_faces,
+                             light_cdf, total_light_area, tex, tex_hdr,
+                             emission, mat_type, mat_rough, mat_ior, env,
+                             env_cdf_m, env_cdf_c, env_pdf, env_w, env_h, spp,
+                             max_bounces, mode);
   TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous() &&
                   grad_image.dtype() == torch::kFloat32 &&
                   grad_image.dim() == 3 && grad_image.size(2) == 3,
   TORCH_CHECK(grad_emission.is_cuda() && grad_emission.is_contiguous() &&
                   grad_emission.sizes() == emission.sizes(),
               "grad_emission must match emission");
+  TORCH_CHECK(grad_env.is_contiguous() && grad_env.sizes() == env.sizes(),
+              "grad_env must match env");
   const at::cuda::CUDAGuard guard(tris.device());
   cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   torch::Tensor cam_h = cam.to(torch::kFloat32).to(torch::kCPU).contiguous();
+  ptd_backward_launch(&a, cam_h.const_data_ptr<float>(),
+                      (int)grad_image.size(0), (int)grad_image.size(1),
+                      (int)spp, (int)max_bounces, (int)mode, (long long)seed,
+                      grad_image.const_data_ptr<float>(),
+                      grad_tex.data_ptr<float>(),
+                      grad_emission.data_ptr<float>(),
+                      grad_env.numel() ? grad_env.data_ptr<float>() : nullptr,
+                      stream);
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
       "pt_forward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
       " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
       " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emission,"
+      " Tensor mat_type, Tensor mat_rough, Tensor mat_ior, Tensor env,"
+      " Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
+      " int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
       " Tensor! image) -> ()");
   ops.impl("pt_forward", torch::kCUDA, &pt_forward);
       "pt_backward(Tensor tris, Tensor mat_ids, Tensor uvs, Tensor nodes_f,"
       " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
       " float total_light_area, Tensor tex, Tensor tex_hdr, Tensor emission,"
+      " Tensor mat_type, Tensor mat_rough, Tensor mat_ior, Tensor env,"
+      " Tensor env_cdf_m, Tensor env_cdf_c, Tensor env_pdf, int env_w,"
+      " int env_h, Tensor cam, int spp, int max_bounces, int mode, int seed,"
+      " Tensor grad_image, Tensor! grad_tex, Tensor! grad_emission,"
+      " Tensor! grad_env) -> ()");
   ops.impl("pt_backward", torch::kCUDA, &pt_backward);
 }

torch-ext/torch_binding.h CHANGED Viewed

@@ -7,14 +7,23 @@ void pt_forward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
-                torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                int64_t seed, torch::Tensor image);
 void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                  torch::Tensor nodes_f, torch::Tensor nodes_i,
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
-                 torch::Tensor cam, int64_t spp, int64_t max_bounces, bool nee,
-                 int64_t seed, torch::Tensor grad_image,
-                 torch::Tensor grad_tex, torch::Tensor grad_emission);

                 torch::Tensor light_faces, torch::Tensor light_cdf,
                 double total_light_area, torch::Tensor tex,
                 torch::Tensor tex_hdr, torch::Tensor emission,
+                torch::Tensor mat_type, torch::Tensor mat_rough,
+                torch::Tensor mat_ior, torch::Tensor env,
+                torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                int64_t mode, int64_t seed, torch::Tensor image);
 void pt_backward(torch::Tensor tris, torch::Tensor mat_ids, torch::Tensor uvs,
                  torch::Tensor nodes_f, torch::Tensor nodes_i,
                  torch::Tensor light_faces, torch::Tensor light_cdf,
                  double total_light_area, torch::Tensor tex,
                  torch::Tensor tex_hdr, torch::Tensor emission,
+                 torch::Tensor mat_type, torch::Tensor mat_rough,
+                 torch::Tensor mat_ior, torch::Tensor env,
+                 torch::Tensor env_cdf_m, torch::Tensor env_cdf_c,
+                 torch::Tensor env_pdf, int64_t env_w, int64_t env_h,
+                 torch::Tensor cam, int64_t spp, int64_t max_bounces,
+                 int64_t mode, int64_t seed, torch::Tensor grad_image,
+                 torch::Tensor grad_tex, torch::Tensor grad_emission,
+                 torch::Tensor grad_env);