diff param identifiers

LuisaGroup · Dec 6, 2023 · 2c12995 · 2c12995
1 parent c4571ef
commit 2c12995
Show file tree

Hide file tree

Showing 16 changed files with 87 additions and 7 deletions.
diff --git a/src/base/differentiation.cpp b/src/base/differentiation.cpp
@@ -105,10 +105,12 @@ Differentiation::TexturedParameter Differentiation::parameter(const Image<float>
     auto grad_offset = _gradient_buffer_size;
     auto param_offset = _param_buffer_size;
     auto counter_offset = _counter_size;
+    auto param_index = static_cast<uint>(_textured_params.size());
     _counter_size = (_counter_size + pixel_count + 3u) & ~0b11u;
     _param_buffer_size = (_param_buffer_size + param_count + 3u) & ~0b11u;
     _gradient_buffer_size = (_gradient_buffer_size + param_count + 3u) & ~0b11u;
-    return _textured_params.emplace_back(TexturedParameter{image, s, grad_offset, param_offset, counter_offset, range});
+    return _textured_params.emplace_back(TexturedParameter{
+        param_index, image, s, grad_offset, param_offset, counter_offset, range});
 }
 
 void Differentiation::materialize(CommandBuffer &command_buffer) noexcept {

diff --git a/src/base/differentiation.h b/src/base/differentiation.h
@@ -53,29 +53,36 @@ class Differentiation {
             : _index{index}, _channels{channels} {}
         [[nodiscard]] auto index() const noexcept { return _index; }
         [[nodiscard]] auto channels() const noexcept { return _channels; }
+        [[nodiscard]] auto identifier() const noexcept { return luisa::format("diffconst({})", _index); }
     };
 
     class TexturedParameter {
 
     private:
         const Image<float> &_image;
+        uint _index;
         TextureSampler _sampler;
         uint _grad_offset;
         uint _param_offset;
         uint _counter_offset;
         float2 _range;
 
     public:
-        TexturedParameter(const Image<float> &image, TextureSampler sampler, uint grad_offset, uint param_offset,
+        TexturedParameter(uint index,
+                          const Image<float> &image, TextureSampler sampler,
+                          uint grad_offset, uint param_offset,
                           uint counter_offset, float2 range) noexcept
-            : _image{image}, _sampler{sampler}, _grad_offset{grad_offset}, _param_offset{param_offset},
+            : _image{image}, _index{index}, _sampler{sampler},
+              _grad_offset{grad_offset}, _param_offset{param_offset},
               _counter_offset{counter_offset}, _range{range} {}
         [[nodiscard]] auto &image() const noexcept { return _image; }
+        [[nodiscard]] auto index() const noexcept { return _index; }
         [[nodiscard]] auto sampler() const noexcept { return _sampler; }
         [[nodiscard]] auto range() const noexcept { return _range; }
         [[nodiscard]] auto gradient_buffer_offset() const noexcept { return _grad_offset; }
         [[nodiscard]] auto param_offset() const noexcept { return _param_offset; }
         [[nodiscard]] auto counter_offset() const noexcept { return _counter_offset; }
+        [[nodiscard]] auto identifier() const noexcept { return luisa::format("difftex({})", _index); }
     };
 
 private:

diff --git a/src/base/surface.h b/src/base/surface.h
@@ -244,7 +244,9 @@ class OpacitySurfaceWrapper : public BaseSurface {
         [[nodiscard]] luisa::string closure_identifier() const noexcept override {
             auto base_identifier = BaseInstance::closure_identifier();
             if (_opacity == nullptr) { return base_identifier; }
-            return luisa::format("opacity<{}>", base_identifier);
+            return luisa::format("opacity<{}, {}>",
+                                 base_identifier,
+                                 Texture::Instance::diff_param_identifier(_opacity));
         }
         [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(
             const SampledWavelengths &swl, Expr<float> time) const noexcept override {

diff --git a/src/base/texture.h b/src/base/texture.h
@@ -68,6 +68,15 @@ class Texture : public SceneNode {
         void backward_unbounded_spectrum(
             const Interaction &it, const SampledWavelengths &swl,
             Expr<float> time, const SampledSpectrum &dSpec) const noexcept;
+
+    public:
+        inline static const luisa::string non_differrentiable_identifier = "nodiff";
+        [[nodiscard]] virtual luisa::string diff_param_identifier() const noexcept {
+            return non_differrentiable_identifier;
+        }
+        [[nodiscard]] static auto diff_param_identifier(const Texture::Instance *t) noexcept {
+            return t ? t->diff_param_identifier() : non_differrentiable_identifier;
+        }
     };
 
 private:

diff --git a/src/surfaces/disney.cpp b/src/surfaces/disney.cpp
@@ -895,6 +895,7 @@ class ThinDisneyClosureImpl : public DisneyClosureImplBase {
         return {.eval = eval, .wi = wi, .event = event};
     }
 };
+
 class DisneySurfaceInstance : public Surface::Instance {
 
 private:

diff --git a/src/surfaces/glass.cpp b/src/surfaces/glass.cpp
@@ -114,6 +114,13 @@ class GlassInstance : public Surface::Instance {
         const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it,
                           Expr<float3> wo, Expr<float> eta_i) const noexcept override;
+    [[nodiscard]] luisa::string closure_identifier() const noexcept override {
+        return luisa::format("glass<{}, {}, {}, {}>",
+                             Texture::Instance::diff_param_identifier(_kr),
+                             Texture::Instance::diff_param_identifier(_kt),
+                             Texture::Instance::diff_param_identifier(_roughness),
+                             Texture::Instance::diff_param_identifier(_eta));
+    }
 };
 
 luisa::unique_ptr<Surface::Instance> GlassSurface::_build(

diff --git a/src/surfaces/layered.cpp b/src/surfaces/layered.cpp
@@ -170,9 +170,12 @@ class LayeredSurfaceInstance : public Surface::Instance {
 
 public:
     [[nodiscard]] luisa::string closure_identifier() const noexcept override {
-        return luisa::format("layered<{}, {}>",
+        return luisa::format("layered<{}, {}, {}, {}, {}>",
                              _top->closure_identifier(),
-                             _bottom->closure_identifier());
+                             _bottom->closure_identifier(),
+                             Texture::Instance::diff_param_identifier(_thickness),
+                             Texture::Instance::diff_param_identifier(_g),
+                             Texture::Instance::diff_param_identifier(_albedo));
     }
     [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it, Expr<float3> wo, Expr<float> eta_i) const noexcept override;

diff --git a/src/surfaces/matte.cpp b/src/surfaces/matte.cpp
@@ -49,6 +49,11 @@ class MatteInstance : public Surface::Instance {
 public:
     [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it, Expr<float3> wo, Expr<float> eta_i) const noexcept override;
+    [[nodiscard]] luisa::string closure_identifier() const noexcept override {
+        return luisa::format("matte<{}, {}>",
+                             Texture::Instance::diff_param_identifier(_kd),
+                             Texture::Instance::diff_param_identifier(_sigma));
+    }
 };
 
 luisa::unique_ptr<Surface::Instance> MatteSurface::_build(

diff --git a/src/surfaces/metal.cpp b/src/surfaces/metal.cpp
@@ -179,6 +179,11 @@ class MetalInstance : public Surface::Instance {
 public:
     [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it, Expr<float3> wo, Expr<float> eta_i) const noexcept override;
+    [[nodiscard]] luisa::string closure_identifier() const noexcept override {
+        return luisa::format("metal<{}, {}>",
+                             Texture::Instance::diff_param_identifier(_roughness),
+                             Texture::Instance::diff_param_identifier(_kd));
+    }
 };
 
 luisa::unique_ptr<Surface::Instance> MetalSurface::_build(

diff --git a/src/surfaces/mirror.cpp b/src/surfaces/mirror.cpp
@@ -53,6 +53,11 @@ class MirrorInstance : public Surface::Instance {
 public:
     [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it, Expr<float3> wo, Expr<float> eta_i) const noexcept override;
+    [[nodiscard]] luisa::string closure_identifier() const noexcept override {
+        return luisa::format("mirror<{}, {}>",
+                             Texture::Instance::diff_param_identifier(_color),
+                             Texture::Instance::diff_param_identifier(_roughness));
+    }
 };
 
 luisa::unique_ptr<Surface::Instance> MirrorSurface::_build(

diff --git a/src/surfaces/mix.cpp b/src/surfaces/mix.cpp
@@ -51,7 +51,10 @@ class MixSurfaceInstance : public Surface::Instance {
 
 public:
     [[nodiscard]] luisa::string closure_identifier() const noexcept override {
-        return luisa::format("mix<{}, {}>", _a->closure_identifier(), _b->closure_identifier());
+        return luisa::format("mix<{}, {}, {}>",
+                             _a->closure_identifier(),
+                             _b->closure_identifier(),
+                             Texture::Instance::diff_param_identifier(_ratio));
     }
     [[nodiscard]] luisa::unique_ptr<Surface::Closure> create_closure(const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it, Expr<float3> wo, Expr<float> eta_i) const noexcept override;

diff --git a/src/surfaces/plastic.cpp b/src/surfaces/plastic.cpp
@@ -91,6 +91,14 @@ class PlasticInstance : public Surface::Instance {
         const SampledWavelengths &swl, Expr<float> time) const noexcept override;
     void populate_closure(Surface::Closure *closure, const Interaction &it,
                           Expr<float3> wo, Expr<float> eta_i) const noexcept override;
+    [[nodiscard]] luisa::string closure_identifier() const noexcept override {
+        return luisa::format("plastic<{}, {}, {}, {}, {}>",
+                             Texture::Instance::diff_param_identifier(_kd),
+                             Texture::Instance::diff_param_identifier(_roughness),
+                             Texture::Instance::diff_param_identifier(_sigma_a),
+                             Texture::Instance::diff_param_identifier(_eta),
+                             Texture::Instance::diff_param_identifier(_thickness));
+    }
 };
 
 luisa::unique_ptr<Surface::Instance> PlasticSurface::_build(

diff --git a/src/textures/checkerboard.cpp b/src/textures/checkerboard.cpp
@@ -117,6 +117,16 @@ class CheckerboardTextureInstance final : public Texture::Instance {
             LUISA_ERROR_WITH_LOCATION("Not supported.");
         }
     }
+
+    [[nodiscard]] luisa::string diff_param_identifier() const noexcept override {
+        auto on_ident = Instance::diff_param_identifier(_on);
+        auto off_ident = Instance::diff_param_identifier(_off);
+        if (on_ident == non_differrentiable_identifier &&
+            off_ident == non_differrentiable_identifier) {
+            return non_differrentiable_identifier;
+        }
+        return luisa::format("diffcheckerboard<{}, {}>", on_ident, off_ident);
+    }
 };
 
 luisa::unique_ptr<Texture::Instance> CheckerboardTexture::build(

diff --git a/src/textures/constant.cpp b/src/textures/constant.cpp
@@ -115,6 +115,9 @@ class ConstantTextureInstance final : public Texture::Instance {
             pipeline().differentiation()->accumulate(*_diff_param, grad, slot_seed);
         }
     }
+    [[nodiscard]] luisa::string diff_param_identifier() const noexcept override {
+        return _diff_param ? _diff_param->identifier() : non_differrentiable_identifier;
+    }
 };
 
 luisa::unique_ptr<Texture::Instance> ConstantTexture::build(

diff --git a/src/textures/image.cpp b/src/textures/image.cpp
@@ -210,6 +210,9 @@ class ImageTextureInstance final : public Texture::Instance {
             pipeline().differentiation()->accumulate(*_diff_param, uv, grad);
         }
     }
+    [[nodiscard]] luisa::string diff_param_identifier() const noexcept override {
+        return _diff_param ? _diff_param->identifier() : non_differrentiable_identifier;
+    }
 };
 
 luisa::unique_ptr<Texture::Instance> ImageTexture::build(Pipeline &pipeline, CommandBuffer &command_buffer) const noexcept {

diff --git a/src/textures/swizzle.cpp b/src/textures/swizzle.cpp
@@ -111,6 +111,13 @@ class SwizzleTextureInstance final : public Texture::Instance {
             _base->backward(it, swl, time, g);
         }
     }
+    [[nodiscard]] luisa::string diff_param_identifier() const noexcept override {
+        auto base_ident = Instance::diff_param_identifier(_base);
+        return base_ident == non_differrentiable_identifier ?
+                   non_differrentiable_identifier :
+                   luisa::format("diffswizzle<{}, {}>",
+                                 base_ident, node<SwizzleTexture>()->channels());
+    }
 };
 
 luisa::unique_ptr<Texture::Instance> SwizzleTexture::build(
-Original file line number
+Diff line change
@@ Expand Up @@
             return {.eval = eval, .wi = wi, .event = event};
         }
     };
     class DisneySurfaceInstance : public Surface::Instance {
     private:
@@ Expand Down @@