diff --git a/include/mitsuba/core/bitmap.h b/include/mitsuba/core/bitmap.h index 2bc50ecc2..2feb41691 100644 --- a/include/mitsuba/core/bitmap.h +++ b/include/mitsuba/core/bitmap.h @@ -652,6 +652,8 @@ class MI_EXPORT_LIB Bitmap : public Object { bool m_premultiplied_alpha; bool m_owns_data; Properties m_metadata; + + DR_TRAVERSE_CB(Object, m_size); }; diff --git a/include/mitsuba/core/bsphere.h b/include/mitsuba/core/bsphere.h index 63d70f446..4f33ec4cb 100644 --- a/include/mitsuba/core/bsphere.h +++ b/include/mitsuba/core/bsphere.h @@ -5,7 +5,7 @@ NAMESPACE_BEGIN(mitsuba) /// Generic n-dimensional bounding sphere data structure -template struct BoundingSphere { +template struct BoundingSphere: drjit::TraversableBase { static constexpr size_t Size = Point_::Size; using Point = Point_; using Float = dr::value_t; @@ -74,6 +74,8 @@ template struct BoundingSphere { dr::squared_norm(o) - dr::square(radius) ); } + + DR_TRAVERSE_CB(drjit::TraversableBase, center, radius); }; /// Print a string representation of the bounding sphere diff --git a/include/mitsuba/core/class.h b/include/mitsuba/core/class.h index 60763ed62..e863db350 100644 --- a/include/mitsuba/core/class.h +++ b/include/mitsuba/core/class.h @@ -247,7 +247,7 @@ NAMESPACE_END(detail) #define MI_REGISTRY_PUT(name, ptr) \ if constexpr (dr::is_jit_v) { \ - jit_registry_put(::mitsuba::detail::get_variant(), \ + drjit::registry_put(::mitsuba::detail::get_variant(), \ "mitsuba::" name, ptr); \ } diff --git a/include/mitsuba/core/distr_1d.h b/include/mitsuba/core/distr_1d.h index deec253b0..abfe553cc 100644 --- a/include/mitsuba/core/distr_1d.h +++ b/include/mitsuba/core/distr_1d.h @@ -4,6 +4,7 @@ #include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -17,7 +18,7 @@ NAMESPACE_BEGIN(mitsuba) * initialization. The associated scale factor can be retrieved using the * function \ref normalization(). */ -template struct DiscreteDistribution { +template struct DiscreteDistribution: drjit::TraversableBase { using Float = std::conditional_t, dr::value_t, Value>; using FloatStorage = DynamicBuffer; @@ -269,6 +270,9 @@ template struct DiscreteDistribution { Float m_sum = 0.f; Float m_normalization = 0.f; Vector2u m_valid; + + DR_TRAVERSE_CB(drjit::TraversableBase, m_pmf, m_cdf, m_sum, m_normalization, + m_valid); }; /** @@ -283,7 +287,7 @@ template struct DiscreteDistribution { * initialization. The associated scale factor can be retrieved using the * function \ref normalization(). */ -template struct ContinuousDistribution { +template struct ContinuousDistribution: drjit::TraversableBase { using Float = std::conditional_t, dr::value_t, Value>; using FloatStorage = DynamicBuffer; @@ -601,6 +605,10 @@ template struct ContinuousDistribution { ScalarVector2f m_range { 0.f, 0.f }; Vector2u m_valid; ScalarFloat m_max = 0.f; + + DR_TRAVERSE_CB(drjit::TraversableBase, m_pdf, m_cdf, m_integral, + m_normalization, m_interval_size, m_inv_interval_size, + m_valid); }; /** @@ -615,7 +623,7 @@ template struct ContinuousDistribution { * initialization. The associated scale factor can be retrieved using the * function \ref normalization(). */ -template struct IrregularContinuousDistribution { +template struct IrregularContinuousDistribution : public drjit::TraversableBase{ using Float = std::conditional_t, dr::value_t, Value>; using FloatStorage = DynamicBuffer; @@ -973,6 +981,9 @@ template struct IrregularContinuousDistribution { Vector2u m_valid; ScalarFloat m_interval_size = 0.f; ScalarFloat m_max = 0.f; + + DR_TRAVERSE_CB(drjit::TraversableBase, m_nodes, m_pdf, m_cdf, m_integral, + m_normalization, m_valid); }; template diff --git a/include/mitsuba/core/distr_2d.h b/include/mitsuba/core/distr_2d.h index 955111f41..ef3995458 100644 --- a/include/mitsuba/core/distr_2d.h +++ b/include/mitsuba/core/distr_2d.h @@ -13,6 +13,7 @@ #include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -72,7 +73,7 @@ NAMESPACE_BEGIN(mitsuba) */ template -class DiscreteDistribution2D { +class DiscreteDistribution2D : drjit::TraversableBase{ public: using Float = Float_; using UInt32 = dr::uint32_array_t; @@ -201,10 +202,14 @@ class DiscreteDistribution2D { Float m_inv_normalization; Float m_normalization; + + DR_TRAVERSE_CB(drjit::TraversableBase, m_data, m_marg_cdf, m_cond_cdf, + m_inv_normalization, m_normalization) }; /// Base class of Hierarchical2D and Marginal2D with common functionality -template class Distribution2D { +template +class Distribution2D : drjit::TraversableBase { public: static constexpr size_t Dimension = Dimension_; using Float = Float_; @@ -308,6 +313,28 @@ template class Distribution2D { /// Total number of slices (in case Dimension > 1) uint32_t m_slices; + +public: + void + traverse_1_cb_ro(void *payload, + drjit::detail::traverse_callback_ro fn) const override { + if constexpr (!std ::is_same_v) + drjit ::TraversableBase ::traverse_1_cb_ro(payload, fn); + for (const auto ¶m_value : m_param_values) { + drjit ::traverse_1_fn_ro(param_value, payload, fn); + } + } + void traverse_1_cb_rw(void *payload, + drjit::detail::traverse_callback_rw fn) override { + if constexpr (!std ::is_same_v) + drjit ::TraversableBase ::traverse_1_cb_rw(payload, fn); + + for (auto ¶m_value : m_param_values) { + drjit ::traverse_1_fn_rw(param_value, payload, fn); + } + } }; /** @@ -788,6 +815,8 @@ class Hierarchical2D : public Distribution2D { return dr::gather(data, i0, active); } } + + DRJIT_STRUCT_NODEF(Level, data) }; /// MIP hierarchy over linearly interpolated patches @@ -795,6 +824,8 @@ class Hierarchical2D : public Distribution2D { /// Number of bilinear patches in the X/Y dimension - 1 ScalarVector2u m_max_patch_index; + + DR_TRAVERSE_CB(Base, m_levels) }; /** @@ -1454,6 +1485,8 @@ class Marginal2D : public Distribution2D { /// Are the probability values normalized? bool m_normalized; + + DR_TRAVERSE_CB(Base, m_data, m_marg_cdf, m_cond_cdf) }; //! @} diff --git a/include/mitsuba/core/field.h b/include/mitsuba/core/field.h index 42db1edaa..2baaadf48 100644 --- a/include/mitsuba/core/field.h +++ b/include/mitsuba/core/field.h @@ -5,6 +5,7 @@ #include #include +#include namespace dr = drjit; NAMESPACE_BEGIN(mitsuba) @@ -62,6 +63,12 @@ struct field @@ -105,6 +112,17 @@ struct field::traverse_1_cb_ro( \ + void *payload, drjit::detail::traverse_callback_ro fn) const { \ + if constexpr (!std ::is_same_v) \ + Base ::traverse_1_cb_ro(payload, fn); \ + DRJIT_MAP(DR_TRAVERSE_MEMBER_RO, __VA_ARGS__) \ + } \ + MI_VARIANT \ + void Type::traverse_1_cb_rw( \ + void *payload, drjit::detail::traverse_callback_rw fn) { \ + if constexpr (!std ::is_same_v) \ + Base ::traverse_1_cb_rw(payload, fn); \ + DRJIT_MAP(DR_TRAVERSE_MEMBER_RW, __VA_ARGS__) \ + } + //! @} // ============================================================= diff --git a/include/mitsuba/core/object.h b/include/mitsuba/core/object.h index 50338b3b2..c22177946 100644 --- a/include/mitsuba/core/object.h +++ b/include/mitsuba/core/object.h @@ -3,6 +3,7 @@ #include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -29,8 +30,10 @@ NAMESPACE_BEGIN(mitsuba) * Python, this counter is shared with Python such that the ownerhsip and * lifetime of any ``Object`` instance across C++ and Python is managed by it. */ -class MI_EXPORT_LIB Object : public nanobind::intrusive_base { +class MI_EXPORT_LIB Object : public drjit::TraversableBase { public: + DR_TRAVERSE_CB(drjit::TraversableBase) + /// Default constructor Object() { } diff --git a/include/mitsuba/render/bsdf.h b/include/mitsuba/render/bsdf.h index f8270a9dc..fb27ecb15 100644 --- a/include/mitsuba/render/bsdf.h +++ b/include/mitsuba/render/bsdf.h @@ -609,6 +609,8 @@ class MI_EXPORT_LIB BSDF : public Object { /// Identifier (if available) std::string m_id; + + DR_TRAVERSE_CB(Object); }; // ----------------------------------------------------------------------- diff --git a/include/mitsuba/render/emitter.h b/include/mitsuba/render/emitter.h index df6370e62..7cc06bbe5 100644 --- a/include/mitsuba/render/emitter.h +++ b/include/mitsuba/render/emitter.h @@ -94,6 +94,8 @@ class MI_EXPORT_LIB Emitter : public Endpoint { /// True if the emitters's parameters have changed bool m_dirty = false; + + DR_TRAVERSE_CB(Base); }; MI_EXTERN_CLASS(Emitter) diff --git a/include/mitsuba/render/endpoint.h b/include/mitsuba/render/endpoint.h index ca74d827a..0f487f0e2 100644 --- a/include/mitsuba/render/endpoint.h +++ b/include/mitsuba/render/endpoint.h @@ -397,6 +397,8 @@ class MI_EXPORT_LIB Endpoint : public Object { bool m_needs_sample_2 = true; bool m_needs_sample_3 = true; std::string m_id; + + MI_DECLARE_TRAVERSE_CB() }; MI_EXTERN_CLASS(Endpoint) diff --git a/include/mitsuba/render/film.h b/include/mitsuba/render/film.h index b58955ecd..9f52b5335 100644 --- a/include/mitsuba/render/film.h +++ b/include/mitsuba/render/film.h @@ -224,6 +224,8 @@ class MI_EXPORT_LIB Film : public Object { bool m_sample_border; ref m_filter; ref m_srf; + + MI_DECLARE_TRAVERSE_CB() }; MI_EXTERN_CLASS(Film) diff --git a/include/mitsuba/render/imageblock.h b/include/mitsuba/render/imageblock.h index 785db9bdf..3f04f8834 100644 --- a/include/mitsuba/render/imageblock.h +++ b/include/mitsuba/render/imageblock.h @@ -346,6 +346,8 @@ class MI_EXPORT_LIB ImageBlock : public Object { bool m_compensate; bool m_warn_negative; bool m_warn_invalid; + + DR_TRAVERSE_CB(Object, m_tensor, m_tensor_compensation) }; MI_EXTERN_CLASS(ImageBlock) diff --git a/include/mitsuba/render/integrator.h b/include/mitsuba/render/integrator.h index debeccd4b..ac35807e8 100644 --- a/include/mitsuba/render/integrator.h +++ b/include/mitsuba/render/integrator.h @@ -343,6 +343,8 @@ class MI_EXPORT_LIB Integrator : public Object { /// Identifier (if available) std::string m_id; + + DR_TRAVERSE_CB(Object) }; /** \brief Abstract integrator that performs Monte Carlo sampling starting from @@ -460,6 +462,8 @@ class MI_EXPORT_LIB SamplingIntegrator : public Integrator { * If set to (uint32_t) -1, all the work is done in a single pass (default). */ uint32_t m_samples_per_pass; + + DR_TRAVERSE_CB(Base) }; /** \brief Abstract integrator that performs *recursive* Monte Carlo sampling @@ -487,6 +491,8 @@ class MI_EXPORT_LIB MonteCarloIntegrator protected: uint32_t m_max_depth; uint32_t m_rr_depth; + + DR_TRAVERSE_CB(Base) }; /** \brief Abstract adjoint integrator that performs Monte Carlo sampling @@ -574,6 +580,8 @@ class MI_EXPORT_LIB AdjointIntegrator : public Integrator { /// Depth to begin using russian roulette int m_rr_depth; + + DR_TRAVERSE_CB(Base) }; diff --git a/include/mitsuba/render/medium.h b/include/mitsuba/render/medium.h index 8aa8450ed..b08667b8e 100644 --- a/include/mitsuba/render/medium.h +++ b/include/mitsuba/render/medium.h @@ -112,6 +112,8 @@ class MI_EXPORT_LIB Medium : public Object { /// Identifier (if available) std::string m_id; + + MI_DECLARE_TRAVERSE_CB() }; MI_EXTERN_CLASS(Medium) diff --git a/include/mitsuba/render/mesh.h b/include/mitsuba/render/mesh.h index 3bb797101..934fe02b1 100644 --- a/include/mitsuba/render/mesh.h +++ b/include/mitsuba/render/mesh.h @@ -10,6 +10,7 @@ #include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -495,6 +496,8 @@ class MI_EXPORT_LIB Mesh : public Shape { MeshAttribute migrate(AllocType at) const { return MeshAttribute { size, type, dr::migrate(buf, at) }; } + + DRJIT_STRUCT_NODEF(MeshAttribute, buf); }; template @@ -591,6 +594,8 @@ class MI_EXPORT_LIB Mesh : public Shape { /// Pointer to the scene that owns this mesh Scene* m_scene = nullptr; + + MI_DECLARE_TRAVERSE_CB() }; MI_EXTERN_CLASS(Mesh) diff --git a/include/mitsuba/render/sampler.h b/include/mitsuba/render/sampler.h index 688f87f4f..275cb169f 100644 --- a/include/mitsuba/render/sampler.h +++ b/include/mitsuba/render/sampler.h @@ -130,11 +130,6 @@ class MI_EXPORT_LIB Sampler : public Object { /// dr::schedule() variables that represent the internal sampler state virtual void schedule_state(); - /// Traversal callback mechanism for symbolic loops - virtual void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const; - /// Traversal callback mechanism for symbolic loops - virtual void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)); - MI_DECLARE_CLASS() protected: @@ -160,6 +155,8 @@ class MI_EXPORT_LIB Sampler : public Object { UInt32 m_dimension_index; /// Index of the current sample in the sequence UInt32 m_sample_index; + + DR_TRAVERSE_CB(Object, m_dimension_index, m_sample_index); }; /// Interface for sampler plugins based on the PCG32 random number generator @@ -172,8 +169,6 @@ class MI_EXPORT_LIB PCG32Sampler : public Sampler { void seed(UInt32 seed, uint32_t wavefront_size = (uint32_t) -1) override; void schedule_state() override; - void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const override; - void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)) override; MI_DECLARE_CLASS() protected: @@ -183,6 +178,8 @@ class MI_EXPORT_LIB PCG32Sampler : public Sampler { PCG32Sampler(const PCG32Sampler &sampler); protected: PCG32 m_rng; + + DR_TRAVERSE_CB(Base, m_rng); }; MI_EXTERN_CLASS(Sampler) diff --git a/include/mitsuba/render/scene.h b/include/mitsuba/render/scene.h index 562a75e72..a8a02afda 100644 --- a/include/mitsuba/render/scene.h +++ b/include/mitsuba/render/scene.h @@ -3,9 +3,9 @@ #include #include #include -#include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -636,6 +636,13 @@ class MI_EXPORT_LIB Scene : public Object { std::unique_ptr> m_silhouette_distr = nullptr; bool m_shapes_grad_enabled; + + void traverse_1_cb_ro_cpu(void *payload, + drjit::detail::traverse_callback_ro fn) const; + void traverse_1_cb_rw_cpu(void *payload, + drjit::detail::traverse_callback_rw fn); + + MI_DECLARE_TRAVERSE_CB() }; /// Dummy function which can be called to ensure that the librender shared library is loaded diff --git a/include/mitsuba/render/sensor.h b/include/mitsuba/render/sensor.h index 3d931f762..2b6f2e52e 100644 --- a/include/mitsuba/render/sensor.h +++ b/include/mitsuba/render/sensor.h @@ -163,6 +163,8 @@ class MI_EXPORT_LIB Sensor : public Endpoint { ScalarFloat m_shutter_open_time; ref m_srf; bool m_alpha; + + DR_TRAVERSE_CB(Base, m_film, m_sampler, m_srf); }; //! @} @@ -216,6 +218,8 @@ class MI_EXPORT_LIB ProjectiveCamera : public Sensor { ScalarFloat m_near_clip; ScalarFloat m_far_clip; Float m_focus_distance; + + DR_TRAVERSE_CB(Base, m_focus_distance); }; // ======================================================================== diff --git a/include/mitsuba/render/shape.h b/include/mitsuba/render/shape.h index ed31ab1e7..b88273356 100644 --- a/include/mitsuba/render/shape.h +++ b/include/mitsuba/render/shape.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -990,6 +991,8 @@ class MI_EXPORT_LIB Shape : public Object { /// True if the shape has called iniatlize() at least once bool m_initialized = false; + + MI_DECLARE_TRAVERSE_CB() }; // ----------------------------------------------------------------------- diff --git a/include/mitsuba/render/shapegroup.h b/include/mitsuba/render/shapegroup.h index 50056c6bc..bec0a4267 100644 --- a/include/mitsuba/render/shapegroup.h +++ b/include/mitsuba/render/shapegroup.h @@ -105,7 +105,11 @@ class MI_EXPORT_LIB ShapeGroup : public Shape { uint32_t m_sbt_offset; #endif + std::vector m_accel_handles; + bool m_has_meshes, m_has_bspline_curves, m_has_linear_curves, m_has_others; + + MI_DECLARE_TRAVERSE_CB() }; MI_EXTERN_CLASS(ShapeGroup) diff --git a/include/mitsuba/render/texture.h b/include/mitsuba/render/texture.h index 48c8bf3dd..4dbef95ef 100644 --- a/include/mitsuba/render/texture.h +++ b/include/mitsuba/render/texture.h @@ -236,6 +236,8 @@ class MI_EXPORT_LIB Texture : public Object { protected: std::string m_id; + + DR_TRAVERSE_CB(Object); }; MI_EXTERN_CLASS(Texture) diff --git a/include/mitsuba/render/volume.h b/include/mitsuba/render/volume.h index 57ff627d5..15585bef5 100644 --- a/include/mitsuba/render/volume.h +++ b/include/mitsuba/render/volume.h @@ -119,6 +119,8 @@ class MI_EXPORT_LIB Volume : public Object { ScalarBoundingBox3f m_bbox; /// Number of channels stored in the volume uint32_t m_channel_count; + + DR_TRAVERSE_CB(Object); }; MI_EXTERN_CLASS(Volume) diff --git a/include/mitsuba/render/volumegrid.h b/include/mitsuba/render/volumegrid.h index e7119abf0..b8e8c005e 100644 --- a/include/mitsuba/render/volumegrid.h +++ b/include/mitsuba/render/volumegrid.h @@ -120,6 +120,8 @@ class MI_EXPORT_LIB VolumeGrid : public Object { ScalarBoundingBox3f m_bbox; ScalarFloat m_max; std::vector m_max_per_channel; + + DR_TRAVERSE_CB(Object) }; MI_EXTERN_CLASS(VolumeGrid) diff --git a/src/bsdfs/blendbsdf.cpp b/src/bsdfs/blendbsdf.cpp index 9dafb5d0b..714179b84 100644 --- a/src/bsdfs/blendbsdf.cpp +++ b/src/bsdfs/blendbsdf.cpp @@ -237,6 +237,8 @@ class BlendBSDF final : public BSDF { protected: ref m_weight; ref m_nested_bsdf[2]; + + DR_TRAVERSE_CB(Base, m_weight, m_nested_bsdf[0], m_nested_bsdf[1]) }; MI_IMPLEMENT_CLASS_VARIANT(BlendBSDF, BSDF) diff --git a/src/bsdfs/bumpmap.cpp b/src/bsdfs/bumpmap.cpp index 39d7ddee5..7a5ac4c90 100644 --- a/src/bsdfs/bumpmap.cpp +++ b/src/bsdfs/bumpmap.cpp @@ -83,6 +83,8 @@ class BumpMap final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture) + DR_TRAVERSE_CB(Base, m_nested_texture, m_nested_bsdf); + BumpMap(const Properties &props) : Base(props) { for (auto &[name, obj] : props.objects(false)) { auto bsdf = dynamic_cast(obj.get()); diff --git a/src/bsdfs/circular.cpp b/src/bsdfs/circular.cpp index 86022e429..f09770991 100644 --- a/src/bsdfs/circular.cpp +++ b/src/bsdfs/circular.cpp @@ -61,6 +61,8 @@ class CircularPolarizer final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture) + DR_TRAVERSE_CB(Base, m_transmittance); + CircularPolarizer(const Properties &props) : Base(props) { m_transmittance = props.texture("transmittance", 1.f); m_left_handed = props.get("left_handed", false); diff --git a/src/bsdfs/conductor.cpp b/src/bsdfs/conductor.cpp index b85e49472..83d2bbb6c 100644 --- a/src/bsdfs/conductor.cpp +++ b/src/bsdfs/conductor.cpp @@ -330,6 +330,8 @@ class SmoothConductor final : public BSDF { private: ref m_specular_reflectance; ref m_eta, m_k; + + DR_TRAVERSE_CB(Base, m_specular_reflectance, m_eta, m_k); }; MI_IMPLEMENT_CLASS_VARIANT(SmoothConductor, BSDF) diff --git a/src/bsdfs/dielectric.cpp b/src/bsdfs/dielectric.cpp index be81ea059..05038b5c3 100644 --- a/src/bsdfs/dielectric.cpp +++ b/src/bsdfs/dielectric.cpp @@ -396,6 +396,9 @@ class SmoothDielectric final : public BSDF { ScalarFloat m_eta; ref m_specular_reflectance; ref m_specular_transmittance; + + DR_TRAVERSE_CB(Base, m_eta, m_specular_reflectance, + m_specular_transmittance); }; MI_IMPLEMENT_CLASS_VARIANT(SmoothDielectric, BSDF) diff --git a/src/bsdfs/diffuse.cpp b/src/bsdfs/diffuse.cpp index bdd8df6e1..f3e41819a 100644 --- a/src/bsdfs/diffuse.cpp +++ b/src/bsdfs/diffuse.cpp @@ -194,6 +194,8 @@ class SmoothDiffuse final : public BSDF { MI_DECLARE_CLASS() private: ref m_reflectance; + + DR_TRAVERSE_CB(Base, m_reflectance); }; MI_IMPLEMENT_CLASS_VARIANT(SmoothDiffuse, BSDF) diff --git a/src/bsdfs/hair.cpp b/src/bsdfs/hair.cpp index 961fc987f..436c15319 100755 --- a/src/bsdfs/hair.cpp +++ b/src/bsdfs/hair.cpp @@ -785,6 +785,12 @@ class Hair final : public BSDF { Float m_v[P_MAX + 1]; /// Longitudinal variance due to roughness Float m_s; /// Azimuthal roughness scaling factor Float m_sin_2k_alpha[3], m_cos_2k_alpha[3]; + + DR_TRAVERSE_CB(Base, m_longitudinal_roughness, m_azimuthal_roughness, + m_alpha, m_eta, m_eumelanin, m_pheomelanin, m_sigma_a, + m_v[0], m_v[1], m_v[2], m_v[3], m_s, m_sin_2k_alpha[0], + m_sin_2k_alpha[1], m_sin_2k_alpha[2], m_cos_2k_alpha[0], + m_cos_2k_alpha[1], m_cos_2k_alpha[2]) }; MI_IMPLEMENT_CLASS_VARIANT(Hair, BSDF) diff --git a/src/bsdfs/mask.cpp b/src/bsdfs/mask.cpp index 91e9d4050..4819a8e1b 100644 --- a/src/bsdfs/mask.cpp +++ b/src/bsdfs/mask.cpp @@ -246,6 +246,8 @@ class MaskBSDF final : public BSDF { private: ref m_opacity; ref m_nested_bsdf; + + DR_TRAVERSE_CB(Base, m_opacity, m_nested_bsdf); }; MI_IMPLEMENT_CLASS_VARIANT(MaskBSDF, BSDF) diff --git a/src/bsdfs/measured.cpp b/src/bsdfs/measured.cpp index 4a638bfb9..c4bb62fca 100644 --- a/src/bsdfs/measured.cpp +++ b/src/bsdfs/measured.cpp @@ -74,6 +74,8 @@ class Measured final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES() + DR_TRAVERSE_CB(Base, m_ndf, m_sigma, m_vndf, m_luminance, m_spectra); + using Warp2D0 = Marginal2D; using Warp2D2 = Marginal2D; using Warp2D3 = Marginal2D; diff --git a/src/bsdfs/measured_polarized.cpp b/src/bsdfs/measured_polarized.cpp index 887444c6d..840e471b6 100644 --- a/src/bsdfs/measured_polarized.cpp +++ b/src/bsdfs/measured_polarized.cpp @@ -108,6 +108,8 @@ class MeasuredPolarized final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture, MicrofacetDistribution) + DR_TRAVERSE_CB(Base, m_interpolator); + using Interpolator = Marginal2D; MeasuredPolarized(const Properties &props) : Base(props) { diff --git a/src/bsdfs/normalmap.cpp b/src/bsdfs/normalmap.cpp index 2e9c6ac88..23364d7bd 100644 --- a/src/bsdfs/normalmap.cpp +++ b/src/bsdfs/normalmap.cpp @@ -81,6 +81,8 @@ class NormalMap final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture) + DR_TRAVERSE_CB(Base, m_nested_bsdf, m_normalmap); + NormalMap(const Properties &props) : Base(props) { for (auto &[name, obj] : props.objects(false)) { auto bsdf = dynamic_cast(obj.get()); diff --git a/src/bsdfs/plastic.cpp b/src/bsdfs/plastic.cpp index 530eeed0a..9366172b6 100644 --- a/src/bsdfs/plastic.cpp +++ b/src/bsdfs/plastic.cpp @@ -390,6 +390,10 @@ class SmoothPlastic final : public BSDF { ScalarFloat m_fdr_ext; Float m_specular_sampling_weight; bool m_nonlinear; + + DR_TRAVERSE_CB(Base, m_diffuse_reflectance, m_specular_reflectance, + m_specular_sampling_weight); + }; MI_IMPLEMENT_CLASS_VARIANT(SmoothPlastic, BSDF) diff --git a/src/bsdfs/polarizer.cpp b/src/bsdfs/polarizer.cpp index c4cb0ba3b..0573efdef 100644 --- a/src/bsdfs/polarizer.cpp +++ b/src/bsdfs/polarizer.cpp @@ -81,6 +81,8 @@ class LinearPolarizer final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture) + DR_TRAVERSE_CB(Base, m_theta, m_transmittance); + LinearPolarizer(const Properties &props) : Base(props) { m_theta = props.texture("theta", 0.f); m_transmittance = props.texture("transmittance", 1.f); diff --git a/src/bsdfs/pplastic.cpp b/src/bsdfs/pplastic.cpp index b4ea9bc7c..e5d1cd8de 100644 --- a/src/bsdfs/pplastic.cpp +++ b/src/bsdfs/pplastic.cpp @@ -149,6 +149,9 @@ class PolarizedPlastic final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture, MicrofacetDistribution) + DR_TRAVERSE_CB(Base, m_diffuse_reflectance, m_specular_reflectance, + m_alpha_u, m_alpha_v, m_eta, m_specular_sampling_weight); + PolarizedPlastic(const Properties &props) : Base(props) { m_diffuse_reflectance = props.texture("diffuse_reflectance", .5f); diff --git a/src/bsdfs/principled.cpp b/src/bsdfs/principled.cpp index e4ef4e409..36531839d 100644 --- a/src/bsdfs/principled.cpp +++ b/src/bsdfs/principled.cpp @@ -184,6 +184,12 @@ class Principled final : public BSDF { public: MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture, MicrofacetDistribution) + + DR_TRAVERSE_CB(Base, m_base_color, m_roughness, m_anisotropic, m_sheen, + m_sheen_tint, m_spec_trans, m_flatness, m_spec_tint, + m_clearcoat, m_clearcoat_gloss, m_metallic, m_eta, + m_specular); + using GTR1 = GTR1Isotropic; Principled(const Properties &props) : Base(props) { diff --git a/src/bsdfs/principledthin.cpp b/src/bsdfs/principledthin.cpp index 8b3e1ad98..a4beb3574 100644 --- a/src/bsdfs/principledthin.cpp +++ b/src/bsdfs/principledthin.cpp @@ -159,6 +159,10 @@ class PrincipledThin final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture, MicrofacetDistribution) + DR_TRAVERSE_CB(Base, m_base_color, m_roughness, m_anisotropic, m_sheen, + m_sheen_tint, m_spec_trans, m_flatness, m_spec_tint, + m_diff_trans, m_eta_thin); + PrincipledThin(const Properties &props) : Base(props) { m_base_color = props.texture("base_color", 0.5f); diff --git a/src/bsdfs/retarder.cpp b/src/bsdfs/retarder.cpp index 8aba0d36e..a3ce39dd0 100644 --- a/src/bsdfs/retarder.cpp +++ b/src/bsdfs/retarder.cpp @@ -72,6 +72,8 @@ class LinearRetarder final : public BSDF { MI_IMPORT_BASE(BSDF, m_flags, m_components) MI_IMPORT_TYPES(Texture) + DR_TRAVERSE_CB(Base, m_theta, m_delta, m_transmittance); + LinearRetarder(const Properties &props) : Base(props) { m_theta = props.texture("theta", 0.f); // As default, instantiate as a quarter-wave plate diff --git a/src/bsdfs/roughconductor.cpp b/src/bsdfs/roughconductor.cpp index 24de2cc67..e8c793e9e 100755 --- a/src/bsdfs/roughconductor.cpp +++ b/src/bsdfs/roughconductor.cpp @@ -544,6 +544,9 @@ class RoughConductor final : public BSDF { ref m_k; /// Specular reflectance component ref m_specular_reflectance; + + DR_TRAVERSE_CB(Base, m_alpha_u, m_alpha_v, m_eta, m_k, + m_specular_reflectance); }; MI_IMPLEMENT_CLASS_VARIANT(RoughConductor, BSDF) diff --git a/src/bsdfs/roughdielectric.cpp b/src/bsdfs/roughdielectric.cpp index fa916e83f..23d843c39 100644 --- a/src/bsdfs/roughdielectric.cpp +++ b/src/bsdfs/roughdielectric.cpp @@ -636,6 +636,9 @@ class RoughDielectric final : public BSDF { ref m_alpha_u, m_alpha_v; Float m_eta, m_inv_eta; bool m_sample_visible; + + DR_TRAVERSE_CB(Base, m_specular_reflectance, m_specular_transmittance, + m_alpha_u, m_alpha_v, m_eta, m_inv_eta); }; MI_IMPLEMENT_CLASS_VARIANT(RoughDielectric, BSDF) diff --git a/src/bsdfs/roughplastic.cpp b/src/bsdfs/roughplastic.cpp index 0a207d91b..9a0a0d5d7 100644 --- a/src/bsdfs/roughplastic.cpp +++ b/src/bsdfs/roughplastic.cpp @@ -536,6 +536,10 @@ class RoughPlastic final : public BSDF { bool m_sample_visible; DynamicBuffer m_external_transmittance; Float m_internal_reflectance; + + DR_TRAVERSE_CB(Base, m_diffuse_reflectance, m_specular_reflectance, m_eta, + m_inv_eta_2, m_alpha, m_specular_sampling_weight, + m_external_transmittance, m_internal_reflectance); }; MI_IMPLEMENT_CLASS_VARIANT(RoughPlastic, BSDF); diff --git a/src/bsdfs/thindielectric.cpp b/src/bsdfs/thindielectric.cpp index a3738700a..4d7a12daf 100644 --- a/src/bsdfs/thindielectric.cpp +++ b/src/bsdfs/thindielectric.cpp @@ -230,6 +230,8 @@ class ThinDielectric final : public BSDF { Float m_eta; ref m_specular_transmittance; ref m_specular_reflectance; + + DR_TRAVERSE_CB(Base, m_specular_reflectance, m_specular_transmittance); }; MI_IMPLEMENT_CLASS_VARIANT(ThinDielectric, BSDF) diff --git a/src/bsdfs/twosided.cpp b/src/bsdfs/twosided.cpp index 22f07893e..256abc35f 100644 --- a/src/bsdfs/twosided.cpp +++ b/src/bsdfs/twosided.cpp @@ -376,6 +376,8 @@ class TwoSidedBRDF final : public BSDF { MI_DECLARE_CLASS() protected: ref m_brdf[2]; + + DR_TRAVERSE_CB(Base, m_brdf[0], m_brdf[1]); }; MI_IMPLEMENT_CLASS_VARIANT(TwoSidedBRDF, BSDF) diff --git a/src/core/python/object.cpp b/src/core/python/object.cpp index 315ca596e..ad3632443 100644 --- a/src/core/python/object.cpp +++ b/src/core/python/object.cpp @@ -4,6 +4,7 @@ #include #include #include +#include extern nb::object cast_object(Object *o); @@ -40,7 +41,7 @@ MI_PY_EXPORT(Object) { return cast_object(pmgr.create_object(props, class_)); }, D(PluginManager, create_object)); - nb::class_( + auto obj = nb::class_( m, "Object", nb::intrusive_ptr( [](Object *o, PyObject *po) noexcept { o->set_self_py(po); }), diff --git a/src/emitters/area.cpp b/src/emitters/area.cpp index 247c14213..003904352 100644 --- a/src/emitters/area.cpp +++ b/src/emitters/area.cpp @@ -5,6 +5,7 @@ #include #include #include +#include NAMESPACE_BEGIN(mitsuba) @@ -269,6 +270,8 @@ class AreaLight final : public Emitter { MI_DECLARE_CLASS() private: ref m_radiance; + + DR_TRAVERSE_CB(Base, m_radiance); }; MI_IMPLEMENT_CLASS_VARIANT(AreaLight, Emitter) diff --git a/src/emitters/constant.cpp b/src/emitters/constant.cpp index 7100fcc12..0affb8abf 100644 --- a/src/emitters/constant.cpp +++ b/src/emitters/constant.cpp @@ -204,6 +204,8 @@ class ConstantBackgroundEmitter final : public Emitter { /// Surface area of the bounding sphere Float m_surface_area; + + DR_TRAVERSE_CB(Base, m_radiance, m_bsphere, m_surface_area) }; MI_IMPLEMENT_CLASS_VARIANT(ConstantBackgroundEmitter, Emitter) diff --git a/src/emitters/directional.cpp b/src/emitters/directional.cpp index 8c46e58b0..444deeda5 100644 --- a/src/emitters/directional.cpp +++ b/src/emitters/directional.cpp @@ -231,6 +231,8 @@ MI_VARIANT class DirectionalEmitter final : public Emitter { protected: ref m_irradiance; ScalarBoundingSphere3f m_bsphere; + + DR_TRAVERSE_CB(Base, m_irradiance); }; MI_IMPLEMENT_CLASS_VARIANT(DirectionalEmitter, Emitter) diff --git a/src/emitters/directionalarea.cpp b/src/emitters/directionalarea.cpp index 356d5df5e..4688b90ee 100644 --- a/src/emitters/directionalarea.cpp +++ b/src/emitters/directionalarea.cpp @@ -190,6 +190,8 @@ class DirectionalArea final : public Emitter { private: ref m_radiance; Float m_area = 0.f; + + DR_TRAVERSE_CB(Base, m_radiance, m_area) }; MI_IMPLEMENT_CLASS_VARIANT(DirectionalArea, Emitter) diff --git a/src/emitters/envmap.cpp b/src/emitters/envmap.cpp index ef89a9ccf..0b334fa93 100644 --- a/src/emitters/envmap.cpp +++ b/src/emitters/envmap.cpp @@ -594,6 +594,8 @@ class EnvironmentMapEmitter final : public Emitter { Warp m_warp; ref m_d65; Float m_scale; + + DR_TRAVERSE_CB(Base, m_bsphere, m_data, m_warp, m_d65, m_scale) }; MI_IMPLEMENT_CLASS_VARIANT(EnvironmentMapEmitter, Emitter) diff --git a/src/emitters/point.cpp b/src/emitters/point.cpp index aa5254248..f550e35bf 100644 --- a/src/emitters/point.cpp +++ b/src/emitters/point.cpp @@ -206,6 +206,8 @@ class PointLight final : public Emitter { private: ref m_intensity; field m_position; + + DR_TRAVERSE_CB(Base, m_intensity, m_position); }; diff --git a/src/emitters/projector.cpp b/src/emitters/projector.cpp index 0c72b3914..eeb093641 100644 --- a/src/emitters/projector.cpp +++ b/src/emitters/projector.cpp @@ -319,6 +319,9 @@ MI_VARIANT class Projector final : public Emitter { Transform4f m_sample_to_camera; ScalarFloat m_x_fov; Float m_sensor_area; + + DR_TRAVERSE_CB(Base, m_irradiance, m_intensity_scale, m_camera_to_sample, + m_sample_to_camera, m_x_fov, m_sensor_area) }; MI_IMPLEMENT_CLASS_VARIANT(Projector, Emitter) diff --git a/src/emitters/spot.cpp b/src/emitters/spot.cpp index 7bb48fa1a..7d47293de 100644 --- a/src/emitters/spot.cpp +++ b/src/emitters/spot.cpp @@ -300,8 +300,11 @@ class SpotLight final : public Emitter { ref m_texture; Float m_beam_width, m_cutoff_angle, m_uv_factor; Float m_cos_beam_width, m_cos_cutoff_angle, m_inv_transition_width; -}; + DR_TRAVERSE_CB(Base, m_intensity, m_texture, m_beam_width, m_cutoff_angle, + m_uv_factor, m_cos_beam_width, m_cos_cutoff_angle, + m_inv_transition_width) +}; MI_IMPLEMENT_CLASS_VARIANT(SpotLight, Emitter) MI_EXPORT_PLUGIN(SpotLight, "Spot emitter") diff --git a/src/films/hdrfilm.cpp b/src/films/hdrfilm.cpp index de9a4757d..d84ee47b8 100644 --- a/src/films/hdrfilm.cpp +++ b/src/films/hdrfilm.cpp @@ -613,6 +613,8 @@ class HDRFilm final : public Film { ref m_storage; mutable std::mutex m_mutex; std::vector m_channels; + + DR_TRAVERSE_CB(Base, m_storage); }; MI_IMPLEMENT_CLASS_VARIANT(HDRFilm, Film) diff --git a/src/integrators/direct.cpp b/src/integrators/direct.cpp index 16466b554..7d42a4567 100644 --- a/src/integrators/direct.cpp +++ b/src/integrators/direct.cpp @@ -229,6 +229,8 @@ class DirectIntegrator : public SamplingIntegrator { size_t m_bsdf_samples; ScalarFloat m_frac_bsdf, m_frac_lum; ScalarFloat m_weight_bsdf, m_weight_lum; + + DR_TRAVERSE_CB(Base, m_frac_bsdf, m_frac_lum, m_weight_bsdf, m_weight_lum) }; MI_IMPLEMENT_CLASS_VARIANT(DirectIntegrator, SamplingIntegrator) diff --git a/src/integrators/moment.cpp b/src/integrators/moment.cpp index 165da6de2..945ba4ef7 100644 --- a/src/integrators/moment.cpp +++ b/src/integrators/moment.cpp @@ -141,6 +141,8 @@ class MomentIntegrator final : public SamplingIntegrator { private: std::vector m_aov_names; std::vector, size_t>> m_integrators; + + DR_TRAVERSE_CB(Base, m_integrators); }; MI_IMPLEMENT_CLASS_VARIANT(MomentIntegrator, SamplingIntegrator) diff --git a/src/media/heterogeneous.cpp b/src/media/heterogeneous.cpp index db51e5650..22c610524 100644 --- a/src/media/heterogeneous.cpp +++ b/src/media/heterogeneous.cpp @@ -216,6 +216,8 @@ class HeterogeneousMedium final : public Medium { ScalarFloat m_scale; Float m_max_density; + + DR_TRAVERSE_CB(Base, m_sigmat, m_albedo, m_max_density); }; MI_IMPLEMENT_CLASS_VARIANT(HeterogeneousMedium, Medium) diff --git a/src/media/homogeneous.cpp b/src/media/homogeneous.cpp index a039b2e03..c278f4e78 100644 --- a/src/media/homogeneous.cpp +++ b/src/media/homogeneous.cpp @@ -194,6 +194,8 @@ class HomogeneousMedium final : public Medium { private: ref m_sigmat, m_albedo; ScalarFloat m_scale; + + DR_TRAVERSE_CB(Base, m_sigmat, m_albedo, m_scale); }; MI_IMPLEMENT_CLASS_VARIANT(HomogeneousMedium, Medium) diff --git a/src/python/python/ad/integrators/prb.py b/src/python/python/ad/integrators/prb.py index 79653bba9..5b18f9c88 100644 --- a/src/python/python/ad/integrators/prb.py +++ b/src/python/python/ad/integrators/prb.py @@ -98,9 +98,13 @@ def sample(self, prev_bsdf_pdf = mi.Float(1.0) prev_bsdf_delta = mi.Bool(True) - while dr.hint(active, - max_iterations=self.max_depth, - label="Path Replay Backpropagation (%s)" % mode.name): + while dr.hint( + active, + max_iterations=self.max_depth, + label="Path Replay Backpropagation (%s)" % mode.name, + # include=[ sampler, ray, depth, L, δL, β, active ], + exclude=[ scene, ], + ): active_next = mi.Bool(active) # Compute a surface interaction that tracks derivatives arising diff --git a/src/python/python/ad/integrators/prb_basic.py b/src/python/python/ad/integrators/prb_basic.py index 7058d102d..3f17a5b3b 100644 --- a/src/python/python/ad/integrators/prb_basic.py +++ b/src/python/python/ad/integrators/prb_basic.py @@ -81,7 +81,8 @@ def sample(self, while dr.hint(active, max_iterations=self.max_depth, - label="Path Replay Backpropagation (%s)" % mode.name): + label="Path Replay Backpropagation (%s)" % mode.name, + exclude = [scene]): active_next = mi.Bool(active) # ---------------------- Direct emission ---------------------- @@ -162,7 +163,7 @@ def sample(self, return ( L if primal else δL, # Radiance/differential radiance - depth != 0, # Ray validity flag for alpha blending + depth != 0, # Ray validity flag for alpha blending [], # Empty typle of AOVs L # State the for differential phase ) diff --git a/src/python/python/ad/integrators/prb_projective.py b/src/python/python/ad/integrators/prb_projective.py index 70e92eaf0..a0246bce4 100644 --- a/src/python/python/ad/integrators/prb_projective.py +++ b/src/python/python/ad/integrators/prb_projective.py @@ -176,7 +176,9 @@ def sample(self, while dr.hint(active, max_iterations=self.max_depth, - label="PRB Projective (%s)" % mode.name): + label="PRB Projective (%s)" % mode.name, + exclude = [scene], + ): active_next = mi.Bool(active) # Compute a surface interaction that tracks derivatives arising @@ -480,7 +482,8 @@ def sample_importance(self, scene, sensor, ss, max_depth, sampler, preprocess, a bsdf_ctx = mi.BSDFContext(mi.TransportMode.Importance) while dr.hint(active_loop, - label="Estimate Importance"): + label="Estimate Importance", + exclude = [scene]): # Is it possible to connect the current vertex to the sensor? bsdf = si_loop.bsdf() active_connect = active_loop & mi.has_flag(bsdf.flags(), mi.BSDFFlags.Smooth) diff --git a/src/python/python/ad/optimizers.py b/src/python/python/ad/optimizers.py index c8bcdae3b..51ed3f6fa 100644 --- a/src/python/python/ad/optimizers.py +++ b/src/python/python/ad/optimizers.py @@ -15,8 +15,8 @@ def __init__(self, lr, params: dict): Parameter ``params`` (:py:class:`dict`): Dictionary-like object containing parameters to optimize. """ - self.lr = defaultdict(lambda: self.lr_default) - self.lr_v = defaultdict(lambda: self.lr_default_v) + self.lr = {} + self.lr_v = {} self.set_learning_rate(lr) self.variables = {} @@ -128,7 +128,14 @@ class SGD(Optimizer): :math:`\\varepsilon` is the learning rate, and :math:`\\mu` is the momentum parameter. """ - def __init__(self, lr, momentum=0, mask_updates=False, params:dict=None): + + DRJIT_STRUCT = { + "state": dict, + "variables": dict, + "lr_default_v": None, + } + + def __init__(self, lr, momentum=0, mask_updates=False, params: dict = None): """ Parameter ``lr``: learning rate @@ -154,6 +161,12 @@ def __init__(self, lr, momentum=0, mask_updates=False, params:dict=None): def step(self): """Take a gradient step""" for k, p in self.variables.items(): + # Ensure that self.lr and self.lr_v contain key + if k not in self.lr: + self.lr[k] = self.lr_default + if k not in self.lr_v: + self.lr_v[k] = self.lr_default_v + g_p = dr.grad(p) shape = dr.shape(g_p) if shape == 0: @@ -224,6 +237,13 @@ class Adam(Optimizer): Enabling ``mask_updates`` avoids these two issues. This is similar to `PyTorch's SparseAdam optimizer `_. """ + DRJIT_STRUCT = { + "state": dict, + "variables": dict, + "lr_default_v": None, + "lr_v": dict, + "t": dict, + } def __init__(self, lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8, mask_updates=False, uniform=False, params: dict=None): """ @@ -257,15 +277,23 @@ def __init__(self, lr, beta_1=0.9, beta_2=0.999, epsilon=1e-8, self.epsilon = epsilon self.mask_updates = mask_updates self.uniform = uniform - self.t = defaultdict(lambda: 0) + self.t = {} super().__init__(lr, params) def step(self): """Take a gradient step""" for k, p in self.variables.items(): + # Ensure that self.t, self.lr and self.lr_v contain key + if k not in self.t: + self.t[k] = mi.UInt(0) + if k not in self.lr: + self.lr[k] = self.lr_default + if k not in self.lr_v: + self.lr_v[k] = self.lr_default_v + self.t[k] += 1 lr_scale = dr.sqrt(1 - self.beta_2 ** self.t[k]) / (1 - self.beta_1 ** self.t[k]) - lr_scale = dr.opaque(dr.detached_t(mi.Float), lr_scale, shape=1) + dr.make_opaque(lr_scale) lr_t = self.lr_v[k] * lr_scale g_p = dr.grad(p) @@ -321,7 +349,7 @@ def reset(self, key): shape = dr.shape(p) if dr.is_tensor_v(p) else dr.width(p) self.state[key] = (dr.zeros(dr.detached_t(p), shape), dr.zeros(dr.detached_t(p), shape)) - self.t[key] = 0 + self.t[key] = mi.UInt(0) def __repr__(self): return ('Adam[\n' diff --git a/src/python/python/util.py b/src/python/python/util.py index b26359afe..208752c9c 100644 --- a/src/python/python/util.py +++ b/src/python/python/util.py @@ -354,7 +354,7 @@ def eval(self, scene, sensor, _, params, integrator, seed, spp): self.spp = spp with dr.suspend_grad(): - return self.integrator.render( + res = self.integrator.render( scene=self.scene, sensor=sensor, seed=seed[0], @@ -362,6 +362,8 @@ def eval(self, scene, sensor, _, params, integrator, seed, spp): develop=True, evaluate=False ) + sensor.sampler().seed(0, 1) + return res def forward(self): self.set_grad_out( diff --git a/src/render/endpoint.cpp b/src/render/endpoint.cpp index fbcae3120..a33bb6a12 100644 --- a/src/render/endpoint.cpp +++ b/src/render/endpoint.cpp @@ -118,6 +118,7 @@ MI_VARIANT void Endpoint::parameters_changed(const std::vector< } } +MI_IMPLEMENT_TRAVERSE_CB(Endpoint, Object, m_to_world, m_medium) MI_IMPLEMENT_CLASS_VARIANT(Endpoint, Object) MI_INSTANTIATE_CLASS(Endpoint) NAMESPACE_END(mitsuba) diff --git a/src/render/film.cpp b/src/render/film.cpp index af6101204..450921466 100644 --- a/src/render/film.cpp +++ b/src/render/film.cpp @@ -117,7 +117,7 @@ MI_VARIANT std::string Film::to_string() const { return oss.str(); } - +MI_IMPLEMENT_TRAVERSE_CB(Film, Object, m_srf) MI_IMPLEMENT_CLASS_VARIANT(Film, Object, "film") MI_INSTANTIATE_CLASS(Film) NAMESPACE_END(mitsuba) diff --git a/src/render/medium.cpp b/src/render/medium.cpp index 6a633334e..2ae91c22a 100644 --- a/src/render/medium.cpp +++ b/src/render/medium.cpp @@ -102,6 +102,7 @@ Medium::transmittance_eval_pdf(const MediumInteraction3f &mi, return { tr, pdf }; } +MI_IMPLEMENT_TRAVERSE_CB(Medium, Object, m_phase_function) MI_IMPLEMENT_CLASS_VARIANT(Medium, Object, "medium") MI_INSTANTIATE_CLASS(Medium) NAMESPACE_END(mitsuba) diff --git a/src/render/mesh.cpp b/src/render/mesh.cpp index 1519ceebe..03548fb86 100644 --- a/src/render/mesh.cpp +++ b/src/render/mesh.cpp @@ -1858,6 +1858,9 @@ MI_VARIANT bool Mesh::parameters_grad_enabled() const { return dr::grad_enabled(m_vertex_positions); } +MI_IMPLEMENT_TRAVERSE_CB(Mesh, Base, m_vertex_positions, m_vertex_normals, + m_vertex_texcoords, m_faces, m_E2E, m_sil_dedge_pmf, + m_mesh_attributes, m_area_pmf, m_parameterization) MI_IMPLEMENT_CLASS_VARIANT(Mesh, Shape) MI_INSTANTIATE_CLASS(Mesh) NAMESPACE_END(mitsuba) diff --git a/src/render/python/bsdf_v.cpp b/src/render/python/bsdf_v.cpp index 7ff5cb44f..4bd7d7985 100644 --- a/src/render/python/bsdf_v.cpp +++ b/src/render/python/bsdf_v.cpp @@ -8,6 +8,7 @@ #include #include #include +#include MI_PY_EXPORT(BSDFSample) { MI_PY_IMPORT_TYPES() @@ -105,6 +106,8 @@ MI_VARIANT class PyBSDF : public BSDF { using BSDF::m_flags; using BSDF::m_components; + + DR_TRAMPOLINE_TRAVERSE_CB(BSDF); }; template void bind_bsdf_generic(Cls &cls) { @@ -190,6 +193,8 @@ MI_PY_EXPORT(BSDF) { .def_field(PyBSDF, m_components, D(BSDF, m_components)) .def("__repr__", &BSDF::to_string, D(BSDF, to_string)); + drjit::bind_traverse(bsdf); + bind_bsdf_generic(bsdf); if constexpr (dr::is_array_v) { diff --git a/src/render/python/emitter_v.cpp b/src/render/python/emitter_v.cpp index 9ca26c4db..94810833b 100644 --- a/src/render/python/emitter_v.cpp +++ b/src/render/python/emitter_v.cpp @@ -84,6 +84,8 @@ MI_VARIANT class PyEmitter : public Emitter { using Emitter::m_flags; using Emitter::m_needs_sample_2; using Emitter::m_needs_sample_3; + + DR_TRAMPOLINE_TRAVERSE_CB(Emitter); }; template void bind_emitter_generic(Cls &cls) { @@ -164,6 +166,10 @@ MI_PY_EXPORT(Emitter) { .def_field(PyEmitter, m_needs_sample_3, D(Endpoint, m_needs_sample_3)) .def_field(PyEmitter, m_flags, D(Emitter, m_flags)); + drjit::bind_traverse(emitter); + + bind_emitter_generic(emitter); + if constexpr (dr::is_array_v) { dr::ArrayBinding b; auto emitter_ptr = dr::bind_array_t(b, m, "EmitterPtr"); diff --git a/src/render/python/endpoint_v.cpp b/src/render/python/endpoint_v.cpp index b0a5ff5d6..0bad8779a 100644 --- a/src/render/python/endpoint_v.cpp +++ b/src/render/python/endpoint_v.cpp @@ -8,10 +8,11 @@ #include #include #include +#include MI_PY_EXPORT(Endpoint) { MI_PY_IMPORT_TYPES() - MI_PY_CLASS(Endpoint, Object) + auto endpoint = MI_PY_CLASS(Endpoint, Object) .def_method(Endpoint, sample_ray, "time"_a, "sample1"_a, "sample2"_a, "sample3"_a, "active"_a = true) .def_method(Endpoint, sample_direction, "it"_a, "sample"_a, "active"_a = true) .def_method(Endpoint, pdf_direction, "it"_a, "ds"_a, "active"_a = true) @@ -29,4 +30,6 @@ MI_PY_EXPORT(Endpoint) { .def_method(Endpoint, set_medium, "medium"_a) .def_method(Endpoint, set_scene, "scene"_a) .def_method(Endpoint, bbox); + + drjit::bind_traverse(endpoint); } diff --git a/src/render/python/film_v.cpp b/src/render/python/film_v.cpp index a1f03ee6e..85ebdcc18 100644 --- a/src/render/python/film_v.cpp +++ b/src/render/python/film_v.cpp @@ -11,6 +11,7 @@ #include #include #include +#include /// Trampoline for derived types implemented in Python MI_VARIANT class PyFilm : public Film { @@ -76,6 +77,8 @@ MI_VARIANT class PyFilm : public Film { using Film::m_sample_border; using Film::m_filter; using Film::m_srf; + + DR_TRAMPOLINE_TRAVERSE_CB(Film) }; MI_PY_EXPORT(Film) { @@ -127,5 +130,7 @@ MI_PY_EXPORT(Film) { .def_method(Film, flags) .def_field(PyFilm, m_flags, D(Film, m_flags)); + drjit::bind_traverse(film); + MI_PY_REGISTER_OBJECT("register_film", Film) } diff --git a/src/render/python/imageblock_v.cpp b/src/render/python/imageblock_v.cpp index e4d5842cf..4b97b0ea7 100644 --- a/src/render/python/imageblock_v.cpp +++ b/src/render/python/imageblock_v.cpp @@ -2,10 +2,11 @@ #include #include #include +#include MI_PY_EXPORT(ImageBlock) { MI_PY_IMPORT_TYPES(ImageBlock, ReconstructionFilter) - MI_PY_CLASS(ImageBlock, Object) + auto image_block = MI_PY_CLASS(ImageBlock, Object) .def(nb::init(), @@ -67,4 +68,6 @@ MI_PY_EXPORT(ImageBlock) { .def("tensor", nb::overload_cast<>(&ImageBlock::tensor), nb::rv_policy::reference_internal, D(ImageBlock, tensor)); + + drjit::bind_traverse(image_block); } diff --git a/src/render/python/integrator_v.cpp b/src/render/python/integrator_v.cpp index 024e4ee25..96c4f7e4b 100644 --- a/src/render/python/integrator_v.cpp +++ b/src/render/python/integrator_v.cpp @@ -7,6 +7,8 @@ #include #include #include "signal.h" +#include +#include #if defined(__APPLE__) || defined(__linux__) # define MI_HANDLE_SIGINT 1 @@ -117,6 +119,8 @@ MI_VARIANT class PySamplingIntegrator : public SamplingIntegrator std::string to_string() const override { NB_OVERRIDE(to_string); } + + DR_TRAMPOLINE_TRAVERSE_CB(AdjointIntegrator) }; /** @@ -267,6 +273,8 @@ MI_VARIANT class PyADIntegrator : public CppADIntegrator { } using Base::m_hide_emitters; + + DR_TRAMPOLINE_TRAVERSE_CB(Base) }; MI_PY_EXPORT(Integrator) { @@ -277,7 +285,7 @@ MI_PY_EXPORT(Integrator) { using PyADIntegrator = PyADIntegrator; using Properties = PropertiesV; - MI_PY_CLASS(Integrator, Object) + auto cls = MI_PY_CLASS(Integrator, Object) .def( "render", [&](Integrator *integrator, Scene *scene, Sensor *sensor, @@ -305,7 +313,9 @@ MI_PY_EXPORT(Integrator) { .def_method(Integrator, should_stop) .def_method(Integrator, aov_names); - MI_PY_TRAMPOLINE_CLASS(PySamplingIntegrator, SamplingIntegrator, Integrator) + drjit::bind_traverse(cls); + + auto sampling_integrator = MI_PY_TRAMPOLINE_CLASS(PySamplingIntegrator, SamplingIntegrator, Integrator) .def(nb::init()) .def( "sample", @@ -364,15 +374,19 @@ MI_PY_EXPORT(Integrator) { "spp"_a = 0) .def_rw("hide_emitters", &PySamplingIntegrator::m_hide_emitters); + drjit::bind_traverse(sampling_integrator); + MI_PY_REGISTER_OBJECT("register_integrator", Integrator) MI_PY_CLASS(MonteCarloIntegrator, SamplingIntegrator); - nb::class_( + auto cpp_ad_integrator = nb::class_( m, "CppADIntegrator") .def(nb::init()); - MI_PY_TRAMPOLINE_CLASS(PyAdjointIntegrator, AdjointIntegrator, Integrator) + drjit::bind_traverse(cpp_ad_integrator); + + auto adjoint_integrator = MI_PY_TRAMPOLINE_CLASS(PyAdjointIntegrator, AdjointIntegrator, Integrator) .def(nb::init()) .def( "render_forward", @@ -418,4 +432,6 @@ MI_PY_EXPORT(Integrator) { "spp"_a = 0) .def_method(AdjointIntegrator, sample, "scene"_a, "sensor"_a, "sampler"_a, "block"_a, "sample_scale"_a); + + drjit::bind_traverse(adjoint_integrator); } diff --git a/src/render/python/interaction_v.cpp b/src/render/python/interaction_v.cpp index 7f68c324d..9af1e7685 100644 --- a/src/render/python/interaction_v.cpp +++ b/src/render/python/interaction_v.cpp @@ -6,6 +6,7 @@ #include #include #include +#include MI_PY_EXPORT(Interaction) { MI_PY_IMPORT_TYPES() diff --git a/src/render/python/medium_v.cpp b/src/render/python/medium_v.cpp index 184925483..d9d8981f9 100644 --- a/src/render/python/medium_v.cpp +++ b/src/render/python/medium_v.cpp @@ -45,6 +45,8 @@ MI_VARIANT class PyMedium : public Medium { using Medium::m_sample_emitters; using Medium::m_is_homogeneous; using Medium::m_has_spectral_extinction; + + DR_TRAMPOLINE_TRAVERSE_CB(Medium) }; template void bind_medium_generic(Cls &cls) { @@ -106,6 +108,8 @@ MI_PY_EXPORT(Medium) { .def_field(PyMedium, m_has_spectral_extinction, D(Medium, m_has_spectral_extinction)) .def("__repr__", &Medium::to_string, D(Medium, to_string)); + drjit::bind_traverse(medium); + bind_medium_generic(medium); if constexpr (dr::is_array_v) { diff --git a/src/render/python/phase_v.cpp b/src/render/python/phase_v.cpp index fd05b11a8..14612ea95 100644 --- a/src/render/python/phase_v.cpp +++ b/src/render/python/phase_v.cpp @@ -55,6 +55,8 @@ MI_VARIANT class PyPhaseFunction : public PhaseFunction { using PhaseFunction::m_flags; using PhaseFunction::m_components; + + DR_TRAMPOLINE_TRAVERSE_CB(PhaseFunction) }; template void bind_phase_generic(Cls &cls) { @@ -113,6 +115,8 @@ MI_PY_EXPORT(PhaseFunction) { .def_field(PyPhaseFunction, m_flags, D(PhaseFunction, m_flags)) .def("__repr__", &PhaseFunction::to_string); + drjit::bind_traverse(phase); + bind_phase_generic(phase); if constexpr (dr::is_array_v) { diff --git a/src/render/python/sampler_v.cpp b/src/render/python/sampler_v.cpp index 1d763ffd5..9f2bcc80e 100644 --- a/src/render/python/sampler_v.cpp +++ b/src/render/python/sampler_v.cpp @@ -45,6 +45,8 @@ MI_VARIANT class PySampler : public Sampler { std::string to_string() const override { NB_OVERRIDE(to_string); } + + DR_TRAMPOLINE_TRAVERSE_CB(Sampler); }; MI_PY_EXPORT(Sampler) { diff --git a/src/render/python/scene_v.cpp b/src/render/python/scene_v.cpp index 10188d910..db28e84a0 100644 --- a/src/render/python/scene_v.cpp +++ b/src/render/python/scene_v.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #if !defined(MI_ENABLE_EMBREE) # include @@ -47,7 +48,7 @@ MI_PY_EXPORT(ShapeKDTree) { MI_PY_EXPORT(Scene) { MI_PY_IMPORT_TYPES(Scene, Integrator, SamplingIntegrator, MonteCarloIntegrator, Sensor) - MI_PY_CLASS(Scene, Object) + auto scene = MI_PY_CLASS(Scene, Object) .def(nb::init()) .def("ray_intersect_preliminary", nb::overload_cast(&Scene::ray_intersect_preliminary, nb::const_), @@ -152,4 +153,5 @@ MI_PY_EXPORT(Scene) { D(Scene, integrator)) .def_method(Scene, shapes_grad_enabled) .def("__repr__", &Scene::to_string); + dr::bind_traverse(scene); } diff --git a/src/render/python/sensor_v.cpp b/src/render/python/sensor_v.cpp index b2d217521..9520c5153 100644 --- a/src/render/python/sensor_v.cpp +++ b/src/render/python/sensor_v.cpp @@ -1,5 +1,6 @@ #include // Needs to be first, to get `ref` caster #include +#include // Has to be included, so that Emitter::CallSupport::Variant is instantiated here #include #include #include @@ -91,6 +92,8 @@ MI_VARIANT class PySensor : public Sensor { using Sensor::m_needs_sample_2; using Sensor::m_needs_sample_3; using Sensor::m_film; + + DR_TRAMPOLINE_TRAVERSE_CB(Sensor) }; template void bind_sensor_generic(Cls &cls) { @@ -175,6 +178,8 @@ MI_PY_EXPORT(Sensor) { .def_field(PySensor, m_needs_sample_3, D(Endpoint, m_needs_sample_3)) .def_field(PySensor, m_film); + drjit::bind_traverse(sensor); + bind_sensor_generic(sensor); if constexpr (dr::is_array_v) { @@ -193,8 +198,9 @@ MI_PY_EXPORT(Sensor) { m.def("perspective_projection", &perspective_projection, "film_size"_a, "crop_size"_a, "crop_offset"_a, "fov_x"_a, "near_clip"_a, "far_clip"_a, D(perspective_projection)); - m.def("orthographic_projection", &orthographic_projection, "film_size"_a, "crop_size"_a, "crop_offset"_a, "near_clip"_a, "far_clip"_a, D(orthographic_projection)); + + dr::bind_traverse(sensor); } diff --git a/src/render/python/shape_v.cpp b/src/render/python/shape_v.cpp index e7a00056c..41e85cf24 100644 --- a/src/render/python/shape_v.cpp +++ b/src/render/python/shape_v.cpp @@ -60,6 +60,8 @@ MI_VARIANT class PyMesh : public Mesh { std::string to_string() const override { NB_OVERRIDE(to_string); } + + DR_TRAMPOLINE_TRAVERSE_CB(Mesh) }; template void bind_shape_generic(Cls &cls) { @@ -305,6 +307,8 @@ MI_PY_EXPORT(Shape) { .def_method(Shape, effective_primitive_count) .def_method(Shape, precompute_silhouette, "viewpoint"_a); + drjit::bind_traverse(shape); + bind_shape_generic(shape); if constexpr (dr::is_array_v) { @@ -356,5 +360,7 @@ MI_PY_EXPORT(Shape) { bind_mesh_generic(mesh_ptr); } + drjit::bind_traverse(mesh_cls); + MI_PY_REGISTER_OBJECT("register_mesh", Mesh) } diff --git a/src/render/python/texture_v.cpp b/src/render/python/texture_v.cpp index 02226a2d7..8ac001cae 100644 --- a/src/render/python/texture_v.cpp +++ b/src/render/python/texture_v.cpp @@ -6,6 +6,7 @@ #include #include #include +#include /// Trampoline for derived types implemented in Python MI_VARIANT class PyTexture : public Texture { @@ -80,6 +81,8 @@ MI_VARIANT class PyTexture : public Texture { std::string to_string() const override { NB_OVERRIDE(to_string); } + + DR_TRAMPOLINE_TRAVERSE_CB(Texture) }; MI_PY_EXPORT(Texture) { @@ -87,7 +90,7 @@ MI_PY_EXPORT(Texture) { using PyTexture = PyTexture; using Properties = PropertiesV; - MI_PY_TRAMPOLINE_CLASS(PyTexture, Texture, Object) + auto texture = MI_PY_TRAMPOLINE_CLASS(PyTexture, Texture, Object) .def(nb::init(), "props"_a) .def_static("D65", nb::overload_cast(&Texture::D65), "scale"_a = 1.f) .def_method(Texture, mean, D(Texture, mean)) @@ -106,5 +109,7 @@ MI_PY_EXPORT(Texture) { .def_method(Texture, spectral_resolution) .def_method(Texture, wavelength_range); + drjit::bind_traverse(texture); + MI_PY_REGISTER_OBJECT("register_texture", Texture) } diff --git a/src/render/python/volume_v.cpp b/src/render/python/volume_v.cpp index b9629402a..b2815d562 100644 --- a/src/render/python/volume_v.cpp +++ b/src/render/python/volume_v.cpp @@ -8,6 +8,7 @@ #include #include #include +#include /// Trampoline for derived types implemented in Python MI_VARIANT class PyVolume : public Volume { @@ -54,6 +55,8 @@ MI_VARIANT class PyVolume : public Volume { std::string to_string() const override { NB_OVERRIDE(to_string); } + + DR_TRAMPOLINE_TRAVERSE_CB(Volume) }; MI_PY_EXPORT(Volume) { @@ -61,7 +64,7 @@ MI_PY_EXPORT(Volume) { using PyVolume = PyVolume; using Properties = PropertiesV; - MI_PY_TRAMPOLINE_CLASS(PyVolume, Volume, Object) + auto volume = MI_PY_TRAMPOLINE_CLASS(PyVolume, Volume, Object) .def(nb::init(), "props"_a) .def_method(Volume, resolution) .def_method(Volume, bbox) @@ -96,5 +99,7 @@ MI_PY_EXPORT(Volume) { "it"_a, "active"_a = true, D(Volume, eval_n)); + drjit::bind_traverse(volume); + MI_PY_REGISTER_OBJECT("register_volume", Volume) } diff --git a/src/render/python/volumegrid_v.cpp b/src/render/python/volumegrid_v.cpp index ee63c5d2b..c158904c5 100644 --- a/src/render/python/volumegrid_v.cpp +++ b/src/render/python/volumegrid_v.cpp @@ -6,6 +6,7 @@ #include #include #include +#include MI_PY_EXPORT(VolumeGrid) { MI_PY_IMPORT_TYPES(VolumeGrid) @@ -44,7 +45,7 @@ MI_PY_EXPORT(VolumeGrid) { volumegrid->set_max_per_channel(max_per_channel.data()); }; - MI_PY_CLASS(VolumeGrid, Object) + auto volume_grid = MI_PY_CLASS(VolumeGrid, Object) .def(nb::init(), "path"_a, nb::call_guard()) .def(nb::init(), "stream"_a, @@ -118,4 +119,6 @@ MI_PY_EXPORT(VolumeGrid) { result["version"] = 3; return nb::object(result); }); + + drjit::bind_traverse(volume_grid); } diff --git a/src/render/sampler.cpp b/src/render/sampler.cpp index 8efa97ecc..a912dbb0b 100644 --- a/src/render/sampler.cpp +++ b/src/render/sampler.cpp @@ -68,18 +68,6 @@ MI_VARIANT void Sampler::schedule_state() { dr::schedule(m_sample_index, m_dimension_index); } -MI_VARIANT void -Sampler::traverse_1_cb_ro(void * /*payload*/, - void (* /*fn*/)(void *, uint64_t)) const { - NotImplementedError("traverse_1_cb_ro"); -} - -MI_VARIANT void -Sampler::traverse_1_cb_rw(void * /*payload*/, - uint64_t (* /*fn*/)(void *, uint64_t)) { - NotImplementedError("traverse_1_cb_rw"); -} - MI_VARIANT void Sampler::set_samples_per_wavefront(uint32_t samples_per_wavefront) { if constexpr (!dr::is_array_v) @@ -148,18 +136,6 @@ MI_VARIANT void PCG32Sampler::schedule_state() { dr::schedule(m_rng.inc, m_rng.state); } -MI_VARIANT void -PCG32Sampler::traverse_1_cb_ro(void *payload, - void (*fn)(void *, uint64_t)) const { - traverse_1_fn_ro(m_rng, payload, fn); -} - -MI_VARIANT void -PCG32Sampler::traverse_1_cb_rw(void *payload, - uint64_t (*fn)(void *, uint64_t)) { - traverse_1_fn_rw(m_rng, payload, fn); -} - MI_VARIANT PCG32Sampler::PCG32Sampler(const PCG32Sampler &sampler) : Base(sampler) { diff --git a/src/render/scene.cpp b/src/render/scene.cpp index 6229574d5..fb51df057 100644 --- a/src/render/scene.cpp +++ b/src/render/scene.cpp @@ -599,6 +599,37 @@ MI_VARIANT void Scene::static_accel_shutdown_gpu() { } #endif Class *__kdtree_class = new Class("TShapeKDTree", "Object", "", nullptr, nullptr); +MI_VARIANT +void Scene::traverse_1_cb_ro( + void *payload, drjit::detail::traverse_callback_ro fn) const { + if constexpr (!std::is_same_v) + Object::traverse_1_cb_ro(payload, fn); + DRJIT_MAP(DR_TRAVERSE_MEMBER_RO, m_accel_handle, m_emitters, m_emitters_dr, + m_shapes, m_shapes_dr, m_shapegroups, m_sensors, m_sensors_dr, + m_children, m_integrator, m_environment, m_emitter_pmf, + m_emitter_distr, m_silhouette_shapes) + if constexpr (dr::is_cuda_v) { + // Nothing to traverse for now + } else { + traverse_1_cb_ro_cpu(payload, fn); + } +} + +MI_VARIANT +void Scene::traverse_1_cb_rw( + void *payload, drjit::detail::traverse_callback_rw fn) { + if constexpr (!std::is_same_v) + Object::traverse_1_cb_rw(payload, fn); + DRJIT_MAP(DR_TRAVERSE_MEMBER_RW, m_accel_handle, m_emitters, m_emitters_dr, + m_shapes, m_shapes_dr, m_shapegroups, m_sensors, m_sensors_dr, + m_children, m_integrator, m_environment, m_emitter_pmf, + m_emitter_distr, m_silhouette_shapes) + if constexpr (dr::is_cuda_v) { + // Nothing to traverse for now + } else { + traverse_1_cb_rw_cpu(payload, fn); + } +} MI_IMPLEMENT_CLASS_VARIANT(Scene, Object, "scene") MI_INSTANTIATE_CLASS(Scene) diff --git a/src/render/scene_embree.inl b/src/render/scene_embree.inl index fab523628..9a993f10d 100644 --- a/src/render/scene_embree.inl +++ b/src/render/scene_embree.inl @@ -16,6 +16,8 @@ struct EmbreeState { std::vector geometries; DynamicBuffer shapes_registry_ids; bool is_nested_scene = false; + void *func_ptr = nullptr; + UInt64 func_handle; }; static void embree_error_callback(void * /*user_ptr */, RTCError code, const char *str) { @@ -176,7 +178,8 @@ MI_VARIANT void Scene::accel_parameters_changed_cpu() { // Prevents the IAS to be released when updating the scene parameters if (m_accel_handle.index()) jit_var_set_callback(m_accel_handle.index(), nullptr, nullptr); - m_accel_handle = dr::opaque(s.accel); + m_accel_handle = UInt64::map_(s.accel, 1, false); + // m_accel_handle = dr::opaque(s.accel); jit_var_set_callback( m_accel_handle.index(), [](uint32_t /* index */, int free, void *payload) { @@ -200,6 +203,27 @@ MI_VARIANT void Scene::accel_parameters_changed_cpu() { }, (void *) m_accel ); + + // To support frozen functions the func_ptr has to exist as a variable + // when the scene is traversed. + // Since vector width should not change, so we determine the + // intersection function here. + uint32_t jit_width = jit_llvm_vector_width(); + void *func_ptr = nullptr; + switch (jit_width) { + case 1: func_ptr = (void *) rtcIntersect1; break; + case 4: func_ptr = (void *) rtcIntersect4; break; + case 8: func_ptr = (void *) rtcIntersect8; break; + case 16: func_ptr = (void *) rtcIntersect16; break; + case 32: func_ptr = (void *) rtcIntersect32; break; + default: + Throw("accel_init_cpu(): Dr.Jit is configured for vectors of " + "width %u, which is not supported by Embree!", + jit_width); + } + + s.func_ptr = func_ptr; + s.func_handle = UInt64::map_(func_ptr, 1, false); } clear_shapes_dirty(); @@ -281,24 +305,12 @@ Scene::ray_intersect_preliminary_cpu(const Ray3f &ray, uint32_t jit_width = jit_llvm_vector_width(); void *scene_ptr = (void *) s.accel, - *func_ptr = nullptr; - - switch (jit_width) { - case 1: func_ptr = (void *) rtcIntersect1; break; - case 4: func_ptr = (void *) rtcIntersect4; break; - case 8: func_ptr = (void *) rtcIntersect8; break; - case 16: func_ptr = (void *) rtcIntersect16; break; - case 32: func_ptr = (void *) rtcIntersect32; break; - default: - Throw("ray_intersect_preliminary_cpu(): Dr.Jit is " - "configured for vectors of width %u, which is not " - "supported by Embree!", jit_width); - } + *func_ptr = s.func_ptr; UInt64 func_v = UInt64::steal(jit_var_pointer( - JitBackend::LLVM, func_ptr, m_accel_handle.index(), 0)), + JitBackend::LLVM, func_ptr, s.func_handle.index(), 0)), scene_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, scene_ptr, 0, 0)); + jit_var_pointer(JitBackend::LLVM, scene_ptr, m_accel_handle.index(), 0)); UInt32 zero = dr::zeros(); @@ -399,24 +411,13 @@ Scene::ray_test_cpu(const Ray3f &ray, Mask coherent, Mask activ uint32_t jit_width = jit_llvm_vector_width(); void *scene_ptr = (void *) s.accel, - *func_ptr = nullptr; - - switch (jit_width) { - case 1: func_ptr = (void *) rtcOccluded1; break; - case 4: func_ptr = (void *) rtcOccluded4; break; - case 8: func_ptr = (void *) rtcOccluded8; break; - case 16: func_ptr = (void *) rtcOccluded16; break; - case 32: func_ptr = (void *) rtcOccluded32; break; - default: - Throw("ray_test_cpu(): Dr.Jit is configured for vectors of " - "width %u, which is not supported by Embree!", jit_width); - } + *func_ptr = s.func_ptr; UInt64 func_v = UInt64::steal(jit_var_pointer( - JitBackend::LLVM, func_ptr, m_accel_handle.index(), 0)), + JitBackend::LLVM, func_ptr, s.func_handle.index(), 0)), scene_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, scene_ptr, 0, 0)); - + jit_var_pointer(JitBackend::LLVM, scene_ptr, m_accel_handle.index(), 0)); + UInt32 zero = dr::zeros(); // Conversion, in case this is a double precision build @@ -450,4 +451,20 @@ Scene::ray_intersect_naive_cpu(const Ray3f &ray, return ray_intersect_cpu(ray, +RayFlags::All, false, active); } +MI_VARIANT void Scene::traverse_1_cb_ro_cpu( + void *payload, drjit::detail::traverse_callback_ro fn) const { + + EmbreeState &s = *(EmbreeState *) m_accel; + drjit ::traverse_1_fn_ro(s.shapes_registry_ids, payload, fn); + drjit ::traverse_1_fn_ro(s.func_handle, payload, fn); +} + +MI_VARIANT void Scene::traverse_1_cb_rw_cpu( + void *payload, drjit::detail::traverse_callback_rw fn) { + + EmbreeState &s = *(EmbreeState *) m_accel; + drjit ::traverse_1_fn_rw(s.shapes_registry_ids, payload, fn); + drjit ::traverse_1_fn_rw(s.func_handle, payload, fn); +} + NAMESPACE_END(mitsuba) diff --git a/src/render/scene_native.inl b/src/render/scene_native.inl index 2d52ea35d..f0c47468f 100644 --- a/src/render/scene_native.inl +++ b/src/render/scene_native.inl @@ -1,3 +1,4 @@ +#include NAMESPACE_BEGIN(mitsuba) template @@ -5,6 +6,8 @@ struct NativeState { MI_IMPORT_CORE_TYPES() ShapeKDTree *accel; DynamicBuffer shapes_registry_ids; + void *func_ptr = nullptr; + UInt64 func_handle; }; MI_VARIANT void Scene::accel_init_cpu(const Properties &props) { @@ -33,6 +36,10 @@ MI_VARIANT void Scene::accel_init_cpu(const Properties &props) accel_parameters_changed_cpu(); } +template +void kdtree_trace_func_wrapper(const int *valid, void *ptr, + void* /* context */, uint8_t *args); + MI_VARIANT void Scene::accel_parameters_changed_cpu() { // Ensure all ray tracing kernels are terminated before releasing the scene if constexpr (dr::is_llvm_v) @@ -59,7 +66,8 @@ MI_VARIANT void Scene::accel_parameters_changed_cpu() { // Prevents the IAS to be released when updating the scene parameters if (m_accel_handle.index()) jit_var_set_callback(m_accel_handle.index(), nullptr, nullptr); - m_accel_handle = dr::opaque(kdtree); + // m_accel_handle = dr::opaque(kdtree); + m_accel_handle = UInt64::map_(m_accel, 1, false); jit_var_set_callback( m_accel_handle.index(), [](uint32_t /* index */, int free, void *payload) { @@ -78,6 +86,27 @@ MI_VARIANT void Scene::accel_parameters_changed_cpu() { }, (void *) m_accel ); + + NativeState &s = *(NativeState *) m_accel; + // To support frozen functions the func_ptr has to exist as a variable + // when the scene is traversed. + // Since vector width should not change, so we determine the + // intersection function here. + int jit_width = jit_llvm_vector_width(); + void *func_ptr = nullptr; + switch (jit_width) { + case 1: func_ptr = (void *) kdtree_trace_func_wrapper; break; + case 4: func_ptr = (void *) kdtree_trace_func_wrapper; break; + case 8: func_ptr = (void *) kdtree_trace_func_wrapper; break; + case 16: func_ptr = (void *) kdtree_trace_func_wrapper; break; + default: + Throw("ray_intersect_preliminary_cpu(): Dr.Jit is " + "configured for vectors of width %u, which is not " + "supported by the kd-tree ray tracing backend!", jit_width); + } + + s.func_ptr = func_ptr; + s.func_handle = UInt64::map_(func_ptr, 1, false); } clear_shapes_dirty(); @@ -181,26 +210,14 @@ Scene::ray_intersect_preliminary_cpu(const Ray3f &ray, const ShapeKDTree *kdtree = (const ShapeKDTree *) m_accel; return kdtree->template ray_intersect_preliminary(ray, active); } else { - NativeState *s = (NativeState *) m_accel; - void *func_ptr = nullptr, + NativeState &s = *(NativeState *) m_accel; + void *func_ptr = s.func_ptr, *scene_ptr = m_accel; - int jit_width = jit_llvm_vector_width(); - switch (jit_width) { - case 1: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 4: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 8: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 16: func_ptr = (void *) kdtree_trace_func_wrapper; break; - default: - Throw("ray_intersect_preliminary_cpu(): Dr.Jit is " - "configured for vectors of width %u, which is not " - "supported by the kd-tree ray tracing backend!", jit_width); - } - UInt64 func_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, func_ptr, m_accel_handle.index(), 0)), + jit_var_pointer(JitBackend::LLVM, func_ptr, s.func_handle.index(), 0)), scene_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, scene_ptr, 0, 0)); + jit_var_pointer(JitBackend::LLVM, scene_ptr, m_accel_handle.index(), 0)); UInt32 zero = dr::zeros(); Float ray_mint = dr::zeros(); @@ -236,7 +253,7 @@ Scene::ray_intersect_preliminary_cpu(const Ray3f &ray, Mask hit_inst = hit && (inst_index != ((uint32_t)-1)); UInt32 index = dr::select(hit_inst, inst_index, pi.shape_index); - ShapePtr shape = dr::gather(s->shapes_registry_ids, index, hit); + ShapePtr shape = dr::gather(s.shapes_registry_ids, index, hit); pi.instance = shape & hit_inst; pi.shape = shape & !hit_inst; @@ -269,24 +286,14 @@ Scene::ray_test_cpu(const Ray3f &ray, const ShapeKDTree *kdtree = (const ShapeKDTree *) m_accel; return kdtree->template ray_intersect_preliminary(ray, active).is_valid(); } else { - void *func_ptr = nullptr, *scene_ptr = m_accel; - - int jit_width = jit_llvm_vector_width(); - switch (jit_width) { - case 1: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 4: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 8: func_ptr = (void *) kdtree_trace_func_wrapper; break; - case 16: func_ptr = (void *) kdtree_trace_func_wrapper; break; - default: - Throw("ray_test_cpu(): Dr.Jit is configured for vectors of " - "width %u, which is not supported by the kd-tree ray " - "tracing backend!", jit_width); - } + NativeState &s = *(NativeState *) m_accel; + void *func_ptr = s.func_ptr, + *scene_ptr = m_accel; UInt64 func_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, func_ptr, m_accel_handle.index(), 0)), + jit_var_pointer(JitBackend::LLVM, func_ptr, s.func_handle.index(), 0)), scene_v = UInt64::steal( - jit_var_pointer(JitBackend::LLVM, scene_ptr, 0, 0)); + jit_var_pointer(JitBackend::LLVM, scene_ptr, m_accel_handle.index(), 0)); UInt32 zero = dr::zeros(); Float ray_mint = dr::zeros(); @@ -320,4 +327,26 @@ Scene::ray_intersect_naive_cpu(const Ray3f &ray, Mask active) c return pi.compute_surface_interaction(ray, +RayFlags::All, active); } +MI_VARIANT void Scene::traverse_1_cb_ro_cpu( + void *payload, drjit::detail::traverse_callback_ro fn) const { + + if constexpr (dr::is_llvm_v) { + NativeState &s = + *(NativeState *) m_accel; + drjit ::traverse_1_fn_ro(s.shapes_registry_ids, payload, fn); + drjit ::traverse_1_fn_ro(s.func_handle, payload, fn); + } +} + +MI_VARIANT void Scene::traverse_1_cb_rw_cpu( + void *payload, drjit::detail::traverse_callback_rw fn) { + + if constexpr (dr::is_llvm_v) { + NativeState &s = + *(NativeState *) m_accel; + drjit ::traverse_1_fn_rw(s.shapes_registry_ids, payload, fn); + drjit ::traverse_1_fn_rw(s.func_handle, payload, fn); + } +} + NAMESPACE_END(mitsuba) diff --git a/src/render/shape.cpp b/src/render/shape.cpp index 2a7734a35..96fa9f023 100644 --- a/src/render/shape.cpp +++ b/src/render/shape.cpp @@ -662,6 +662,10 @@ MI_VARIANT std::string Shape::get_children_string() const { return oss.str(); } +MI_IMPLEMENT_TRAVERSE_CB(Shape, Object, m_bsdf, m_emitter, m_sensor, + m_interior_medium, m_exterior_medium, + m_texture_attributes, m_to_world, m_to_object); + MI_IMPLEMENT_CLASS_VARIANT(Shape, Object, "shape") MI_INSTANTIATE_CLASS(Shape) NAMESPACE_END(mitsuba) diff --git a/src/render/shapegroup.cpp b/src/render/shapegroup.cpp index a288e72f5..77ada5a45 100644 --- a/src/render/shapegroup.cpp +++ b/src/render/shapegroup.cpp @@ -178,6 +178,12 @@ MI_VARIANT void ShapeGroup::optix_build_gas(const OptixDeviceCo build_gas(context, m_shapes, m_accel); for (auto &s : m_shapes) s->m_dirty = false; + + m_accel_handles.clear(); + m_accel_handles.push_back(dr::opaque(m_accel.meshes.handle)); + m_accel_handles.push_back(dr::opaque(m_accel.bspline_curves.handle)); + m_accel_handles.push_back(dr::opaque(m_accel.linear_curves.handle)); + m_accel_handles.push_back(dr::opaque(m_accel.custom_shapes.handle)); } } #endif @@ -256,6 +262,8 @@ MI_VARIANT std::string ShapeGroup::to_string() const { return oss.str(); } +MI_IMPLEMENT_TRAVERSE_CB(ShapeGroup, Base, m_shapes, m_shapes_registry_ids, + m_accel_handles) MI_IMPLEMENT_CLASS_VARIANT(ShapeGroup, Shape) MI_INSTANTIATE_CLASS(ShapeGroup) NAMESPACE_END(mitsuba) diff --git a/src/render/tests/test_freeze.py b/src/render/tests/test_freeze.py new file mode 100644 index 000000000..af7b1a03c --- /dev/null +++ b/src/render/tests/test_freeze.py @@ -0,0 +1,1170 @@ +import os +from typing import Any, Callable, List +import pytest +import drjit as dr +import mitsuba as mi +import glob +import gc +import numpy as np +from os.path import join, realpath, dirname, basename, splitext, exists + +from mitsuba.scalar_rgb.test.util import find_resource + +# dr.set_log_level(dr.LogLevel.Trace) +# dr.set_flag(dr.JitFlag.ReuseIndices, False) + + +def test01_cornell_box(variants_vec_rgb): + w = 16 + h = 16 + + n = 5 + + k = "light.emitter.radiance.value" + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def load_scene() -> mi.Scene: + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + scene = mi.load_dict(scene, parallel=False) + return scene + + def run( + scene: mi.Scene, n: int, func: Callable[[mi.Scene, Any], mi.TensorXf] + ) -> List[mi.TensorXf]: + params = mi.traverse(scene) + value = mi.Float(params[k].x) + + images = [] + for i in range(n): + params[k].x = value + 10.0 * i + params.update() + + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + scene = load_scene() + images_ref = run(scene, n, func) + scene2 = load_scene() + images_frozen = run(scene2, n, frozen) + + assert frozen.n_recordings < n + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +def test02_cornell_box_native(variants_vec_rgb): + if mi.MI_ENABLE_EMBREE: + pytest.skip("EMBREE enabled") + + w = 16 + h = 16 + + n = 5 + + k = "light.emitter.radiance.value" + + def func(scene: mi.Scene, x) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def run(n: int, func: Callable[[mi.Scene, Any], mi.TensorXf]) -> List[mi.TensorXf]: + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + scene = mi.load_dict(scene, parallel=False) + + params = mi.traverse(scene) + value = mi.Float(params[k].x) + + images = [] + for i in range(n): + params[k].x = value + 10.0 * i + params.update() + + img = func(scene, params[k].x) + dr.eval(img) + + images.append(img) + + return images + + images_ref = run(n, func) + images_frozen = run(n, frozen) + + assert frozen.n_recordings < n + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +tutorials_dir = realpath(join(dirname(__file__), "../../../tutorials")) + + +@pytest.mark.parametrize( + "integrator", + [ + "direct", + "prb", + # "prb_basic", + "direct_projective", + "prb_projective", + ], +) +def test02_pose_estimation(variants_vec_rgb, integrator): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + # dr.set_flag(dr.JitFlag.LaunchBlocking, True) + # dr.set_flag(dr.JitFlag.OptimizeCalls, False) + w = 16 + h = 16 + + def apply_transformation(initial_vertex_positions, opt, params): + opt["trans"] = dr.clip(opt["trans"], -0.5, 0.5) + opt["angle"] = dr.clip(opt["angle"], -0.5, 0.5) + + trafo = ( + mi.Transform4f() + .translate([opt["trans"].x, opt["trans"].y, 0.0]) + .rotate([0, 1, 0], opt["angle"] * 100.0) + ) + + params["bunny.vertex_positions"] = dr.ravel(trafo @ initial_vertex_positions) + + def mse(image, image_ref): + return dr.sum(dr.square(image - image_ref), axis=None) + + def optimize(scene, ref, initial_vertex_positions, other): + params = mi.traverse(scene) + + image = mi.render(scene, params, spp=1, seed=1, seed_grad=2) + + # Evaluate the objective function from the current rendered image + loss = mse(image, ref) + + # Backpropagate through the rendering process + dr.backward(loss) + + return image, loss + + frozen = dr.freeze(optimize) + + def load_scene(): + from mitsuba.scalar_rgb import Transform4f as T + + scene = mi.cornell_box() + del scene["large-box"] + del scene["small-box"] + del scene["green-wall"] + del scene["red-wall"] + del scene["floor"] + del scene["ceiling"] + scene["bunny"] = { + "type": "ply", + "filename": f"{tutorials_dir}/scenes/meshes/bunny.ply", + "to_world": T().scale(6.5), + "bsdf": { + "type": "diffuse", + "reflectance": {"type": "rgb", "value": (0.3, 0.3, 0.75)}, + }, + } + scene["integrator"] = { + "type": "prb", + } + scene["sensor"]["film"] = { + "type": "hdrfilm", + "width": w, + "height": h, + "rfilter": {"type": "gaussian"}, + "sample_border": True, + } + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run(scene: mi.Scene, optimize, n) -> tuple[mi.TensorXf, mi.Point3f, mi.Float]: + params = mi.traverse(scene) + + params.keep("bunny.vertex_positions") + initial_vertex_positions = dr.unravel( + mi.Point3f, params["bunny.vertex_positions"] + ) + + image_ref = mi.render(scene, spp=4) + + opt = mi.ad.Adam(lr=0.025) + opt["angle"] = mi.Float(0.25) + opt["trans"] = mi.Point3f(0.1, -0.25, 0.0) + + for i in range(n): + params = mi.traverse(scene) + params.keep("bunny.vertex_positions") + + apply_transformation(initial_vertex_positions, opt, params) + + with dr.profile_range("optimize"): + image, loss = optimize( + scene, + image_ref, + initial_vertex_positions, + [ + params["bunny.vertex_positions"], + ], + ) + + opt.step() + + image_final = mi.render(scene, spp=4, seed=1, seed_grad=2) + + return image_final, opt["trans"], opt["angle"] + + n = 10 + + # NOTE: + # In this case, we have to use the same scene object + # for the frozen and non-frozen case, as re-loading + # the scene causes mitsuba to render different images, + # leading to diverging descent traijectories. + + scene = load_scene() + params = mi.traverse(scene) + initial_vertex_positions = mi.Float(params["bunny.vertex_positions"]) + + img_ref, trans_ref, angle_ref = run(scene, optimize, n) + + # Reset parameters: + params["bunny.vertex_positions"] = initial_vertex_positions + params.update() + + img_frozen, trans_frozen, angle_frozen = run(scene, frozen, n) + + # NOTE: cannot compare results as errors accumulate and the result will never be the same. + + assert dr.allclose(trans_ref, trans_frozen) + assert dr.allclose(angle_ref, angle_frozen) + assert frozen.n_recordings < n + if integrator != "prb_projective": + assert dr.allclose(img_ref, img_frozen, atol=1e-4) + + +def test03_optimize_color(variants_vec_rgb): + k = "red.reflectance.value" + w = 16 + h = 16 + n = 10 + + def mse(image, image_ref): + return dr.sum(dr.square(image - image_ref), axis=None) + + def optimize(scene, image_ref): + params = mi.traverse(scene) + params.keep(k) + + image = mi.render(scene, params, spp=1) + + loss = mse(image, image_ref) + + dr.backward(loss) + + return image, loss + + frozen = dr.freeze(optimize) + + def run(n: int, optimize): + scene = mi.cornell_box() + scene["integrator"] = { + "type": "prb", + } + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + scene = mi.load_dict(scene, parallel=False) + + image_ref = mi.render(scene, spp=512) + + params = mi.traverse(scene) + params.keep(k) + + opt = mi.ad.Adam(lr=0.05) + opt[k] = mi.Color3f(0.01, 0.2, 0.9) + + for i in range(n): + params = mi.traverse(scene) + params.keep(k) + params.update(opt) + + image, loss = optimize(scene, image_ref) + + opt.step() + + return image, opt[k] + + image_ref, param_ref = run(n, optimize) + + image_frozen, param_frozen = run(n, frozen) + + assert frozen.n_recordings < n + # Optimizing the reflectance is not as prone to divergence, + # therefore we can test if the two methods produce the same results + assert dr.allclose(param_ref, param_frozen) + + +@pytest.mark.parametrize( + "bsdf", + [ + "diffuse", + "dielectric", + "thindielectric", + "roughdielectric", + "conductor", + "roughconductor", + "hair", + "plastic", + "roughplastic", + "bumpmap", + "normalmap", + "blendbsdf", + "mask", + "twosided", + "principled", + "principledthin", + "custom", + ], +) +def test04_bsdf(variants_vec_rgb, bsdf): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + + if bsdf == "custom": + + class MyBSDF(mi.BSDF): + def __init__(self, props): + mi.BSDF.__init__(self, props) + + # Read 'eta' and 'tint' properties from `props` + self.eta = 1.33 + if props.has_property("eta"): + self.eta = props["eta"] + + self.tint = props["tint"] + + # Set the BSDF flags + reflection_flags = ( + mi.BSDFFlags.DeltaReflection + | mi.BSDFFlags.FrontSide + | mi.BSDFFlags.BackSide + ) + transmission_flags = ( + mi.BSDFFlags.DeltaTransmission + | mi.BSDFFlags.FrontSide + | mi.BSDFFlags.BackSide + ) + self.m_components = [reflection_flags, transmission_flags] + self.m_flags = reflection_flags | transmission_flags + + def sample(self, ctx, si, sample1, sample2, active): + # Compute Fresnel terms + cos_theta_i = mi.Frame3f.cos_theta(si.wi) + r_i, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, self.eta) + t_i = dr.maximum(1.0 - r_i, 0.0) + + # Pick between reflection and transmission + selected_r = (sample1 <= r_i) & active + + # Fill up the BSDFSample struct + bs = mi.BSDFSample3f() + bs.pdf = dr.select(selected_r, r_i, t_i) + bs.sampled_component = dr.select(selected_r, mi.UInt32(0), mi.UInt32(1)) + bs.sampled_type = dr.select( + selected_r, + mi.UInt32(+mi.BSDFFlags.DeltaReflection), + mi.UInt32(+mi.BSDFFlags.DeltaTransmission), + ) + bs.wo = dr.select( + selected_r, + mi.reflect(si.wi), + mi.refract(si.wi, cos_theta_t, eta_ti), + ) + bs.eta = dr.select(selected_r, 1.0, eta_it) + + # For reflection, tint based on the incident angle (more tint at grazing angle) + value_r = dr.lerp( + mi.Color3f(self.tint), + mi.Color3f(1.0), + dr.clip(cos_theta_i, 0.0, 1.0), + ) + + # For transmission, radiance must be scaled to account for the solid angle compression + value_t = mi.Color3f(1.0) * dr.square(eta_ti) + + value = dr.select(selected_r, value_r, value_t) + + return (bs, value) + + def eval(self, ctx, si, wo, active): + return 0.0 + + def pdf(self, ctx, si, wo, active): + return 0.0 + + def eval_pdf(self, ctx, si, wo, active): + return 0.0, 0.0 + + def traverse(self, callback): + callback.put_parameter("tint", self.tint, mi.ParamFlags.Differentiable) + + def parameters_changed(self, keys): + print("🏝️ there is nothing to do here 🏝️") + + def to_string(self): + return ( + "MyBSDF[\n" + " eta=%s,\n" + " tint=%s,\n" + "]" + % ( + self.eta, + self.tint, + ) + ) + + mi.register_bsdf("mybsdf", lambda props: MyBSDF(props)) + + w = 16 + h = 16 + + n = 5 + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def run( + scene: mi.Scene, n: int, func: Callable[[mi.Scene], mi.TensorXf] + ) -> List[mi.TensorXf]: + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + def load_scene(bsdf: str): + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + if bsdf == "twosided": + scene["white"] = { + "type": "twosided", + "material": { + "type": "diffuse", + "reflectance": {"type": "rgb", "value": 0.4}, + }, + } + elif bsdf == "mask": + scene["white"] = { + "type": "mask", + # Base material: a two-sided textured diffuse BSDF + "material": { + "type": "twosided", + "bsdf": { + "type": "diffuse", + "reflectance": { + "type": "bitmap", + "filename": find_resource( + "resources/data/common/textures/wood.jpg" + ), + }, + }, + }, + # Fetch the opacity mask from a monochromatic texture + "opacity": { + "type": "bitmap", + "filename": find_resource( + "resources/data/common/textures/leaf_mask.png" + ), + }, + } + elif bsdf == "bumpmap": + scene["white"] = { + "type": "bumpmap", + "arbitrary": { + "type": "bitmap", + "raw": True, + "filename": find_resource( + "resources/data/common/textures/floor_tiles_bumpmap.png" + ), + }, + "bsdf": {"type": "roughplastic"}, + } + elif bsdf == "normalmap": + scene["white"] = { + "type": "normalmap", + "normalmap": { + "type": "bitmap", + "raw": True, + "filename": find_resource( + "resources/data/common/textures/floor_tiles_normalmap.jpg" + ), + }, + "bsdf": {"type": "roughplastic"}, + } + elif bsdf == "blendbsdf": + scene["white"] = { + "type": "blendbsdf", + "weight": { + "type": "bitmap", + "filename": find_resource( + "resources/data/common/textures/noise_01.jpg" + ), + }, + "bsdf_0": {"type": "conductor"}, + "bsdf_1": {"type": "roughplastic", "diffuse_reflectance": 0.1}, + } + elif bsdf == "diffuse": + scene["white"] = { + "type": "diffuse", + "reflectance": { + "type": "bitmap", + "filename": find_resource( + "resources/data/common/textures/wood.jpg" + ), + }, + } + elif bsdf == "custom": + scene["white"] = { + "type": "mybsdf", + "tint": [0.2, 0.9, 0.2], + "eta": 1.33, + } + else: + scene["white"] = { + "type": bsdf, + } + scene = mi.load_dict(scene, parallel=False) + return scene + + scene = load_scene(bsdf) + images_ref = run(scene, n, func) + scene = load_scene(bsdf) + images_frozen = run(scene, n, frozen) + + assert frozen.n_recordings < n + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +@pytest.mark.parametrize( + "emitter", + [ + "area", + "point", + "constant", + "envmap", + "spot", + "projector", + "directional", + "directionalarea", + ], +) +def test05_emitter(variants_vec_rgb, emitter): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + + w = 16 + h = 16 + + n = 5 + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def load_scene(emitter): + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + + if emitter == "point": + scene["emitter"] = { + "type": "point", + "position": [0.0, 0.0, 0.0], + "intensity": { + "type": "rgb", + "value": mi.ScalarColor3d(50, 50, 50), + }, + } + elif emitter == "constant": + scene["emitter"] = { + "type": "constant", + "radiance": { + "type": "rgb", + "value": 1.0, + }, + } + elif emitter == "envmap": + scene["emitter"] = { + "type": "envmap", + "filename": find_resource( + "resources/data/scenes/matpreview/envmap.exr" + ), + } + elif emitter == "spot": + scene["emitter"] = { + "type": "spot", + "to_world": mi.ScalarTransform4f().look_at( + origin=[0, 0, 0], + target=[0, -1, 0], + up=[0, 0, 1], + ), + "intensity": { + "type": "rgb", + "value": 1.0, + }, + } + elif emitter == "projector": + scene["emitter"] = { + "type": "projector", + "to_world": mi.ScalarTransform4f().look_at( + origin=[0, 0, 0], + target=[0, -1, 0], + up=[0, 0, 1], + ), + "fov": 45, + "irradiance": { + "type": "bitmap", + "filename": find_resource( + "resources/data/common/textures/flower.bmp" + ), + }, + } + elif emitter == "directional": + scene["emitter"] = { + "type": "directional", + "direction": [0, 0, -1], + "irradiance": { + "type": "rgb", + "value": 1.0, + }, + } + elif emitter == "directionalarea": + scene["light"]["emitter"] = { + "type": "directionalarea", + "radiance": { + "type": "rgb", + "value": 1.0, + }, + } + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run(n: int, func: Callable[[mi.Scene], mi.TensorXf]) -> List[mi.TensorXf]: + scene = load_scene(emitter) + + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + images_ref = run(n, func) + # scene = load_scene(emitter) + # del scene + # gc.collect() + # gc.collect() + images_frozen = run(n, frozen) + + assert frozen.n_recordings < n + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +@pytest.mark.parametrize( + "integrator", + [ + "direct", + "path", + "prb", + "prb_basic", + "direct_projective", + "prb_projective", + "moment", + "ptracer", + "depth", + ], +) +def test06_integrators(variants_vec_rgb, integrator): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + + w = 16 + h = 16 + + n = 5 + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def load_scene(): + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + + if integrator == "path": + scene["integrator"] = { + "type": "path", + "max_depth": 4, + } + elif integrator == "direct": + scene["integrator"] = { + "type": "direct", + } + elif integrator == "prb": + scene["integrator"] = { + "type": "prb", + } + elif integrator == "prb_basic": + scene["integrator"] = { + "type": "prb_basic", + } + elif integrator == "direct_projective": + scene["integrator"] = { + "type": "direct_projective", + } + elif integrator == "prb_projective": + scene["integrator"] = { + "type": "prb_projective", + } + elif integrator == "moment": + scene["integrator"] = { + "type": "moment", + "nested": { + "type": "path", + }, + } + elif integrator == "ptracer": + scene["integrator"] = { + "type": "ptracer", + "max_depth": 8, + } + elif integrator == "depth": + scene["integrator"] = {"type": "depth"} + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run(n: int, func: Callable[[mi.Scene], mi.TensorXf]) -> List[mi.TensorXf]: + scene = load_scene() + + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + # scene = load_scene() + images_ref = run(n, func) + images_frozen = run(n, frozen) + + assert frozen.n_recordings < n + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +@pytest.mark.parametrize( + "shape", + [ + "mesh", + "disk", + "cylinder", + "bsplinecurve", + "linearcurve", + "sdfgrid", + # "instance", + "sphere", + ], +) +def test07_shape(variants_vec_rgb, shape): + w = 16 + h = 16 + + n = 5 + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + # dr.set_flag(dr.JitFlag.LaunchBlocking, True) + # dr.set_flag(dr.JitFlag.OptimizeCalls, False) + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + + frozen = dr.freeze(func) + + def load_scene(): + from mitsuba.scalar_rgb import Transform4f as T + + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + del scene["small-box"] + del scene["large-box"] + + if shape == "mesh": + scene["shape"] = { + "type": "ply", + "filename": find_resource("resources/data/common/meshes/teapot.ply"), + "to_world": T().scale(0.1), + } + elif shape == "disk": + scene["shape"] = { + "type": "disk", + "material": { + "type": "diffuse", + "reflectance": { + "type": "checkerboard", + "to_uv": mi.ScalarTransform4f().rotate( + axis=[1, 0, 0], angle=45 + ), + }, + }, + } + elif shape == "cylinder": + scene["shape"] = { + "type": "cylinder", + "radius": 0.3, + "material": {"type": "diffuse"}, + "to_world": mi.ScalarTransform4f().rotate(axis=[1, 0, 0], angle=10), + } + elif shape == "bsplinecurve": + scene["shape"] = { + "type": "bsplinecurve", + "to_world": mi.ScalarTransform4f() + .scale(1.0) + .rotate(axis=[0, 1, 0], angle=45), + "filename": find_resource("resources/data/common/meshes/curve.txt"), + "silhouette_sampling_weight": 1.0, + } + elif shape == "linearcurve": + scene["shape"] = { + "type": "linearcurve", + "to_world": mi.ScalarTransform4f() + .translate([0, -1, 0]) + .scale(1) + .rotate(axis=[0, 1, 0], angle=45), + "filename": find_resource("resources/data/common/meshes/curve.txt"), + } + elif shape == "sdfgrid": + scene["shape"] = { + "type": "sdfgrid", + "bsdf": {"type": "diffuse"}, + "filename": find_resource( + "resources/data/docs/scenes/sdfgrid/torus_sdfgrid.vol" + ), + } + elif shape == "instance": + scene["sg"] = { + "type": "shapegroup", + "first_object": { + "type": "ply", + "filename": find_resource( + "resources/data/common/meshes/teapot.ply" + ), + "bsdf": { + "type": "roughconductor", + }, + "to_world": mi.ScalarTransform4f() + .translate([0.5, 0, 0]) + .scale([0.2, 0.2, 0.2]), + }, + "second_object": { + "type": "sphere", + "to_world": mi.ScalarTransform4f() + .translate([-0.5, 0, 0]) + .scale([0.2, 0.2, 0.2]), + "bsdf": { + "type": "diffuse", + }, + }, + } + scene["first_instance"] = { + "type": "instance", + "shapegroup": {"type": "ref", "id": "sg"}, + } + elif shape == "sphere": + scene["shape"] = { + "type": "sphere", + "center": [0, 0, 0], + "radius": 0.5, + "bsdf": {"type": "diffuse"}, + } + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run( + scene, n: int, func: Callable[[mi.Scene], mi.TensorXf] + ) -> List[mi.TensorXf]: + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + scene = load_scene() + images_ref = run(scene, n, func) + images_frozen = run(scene, n, frozen) + + assert frozen.n_recordings < n + for i, (ref, frozen) in enumerate(zip(images_ref, images_frozen)): + os.makedirs(f"out/{shape}", exist_ok=True) + mi.util.write_bitmap(f"out/{shape}/ref{i}.jpg", ref) + mi.util.write_bitmap(f"out/{shape}/frozen{i}.jpg", frozen) + + for i, (ref, frozen) in enumerate(zip(images_ref, images_frozen)): + assert dr.allclose(ref, frozen) + + +@pytest.mark.parametrize("optimizer", ["sgd", "adam"]) +def test07_optimizer(variants_vec_rgb, optimizer): + k = "red.reflectance.value" + w = 16 + h = 16 + n = 10 + + def mse(image, image_ref): + return dr.sum(dr.square(image - image_ref), axis=None) + + def optimize(scene, opt, image_ref): + params = mi.traverse(scene) + params.update(opt) + + image = mi.render(scene, params, spp=1) + + loss = mse(image, image_ref) + + dr.backward(loss) + + opt.step() + + return image, loss + + frozen = dr.freeze(optimize) + + def run(n: int, optimize): + scene = mi.cornell_box() + scene["integrator"] = { + "type": "prb", + } + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + scene = mi.load_dict(scene, parallel=False) + + image_ref = mi.render(scene, spp=512) + + params = mi.traverse(scene) + params.keep(k) + + if optimizer == "adam": + opt = mi.ad.Adam(lr=0.05) + elif optimizer == "sgd": + opt = mi.ad.SGD(lr=0.005) + opt[k] = mi.Color3f(0.01, 0.2, 0.9) + + for i in range(n): + image, loss = optimize(scene, opt, image_ref) + + return image, opt[k] + + image_ref, param_ref = run(n, optimize) + + image_frozen, param_frozen = run(n, frozen) + assert frozen.n_recordings == 2 + + # Optimizing the reflectance is not as prone to divergence, + # therefore we can test if the two methods produce the same results + assert dr.allclose(param_ref, param_frozen) + + +@pytest.mark.parametrize( + "medium", + [ + "homogeneous", + "heterogeneous", + ], +) +def test08_medium(variants_vec_rgb, medium): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + + w = 16 + h = 16 + + n = 5 + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + result = mi.render(scene, spp=1) + return result + frozen = dr.freeze(func) + + def load_scene(): + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + scene["integrator"] = {"type": "volpath"} + scene["sensor"]["medium"] = {"type": "ref", "id": "fog"} + + if medium == "homogeneous": + scene["sensor"]["medium"] = { + "type": "homogeneous", + "albedo": {"type": "rgb", "value": [0.99, 0.9, 0.96]}, + "sigma_t": { + "type": "rgb", + "value": [0.5, 0.25, 0.8], + }, + } + elif medium == "heterogeneous": + scene["sensor"]["medium"] = { + "type": "heterogeneous", + "albedo": {"type": "rgb", "value": [0.99, 0.9, 0.96]}, + "sigma_t": { + "type": "gridvolume", + "filename": find_resource( + "resources/data/docs/scenes/textures/albedo.vol" + ), + }, + } + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run(n: int, func: Callable[[mi.Scene], mi.TensorXf]) -> List[mi.TensorXf]: + scene = load_scene() + + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + # scene = load_scene() + images_ref = run(n, func) + images_frozen = run(n, frozen) + + assert frozen.n_recordings < n + + for i, (ref, frozen) in enumerate(zip(images_ref, images_frozen)): + os.makedirs(f"out/{medium}", exist_ok=True) + mi.util.write_bitmap(f"out/{medium}/ref{i}.jpg", ref) + mi.util.write_bitmap(f"out/{medium}/frozen{i}.jpg", frozen) + + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) + + +@pytest.mark.parametrize( + "sampler", + [ + "independent", + "stratified", + "multijitter", + "orthogonal", + "ldsampler", + ], +) +def test09_sampler(variants_vec_rgb, sampler): + # dr.set_log_level(dr.LogLevel.Trace) + # dr.set_flag(dr.JitFlag.ReuseIndices, False) + # dr.set_flag(dr.JitFlag.Debug, True) + + w = 16 + h = 16 + + n = 5 + + def func(scene: mi.Scene) -> mi.TensorXf: + with dr.profile_range("render"): + # spp of 4 to suppress warning + result = mi.render(scene, spp=4) + return result + frozen = dr.freeze(func) + + def load_scene(): + scene = mi.cornell_box() + scene["sensor"]["film"]["width"] = w + scene["sensor"]["film"]["height"] = h + + scene["sensor"]["sampler"] = {"type": sampler} + + scene = mi.load_dict(scene, parallel=False) + return scene + + def run(n: int, func: Callable[[mi.Scene], mi.TensorXf]) -> List[mi.TensorXf]: + scene = load_scene() + + images = [] + for i in range(n): + img = func(scene) + dr.eval(img) + + images.append(img) + + return images + + # scene = load_scene() + images_ref = run(n, func) + images_frozen = run(n, frozen) + + assert frozen.n_recordings < n + + for i, (ref, frozen) in enumerate(zip(images_ref, images_frozen)): + os.makedirs(f"out/{sampler}", exist_ok=True) + mi.util.write_bitmap(f"out/{sampler}/ref{i}.jpg", ref) + mi.util.write_bitmap(f"out/{sampler}/frozen{i}.jpg", frozen) + + for ref, frozen in zip(images_ref, images_frozen): + assert dr.allclose(ref, frozen) diff --git a/src/samplers/independent.cpp b/src/samplers/independent.cpp index 519261638..500d7fd7a 100644 --- a/src/samplers/independent.cpp +++ b/src/samplers/independent.cpp @@ -111,6 +111,8 @@ class IndependentSampler final : public PCG32Sampler { private: IndependentSampler(const IndependentSampler &sampler) : Base(sampler) {} + + DR_TRAVERSE_CB(Base) }; MI_IMPLEMENT_CLASS_VARIANT(IndependentSampler, Sampler) diff --git a/src/samplers/ldsampler.cpp b/src/samplers/ldsampler.cpp index 26e6cf28a..b95b00d30 100644 --- a/src/samplers/ldsampler.cpp +++ b/src/samplers/ldsampler.cpp @@ -151,16 +151,6 @@ class LowDiscrepancySampler final : public Sampler { dr::schedule(m_scramble_seed); } - void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const override { - auto fields = dr::make_tuple(m_scramble_seed, m_dimension_index); - dr::traverse_1_fn_ro(fields, payload, fn); - } - - void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)) override { - auto fields = dr::tie(m_scramble_seed, m_dimension_index); - dr::traverse_1_fn_rw(fields, payload, fn); - } - std::string to_string() const override { std::ostringstream oss; oss << "LowDiscrepancySampler [" << std::endl @@ -178,6 +168,8 @@ class LowDiscrepancySampler final : public Sampler { /// Per-sequence scramble seed UInt32 m_scramble_seed; + + DR_TRAVERSE_CB(Base, m_scramble_seed); }; MI_IMPLEMENT_CLASS_VARIANT(LowDiscrepancySampler , Sampler) diff --git a/src/samplers/multijitter.cpp b/src/samplers/multijitter.cpp index 152a62e77..1c47ebdcd 100644 --- a/src/samplers/multijitter.cpp +++ b/src/samplers/multijitter.cpp @@ -170,16 +170,6 @@ class MultijitterSampler final : public PCG32Sampler { dr::schedule(m_permutation_seed); } - void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const override { - auto fields = dr::make_tuple(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_ro(fields, payload, fn); - } - - void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)) override { - auto fields = dr::tie(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_rw(fields, payload, fn); - } - std::string to_string() const override { std::ostringstream oss; oss << "MultijitterSampler[" << std::endl @@ -211,6 +201,8 @@ class MultijitterSampler final : public PCG32Sampler { /// Per-sequence permutation seed UInt32 m_permutation_seed; + + DR_TRAVERSE_CB(Base, m_permutation_seed); }; MI_IMPLEMENT_CLASS_VARIANT(MultijitterSampler, Sampler) diff --git a/src/samplers/orthogonal.cpp b/src/samplers/orthogonal.cpp index 3839c4064..94e37bea6 100644 --- a/src/samplers/orthogonal.cpp +++ b/src/samplers/orthogonal.cpp @@ -157,16 +157,6 @@ class OrthogonalSampler final : public PCG32Sampler { dr::schedule(m_permutation_seed); } - void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const override { - auto fields = dr::make_tuple(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_ro(fields, payload, fn); - } - - void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)) override { - auto fields = dr::tie(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_rw(fields, payload, fn); - } - std::string to_string() const override { std::ostringstream oss; oss << "OrthogonalSampler[" << std::endl @@ -270,6 +260,8 @@ class OrthogonalSampler final : public PCG32Sampler { /// Per-sequence permutation seed UInt32 m_permutation_seed; + + DR_TRAVERSE_CB(Base, m_permutation_seed); }; MI_IMPLEMENT_CLASS_VARIANT(OrthogonalSampler, Sampler) diff --git a/src/samplers/stratified.cpp b/src/samplers/stratified.cpp index cd1ef9c75..203a9bee3 100644 --- a/src/samplers/stratified.cpp +++ b/src/samplers/stratified.cpp @@ -157,16 +157,6 @@ class StratifiedSampler final : public PCG32Sampler { dr::schedule(m_permutation_seed); } - void traverse_1_cb_ro(void *payload, void (*fn)(void *, uint64_t)) const override { - auto fields = dr::make_tuple(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_ro(fields, payload, fn); - } - - void traverse_1_cb_rw(void *payload, uint64_t (*fn)(void *, uint64_t)) override { - auto fields = dr::tie(m_rng, m_dimension_index, m_permutation_seed); - dr::traverse_1_fn_rw(fields, payload, fn); - } - std::string to_string() const override { std::ostringstream oss; oss << "StratifiedSampler[" << std::endl @@ -198,6 +188,8 @@ class StratifiedSampler final : public PCG32Sampler { /// Per-sequence permutation seed UInt32 m_permutation_seed; + + DR_TRAVERSE_CB(Base, m_permutation_seed); }; MI_IMPLEMENT_CLASS_VARIANT(StratifiedSampler, Sampler) diff --git a/src/sensors/perspective.cpp b/src/sensors/perspective.cpp index 6a4758ddc..2d537a1e8 100644 --- a/src/sensors/perspective.cpp +++ b/src/sensors/perspective.cpp @@ -412,6 +412,10 @@ class PerspectiveCamera final : public ProjectiveCamera { Float m_x_fov; Vector3f m_dx, m_dy; Vector2f m_principal_point_offset; + + DR_TRAVERSE_CB(Base, m_camera_to_sample, m_sample_to_camera, m_x_fov, m_dx, + m_dy, m_image_rect, m_normalization, m_x_fov, + m_principal_point_offset); }; MI_IMPLEMENT_CLASS_VARIANT(PerspectiveCamera, ProjectiveCamera) diff --git a/src/shapes/bsplinecurve.cpp b/src/shapes/bsplinecurve.cpp index fb63a28a4..c1c5f52d5 100644 --- a/src/shapes/bsplinecurve.cpp +++ b/src/shapes/bsplinecurve.cpp @@ -1310,6 +1310,8 @@ class BSplineCurve final : public Shape { mutable CUdeviceptr* m_radius_buffer_ptr = nullptr; mutable CUdeviceptr* m_index_buffer_ptr = nullptr; #endif + + DR_TRAVERSE_CB(Base, m_curves_prim_idx, m_indices, m_control_points) }; MI_IMPLEMENT_CLASS_VARIANT(BSplineCurve, Shape) diff --git a/src/shapes/cube.cpp b/src/shapes/cube.cpp index c6c595745..3888f4151 100644 --- a/src/shapes/cube.cpp +++ b/src/shapes/cube.cpp @@ -163,6 +163,8 @@ MI_VARIANT class Cube final : public Mesh { } MI_DECLARE_CLASS() + + DR_TRAVERSE_CB(Base); }; MI_IMPLEMENT_CLASS_VARIANT(Cube, Mesh) diff --git a/src/shapes/disk.cpp b/src/shapes/disk.cpp index 2f859a226..604fb1c4e 100644 --- a/src/shapes/disk.cpp +++ b/src/shapes/disk.cpp @@ -552,6 +552,8 @@ class Disk final : public Shape { Frame3f m_frame; Float m_du, m_dv; Float m_inv_surface_area; + + DR_TRAVERSE_CB(Base, m_frame, m_du, m_dv, m_inv_surface_area) }; MI_IMPLEMENT_CLASS_VARIANT(Disk, Shape) diff --git a/src/shapes/instance.cpp b/src/shapes/instance.cpp index 3c13e1783..6d5d300de 100644 --- a/src/shapes/instance.cpp +++ b/src/shapes/instance.cpp @@ -300,6 +300,8 @@ class Instance final: public Shape { MI_DECLARE_CLASS() private: ref m_shapegroup; + + DR_TRAVERSE_CB(Base, m_shapegroup) }; MI_IMPLEMENT_CLASS_VARIANT(Instance, Shape) diff --git a/src/shapes/linearcurve.cpp b/src/shapes/linearcurve.cpp index 2207e5402..31c59e170 100644 --- a/src/shapes/linearcurve.cpp +++ b/src/shapes/linearcurve.cpp @@ -517,6 +517,8 @@ class LinearCurve final : public Shape { mutable void* m_radius_buffer_ptr = nullptr; mutable void* m_index_buffer_ptr = nullptr; #endif + + DR_TRAVERSE_CB(Base, m_indices, m_control_points) }; MI_IMPLEMENT_CLASS_VARIANT(LinearCurve, Shape) diff --git a/src/shapes/obj.cpp b/src/shapes/obj.cpp index c3cbf3adb..4a9913930 100644 --- a/src/shapes/obj.cpp +++ b/src/shapes/obj.cpp @@ -409,6 +409,8 @@ class OBJMesh final : public Mesh { } MI_DECLARE_CLASS() + + DR_TRAVERSE_CB(Base); }; MI_IMPLEMENT_CLASS_VARIANT(OBJMesh, Mesh) diff --git a/src/shapes/rectangle.cpp b/src/shapes/rectangle.cpp index 3b41897cb..df0e97444 100644 --- a/src/shapes/rectangle.cpp +++ b/src/shapes/rectangle.cpp @@ -600,6 +600,8 @@ class Rectangle final : public Shape { private: Frame3f m_frame; Float m_inv_surface_area; + + DR_TRAVERSE_CB(Base, m_frame, m_inv_surface_area); }; MI_IMPLEMENT_CLASS_VARIANT(Rectangle, Shape) diff --git a/src/shapes/sdfgrid.cpp b/src/shapes/sdfgrid.cpp index d16251cd3..bb62d5845 100644 --- a/src/shapes/sdfgrid.cpp +++ b/src/shapes/sdfgrid.cpp @@ -1139,6 +1139,9 @@ class SDFGrid final : public Shape { uint32_t m_filled_voxel_count = 0; NormalMethod m_normal_method; + + DR_TRAVERSE_CB(Base, m_grid_texture, m_inv_shape, m_voxel_size, + m_jit_bboxes, m_jit_voxel_indices) }; MI_IMPLEMENT_CLASS_VARIANT(SDFGrid, Shape) diff --git a/src/shapes/sphere.cpp b/src/shapes/sphere.cpp index 33c2441ae..241c170a6 100644 --- a/src/shapes/sphere.cpp +++ b/src/shapes/sphere.cpp @@ -766,6 +766,8 @@ class Sphere final : public Shape { Float m_inv_surface_area; bool m_flip_normals; + + DR_TRAVERSE_CB(Base, m_center, m_radius, m_inv_surface_area) }; MI_IMPLEMENT_CLASS_VARIANT(Sphere, Shape) diff --git a/src/spectra/d65.cpp b/src/spectra/d65.cpp index 01e99322d..42d6f3742 100644 --- a/src/spectra/d65.cpp +++ b/src/spectra/d65.cpp @@ -313,6 +313,8 @@ class D65Spectrum final : public Texture { ScalarFloat m_scale; bool m_has_value = false; + + DR_TRAVERSE_CB(Base, m_value, m_nested_texture, m_d65); }; MI_IMPLEMENT_CLASS_VARIANT(D65Spectrum, Texture) diff --git a/src/spectra/irregular.cpp b/src/spectra/irregular.cpp index 38204e46a..5704b0526 100644 --- a/src/spectra/irregular.cpp +++ b/src/spectra/irregular.cpp @@ -171,6 +171,8 @@ class IrregularSpectrum final : public Texture { MI_DECLARE_CLASS() private: IrregularContinuousDistribution m_distr; + + DR_TRAVERSE_CB(Texture, m_distr); }; MI_IMPLEMENT_CLASS_VARIANT(IrregularSpectrum, Texture) diff --git a/src/spectra/regular.cpp b/src/spectra/regular.cpp index 55f966009..2ef3a7947 100644 --- a/src/spectra/regular.cpp +++ b/src/spectra/regular.cpp @@ -168,6 +168,8 @@ class RegularSpectrum final : public Texture { MI_DECLARE_CLASS() private: ContinuousDistribution m_distr; + + DR_TRAVERSE_CB(Texture, m_distr); }; MI_IMPLEMENT_CLASS_VARIANT(RegularSpectrum, Texture) diff --git a/src/spectra/srgb.cpp b/src/spectra/srgb.cpp index e1d9ec070..a5ad61e5f 100644 --- a/src/spectra/srgb.cpp +++ b/src/spectra/srgb.cpp @@ -145,6 +145,8 @@ class SRGBReflectanceSpectrum final : public Texture { static constexpr size_t ChannelCount = is_monochromatic_v ? 1 : 3; Color m_value; + + DR_TRAVERSE_CB(Texture, m_value); }; MI_IMPLEMENT_CLASS_VARIANT(SRGBReflectanceSpectrum, Texture) diff --git a/src/spectra/uniform.cpp b/src/spectra/uniform.cpp index 5e09c780e..69485e28b 100644 --- a/src/spectra/uniform.cpp +++ b/src/spectra/uniform.cpp @@ -121,6 +121,8 @@ class UniformSpectrum final : public Texture { private: Float m_value; ScalarVector2f m_range; + + DR_TRAVERSE_CB(Texture, m_value); }; MI_IMPLEMENT_CLASS_VARIANT(UniformSpectrum, Texture) diff --git a/src/textures/bitmap.cpp b/src/textures/bitmap.cpp index 30ce0a635..eaeb2844f 100644 --- a/src/textures/bitmap.cpp +++ b/src/textures/bitmap.cpp @@ -355,6 +355,8 @@ class BitmapTexture final : public Texture { dr::WrapMode m_wrap_mode; mutable ref m_bitmap; TensorXf* m_tensor; + + DR_TRAVERSE_CB(Texture, *m_tensor, m_bitmap); }; template @@ -904,6 +906,8 @@ class BitmapTextureImpl : public Texture { // Optional: distribution for importance sampling mutable std::mutex m_mutex; std::unique_ptr> m_distr2d; + + DR_TRAVERSE_CB(Texture, m_texture, m_mean, m_distr2d); }; MI_IMPLEMENT_CLASS_VARIANT(BitmapTexture, Texture) diff --git a/src/textures/checkerboard.cpp b/src/textures/checkerboard.cpp index 843e3612d..e68c084ff 100644 --- a/src/textures/checkerboard.cpp +++ b/src/textures/checkerboard.cpp @@ -130,6 +130,8 @@ class Checkerboard final : public Texture { ref m_color0; ref m_color1; ScalarTransform3f m_transform; + + DR_TRAVERSE_CB(Texture, m_color0, m_color1) }; MI_IMPLEMENT_CLASS_VARIANT(Checkerboard, Texture) diff --git a/src/textures/volume.cpp b/src/textures/volume.cpp index 15a85c2e9..58753efd8 100644 --- a/src/textures/volume.cpp +++ b/src/textures/volume.cpp @@ -92,8 +92,10 @@ class VolumeAdapter final : public Texture { MI_DECLARE_CLASS() protected: ref m_volume; + + DR_TRAVERSE_CB(Texture, m_volume); }; MI_IMPLEMENT_CLASS_VARIANT(VolumeAdapter, Texture) MI_EXPORT_PLUGIN(VolumeAdapter, "Volumetric texture") -NAMESPACE_END(mitsuba) \ No newline at end of file +NAMESPACE_END(mitsuba) diff --git a/src/volumes/const.cpp b/src/volumes/const.cpp index 22a8be1ed..e5abd6075 100644 --- a/src/volumes/const.cpp +++ b/src/volumes/const.cpp @@ -95,6 +95,8 @@ class ConstVolume final : public Volume { MI_DECLARE_CLASS() protected: ref m_value; + + DR_TRAVERSE_CB(Base, m_value); }; MI_IMPLEMENT_CLASS_VARIANT(ConstVolume, Volume) diff --git a/src/volumes/grid.cpp b/src/volumes/grid.cpp index 2ba10f5b6..37f8a2341 100644 --- a/src/volumes/grid.cpp +++ b/src/volumes/grid.cpp @@ -625,6 +625,8 @@ class GridVolume final : public Volume { bool m_fixed_max = false; ScalarFloat m_max; std::vector m_max_per_channel; + + DR_TRAVERSE_CB(Base, m_texture); }; MI_IMPLEMENT_CLASS_VARIANT(GridVolume, Volume)