Frozen ThreadState #107

DoeringChristian · 2024-10-23T10:48:12Z

This PR implements the Freezing feature in drjit-core (depends on #106).

It makes the following changes:

Adds the freezing api to jit.h, including the following functions: jit_freeze_start, jit_freeze_stop, jit_freeze_replay, jit_freeze_dry_run, jit_freeze_pause, jit_freeze_resume, jit_freeze_abort and jit_freeze_destroy as well as the Recording struct.
Adds the implementation of the RecordThreadState, a wrapper around either the CUDA or LLVM ThreadState for freezing them.

DoeringChristian · 2024-10-25T11:47:35Z

@merlinND, could you please do a preliminary pass over this PR, to catch some obvious mistakes of mine. Note, this also includes the commit from #106, which I'll remove when that gets merged.

merlinND

Review part 1. It's mostly about documentation / comments / docstrings to make everything super clear.

It's great that you already created unit tests at the drjit-core level!

I still have to look at the big record_ts.* files.

merlinND · 2024-11-01T09:38:48Z

src/init.cpp

@@ -10,6 +10,7 @@
 #include "internal.h"
 #include "cuda_ts.h"
 #include "llvm_ts.h"
+#include "record_ts.h"


Is this include used?

merlinND · 2024-11-01T10:43:45Z

tests/vcall.cpp

+        jit_var_inc_ref(o0.index());
+
+        outputs[0] = o0.index();


Maybe this can be done with UInt32::steal()? (Not sure)

merlinND · 2024-11-01T10:46:15Z

tests/vcall.cpp

+
+    jit_log(LogLevel::Info, "Replay:");
+    {
+        BasePtr self = (arange<UInt32>(10) + 1) % 3;


Could you try replaying with a different width?

merlinND · 2024-11-01T10:47:21Z

tests/vcall.cpp

+    };
+
+    {
+        BasePtr self = arange<UInt32>(10) % 3;


Could you please add a comment explaining what happens if some instances are not used during recording, but we try to use it during replay (i.e. whether or not it's expected to work).

merlinND · 2024-11-01T10:48:41Z

tests/vcall.cpp

+            i0.index(),
+        };
+
+        jit_freeze_start(Backend, inputs, 2);


Please add a comment specifying whether or not this will evaluate some things (e.g. inputs) before starting the recording.
If it doesn't do any evals, then maybe you could move the self.eval(); i0.eval(); just above to make clear that they're needed.

merlinND · 2024-11-01T13:58:10Z

tests/record.cpp

+}
+
+/**
+ * This tests, weather it is possible to record multiple kernels in parallel.


Suggested change

* This tests, weather it is possible to record multiple kernels in parallel.

* This tests, weather it is possible to record multiple independent kernels in the same recording.

merlinND · 2024-11-01T13:59:48Z

tests/record.cpp

+}
+
+/**
+ * This tests recording and replay of a horizontal reduction operation (hsum).


Suggested change

* This tests recording and replay of a horizontal reduction operation (hsum).

* This tests the recording and replay of a horizontal reduction operation (hsum).

merlinND · 2024-11-01T14:00:18Z

tests/record.cpp

+
+        jit_freeze_start(Backend, inputs, 1);
+
+        UInt32 o0 = hsum(i0);


It would be more interesting if there was one operation creating a first kernel, and then hsum() is applied on the output of the first kernel.

merlinND · 2024-11-01T14:02:23Z

tests/record.cpp

+/**
+ * Basic addition test.
+ * Supplying a different input should replay the operation, with this input.
+ * In this case, the input at replay is incremented and should result in an
+ * incremented output.
+ */
+TEST_BOTH(9_resized_input) {


Wouldn't be needed anymore if you change the width during replay for all of the previous tests.

merlinND · 2024-11-01T14:02:42Z

tests/record.cpp

+    jit_freeze_destroy(recording);
+}
+
+TEST_BOTH(10_input_passthrough) {


Comment describing the test case, like for the previous tests.

merlinND

Review part 2

merlinND · 2024-11-01T14:37:51Z

src/record_ts.h

+///
+/// Output variables are only tracked through the outputs array, as this
+/// information is only needed when constructing the output variables.
+///


Suggested change

///

merlinND · 2024-11-01T14:39:10Z

src/record_ts.h

+            bool reverse;
+        } prefix_reduce;
+        /// The bucket count for the mkperm operation
+        uint32_t bucket_count;


If the bucket count of an mkperm operation needs to change during replay, would that be correctly detected & handled (maybe by re-recording)?

Good point. As long as it comes from python it should be fine. If it is determined by the width of some variable this might be a problem.

merlinND · 2024-11-01T14:39:20Z

src/record_ts.h

+        struct {
+            /// The reduce type of a prefix reduction operation
+            ReduceOp rtype;
+            /// Weather a prefix sum operation is exclusive


Suggested change

/// Weather a prefix sum operation is exclusive

/// Whether a prefix sum operation is exclusive

Search & replace for this typo

merlinND · 2024-11-01T14:41:27Z

src/record_ts.h

+    /// Does this operation use optix?
+    bool uses_optix = false;
+    /// A copy of the shader binding table, used by the kernel.
+    OptixShaderBindingTable *sbt;


Is it a deep copy? What happens if the corresponding scene gets destroyed, could this or any elements of the SBT become dangling pointers?

merlinND · 2024-11-01T14:42:21Z

src/record_ts.h

+    OpOutput,
+    /// This variable is part of the function input
+    Input,
+    /// This variable has been captured


Specify what "captured" means. Is it "captured" in the sense of a closure?

merlinND · 2024-11-01T16:35:51Z

src/record_ts.h

+     * This is used by the input variables of a kernel.
+     */
+    uint32_t add_variable(const void *ptr, RecordVariable rv) {
+


Suggested change

merlinND · 2024-11-01T16:37:54Z

src/record_ts.h

+    }
+
+    /// Return the slot index given the data pointer of a variable.
+    /// This fails if the variable has not been added.


Suggested change

/// This fails if the variable has not been added.

/// This fails if the variable had not been previously added.

merlinND · 2024-11-01T16:39:10Z

src/record_ts.h

+                jitc_raise("record(): Varaible at slot s%u was read from by "
+                           "operation o%u, but has not yet been initialized! "
+                           "This can happen if the variable was not part of "
+                           "the input but is used by an recorded operation.",


Explain how the user could make it "part of the input", if possible.

Not sure, does this make sense?

jitc_raise("record(): Varaible at slot s%u was read from by " "operation o%u, but has 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.", info.slot, (uint32_t) this->m_recording.operations.size());

Okay! Fixed a couple of typos below:

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 " "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.", info.slot, (uint32_t) this->m_recording.operations.size());

merlinND · 2024-11-01T16:40:39Z

src/record_ts.h

+    /// Helper function recording an output access, given the slot and \ref
+    /// VarType


Suggested change

/// Helper function recording an output access, given the slot and \ref

/// VarType

/// Helper function recording an output access, given the slot and \ref VarType

and same below

merlinND · 2024-11-01T16:41:36Z

src/record_ts.h

+void jitc_freeze_start(JitBackend backend, const uint32_t *inputs,
+                       uint32_t n_inputs);
+
+Recording *jitc_freeze_stop(JitBackend backend, const uint32_t *outputs,
+                            uint32_t n_outputs);
+
+void jitc_freeze_abort(JitBackend backend);
+
+void jitc_freeze_destroy(Recording *recording);
+
+bool jitc_freeze_pause(JitBackend backend);
+
+bool jitc_freeze_resume(JitBackend backend);
+
+void jitc_freeze_replay(Recording *recording, const uint32_t *inputs,
+                        uint32_t *outputs);
+
+int jitc_freeze_dry_run(Recording *recording, const uint32_t *inputs,
+                        uint32_t *outputs);


Maybe move these to the top to make them more noticeable, since RecordThreadState is so long.

merlinND

Review part 3.
It's a partial review of record_ts.cpp, but unfortunately I won't have time to review the rest of the file, there's too much code.

merlinND · 2024-11-01T16:43:43Z

src/record_ts.cpp

+        this->init = rv.init;
+
+        if (init == RecordVarInit::Captured) {
+            // copy the variable, so that it isn't changed


Would that always match the behavior of no-recording?
I.e. if there are captured variables and I run the code with / without recording, would I always get the same result?

merlinND · 2024-11-04T09:09:15Z

src/record_ts.h

+                jitc_raise("record(): Varaible at slot s%u was read from by "
+                           "operation o%u, but has not yet been initialized! "
+                           "This can happen if the variable was not part of "
+                           "the input but is used by an recorded operation.",


Okay! Fixed a couple of typos below:

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 " "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.", info.slot, (uint32_t) this->m_recording.operations.size());

merlinND · 2024-11-04T09:13:13Z

src/record_ts.cpp

+/// Temporary variables used for replaying a recording.
+static std::vector<ReplayVariable> replay_variables;


What happens if someone tries replaying two different recordings from two different threads?
(I'm not sure what thread-safety the new DrJit version provides)

merlinND · 2024-11-04T09:15:29Z

src/record_ts.cpp

+
+                // Reconstruct the \ref kernel_params for this launch given the
+                // allocations when replaying.
+                kernel_params.clear();


Is there any risk this interferes with an unrelated pending kernel launch?

merlinND · 2024-11-04T09:15:45Z

src/record_ts.cpp

+                    // First 3 parameters reserved for: kernel ptr, size, ITT
+                    // identifier


Suggested change

// First 3 parameters reserved for: kernel ptr, size, ITT

// identifier

// First 3 parameters reserved for: kernel ptr, size, ITT identifier

merlinND · 2024-11-04T09:28:38Z

src/record_ts.cpp

+                jitc_log(LogLevel::Debug, "    src.data=%p", src_var.data);
+                jitc_log(LogLevel::Debug, "    dst.data=%p", dst_var.data);


Reorganize a bit so that the 3 logs can be combined?

merlinND · 2024-11-04T09:29:08Z

src/record_ts.cpp

+                offsets_var.alloc(backend, bucket_count * 4 + 1,
+                                  offsets_info.vtype);
+
+                jitc_log(LogLevel::Debug,


Combine all the log calls

merlinND · 2024-11-04T09:31:19Z

src/record_ts.cpp

+                         "replay(): MemcpyAsync s%u <- s%u [%zu]",
+                         dst_info.slot, src_info.slot, src_var.data_size);
+
+                dst_var.alloc(backend, src_var.data_size);


Could it be a problem that the memory is allocated even during a dry-run? (for this and other ops)
Will it be quickly freed again? We don't want to accidentally double memory usage because of dry runs.

merlinND · 2024-11-04T09:32:58Z

src/record_ts.cpp

+
+            };
+            break;
+            case OpType::Aggregate: {


Move contents to a dedicated function since it's a complex case

merlinND · 2024-11-04T09:33:30Z

src/record_ts.cpp

+            default:
+                jitc_fail(
+                    "An operation has been recorded, that is not known to "
+                    "the replay functionality!");


Print the op type to help with debugging

Added registry_fill_ptrs function

1f73258

DoeringChristian force-pushed the frozen-threadstate branch 3 times, most recently from 288194d to bb38a86 Compare October 23, 2024 13:23

DoeringChristian marked this pull request as ready for review October 25, 2024 10:58

merlinND reviewed Nov 1, 2024

View reviewed changes

DoeringChristian added 7 commits November 2, 2024 15:53

Addressed comments from @merlinND

07b77a8

Added the jit_freeze api

873e47a

Added freezing feature

115ac4e

Added tests for ThreadState freezing

c010ae6

Minor formatting and logging changes

a1bd1d0

Formatting break statements

184e670

Added semicolons after case curly brackets

6a53f8f

DoeringChristian force-pushed the frozen-threadstate branch from 489f314 to 6a53f8f Compare November 2, 2024 14:53

DoeringChristian added 2 commits November 4, 2024 09:34

Addressed most comments from @merlinND review 1st part

5d97c79

Addressed most comments from @merlinND review 2nd part

e01f19a

DoeringChristian force-pushed the frozen-threadstate branch from 7ea730c to e01f19a Compare November 4, 2024 09:34

merlinND reviewed Nov 4, 2024

View reviewed changes

	* This tests, weather it is possible to record multiple kernels in parallel.
	* This tests, weather it is possible to record multiple independent kernels in the same recording.

	* This tests recording and replay of a horizontal reduction operation (hsum).
	* This tests the recording and replay of a horizontal reduction operation (hsum).

	/// Weather a prefix sum operation is exclusive
	/// Whether a prefix sum operation is exclusive

	/// This fails if the variable has not been added.
	/// This fails if the variable had not been previously added.

		/// Helper function recording an output access, given the slot and \ref
		/// VarType

		/// Temporary variables used for replaying a recording.
		static std::vector<ReplayVariable> replay_variables;

		// First 3 parameters reserved for: kernel ptr, size, ITT
		// identifier

		jitc_log(LogLevel::Debug, " src.data=%p", src_var.data);
		jitc_log(LogLevel::Debug, " dst.data=%p", dst_var.data);

Frozen ThreadState #107

Are you sure you want to change the base?

Frozen ThreadState #107

Conversation

DoeringChristian commented Oct 23, 2024 • edited Loading

DoeringChristian commented Oct 25, 2024

merlinND left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merlinND left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

merlinND left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DoeringChristian commented Oct 23, 2024 •

edited

Loading