drjit
drjit copied to clipboard
Kernel freezing frontend
(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
dataclass
es - 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