add flag to specify the number of cpus

7 files changed, 28 insertions, 6 deletions

diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index 7d32ee42573..5b16fc2948c 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -550,6 +550,13 @@ fn main() {
             } else {
                 show_error!("-Zmiri-extern-so-file `{}` does not exist", filename);
             }
+        } else if let Some(param) = arg.strip_prefix("-Zmiri-num-cpus=") {
+            let num_cpus = match param.parse::<u32>() {
+                Ok(i) => i,
+                Err(err) => show_error!("-Zmiri-num-cpus requires a `u32`: {}", err),
+            };
+
+            miri_config.num_cpus = num_cpus;
         } else {
             // Forward to rustc.
             rustc_args.push(arg);
diff --git a/src/tools/miri/src/eval.rs b/src/tools/miri/src/eval.rs
index 91a2ac13b1b..4734d488d04 100644
--- a/src/tools/miri/src/eval.rs
+++ b/src/tools/miri/src/eval.rs
@@ -132,6 +132,8 @@ pub struct MiriConfig {
     pub external_so_file: Option<PathBuf>,
     /// Run a garbage collector for SbTags every N basic blocks.
     pub gc_interval: u32,
+    /// FIXME: add docs.
+    pub num_cpus: u32,
 }
 
 impl Default for MiriConfig {
@@ -164,6 +166,7 @@ impl Default for MiriConfig {
             retag_fields: false,
             external_so_file: None,
             gc_interval: 10_000,
+            num_cpus: 1,
         }
     }
 }
diff --git a/src/tools/miri/src/lib.rs b/src/tools/miri/src/lib.rs
index 6006d6c89db..43455de8c6b 100644
--- a/src/tools/miri/src/lib.rs
+++ b/src/tools/miri/src/lib.rs
@@ -104,7 +104,7 @@ pub use crate::helpers::{CurrentSpan, EvalContextExt as HelpersEvalContextExt};
 pub use crate::intptrcast::ProvenanceMode;
 pub use crate::machine::{
     AllocExtra, FrameData, MiriInterpCx, MiriInterpCxExt, MiriMachine, MiriMemoryKind, Provenance,
-    ProvenanceExtra, NUM_CPUS, PAGE_SIZE, STACK_ADDR, STACK_SIZE,
+    ProvenanceExtra, PAGE_SIZE, STACK_ADDR, STACK_SIZE,
 };
 pub use crate::mono_hash_map::MonoHashMap;
 pub use crate::operator::EvalContextExt as OperatorEvalContextExt;
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index dcfb998c564..bd51207bbb1 100644
--- a/src/tools/miri/src/machine.rs
+++ b/src/tools/miri/src/machine.rs
@@ -35,7 +35,6 @@ use crate::{
 pub const PAGE_SIZE: u64 = 4 * 1024; // FIXME: adjust to target architecture
 pub const STACK_ADDR: u64 = 32 * PAGE_SIZE; // not really about the "stack", but where we start assigning integer addresses to allocations
 pub const STACK_SIZE: u64 = 16 * PAGE_SIZE; // whatever
-pub const NUM_CPUS: u64 = 1;
 
 /// Extra data stored with each stack frame
 pub struct FrameData<'tcx> {
@@ -407,6 +406,8 @@ pub struct MiriMachine<'mir, 'tcx> {
     pub(crate) gc_interval: u32,
     /// The number of blocks that passed since the last SbTag GC pass.
     pub(crate) since_gc: u32,
+    /// FIXME: docs,
+    pub(crate) num_cpus: u32,
 }
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
@@ -486,6 +487,7 @@ impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
             }),
             gc_interval: config.gc_interval,
             since_gc: 0,
+            num_cpus: config.num_cpus,
         }
     }
 
diff --git a/src/tools/miri/src/shims/unix/foreign_items.rs b/src/tools/miri/src/shims/unix/foreign_items.rs
index 153e5852dcc..c21e0441cac 100644
--- a/src/tools/miri/src/shims/unix/foreign_items.rs
+++ b/src/tools/miri/src/shims/unix/foreign_items.rs
@@ -225,13 +225,12 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             "sysconf" => {
                 let [name] = this.check_shim(abi, Abi::C { unwind: false }, link_name, args)?;
                 let name = this.read_scalar(name)?.to_i32()?;
-
                 // FIXME: Which of these are POSIX, and which are GNU/Linux?
                 // At least the names seem to all also exist on macOS.
                 let sysconfs: &[(&str, fn(&MiriInterpCx<'_, '_>) -> Scalar<Provenance>)] = &[
                     ("_SC_PAGESIZE", |this| Scalar::from_int(PAGE_SIZE, this.pointer_size())),
-                    ("_SC_NPROCESSORS_CONF", |this| Scalar::from_int(NUM_CPUS, this.pointer_size())),
-                    ("_SC_NPROCESSORS_ONLN", |this| Scalar::from_int(NUM_CPUS, this.pointer_size())),
+                    ("_SC_NPROCESSORS_CONF", |this| Scalar::from_int(this.machine.num_cpus, this.pointer_size())),
+                    ("_SC_NPROCESSORS_ONLN", |this| Scalar::from_int(this.machine.num_cpus, this.pointer_size())),
                     // 512 seems to be a reasonable default. The value is not critical, in
                     // the sense that getpwuid_r takes and checks the buffer length.
                     ("_SC_GETPW_R_SIZE_MAX", |this| Scalar::from_int(512, this.pointer_size()))
diff --git a/src/tools/miri/src/shims/windows/foreign_items.rs b/src/tools/miri/src/shims/windows/foreign_items.rs
index 53ab97b255e..c5f0de4307c 100644
--- a/src/tools/miri/src/shims/windows/foreign_items.rs
+++ b/src/tools/miri/src/shims/windows/foreign_items.rs
@@ -163,7 +163,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 )?;
                 // Set number of processors.
                 let num_cpus = system_info.offset(field_offsets[6], dword_layout, &this.tcx)?;
-                this.write_scalar(Scalar::from_int(NUM_CPUS, dword_layout.size), &num_cpus.into())?;
+                this.write_scalar(
+                    Scalar::from_int(this.machine.num_cpus, dword_layout.size),
+                    &num_cpus.into(),
+                )?;
             }
 
             // Thread-local storage
diff --git a/src/tools/miri/tests/pass/available-parallelism-miri-num-cpus.rs b/src/tools/miri/tests/pass/available-parallelism-miri-num-cpus.rs
new file mode 100644
index 00000000000..137fa510249
--- /dev/null
+++ b/src/tools/miri/tests/pass/available-parallelism-miri-num-cpus.rs
@@ -0,0 +1,8 @@
+//@compile-flags: -Zmiri-num-cpus=1024
+
+use std::num::NonZeroUsize;
+use std::thread::available_parallelism;
+
+fn main() {
+    assert_eq!(available_parallelism().unwrap(), NonZeroUsize::new(1024).unwrap());
+}