using si.wi values in sample/eval methods of a custom BSDF

[rajeshsharma-ai]

Hello!
I am trying to make use of si.wi values to do some computations in a custom BSDF but I get the
RuntimeError: jitc_var_eval(r1848): placeholder variables are used to record computation symbolically and cannot be scheduled for evaluation! This error message could appear for the following reasons error.

Attached is a working example. In this custom BSDF, I use a sidecar diffuseBSDF to get wo and pdf but want to do some custom math with the si.wi for the values to return. Thanks!

import drjit as dr
import mitsuba as mi
import matplotlib.pyplot as plt
#mi.set_variant('llvm_ad_rgb')
mi.set_variant('cuda_ad_rgb')

class MyBSDF(mi.BSDF):
    """
    Simple Custom bsdf 
    """
    def __init__(self, props):
        mi.BSDF.__init__(self, props)
        self.bsdf_dict = {
            'type': 'diffuse',
            'reflectance' : {
                'type': 'rgb',
                'value': (0.3, 0.1, 0.1)
            }
        }
        reflection_flags   = mi.BSDFFlags.DiffuseReflection|mi.BSDFFlags.FrontSide 
        self.m_components  = [reflection_flags]
        self.m_flags = reflection_flags
        # use a sidecar diffuse bsdf for eval/sample/pdf
        self.m_diffusebsdf = mi.load_dict(self.bsdf_dict)
        self.m_params = mi.traverse(self.m_diffusebsdf)

    def sample(self, ctx, si, sample1, sample2, active):
        """
        Get the wo and the pdf from the diffusebsdf but the value from the data source
        """
        bsp, val = self.m_diffusebsdf.sample(ctx, si, sample1, sample2, active)

        # populate the bs struct 
        bs = mi.BSDFSample3f()
        bs.pdf = bsp.pdf
        bs.wo = bsp.wo
        bs.eta = 1.0
        bs.sampled_type = mi.UInt32(+mi.BSDFFlags.DiffuseReflection)
        bs.sampled_component = 0

        cos_theta_o = mi.Frame3f.cos_theta(bs.wo)
        active &= cos_theta_o > 0.0

        value = self.eval(ctx, si, bs.wo, active)
        return (bs, value/bs.pdf & active)

    def eval(self, ctx, si, wo, active):
        """
        Eval the diffuse bsdf instead
        """
        cos_theta_i = mi.Frame3f.cos_theta(si.wi)
        cos_theta_o = mi.Frame3f.cos_theta(wo)

        active &= (cos_theta_i > 0.0)
        active &= (cos_theta_o > 0.0)

        # do some custom computation 
        val = self.evalCustom(ctx, si, wo, active)

        return dr.select(active, val,  mi.Color3f(0))

    def pdf(self, ctx, si, wo, active):
        # get the pdf from the diffusebsdf bsdf
        pdf = self.m_diffusebsdf.pdf(ctx, si, wo, active)
        return pdf

    def eval_pdf(self, ctx, si, wo, active):
        f = self.eval(ctx,si,wo,active)
        pdf = self.pdf(ctx,si,wo,active)
        return (dr.select(active, f, mi.Vector3f(0.)), dr.select(active, pdf, 0.))
   
    def evalCustom(self, ctx, si, wo, active):

        values = dr.zeros(mi.Color3f, dr.width(si.wi))

        # this seems to work fine
        values[0] = 0.7 
        values[1] = 0.2
        values[2] = 0.2

        # when I try to make use of si.wi, I get the 
        '''
        RuntimeError: jitc_var_eval(r1848): placeholder variables are used to record computation symbolically
        and cannot be scheduled for evaluation! This error message could appear for
        the following reason...
        Ultimately, I would like to do Tensor-based operations involving si.wi and even use data read from file
        '''
        #values[0] = si.wi.x * 0.7
        #values[1] = si.wi.y * 0.2
        #values[2] = si.wi.z * 0.2

        self.m_params['reflectance.value'] = values 
        return self.m_diffusebsdf.eval(ctx, si, wo, active)
        
    def to_string(self,):
        return 'MyBSDF'

# register and load the custom bsdf
mi.register_bsdf("mybsdf", lambda props: MyBSDF(props))
# load it
mybsdf = mi.load_dict({
    'type' : 'mybsdf',
})

# simple scene
scene = mi.load_dict({
    'type': 'scene',
    'integrator': {
        'type': 'path'
    },
    'light': {
        'type': 'constant',
    },
    'sphere' : {
        'type': 'sphere',
        'bsdf': mybsdf
    },
    'sensor': {
        'type': 'perspective',
        'to_world': mi.ScalarTransform4f.look_at(origin=[0, -5, 5],
                                                 target=[0, 0, 0],
                                                 up=[0, 0, 1]),
    }
})

img = mi.render(scene, spp=64)
print(img)
plt.imshow(img ** (1.0 / 2.2))
plt.show()

[njroussel]

Hi @rajeshsharma-ai

The issue here is that you're using the traverse mechanism to update a plugin during a render call. Plugins' are typically considered immutable during a render. There are too many potential race conditions that could happen here otherwise: what if one thread updates the parameters, but by the time it reaches the eval() statement some other thread updates the plugin.

Depending on the operation that you're doing in evalCustom, you might be able to just modify the input or output of the inner BSDF evaluation without modifying its parameters. However, if the operation requires changing the BSDF's state, you will need to get rid of the inner BSDF and re-implement that logic in your custom BSDF (without the dependency on the state).