Add the option to generate performance flamegraphs

Cargo.toml

@@ -32,6 +32,12 @@ path = "src/bin/cargo_kani.rs"
 [profile.release]
 strip = "debuginfo"
 
+[profile.profiling]
+inherits = "release"
+debug = true
+strip = "none"
+lto = "off"
+
 # Below is the workspace (vs above is "kani-verifier" crate) config:
 
 [workspace]

docs/src/SUMMARY.md

@@ -44,6 +44,7 @@
     - [`benchcomp` command line](./benchcomp-cli.md)
     - [`benchcomp` configuration file](./benchcomp-conf.md)
     - [Custom parsers](./benchcomp-parse.md)
+  - [Profiling Kani](./profiling.md)
 
 - [Limitations](./limitations.md)
   - [Undefined behaviour](./undefined-behaviour.md)

docs/src/dev-documentation.md

@@ -15,6 +15,7 @@ developers (including external contributors):
  4. [Development setup recommendations for working with `rustc`](./rustc-hacks.md).
  5. [Guide for testing in Kani](./testing.md).
  6. [Transition to StableMIR](./stable-mir.md).
+ 7. [Profiling Kani's performance](./profiling.md)
 
 > **NOTE**: The developer documentation is intended for Kani developers and not
 users. At present, the project is under heavy development and some items

docs/src/profiling.md

@@ -0,0 +1,25 @@
+# Profiling Kani's Performance
+
+To profile Kani's performance at a fine-grained level, we use a tool called [`samply`](https://github.com/mstange/samply) that allows the compiler & driver to periodically record the current stack trace, allowing us to construct flamegraphs of where they are spending most of their time.
+
+## Install samply
+First, install `samply` using [the instructions](https://github.com/mstange/samply?tab=readme-ov-file#installation) from their repo. The easier methods include installing a prebuilt binary or installing from crates.io.
+
+
+## Running Kani for profiling output
+1. First, build Kani from source with `cargo build-dev --profile profiling` to ensure you are getting all release mode optimizations without stripping useful debug info.
+2. Then, you can profile the Kani compiler on a crate of your choice by running `FLAMEGRAPH=[OPTION] cargo kani` within the crate.
+
+The `FLAMEGRAPH` environment variable can be set to `driver` (to profile the complete `kani-driver` execution), `compiler` (to profile each time the `kani-compiler` is called) or `all` (to profile both).
+
+We have to instrument the driver and compiler separately because samply's instrumentation usually cannot handle detecting the subprocess the driver uses to call the compiler.
+
+> Note: Specifically when profiling the compiler, ensure you are running `cargo clean` immediately before `cargo kani`, or parts of the workspace may not be recompiled by the Kani compiler.
+
+
+## Displaying profiling output
+This will create a new `flamegraph` directory in the crate which will contain a single `driver.json.gz` output file and one `compiler-{crate_name}.json.gz` file for each crate in the workspace. Run `samply load flamegraphs/XXX.json.gz` on any of these to open a local server that will display the file's flamegraph.
+
+Once the server has opened, you'll see a display with the list of threads in rows at the top, and a flamegraph for the currently selected thread at the bottom. There is typically only one process when profiling the driver. When profiling the compiler, the process that runs the `kani-compiler` and handles all codegen is usually at the very bottom of the thread window.
+
+In the flamegraph view, I've found it very useful to right click on a function of interest and select "focus on subtree only" so that it zooms in and you can more clearly see the callees it uses. This can then be undone with the breadcrumb trail at the top of the flamegraph.

kani-driver/src/call_cargo.rs

@@ -6,6 +6,7 @@ use crate::call_single_file::{LibConfig, to_rustc_arg};
 use crate::project::Artifact;
 use crate::session::{
     KaniSession, get_cargo_path, lib_folder, lib_no_core_folder, setup_cargo_command,
+    setup_cargo_command_inner,
 };
 use crate::util;
 use anyhow::{Context, Result, bail};
@@ -204,7 +205,8 @@ crate-type = ["lib"]
         let mut failed_targets = vec![];
         for package in packages {
             for verification_target in package_targets(&self.args, package) {
-                let mut cmd = setup_cargo_command()?;
+                let mut cmd =
+                    setup_cargo_command_inner(Some(verification_target.target().name.clone()))?;
                 cmd.args(&cargo_args)
                     .args(vec!["-p", &package.id.to_string()])
                     .args(verification_target.to_args())

kani-driver/src/session.rs

@@ -418,18 +418,46 @@ fn init_logger(args: &VerificationArgs) {
     tracing::subscriber::set_global_default(subscriber).unwrap();
 }
 
-// Setup the default version of cargo being run, based on the type/mode of installation for kani
-// If kani is being run in developer mode, then we use the one provided by rustup as we can assume that the developer will have rustup installed
+pub fn setup_cargo_command() -> Result<Command> {
+    setup_cargo_command_inner(None)
+}
+
+// Setup the default version of cargo being run, based on the type/mode of installation for kani.
+// Optionally takes a path to output compiler profiling info to.
+// If kani is being run in developer mode, then we use the one provided by rustup as we can assume that the developer will have rustup installed.
 // For release versions of Kani, we use a version of cargo that's in the toolchain that's been symlinked during `cargo-kani` setup. This will allow
 // Kani to remove the runtime dependency on rustup later on.
-pub fn setup_cargo_command() -> Result<Command> {
+pub fn setup_cargo_command_inner(profiling_out_path: Option<String>) -> Result<Command> {
     let install_type = InstallType::new()?;
 
     let cmd = match install_type {
         InstallType::DevRepo(_) => {
-            let mut cmd = Command::new("cargo");
-            cmd.arg(self::toolchain_shorthand());
-            cmd
+            // check if we should instrument the compiler for a flamegraph
+            let instrument_compiler = matches!(
+                std::env::var("FLAMEGRAPH"),
+                Ok(ref s) if s == "all" || s == "compiler"
+            );
+            let flamegraph_dir = "flamegraphs";
+
+            if let Some(profiler_out_path) = profiling_out_path
+                && instrument_compiler
+            {
+                let _ = std::fs::create_dir(flamegraph_dir); // create temporary flamegraph directory
+
+                let mut cmd = Command::new("samply");
+                cmd.arg("record");
+                cmd.arg("-r").arg("8000"); // adjust the sampling rate (in Hz)
+                cmd.arg("-o")
+                    .arg(format!("{flamegraph_dir}/compiler-{profiler_out_path}.json.gz",)); // set the output file
+                cmd.arg("--save-only"); // just save the output and don't open the interactive UI.
+                cmd.arg("cargo");
+                cmd.arg(self::toolchain_shorthand());
+                cmd
+            } else {
+                let mut cmd = Command::new("cargo");
+                cmd.arg(self::toolchain_shorthand());
+                cmd
+            }
         }
         InstallType::Release(kani_dir) => {
             let cargo_path = kani_dir.join("toolchain").join("bin").join("cargo");

scripts/cargo-kani

@@ -18,4 +18,16 @@ then
     exit 1
 fi
 
-exec -a cargo-kani "${KANI_DRIVER}" "$@"
+FLAMEGRAPH_DRIVER=0
+FLAMEGRAPH_DIR="flamegraphs"
+
+if [[ ( -n "${FLAMEGRAPH:-}" ) && ( "${FLAMEGRAPH}" == "driver" || "${FLAMEGRAPH}" == "all" ) ]]; then
+    FLAMEGRAPH_DRIVER=1
+fi
+
+if (( FLAMEGRAPH_DRIVER == 0 )); then 
+    exec -a cargo-kani "${KANI_DRIVER}" "$@"
+else 
+    mkdir -p "${FLAMEGRAPH_DIR}"
+    exec -a cargo-kani samply record -r 8000 -o "${FLAMEGRAPH_DIR}/driver.json.gz" --save-only "${KANI_DRIVER}" "$@"
+fi