Auto merge of #2562 - pvdrz:miri-num-cpus, r=RalfJung

Add flag to specify the number of cpus

Apparently you can't rename a branch from github's website without it closing all your PRs with that branch. So this is  just #2545

8 files changed, 33 insertions, 8 deletions

diff --git a/src/tools/miri/README.md b/src/tools/miri/README.md
index 2e27bb39afd..bb9b002c689 100644
--- a/src/tools/miri/README.md
+++ b/src/tools/miri/README.md
@@ -293,6 +293,9 @@ environment variable. We first document the most relevant and most commonly used
   value of forwarded variables stays the same. Has no effect if `-Zmiri-disable-isolation` is set.
 * `-Zmiri-ignore-leaks` disables the memory leak checker, and also allows some
   remaining threads to exist when the main thread exits.
+* `-Zmiri-num-cpus` states the number of available CPUs to be reported by miri. By default, the
+  number of available CPUs is `1`. Note that this flag does not affect how miri handles threads in
+  any way.
 * `-Zmiri-permissive-provenance` disables the warning for integer-to-pointer casts and
   [`ptr::from_exposed_addr`](https://doc.rust-lang.org/nightly/std/ptr/fn.from_exposed_addr.html).
   This will necessarily miss some bugs as those operations are not efficiently and accurately
@@ -357,7 +360,7 @@ to Miri failing to detect cases of undefined behavior in a program.
   This is **work in progress**; currently, only integer arguments and return values are
   supported (and no, pointer/integer casts to work around this limitation will not work;
   they will fail horribly). It also only works on unix hosts for now.
-  Follow [the discussion on supporting other types](https://github.com/rust-lang/miri/issues/2365). 
+  Follow [the discussion on supporting other types](https://github.com/rust-lang/miri/issues/2365).
 * `-Zmiri-measureme=<name>` enables `measureme` profiling for the interpreted program.
    This can be used to find which parts of your program are executing slowly under Miri.
    The profile is written out to a file with the prefix `<name>`, and can be processed
@@ -387,7 +390,7 @@ to Miri failing to detect cases of undefined behavior in a program.
   Borrows "protectors". Specifying this argument multiple times does not overwrite the previous
   values, instead it appends its values to the list. Listing an id multiple times has no effect.
 * `-Zmiri-track-pointer-tag=<tag1>,<tag2>,...` shows a backtrace when a given pointer tag
-  is created and when (if ever) it is popped from a borrow stack (which is where the tag becomes invalid 
+  is created and when (if ever) it is popped from a borrow stack (which is where the tag becomes invalid
   and any future use of it will error).  This helps you in finding out why UB is
   happening and where in your code would be a good place to look for it.
   Specifying this argument multiple times does not overwrite the previous
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..b211f3c5f71 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,
+    /// The number of CPUs to be reported by miri.
+    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 0fe21c9243a..bf4e21319b5 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 9a1621dadc3..03df78c1af7 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> {
@@ -410,6 +409,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,
+    /// The number of CPUs to be reported by miri.
+    pub(crate) num_cpus: u32,
 }
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {
@@ -489,6 +490,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());
+}