i#7157 dyn inject: Fix extra switch injections (#7302)

Fixes a bug in the dynamic injection logic in the drmemtrace scheduler
that caused the context switch sequence to be injected more times than
actual context switches. It was erroneously injected also when
set_cur_input() was invoked without any change to the input.

update_switch_stats() is a better control point for switch sequence
injection, which is renamed to on_context_switch() here. Now, in the
switch_insertion offline test, we have 11 context switch injections (==
input-to-input + idle-to-input switches), instead of the prior 363.

We deliberately do not inject the context switch sequence on
input-to-idle transitions. Note that we already inject on idle-to-input
transitions, and it is not clear yet whether we want to repeat it at
input-to-idle also or maybe inject some different sequence.

The above move also removes the use of get_instruction_ordinal() to
determine whether switch records should be inserted. Existing logic
checked for get_instruction_ordinal() > 0 to avoid inserting switch
records before the very first input on an output. This does not work as
required when USE_INPUT_ORDINALS is set, as then the input-local
instruction ordinal instead (which is non-zero as it is not adjusted for
the read-ahead instructions) is returned, which ends up inserting the
switch sequence before the very first input also. A scheduler unit_test
that uses USE_INPUT_ORDINALS is coming up in #7299 with more fixes.

Issue: #7157

Author: abhinav92003

clients/drcachesim/common/memtrace_stream.h

@@ -113,7 +113,7 @@ class memtrace_stream_t {
          */
         SCHED_STAT_HIT_OUTPUT_LIMIT,
         /**
-         * Counts the instances when the context switch sequence was injected.
+         * Counts the instances when the kernel context switch sequence was injected.
          */
         SCHED_STAT_KERNEL_SWITCH_SEQUENCE_INJECTIONS,
         /** Count of statistic types. */

clients/drcachesim/scheduler/scheduler_impl.cpp

@@ -1,5 +1,5 @@
 /* **********************************************************
- * Copyright (c) 2023-2024 Google, Inc.  All rights reserved.
+ * Copyright (c) 2023-2025 Google, Inc.  All rights reserved.
  * **********************************************************/
 
 /*
@@ -2288,41 +2288,6 @@ scheduler_impl_tmpl_t<RecordType, ReaderType>::set_cur_input(
         stream->page_size_ = inputs_[input].reader->get_page_size();
     }
 
-    if (inputs_[input].pid != INVALID_PID) {
-        insert_switch_tid_pid(inputs_[input]);
-    }
-
-    if (!switch_sequence_.empty() &&
-        outputs_[output].stream->get_instruction_ordinal() > 0) {
-        switch_type_t switch_type = sched_type_t::SWITCH_INVALID;
-        if (prev_workload != inputs_[input].workload)
-            switch_type = sched_type_t::SWITCH_PROCESS;
-        else
-            switch_type = sched_type_t::SWITCH_THREAD;
-        // Inject kernel context switch code.  Since the injected records belong to this
-        // input (the kernel is acting on behalf of this input) we insert them into the
-        // input's queue, but ahead of any prior queued items.  This is why we walk in
-        // reverse, for the push_front calls to the deque.  We update the tid of the
-        // records here to match.  They are considered as is_record_synthetic() and do
-        // not affect input stream ordinals.
-        // XXX: These will appear before the top headers of a new thread which is slightly
-        // odd to have regular records with the new tid before the top headers.
-        if (!switch_sequence_[switch_type].empty()) {
-            ++outputs_[output]
-                  .stats[memtrace_stream_t::SCHED_STAT_KERNEL_SWITCH_SEQUENCE_INJECTIONS];
-            for (int i = static_cast<int>(switch_sequence_[switch_type].size()) - 1;
-                 i >= 0; --i) {
-                RecordType record = switch_sequence_[switch_type][i];
-                record_type_set_tid(record, inputs_[input].tid);
-                inputs_[input].queue.push_front(record);
-            }
-            VPRINT(this, 3,
-                   "Inserted %zu switch records for type %d from %d.%d to %d.%d\n",
-                   switch_sequence_[switch_type].size(), switch_type, prev_workload,
-                   outputs_[output].prev_input, inputs_[input].workload, input);
-        }
-    }
-
     inputs_[input].prev_time_in_quantum =
         outputs_[output].cur_time->load(std::memory_order_acquire);
 
@@ -2415,28 +2380,80 @@ scheduler_impl_tmpl_t<RecordType, ReaderType>::pick_next_input(output_ordinal_t
     }
     // We can't easily place these stats inside set_cur_input() as we call that to
     // temporarily give up our input.
-    update_switch_stats(output, prev_index, index);
+    on_context_switch(output, prev_index, index);
     set_cur_input(output, index);
     return res;
 }
 
 template <typename RecordType, typename ReaderType>
 void
-scheduler_impl_tmpl_t<RecordType, ReaderType>::update_switch_stats(
+scheduler_impl_tmpl_t<RecordType, ReaderType>::on_context_switch(
     output_ordinal_t output, input_ordinal_t prev_input, input_ordinal_t new_input)
 {
+    bool do_inject_switch_seq = false;
     if (prev_input == new_input)
         ++outputs_[output].stats[memtrace_stream_t::SCHED_STAT_SWITCH_NOP];
     else if (prev_input != sched_type_t::INVALID_INPUT_ORDINAL &&
-             new_input != sched_type_t::INVALID_INPUT_ORDINAL)
+             new_input != sched_type_t::INVALID_INPUT_ORDINAL) {
         ++outputs_[output].stats[memtrace_stream_t::SCHED_STAT_SWITCH_INPUT_TO_INPUT];
-    else if (new_input == sched_type_t::INVALID_INPUT_ORDINAL)
+        do_inject_switch_seq = true;
+    } else if (new_input == sched_type_t::INVALID_INPUT_ORDINAL) {
+        // XXX: For now, we do not inject a kernel context switch sequence on
+        // input-to-idle transitions (note that we do so on idle-to-input though).
+        // However, we may want to inject some other suitable sequence, but we're not
+        // sure yet.
         ++outputs_[output].stats[memtrace_stream_t::SCHED_STAT_SWITCH_INPUT_TO_IDLE];
-    else {
+    } else {
         ++outputs_[output].stats[memtrace_stream_t::SCHED_STAT_SWITCH_IDLE_TO_INPUT];
         // Reset the flag so we'll try to steal if we go idle again.
         outputs_[output].tried_to_steal_on_idle = false;
+        do_inject_switch_seq = true;
     }
+
+    // We want to insert the context switch sequence on input-to-input and
+    // idle-to-input cases. This is a better control point to do that than
+    // set_cur_input. Here we get the stolen input events too, and we don't have
+    // to filter out the init-time set_cur_input cases.
+    if (!do_inject_switch_seq)
+        return;
+    if (inputs_[new_input].pid != INVALID_PID) {
+        insert_switch_tid_pid(inputs_[new_input]);
+    }
+    if (switch_sequence_.empty())
+        return;
+    switch_type_t switch_type = sched_type_t::SWITCH_INVALID;
+    if ( // XXX: idle-to-input transitions are assumed to be process switches
+         // for now. But we may want to improve this heuristic.
+        prev_input == sched_type_t::INVALID_INPUT_ORDINAL ||
+        inputs_[prev_input].workload != inputs_[new_input].workload)
+        switch_type = sched_type_t::SWITCH_PROCESS;
+    else
+        switch_type = sched_type_t::SWITCH_THREAD;
+    if (switch_sequence_[switch_type].empty())
+        return;
+    // Inject kernel context switch code.  Since the injected records belong to
+    // this input (the kernel is acting on behalf of this input) we insert them
+    // into the input's queue, but ahead of any prior queued items.  This is why
+    // we walk in reverse, for the push_front calls to the deque.  We update the
+    // tid of the records here to match.  They are considered as
+    // is_record_synthetic() and do not affect input stream ordinals.
+    // XXX: These will appear before the top headers of a new thread which is
+    // slightly odd to have regular records with the new tid before the top
+    // headers.
+    ++outputs_[output]
+          .stats[memtrace_stream_t::SCHED_STAT_KERNEL_SWITCH_SEQUENCE_INJECTIONS];
+    for (int i = static_cast<int>(switch_sequence_[switch_type].size()) - 1; i >= 0;
+         --i) {
+        RecordType record = switch_sequence_[switch_type][i];
+        record_type_set_tid(record, inputs_[new_input].tid);
+        inputs_[new_input].queue.emplace_front(record);
+    }
+    VPRINT(this, 3, "Inserted %zu switch for type %d from %d.%d to %d.%d\n",
+           switch_sequence_[switch_type].size(), switch_type,
+           prev_input != sched_type_t::INVALID_INPUT_ORDINAL
+               ? inputs_[prev_input].workload
+               : -1,
+           prev_input, inputs_[new_input].workload, new_input);
 }
 
 template <typename RecordType, typename ReaderType>
@@ -2655,7 +2672,7 @@ scheduler_impl_tmpl_t<RecordType, ReaderType>::next_record(output_ordinal_t outp
             } else if (res == sched_type_t::STATUS_STOLE) {
                 // We need to loop to get the new record.
                 input = &inputs_[outputs_[output].cur_input];
-                update_switch_stats(output, prev_input, input->index);
+                on_context_switch(output, prev_input, input->index);
                 lock.unlock();
                 lock = std::unique_lock<mutex_dbg_owned>(*input->lock);
                 lock.lock();

clients/drcachesim/scheduler/scheduler_impl.h

@@ -406,8 +406,8 @@ template <typename RecordType, typename ReaderType> class scheduler_impl_tmpl_t
     scale_blocked_time(uint64_t initial_time) const;
 
     void
-    update_switch_stats(output_ordinal_t output, input_ordinal_t prev_input,
-                        input_ordinal_t new_input);
+    on_context_switch(output_ordinal_t output, input_ordinal_t prev_input,
+                      input_ordinal_t new_input);
 
     ///
     ///////////////////////////////////////////////////////////////////////////

clients/drcachesim/tests/switch_insertion.templatex

@@ -2,25 +2,25 @@ Schedule stats tool results:
 Total counts:
            4 cores
 .*
-      639664 instructions
+      638960 instructions
           11 total context switches
 .*
            6 voluntary context switches
            0 direct context switches
-         363 context switch sequence injections
+          11 context switch sequence injections
 .*
            9 switches input-to-input
            5 switches input-to-idle
            2 switches idle-to-input
 .*
 Core #0 counts:
 .*
-      117223 instructions
+      117031 instructions
            8 total context switches
 .*
            3 voluntary context switches
            0 direct context switches
-         104 context switch sequence injections
+           8 context switch sequence injections
 .*
            6 switches input-to-input
            3 switches input-to-idle