i#7157 dyn inject: Fix extra switch injections

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