bug in autodiff TBD

.gitignore

@@ -299,6 +299,7 @@ CMakeSettings.json
 
 *.png
 /outputs/
+/outputs_ref/
 /imgui.ini
 /set_accel_kernel
 /results.txt

cbox_diff/cbox-diff/cbox-diff-2.luisa → cbox-diff/scenes/cbox-diff/cbox-diff-2.luisa

@@ -124,11 +124,15 @@ Film film : Color {
   exposure { 0 }
 }
 
+FIlm display : Display {
+  base { @film }
+}
+
 Camera camera : Pinhole {
   position { 0.0, 1.0, 5.0 }
   fov { 27.8 }
-  spp { 1024 }
-  film { @film }
+  spp { 64 }
+  film { @display }
   file { "../../renders/cbox-diff.exr" }
   filter : Gaussian {
     radius { 1 }
@@ -150,7 +154,7 @@ Integrator pt : MegaReplayDiff {
   iterations { 100 }
 
   optimizer : Adam {
-    learning_rate { 0.003 }
+    learning_rate { 0.03 }
   }
 }
 

data/scenes/cbox-diff/cbox-diff-2.luisa

@@ -151,10 +151,10 @@ Integrator pt : MegaReplayDiff {
   display_camera_index { 0 }
   save_process { true }
   loss : L2 {}
-  iterations { 1000 }
+  iterations { 100 }
 
   optimizer : Adam {
-    learning_rate { 0.3 }
+    learning_rate { 0.03 }
   }
 }
 

src/base/differentiation.cpp

@@ -50,7 +50,7 @@ Differentiation::Differentiation(Pipeline &pipeline) noexcept
 #ifdef LUISA_RENDER_USE_BP_TIMES_NORMALIZATION
         grad /= Float(max(count, 1u));
 #endif
-        device_log("grad accumulated: ({}, {}, {})", grad[0u], grad[1u], grad[2u]);
+        // device_log("grad accumulated: ({}, {}, {})", grad[0u], grad[1u], grad[2u]);
 
         auto param_offset = thread * 4u;
         param_gradients.write(param_offset + 0u, grad.x);
@@ -77,7 +77,6 @@ Differentiation::Differentiation(Pipeline &pipeline) noexcept
 
 Differentiation::ConstantParameter Differentiation::parameter(float4 x, uint channels, float2 range) noexcept {
     auto index = static_cast<uint>(_constant_params.size());
-    LUISA_INFO("index: {}", index);
     _constant_params.emplace_back(x);
     _constant_ranges.emplace_back(range);
     return {index, channels};

src/base/surface.cpp

@@ -70,6 +70,7 @@ Surface::Sample Surface::Closure::sample(Expr<float3> wo,
 void Surface::Closure::backward(
     Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df,
     TransportMode mode) const noexcept {
+    compute::device_log("ok till before backward surface");
     if (instance()->node()->is_differentiable()) {
         _backward(wo, wi, df, mode);
     }

src/integrators/mega_replay_diff.cpp

@@ -2,6 +2,7 @@
 // Created by ChenXin on 2022/2/23.
 //
 // #include <iostream>
+#include "core/logging.h"
 #include <luisa-compute.h>
 #include <util/imageio.h>
 #include <util/sampling.h>
@@ -14,7 +15,7 @@
 namespace luisa::render {
 
 // #define LUISA_RENDER_PATH_REPLAY_DEBUG
-#define LUISA_RENDER_PATH_REPLAY_DEBUG_2
+// #define LUISA_RENDER_PATH_REPLAY_DEBUG_2
 
 using namespace luisa::compute;
 

src/surfaces/glass.cpp

@@ -252,7 +252,7 @@ class GlassClosure : public Surface::Closure {
         $if(same_hemisphere(wo_local, wi_local)) {
             // Kr
             if (_instance->Kr() && _instance->Kr()->node()->requires_gradients()) {
-                auto d_f = refl.backward(wo_local, wi_local, df);
+                auto d_f = refl.backward(wo_local, wi_local, df, mode);
                 d_alpha = d_f.dAlpha;
                 _instance->Kr()->backward_albedo_spectrum(it, swl(), time(), zero_if_any_nan(d_f.dR));
             }

src/surfaces/matte.cpp

@@ -125,7 +125,8 @@ class MatteClosure : public Surface::Closure {
         auto wo_local = ctx.it.shading().world_to_local(wo);
         auto wi_local = ctx.it.shading().world_to_local(wi);
         auto df = df_in * abs_cos_theta(wi_local);
-        auto grad = _refl->backward(wo_local, wi_local, df);
+        compute::device_log("ok till before backward matte");
+        auto grad = _refl->backward(wo_local, wi_local, df, mode);
         // device_log("grad in matte: ({}, {}, {})", grad.dR[0u], grad.dR[1u], grad.dR[2u]);
         if (auto kd = _instance->Kd()) {
             kd->backward_albedo_spectrum(ctx.it, swl(), time(), zero_if_any_nan(grad.dR));

src/surfaces/mirror.cpp

@@ -159,7 +159,7 @@ class MirrorClosure : public Surface::Closure {
         auto wo_local = it.shading().world_to_local(wo);
         auto wi_local = it.shading().world_to_local(wi);
         auto df = df_in * abs_cos_theta(wi_local);
-        auto grad = refl.backward(wo_local, wi_local, df);
+        auto grad = refl.backward(wo_local, wi_local, df, mode);
         auto d_fresnel = dynamic_cast<SchlickFresnel::Gradient *>(grad.dFresnel.get());
         if (auto color = _instance->color()) {
             color->backward_albedo_spectrum(it, swl(), time(), zero_if_any_nan(grad.dR + d_fresnel->dR0));

src/util/scattering.cpp

@@ -2,14 +2,20 @@
 // Created by Mike Smith on 2022/1/31.
 //
 
+#include "dsl/builtin.h"
+#include "dsl/stmt.h"
+#include "util/spec.h"
 #include <dsl/sugar.h>
 #include <util/frame.h>
 #include <util/sampling.h>
 #include <util/scattering.h>
 
 namespace luisa::render {
 
+using compute::backward;
 using compute::Callable;
+using compute::grad;
+using compute::requires_grad;
 
 Bool refract(Float3 wi, Float3 n, Float eta, Float3 *wt) noexcept {
     static Callable impl = [](Float3 wi, Float3 n, Float eta) noexcept {
@@ -317,7 +323,7 @@ SampledSpectrum LambertianReflection::evaluate(
 }
 
 LambertianReflection::Gradient LambertianReflection::backward(
-    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept {
+    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept {
     return {.dR = df * ite(same_hemisphere(wo, wi), inv_pi, 0.f)};
 }
 
@@ -337,7 +343,7 @@ Float LambertianTransmission::pdf(Expr<float3> wo, Expr<float3> wi, TransportMod
 }
 
 LambertianTransmission::Gradient LambertianTransmission::backward(
-    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) noexcept {
+    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) noexcept {
     return {.dT = df * ite(!same_hemisphere(wo, wi), inv_pi, 0.f)};
 }
 
@@ -378,7 +384,7 @@ Float MicrofacetReflection::pdf(Expr<float3> wo, Expr<float3> wi, TransportMode
 }
 
 MicrofacetReflection::Gradient MicrofacetReflection::backward(
-    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept {
+    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept {
     using compute::any;
     using compute::normalize;
     auto wh = wi + wo;
@@ -500,6 +506,11 @@ OrenNayar::OrenNayar(const SampledSpectrum &R, Expr<float> sigma) noexcept
 
 SampledSpectrum OrenNayar::evaluate(
     Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept {
+    return forward_compute(wo, wi, mode, _a, _b, _r);
+}
+
+SampledSpectrum OrenNayar::forward_compute(
+    Expr<float3> wo, Expr<float3> wi, TransportMode mode, Float a, Float b, SampledSpectrum r) const noexcept {
     auto valid = same_hemisphere(wo, wi);
     auto s = ite(valid, inv_pi, 0.f);
     static Callable scale = [](Float3 wo, Float3 wi, Float a, Float b) noexcept {
@@ -520,39 +531,28 @@ SampledSpectrum OrenNayar::evaluate(
                            sinThetaI / absCosThetaI, sinThetaO / absCosThetaO);
         return (a + b * maxCos * sinAlpha * tanBeta);
     };
-    return s * scale(wo, wi, _a, _b) * _r;
+    return s * scale(wo, wi, a, b) * r;
 }
 
 OrenNayar::Gradient OrenNayar::backward(
-    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept {
-    auto sinThetaI = sin_theta(wi);
-    auto sinThetaO = sin_theta(wo);
-    // Compute cosine term of Oren-Nayar model
-    auto sinPhiI = sin_phi(wi);
-    auto cosPhiI = cos_phi(wi);
-    auto sinPhiO = sin_phi(wo);
-    auto cosPhiO = cos_phi(wo);
-    auto dCos = cosPhiI * cosPhiO + sinPhiI * sinPhiO;
-    auto maxCos = ite(sinThetaI > 1e-4f & sinThetaO > 1e-4f, max(0.f, dCos), 0.f);
-    // Compute sine and tangent terms of Oren-Nayar model
-    auto absCosThetaI = abs_cos_theta(wi);
-    auto absCosThetaO = abs_cos_theta(wo);
-    auto sinAlpha = ite(absCosThetaI > absCosThetaO, sinThetaO, sinThetaI);
-    auto tanBeta = ite(absCosThetaI > absCosThetaO,
-                       sinThetaI / absCosThetaI, sinThetaO / absCosThetaO);
-    auto sigma2 = sqr(radians(_sigma));
-
-    // backward
-    auto sigma2_sigma = radians(_sigma) * pi / 90.f;
-    auto a_sigma2 = -0.165f / sqr(sigma2 + 0.33f);
-    auto b_sigma2 = 0.0405f / sqr(sigma2 + 0.09f);
-    auto k0 = maxCos * sinAlpha * tanBeta;
-    auto d_r = df * inv_pi * (_a + _b * k0);
-    auto k1 = inv_pi * (df * _r).sum();
-    auto d_a = k1;
-    auto d_b = k1 * k0;
-    auto d_sigma2 = d_a * a_sigma2 + d_b * b_sigma2;
-    auto d_sigma = d_sigma2 * sigma2_sigma;
+    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept {
+    auto d_r = _r;
+    auto d_sigma = _sigma;
+    compute::device_log("ok till before backward");
+    $autodiff {
+        auto r = _r;
+        auto sigma = _sigma;
+        r.requires_grad();
+        requires_grad(sigma);
+        auto sigma2 = sqr(radians(sigma));
+        auto a = 1.f - (sigma2 / (2.f * sigma2 + 0.66f));
+        auto b = 0.45f * sigma2 / (sigma2 + 0.09f);
+        auto y = forward_compute(wo, wi, mode, a, b, r);
+        y.backward(df);
+        d_r = r.grad();
+        d_sigma = grad(sigma);
+    };
+    compute::device_log("ok after backward");
     return {.dR = d_r, .dSigma = d_sigma};
 }
 
@@ -604,7 +604,7 @@ Float FresnelBlend::pdf(Expr<float3> wo, Expr<float3> wi, TransportMode mode) co
 }
 
 FresnelBlend::Gradient FresnelBlend::backward(
-    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept {
+    Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept {
 
     auto pow5 = [](auto &&v) noexcept { return sqr(sqr(v)) * v; };
     auto absCosThetaI = abs_cos_theta(wi);

src/util/scattering.h

@@ -152,7 +152,7 @@ class LambertianReflection : public BxDF {
 public:
     explicit LambertianReflection(const SampledSpectrum &R) noexcept : _r{R} {}
     [[nodiscard]] SampledSpectrum evaluate(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
-    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept;
+    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept;
     [[nodiscard]] SampledSpectrum albedo() const noexcept override { return _r; }
 };
 
@@ -171,7 +171,7 @@ class LambertianTransmission : public BxDF {
     [[nodiscard]] SampledSpectrum evaluate(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
     [[nodiscard]] SampledDirection sample_wi(Expr<float3> wo, Expr<float2> u, TransportMode mode) const noexcept override;
     [[nodiscard]] Float pdf(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
-    [[nodiscard]] static Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) noexcept;
+    [[nodiscard]] static Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) noexcept;
     [[nodiscard]] SampledSpectrum albedo() const noexcept override { return SampledSpectrum(0.f); }
 };
 
@@ -196,7 +196,7 @@ class MicrofacetReflection : public BxDF {
     [[nodiscard]] SampledSpectrum evaluate(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
     [[nodiscard]] SampledDirection sample_wi(Expr<float3> wo, Expr<float2> u, TransportMode mode) const noexcept override;
     [[nodiscard]] Float pdf(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
-    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept;
+    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept;
     [[nodiscard]] SampledSpectrum albedo() const noexcept override { return _r; }
 };
 
@@ -245,7 +245,8 @@ class OrenNayar : public BxDF {
 public:
     OrenNayar(const SampledSpectrum &R, Expr<float> sigma) noexcept;
     [[nodiscard]] SampledSpectrum evaluate(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
-    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept;
+    [[nodiscard]] SampledSpectrum forward_compute(Expr<float3> wo, Expr<float3> wi, TransportMode mode, Float a, Float b, SampledSpectrum r) const noexcept;
+    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept;
     [[nodiscard]] SampledSpectrum albedo() const noexcept override { return _r; }
 };
 
@@ -275,7 +276,7 @@ class FresnelBlend : public BxDF {
     [[nodiscard]] SampledSpectrum evaluate(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
     [[nodiscard]] SampledDirection sample_wi(Expr<float3> wo, Expr<float2> u, TransportMode mode) const noexcept override;
     [[nodiscard]] Float pdf(Expr<float3> wo, Expr<float3> wi, TransportMode mode) const noexcept override;
-    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df) const noexcept;
+    [[nodiscard]] Gradient backward(Expr<float3> wo, Expr<float3> wi, const SampledSpectrum &df, TransportMode mode) const noexcept;
     [[nodiscard]] SampledSpectrum albedo() const noexcept override { return _rd; }
 };