drjit icon indicating copy to clipboard operation
drjit copied to clipboard

Kernel freezing frontend

Open merlinND opened this issue 11 months ago • 1 comments

(As discussed with @wjakob and @DoeringChristian)

This code is not intended to work as-is, but rather to inform the kind of cases we need to check for, and how to do the Python function --> frozen function inputs / outputs mapping while being aware of aliasing, literal, etc. A fair bit of the code is there to detect and prevent unsupported cases.

I opened this PR mostly to make it easy to read & comment on the code, but we can also use it to iterate on the frontend.

Some assumptions made by this implementation and that would be nice to have in the final version as well:

  • PyTrees can be and / or contain dataclasses
  • There is a dr.JitFlag.KernelFreezing to globally disable any kernel freezing (very useful, if only to make sure that we are not recording frozen kernels launches while recording frozen kernels)

A few things that will change for sure in the final implementation:

  • Move the PyTree walking to C++
  • No longer rely on KernelHistory, but rather use the virtualized backend from https://github.com/mitsuba-renderer/drjit-core/pull/83

merlinND avatar Mar 25 '24 09:03 merlinND