Storing each bounce in a buffer in vectorized variants

[ZenithGD]

Hello everyone,

I'm currently implementing an integrator on Mitsuba 3 (stable branch) where I need to store each of the individual bounces of each sampled path and then access that information afterwards. I'm having a bit of trouble understanding how to store this information and use it, specially in the context of recorded loops, and I've thought of some possible solutions but I thought I needed some external help as well.

Basically, the integrator's sample method consists of two phases that I implemented in two separate, private methods. I will gloss over the details of the algorithm to keep it simple, but if you need more details let me know (I'm currently trying to implement the sample-solve algorithm from Steinberg 2023 if this helps).

• The first phase is essentially identical to a path tracer integrator without accounting for NEE. Here, I just need to store each of the bounces in some data structure, which i will denote as the bounce buffer. The paths will be terminated on a light source as usual.
• The second phase uses the information stored in the bounce buffer, retracing the path from the light source, so I need to read the bounces in the reverse order at which they were added.

I implemented the following custom structure in C++ for storing each individual bounce, keeping in mind that the underlying data is vectorized, so not all of the components may be active, therefore I included a Mask as well to indicate which lanes are active and store a valid bounce. I will also omit some of the fields for brevity:

/**
 * \brief Structure that stores one bounce of some path traversing the scene.
 * 
 * A bounce data object contains all the necessary information to compute both forwards and backwards light transport.
 */
template <typename Float_, typename Spectrum_>
struct BounceData {
    using Float    = Float_;
    using Spectrum = Spectrum_;
    
    MI_IMPORT_TYPES()
    MI_IMPORT_OBJECT_TYPES()

    /// \brief The id of the vertex in the path (mainly for debug purposes)
    UInt32 id;

    /// \brief Information about the interaction at this vertex
    SurfaceInteraction3f interaction;
    
    //...
    
    /// \brief Whether this vertex of the path belongs to an active path.
    Mask active;
    
    BounceData(
        const UInt32 id_,
        const SurfaceInteraction3f& it_, ... , const Mask& active_)

        : id(id_), interaction(it_), ..., active(active_) {}

    DRJIT_STRUCT(BounceData, id, interaction, ..., active);
};

I've also exposed the data structure to the Python API with Pybind11. The integrator's sampling function and the two phases can be summarized in the following schema:

def plt_sample_phase(self, ...):
        
        # define internal state
        # Here's where the bounces are stored.
        bounce_buffer = ?
        ...

        # Note: I'm not sure if the bounce_buffer should be part of the loop state...
        loop = mi.Loop("PLT Path tracer (sample phase)", 
                        state = lambda : (...))
        
        loop.set_max_iterations(self.max_depth)

        while loop(active):
            # Sampling logic is identical to the path tracer integrator, omitting NEE and MIS.
            # It also updates the active mask as the path tracer does.
            # ...

            # store bounce data
            bounce = mi.BounceData(id=i, interaction=si, active=active, ...)
            
            # The buffer could have this signature for adding bounces
            bounce_buffer.add_bounce(depth, bounce)
        
        return bounce_buffer, ...
        
 def plt_solve_phase(self, ..., bounce_buffer):
        
        # Solve data
        L = mi.Spectrum(0.0)
        i = mi.UInt(0)

        loop = mi.Loop(name="PLT Path tracer (solve phase)", state=lambda:(?))
        
        loop.set_max_iterations(self.max_depth)

        while loop(i < self.max_depth):

            # read a bounce here
            bounce = bounce_buffer.read_bounce(i)
            
            # do something with the recovered bounce
            # ...

            # next bounce
            i += 1

        return (L, ...)

def sample(...):
    bounce_buffer, ... = self.plt_sample_phase(...)
  
    L = self.plt_solve_phase(..., bounce_buffer)
    return (L, ...)

I'm having trouble implementing a structure that stores these bounces. I've tried to use conventional Python lists, or DynamicArray in the C++ side, but I haven't found a solution that works.

Sorry for the big snippets of code, if you have any ideas, let me know, i will highly appreciate it!

[mcrescas]

Hi @ZenithGD ,

Usually for creating these types of buffers, you need to reserve an amount of memory beforehand. In your case, that could be the maximum number of bounces.

After deciding, I would recommend you take a look into a new feature of Dr.JIT named "Local Memory" which allows to have a buffer of fixed size per thread. You can imagine that it is like local storage per thread (in the CPU context).
With that, you could use the depth of the integrator to index in that buffer to retrieve the information saved on it and it will be more performant than previous solutions that had to reserve memory for the entire wavefront.
Just a final note: I am assuming that both the first and second phases are computed in the same kernel of Dr.JIT.

[ZenithGD]

Hello @mcrescas ,

First of all, thanks for the quick reply! I appreciate it a lot :)

Concerning your recommendation, I will consider updating and migrating the changes I integrated already into the stable version of Mitsuba into the new 3.6.0 version. Local memory seems like the best solution here, avoiding the allocation of the buffer for the entire wavefront as you said, and I probably will also benefit from the new dr.JIT features as well. And yes, the first and second phases are within the same kernel, i'm running both phases inside the sampling function of the integrator.

However, i'm still on the stable branch of Mitsuba (version 3.5.2), and the dr.JIT version is still 0.4.6, which I think doesn't have the local memory features, so I was a bit hesitant of updating to the latest development version. Do you think there is any other way that doesn't really use local memory, at the expense of allocating memory for the whole wavefront, at least as a provisional solution while I work on the migration to Mitsuba 3.6.0? Anyways, I will keep you updated if I manage to implement the local memory approach.

Thanks!

[ZenithGD]

Hello again,

I finally managed to implement and migrate all the changes into the latest Mitsuba 3.6 dev version, and the local memory approach worked like a charm. Thank you again for your help!