Invoke PyTorch function in a plugin method

[gerwang]

Hi,

I would like to convert a si.wi (which is a mi.Vector3f) to a mi.TensorXf so that I can pass it to a dr.wrap_ad function. However, I am facing a problem: I want to do this conversion in a method of a custom plugin, where the number of elements in si.wi is not fixed.

The following code describes my idea:

import mitsuba as mi
import drjit as dr

class PyTorchEmitter(mi.Emitter):

    @dr.wrap_ad(source='drjit', target='torch')
    def torch_func(self, v):
        # ... evaluate some differentiable functions in PyTorch
        return v

    def eval(self, si: mi.SurfaceInteraction3f, active: bool = True) -> mi.Color3f:
        v = si.wi
        v_tensor = dr.empty(mi.TensorXf, shape=(3, len(v[0])))
        v_tensor[0] = v.x
        v_tensor[1] = v.y
        v_tensor[2] = v.z
        res_tensor = self.torch_func(v_tensor)
        return dr.unravel(mi.Color3f, res_tensor)

The code actually runs, but each time the method is invoked, len(v[0]) is 1. It seems that the code picks the first element, converts it to a length 1 tensor, and broadcasts the return value.

PS: I 'm using the cuda_ad_rgb variant.

I would like to know if there is any way to invoke a PyTorch function in a plugin using data with variable lengths. Thank you in advance!

[njroussel]

Hi @gerwang

There are some limitations to using wrap_ad inside of plugins. Please refer to Dr.Jit's documentation first, to roughly understand how it works.

DrJit is a just-in-time compiler and records the computation rather than executing it sequentially. Especially for virtual function calls there is an important recording mechanism. In your example, eval will be some target of a virtual function call in your integrator. In fact, there might be thousands of emitters in your scene and we need to record every possibility.

The recorded computation is only evaluated once it needs to be used by some function outside of Dr.Jit/Mitsuba as this function would not know what to do with DrJit's recorded computation graph. Unfortunately, it is not possible to evaluate a variable everywhere, namely virtual function calls are restricted. Calling len() and wrap_ad would require an evaluation in your example code.

I am surprised you didn't get an error/warnming message regarding this. Were you just calling mi.render() or were you using your custom Emitter differently?

Instead of len() you'll want to use dr.width() it will give you the vectorization width.

[gerwang]

Thank you for your reply!

I just called mi.render and did not receive any error messages in the output. Here is a runnable code that can replicate a similar situation that I encountered. I replaced len() with dr.width() in this example.

import mitsuba as mi

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

class TinyTorchNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = torch.nn.Linear(in_features=3, out_features=64)
        self.lin2 = torch.nn.Linear(in_features=64, out_features=3)
        self.act = torch.nn.ReLU(inplace=True)
        self.act_out = torch.nn.Softplus()

    def forward(self, x: torch.Tensor):
        x = self.lin1(x)
        x = self.act(x)
        x = self.lin2(x)
        x = self.act_out(x)
        return x

class PyTorchEmitter(mi.Emitter):

    def __init__(self, props: mi.Properties):
        super().__init__(props)
        self.tiny_nn = TinyTorchNN().cuda()

    @dr.wrap_ad(source='drjit', target='torch')
    def torch_func(self, v):
        print(v.shape)
        v = v.t()
        res = self.tiny_nn(v)
        res = res.t().contiguous()
        return res * 20  # brighter

    def eval(self, si: mi.SurfaceInteraction3f, active: bool = True) -> mi.Color3f:
        v = si.wi
        v_tensor = dr.empty(mi.TensorXf, shape=dr.shape(v))
        v_tensor[0] = v.x
        v_tensor[1] = v.y
        v_tensor[2] = v.z
        res_tensor = self.torch_func(v_tensor)
        return dr.unravel(mi.Color3f, res_tensor)

    def pdf_direction(self, it, ds, active: bool = True) -> float:
        return 0.  # mock function

if __name__ == '__main__':
    mi.register_emitter('mytorch', lambda props: PyTorchEmitter(props))
    scene_dict = mi.cornell_box()
    scene_dict['light']['emitter'] = {
        'type': 'mytorch'  # replace area emitter with my emitter
    }
    scene_dict['integrator'] = {
        'type': 'direct',
        'emitter_samples': 0,  # so that we only need to implement `eval()` in emitter
        'bsdf_samples': 1,
    }
    scene = mi.load_dict(scene_dict)
    img = mi.render(scene)
    mi.util.write_bitmap('./out.png', img)

The console output is below, with no errors:

torch.Size([3, 1])
torch.Size([3, 1])

The out.png image looks like:

So the emitter seems uniform instead of spatially varying.

After reading drjit’s documentation and its paper, I guess that my plugin method is called during tracing, and then it is compiled into PTX code.

However, I really want to implement the feature of invoking PyTorch in a plugin method because it has a broad application of introducing tiny-cuda-nn into rendering. In the above example the emitter is just a planar area light, but with tiny-cuda-nn we can achieve much more.

I came up with some ideas but no success yet. Do you think any one of them is feasible?

• dr.CustomOp: I tried to wrap the vector to tensor operation and PyTorch into a dr.CustomOp, but it seems that it is just evaluated during tracing.
• Turning off megakernels and use wavefronts instead: Can it help if I turn off megakernels? I searched the drjit documentation but did not find an API.
• Move the emitter code right inside the integrator’s render or sample method. I assume there is only one emitter. I want to sidestep the polymorphism of EmitterPtr by directly inserting the PyTorch function in mitsuba3/src/integrators/direct.cpp. However, if render is another virtual function, the method seems not to work.

Would you please tell me how to achieve this feature? Thank you!

[njroussel]

Turning off megakernels is the way to go. In order to do it you need to run the two following lines before calling the render function:

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

This should hopefully work.

Thank you for the reproducer, I am seeing the same behaviour. This is really suprising, as I said before, we should have generated an error message during the call to torch_func(v_tensor). I'll look into this.