Invoke pytorch code in recorded loop

[KoykL]

With current version of mitsuba 3, is it possible to call pytorch code in recorded loop? This is useful for e.g. performing sphere tracing on SDF represented with a pytorch neural network. Thank you!

[njroussel]

Hi @KoykL

You'll want to take a look at this thread too: #586
That discussion is specific to virtual function calls but they share the same limitations to recorded loop. It might give you some more intuition.

In short, the answer is no, it is not possible to call pytorch code in a drjit recorded loop. As its name suggests, the purpose of those loops is to record the computation in their body. However any operation which doesn't need drjit types or function will not be recorded as drjit is simply unaware of it.

The only workaround is to disable all recording features with:

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

or to use plain Python loops.

[KoykL]

Thank you for your reply! I am just thinking would it be possible to implement something like https://www.tensorflow.org/api_docs/python/tf/compat/v1/py_func from tensorflow 1.x into mitsuba? It would be nice to be able to incorporate python code into mitsuba computation graph for neural network based graphics primitives.

The generated code can call python c api to evaluate the wrapped python code.

[KoykL]

By the way, if I turn off all recording features, it won't support dynamic criteria for termination (e.g. when all rays have very small sdf value.) Is that correct?

[njroussel]

Theoretically it seems possible to weave a function pointer into a computations graph, I agree. However, in practice, the Optix/CUDA kernels we launch have certain limitations. We cannot arbitrarily insert any function into such a kernel.
The workaround is to split the kernels: have it stop right before the pytorch/tensorflow/etc. function and start again right after the function returned. This is what "wavefront mode" does and it is enabled by disabling recording features with the two lines I mentioned above.

No, that is incorrect. Recording has no influence on the result of any computation, it is simply a performance consideration. With recording features, we are able to generate a single (big) kernel for an entire rendering pipeline (we call these "megakernels"). Without recording features, multiple smaller kernels are generated which typically require more memory and computation time than a single megakernel.

[KoykL]

Thank you so much for your detailed explanation! It is very helpful.

[wjakob]

@njroussel Since it seems to be a recurring question, it might make sense to have a thorough FAQ entry explaining why the default "PyTorch in recorded loop" doesn't work, and what the alternatives are.