Implemented auto opaque feature

Author: DoeringChristian

drjit/__init__.py

@@ -2414,6 +2414,7 @@ def freeze(
     max_cache_size: Optional[int] = None,
     warn_recording_count: int = 10,
     backend: Optional[JitBackend] = None,
+    auto_opaque: bool = True,
 ) -> Callable[[F], F]:
     """
     Decorator to freeze a function for replaying kernels without re-tracing.
@@ -2460,11 +2461,6 @@ def freeze(
     y = dr.gather(type(x), x, dr.width(x)//2)
     ```
 
-    Similarly, calculating the mean of a variable relies on the number of entries,
-    which will be baked into the frozen function. To avoid this, we suggest
-    supplying the number of entries as a Dr.Jit literal in the arguments to the
-    function.
-
     Args:
         f: The function to be frozen
 
@@ -2485,6 +2481,12 @@ def freeze(
           frozen function, the backend used has to be specified using this argument.
           It has to be the same backend that is used for variables inside the function,
           otherwise recording will fail and variables may be leaked.
+
+          auto_opaque: (bool): If this flag is set true and only literal values
+          or their size changes between calls to the function, these variables
+          will be marked and made opaque. This reduces the memory usage, traversal
+          overhead, and can improved the performance of generated kernels.
+          If the flag is set to false, all input variables will be made opaque.
     """
 
 
@@ -2496,6 +2498,7 @@ def freeze(
     max_cache_size: Optional[int] = None,
     warn_recording_count: int = 10,
     backend: Optional[JitBackend] = None,
+    auto_opaque: bool = True,
 ) -> F: ...
 
 
@@ -2506,6 +2509,7 @@ def freeze(
     max_cache_size: Optional[int] = None,
     warn_recording_count: int = 10,
     backend: Optional[JitBackend] = None,
+    auto_opaque: bool = True,
 ) -> Union[F, Callable[[F2], F2]]:
     max_cache_size = max_cache_size if max_cache_size is not None else -1
     backend = backend if backend is not None else JitBackend.Invalid
@@ -2530,10 +2534,7 @@ def __init__(self, f) -> None:
                 closure = inspect.getclosurevars(f)
                 self.closure = (closure.nonlocals, closure.globals)
                 self.frozen = detail.FrozenFunction(
-                    inner,
-                    max_cache_size,
-                    warn_recording_count,
-                    backend,
+                    inner, max_cache_size, warn_recording_count, backend, auto_opaque
                 )
 
             def __call__(self, *args, **kwargs):