Addressed comments from Wenzel

Author: DoeringChristian

include/drjit-core/jit.h

@@ -1176,8 +1176,7 @@ extern JIT_EXPORT uint32_t jit_var_data(uint32_t index, void **ptr_out);
 /// Query the size of a given variable
 extern JIT_EXPORT size_t jit_var_size(uint32_t index);
 
-/// Query the size of a given variable, as an opaque variable. This also
-/// notifies the ThreadState, allowing us to record this for frozen functions.
+/// Query the size of a given variable, as an opaque variable.
 extern JIT_EXPORT uint32_t jit_var_opaque_width(uint32_t index);
 
 /// Query the type of a given variable

src/api.cpp

@@ -616,16 +616,17 @@ uint32_t jit_var_opaque_width(uint32_t index) {
 
     lock_guard guard(state.lock);
 
-    Variable *var = jitc_var(index);
-    uint32_t var_size = var->size;
+    Variable *var      = jitc_var(index);
+    JitBackend backend = (JitBackend) var->backend;
+    uint32_t var_size  = var->size;
 
     // The variable has to be evaluated, to notify the ThreadState
     jitc_var_eval(index);
 
-    uint32_t width_index = jitc_var_literal(
-        (JitBackend) var->backend, VarType::UInt32, &var_size, 1, true);
+    uint32_t width_index =
+        jitc_var_literal(backend, VarType::UInt32, &var_size, 1, true);
 
-    ThreadState *ts = thread_state(var->backend);
+    ThreadState *ts = thread_state(backend);
     ts->notify_opaque_width(index, width_index);
 
     return width_index;

src/internal.h

@@ -720,6 +720,11 @@ struct ThreadState : public ThreadStateBase {
     virtual void reduce_expanded(VarType vt, ReduceOp op, void *data,
                                  uint32_t exp, uint32_t size) = 0;
 
+    /// Some kernels use the width of an array in a computation. When using the
+    /// kernel freezing feature, this requires special precautions to ensure
+    /// that the resulting capture remains usable with different array sizes.
+    /// This notification function exists so that this special-case handling can
+    /// be realized.
     virtual void notify_opaque_width(uint32_t index, uint32_t width_index);
 
     /// Notify the \c ThreadState that \c jitc_free has been called on a pointer.

src/record_ts.cpp

@@ -476,11 +476,11 @@ void RecordThreadState::barrier() {
 }
 
 void RecordThreadState::notify_opaque_width(uint32_t index,
-                                              uint32_t width_index) {
+                                            uint32_t width_index) {
     if (!paused()) {
         uint32_t start = m_recording.dependencies.size();
-        Variable *v1 = jitc_var(index);
-        Variable *v2 = jitc_var(width_index);
+        Variable *v1   = jitc_var(index);
+        Variable *v2   = jitc_var(width_index);
         add_in_param(v1->data, (VarType) v1->type);
         add_out_param(v2->data, VarType::UInt32);
         uint32_t end = m_recording.dependencies.size();
@@ -492,13 +492,13 @@ void RecordThreadState::notify_opaque_width(uint32_t index,
     }
 }
 
-int Recording::replay_opaque_width(Operation &op){
+int Recording::replay_opaque_width(Operation &op) {
 
     uint32_t dependency_index = op.dependency_range.first;
-    AccessInfo in_info = dependencies[dependency_index];
-    AccessInfo out_info = dependencies[dependency_index + 1];
+    AccessInfo in_info        = dependencies[dependency_index];
+    AccessInfo out_info       = dependencies[dependency_index + 1];
 
-    ReplayVariable &in_var = replay_variables[in_info.slot];
+    ReplayVariable &in_var  = replay_variables[in_info.slot];
     ReplayVariable &out_var = replay_variables[out_info.slot];
 
     out_var.alloc(backend, 1, out_info.vtype);