@@ -6,8 +6,11 @@ Function Freezing
6
6
=================
7
7
8
8
This feature is still experimental, and we list a number of unsupported cases
9
- in the :ref: `pitfalls ` section. If you encounter any issues please feel free to
10
- open an issue `here <https://github.com/mitsuba-renderer/drjit/issues >`__.
9
+ in the :ref: `pitfalls ` section. This feature also only supports a subset of the
10
+ operations, that can be performed with Dr.Jit, we list them in the
11
+ :ref: `unsupported_operations ` section. If you encounter any issues please feel
12
+ free to open an issue `here
13
+ <https://github.com/mitsuba-renderer/drjit/issues> `__.
11
14
12
15
Introduction
13
16
------------
@@ -23,16 +26,19 @@ default using a hash of the assembled IR code. As mentioned in the :ref:`_eval`
23
26
page, changing literal values can cause re-compilation of the kernel and result
24
27
in a significant performance bottleneck. However, the first two steps of
25
28
tracing the Python code and generating the intermediary representation can
26
- still be expensive. This feature tries to address this performance bottleneck,
27
- by introducing the :py:func: `drjit.freeze ` decorator. If a function is
28
- annotated with this decorator, Dr.Jit will try to cache the tracing and
29
- assembly steps as well. When a frozen function is called the first time, Dr.Jit
30
- will analyze the inputs, and then trace the function once, capturing all
31
- kernels lauched. On subsequent calls to the function Dr.Jit will try to find
32
- previous recordings with compatible input layouts. If such a recording is
33
- found, it will be launched instead of re-tracing the function. This skips
34
- tracing and assembly of kernels, as well as compilation, reducing the time
35
- spent not executing kernels.
29
+ still be expensive. When a lot of Python code has to be traced, such as custom
30
+ Python functions, the GIL has to be locked multiple times. Similarely, when
31
+ tracing virtual function calls of many instances of custom plugins, these
32
+ functions can cause a large performance overhead. This feature tries to address
33
+ this performance bottleneck, by introducing the :py:func: `drjit.freeze `
34
+ decorator. If a function is annotated with this decorator, Dr.Jit will try to
35
+ cache the tracing and assembly steps as well. When a frozen function is called
36
+ the first time, Dr.Jit will analyze the inputs, and then trace the function
37
+ once, capturing all kernels lauched. On subsequent calls to the function Dr.Jit
38
+ will try to find previous recordings with compatible input layouts. If such a
39
+ recording is found, it will be launched instead of re-tracing the function.
40
+ This skips tracing and assembly of kernels, as well as compilation, reducing
41
+ the time spent not executing kernels.
36
42
37
43
.. code-block :: python
38
44
@@ -124,11 +130,11 @@ by saving the layout of the output returned when recording the function. Since
124
130
the output has to be constructed, only a subset of traversable variables can be
125
131
returned from frozen functions. This includes:
126
132
127
- - JIT and AD variables
128
- - Dr.Jit Tensors and Arrays
129
- - Python lists, tuples and dictionaries
130
- - Dataclasses
131
- - ``DRJIT_STRUCT `` annotated classes with a default constructor
133
+ - JIT and AD variables.
134
+ - Dr.Jit Tensors and Arrays.
135
+ - Python lists, tuples and dictionaries.
136
+ - Dataclasses i.e. classes annotated with `` @dataclass ``.
137
+ - ``DRJIT_STRUCT `` annotated classes with a default constructor.
132
138
133
139
The following example shows an unsupported return type, because the constructor
134
140
of ``MyClass `` expects a variable.
@@ -197,17 +203,20 @@ then equivalent to the following function.
197
203
.. code-block :: python
198
204
199
205
def func (y ):
206
+ # The isolate grad scope is added implicitly by the freezing decorator
200
207
with dr.isolate_grad():
201
208
# Some differentiable operation...
202
209
z = dr.mean(y)
203
210
# Propagate the gradients to the input of the function...
204
211
dr.backward(z)
205
212
213
+ .. _unsupported_operations :
214
+
206
215
Unsupported Operations
207
216
----------------------
208
217
209
218
Since frozen functions record kernel launches and have to be able to replay
210
- them later, certian operations are not supported inside frozen functions .
219
+ them later, certian operations are not supported by them .
211
220
212
221
Array Access
213
222
~~~~~~~~~~~~
@@ -566,6 +575,26 @@ tensor array can be calculated without involving the first dimension.
566
575
Textures
567
576
~~~~~~~~
568
577
578
+ Textures can be used inside of frozen functions for lookups, as well as for
579
+ gradient calculations. However because they require special memory operations
580
+ on CUDA , it is not possible to update or initialize CUDA textures inside of
581
+ frozen functions.
582
+
583
+ .. code- block:: python
584
+
585
+ @ dr.freeze
586
+ def func(tex: Texture1f, pos: Float):
587
+ return tex.eval(pos)
588
+
589
+ tex = Texture1f([2 ], 1 )
590
+ tex.set_value(t(0 , 1 ))
591
+
592
+ pos = dr.arange(Float, 4 ) / 4
593
+
594
+ # The texture can be evaluated inside the frozen function.
595
+ func(tex, pos)
596
+
597
+
569
598
Virtual Function Calls
570
599
~~~~~~~~~~~~~~~~~~~~~
571
600
0 commit comments