Wavefronts in Texture.eval()

[saeedhd96]

Hello,

I realized that the wavefront size in Texture.eval() shrinks, in other words, the vectors that are used in an integrator, when going through m_reflectance.eval(), shrink to smaller vectors. Although my guess is that it is shrunk due to removing inactive lanes, I cannot find where this happens in the code. Can you please refer me to the peice of code that causes this behavior?

Thanks,
-Saeed

[njroussel]

Hi @saeedhd96

Have you disabled megakernels (VCallRecord specifically)? The kernel should not split and change sizes.
If you can send a tiny reproducer I will hopefully be able to point you to what is happening in the code.

[saeedhd96]

Code to reproduce:

import mitsuba as mi
import drjit as dr
import torch
mi.set_variant('cuda_ad_rgb')

class MyTexture(mi.Texture):
    def __init__(self, props: mi.Properties) -> None:
        mi.Texture.__init__(self, props)

    def eval(self, si, active):
        print(si.p.torch().shape)
        return si.p*0
    
mi.register_texture('my_texture', MyTexture)


bsdf = {'type': 'diffuse', 'reflectance': {'type':'my_texture'}}
integrator = {'type': 'path', 'max_depth': 2}

my_cornell = mi.cornell_box().copy()
my_cornell['red'] = bsdf
my_cornell['integrator'] = integrator
my_cornell['sensor']['film']['width'] = 256
my_cornell['sensor']['film']['height'] = 256

scene = mi.load_dict(my_cornell)

dr.set_flag(dr.JitFlag.LoopRecord, False)
dr.set_flag(dr.JitFlag.VCallRecord, False)
dr.set_flag(dr.JitFlag.VCallOptimize, False)


img = mi.render(scene, spp = 1)
#The printed numbers do not match the expected size: 256*256

[njroussel]

Thank you for the fantastic reproducer!

Indeed, this is expected. As you mentionned, any lane/thread that is not evaluating the red BSDF is set to be inactive hence the changing sizes. (This is only true with VCallRecord turned off).

There's a lot going on here, so let me point you to two key areas:

1. This line (BSDFPtr bsdf = si.bsdf(ray);) in the integrator will essentially fetch a pointer per lane. Some of the lanes will obviously have nullptrs as they didn't hit anything.
2. This snippet in DrJit is the code which is run when a non-recorded virtual function call is used. The idea is that we run each BSDF's eval method and we gather() the parameters from the appropriate indices where all lanes have the same BSDF pointer. The notion of bucket is used to indicate all lanes that have the same target function.

Hope this helps.

[saeedhd96]

Is there a way to disable this feature while still being in wavefront mode? Somehow make sure that the vector shape remains the same always? I need to be able to map vector entires outside the texture to the entires inside, 1 to 1, but because of this shrinking I can't. I tried to make a mask using dr.eq(bsdf, None) to replace every inactive bsdf with an active one, but this effect still seems to happen, maybe in other layers-- for example insie bsdf.sample() or bsdf.eval_pdf() something happens that even when all bsdfs are active, the vectors shrink.

[njroussel]

Sorry for the delay.
Unfortunately, I don't think this can be changed. My suggestion would be to carry some extra metadata per ray (you can maybe even use some unused field like ray.time so that you don't have to change any interface) which you'll use to do the mapping inside the texture.