feat: add Profile-Guided Optimization (PGO) support

src/build_context/builder.rs

@@ -29,6 +29,7 @@ pub struct BuildContextBuilder {
     editable: bool,
     sdist_only: bool,
     pyproject_toml_path: Option<PathBuf>,
+    pgo: bool,
 }
 
 impl BuildContextBuilder {
@@ -39,6 +40,7 @@ impl BuildContextBuilder {
             editable: false,
             sdist_only: false,
             pyproject_toml_path: None,
+            pgo: false,
         }
     }
 
@@ -62,6 +64,11 @@ impl BuildContextBuilder {
         self
     }
 
+    pub fn pgo(mut self, pgo: bool) -> Self {
+        self.pgo = pgo;
+        self
+    }
+
     #[instrument(skip_all)]
     pub fn build(self) -> Result<BuildContext> {
         let Self {
@@ -70,6 +77,7 @@ impl BuildContextBuilder {
             editable,
             sdist_only,
             pyproject_toml_path: explicit_pyproject_path,
+            pgo,
         } = self;
         build_options.compression.validate();
         let ProjectResolver {
@@ -215,6 +223,21 @@ impl BuildContextBuilder {
             Vec::new()
         };
 
+        let pgo_command = if pgo {
+            let cmd = pyproject
+                .and_then(|p| p.pgo_command())
+                .map(|s| s.trim().to_string())
+                .filter(|s| !s.is_empty());
+            if cmd.is_none() {
+                bail!(
+                    "--pgo requires a non-empty `pgo-command` to be set in `[tool.maturin]` in pyproject.toml"
+                );
+            }
+            cmd
+        } else {
+            None
+        };
+
         Ok(BuildContext {
             target,
             compile_targets,
@@ -243,6 +266,8 @@ impl BuildContextBuilder {
             include_import_lib,
             include_debuginfo,
             conditional_features,
+            pgo_phase: None,
+            pgo_command,
         })
     }
 

src/build_context/mod.rs

@@ -11,6 +11,7 @@ use crate::build_options::CargoOptions;
 use crate::compile::CompileTarget;
 use crate::compression::CompressionOptions;
 use crate::module_writer::{WheelWriter, write_pth};
+use crate::pgo::{PgoContext, PgoPhase};
 use crate::project_layout::ProjectLayout;
 use crate::pyproject_toml::ConditionalFeature;
 use crate::sbom::generate_sbom_data;
@@ -87,6 +88,10 @@ pub struct BuildContext {
     pub include_debuginfo: bool,
     /// Cargo features conditionally enabled based on the target Python version/implementation
     pub conditional_features: Vec<ConditionalFeature>,
+    /// Current PGO build phase (if PGO is enabled)
+    pub pgo_phase: Option<PgoPhase>,
+    /// PGO training command from pyproject.toml (only set when --pgo is passed)
+    pub pgo_command: Option<String>,
 }
 
 /// The wheel file location and its Python version tag (e.g. `py3`).
@@ -100,6 +105,14 @@ impl BuildContext {
     /// correct builder.
     #[instrument(skip_all)]
     pub fn build_wheels(&self) -> Result<Vec<BuiltWheelMetadata>> {
+        if self.pgo_command.is_some() {
+            return self.build_wheels_pgo();
+        }
+        self.build_wheels_inner()
+    }
+
+    /// Standard wheel build pipeline (no PGO).
+    fn build_wheels_inner(&self) -> Result<Vec<BuiltWheelMetadata>> {
         fs::create_dir_all(&self.out)
             .context("Failed to create the target directory for the wheels")?;
 
@@ -139,6 +152,177 @@ impl BuildContext {
         Ok(wheels)
     }
 
+    /// PGO three-phase build: instrumented → instrumentation → optimized.
+    ///
+    /// For non-abi3 PyO3 builds each interpreter gets its own
+    /// instrument → train → optimize cycle, because different Python
+    /// versions produce different compiled code via PyO3's build-time
+    /// configuration, so sharing a single profile across interpreters
+    /// causes "function control flow change detected (hash mismatch)"
+    /// warnings and defeats the purpose of PGO.
+    ///
+    /// For single-artifact builds (abi3, cffi, uniffi, bin) the compiled
+    /// code is identical across interpreters, so one PGO cycle suffices.
+    fn build_wheels_pgo(&self) -> Result<Vec<BuiltWheelMetadata>> {
+        let pgo_command = self
+            .pgo_command
+            .as_ref()
+            .expect("pgo_command must be set when build_wheels_pgo is called");
+
+        let needs_per_interpreter_pgo = matches!(
+            self.bridge(),
+            BridgeModel::PyO3(crate::PyO3 { abi3: None, .. })
+        );
+
+        eprintln!("🚀 Starting PGO build...");
+
+        // Verify llvm-profdata is available before starting
+        PgoContext::find_llvm_profdata()?;
+
+        if needs_per_interpreter_pgo {
+            self.build_wheels_pgo_per_interpreter(pgo_command)
+        } else {
+            self.build_wheels_pgo_single_pass(pgo_command)
+        }
+    }
+
+    /// Clone this context with PGO disabled (to prevent recursion) and
+    /// the given PGO phase set.
+    fn clone_for_pgo(&self, phase: PgoPhase) -> Self {
+        let mut ctx = self.clone();
+        ctx.pgo_command = None;
+        ctx.pgo_phase = Some(phase);
+        ctx
+    }
+
+    /// Single-pass PGO for abi3, cffi, uniffi, and bin builds where the
+    /// compiled artifact is the same regardless of the Python interpreter.
+    fn build_wheels_pgo_single_pass(&self, pgo_command: &str) -> Result<Vec<BuiltWheelMetadata>> {
+        let instrumentation_python = self
+            .interpreter
+            .first()
+            .context(
+                "PGO builds require a Python interpreter. \
+                 Please specify one with `--interpreter`.",
+            )?
+            .executable
+            .clone();
+
+        let pgo_ctx = PgoContext::new(pgo_command.to_owned())?;
+
+        // Phase 1: Build a single instrumented wheel for training.
+        // We only need one wheel for profiling — the compiled native code is
+        // identical across interpreters for abi3/cffi/uniffi/bin builds.
+        eprintln!("📊 Phase 1/3: Building instrumented wheel...");
+        let mut instrumented_ctx = self.clone_for_pgo(PgoPhase::Generate(
+            pgo_ctx.profdata_dir_path().to_path_buf(),
+        ));
+        instrumented_ctx.interpreter = vec![self.interpreter[0].clone()];
+        let instrumented_out =
+            tempfile::TempDir::new().context("Failed to create temp dir for instrumented wheel")?;
+        instrumented_ctx.out = instrumented_out.path().to_path_buf();
+        let instrumented_wheels = instrumented_ctx.build_wheels_inner()?;
+
+        // Phase 2: Instrumentation
+        eprintln!("🔬 Phase 2/3: Running PGO instrumentation...");
+        let instrumented_wheel_path = &instrumented_wheels
+            .first()
+            .context("No instrumented wheel was built")?
+            .0;
+        pgo_ctx.run_instrumentation(&instrumentation_python, instrumented_wheel_path, self)?;
+        pgo_ctx.merge_profiles()?;
+
+        // Phase 3: Optimized build
+        eprintln!("⚡ Phase 3/3: Building PGO-optimized wheel...");
+        let optimized_ctx =
+            self.clone_for_pgo(PgoPhase::Use(pgo_ctx.merged_profdata_path().to_path_buf()));
+        let wheels = optimized_ctx.build_wheels_inner()?;
+
+        eprintln!("🎉 PGO build complete!");
+        Ok(wheels)
+    }
+
+    /// Per-interpreter PGO for non-abi3 PyO3 builds. Each interpreter gets
+    /// its own instrument → train → optimize cycle so that profiles match
+    /// the exact compiled code for that Python version.
+    fn build_wheels_pgo_per_interpreter(
+        &self,
+        pgo_command: &str,
+    ) -> Result<Vec<BuiltWheelMetadata>> {
+        fs::create_dir_all(&self.out)
+            .context("Failed to create the target directory for the wheels")?;
+
+        let sbom_data = generate_sbom_data(self)?;
+        let mut wheels = Vec::new();
+
+        for (i, python_interpreter) in self.interpreter.iter().enumerate() {
+            eprintln!(
+                "📊 [{}/{}] PGO cycle for {} {}.{}...",
+                i + 1,
+                self.interpreter.len(),
+                python_interpreter.interpreter_kind,
+                python_interpreter.major,
+                python_interpreter.minor,
+            );
+
+            let pgo_ctx = PgoContext::new(pgo_command.to_owned())?;
+
+            // Phase 1: Build instrumented wheel for this interpreter
+            eprintln!("  📊 Phase 1/3: Building instrumented wheel...");
+            let mut instrumented_ctx = self.clone_for_pgo(PgoPhase::Generate(
+                pgo_ctx.profdata_dir_path().to_path_buf(),
+            ));
+            let instrumented_out = tempfile::TempDir::new()
+                .context("Failed to create temp dir for instrumented wheel")?;
+            instrumented_ctx.out = instrumented_out.path().to_path_buf();
+            let (instrumented_wheel_path, _) =
+                instrumented_ctx.build_single_pyo3_wheel(python_interpreter, &sbom_data)?;
+
+            // Phase 2: Run instrumentation with this interpreter
+            eprintln!("  🔬 Phase 2/3: Running PGO instrumentation...");
+            pgo_ctx.run_instrumentation(
+                &python_interpreter.executable,
+                &instrumented_wheel_path,
+                self,
+            )?;
+            pgo_ctx.merge_profiles()?;
+
+            // Phase 3: Build optimized wheel for this interpreter
+            eprintln!("  ⚡ Phase 3/3: Building PGO-optimized wheel...");
+            let optimized_ctx =
+                self.clone_for_pgo(PgoPhase::Use(pgo_ctx.merged_profdata_path().to_path_buf()));
+            let (wheel_path, tag) =
+                optimized_ctx.build_single_pyo3_wheel(python_interpreter, &sbom_data)?;
+
+            eprintln!(
+                "  📦 Built PGO-optimized wheel for {} {}.{}{} to {}",
+                python_interpreter.interpreter_kind,
+                python_interpreter.major,
+                python_interpreter.minor,
+                python_interpreter.abiflags,
+                wheel_path.display()
+            );
+            wheels.push((wheel_path, tag));
+        }
+
+        // Validate wheel filenames against PyPI platform tag rules if requested
+        if self.pypi_validation {
+            for (wheel_path, _) in &wheels {
+                let filename = wheel_path
+                    .file_name()
+                    .and_then(|name| name.to_str())
+                    .ok_or_else(|| anyhow!("Invalid wheel filename: {:?}", wheel_path))?;
+
+                if let Err(error) = validate_wheel_filename_for_pypi(filename) {
+                    bail!("PyPI validation failed: {}", error);
+                }
+            }
+        }
+
+        eprintln!("🎉 PGO build complete!");
+        Ok(wheels)
+    }
+
     /// Bridge model
     pub fn bridge(&self) -> &BridgeModel {
         // FIXME: currently we only allow multiple bin targets so bridges are all the same

src/build_context/wheels.rs

@@ -214,6 +214,34 @@ impl BuildContext {
         )
     }
 
+    /// Compile, audit, and write a single PyO3 wheel for one interpreter.
+    ///
+    /// This is the core pipeline shared by [`build_pyo3_wheels`] and the
+    /// per-interpreter PGO path in [`build_wheels_pgo_per_interpreter`].
+    pub(super) fn build_single_pyo3_wheel(
+        &self,
+        python_interpreter: &PythonInterpreter,
+        sbom_data: &Option<SbomData>,
+    ) -> Result<BuiltWheelMetadata> {
+        let (artifact, out_dirs) = self.compile_cdylib(
+            Some(python_interpreter),
+            Some(&self.project_layout.extension_name),
+        )?;
+        let (policy, external_libs) =
+            self.auditwheel(&artifact, &self.platform_tag, Some(python_interpreter))?;
+        let platform_tags = self.resolve_platform_tags(&policy);
+        let wheel_path = self.write_pyo3_wheel(
+            python_interpreter,
+            artifact,
+            &platform_tags,
+            external_libs,
+            sbom_data,
+            &out_dirs,
+        )?;
+        let tag = format!("cp{}{}", python_interpreter.major, python_interpreter.minor);
+        Ok((wheel_path, tag))
+    }
+
     /// Builds wheels for a pyo3 extension for all given python versions.
     /// Return type is the same as [BuildContext::build_wheels()]
     ///
@@ -228,21 +256,7 @@ impl BuildContext {
     ) -> Result<Vec<BuiltWheelMetadata>> {
         let mut wheels = Vec::new();
         for python_interpreter in interpreters {
-            let (artifact, out_dirs) = self.compile_cdylib(
-                Some(python_interpreter),
-                Some(&self.project_layout.extension_name),
-            )?;
-            let (policy, external_libs) =
-                self.auditwheel(&artifact, &self.platform_tag, Some(python_interpreter))?;
-            let platform_tags = self.resolve_platform_tags(&policy);
-            let wheel_path = self.write_pyo3_wheel(
-                python_interpreter,
-                artifact,
-                &platform_tags,
-                external_libs,
-                sbom_data,
-                &out_dirs,
-            )?;
+            let (wheel_path, tag) = self.build_single_pyo3_wheel(python_interpreter, sbom_data)?;
             eprintln!(
                 "📦 Built wheel for {} {}.{}{} to {}",
                 python_interpreter.interpreter_kind,
@@ -251,8 +265,6 @@ impl BuildContext {
                 python_interpreter.abiflags,
                 wheel_path.display()
             );
-
-            let tag = format!("cp{}{}", python_interpreter.major, python_interpreter.minor);
             wheels.push((wheel_path, tag));
         }
 

src/compile.rs

@@ -249,6 +249,22 @@ fn cargo_build_command(
         .unwrap_or_default();
     let original_rustflags = rustflags.flags.clone();
 
+    // Inject PGO flags if a PGO build phase is active
+    if let Some(ref pgo_phase) = context.pgo_phase {
+        match pgo_phase {
+            crate::pgo::PgoPhase::Generate(profdata_dir) => {
+                rustflags
+                    .flags
+                    .push(format!("-Cprofile-generate={}", profdata_dir.display()));
+            }
+            crate::pgo::PgoPhase::Use(profdata_path) => {
+                rustflags
+                    .flags
+                    .push(format!("-Cprofile-use={}", profdata_path.display()));
+            }
+        }
+    }
+
     let bridge_model = &compile_target.bridge_model;
     configure_bin_lib_flags(
         &mut cargo_rustc,

src/develop/mod.rs

@@ -1,4 +1,4 @@
-mod install_backend;
+pub(crate) mod install_backend;
 
 use crate::BuildContext;
 use crate::BuildOptions;

src/lib.rs

@@ -59,12 +59,14 @@ pub mod ci;
 mod compile;
 mod compression;
 mod cross_compile;
-mod develop;
+pub(crate) mod develop;
 mod generate_json_schema;
 mod metadata;
 mod module_writer;
 #[cfg(feature = "scaffolding")]
 mod new_project;
+/// Profile-Guided Optimization (PGO) orchestration
+pub(crate) mod pgo;
 mod project_layout;
 pub mod pyproject_toml;
 mod python_interpreter;