perf(cpu-ep): build the CPU kernels as one codegen unit (1.2-2.1x across the activation family)

Cargo.toml

@@ -133,6 +133,27 @@ license = "MIT"
 authors = ["Justin Chu"]
 repository = "https://github.com/justinchuby/onnx-genai"
 
+# The CPU EP's elementwise kernels are AVX2 `#[target_feature]` functions whose
+# hot loops only vectorise when the compiler can see the dispatcher, the chunk
+# loop and the per-vector body together. The default sixteen codegen units split
+# that chain and the loops come out scalar, at a cost of up to 2.65x on the
+# elementwise family (`Sqrt` f32 at 4096 elements: 0.639 ns/elem against 0.241;
+# 77 of 105 measured cases move by more than 1.15x).
+# The failure is silent -- identical source, identical tests, no runtime-path
+# difference -- so it is pinned here rather than left to whichever partition the
+# compiler happens to pick. See `docs/performance/CPU_ACTIVATION_GAPS.md`.
+#
+# Scoped to this package alone: the rest of the workspace keeps parallel codegen
+# and only this crate pays the serial build (~19 s -> ~65 s).
+[profile.release.package.onnx-runtime-ep-cpu]
+codegen-units = 1
+
+# `bench` does not inherit a per-package override from `release`, so `cargo
+# bench` would rebuild the crate as sixteen units and measure the de-vectorized
+# code -- the one thing a benchmark must not do.
+[profile.bench.package.onnx-runtime-ep-cpu]
+codegen-units = 1
+
 # Central workspace lint policy (rico audit #11). Keep this SMALL, WARN-only,
 # and defensible: it is a shared, extensible anchor, not a cleanup crusade.
 # Member crates opt in via `[lints] workspace = true`. Do NOT add `deny`/`forbid`

crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs

@@ -5705,6 +5705,225 @@ fn percentile(samples: &mut [f64], p: f64) -> f64 {
     samples[idx]
 }
 
+/// A single-input elementwise case: one node, one tensor in, same shape out.
+///
+/// Reuses [`MatmulFamilyCase`], which is a single-node generated model plus its
+/// runtime inputs and nothing matmul-specific beyond its name.
+///
+/// `positive` maps the generated activations into `[0.5, 1.5]` for the ops that
+/// are undefined or uninteresting on negatives (`Sqrt`, `Log`); everything else
+/// gets the two-sided `[-1, 1]` range so saturating branches are exercised on
+/// both sides.
+///
+/// `tolerance` is inert for these: the only consumer of the unary cases is the
+/// timing A/B, which asserts assignment and then measures. Elementwise numerics
+/// are covered against closed forms in `onnx-runtime-ep-cpu`'s own kernel
+/// tests, not here.
+#[allow(clippy::too_many_arguments)]
+fn unary_case(
+    name: &'static str,
+    op: &'static str,
+    domain: &'static str,
+    elem: ort::ONNXTensorElementDataType,
+    len: usize,
+    opset: u64,
+    positive: bool,
+    tolerance: f32,
+) -> MatmulFamilyCase {
+    let dims = vec![1, len as i64];
+    let opset_imports = if domain.is_empty() {
+        format!("[{{ version: {opset} }}]")
+    } else {
+        format!("[{{ version: {opset} }}, {{ domain: \"{domain}\" version: 1 }}]")
+    };
+    let attrs = match op {
+        // `Gelu`'s default is already "none" (exact); spell whichever variant
+        // the case name promises so the model cannot silently drift.
+        "Gelu" if name.contains("_tanh") => {
+            " attribute: [{ name: \"approximate\" type: STRING s: \"tanh\" }]"
+        }
+        "Gelu" => " attribute: [{ name: \"approximate\" type: STRING s: \"none\" }]",
+        _ => "",
+    };
+    let model = format!(
+        r#"ir_version: 10
+graph {{
+  node: [{{ input: ["X"] output: ["Y"] op_type: "{op}" domain: "{domain}"{attrs} }}]
+  name: "{name}"
+  input: [{}]
+  output: [{}]
+}}
+opset_import: {opset_imports}
+"#,
+        textproto_value_info("X", elem, &dims),
+        textproto_value_info("Y", elem, &dims),
+    );
+    let mut x = activation_f32(1, len);
+    if positive {
+        for v in &mut x {
+            *v = v.abs() + 0.5;
+        }
+    }
+    let data = if elem == ELEM_F16 {
+        f16_slice_to_bytes(&x)
+    } else {
+        f32_slice_to_bytes(&x)
+    };
+    MatmulFamilyCase {
+        name,
+        op,
+        model,
+        inputs: vec![GeneratedInput {
+            name: "X",
+            elem,
+            dims,
+            data,
+        }],
+        output: "Y",
+        output_elem: elem,
+        ort_can_build: true,
+        tolerance,
+    }
+}
+
+/// The elementwise grid `CPU_ACTIVATION_GAPS.md` is written against, as
+/// session-level cases.
+///
+/// Two sizes: a decode-width row (4 Ki) where per-node dispatch overhead is
+/// still visible, and 1 Mi where the kernel is the whole cost. Both are the
+/// sizes the published ratio table uses, so a number from here is directly
+/// comparable to a number in that file.
+fn unary_bench_cases() -> Vec<MatmulFamilyCase> {
+    const K4: usize = 4096;
+    const M1: usize = 1 << 20;
+    vec![
+        unary_case(
+            "bench_tanh_f32_4k",
+            "Tanh",
+            "",
+            ELEM_F32,
+            K4,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_tanh_f32_1m",
+            "Tanh",
+            "",
+            ELEM_F32,
+            M1,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_sigmoid_f32_4k",
+            "Sigmoid",
+            "",
+            ELEM_F32,
+            K4,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_sigmoid_f32_1m",
+            "Sigmoid",
+            "",
+            ELEM_F32,
+            M1,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case("bench_erf_f32_4k", "Erf", "", ELEM_F32, K4, 17, false, 1e-4),
+        unary_case("bench_erf_f32_1m", "Erf", "", ELEM_F32, M1, 17, false, 1e-4),
+        unary_case(
+            "bench_gelu_tanh_f32_1m",
+            "Gelu",
+            "",
+            ELEM_F32,
+            M1,
+            20,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_gelu_exact_f32_1m",
+            "Gelu",
+            "",
+            ELEM_F32,
+            M1,
+            20,
+            false,
+            1e-4,
+        ),
+        unary_case("bench_exp_f32_1m", "Exp", "", ELEM_F32, M1, 17, false, 1e-4),
+        unary_case(
+            "bench_relu_f32_1m",
+            "Relu",
+            "",
+            ELEM_F32,
+            M1,
+            17,
+            false,
+            0.0,
+        ),
+        unary_case(
+            "bench_sqrt_f32_4k",
+            "Sqrt",
+            "",
+            ELEM_F32,
+            K4,
+            17,
+            true,
+            1e-4,
+        ),
+        unary_case(
+            "bench_sqrt_f32_1m",
+            "Sqrt",
+            "",
+            ELEM_F32,
+            M1,
+            17,
+            true,
+            1e-4,
+        ),
+        unary_case(
+            "bench_quickgelu_f32_1m",
+            "QuickGelu",
+            "com.microsoft",
+            ELEM_F32,
+            M1,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_fastgelu_f32_1m",
+            "FastGelu",
+            "com.microsoft",
+            ELEM_F32,
+            M1,
+            17,
+            false,
+            1e-4,
+        ),
+        unary_case(
+            "bench_tanh_f16_1m",
+            "Tanh",
+            "",
+            ELEM_F16,
+            M1,
+            17,
+            false,
+            5e-3,
+        ),
+        unary_case("bench_exp_f16_1m", "Exp", "", ELEM_F16, M1, 17, false, 5e-3),
+    ]
+}
+
 /// Shapes worth measuring on the plugin path, at sizes a real decode/prefill
 /// step reaches rather than the smallest size that proves correctness.
 fn matmul_bench_cases() -> Vec<MatmulFamilyCase> {
@@ -5819,7 +6038,7 @@ fn plugin_path_ab_vs_plain_ort() {
     println!(
         "case,ours_p50_ms,ort_p50_ms,ratio_p50,ours_p90_ms,ort_p90_ms,ratio_p90,cold_ours_ms,cold_ort_ms"
     );
-    for case in matmul_bench_cases() {
+    for case in matmul_bench_cases().into_iter().chain(unary_bench_cases()) {
         if !filter.is_empty() && !case.name.contains(&filter) {
             continue;
         }

crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs

@@ -4221,6 +4221,79 @@ mod chunking_instantiation_is_local {
     }
 }
 
+/// The same hazard from the other side: keeping `run_chunked` private stops one
+/// *known* way of repartitioning the crate, but the compiler is free to split it
+/// sixteen ways regardless, and by default it does. Measured on the release
+/// profile, this crate built as one codegen unit against the default sixteen:
+///
+/// | case (1 thread, p50 of 5 interleaved runs) | 16 CGUs | 1 CGU | ratio |
+/// |---|---|---|---|
+/// | `Sqrt` f32, 4096 elems | 0.639 ns/elem | 0.241 | 2.65x |
+/// | `Tanh` f32, 4096 elems | 0.728 | 0.381 | 1.91x |
+/// | `Sigmoid` f32, 4096 elems | 0.729 | 0.392 | 1.86x |
+/// | `Sqrt` f16, 4096 elems | 0.907 | 0.511 | 1.78x |
+/// | `Tanh` f32, 2 Mi elems | 0.202 | 0.099 | 2.04x |
+/// | `QuickGelu` f32, 16 elems (control) | 8.56 | 8.66 | 0.99x |
+///
+/// 77 of 105 measured cases move by more than 1.15x; the ones that do not are
+/// the 16-element shapes, where dispatch rather than the loop is the cost. That
+/// is the signature of a de-vectorized inner loop, not of noise.
+///
+/// So the workspace pins `codegen-units = 1` for this package. That is a build
+/// setting a rebase can drop without any test noticing, which is exactly how the
+/// same class of regression reached `main` before, so this test reads it back.
+#[cfg(test)]
+mod codegen_units_are_pinned {
+    /// Both shipping profiles must still build this crate as a single codegen
+    /// unit.
+    ///
+    /// `bench` is checked too because a per-package override does not inherit
+    /// from `release`, so dropping it would leave `cargo bench` measuring the
+    /// de-vectorized build -- the one place a wrong number is most expensive.
+    #[test]
+    fn the_workspace_pins_one_codegen_unit_for_this_crate() {
+        for profile in ["release", "bench"] {
+            assert_pinned(profile);
+        }
+    }
+
+    fn assert_pinned(profile: &str) {
+        let manifest = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
+            .parent()
+            .and_then(std::path::Path::parent)
+            .expect("crates/<pkg> has a workspace root two levels up")
+            .join("Cargo.toml");
+        let text = std::fs::read_to_string(&manifest)
+            .unwrap_or_else(|e| panic!("read {}: {e}", manifest.display()));
+        let header = format!("[profile.{profile}.package.onnx-runtime-ep-cpu]");
+        let Some(start) = text.find(&header) else {
+            panic!(
+                "{} no longer contains {header}; building this crate as the \
+                 default sixteen codegen units de-vectorizes the AVX2 \
+                 elementwise kernels and costs up to 2.65x",
+                manifest.display()
+            );
+        };
+        let section = &text[start + header.len()..];
+        let section = &section[..section.find("\n[").unwrap_or(section.len())];
+        let pinned = section.lines().any(|line| {
+            let line = line.trim();
+            line.starts_with("codegen-units")
+                && line.split_once('=').is_some_and(|(_, value)| {
+                    // Tolerate an inline comment: the most likely future edit to
+                    // this line is someone explaining it, and a guard that fails
+                    // on its own documentation teaches people to delete it.
+                    let value = value.split('#').next().unwrap_or(value);
+                    value.trim().trim_end_matches(',') == "1"
+                })
+        });
+        assert!(
+            pinned,
+            "{header} exists but does not set `codegen-units = 1`; found:{section}"
+        );
+    }
+}
+
 /// The host-pool split: what happens when ORT lends us its intra-op threads.
 ///
 /// The production host is ORT's `KernelContext_ParallelFor`, which needs a