cap thread pool size at # of harnesses

* thread pools are now dynamically sized
* that size is set based on a const var capped by the # of harnesses

Author: AlexanderPortland

kani-compiler/src/codegen_cprover_gotoc/compiler_interface.rs

@@ -43,6 +43,7 @@ use stable_mir::CrateDef;
 use stable_mir::mir::mono::{Instance, MonoItem};
 use stable_mir::rustc_internal;
 use std::any::Any;
+use std::cmp::min;
 use std::collections::BTreeMap;
 use std::fmt::Write;
 use std::fs::File;
@@ -74,15 +75,15 @@ impl GotocCodegenBackend {
     ///
     /// Invariant: iff `check_contract.is_some()` then `return.2.is_some()`
     #[allow(clippy::too_many_arguments)]
-    fn codegen_items<'tcx, const NUM_THREADS: usize>(
+    fn codegen_items<'tcx>(
         &self,
         tcx: TyCtxt<'tcx>,
         starting_items: &[MonoItem],
         symtab_goto: &Path,
         machine_model: &MachineModel,
         check_contract: Option<InternalDefId>,
         mut transformer: BodyTransformation,
-        thread_pool: &ThreadPool<NUM_THREADS>,
+        thread_pool: &ThreadPool,
     ) -> (GotocCtx<'tcx>, Vec<MonoItem>, Option<AssignsContract>) {
         // This runs reachability analysis before global passes are applied.
         //
@@ -349,13 +350,24 @@ impl CodegenBackend for GotocCodegenBackend {
                 return codegen_results(tcx, &results.machine_model);
             }
 
-            let export_thread_pool = ThreadPool::<NUM_FILE_EXPORT_THREADS>::new();
+            // Create an empty thread pool. We will set the size later once we
+            // concretely know the # of harnesses we need to analyze.
+            let mut export_thread_pool = ThreadPool::empty();
 
             match reachability {
                 ReachabilityType::AllFns | ReachabilityType::Harnesses => {
                     let mut units = CodegenUnits::new(&queries, tcx);
                     let mut modifies_instances = vec![];
                     let mut loop_contracts_instances = vec![];
+
+                    // If we are verifying all harnesses, cap the number of threads at `[# of harnesses] - 1`.
+                    // This is just a common sense rule to keep us from spinning up threads that couldn't possibly
+                    // provide additional parallelism over existing threads. One thread can potentially work on each
+                    // harness' goto file but the main compiler thread can just handle the last file itself.
+                    let num_harnesses: usize = units.iter().map(|unit| unit.harnesses.len()).sum();
+                    export_thread_pool
+                        .add_workers(min(NUM_FILE_EXPORT_THREADS, num_harnesses.saturating_sub(1)));
+
                     // Cross-crate collecting of all items that are reachable from the crate harnesses.
                     for unit in units.iter() {
                         // We reset the body cache for now because each codegen unit has different
@@ -391,6 +403,11 @@ impl CodegenBackend for GotocCodegenBackend {
                 ReachabilityType::None => unreachable!(),
                 ReachabilityType::PubFns => {
                     let unit = CodegenUnit::default();
+
+                    // Here, it's more difficult to determine which functions will be analyzed,
+                    // so just use the max NUM_FILE_EXPORT_THREADS.
+                    export_thread_pool.add_workers(NUM_FILE_EXPORT_THREADS);
+
                     let transformer = BodyTransformation::new(&queries, tcx, &unit);
                     let main_instance =
                         stable_mir::entry_fn().map(|main_fn| Instance::try_from(main_fn).unwrap());

kani-compiler/src/codegen_cprover_gotoc/utils/file_writing_pool.rs

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 use std::collections::BTreeMap;
-use std::convert::identity;
 use std::path::PathBuf;
 use std::sync::mpmc::{Receiver, Sender, channel};
 use std::sync::mpsc::TryRecvError;
@@ -50,34 +49,37 @@ impl FileDataToWrite {
 ///
 /// File data can be sent to the `work_queue`. This will wake a worker thread which will then serialize and write
 /// it to disk in parallel, allowing the main compiler thread to continue codegen.
-pub struct ThreadPool<const N: usize> {
-    pub(crate) work_queue: Sender<WithInterner<FileDataToWrite>>,
-    join_handles: [JoinHandle<WorkerReturn>; N],
+pub struct ThreadPool {
+    pub(crate) work_queue: Sender<WorkToSend>,
+    work_queue_recv: Receiver<WorkToSend>,
+    join_handles: Vec<JoinHandle<WorkerReturn>>,
 }
 
 type WorkerReturn = ();
 
 type WorkToSend = WithInterner<FileDataToWrite>;
-impl<const N: usize> ThreadPool<N> {
-    pub fn new() -> Self {
+impl ThreadPool {
+    pub fn empty() -> Self {
         let (work_queue_send, work_queue_recv) = channel();
+        ThreadPool { work_queue: work_queue_send, work_queue_recv, join_handles: Vec::new() }
+    }
 
-        // Spawn a thread for each worker, and pass it the recv end of the work queue.
-        let join_handles: [JoinHandle<()>; N] = core::array::from_fn(identity).map(|_| {
-            let new_work_queue_recv = work_queue_recv.clone();
-            std::thread::spawn(move || {
-                worker_loop(new_work_queue_recv);
-            })
-        });
+    fn new_worker(work_queue: &Receiver<WorkToSend>) -> JoinHandle<()> {
+        let new_work_queue = work_queue.clone();
+        std::thread::spawn(move || {
+            worker_loop(new_work_queue);
+        })
+    }
 
-        ThreadPool { work_queue: work_queue_send, join_handles }
+    pub fn add_workers(&mut self, count: usize) {
+        self.join_handles.extend((0..count).map(|_| Self::new_worker(&self.work_queue_recv)));
     }
 
     /// Try to send work to the work queue, or do it yourself if there's no worker threads.
     /// Will only fail if all recievers have disconnected.
     pub fn send_work(&self, work: WorkToSend) -> Result<(), &str> {
         // If we don't have any workers, just synchronously handle the work ourselves.
-        if N == 0 {
+        if self.join_handles.is_empty() {
             handle_file(work.into_inner());
             return Ok(());
         }