Added symbolic width for freezing means

Author: DoeringChristian

include/drjit-core/array.h

@@ -253,6 +253,11 @@ template <JitBackend Backend_, typename Value_> struct JitArray {
         return jit_var_size(m_index);
     }
 
+    auto symbolic_width() {
+        using UInt32 = JitArray<Backend_, uint32_t>;
+        return UInt32::steal(jit_var_symbolic_width(m_index));
+    }
+
     void resize(size_t size) {
         uint32_t index = jit_var_resize(m_index, size);
         jit_var_dec_ref(m_index);

include/drjit-core/jit.h

@@ -1176,6 +1176,8 @@ 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 type of a given variable
 extern JIT_EXPORT JIT_ENUM VarType jit_var_type(uint32_t index);
 

src/api.cpp

@@ -610,6 +610,24 @@ size_t jit_var_size(uint32_t index) {
     return (size_t) jitc_var(index)->size;
 }
 
+uint32_t jit_var_symbolic_width(uint32_t index){
+    if(index == 0)
+        return 0;
+
+    lock_guard guard(state.lock);
+
+    Variable *var = jitc_var(index);
+    uint32_t var_size = var->size;
+
+    uint32_t width_index = jitc_var_literal(
+        (JitBackend) var->backend, VarType::UInt32, &var_size, 1, true);
+
+    ThreadState *ts = thread_state(var->backend);
+    ts->notify_symbolic_width(index, width_index);
+
+    return width_index;
+}
+
 VarState jit_var_state(uint32_t index) {
     if (index == 0)
         return VarState::Invalid;

src/init.cpp

@@ -768,3 +768,4 @@ void ThreadState::reset_state() {
 }
 void ThreadState::notify_free(const void *) { }
 void ThreadState::notify_expand(uint32_t) { }
+void ThreadState::notify_symbolic_width(uint32_t, uint32_t) { }

src/internal.h

@@ -720,6 +720,8 @@ struct ThreadState : public ThreadStateBase {
     virtual void reduce_expanded(VarType vt, ReduceOp op, void *data,
                                  uint32_t exp, uint32_t size) = 0;
 
+    virtual void notify_symbolic_width(uint32_t index, uint32_t width_index);
+
     /// Notify the \c ThreadState that \c jitc_free has been called on a pointer.
     /// This is required for kernel freezing.
     virtual void notify_free(const void *ptr);

src/record_ts.cpp

@@ -373,6 +373,10 @@ int Recording::replay(const uint32_t *replay_inputs, uint32_t *replay_outputs) {
                 if (!replay_aggregate(op))
                     return false;
                 break;
+            case OpType::SymbolicWidth:
+                if(!replay_symbolic_width(op))
+                    return false;
+                break;
             case OpType::Free: {
                 ProfilerPhase profiler2(pr_free);
 
@@ -471,6 +475,41 @@ void RecordThreadState::barrier() {
     return m_internal->barrier();
 }
 
+void RecordThreadState::notify_symbolic_width(uint32_t 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);
+        add_in_param(v1->data, (VarType) v1->type);
+        add_out_param(v2->data, VarType::UInt32);
+        uint32_t end = m_recording.dependencies.size();
+
+        Operation op;
+        op.type             = OpType::SymbolicWidth;
+        op.dependency_range = std::pair(start, end);
+        m_recording.operations.push_back(op);
+    }
+}
+
+int Recording::replay_symbolic_width(Operation &op){
+
+    uint32_t dependency_index = op.dependency_range.first;
+    AccessInfo in_info = dependencies[dependency_index];
+    AccessInfo out_info = dependencies[dependency_index + 1];
+
+    ReplayVariable &in_var = replay_variables[in_info.slot];
+    ReplayVariable &out_var = replay_variables[out_info.slot];
+
+    out_var.alloc(backend, 1, out_info.vtype);
+    uint32_t size = in_var.size(in_info.vtype);
+
+    if (!dry_run)
+        jitc_memcpy(backend, out_var.data, &size, sizeof(uint32_t));
+
+    return true;
+}
+
 /**
  * This function is called every time a pointer is freed using \ref
  * jitc_free. It records the operation and removes the mapping from that

src/record_ts.h

@@ -38,6 +38,7 @@ enum class OpType {
     BlockPrefixReduce,
     ReduceDot,
     Aggregate,
+    SymbolicWidth,
     Free,
     Count,
 };
@@ -296,6 +297,8 @@ struct Recording {
 
     int replay_aggregate(Operation &op);
 
+    int replay_symbolic_width(Operation &op);
+
     /// This function is called after recording and checks that the recording is
     /// valid i.e. that no variables where left uninitialized.
     void validate();
@@ -428,6 +431,8 @@ struct RecordThreadState : ThreadState {
     void reduce_expanded(VarType vt, ReduceOp reduce_op, void *data,
                          uint32_t exp, uint32_t size) override;
 
+    void notify_symbolic_width(uint32_t index, uint32_t width_index) override;
+
     /**
      * This function is called every time a pointer is freed using \ref
      * jitc_free. It records the operation and removes the mapping from that

tests/record.cpp

@@ -248,3 +248,22 @@ TEST_LLVM(10_scatter) {
         jit_assert(all(eq(y, arange<UInt32>(10 + i) + 1)));
     }
 }
+
+TEST_BOTH(11_symbolic_width) {
+    auto func = [](UInt32 x) {
+        auto y = block_prefix_sum(x, x.size());
+        y = y / x.symbolic_width();
+        return y;
+    };
+
+    FrozenFunction frozen(Backend, func);
+
+    for (uint32_t i = 0; i < 4; i++) {
+        auto x = arange<UInt32>(10 + i);
+
+        auto res = frozen(x);
+        auto ref = func(x);
+
+        jit_assert(all(eq(res, ref)));
+    }
+}