fix: flamegraph script (and enable > 1 circuit)

noir-projects/noir-protocol-circuits/scripts/flamegraph.sh

@@ -1,34 +1,85 @@
 #!/usr/bin/env bash
 set -eu
 
-EXAMPLE_CMD="$0 private_kernel_init"
+EXAMPLE_CMD="$0 private_kernel_init rollup_merge"
 
-# First arg is the circuit name.
-if [[ $# -eq 0 || ($1 == -* && $1 != "-h") ]]; then
-    echo "Please specify the name of the circuit."
-    echo "e.g.: $EXAMPLE_CMD"
-    exit 1
-fi
-
-CIRCUIT_NAME=$1
+# Parse global options.
+CIRCUIT_NAMES=()
 SERVE=false
 PORT=5000
+ALLOW_NO_CIRCUIT_NAMES=false
+
+# Get the directory of the script.
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+# and of the artifact
+ARTIFACT_DIR="$SCRIPT_DIR/../target"
+
+# Function to get filenames from a directory
+get_filenames() {
+    local dir="$1"
+    # Return filenames (without extensions) from the directory
+    for file in "$dir"/*; do
+        if [[ -f "$file" ]]; then
+            filename="$(basename "$file" .${file##*.})"
+            echo "$filename"
+        fi
+    done
+}
+
+NAUGHTY_LIST=("empty_nested") # files with no opcodes, which break the flamegraph tool.
+
+get_valid_circuit_names() {
+    # Capture the output of function call in an array:
+    ALL_CIRCUIT_NAMES=($(get_filenames "$ARTIFACT_DIR"))
+    for circuit_name in "${ALL_CIRCUIT_NAMES[@]}"; do
+        # Skip files that include the substring "simulated"
+        if [[ "$circuit_name" == *"simulated"* ]]; then
+            continue
+        fi
+        # Skip the file if it's on the naughty list:
+        if [[ " ${NAUGHTY_LIST[@]} " =~ " ${circuit_name} " ]]; then
+            continue
+        fi
+        CIRCUIT_NAMES+=("$circuit_name")
+    done
+}
+
 while [[ $# -gt 0 ]]; do
     case $1 in
         -h|--help)
-            echo "Generates a flamegraph for the specified protocol circuit."
+            echo "Generates flamegraphs for the specified protocol circuits."
             echo ""
             echo "Usage:"
-            echo "    $0 <CIRCUIT_NAME>"
+            echo "    $0 <CIRCUIT_NAME> [<CIRCUIT_NAME> ...] [options]"
             echo ""
-            echo "    e.g.: $EXAMPLE_CMD"
+            echo "    e.g.: $EXAMPLE_CMD -s -p 8080"
             echo ""
-            echo "Arguments:"
-            echo "    -s    Serve the file over http"
+            echo "Options:"
+            echo "    -s    Serve the file(s) over http"
             echo "    -p    Specify custom port. Default: ${PORT}"
             echo ""
+            echo "If you're feeling lazy, you can also just list available (compiled) circuit names with:"
+            echo "    $0 -l"
+            exit 0
+            ;;
+        -l|--list)
+            echo "Available circuits (that have been compiled):"
+            get_valid_circuit_names
+            for circuit_name in "${CIRCUIT_NAMES[@]}"; do
+                echo "$circuit_name"
+            done
             exit 0
             ;;
+        -a|--all)
+            echo "This will probably take a while..."
+            get_valid_circuit_names
+            shift
+            ;;
+        -n|--allow-no-circuit-names)
+            # Enables the existing flamegraphs to be served quickly.
+            ALLOW_NO_CIRCUIT_NAMES=true
+            shift
+            ;;
         -s|--serve)
             SERVE=true
             shift
@@ -43,22 +94,23 @@ while [[ $# -gt 0 ]]; do
             shift 2
             ;;
         *)
+            # Treat any argument not matching an option as a CIRCUIT_NAME.
+            CIRCUIT_NAMES+=("$1")
             shift
-        ;;
+            ;;
     esac
 done
 
-# Get the directory of the script.
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-
-# Check if the artifact exists.
-ARTIFACT="$SCRIPT_DIR/../target/$CIRCUIT_NAME.json"
-if [[ ! -f $ARTIFACT ]]; then
-    echo "Cannot find artifact: ${ARTIFACT}"
-    exit 1
+# Ensure at least one CIRCUIT_NAME was specified.
+if [[ ! $ALLOW_NO_CIRCUIT_NAMES ]]; then
+    if [[ ${#CIRCUIT_NAMES[@]} -eq 0 ]]; then
+        echo "Please specify at least one circuit name."
+        echo "e.g.: $EXAMPLE_CMD"
+        exit 1
+    fi
 fi
 
-# Build profier if it's not available.
+# Build profiler if it's not available.
 PROFILER="$SCRIPT_DIR/../../../noir/noir-repo/target/release/noir-profiler"
 if [ ! -f $PROFILER ]; then
     echo "Profiler not found, building profiler"
@@ -67,33 +119,49 @@ if [ ! -f $PROFILER ]; then
     cd "$SCRIPT_DIR"
 fi
 
-# We create dest directory and use it as an output for the generated main.svg file.
+# Create the output directory.
 DEST="$SCRIPT_DIR/../dest"
 mkdir -p $DEST
 
 MEGA_HONK_CIRCUIT_PATTERNS=$(jq -r '.[]' "$SCRIPT_DIR/../../mega_honk_circuits.json")
 
-# Check if the target circuit is a mega honk circuit.
-ARTIFACT_FILE_NAME=$(basename -s .json "$ARTIFACT")
+# Process each CIRCUIT_NAME.
+for CIRCUIT_NAME in "${CIRCUIT_NAMES[@]}"; do
+    (
+        echo ""
+        echo "Doing $CIRCUIT_NAME..."
+        # Check if the artifact exists.
+        ARTIFACT="$ARTIFACT_DIR/$CIRCUIT_NAME.json"
+        if [[ ! -f $ARTIFACT ]]; then
+            artifact_error="Cannot find artifact: ${ARTIFACT}"
+            echo "$artifact_error"
+        fi
 
-IS_MEGA_HONK_CIRCUIT="false"
-for pattern in $MEGA_HONK_CIRCUIT_PATTERNS; do
-    if echo "$ARTIFACT_FILE_NAME" | grep -qE "$pattern"; then
-        IS_MEGA_HONK_CIRCUIT="true"
-        break
-    fi
-done
+        ARTIFACT_FILE_NAME=$(basename -s .json "$ARTIFACT")
 
-# At last, generate the flamegraph.
-# If it's a mega honk circuit, we need to set the backend_gates_command argument to "gates_mega_honk".
-if [ "$IS_MEGA_HONK_CIRCUIT" = "true" ]; then
-    $PROFILER gates-flamegraph --artifact-path "${ARTIFACT}" --backend-path "$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb"  --output "$DEST" --backend-gates-command "gates_mega_honk" -- -h
-else
-    $PROFILER gates-flamegraph --artifact-path "${ARTIFACT}" --backend-path "$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb"  --output "$DEST" -- -h
-fi
+        # Determine if the circuit is a mega honk circuit.
+        IS_MEGA_HONK_CIRCUIT="false"
+        for pattern in $MEGA_HONK_CIRCUIT_PATTERNS; do
+            if echo "$ARTIFACT_FILE_NAME" | grep -qE "$pattern"; then
+                IS_MEGA_HONK_CIRCUIT="true"
+                break
+            fi
+        done
+
+        # Generate the flamegraph.
+        if [ "$IS_MEGA_HONK_CIRCUIT" = "true" ]; then
+            $PROFILER gates --artifact-path "${ARTIFACT}" --backend-path "$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb" --output "$DEST" --output-filename "$CIRCUIT_NAME" --backend-gates-command "gates_mega_honk" -- -h
+        else
+            $PROFILER gates --artifact-path "${ARTIFACT}" --backend-path "$SCRIPT_DIR/../../../barretenberg/cpp/build/bin/bb" --output "$DEST" --output-filename "$CIRCUIT_NAME" -- -h
+        fi
 
-# Serve the file over http if -s is set.
+        echo "Flamegraph generated for circuit: $CIRCUIT_NAME"
+    ) & # These parenthesis `( stuff ) &` mean "do all this in parallel"
+done
+wait # wait for parallel processes to finish
+
+# Serve the files over HTTP if -s is set.
 if $SERVE; then
-    echo "Serving flamegraph at http://0.0.0.0:${PORT}/main.svg"
-    python3 -m http.server --directory "$SCRIPT_DIR/../dest" $PORT
-fi
+    echo "Serving flamegraphs at http://0.0.0.0:${PORT}/"
+    python3 -m http.server --directory "$DEST" $PORT
+fi

noir/noir-repo/tooling/profiler/src/cli/gates_flamegraph_cmd.rs

@@ -31,6 +31,10 @@ pub(crate) struct GatesFlamegraphCommand {
     /// The output folder for the flamegraph svg files
     #[clap(long, short)]
     output: String,
+
+    /// The output name for the flamegraph svg files
+    #[clap(long, short = 'f', default_value = "main")]
+    output_filename: String
 }
 
 pub(crate) fn run(args: GatesFlamegraphCommand) -> eyre::Result<()> {
@@ -43,6 +47,7 @@ pub(crate) fn run(args: GatesFlamegraphCommand) -> eyre::Result<()> {
         },
         &InfernoFlamegraphGenerator { count_name: "gates".to_string() },
         &PathBuf::from(args.output),
+        &args.output_filename
     )
 }
 
@@ -51,6 +56,7 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
     gates_provider: &Provider,
     flamegraph_generator: &Generator,
     output_path: &Path,
+    output_filename: &String
 ) -> eyre::Result<()> {
     let mut program =
         read_program_from_file(artifact_path).context("Error reading program from file")?;
@@ -97,7 +103,7 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
             &debug_artifact,
             artifact_path.to_str().unwrap(),
             &func_name,
-            &Path::new(&output_path).join(Path::new(&format!("{}_gates.svg", &func_name))),
+            &Path::new(&output_path).join(Path::new(&format!("{}_gates.svg", output_filename))),
         )?;
     }
 
@@ -192,7 +198,8 @@ mod tests {
         };
         let flamegraph_generator = TestFlamegraphGenerator::default();
 
-        super::run_with_provider(&artifact_path, &provider, &flamegraph_generator, temp_dir.path())
+        let null_str: Option<String> = None;
+        super::run_with_provider(&artifact_path, &provider, &flamegraph_generator, temp_dir.path(), &null_str)
             .expect("should run without errors");
 
         // Check that the output file was written to