Performance issue with custom plugin

[imanolooo]

Hi all,

I want to implement a version of the normal map plugin to handle differently the primary rays converted to the back face after perturbing the shading frame with the normal map. I have this code working in the cpp version of mitsuba3, but I want it to have it as a custom plugin to use it without recompiling mitsuba.

As a first step, I have reimplemented the Mitsuba3 normalmap plugin directly in python without any modification. I think that my implementation is correct (not completely sure). The results seem correct too. However, I have a huge performance difference between using my normal map plugin or the provided by mitsuba. For a simple scene with a textured square, the difference are 65.01 seconds for my plugin, 0.14 seconds for mitsuba's one... It seems too much for me, maybe my code is not parallelizing well or is this difference normal?

Thanks a lot!!!

Imanol

Version: 3.6.0 installed using pip
System: Ubuntu 22.04 through WSL2 and Windows 10

I attach also the code:
TestNormalMap.py

import mitsuba as mi
import drjit as dr

import time

from MyNormalMap import MyNormalMap

mi.set_variant("scalar_rgb")

mi.register_bsdf("mynormalmap", lambda props: MyNormalMap(props))

scene = mi.load_dict({
    'type': 'scene',
    'integrator': {
        'type': 'path'
    },
    'light': {
        'type': 'constant',
        'radiance': 0.99,
    },
    'rectangle': {
        'type': 'rectangle',
        'bsdf': {
            'type': 'mynormalmap',#'normalmap',
            'normalmap': {
                'type': 'bitmap',
                'raw': True,
                'filename': 'textures/normalmap.png'
            },
            'bsdf': {
                'type': 'diffuse',
                'reflectance': {
                    'type': 'rgb',
                    'value': [0.8, 0.2, 0.2]  # Diffuse red color
                }
            }
        }
    },
    'sensor': {
        'type': 'perspective',
        'to_world': mi.ScalarTransform4f().look_at(origin=[0, 0, 5],
                                                 target=[0, 0, 0],
                                                 up=[0, 1, 0]),
    }
})

print("Start rendering...")
start_time = time.perf_counter()
image = mi.render(scene)
end_time = time.perf_counter()
elapsed_time = end_time-start_time
print(f"Done in {elapsed_time} seconds.")

mi.util.write_bitmap('rectangle_scene.png', image)

MyNormalMap.py

import mitsuba as mi
import drjit as dr

mi.set_variant("scalar_rgb")

class MyNormalMap(mi.BSDF):
    def __init__(self, props):
        mi.BSDF.__init__(self, props)

        # Retrieve the nested BSDF child object
        self.m_nested_bsdf = props.get('bsdf', None)
        if self.m_nested_bsdf is None or not isinstance(self.m_nested_bsdf, mi.BSDF):
            raise RuntimeError("Exactly one BSDF child object must be specified.")

        # Ensure the normal map is a texture
        self.m_normalmap = props.get('normalmap', None)
        if self.m_normalmap is None or not isinstance(self.m_normalmap, mi.Texture):
            raise RuntimeError("A 'normalmap' texture must be specified.")

        # Add all nested components
        self.m_flags = 0
        self.m_components = []
        for i in range(self.m_nested_bsdf.component_count()):
            flags = self.m_nested_bsdf.flags(i)
            self.m_components.append(flags)
            self.m_flags |= flags

    def sample(self, ctx, si, sample1, sample2, active):
        # Sample the nested BSDF with a perturbed shading frame
        perturbed_si = mi.SurfaceInteraction3f(si)
        perturbed_si.sh_frame = self.frame(si, active)
        perturbed_si.wi = perturbed_si.to_local(si.wi)

        bs, weight = self.m_nested_bsdf.sample(ctx, perturbed_si, sample1, sample2, active)

        # Update the active mask based on the weight
        active &= dr.any(mi.unpolarized_spectrum(weight) != 0.0)
        if not dr.any(active):
            return bs, 0.0

        # Transform the sampled 'wo' back to the original frame and verify orientation
        perturbed_wo = perturbed_si.to_world(bs.wo)
        active &= (mi.Frame3f.cos_theta(bs.wo) * 
                mi.Frame3f.cos_theta(perturbed_wo)) > 0.0
        bs.wo = perturbed_wo

        return bs, weight & active

    def eval(self, ctx, si, wo, active):
        # Evaluate nested BSDF with perturbed shading frame
        perturbed_si = mi.SurfaceInteraction3f(si)
        perturbed_si.sh_frame = self.frame(si, active)
        perturbed_si.wi = perturbed_si.to_local(si.wi)
        perturbed_wo = perturbed_si.to_local(wo)

        # Adjust active mask based on cosine terms
        active &= (mi.Frame3f.cos_theta(wo) * 
                mi.Frame3f.cos_theta(perturbed_wo)) > 0.0

        # Evaluate the nested BSDF
        return self.m_nested_bsdf.eval(ctx, perturbed_si, perturbed_wo, active) & active
    
    def pdf(self, ctx, si, wo, active):
        # Evaluate nested BSDF with perturbed shading frame
        perturbed_si = mi.SurfaceInteraction3f(si)
        perturbed_si.sh_frame = self.frame(si, active)
        perturbed_si.wi = perturbed_si.to_local(si.wi)
        perturbed_wo = perturbed_si.to_local(wo)

        # Adjust active mask based on cosine terms
        active &= (mi.Frame3f.cos_theta(wo) * 
                mi.Frame3f.cos_theta(perturbed_wo)) > 0.0

        # Evaluate and return the PDF of the nested BSDF
        return mi.select(active, self.m_nested_bsdf.pdf(ctx, perturbed_si, perturbed_wo, active), 0.0)

    def eval_pdf(self, ctx, si, wo, active):
        # Profiler phase (can be omitted if unnecessary in Python)
        #mi.MaskProfiler.phase(mi.ProfilerPhase.BSDFEvaluate, active)

        # Evaluate nested BSDF with perturbed shading frame
        perturbed_si = mi.SurfaceInteraction3f(si)
        perturbed_si.sh_frame = self.frame(si, active)
        perturbed_si.wi = perturbed_si.to_local(si.wi)
        perturbed_wo = perturbed_si.to_local(wo)

        # Adjust active mask based on cosine terms
        active &= (mi.Frame3f.cos_theta(wo) * 
                mi.Frame3f.cos_theta(perturbed_wo)) > 0.0

        # Evaluate the nested BSDF's eval_pdf
        value, pdf = self.m_nested_bsdf.eval_pdf(ctx, perturbed_si, perturbed_wo, active)

        # Apply the active mask to the output
        value = value & active
        pdf = dr.select(active, pdf, 0.0)

        return value, pdf

    def traverse(self, callback):
        callback.put_parameter('nested_bsdf', self.m_nested_bsdf, mi.ParamFlags.Differentiable)
        callback.put_parameter('normalmap', self.m_normalmap, mi.ParamFlags.Differentiable | mi.ParamFlags.Discontinuous)

    def parameters_changed(self, keys):
        print("🏝️ there is nothing to do here 🏝️")

    def to_string(self):
        return ('MyNormalMap[\n'
                '    nested_bsdf=%s,\n'
                '    normal_map=%s,\n'
                ']' % (self.m_nested_bsdf, self.m_normalmap))

    def frame(self, si, active):
        # Evaluate the normal map texture and scale it to the range [-1, 1]
        n = dr.fma(self.m_normalmap.eval_3(si, active), 2.0, -1.0)

        # Create a new frame and normalize the perturbed normal
        result = mi.Frame3f()
        result.n = dr.normalize(n)

        # Compute the tangent and bitangent vectors
        result.s = dr.normalize(dr.fma(-result.n, dr.dot(result.n, si.dp_du), si.dp_du))
        result.t = dr.cross(result.n, result.s)
        
        return result

    def eval_diffuse_reflectance(self, si, active):
        return self.m_nested_bsdf.eval_diffuse_reflectance(si, active)```

[merlinND]

Hello @imanolooo,

Custom Python plugins require using vectorized variants (e.g. llvm_ad_rgb, cuda_ad_rgb) for good performance. Please see the documentation:

Custom Mitsuba plugins written in Python work best with Just-In-Time (JIT) variants. This is because it would be pretty inefficient to execute Python BSDF code for millions of scalar light paths. JIT variants, on the other hand, only execute a few calls to those methods on arrays containing millions of entries at once, mitigating the overhead coming from the the Python layer.