perf: improve parser performance

Signed-off-by: Henry <mail@henrygressmann.de>

1 files changed, 83 insertions, 87 deletions

diff --git a/crates/parser/src/parallel.rs b/crates/parser/src/parallel.rs
index 137eb2f..88d1035 100644
--- a/crates/parser/src/parallel.rs
+++ b/crates/parser/src/parallel.rs
@@ -4,7 +4,7 @@ use alloc::sync::Arc;
 use alloc::vec::Vec;
 use core::ops::Range;
 use tinywasm_types::ValueCounts;
-use wasmparser::{FuncValidatorAllocations, ValidatorResources};
+use wasmparser::{FuncValidatorAllocations, OperatorsReaderAllocations, ValidatorResources};
 
 pub(crate) enum FunctionBodyInput<'a> {
     Borrowed(wasmparser::FunctionBody<'a>),
@@ -22,25 +22,29 @@ pub(crate) struct OwnedFunctionBody {
 pub(crate) struct PendingFunction<'a> {
     pub ordinal: usize,
     pub results: ValueCounts,
-    pub ty_idx: u32,
     pub func_to_validate: Option<wasmparser::FuncToValidate<ValidatorResources>>,
+    pub ty_idx: u32,
     pub body: FunctionBodyInput<'a>,
 }
 
-pub(crate) const MIN_FUNCTIONS: usize = 8;
-const MIN_CODE_SECTION_BYTES: usize = 32 * 1024;
-const MIN_FUNCTION_BODY_BYTES: usize = 4;
+const MIN_FUNCTIONS: usize = 8;
+const MIN_CODE_SECTION_BYTES: usize = 16 * 1024;
+const MIN_FUNCTION_BODY_BYTES: usize = 16;
+const MIN_FUNCTIONS_PER_WORKER: usize = 4;
+const MAX_WORKERS: usize = 12;
 
-pub(crate) fn should_parallelize_function(body_len: usize) -> bool {
-    body_len >= MIN_FUNCTION_BODY_BYTES
+fn body_len(body: &FunctionBodyInput<'_>) -> usize {
+    match body {
+        FunctionBodyInput::Borrowed(func) => func.as_bytes().len(),
+        FunctionBodyInput::Owned(body) => body.body_range.len(),
+    }
 }
 
 pub(crate) fn should_use_parallel(options: &ParserOptions, num_functions: usize, code_section_bytes: usize) -> bool {
-    if num_functions < MIN_FUNCTIONS || code_section_bytes < MIN_CODE_SECTION_BYTES {
-        return false;
-    }
-
-    worker_count(options, num_functions) > 1
+    num_functions >= MIN_FUNCTIONS
+        && code_section_bytes >= MIN_CODE_SECTION_BYTES
+        && code_section_bytes / num_functions >= MIN_FUNCTION_BODY_BYTES
+        && worker_count(options, num_functions) > 1
 }
 
 fn worker_count(options: &ParserOptions, num_functions: usize) -> usize {
@@ -48,15 +52,7 @@ fn worker_count(options: &ParserOptions, num_functions: usize) -> usize {
         .parser_threads()
         .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1))
         .max(1);
-
-    requested.min(num_functions).max(1)
-}
-
-fn body_len(body: &FunctionBodyInput<'_>) -> usize {
-    match body {
-        FunctionBodyInput::Borrowed(func) => func.as_bytes().len(),
-        FunctionBodyInput::Owned(body) => body.body_range.len(),
-    }
+    requested.min(MAX_WORKERS).min(num_functions.div_ceil(MIN_FUNCTIONS_PER_WORKER)).max(1)
 }
 
 fn process_function_job(
@@ -65,16 +61,18 @@ fn process_function_job(
     options: &ParserOptions,
     imported_func_count: usize,
     imported_memory_count: u32,
-) -> Result<(usize, FunctionCode)> {
-    let validator = job.func_to_validate.map(|func| func.into_validator(FuncValidatorAllocations::default()));
-    let (code, _, _) = match job.body {
+    validator_allocs: Option<FuncValidatorAllocations>,
+    reader_allocs: OperatorsReaderAllocations,
+) -> Result<(FunctionCode, Option<FuncValidatorAllocations>, OperatorsReaderAllocations)> {
+    let validator = job.func_to_validate.map(|func| func.into_validator(validator_allocs.unwrap_or_default()));
+    let (code, validator_allocs, reader_allocs) = match job.body {
         FunctionBodyInput::Borrowed(func) => {
-            conversion::convert_module_code(func, validator, Default::default(), metadata, job.ty_idx)?
+            conversion::convert_module_code(func, validator, reader_allocs, metadata, job.ty_idx)?
         }
         FunctionBodyInput::Owned(body) => {
             let reader = wasmparser::BinaryReader::new(&body.section_bytes[body.body_range], body.body_offset);
             let func = wasmparser::FunctionBody::new(reader);
-            conversion::convert_module_code(func, validator, Default::default(), metadata, job.ty_idx)?
+            conversion::convert_module_code(func, validator, reader_allocs, metadata, job.ty_idx)?
         }
     };
 
@@ -86,79 +84,77 @@ fn process_function_job(
         imported_memory_count,
     )?;
 
-    Ok((job.ordinal, code))
+    Ok((code, validator_allocs, reader_allocs))
 }
 
-pub(crate) fn process_pending(
-    pending: Vec<PendingFunction<'_>>,
+fn process_chunk<'a>(
+    jobs: impl IntoIterator<Item = PendingFunction<'a>>,
     metadata: &crate::visit::ModuleMetadata,
     options: &ParserOptions,
     imported_func_count: usize,
     imported_memory_count: u32,
 ) -> Result<Vec<FunctionCode>> {
-    let (small_jobs, large_jobs): (Vec<_>, Vec<_>) =
-        pending.into_iter().partition(|job| !should_parallelize_function(body_len(&job.body)));
+    let mut validator_allocs = None;
+    let mut reader_allocs = OperatorsReaderAllocations::default();
+    let jobs = jobs.into_iter();
+    let mut codes = Vec::with_capacity(jobs.size_hint().0);
 
-    let mut codes = small_jobs
-        .into_iter()
-        .map(|job| process_function_job(job, metadata, options, imported_func_count, imported_memory_count))
-        .collect::<Result<Vec<_>>>()?;
+    for job in jobs {
+        let (code, next_validator_allocs, next_reader_allocs) = process_function_job(
+            job,
+            metadata,
+            options,
+            imported_func_count,
+            imported_memory_count,
+            validator_allocs,
+            reader_allocs,
+        )?;
+        codes.push(code);
+        validator_allocs = next_validator_allocs;
+        reader_allocs = next_reader_allocs;
+    }
 
-    let num_workers = worker_count(options, large_jobs.len());
+    Ok(codes)
+}
+
+pub(crate) fn process_pending(
+    pending: Vec<PendingFunction<'_>>,
+    metadata: &crate::visit::ModuleMetadata,
+    options: &ParserOptions,
+    imported_func_count: usize,
+    imported_memory_count: u32,
+) -> Result<Vec<FunctionCode>> {
+    let num_workers = worker_count(options, pending.len());
     if num_workers == 1 {
-        codes.extend(
-            large_jobs
-                .into_iter()
-                .map(|job| process_function_job(job, metadata, options, imported_func_count, imported_memory_count))
-                .collect::<Result<Vec<_>>>()?,
-        );
-    } else {
-        let chunk_size = large_jobs.len().div_ceil(num_workers);
-        let chunks = {
-            let mut chunks = Vec::with_capacity(num_workers);
-            let mut iter = large_jobs.into_iter();
-            while let Some(first) = iter.next() {
-                let mut chunk = alloc::vec![first];
-                for _ in 1..chunk_size {
-                    match iter.next() {
-                        Some(job) => chunk.push(job),
-                        None => break,
-                    }
-                }
-                chunks.push(chunk);
+        return process_chunk(pending, metadata, options, imported_func_count, imported_memory_count);
+    }
+    let code_count = pending.len();
+    let chunk_size = pending.len().div_ceil(num_workers);
+    let chunk_bytes = pending.iter().map(|job| body_len(&job.body)).sum::<usize>().div_ceil(num_workers);
+    std::thread::scope(|scope| {
+        let mut jobs = pending.into_iter();
+        let mut handles = Vec::with_capacity(num_workers);
+        while let Some(first) = jobs.next() {
+            let mut chunk = Vec::with_capacity(chunk_size);
+            let mut bytes = body_len(&first.body);
+            chunk.push(first);
+            while bytes < chunk_bytes
+                && let Some(job) = jobs.next()
+            {
+                bytes += body_len(&job.body);
+                chunk.push(job);
             }
-            chunks
-        };
-
-        let results: Vec<Result<(usize, FunctionCode)>> = std::thread::scope(|s| {
-            let handles: Vec<_> = chunks
-                .into_iter()
-                .map(|chunk| {
-                    s.spawn(move || {
-                        chunk
-                            .into_iter()
-                            .map(|job| {
-                                process_function_job(job, metadata, options, imported_func_count, imported_memory_count)
-                            })
-                            .collect::<Vec<_>>()
-                    })
-                })
-                .collect();
-
             handles
-                .into_iter()
-                .flat_map(|handle| match handle.join() {
-                    Ok(results) => results,
-                    Err(_) => alloc::vec![Err(ParseError::Other("worker thread panicked".into()))],
-                })
-                .collect()
-        });
-
-        for result in results {
-            codes.push(result?);
+                .push(scope.spawn(move || {
+                    process_chunk(chunk, metadata, options, imported_func_count, imported_memory_count)
+                }));
         }
-    }
 
-    codes.sort_by_key(|(ordinal, _)| *ordinal);
-    Ok(codes.into_iter().map(|(_, code)| code).collect())
+        let mut codes = Vec::with_capacity(code_count);
+        for handle in handles {
+            let chunk = handle.join().map_err(|_| ParseError::Other("worker thread panicked".into()))??;
+            codes.extend(chunk);
+        }
+        Ok(codes)
+    })
 }