Added opaque width for freezing mean

Author: DoeringChristian

include/drjit-core/array.h

@@ -255,7 +255,7 @@ template <JitBackend Backend_, typename Value_> struct JitArray {
 
     auto symbolic_width() {
         using UInt32 = JitArray<Backend_, uint32_t>;
-        return UInt32::steal(jit_var_symbolic_width(m_index));
+        return UInt32::steal(jit_var_opaque_width(m_index));
     }
 
     void resize(size_t size) {

include/drjit-core/jit.h

@@ -1176,7 +1176,9 @@ 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);
 
-extern JIT_EXPORT uint32_t jit_var_symbolic_width(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.
+extern JIT_EXPORT uint32_t jit_var_opaque_width(uint32_t index);
 
 /// Query the type of a given variable
 extern JIT_EXPORT JIT_ENUM VarType jit_var_type(uint32_t index);

src/api.cpp

@@ -610,7 +610,7 @@ size_t jit_var_size(uint32_t index) {
     return (size_t) jitc_var(index)->size;
 }
 
-uint32_t jit_var_symbolic_width(uint32_t index){
+uint32_t jit_var_opaque_width(uint32_t index){
     if(index == 0)
         return 0;
 
@@ -619,6 +619,9 @@ uint32_t jit_var_symbolic_width(uint32_t index){
     Variable *var = jitc_var(index);
     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);
 

tests/record.cpp

@@ -251,8 +251,11 @@ TEST_LLVM(10_scatter) {
 
 TEST_BOTH(11_symbolic_width) {
     auto func = [](UInt32 x) {
-        auto y = block_prefix_sum(x, x.size());
-        y = y / x.symbolic_width();
+        auto y = block_prefix_sum(x+1, x.size());
+        if (jit_flag(JitFlag::FreezingScope))
+            y = y / x.symbolic_width();
+        else
+            y = y / x.size();
         return y;
     };