Implement opt-bisect-limit for mir

Author: sgasho

compiler/rustc_mir_transform/src/pass_manager.rs

@@ -1,5 +1,6 @@
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
+use std::sync::atomic::Ordering;
 
 use rustc_data_structures::fx::{FxHashMap, FxIndexSet};
 use rustc_middle::mir::{Body, MirDumper, MirPhase, RuntimePhase};
@@ -285,6 +286,19 @@ fn run_passes_inner<'tcx>(
                 continue;
             };
 
+            if is_optimization_stage(body, phase_change, optimizations)
+                && let Some(limit) = &tcx.sess.opts.unstable_opts.mir_opt_bisect_limit
+            {
+                if limited_by_opt_bisect(
+                    tcx,
+                    tcx.def_path_debug_str(body.source.def_id()),
+                    *limit,
+                    *pass,
+                ) {
+                    continue;
+                }
+            }
+
             let dumper = if pass.is_mir_dump_enabled()
                 && let Some(dumper) = MirDumper::new(tcx, pass_name, body)
             {
@@ -356,3 +370,46 @@ pub(super) fn dump_mir_for_phase_change<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tc
         dumper.set_show_pass_num().set_disambiguator(&"after").dump_mir(body)
     }
 }
+
+fn is_optimization_stage(
+    body: &Body<'_>,
+    phase_change: Option<MirPhase>,
+    optimizations: Optimizations,
+) -> bool {
+    optimizations == Optimizations::Allowed
+        && body.phase == MirPhase::Runtime(RuntimePhase::PostCleanup)
+        && phase_change == Some(MirPhase::Runtime(RuntimePhase::Optimized))
+}
+
+fn limited_by_opt_bisect<'tcx, P>(
+    tcx: TyCtxt<'tcx>,
+    def_path: String,
+    limit: usize,
+    pass: &P,
+) -> bool
+where
+    P: MirPass<'tcx> + ?Sized,
+{
+    let current_opt_bisect_count =
+        tcx.sess.mir_opt_bisect_eval_count.fetch_add(1, Ordering::Relaxed);
+
+    let can_run = current_opt_bisect_count < limit;
+
+    if can_run {
+        eprintln!(
+            "BISECT: running pass ({}) {} on {}",
+            current_opt_bisect_count + 1,
+            pass.name(),
+            def_path
+        );
+    } else {
+        eprintln!(
+            "BISECT: NOT running pass ({}) {} on {}",
+            current_opt_bisect_count + 1,
+            pass.name(),
+            def_path
+        );
+    }
+
+    !can_run
+}

compiler/rustc_session/src/options.rs

@@ -2481,6 +2481,9 @@ options! {
     mir_include_spans: MirIncludeSpans = (MirIncludeSpans::default(), parse_mir_include_spans, [UNTRACKED],
         "include extra comments in mir pretty printing, like line numbers and statement indices, \
          details about types, etc. (boolean for all passes, 'nll' to enable in NLL MIR only, default: 'nll')"),
+    mir_opt_bisect_limit: Option<usize> = (None, parse_opt_number, [TRACKED],
+        "limit the number of MIR optimization pass executions (global across all bodies). \
+        Pass executions after this limit are skipped and reported. (default: no limit)"),
     #[rustc_lint_opt_deny_field_access("use `Session::mir_opt_level` instead of this field")]
     mir_opt_level: Option<usize> = (None, parse_opt_number, [TRACKED],
         "MIR optimization level (0-4; default: 1 in non optimized builds and 2 in optimized builds)"),

compiler/rustc_session/src/session.rs

@@ -2,7 +2,7 @@ use std::any::Any;
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::Arc;
-use std::sync::atomic::AtomicBool;
+use std::sync::atomic::{AtomicBool, AtomicUsize};
 use std::{env, io};
 
 use rand::{RngCore, rng};
@@ -158,6 +158,12 @@ pub struct Session {
     /// The names of intrinsics that the current codegen backend replaces
     /// with its own implementations.
     pub replaced_intrinsics: FxHashSet<Symbol>,
+
+    /// Global per-session counter for MIR optimization pass applications.
+    ///
+    /// Used by `-Zmir-opt-bisect-limit` to assign an index to each
+    /// optimization-pass execution candidate during this compilation.
+    pub mir_opt_bisect_eval_count: AtomicUsize,
 }
 
 #[derive(Clone, Copy)]
@@ -1097,6 +1103,7 @@ pub fn build_session(
         host_filesearch,
         invocation_temp,
         replaced_intrinsics: FxHashSet::default(), // filled by `run_compiler`
+        mir_opt_bisect_eval_count: AtomicUsize::new(0),
     };
 
     validate_commandline_args_with_session_available(&sess);

tests/run-make/mir-opt-bisect-limit/main.rs

@@ -0,0 +1,10 @@
+#![crate_type = "lib"]
+
+#[inline(never)]
+pub fn callee(x: u64) -> u64 {
+    x.wrapping_mul(3).wrapping_add(7)
+}
+
+pub fn caller(a: u64, b: u64) -> u64 {
+    callee(a) + callee(b)
+}

tests/run-make/mir-opt-bisect-limit/rmake.rs

@@ -0,0 +1,119 @@
+use std::path::Path;
+
+use run_make_support::{CompletedProcess, rfs, rustc};
+
+struct Case {
+    name: &'static str,
+    flags: &'static [&'static str],
+    expect_inline_dump: bool,
+    expect_running: ExpectedCount,
+    expect_not_running: ExpectedCount,
+}
+
+enum ExpectedCount {
+    Exactly(usize),
+    AtLeastOne,
+    Zero,
+}
+
+fn main() {
+    let cases = [
+        Case {
+            name: "limit0",
+            flags: &["-Zmir-opt-bisect-limit=0"],
+            expect_inline_dump: false,
+            expect_running: ExpectedCount::Exactly(0),
+            expect_not_running: ExpectedCount::AtLeastOne,
+        },
+        Case {
+            name: "limit1",
+            flags: &["-Zmir-opt-bisect-limit=1"],
+            expect_inline_dump: false,
+            expect_running: ExpectedCount::Exactly(1),
+            expect_not_running: ExpectedCount::AtLeastOne,
+        },
+        Case {
+            name: "huge_limit",
+            flags: &["-Zmir-opt-bisect-limit=1000000000"],
+            expect_inline_dump: true,
+            expect_running: ExpectedCount::AtLeastOne,
+            expect_not_running: ExpectedCount::Zero,
+        },
+        Case {
+            name: "limit0_with_force_enable_inline",
+            flags: &["-Zmir-opt-bisect-limit=0", "-Zmir-enable-passes=+Inline"],
+            expect_inline_dump: false,
+            expect_running: ExpectedCount::Exactly(0),
+            expect_not_running: ExpectedCount::AtLeastOne,
+        },
+    ];
+
+    for case in cases {
+        let (inline_dumped, running_count, not_running_count, output) =
+            compile_case(case.name, case.flags);
+
+        assert_eq!(
+            inline_dumped, case.expect_inline_dump,
+            "{}: unexpected Inline dump presence",
+            case.name
+        );
+
+        assert_expected_count(
+            running_count,
+            case.expect_running,
+            &format!("{}: running count", case.name),
+        );
+        assert_expected_count(
+            not_running_count,
+            case.expect_not_running,
+            &format!("{}: NOT running count", case.name),
+        );
+    }
+}
+
+fn compile_case(dump_dir: &str, extra_flags: &[&str]) -> (bool, usize, usize, CompletedProcess) {
+    if Path::new(dump_dir).exists() {
+        rfs::remove_dir_all(dump_dir);
+    }
+    rfs::create_dir_all(dump_dir);
+
+    let mut cmd = rustc();
+    cmd.input("main.rs")
+        .arg("--emit=mir")
+        .arg("-Zmir-opt-level=2")
+        .arg("-Copt-level=2")
+        .arg("-Zthreads=1")
+        .arg("-Zdump-mir=Inline")
+        .arg(format!("-Zdump-mir-dir={dump_dir}"));
+
+    for &flag in extra_flags {
+        cmd.arg(flag);
+    }
+
+    let output = cmd.run();
+    let (running_count, not_running_count) = bisect_line_counts(&output);
+    (has_inline_dump_file(dump_dir), running_count, not_running_count, output)
+}
+
+fn assert_expected_count(actual: usize, expected: ExpectedCount, context: &str) {
+    match expected {
+        ExpectedCount::Exactly(n) => assert_eq!(actual, n, "{context}"),
+        ExpectedCount::AtLeastOne => assert!(actual > 0, "{context}"),
+        ExpectedCount::Zero => assert_eq!(actual, 0, "{context}"),
+    }
+}
+
+fn has_inline_dump_file(dir: &str) -> bool {
+    rfs::read_dir(dir)
+        .flatten()
+        .any(|entry| entry.file_name().to_string_lossy().contains(".Inline."))
+}
+
+fn bisect_line_counts(output: &CompletedProcess) -> (usize, usize) {
+    let stderr = output.stderr_utf8();
+    let running_count =
+        stderr.lines().filter(|line| line.starts_with("BISECT: running pass (")).count();
+    let not_running_count =
+        stderr.lines().filter(|line| line.starts_with("BISECT: NOT running pass (")).count();
+    (running_count, not_running_count)
+}