Fixed race condition in jit_freeze_discard

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.

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) {

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<RecordThreadState *>(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;
 }

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<const void *, uint32_t, PointerHasher>;
@@ -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;