Multiple Discrete1D distributions

[saeedhd96]

Hello,

Is there a way to define a data structure for multiple 1D discrete distirbutions and sample from them without looping through the distributions? Essentially, that means in the wavefront mode, per ray, we have a 1D distribution with a fixed number of elements that we want to sample from, again per ray.

[DoeringChristian]

Hi @saeedhd96,

I think that should be possible. Here is a little reproducer. It only supports two distributions at the moment but if you can stack your distribution it should work as well.

if __name__ == "__main__":
    cdf1 = mi.Float(0.1, 0.3, 0.4, 0.6, 0.8, 1.0)
    cdf2 = mi.Float(0.1, 0.2, 0.6, 0.7, 0.9, 1.0)

    sample = mi.Float(0.6, 0.6)

    cdf = dr.zeros(mi.Float, len(cdf1) + len(cdf2))

    dr.scatter(cdf, cdf1, dr.arange(mi.UInt, len(cdf1)))
    dr.scatter(cdf, cdf2, dr.arange(mi.UInt, len(cdf2)) + len(cdf1))

    start = mi.UInt(0, len(cdf1))

    idx = dr.binary_search(
        0, len(cdf1), lambda idx: dr.gather(mi.Float, cdf, idx + start) < sample
    )

    print(f"{idx=}")

Another option would be to implement your own binary search, where you can specify the start and end as Dr.Jit arrays.
This way it could also support differently sized distributions.
Something like this:

import mitsuba as mi
import drjit as dr


if __name__ == "__main__":
    mi.set_variant("llvm_ad_rgb")


def binary_search(start: mi.UInt, end: mi.UInt, pred):
    iterations = dr.select(start < end, dr.log2i(end - start) + 1, 0)
    index = dr.zeros(mi.UInt, dr.width(pred(start)))
    start_ = mi.UInt(start)

    loop = mi.Loop("Binary Search", lambda: (start_, end, index))

    while loop(index < iterations):
        middle = (start_ + end) >> 1
        cond = dr.detach(pred(middle))

        start_ = dr.select(cond, dr.minimum(middle + 1, end), start_)
        end = dr.select(cond, end, middle)
        index += 1

    return start_ - start


if __name__ == "__main__":
    cdf1 = mi.Float(0.1, 0.3, 0.4, 0.6, 1.0)
    cdf2 = mi.Float(0.1, 0.2, 0.6, 0.7, 0.9, 1.0)

    sample = mi.Float(0.6, 0.6)

    cdf = dr.zeros(mi.Float, len(cdf1) + len(cdf2))

    dr.scatter(cdf, cdf1, dr.arange(mi.UInt, len(cdf1)))
    dr.scatter(cdf, cdf2, dr.arange(mi.UInt, len(cdf2)) + len(cdf1))

    start = mi.UInt(0, len(cdf1))
    end = mi.UInt(len(cdf1), len(cdf1) + len(cdf2))

    idx = binary_search(start, end, lambda idx: dr.gather(mi.Float, cdf, idx) < sample)

    print(f"{idx=}")

Hope this could help.

[njroussel]

💯

For completeness, let me add that you can retrieve the individual CDFs directly from constructed DiscreteDistribution objects:
https://mitsuba.readthedocs.io/en/stable/src/api_reference.html#mitsuba.DiscreteDistribution.cdf

[saeedhd96]

Thanks for your response.
I see how this could work but I have three issues:

1. The second solution still has a recorded python loop which is going to be slow in the wavefront mode, no?
2. I could not compute the CDFs from my PDFs without having another python loop, because drjit has no cumsum. I cannot write a recorded loop to compute cumsum:

import mitsuba as mi
import drjit as dr
mi.set_variant("cuda_ad_rgb")


size = 5
pdf = mi.Float([0.1, 0.2, 0.3, 0.4, 0.5, 0.1, 0.2, 0.3, 0.4, 0.5])

iterations = dr.width(pdf)
index = dr.zeros(mi.UInt32, iterations)
size = mi.UInt32(size)
cdf = dr.zeros(mi.Float, dr.width(pdf))
loop = mi.Loop("Cumsum", lambda: (index, cdf))

while loop(index<iterations):
#    d = data[index]
    d = dr.gather(mi.Float, pdf, index)
    prev_sum = dr.gather(mi.Float, cdf, index-1)
    a = dr.select(index%size!=0,prev_sum, 0)
    dr.scatter(cdf, d+prev_sum, index)
    index+=1

This raises an error 'jit_var_gather(): cannot gather from a placeholder variable!'

3. @njroussel I know DiscreteDistribution, but how can I use a million of DiscreteDistribution for a million of distributions which is the case here? I don't think that is supported currently.

[njroussel]

1. You can locally enable/disable recorded loops. Just call jit_set_flag() before your loop and after it. It does not have to be a project-wide constraint. For this specific snippet (the binary search), I don't see how/why you would need this be in wavefront-mode.

2. Do you really need to do this in drjit ? This is a sequential operation anyway, you can't really parallelize it. Just use numpy if it can make your life easier. Anyway, here's a fixed snippet:

import mitsuba as mi
import drjit as dr
mi.set_variant("cuda_ad_rgb")


size = 5
pdf = mi.Float([0.1, 0.2, 0.3, 0.4, 0.5, 0.1, 0.2, 0.3, 0.4, 0.5])

iterations = dr.width(pdf)
index = dr.zeros(mi.UInt32, iterations)
size = mi.UInt32(size)
cdf = dr.zeros(mi.Float, dr.width(pdf))

prev_sum = mi.Float(0)
loop = mi.Loop("Cumsum", lambda: (index, prev_sum))
while loop(index<iterations):
    d = dr.gather(mi.Float, pdf, index)

    prev_sum = dr.select(dr.neq(index % size, 0), prev_sum, 0)
    cum_sum = d + prev_sum
    dr.scatter(cdf, cum_sum, index)
    prev_sum = cum_sum

    index+=1

print(f"{cdf=}")

If all of your distributions have the same size you could actually compute their cumsum in parallel.

3. Oh, sorry, I didn't mean that this was possible. @DoeringChristian's suggestion are still the appropriate solution. My point was simply that you could have your individual distributions stored in DiscreteDisctribution objects such that they will handle your CDF computation for example. In a second step, you would then merge them into a large flat buffer by accesssing each one with DiscreteDistribution.cdf.

[saeedhd96]

Thanks for your great response. I really appreciate it. You are amazing @njroussel and @DoeringChristian.

The reason I want to do this in drjit and not numpy or torch is because I think the sandwiching (interoperability of Mitsuba and these) is time consuming.

Here is a notebook in which I tried to measure this operation using torch, numpy, drjit sequentially, drjit parellel, and using the mega kernel on and off:

Mega off:

jit_parallel: 0.003554387092590332
jit_sequential: forever
torch: 0.0005819940567016601
numpy: 0.04193118572235108

Mega on:

jit_parallel: 0.00019636154174804688
jit_sequential: 0.0002063131332397461
torch: 0.0005804347991943359
numpy: 0.04199970245361328

Not sure if my measurements are right, but I tried to synchronize everything and then measure.
Notebook: Link
If you tried to run this, I think a fair practice is to measure each in a separate run to avoid using cached stuff.

[DoeringChristian]

Hi @saeedhd96,
I think I understand what you mean, so I'm not sure if this is going to help but I have implemented the fast prefix-sum/cumsum in mitsuba3 some time ago. You might want to test it a bit to make sure it works.

def prefix_sum(x: mi.Float) -> mi.Float:
    """
    Implementation of a paralell prefix-sum described in

    W. Daniel Hillis and Guy L. Steele. 1986. Data parallel algorithms. Commun. ACM 29, 12 (Dec. 1986), 1170–1183. https://doi.org/10.1145/7902.7903
    """
    x = type(x)(x)

    loop_record = dr.flag(dr.JitFlag.LoopRecord)
    dr.set_flag(dr.JitFlag.LoopRecord, False)

    n = dr.shape(x)[-1]
    i = 0

    loop = mi.Loop("prefix-sum", lambda: ())

    while loop(i <= math.floor(math.log2(n))):
        j = dr.arange(mi.UInt, 2**i, n)
        if dr.shape(j)[-1] == 0:
            break
        res = dr.gather(type(x), x, j) + dr.gather(type(x), x, j - 2**i)
        dr.scatter(x, res, j)

        i += 1

    dr.set_flag(dr.JitFlag.LoopRecord, loop_record)

    return x

To be able to calculate multiple prefix sums you could have multiple pdfs in one array and play around with the indices. It also has to run in wavefront mode unlike the solution you and @njroussel provided. Maybe you could also implement something like that inline if you are trying to compute millions of cdfs.