From d474859a7da3216df39aa0de7ec485a8f6c2cc5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20D=C3=B6ring?= Date: Thu, 9 Oct 2025 17:10:27 +0200 Subject: [PATCH 1/2] Added mapping from pointer to variable for debugging --- src/eval.cpp | 4 +++ src/internal.h | 9 +++++++ src/record_ts.cpp | 66 +++++++++++++++++++++++++++++++++-------------- src/record_ts.h | 2 +- src/var.cpp | 21 ++++++++++++++- 5 files changed, 81 insertions(+), 21 deletions(-) diff --git a/src/eval.cpp b/src/eval.cpp index 09d5cc52..96145c6c 100644 --- a/src/eval.cpp +++ b/src/eval.cpp @@ -817,6 +817,10 @@ void jitc_eval_impl(ThreadState *ts) { if (v->is_array()) v->scope = 0; + +#ifndef NDEBUG + state.ptr_to_variable.insert({ v->data, index }); +#endif } uint32_t dep[4], side_effect = v->side_effect; diff --git a/src/internal.h b/src/internal.h index 16a6ea25..e9ff4adc 100644 --- a/src/internal.h +++ b/src/internal.h @@ -840,6 +840,8 @@ struct KernelHistory { using UnusedPQ = std::priority_queue, std::greater>; +using PointerMap = tsl::robin_map; + /// Records the full JIT compiler state (most frequently two used entries at top) struct State { /// Must be held to access members of this data structure @@ -927,6 +929,13 @@ struct State { uint32_t optix_default_sbt_index = 0; #endif +#ifndef NDEBUG + /// Mapping from pointers that are managed by variables to their variable + /// indices. This is used for debugging purposes in frozen functions. + PointerMap ptr_to_variable; +#endif + + State() { variables.resize(1); extra.resize(1); diff --git a/src/record_ts.cpp b/src/record_ts.cpp index c9d9d4e4..f5ad4f78 100644 --- a/src/record_ts.cpp +++ b/src/record_ts.cpp @@ -2034,34 +2034,46 @@ void RecordThreadState::enqueue_host_func(void (*callback)(void *), (void) callback; (void) payload; } -void Recording::validate() { +void Recording::validate(uint32_t scope) { for (uint32_t i = 0; i < recorded_variables.size(); i++) { RecordedVariable &rv = recorded_variables[i]; if (rv.state == RecordedVarState::Uninitialized) { Operation &last_op = operations[rv.last_op]; #ifndef NDEBUG + uint32_t index = 0; + const char *scope_string = "before"; + auto it = state.ptr_to_variable.find(rv.ptr); + if (it != state.ptr_to_variable.end()) { + index = it->second; + Variable *var = jitc_var(index); + if (var->scope >= scope) + scope_string = "inside"; + } if (last_op.type == OpType::Aggregate) { jitc_raise( "validate(): The frozen function included a virtual " - "function call involving variable s%u <%p>, last used by " - "operation o%u. Dr.Jit would normally traverse a registry " - "of all relevant object instances in order to collect " - "their member variables. However, when recording this " - "frozen function, this traversal was skipped because no " - "such object instance was found in the function's inputs. " - "You can trigger traversal by including the relevant " - "objects in the function input, or by specifying them " - "using the state_fn argument. Alternatively, this error " - "might be caused by a nested " + "function call involving Variable r%u at slot s%u <%p>" + "which created %s the frozen function and was last used by " + "operation o%u. Dr.Jit would " + "normally traverse a registry of all relevant object " + "instances in order to collect their member variables. " + "However, when recording this frozen function, this " + "traversal was skipped because no such object instance was " + "found in the function's inputs. You can trigger traversal " + "by including the relevant objects in the function input, " + "or by specifying them using the state_fn argument. " + "Alternatively, this error might be caused by a nested " "virtual function call.", - i, rv.ptr, rv.last_op); + index, i, rv.ptr, scope_string, rv.last_op); } else jitc_raise( - "validate(): Variable at slot s%u <%p> was used by %s operation " - "o%u but left in an uninitialized state! This indicates " - "that the associated variable was used, but not traversed " - "as part of the frozen function input.", - i, rv.ptr, op_type_name[(uint32_t) last_op.type], rv.last_op); + "validate(): Variable r%u at slot s%u <%p> which was " + "created %s the frozen function and was last used by %s " + "operation o%u but left in an uninitialized state! This " + "indicates that the associated variable was used, but not " + "traversed as part of the frozen function input.", + index, i, rv.ptr, scope_string, + op_type_name[(uint32_t) last_op.type], rv.last_op); #else if (last_op.type == OpType::Aggregate) { jitc_raise( @@ -2301,7 +2313,20 @@ void RecordThreadState::add_param(AccessInfo info) { jitc_log(LogLevel::Debug, " -> param s%u", info.slot); RecordedVariable &rv = m_recording.recorded_variables[info.slot]; - if (info.test_uninit && rv.state == RecordedVarState::Uninitialized) + if (info.test_uninit && rv.state == RecordedVarState::Uninitialized){ +#ifndef NDEBUG + uint32_t index = 0; + auto it = state.ptr_to_variable.find(rv.ptr); + if (it != state.ptr_to_variable.end()) + index = it->second; + jitc_raise("record(): Variable r%u at slot s%u was read by " + "operation o%u, but it had not yet been initialized! " + "This can occur if the variable was not part of " + "the input but is used by a recorded operation, for " + "example if it was not specified as a member in a " + "DRJIT_STRUCT but used in the frozen function.", + index, info.slot, (uint32_t) m_recording.operations.size()); +#else jitc_raise("record(): Variable at slot s%u was read by " "operation o%u, but it had not yet been initialized! " "This can occur if the variable was not part of " @@ -2309,6 +2334,8 @@ void RecordThreadState::add_param(AccessInfo info) { "example if it was not specified as a member in a " "DRJIT_STRUCT but used in the frozen function.", info.slot, (uint32_t) m_recording.operations.size()); +#endif + } if (info.vtype == VarType::Void) info.vtype = rv.type; @@ -2544,6 +2571,7 @@ Recording *jitc_freeze_stop(JitBackend backend, const uint32_t *outputs, dynamic_cast(thread_state(backend)); rts != nullptr) { ThreadState *internal = rts->m_internal; + uint32_t scope = internal->scope; // Perform reassignments to internal thread-state of possibly changed // variables @@ -2571,7 +2599,7 @@ Recording *jitc_freeze_stop(JitBackend backend, const uint32_t *outputs, } Recording *recording = new Recording(std::move(rts->m_recording)); try{ - recording->validate(); + recording->validate(scope); } catch (const std::exception &) { recording->destroy(); throw; diff --git a/src/record_ts.h b/src/record_ts.h index 0f6c40ba..917ce6cd 100644 --- a/src/record_ts.h +++ b/src/record_ts.h @@ -321,7 +321,7 @@ struct Recording { /// This function is called after recording and checks that the recording is /// valid i.e. that no variables where left uninitialized. - void validate(); + void validate(uint32_t scope); /// Checks if all recorded kernels are still in the kernel cache. This might /// occur when calling dr.kernel_cache_flush between recording the function /// and replaying it. diff --git a/src/var.cpp b/src/var.cpp index 7c570817..5a8632d4 100644 --- a/src/var.cpp +++ b/src/var.cpp @@ -309,8 +309,22 @@ JIT_NOINLINE void jitc_var_free(uint32_t index, Variable *v) noexcept { if (v->is_evaluated()) { // Release memory referenced by this variable - if (!v->retain_data) + if (!v->retain_data) { jitc_free(v->data); + +#ifndef NDEBUG + // This warning should never be thrown, except if we forgot to + // populate the mapping + if (!state.ptr_to_variable.contains(v->data)) + jitc_log( + LogLevel::Warn, + "Pointr <%p> was mangaged by variable r%u, but this " + "was not recorded in the pointer to variable map!", + v->data, index); + + state.ptr_to_variable.erase(v->data); +#endif + } } else { // Unevaluated variable, drop from CSE cache jitc_lvn_drop(index, v); @@ -784,6 +798,11 @@ uint32_t jitc_var_new(Variable &v, bool disable_lvn) { jitc_sanitation_checkpoint(); #endif +#ifndef NDEBUG + if (v.is_evaluated()) + state.ptr_to_variable.insert({ v.data, index }); +#endif + return index; } From 72140d55d805ac825b2f6fca1f3a1505e541f8ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christian=20D=C3=B6ring?= Date: Thu, 30 Oct 2025 16:32:56 +0100 Subject: [PATCH 2/2] Fixed bug where optix would fail after a recording threw an exception --- src/record_ts.cpp | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/record_ts.cpp b/src/record_ts.cpp index f5ad4f78..2c7e5c7d 100644 --- a/src/record_ts.cpp +++ b/src/record_ts.cpp @@ -668,6 +668,11 @@ Task *RecordThreadState::launch(Kernel kernel, KernelKey *key, } } pause_scope pause(this); + // Forward the OptiX SBT and pipeline to the internal thread state. +#if defined(DRJIT_ENABLE_OPTIX) + m_internal->optix_pipeline = optix_pipeline; + m_internal->optix_sbt = optix_sbt; +#endif return m_internal->launch(kernel, key, hash, size, params, nullptr, kernel_history_entry); } @@ -875,13 +880,6 @@ void RecordThreadState::record_launch( } m_recording.operations.push_back(op); - - // Re-assign optix specific variables to internal thread state since they - // might have changed -#if defined(DRJIT_ENABLE_OPTIX) - m_internal->optix_pipeline = optix_pipeline; - m_internal->optix_sbt = optix_sbt; -#endif } int Recording::replay_launch(Operation &op) {