diff --git a/include/drjit-core/jit.h b/include/drjit-core/jit.h index 34ef2b70..7bb2e402 100644 --- a/include/drjit-core/jit.h +++ b/include/drjit-core/jit.h @@ -2771,18 +2771,26 @@ extern JIT_EXPORT int jit_freeze_dry_run(Recording *recording, /** * \brief Some function result in recordings that cannot be replayed. This - * function marks such recordings, so that Dr.Jit can discard them. + * function marks such recordings, so that Dr.Jit can discard them. If no + * function is in the progress of being recorded this function has no + * effect. + * + * \param backend + * Backend for which the current recording should be discarded. * * \param message * This function accepts an optional message. */ -extern JIT_EXPORT void jit_freeze_discard(const char *message); +extern JIT_EXPORT void jit_freeze_discard(JitBackend backend, const char *message); /** - * \brief This function returns if the last frozen function recording should be + * \brief This function returns if the frozen function recording should be * discarded or not. + * + * \param recording + * Recording for which this function should querry if it should be discarded. */ -extern JIT_EXPORT int jit_freeze_discarded(); +extern JIT_EXPORT int jit_freeze_discarded(const Recording *recording); /** * \brief Pause recording the ThreadState for this backend. diff --git a/src/api.cpp b/src/api.cpp index 5bfa50b8..b5e1934f 100644 --- a/src/api.cpp +++ b/src/api.cpp @@ -1665,14 +1665,14 @@ int jit_freeze_dry_run(Recording *recording, const uint32_t *inputs) { return jitc_freeze_dry_run(recording, inputs); } -void jit_freeze_discard(const char *message) { +void jit_freeze_discard(JitBackend backend, const char *message) { lock_guard guard(state.lock); - jitc_freeze_discard(message); + jitc_freeze_discard(backend, message); } -int jit_freeze_discarded() { +int jit_freeze_discarded(const Recording *recording) { lock_guard guard(state.lock); - return jitc_freeze_discarded(); + return jitc_freeze_discarded(recording); } void jit_freeze_destroy(Recording *recording) { diff --git a/src/record_ts.cpp b/src/record_ts.cpp index 956e85fb..454f3b79 100644 --- a/src/record_ts.cpp +++ b/src/record_ts.cpp @@ -152,12 +152,6 @@ static bool dry_run = false; /// entries. This function will be set by ``jitc_freeze_replay()``. static bool record_kernel_history = false; -/// Indicates if this frozen function recording should be discarded. -/// Some functions can generate recordings that cannot be replayed. These -/// functions can call `jit_freeze_discard` to mark the recording as such. -/// Dr.Jit will then not add it to the recording cache. -static bool discarded = false; - /** * Represents a variable during replay. * It is created from the RecordVariable at the top of the replay function. @@ -2550,9 +2544,6 @@ void jitc_freeze_start(JitBackend backend, const uint32_t *inputs, jitc_fail("Tried to record a thread_state while inside another " "FreezingScope!"); - // Reset the discarded flag. - discarded = false; - // Increment scope, can be used to track missing inputs jitc_new_scope(backend); @@ -2726,10 +2717,14 @@ int jitc_freeze_dry_run(Recording *recording, const uint32_t *inputs) { return result; } -void jitc_freeze_discard(const char *) { - discarded = true; +void jitc_freeze_discard(JitBackend backend, const char *) { + if (RecordThreadState *rts = + dynamic_cast(thread_state(backend)); + rts != nullptr) { + rts->m_recording.discarded = true; + } } -int jitc_freeze_discarded() { - return discarded; +int jitc_freeze_discarded(const Recording *recording) { + return recording->discarded; } diff --git a/src/record_ts.h b/src/record_ts.h index 5c9bcf44..041cdc78 100644 --- a/src/record_ts.h +++ b/src/record_ts.h @@ -30,9 +30,9 @@ void jitc_freeze_replay(Recording *recording, const uint32_t *inputs, int jitc_freeze_dry_run(Recording *recording, const uint32_t *inputs); -void jitc_freeze_discard(const char *message); +void jitc_freeze_discard(JitBackend backend, const char *message); -int jitc_freeze_discarded(); +int jitc_freeze_discarded(const Recording *recording); /// HashMap, mapping an allocation to a recorded variable using PtrToSlot = tsl::robin_map; @@ -280,6 +280,12 @@ struct Recording { /// The backend, which was used while recording. JitBackend backend; + /// Indicates if this frozen function recording should be discarded. + /// Some functions can generate recordings that cannot be replayed. These + /// functions can call `jit_freeze_discard` to mark the recording as such. + /// Dr.Jit will then not add it to the recording cache. + bool discarded = false; + #ifndef NDEBUG /// Counter, counting the number of kernels for debugging. uint32_t n_kernels = 0;