Add various sampler plugins and update sampler API

docs/examples/02_depth_integrator/depth_integrator.py

@@ -32,7 +32,7 @@
 total_sample_count = ek.hprod(film_size) * spp
 
 if sampler.wavefront_size() != total_sample_count:
-    sampler.seed(ek.arange(UInt64, total_sample_count))
+    sampler.seed(0, total_sample_count)
 
 # Enumerate discrete sample & pixel indices, and uniformly sample
 # positions within each pixel.

docs/generate_plugin_doc.py

@@ -51,7 +51,11 @@
                      'srgb_d65',
                      'blackbody']
 
-SAMPLER_ORDERING = ['independent']
+SAMPLER_ORDERING = ['independent',
+                    'stratified',
+                    'multijitter',
+                    'orthogonal',
+                    'ldsampler']
 
 INTEGRATOR_ORDERING = ['direct',
                        'path',

docs/references.bib

@@ -23,6 +23,44 @@ @inproceedings {Mojzik2016BidirectionalPol
   year = {2016},
 }
 
+@article{kensler1967correlated,
+  title={Correlated multi-jittered sampling},
+  author={Kensler, Andrew},
+  journal={Mathematical Physics and applied mathematics},
+  volume={7},
+  pages={86--112},
+  year={1967}
+}
+
+@article {Kollig2002Efficient,
+    author = {Kollig, Thomas and Keller, Alexander},
+    title = {{Efficient Multidimensional Sampling}},
+    journal = {Computer Graphics Forum},
+    year = {2002},
+    editor = {},
+    volume = {21},
+    number = {3},
+    publisher = {Blackwell Publishers, Inc and the Eurographics Association},
+    pages = {557-563},
+    DOI = {10.1111/1467-8659.00706}
+}
+
+@article{jarosz19orthogonal,
+    author = "Jarosz, Wojciech and Enayet, Afnan and Kensler, Andrew and Kilpatrick, Charlie and Christensen, Per",
+    title = "Orthogonal array sampling for {{Monte}} {{Carlo}} rendering",
+    journal = "Computer Graphics Forum (Proceedings of EGSR)",
+    year = "2019",
+    volume = "38",
+    number = "4",
+    month = jul,
+    pages = "135--147",
+    doi = "10/gf6rx5",
+    publisher = "The Eurographics Association",
+    abstract = "We generalize N-rooks, jittered, and (correlated) multi-jittered sampling to higher dimensions by importing and improving upon a class of techniques called orthogonal arrays from the statistics literature. Renderers typically combine or ``pad'' a collection of lower-dimensional (e.g. 2D and 1D) stratified patterns to form higher-dimensional samples for integration. This maintains stratification in the original dimension pairs, but looses it for all other dimension pairs. For truly multi-dimensional integrands like those in rendering, this increases variance and deteriorates its rate of convergence to that of pure random sampling. Care must therefore be taken to assign the primary dimension pairs to the dimensions with most integrand variation, but this complicates implementations. We tackle this problem by developing a collection of practical, in-place multi-dimensional sample generation routines that stratify points on all t-dimensional and 1-dimensional projections simultaneously. For instance, when t=2, any 2D projection of our samples is a (correlated) multi-jittered point set. This property not only reduces variance, but also simplifies implementations since sample dimensions can now be assigned to integrand dimensions arbitrarily while maintaining the same level of stratification. Our techniques reduce variance compared to traditional 2D padding approaches like PBRT's (0,2) and Stratified samplers, and provide quality nearly equal to state-of-the-art QMC samplers like Sobol and Halton while avoiding their structured artifacts as commonly seen when using a single sample set to cover an entire image. While in this work we focus on constructing finite sampling point sets, we also discuss potential avenues for extending our work to progressive sequences (more suitable for incremental rendering) in the future."
+}
+
+
+
 @article{Jakob2019Spectral,
 	author = {Wenzel Jakob and Johannes Hanika},
 	title = {A Low-Dimensional Function Space for Efficient Spectral Upsampling},

include/mitsuba/core/qmc.h

@@ -183,4 +183,47 @@ class MTS_EXPORT_CORE RadicalInverse : public Object {
     int m_scramble;
 };
 
+
+/// Van der Corput radical inverse in base 2
+template <typename UInt, typename Float = float_array_t<UInt>>
+Float radical_inverse_2(UInt index, UInt scramble = 0) {
+    if constexpr (is_double_v<Float>) {
+        index = (index << 32) | (index >> 32);
+        index = ((index & 0x0000ffff0000ffffULL) << 16) | ((index & 0xffff0000ffff0000ULL) >> 16);
+        index = ((index & 0x00ff00ff00ff00ffULL) << 8)  | ((index & 0xff00ff00ff00ff00ULL) >> 8);
+        index = ((index & 0x0f0f0f0f0f0f0f0fULL) << 4)  | ((index & 0xf0f0f0f0f0f0f0f0ULL) >> 4);
+        index = ((index & 0x3333333333333333ULL) << 2)  | ((index & 0xccccccccccccccccULL) >> 2);
+        index = ((index & 0x5555555555555555ULL) << 1)  | ((index & 0xaaaaaaaaaaaaaaaaULL) >> 1);
+
+        /* Generate an uniformly distributed double precision number in [1,2)
+         * from the scrambled index and subtract 1. */
+        return reinterpret_array<Float>(sr<12>(index ^ scramble) | 0x3ff0000000000000ull) - 1.0;
+    } else {
+        index = (index << 16) | (index >> 16);
+        index = ((index & 0x00ff00ff) << 8) | ((index & 0xff00ff00) >> 8);
+        index = ((index & 0x0f0f0f0f) << 4) | ((index & 0xf0f0f0f0) >> 4);
+        index = ((index & 0x33333333) << 2) | ((index & 0xcccccccc) >> 2);
+        index = ((index & 0x55555555) << 1) | ((index & 0xaaaaaaaa) >> 1);
+
+        /* Generate an uniformly distributed single precision number in [1,2)
+         * from the scrambled index and subtract 1. */
+        return reinterpret_array<Float>(sr<9>(index ^ scramble) | 0x3f800000u) - 1.f;
+    }
+}
+
+
+/// Sobol' radical inverse in base 2
+template <typename UInt, typename Float = float_array_t<UInt>>
+Float sobol_2(UInt index, UInt scramble = 0) {
+    if constexpr (is_double_v<Float>) {
+        for (UInt v = 1ULL << 52; index != 0; index >>= 1, v ^= v >> 1)
+            masked(scramble, eq(index & 1U, 1U)) ^= v;
+        return reinterpret_array<Float>(sr<12>(scramble) | 0x3ff0000000000000ull) - 1.0;
+    } else {
+        for (UInt v = 1U << 31; index != 0; index >>= 1, v ^= v >> 1)
+            masked(scramble, eq(index & 1U, 1U)) ^= v;
+        return Float(scramble) / Float(1ULL << 32);
+    }
+}
+
 NAMESPACE_END(mitsuba)

include/mitsuba/core/random.h

@@ -28,9 +28,11 @@
 
 #pragma once
 
+#include <mitsuba/core/math.h>
 #include <mitsuba/core/simd.h>
 #include <mitsuba/core/logger.h>
 #include <mitsuba/core/traits.h>
+#include <mitsuba/core/fwd.h>
 #include <enoki/random.h>
 
 NAMESPACE_BEGIN(enoki)
@@ -171,4 +173,104 @@ auto sample_tea_float(UInt v0, UInt v1, int rounds = 4) {
         return sample_tea_float64(v0, v1, rounds);
 }
 
+/**
+ * \brief Generate pseudorandom permutation vector using a shuffling network and the
+ * \ref sample_tea function. This algorithm has a O(log2(sample_count)) complexity but
+ * only supports permutation vectors whose lengths are a power of 2.
+ *
+ * \param index
+ *     Input index to be mapped
+ * \param sample_count
+ *     Length of the permutation vector
+ * \param seed
+ *     Seed value used as second input to the Tiny Encryption Algorithm. Can be used to
+ *     generate different permutation vectors.
+ * \param rounds
+ *     How many rounds should be executed by the Tiny Encryption Algorithm? The default is 2.
+ * \return
+ *     The index corresponding to the input index in the pseudorandom permutation vector.
+ */
+
+template <typename UInt32>
+UInt32 permute(UInt32 index, uint32_t sample_count, UInt32 seed, int rounds = 2) {
+    uint32_t  n = log2i(sample_count);
+    Assert((1 << n) == sample_count, "sample_count should be a power of 2");
+
+    for (uint32_t  level = 0; level < n; ++level) {
+        UInt32 bit = UInt32(1 << level);
+        // Take a random integer indentical for indices that might be swapped at this level
+        UInt32 rand = sample_tea_32(index | bit, seed, rounds);
+        masked(index, eq(rand & bit, bit)) = index ^ bit;
+    }
+
+    return index;
+}
+
+/**
+ * \brief Generate pseudorandom permutation vector using the algorithm described in Pixar's
+ * technical memo "Correlated Multi-Jittered Sampling":
+ *
+ *     https://graphics.pixar.com/library/MultiJitteredSampling/
+ *
+ *  Unlike \ref permute, this function supports permutation vectors of any length.
+ *
+ * \param index
+ *     Input index to be mapped
+ * \param sample_count
+ *     Length of the permutation vector
+ * \param seed
+ *     Seed value used as second input to the Tiny Encryption Algorithm. Can be used to
+ *     generate different permutation vectors.
+ * \return
+ *     The index corresponding to the input index in the pseudorandom permutation vector.
+ */
+
+template <typename UInt32>
+UInt32 permute_kensler(UInt32 index, uint32_t sample_count, UInt32 seed, mask_t<UInt32> active = true) {
+    if (sample_count == 1)
+        return zero<UInt32>(slices(index));
+
+    UInt32 w = sample_count - 1;
+    w |= w >> 1;
+    w |= w >> 2;
+    w |= w >> 4;
+    w |= w >> 8;
+    w |= w >> 16;
+
+    // Worst case is when the index is sequentially mapped to every invalid numbers (out
+    // of range) before being mapped into the correct range. E.g. decreasing sequence
+    uint32_t max_iter = 0;
+    if constexpr (is_cuda_array_v<UInt32>)
+        max_iter = math::round_to_power_of_two(sample_count) - sample_count + 1;
+
+    uint32_t iter = 0;
+    mask_t<UInt32> invalid = true;
+    UInt32 tmp;
+    do {
+        tmp = index;
+        tmp ^= seed;
+        tmp *= 0xe170893d;
+        tmp ^= seed >> 16;
+        tmp ^= (tmp & w) >> 4;
+        tmp ^= seed >> 8;
+        tmp *= 0x0929eb3f;
+        tmp ^= seed >> 23;
+        tmp ^= (tmp & w) >> 1;
+        tmp *= 1 | seed >> 27;
+        tmp *= 0x6935fa69;
+        tmp ^= (tmp & w) >> 11;
+        tmp *= 0x74dcb303;
+        tmp ^= (tmp & w) >> 2;
+        tmp *= 0x9e501cc3;
+        tmp ^= (tmp & w) >> 2;
+        tmp *= 0xc860a3df;
+        tmp &= w;
+        tmp ^= tmp >> 5;
+        masked(index, invalid) = tmp;
+        invalid = (index >= sample_count);
+    } while (any_or<false>(active && invalid) || (max_iter > ++iter));
+
+    return (index + seed) % sample_count;
+}
+
 NAMESPACE_END(mitsuba)

include/mitsuba/python/docstr.h

@@ -4496,6 +4496,8 @@ static const char *__doc_mitsuba_ProfilerPhase_SampleEmitterDirection = R"doc()d
 
 static const char *__doc_mitsuba_ProfilerPhase_SampleEmitterRay = R"doc()doc";
 
+static const char *__doc_mitsuba_ProfilerPhase_SamplerSeed = R"doc()doc";
+
 static const char *__doc_mitsuba_ProfilerPhase_SamplingIntegratorSample = R"doc()doc";
 
 static const char *__doc_mitsuba_ProfilerPhase_TextureEvaluate = R"doc()doc";
@@ -4921,6 +4923,24 @@ static const char *__doc_mitsuba_RadicalInverse_scramble = R"doc(Return the orig
 
 static const char *__doc_mitsuba_RadicalInverse_to_string = R"doc(Return a human-readable string representation)doc";
 
+static const char *__doc_mitsuba_RandomSampler = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_2 = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_3 = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_RandomSampler = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_check_rng = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_class = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_m_rng = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_seed = R"doc()doc";
+
+static const char *__doc_mitsuba_RandomSampler_wavefront_size = R"doc()doc";
+
 static const char *__doc_mitsuba_Ray =
 R"doc(Simple n-dimensional ray segment data structure
 
@@ -5164,10 +5184,16 @@ static const char *__doc_mitsuba_Sampler_m_base_seed = R"doc()doc";
 
 static const char *__doc_mitsuba_Sampler_m_sample_count = R"doc()doc";
 
+static const char *__doc_mitsuba_Sampler_m_samples_per_wavefront = R"doc()doc";
+
+static const char *__doc_mitsuba_Sampler_m_wavefront_count = R"doc()doc";
+
 static const char *__doc_mitsuba_Sampler_next_1d = R"doc(Retrieve the next component value from the current sample)doc";
 
 static const char *__doc_mitsuba_Sampler_next_2d = R"doc(Retrieve the next two component values from the current sample)doc";
 
+static const char *__doc_mitsuba_Sampler_prepare_wavefront = R"doc(Start the next wavefront)doc";
+
 static const char *__doc_mitsuba_Sampler_sample_count = R"doc(Return the number of samples per pixel)doc";
 
 static const char *__doc_mitsuba_Sampler_seed =
@@ -5177,6 +5203,10 @@ In the context of wavefront ray tracing & dynamic arrays, this
 function must be called with a ``seed_value`` matching the size of the
 wavefront.)doc";
 
+static const char *__doc_mitsuba_Sampler_set_samples_per_wavefront =
+R"doc(Set the number of samples per pass in the wavefront modes (default is
+1))doc";
+
 static const char *__doc_mitsuba_Sampler_wavefront_size = R"doc(Return the size of the wavefront (or 0, if not seeded))doc";
 
 static const char *__doc_mitsuba_SamplingIntegrator =
@@ -8313,6 +8343,29 @@ static const char *__doc_mitsuba_hasher_operator_call = R"doc()doc";
 
 static const char *__doc_mitsuba_ior_from_file = R"doc()doc";
 
+static const char *__doc_mitsuba_kensler_permute =
+R"doc(Generate pseudorandom permutation vector using the algorithm described
+in Pixar's technical memo "Correlated Multi-Jittered Sampling":
+
+https://graphics.pixar.com/library/MultiJitteredSampling/
+
+Unlike permute, this function supports permutation vectors of any
+length.
+
+Parameter ``index``:
+    Input index to be mapped
+
+Parameter ``sample_count``:
+    Length of the permutation vector
+
+Parameter ``seed``:
+    Seed value used as second input to the Tiny Encryption Algorithm.
+    Can be used to generate different permutation vectors.
+
+Returns:
+    The index corresponding to the input index in the pseudorandom
+    permutation vector.)doc";
+
 static const char *__doc_mitsuba_librender_nop =
 R"doc(Dummy function which can be called to ensure that the librender shared
 library is loaded)doc";
@@ -8832,6 +8885,10 @@ Parameter ``w``:
     The (implicitly defined) reference coordinate system basis for the
     Stokes vector travelling along w.)doc";
 
+static const char *__doc_mitsuba_next_float =
+R"doc(Forward next_float call to PCG32 random generator based given type
+size)doc";
+
 static const char *__doc_mitsuba_operator_add = R"doc()doc";
 
 static const char *__doc_mitsuba_operator_add_2 = R"doc(Adding a vector to a point should always yield a point)doc";
@@ -8976,6 +9033,30 @@ static const char *__doc_mitsuba_pdf_uniform_spectrum_2 = R"doc()doc";
 
 static const char *__doc_mitsuba_perspective_projection = R"doc()doc";
 
+static const char *__doc_mitsuba_permute =
+R"doc(Generate pseudorandom permutation vector using a shuffling network and
+the sample_tea function. This algorithm has a O(log2(sample_count))
+complexity but only supports permutation vectors whose lengths are a
+power of 2.
+
+Parameter ``index``:
+    Input index to be mapped
+
+Parameter ``sample_count``:
+    Length of the permutation vector
+
+Parameter ``seed``:
+    Seed value used as second input to the Tiny Encryption Algorithm.
+    Can be used to generate different permutation vectors.
+
+Parameter ``rounds``:
+    How many rounds should be executed by the Tiny Encryption
+    Algorithm? The default is 2.
+
+Returns:
+    The index corresponding to the input index in the pseudorandom
+    permutation vector.)doc";
+
 static const char *__doc_mitsuba_profiler_flags = R"doc()doc";
 
 static const char *__doc_mitsuba_quad_composite_simpson =
@@ -9055,6 +9136,8 @@ Parameter ``n``:
     A tuple (nodes, weights) storing the nodes and weights of the
     quadrature rule.)doc";
 
+static const char *__doc_mitsuba_radical_inverse_2 = R"doc(Van der Corput radical inverse in base 2)doc";
+
 static const char *__doc_mitsuba_ref =
 R"doc(Reference counting helper
 
@@ -9275,6 +9358,8 @@ Parameter ``si``:
 
 \note Defined in scene.h)doc";
 
+static const char *__doc_mitsuba_sobol_2 = R"doc(Sobol' radical inverse in base 2)doc";
+
 static const char *__doc_mitsuba_spectrum_from_file = R"doc()doc";
 
 static const char *__doc_mitsuba_spectrum_from_file_2 = R"doc()doc";

include/mitsuba/render/integrator.h

@@ -158,7 +158,8 @@ class MTS_EXPORT_RENDER SamplingIntegrator : public Integrator<Float, Spectrum>
                               Sampler *sampler,
                               ImageBlock *block,
                               Float *aovs,
-                              size_t sample_count = size_t(-1)) const;
+                              size_t sample_count,
+                              size_t block_id) const;
 
     void render_sample(const Scene *scene,
                        const Sensor *sensor,