pathtracer-diff v1: differentiable path tracer source tree

Browse files

Files changed (14) hide show

CARD.md +130 -0
README.md +130 -0
build.toml +26 -0
dev/report_local.py +180 -0
dev/run_local.py +20 -0
flake.nix +17 -0
load_local.py +41 -0
local_bind.cpp +78 -0
pathtracer_diff_cuda/pathtracer.cu +483 -0
pathtracer_diff_cuda/pathtracer_launch.h +30 -0
tests/test_pathtracer_diff.py +227 -0
torch-ext/pathtracer_diff/__init__.py +212 -0
torch-ext/torch_binding.cpp +149 -0
torch-ext/torch_binding.h +20 -0

CARD.md ADDED Viewed

	@@ -0,0 +1,130 @@

+---
+library_name: kernels
+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, 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 per-material albedo and 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.
+## Usage
+```python
+import torch
+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=albedo,        # [M, 3] float32, may require grad
+                  emission=emission)    # [M, 3] float32, may require grad
+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 scene.albedo / scene.emission
+```
+`version` selects the release branch; `trust_remote_code` is required by
+`kernels` for publishers without the trusted-publisher mark.
+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)` | triangle-mesh scene; the BVH and the emissive-face light list are built 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 `scene.albedo` and `scene.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: cosine-weighted hemisphere sampling makes
+the diffuse throughput multiplier exactly the 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 accumulates into per-material gradient buffers through
+block-shared memory. 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
+Measured on RTX 6000 Ada (sm89) from a source build:
+- **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 a maximum relative error of 3.0e-4 over nine
+  albedo and emission entries in a Cornell scene.
+- **Inverse rendering.** Recovering a wall albedo (0.2, 0.5, 0.7) from a
+  rendered target by Adam through the kernel reaches a maximum absolute
+  error of 0.003 in 80 steps, driving the image loss from 2.6e-4 to 4.1e-8.
+- **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 | 8.7 ms | 1.94 Gpaths/s |
+| backward (replay) | 20.1 ms | 2.32x 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.
+- Differentiable parameters are per-material albedo and emission only.
+  Geometry, camera, and per-vertex quantities are not differentiable inputs,
+  and there are no visibility/silhouette gradients.
+- 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.

README.md ADDED Viewed

	@@ -0,0 +1,130 @@

+---
+library_name: kernels
+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, 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 per-material albedo and 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.
+## Usage
+```python
+import torch
+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=albedo,        # [M, 3] float32, may require grad
+                  emission=emission)    # [M, 3] float32, may require grad
+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 scene.albedo / scene.emission
+```
+`version` selects the release branch; `trust_remote_code` is required by
+`kernels` for publishers without the trusted-publisher mark.
+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)` | triangle-mesh scene; the BVH and the emissive-face light list are built 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 `scene.albedo` and `scene.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: cosine-weighted hemisphere sampling makes
+the diffuse throughput multiplier exactly the 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 accumulates into per-material gradient buffers through
+block-shared memory. 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
+Measured on RTX 6000 Ada (sm89) from a source build:
+- **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 a maximum relative error of 3.0e-4 over nine
+  albedo and emission entries in a Cornell scene.
+- **Inverse rendering.** Recovering a wall albedo (0.2, 0.5, 0.7) from a
+  rendered target by Adam through the kernel reaches a maximum absolute
+  error of 0.003 in 80 steps, driving the image loss from 2.6e-4 to 4.1e-8.
+- **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 | 8.7 ms | 1.94 Gpaths/s |
+| backward (replay) | 20.1 ms | 2.32x 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.
+- Differentiable parameters are per-material albedo and emission only.
+  Geometry, camera, and per-vertex quantities are not differentiable inputs,
+  and there are no visibility/silhouette gradients.
+- 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.

build.toml ADDED Viewed

	@@ -0,0 +1,26 @@

+[general]
+name = "pathtracer-diff"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = [
+  "cuda",
+]
+[general.hub]
+repo-id = "phanerozoic/pathtracer-diff"
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+[kernel.pathtracer-diff]
+backend = "cuda"
+cuda-capabilities = ["8.0", "8.6", "8.9", "9.0", "10.0", "12.0"]
+depends = ["torch"]
+src = [
+  "pathtracer_diff_cuda/pathtracer.cu",
+  "pathtracer_diff_cuda/pathtracer_launch.h",
+]

dev/report_local.py ADDED Viewed

	@@ -0,0 +1,180 @@

+"""Measured numbers for the card, run on the local GPU from source."""
+import os
+import sys
+import time
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+import torch
+import load_local
+ptd = load_local.load()
+print("GPU:", torch.cuda.get_device_name(0))
+def quad(a, b, c, d):
+    return [[a, b, c], [a, c, d]]
+def add_box(verts, faces, mat, lo, hi, mat_id):
+    x0, y0, z0 = lo
+    x1, y1, z1 = hi
+    base = len(verts)
+    verts.extend([(x0, y0, z0), (x1, y0, z0), (x1, y0, z1), (x0, y0, z1),
+                  (x0, y1, z0), (x1, y1, z0), (x1, y1, z1), (x0, y1, z1)])
+    i = lambda k: base + k
+    for q in [quad(i(0), i(1), i(2), i(3)), quad(i(4), i(7), i(6), i(5)),
+              quad(i(0), i(4), i(5), i(1)), quad(i(3), i(2), i(6), i(7)),
+              quad(i(0), i(3), i(7), i(4)), quad(i(1), i(5), i(6), i(2))]:
+        faces.extend(q)
+        mat.extend([mat_id, mat_id])
+def cornell(with_box=True):
+    verts, faces, mat = [], [], []
+    def wall(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])
+    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)
+    if with_box:
+        add_box(verts, faces, mat, (1.3, 0.0, 0.65), (2.4, 1.65, 1.7), 0)
+    albedo = [[0.73, 0.73, 0.73], [0.65, 0.05, 0.05], [0.12, 0.45, 0.15],
+              [0.78, 0.78, 0.78]]
+    emission = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [18.4, 15.6, 8.0]]
+    return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission)
+cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8),
+                 vfov_deg=39.0)
+# ---- furnace exactness
+a, E, B = 0.5, 0.25, 8
+verts, faces, mat = [], [], []
+add_box(verts, faces, mat, (-2, -2, -2), (2, 2, 2), 0)
+fs = ptd.Scene(verts, faces, mat, albedo=[[a] * 3], emission=[[E] * 3])
+fc = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=60.0)
+img = ptd.render(fs, fc, 64, 64, spp=4, max_bounces=B, nee=False, seed=0)
+exp = E * (1 - a ** B) / (1 - a)
+print(f"furnace: analytic {exp:.9f}, max abs err {(img - exp).abs().max().item():.2e}")
+# ---- estimator agreement
+verts, faces, mat = [], [], []
+add_box(verts, faces, mat, (0, 0, 0), (5, 5, 5), 0)
+base = len(verts)
+verts.extend([(1.0, 4.999, 1.0), (4.5, 4.999, 1.0), (4.5, 4.999, 4.5),
+              (1.0, 4.999, 4.5)])
+faces.extend(quad(base, base + 1, base + 2, base + 3))
+mat.extend([1, 1])
+es = ptd.Scene(verts, faces, mat, albedo=[[0.6] * 3, [0.0] * 3],
+               emission=[[0] * 3, [3.0] * 3])
+ec = ptd.Camera(position=(2.5, 2.5, 0.2), look_at=(2.5, 2.5, 5.0), vfov_deg=70.0)
+m_nee = ptd.render(es, ec, 128, 128, spp=128, max_bounces=4, nee=True, seed=1).mean().item()
+m_brdf = ptd.render(es, ec, 128, 128, spp=512, max_bounces=4, nee=False, seed=2).mean().item()
+print(f"estimators: NEE {m_nee:.5f} vs BRDF {m_brdf:.5f}, rel diff {abs(m_nee-m_brdf)/m_brdf:.4%}")
+# ---- FD gradient agreement
+scene = cornell(with_box=False)
+H = W = 48
+torch.manual_seed(0)
+Wr = torch.rand(H, W, 3, device="cuda")
+def loss_of(alb, emi):
+    scene.albedo = alb
+    scene.emission = emi
+    return (ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=7) * Wr).sum()
+alb0 = scene.albedo.clone()
+emi0 = scene.emission.clone()
+alb = alb0.clone().requires_grad_(True)
+emi = emi0.clone().requires_grad_(True)
+loss_of(alb, emi).backward()
+worst = 0.0
+for (i, c) in [(0, 0), (0, 1), (0, 2), (1, 0), (2, 1), (2, 2)]:
+    h = 2e-3
+    ap = alb0.clone(); ap[i, c] += h
+    am = alb0.clone(); am[i, c] -= h
+    fd = (loss_of(ap, emi0) - loss_of(am, emi0)).item() / (2 * h)
+    worst = max(worst, abs(alb.grad[i, c].item() - fd) / abs(fd))
+h = 0.05
+for c in range(3):
+    ep = emi0.clone(); ep[3, c] += h
+    em = emi0.clone(); em[3, c] -= h
+    fd = (loss_of(alb0, ep) - loss_of(alb0, em)).item() / (2 * h)
+    worst = max(worst, abs(emi.grad[3, c].item() - fd) / abs(fd))
+print(f"gradients vs central differences: max rel err {worst:.2e} over 9 entries")
+# ---- inverse rendering
+target_row = torch.tensor([0.2, 0.5, 0.7], device="cuda")
+scene = cornell(with_box=False)
+t_alb = scene.albedo.clone(); t_alb[1] = target_row
+scene.albedo = t_alb
+target = ptd.render(scene, cam, 48, 48, spp=8, max_bounces=3, seed=3).detach()
+alb = t_alb.clone(); alb[1] = torch.tensor([0.5, 0.5, 0.5], device="cuda")
+alb = alb.requires_grad_(True)
+opt = torch.optim.Adam([alb], lr=0.05)
+first = None
+for it in range(80):
+    opt.zero_grad()
+    scene.albedo = alb
+    loss = (ptd.render(scene, cam, 48, 48, spp=8, max_bounces=3, seed=3) - target).square().mean()
+    loss.backward()
+    alb.grad[0] = 0; alb.grad[2] = 0; alb.grad[3] = 0
+    opt.step()
+    with torch.no_grad():
+        alb.clamp_(0.02, 0.98)
+    if first is None:
+        first = loss.item()
+print(f"inverse: target {target_row.tolist()} recovered "
+      f"{[round(x, 3) for x in alb[1].detach().tolist()]}, "
+      f"max abs err {(alb[1].detach() - target_row).abs().max().item():.3f}, "
+      f"loss {first:.2e} -> {loss.item():.2e} in 80 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")

dev/run_local.py ADDED Viewed

	@@ -0,0 +1,20 @@

+"""Build the extension locally and run the test suite against it, stubbing
+kernels.get_kernel so the published-variant tests run natively."""
+import os
+import sys
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+import load_local
+mod = load_local.load(verbose="-v" in sys.argv)
+import kernels
+kernels.get_kernel = lambda *a, **k: mod
+import pytest
+sys.exit(pytest.main([os.path.join(ROOT, "tests"), "-x", "-q", "-p",
+                      "no:cacheprovider"]))

flake.nix ADDED Viewed

	@@ -0,0 +1,17 @@

+{
+  description = "Flake for pathtracer-diff kernel";
+  inputs = {
+    builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    {
+      self,
+      builder,
+    }:
+    builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

load_local.py ADDED Viewed

	@@ -0,0 +1,41 @@

+"""Load pathtracer-diff natively on this machine (Windows/SAURON included).
+The published Hugging Face kernel carries `*-linux` variants only, so on
+Windows `get_kernel` has no matching build. The source is right here, though:
+JIT-build it with `torch.utils.cpp_extension` and expose the identical
+`Scene` / `Camera` / `render` API.
+    import load_local
+    ptd = load_local.load()
+"""
+import os
+import sys
+import types
+_cached = None
+def load(verbose=False):
+    """Return the pathtracer_diff module, JIT-built from the local source."""
+    global _cached
+    if _cached is not None:
+        return _cached
+    from torch.utils.cpp_extension import load as _jit
+    root = os.path.dirname(os.path.abspath(__file__))
+    cuda_dir = os.path.join(root, "pathtracer_diff_cuda")
+    ext = _jit(name="pathtracer_diff_ext",
+               sources=[os.path.join(root, "local_bind.cpp"),
+                        os.path.join(cuda_dir, "pathtracer.cu")],
+               extra_include_paths=[cuda_dir],
+               verbose=verbose)
+    ops_mod = types.ModuleType("pathtracer_diff._ops")
+    ops_mod.ops = type("_Ops", (), {
+        "pt_forward": staticmethod(ext.pt_forward),
+        "pt_backward": staticmethod(ext.pt_backward),
+    })()
+    sys.modules["pathtracer_diff._ops"] = ops_mod
+    sys.path.insert(0, os.path.join(root, "torch-ext"))
+    import pathtracer_diff
+    _cached = pathtracer_diff
+    return pathtracer_diff

local_bind.cpp ADDED Viewed

	@@ -0,0 +1,78 @@

+// Local development binding (torch.utils.cpp_extension JIT; Windows/MSVC).
+// Mirrors the torch-ext binding's argument order exactly.
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "pathtracer_launch.h"
+namespace {
+void pt_forward(torch::Tensor tris, torch::Tensor mat_ids,
+                torch::Tensor nodes_f, torch::Tensor nodes_i,
+                torch::Tensor light_faces, torch::Tensor light_cdf,
+                double total_light_area, torch::Tensor albedo,
+                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(image.is_cuda() && image.is_contiguous() && image.dim() == 3,
+              "image [H, W, 3]");
+  TORCH_CHECK(albedo.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>(),
+      (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, albedo.const_data_ptr<float>(),
+      emission.const_data_ptr<float>(), (int)albedo.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();
+}
+void pt_backward(torch::Tensor tris, torch::Tensor mat_ids,
+                 torch::Tensor nodes_f, torch::Tensor nodes_i,
+                 torch::Tensor light_faces, torch::Tensor light_cdf,
+                 double total_light_area, torch::Tensor albedo,
+                 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_albedo,
+                 torch::Tensor grad_emission) {
+  TORCH_CHECK(grad_image.is_cuda() && grad_image.is_contiguous(), "grad_image");
+  TORCH_CHECK(albedo.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>(),
+      (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, albedo.const_data_ptr<float>(),
+      emission.const_data_ptr<float>(), (int)albedo.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_albedo.data_ptr<float>(), grad_emission.data_ptr<float>(), stream);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+}  // namespace
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("pt_forward", &pt_forward, "pathtracer-diff forward");
+  m.def("pt_backward", &pt_backward, "pathtracer-diff backward");
+}

pathtracer_diff_cuda/pathtracer.cu ADDED Viewed

	@@ -0,0 +1,483 @@

+// pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel.
+//
+// Forward: a megakernel unidirectional path tracer (diffuse + emissive
+// materials, cosine-weighted BRDF sampling, optional next-event estimation
+// against an area-sampled light list, 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 (per-material albedo and emission), so the backward kernel
+// re-traces the identical paths with the identical RNG draws and
+// accumulates dLoss/dAlbedo and dLoss/dEmission analytically, storing no
+// path state between the passes (path-replay backpropagation, restricted
+// to interior terms: no visibility/silhouette gradients).
+//
+// Radiance estimator with NEE on: emission is gathered at the camera hit
+// only; every later vertex estimates direct light by area sampling, so no
+// contribution is double-counted. 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>
+#include <cstdint>
+#include "pathtracer_launch.h"
+namespace {
+constexpr int kThreads = 128;
+constexpr int kMaxBounces = 16;
+constexpr int kMaxMats = 64;
+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) {
+  return make_float3(x, y, z);
+}
+__device__ __forceinline__ float3 operator+(float3 a, float3 b) {
+  return f3(a.x + b.x, a.y + b.y, a.z + b.z);
+}
+__device__ __forceinline__ float3 operator-(float3 a, float3 b) {
+  return f3(a.x - b.x, a.y - b.y, a.z - b.z);
+}
+__device__ __forceinline__ float3 operator*(float3 a, float s) {
+  return f3(a.x * s, a.y * s, a.z * s);
+}
+__device__ __forceinline__ float3 operator*(float3 a, float3 b) {
+  return f3(a.x * b.x, a.y * b.y, a.z * b.z);
+}
+__device__ __forceinline__ float dot(float3 a, float3 b) {
+  return a.x * b.x + a.y * b.y + a.z * b.z;
+}
+__device__ __forceinline__ float3 cross(float3 a, float3 b) {
+  return f3(a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x);
+}
+__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) {
+  float s = copysignf(1.0f, n.z);
+  float a = -1.0f / (s + n.z);
+  float c = n.x * n.y * a;
+  t = f3(1.0f + s * n.x * n.x * a, s * c, -s * n.x);
+  b = f3(c, s + n.y * n.y * a, -n.y);
+}
+// ---------------------------------------------------------------- scene view
+struct DevScene {
+  const float* tris;     // [F, 9] v0 v1 v2
+  const int* mats;       // [F]
+  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* alb;      // [M, 3]
+  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];
+};
+// ------------------------------------------------------------- intersection
+__device__ __forceinline__ bool tri_hit(const float* v, float3 ro, float3 rd,
+                                        float tmin, float tmax, float& t,
+                                        float3& ng) {
+  float3 v0 = f3(v[0], v[1], v[2]);
+  float3 e1 = f3(v[3], v[4], v[5]) - v0;
+  float3 e2 = f3(v[6], v[7], v[8]) - v0;
+  float3 p = cross(rd, e2);
+  float det = dot(e1, p);
+  if (fabsf(det) < 1e-12f) return false;
+  float inv = 1.0f / det;
+  float3 s = ro - v0;
+  float u = dot(s, p) * inv;
+  if (u < -1e-6f || u > 1.0f + 1e-6f) return false;
+  float3 q = cross(s, e1);
+  float w = dot(rd, q) * inv;
+  if (w < -1e-6f || u + w > 1.0f + 1e-6f) return false;
+  float tt = dot(e2, q) * inv;
+  if (tt < tmin || tt > tmax) return false;
+  t = tt;
+  ng = cross(e1, e2);
+  return true;
+}
+__device__ __forceinline__ bool slab(const float* b, const float ro[3],
+                                     const float inv[3], float tmax) {
+  float t0 = kRayEps, t1 = tmax;
+#pragma unroll
+  for (int a = 0; a < 3; ++a) {
+    float lo = (b[a] - ro[a]) * inv[a];
+    float hi = (b[3 + a] - ro[a]) * inv[a];
+    if (lo > hi) {
+      float tmp = lo;
+      lo = hi;
+      hi = tmp;
+    }
+    t0 = fmaxf(t0, lo);
+    t1 = fminf(t1, hi);
+  }
+  return t0 <= t1;
+}
+__device__ __forceinline__ void inv_dir(float3 rd, float inv[3]) {
+  float d;
+  d = rd.x; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[0] = 1.0f / d;
+  d = rd.y; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[1] = 1.0f / d;
+  d = rd.z; if (fabsf(d) < 1e-12f) d = copysignf(1e-12f, d); inv[2] = 1.0f / d;
+}
+__device__ int bvh_closest(const DevScene& sc, float3 ro, float3 rd,
+                           float tmin, float& tbest, float3& ngbest) {
+  float roa[3] = {ro.x, ro.y, ro.z};
+  float inv[3];
+  inv_dir(rd, inv);
+  int stack[kStack];
+  int sp = 0;
+  stack[sp++] = 0;
+  int best = -1;
+  while (sp > 0) {
+    int nid = stack[--sp];
+    if (!slab(&sc.nf[nid * 6], roa, inv, tbest)) continue;
+    const int* n = &sc.ni[nid * 3];
+    if (n[2]) {  // leaf: faces n[0] .. n[0]+n[1]
+      for (int f = n[0]; f < n[0] + n[1]; ++f) {
+        float t;
+        float3 ng;
+        if (tri_hit(&sc.tris[f * 9], ro, rd, tmin, tbest, t, ng)) {
+          tbest = t;
+          ngbest = ng;
+          best = f;
+        }
+      }
+    } else if (sp + 2 <= kStack) {
+      stack[sp++] = n[0];
+      stack[sp++] = n[1];
+    }
+  }
+  return best;
+}
+__device__ bool bvh_occluded(const DevScene& sc, float3 ro, float3 rd,
+                             float tmax) {
+  float roa[3] = {ro.x, ro.y, ro.z};
+  float inv[3];
+  inv_dir(rd, inv);
+  int stack[kStack];
+  int sp = 0;
+  stack[sp++] = 0;
+  while (sp > 0) {
+    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;
+        float3 ng;
+        if (tri_hit(&sc.tris[f * 9], ro, rd, kRayEps, tmax, t, ng)) return true;
+      }
+    } else if (sp + 2 <= kStack) {
+      stack[sp++] = n[0];
+      stack[sp++] = n[1];
+    }
+  }
+  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;
+  }
+  return lo;
+}
+// ------------------------------------------------------------- path tracing
+// One camera path: forward accumulates radiance into acc; backward (GRAD)
+// accumulates parameter gradients into the shared buffers instead. 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* s_ga,
+                           float* s_ge) {
+  float T[3] = {1.0f, 1.0f, 1.0f};      // prefix albedo product
+  float albs[kMaxBounces][3];           // per-bounce albedo (GRAD replay)
+  int mats[kMaxBounces];
+  float suf[kMaxBounces + 1][3];        // suffix products workspace
+  for (int k = 0; k < B; ++k) {
+    float tbest = 3.0e38f;
+    float3 ng;
+    int face = bvh_closest(sc, ro, rd, kRayEps, tbest, 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);
+    const float* am = &sc.alb[m * 3];
+    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 {
+        // dC/dE[m] = T; dC/da[m_j] (j<k) = E * excl_j.
+        for (int c = 0; c < 3; ++c)
+          atomicAdd(&s_ge[m * 3 + c], gs[c] * T[c]);
+        if (k > 0) {
+          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 excl = pref[c] * suf[j + 1][c];
+              atomicAdd(&s_ga[mats[j] * 3 + c], gs[c] * em[c] * excl);
+              pref[c] *= albs[j][c];
+            }
+          }
+        }
+      }
+    }
+    if (GRAD) {
+      albs[k][0] = am[0];
+      albs[k][1] = am[1];
+      albs[k][2] = am[2];
+      mats[k] = m;
+    }
+    // Next-event estimation against the area-sampled light list. NEE at
+    // vertex k reaches the light at hit k+1, so it is skipped at the final
+    // vertex: both estimators then cover camera-to-light paths of at most
+    // max_bounces surface hits and agree in expectation, not just in the
+    // limit.
+    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: v0 + e1*(1-sqrt(r2)) + e2*(r3*sqrt(r2))
+      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);  // face the shading point
+      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; the vertex albedo a_k is
+            // one of the factors (albs[k] is already stored above).
+            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 excl = pref[c] * suf[j + 1][c];
+                atomicAdd(&s_ga[mats[j] * 3 + c], gs[c] * el[c] * S * excl);
+                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];
+  }
+}
+__device__ __forceinline__ float3 camera_ray(const DevCam& cam, int px, int py,
+                                             int W, int H, float jx, float jy) {
+  float nx = 2.0f * ((px + jx) / (float)W) - 1.0f;
+  float ny = 1.0f - 2.0f * ((py + jy) / (float)H);
+  float3 d = f3(cam.f[0] + nx * cam.r[0] + ny * cam.u[0],
+                cam.f[1] + nx * cam.r[1] + ny * cam.u[1],
+                cam.f[2] + nx * cam.r[2] + ny * cam.u[2]);
+  return normalize(d);
+}
+__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;
+  int px = pid % W, py = pid / W;
+  float3 ro = f3(cam.p[0], cam.p[1], cam.p[2]);
+  float3 acc = f3(0.0f, 0.0f, 0.0f);
+  for (int s = 0; s < spp; ++s) {
+    curandStatePhilox4_32_10_t st;
+    curand_init(seed, (unsigned long long)pid * spp + s, 0, &st);
+    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);
+  }
+  float inv_spp = 1.0f / (float)spp;
+  img[pid * 3 + 0] = acc.x * inv_spp;
+  img[pid * 3 + 1] = acc.y * inv_spp;
+  img[pid * 3 + 2] = acc.z * inv_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* ga, float* ge) {
+  __shared__ float s_ga[kMaxMats * 3];
+  __shared__ float s_ge[kMaxMats * 3];
+  int nm3 = sc.nm * 3;
+  for (int i = threadIdx.x; i < nm3; i += blockDim.x) {
+    s_ga[i] = 0.0f;
+    s_ge[i] = 0.0f;
+  }
+  __syncthreads();
+  int pid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (pid < H * W) {
+    int px = pid % W, py = pid / W;
+    float inv_spp = 1.0f / (float)spp;
+    float gs[3] = {gimg[pid * 3 + 0] * inv_spp, gimg[pid * 3 + 1] * inv_spp,
+                   gimg[pid * 3 + 2] * inv_spp};
+    float3 ro = f3(cam.p[0], cam.p[1], cam.p[2]);
+    float3 acc = f3(0.0f, 0.0f, 0.0f);
+    for (int s = 0; s < spp; ++s) {
+      curandStatePhilox4_32_10_t st;
+      curand_init(seed, (unsigned long long)pid * spp + s, 0, &st);
+      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, s_ga, s_ge);
+    }
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < nm3; i += blockDim.x) {
+    atomicAdd(&ga[i], s_ga[i]);
+    atomicAdd(&ge[i], s_ge[i]);
+  }
+}
+DevScene make_scene(const float* tris, const int* mat_ids,
+                    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* albedo,
+                    const float* emission, int n_mats) {
+  DevScene sc;
+  sc.tris = tris;
+  sc.mats = mat_ids;
+  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.alb = albedo;
+  sc.emi = emission;
+  sc.nm = n_mats;
+  return sc;
+}
+DevCam make_cam(const float* cam) {
+  DevCam c;
+  for (int i = 0; i < 3; ++i) {
+    c.p[i] = cam[i];
+    c.f[i] = cam[3 + i];
+    c.r[i] = cam[6 + i];
+    c.u[i] = cam[9 + i];
+  }
+  return c;
+}
+}  // namespace
+extern "C" void ptd_forward_launch(
+    const float* tris, const int* mat_ids, 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* albedo, 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, nodes_f, nodes_i, n_nodes,
+                           light_faces, light_cdf, n_lights, total_light_area,
+                           albedo, 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, 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* albedo, 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_albedo,
+    float* grad_emission, cudaStream_t stream) {
+  DevScene sc = make_scene(tris, mat_ids, nodes_f, nodes_i, n_nodes,
+                           light_faces, light_cdf, n_lights, total_light_area,
+                           albedo, emission, n_mats);
+  DevCam dc = make_cam(cam);
+  int blocks = (H * W + kThreads - 1) / kThreads;
+  k_backward<<<blocks, kThreads, 0, stream>>>(
+      sc, dc, H, W, spp, max_bounces, nee, (unsigned long long)seed,
+      grad_image, grad_albedo, grad_emission);
+}

pathtracer_diff_cuda/pathtracer_launch.h ADDED Viewed

	@@ -0,0 +1,30 @@

+// Raw-pointer launch interface for the pathtracer-diff kernels. Torch-free so
+// the .cu never sees libtorch headers (required for local MSVC builds).
+#pragma once
+#include <cuda_runtime_api.h>
+extern "C" {
+void ptd_forward_launch(
+    const float* tris, const int* mat_ids, 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* albedo, 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, 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* albedo, 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_albedo, float* grad_emission, cudaStream_t stream);
+}  // extern "C"

tests/test_pathtracer_diff.py ADDED Viewed

	@@ -0,0 +1,227 @@

+import pytest
+import torch
+import kernels
+ptd = kernels.get_kernel("phanerozoic/pathtracer-diff", version=1,
+                         trust_remote_code=True)
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(),
+                                   reason="CUDA required")
+def quad(a, b, c, d):
+    return [[a, b, c], [a, c, d]]
+def add_box(verts, faces, mat, lo, hi, mat_id):
+    x0, y0, z0 = lo
+    x1, y1, z1 = hi
+    base = len(verts)
+    verts.extend([(x0, y0, z0), (x1, y0, z0), (x1, y0, z1), (x0, y0, z1),
+                  (x0, y1, z0), (x1, y1, z0), (x1, y1, z1), (x0, y1, z1)])
+    i = lambda k: base + k
+    for q in [quad(i(0), i(1), i(2), i(3)),   # bottom
+              quad(i(4), i(7), i(6), i(5)),   # top
+              quad(i(0), i(4), i(5), i(1)),   # z0
+              quad(i(3), i(2), i(6), i(7)),   # z1
+              quad(i(0), i(3), i(7), i(4)),   # x0
+              quad(i(1), i(5), i(6), i(2))]:  # x1
+        faces.extend(q)
+        mat.extend([mat_id, mat_id])
+def furnace_scene(albedo=0.5, emission=0.25):
+    verts, faces, mat = [], [], []
+    add_box(verts, faces, mat, (-2, -2, -2), (2, 2, 2), 0)
+    return ptd.Scene(verts, faces, mat,
+                     albedo=[[albedo] * 3], emission=[[emission] * 3])
+def cornell(with_box=True, albedo=None, emission=None):
+    verts, faces, mat = [], [], []
+    def wall(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])
+    X, Y, Z = 5.56, 5.488, 5.592
+    wall((0, 0, 0), (X, 0, 0), (X, 0, Z), (0, 0, Z), 0)          # floor
+    wall((0, Y, 0), (0, Y, Z), (X, Y, Z), (X, Y, 0), 0)          # ceiling
+    wall((0, 0, Z), (X, 0, Z), (X, Y, Z), (0, Y, Z), 0)          # back
+    wall((0, 0, 0), (0, 0, Z), (0, Y, Z), (0, Y, 0), 1)          # left (red)
+    wall((X, 0, 0), (X, Y, 0), (X, Y, Z), (X, 0, Z), 2)          # right (green)
+    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)    # light
+    if with_box:
+        add_box(verts, faces, mat, (1.3, 0.0, 0.65), (2.4, 1.65, 1.7), 0)
+    if albedo is None:
+        albedo = [[0.73, 0.73, 0.73], [0.65, 0.05, 0.05],
+                  [0.12, 0.45, 0.15], [0.78, 0.78, 0.78]]
+    if emission is None:
+        emission = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [18.4, 15.6, 8.0]]
+    return ptd.Scene(verts, faces, mat, albedo=albedo, emission=emission)
+def cornell_camera():
+    return ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8),
+                      vfov_deg=39.0)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_furnace_analytic_zero_variance():
+    """Closed uniform box, BRDF sampling only: with cosine sampling the
+    diffuse throughput multiplier is exactly the albedo, so every path
+    returns exactly E * sum_{k<B} a^k and every pixel equals the analytic
+    series regardless of the random walk."""
+    a, E, B = 0.5, 0.25, 8
+    scene = furnace_scene(a, E)
+    cam = ptd.Camera(position=(0, 0, 0), look_at=(0, 0, 1), vfov_deg=60.0)
+    img = ptd.render(scene, cam, 32, 32, spp=2, max_bounces=B, nee=False,
+                     seed=0)
+    expected = E * (1 - a ** B) / (1 - a)
+    assert torch.isfinite(img).all()
+    assert (img - expected).abs().max().item() < 1e-3, \
+        (img.min().item(), img.max().item(), expected)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_forward_bitwise_deterministic():
+    scene = cornell()
+    cam = cornell_camera()
+    i1 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=11)
+    i2 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=11)
+    assert torch.equal(i1, i2)
+    i3 = ptd.render(scene, cam, 96, 96, spp=8, max_bounces=4, seed=12)
+    assert not torch.equal(i1, i3)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_nee_matches_brdf_estimator():
+    """The NEE-only and BRDF-only estimators target the same integral; their
+    image means must agree within Monte Carlo tolerance."""
+    verts, faces, mat = [], [], []
+    add_box(verts, faces, mat, (0, 0, 0), (5, 5, 5), 0)
+    base = len(verts)
+    verts.extend([(1.0, 4.999, 1.0), (4.5, 4.999, 1.0), (4.5, 4.999, 4.5),
+                  (1.0, 4.999, 4.5)])
+    faces.extend(quad(base, base + 1, base + 2, base + 3))
+    mat.extend([1, 1])
+    scene = ptd.Scene(verts, faces, mat,
+                      albedo=[[0.6, 0.6, 0.6], [0.0, 0.0, 0.0]],
+                      emission=[[0, 0, 0], [3.0, 3.0, 3.0]])
+    cam = ptd.Camera(position=(2.5, 2.5, 0.2), look_at=(2.5, 2.5, 5.0),
+                     vfov_deg=70.0)
+    m_nee = ptd.render(scene, cam, 64, 64, spp=64, max_bounces=4, nee=True,
+                       seed=1).mean().item()
+    m_brdf = ptd.render(scene, cam, 64, 64, spp=256, max_bounces=4, nee=False,
+                        seed=2).mean().item()
+    assert abs(m_nee - m_brdf) / m_brdf < 0.03, (m_nee, m_brdf)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_gradients_match_finite_differences():
+    """Sampling directions never depend on albedo/emission, so with a fixed
+    seed the loss is smooth in the parameters along the same paths and
+    central differences must match the analytic replay gradients."""
+    scene = cornell(with_box=False)
+    cam = cornell_camera()
+    H = W = 48
+    torch.manual_seed(0)
+    Wr = torch.rand(H, W, 3, device="cuda")
+    def loss_of(albedo, emission):
+        scene.albedo = albedo
+        scene.emission = emission
+        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, nee=True,
+                         seed=7)
+        return (img * Wr).sum()
+    alb0 = scene.albedo.clone()
+    emi0 = scene.emission.clone()
+    alb = alb0.clone().requires_grad_(True)
+    emi = emi0.clone().requires_grad_(True)
+    loss_of(alb, emi).backward()
+    checks = []
+    for (i, c) in [(0, 1), (1, 0), (2, 2)]:  # white g, red r, green b
+        h = 2e-3
+        ap = alb0.clone(); ap[i, c] += h
+        am = alb0.clone(); am[i, c] -= h
+        fd = (loss_of(ap, emi0) - loss_of(am, emi0)).item() / (2 * h)
+        an = alb.grad[i, c].item()
+        checks.append(("albedo", i, c, an, fd))
+    h = 0.05
+    ep = emi0.clone(); ep[3, 1] += h
+    em = emi0.clone(); em[3, 1] -= h
+    fd = (loss_of(alb0, ep) - loss_of(alb0, em)).item() / (2 * h)
+    an = emi.grad[3, 1].item()
+    checks.append(("emission", 3, 1, an, fd))
+    for kind, i, c, an, fd in checks:
+        assert fd != 0.0, (kind, i, c)
+        assert abs(an - fd) / max(abs(fd), 1e-6) < 2e-2, \
+            (kind, i, c, an, fd)
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_inverse_rendering_recovers_wall_albedo():
+    """Recover the left-wall albedo from a rendered target by gradient
+    descent through the kernel: the flagship no-framework-in-the-loop use."""
+    target_row = torch.tensor([0.2, 0.5, 0.7], device="cuda")
+    scene = cornell(with_box=False)
+    cam = cornell_camera()
+    H = W = 48
+    tgt_alb = scene.albedo.clone()
+    tgt_alb[1] = target_row
+    scene.albedo = tgt_alb
+    target = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=3).detach()
+    alb = scene.albedo.clone()
+    alb[1] = torch.tensor([0.5, 0.5, 0.5], device="cuda")
+    alb = alb.requires_grad_(True)
+    opt = torch.optim.Adam([alb], lr=0.05)
+    losses = []
+    for _ in range(80):
+        opt.zero_grad()
+        scene.albedo = alb
+        img = ptd.render(scene, cam, H, W, spp=8, max_bounces=3, seed=3)
+        loss = (img - target).square().mean()
+        loss.backward()
+        alb.grad[0] = 0
+        alb.grad[2] = 0
+        alb.grad[3] = 0
+        opt.step()
+        with torch.no_grad():
+            alb.clamp_(0.02, 0.98)
+        losses.append(loss.item())
+    err = (alb[1].detach() - target_row).abs().max().item()
+    assert err < 0.05, (alb[1].detach().tolist(), losses[0], losses[-1])
+    assert losses[-1] < losses[0] * 0.1, (losses[0], losses[-1])
+@requires_cuda
+@pytest.mark.kernels_ci
+def test_validation():
+    verts, faces, mat = [], [], []
+    add_box(verts, faces, mat, (0, 0, 0), (1, 1, 1), 0)
+    with pytest.raises(ValueError):
+        ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3] * 65,
+                  emission=[[0.0] * 3] * 65)
+    scene = ptd.Scene(verts, faces, mat, albedo=[[0.5] * 3],
+                      emission=[[1.0] * 3])
+    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()

torch-ext/pathtracer_diff/__init__.py ADDED Viewed

	@@ -0,0 +1,212 @@

+"""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, 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
+per-material albedo and emission, storing no path state between the passes.
+Sampling directions never depend on the differentiable parameters, so the
+replay is exact. No visibility/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)
+    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()
+"""
+import math
+import torch
+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. 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)."""
+    import numpy as np
+    tris = np.asarray(tris, dtype=np.float32)
+    F = tris.shape[0]
+    v = tris.reshape(F, 3, 3)
+    lo = v.min(axis=1)
+    hi = v.max(axis=1)
+    cen = (lo + hi) * 0.5
+    nodes_f, nodes_i, order = [], [], []
+    def rec(idx):
+        nid = len(nodes_f)
+        nodes_f.append(None)
+        nodes_i.append(None)
+        blo = lo[idx].min(axis=0)
+        bhi = hi[idx].max(axis=0)
+        pad = 1e-5 * (1.0 + np.abs(blo) + np.abs(bhi))
+        nodes_f[nid] = np.concatenate([blo - pad, bhi + pad])
+        if len(idx) <= leaf_size:
+            start = len(order)
+            order.extend(idx.tolist())
+            nodes_i[nid] = (start, len(idx), 1)
+            return nid
+        ax = int(np.argmax(bhi - blo))
+        srt = idx[np.argsort(cen[idx, ax], kind="stable")]
+        mid = len(srt) // 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(np.arange(F))
+    finally:
+        sys.setrecursionlimit(old)
+    return (np.stack(nodes_f).astype(np.float32),
+            np.asarray(nodes_i, dtype=np.int32),
+            np.asarray(order, dtype=np.int64))
+class Camera:
+    """Pinhole camera. `tensor(H, W)` packs (pos, forward, right*tan(v/2)*aspect,
+    up*tan(v/2)) as a float32 [12] tensor for the kernel."""
+    def __init__(self, position, look_at, up=(0.0, 1.0, 0.0), vfov_deg=40.0):
+        p = torch.tensor(position, dtype=torch.float64)
+        t = torch.tensor(look_at, dtype=torch.float64)
+        u = torch.tensor(up, dtype=torch.float64)
+        f = t - p
+        f = f / f.norm()
+        r = torch.linalg.cross(f, u)
+        r = r / r.norm()
+        uu = torch.linalg.cross(r, f)
+        self.position, self.forward, self.right, self.up = p, f, r, uu
+        self.vfov_deg = float(vfov_deg)
+    def tensor(self, H, W, device="cuda"):
+        th = math.tan(math.radians(self.vfov_deg) * 0.5)
+        rs = self.right * (th * W / H)
+        us = self.up * th
+        return torch.cat([self.position, self.forward, rs, us]).to(
+            device=device, dtype=torch.float32).contiguous()
+class Scene:
+    """Triangle-mesh scene with per-material albedo and emission.
+    vertices [V, 3] float, faces [F, 3] int, material_ids [F] int,
+    albedo [M, 3] float32 (may require grad), emission [M, 3] float32
+    (may require grad). Faces are reordered by the BVH build; emissive
+    faces become the area-sampled light list."""
+    def __init__(self, vertices, faces, material_ids, albedo, emission,
+                 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 not torch.is_tensor(albedo):
+            albedo = torch.tensor(albedo, dtype=torch.float32, device=device)
+        if not torch.is_tensor(emission):
+            emission = torch.tensor(emission, dtype=torch.float32, device=device)
+        M = albedo.shape[0]
+        if M > MAX_MATERIALS:
+            raise ValueError(f"at most {MAX_MATERIALS} materials, got {M}")
+        if albedo.shape != (M, 3) or emission.shape != (M, 3):
+            raise ValueError("albedo and emission must be [M, 3]")
+        if int(material_ids.max()) >= M or int(material_ids.min()) < 0:
+            raise ValueError("material_ids out of range")
+        tris = vertices[faces].reshape(-1, 9)
+        nodes_f, nodes_i, order = _build_bvh(tris.numpy())
+        order_t = torch.from_numpy(order)
+        tris = tris[order_t]
+        mat_ids = material_ids[order_t].to(torch.int32)
+        emissive = (emission.detach().cpu().amax(dim=1) > 0)[mat_ids.long()]
+        lf = torch.nonzero(emissive, as_tuple=False).flatten().to(torch.int32)
+        if lf.numel() > 0:
+            t = tris[lf.long()].reshape(-1, 3, 3).double()
+            e1 = t[:, 1] - t[:, 0]
+            e2 = t[:, 2] - t[:, 0]
+            areas = 0.5 * torch.linalg.cross(e1, e2).norm(dim=1)
+            total = float(areas.sum())
+            cdf = (areas.cumsum(0) / areas.sum()).float()
+        else:
+            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()
+        self.nodes_f = torch.from_numpy(nodes_f).to(device).contiguous()
+        self.nodes_i = torch.from_numpy(nodes_i).to(device).contiguous()
+        self.light_faces = lf.to(device).contiguous()
+        self.light_cdf = cdf.to(device).contiguous()
+        self.total_light_area = total
+        self.albedo = albedo.to(device)
+        self.emission = emission.to(device)
+    @property
+    def n_faces(self):
+        return self.tris.shape[0]
+class _RenderFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, albedo, emission, scene, cam_t, H, W, spp, max_bounces,
+                nee, seed):
+        image = torch.empty(H, W, 3, device=albedo.device, dtype=torch.float32)
+        ops.pt_forward(scene.tris, scene.mat_ids, scene.nodes_f, scene.nodes_i,
+                       scene.light_faces, scene.light_cdf,
+                       float(scene.total_light_area),
+                       albedo.detach().contiguous(),
+                       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(albedo, emission)
+        return image
+    @staticmethod
+    def backward(ctx, grad_image):
+        albedo, 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(albedo)
+        ge = torch.zeros_like(emission)
+        ops.pt_backward(scene.tris, scene.mat_ids, scene.nodes_f,
+                        scene.nodes_i, scene.light_faces, scene.light_cdf,
+                        float(scene.total_light_area),
+                        albedo.detach().contiguous(),
+                        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 scene.albedo and scene.emission. 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)
+    return _RenderFn.apply(scene.albedo, scene.emission, scene, cam_t, height,
+                           width, int(spp), int(max_bounces), bool(nee),
+                           int(seed))

torch-ext/torch_binding.cpp ADDED Viewed

	@@ -0,0 +1,149 @@

+#include <torch/library.h>
+#include <torch/all.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "registration.h"
+#include "torch_binding.h"
+#include "../pathtracer_diff_cuda/pathtracer_launch.h"
+namespace {
+void check_scene(const torch::Tensor& tris, const torch::Tensor& mat_ids,
+                 const torch::Tensor& nodes_f, const torch::Tensor& nodes_i,
+                 const torch::Tensor& light_faces,
+                 const torch::Tensor& light_cdf, const torch::Tensor& albedo,
+                 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,
+              "tris must be contiguous CUDA f32 [F, 9]");
+  TORCH_CHECK(mat_ids.is_cuda() && mat_ids.is_contiguous() &&
+                  mat_ids.dtype() == torch::kInt32 &&
+                  mat_ids.numel() == tris.size(0),
+              "mat_ids must be contiguous CUDA i32 [F]");
+  TORCH_CHECK(nodes_f.is_cuda() && nodes_f.is_contiguous() &&
+                  nodes_f.dtype() == torch::kFloat32 && nodes_f.dim() == 2 &&
+                  nodes_f.size(1) == 6,
+              "nodes_f must be contiguous CUDA f32 [N, 6]");
+  TORCH_CHECK(nodes_i.is_cuda() && nodes_i.is_contiguous() &&
+                  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(albedo.is_cuda() && albedo.is_contiguous() &&
+                  albedo.dtype() == torch::kFloat32 && albedo.dim() == 2 &&
+                  albedo.size(1) == 3,
+              "albedo must be contiguous CUDA f32 [M, 3]");
+  TORCH_CHECK(emission.is_cuda() && emission.is_contiguous() &&
+                  emission.dtype() == torch::kFloat32 &&
+                  emission.sizes() == albedo.sizes(),
+              "emission must be contiguous CUDA f32 [M, 3]");
+  TORCH_CHECK(albedo.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
+void pt_forward(torch::Tensor tris, torch::Tensor mat_ids,
+                torch::Tensor nodes_f, torch::Tensor nodes_i,
+                torch::Tensor light_faces, torch::Tensor light_cdf,
+                double total_light_area, torch::Tensor albedo,
+                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, nodes_f, nodes_i, light_faces, light_cdf, albedo,
+              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();
+  // The launcher reads the 12 camera floats on the host.
+  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>(),
+      (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, albedo.const_data_ptr<float>(),
+      emission.const_data_ptr<float>(), (int)albedo.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();
+}
+void pt_backward(torch::Tensor tris, torch::Tensor mat_ids,
+                 torch::Tensor nodes_f, torch::Tensor nodes_i,
+                 torch::Tensor light_faces, torch::Tensor light_cdf,
+                 double total_light_area, torch::Tensor albedo,
+                 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_albedo,
+                 torch::Tensor grad_emission) {
+  check_scene(tris, mat_ids, nodes_f, nodes_i, light_faces, light_cdf, albedo,
+              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,
+              "grad_image must be contiguous CUDA f32 [H, W, 3]");
+  TORCH_CHECK(grad_albedo.is_cuda() && grad_albedo.is_contiguous() &&
+                  grad_albedo.sizes() == albedo.sizes(),
+              "grad_albedo must match albedo");
+  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>(),
+      (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, albedo.const_data_ptr<float>(),
+      emission.const_data_ptr<float>(), (int)albedo.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_albedo.data_ptr<float>(), grad_emission.data_ptr<float>(), stream);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def(
+      "pt_forward(Tensor tris, Tensor mat_ids, Tensor nodes_f, Tensor nodes_i,"
+      " Tensor light_faces, Tensor light_cdf, float total_light_area,"
+      " Tensor albedo, Tensor emission, Tensor cam, int spp, int max_bounces,"
+      " bool nee, int seed, Tensor! image) -> ()");
+  ops.impl("pt_forward", torch::kCUDA, &pt_forward);
+  ops.def(
+      "pt_backward(Tensor tris, Tensor mat_ids, Tensor nodes_f,"
+      " Tensor nodes_i, Tensor light_faces, Tensor light_cdf,"
+      " float total_light_area, Tensor albedo, Tensor emission, Tensor cam,"
+      " int spp, int max_bounces, bool nee, int seed, Tensor grad_image,"
+      " Tensor! grad_albedo, Tensor! grad_emission) -> ()");
+  ops.impl("pt_backward", torch::kCUDA, &pt_backward);
+}
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h ADDED Viewed

	@@ -0,0 +1,20 @@

+#pragma once
+#include <torch/torch.h>
+void pt_forward(torch::Tensor tris, torch::Tensor mat_ids,
+                torch::Tensor nodes_f, torch::Tensor nodes_i,
+                torch::Tensor light_faces, torch::Tensor light_cdf,
+                double total_light_area, torch::Tensor albedo,
+                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 nodes_f, torch::Tensor nodes_i,
+                 torch::Tensor light_faces, torch::Tensor light_cdf,
+                 double total_light_area, torch::Tensor albedo,
+                 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_albedo,
+                 torch::Tensor grad_emission);