feat: add gate count profiling for transactions (#9632)

Part of #9299 

- Added `computeGateCountForCircuit` to `bb/execute.s` and use it in
`kernal_prover` to calculate gate-count per function (under a flag).
- Add `dryRun` to `kernal_prover` to skip ClientIVC proof generation
(useful for getting gate count without proof generation)
- Add `--profile` flag to cli-wallet `simulate` command to prints the
gate count per circuit, Example:
<img width="888" alt="Screenshot 2024-10-31 at 9 46 10 PM"
src="https://github.com/user-attachments/assets/4bbb062c-acdf-4889-aa6b-294c030b34d5">
<br/><br/>

- Not reusing `ContractFunctionInteraction.simulate` as it is difficult
to change its return type as its widely used.
- Can also add gas estimation to simulate command (instead of `send`) so
`simulate` does actual simulation of a transaction and display the full
data.

yarn-project/aztec.js/src/contract/contract_function_interaction.ts

@@ -1,4 +1,4 @@
-import type { FunctionCall, TxExecutionRequest } from '@aztec/circuit-types';
+import type { FunctionCall, PrivateKernelProverProfileResult, TxExecutionRequest } from '@aztec/circuit-types';
 import { type AztecAddress, type GasSettings } from '@aztec/circuits.js';
 import {
   type FunctionAbi,
@@ -27,6 +27,14 @@ export type SimulateMethodOptions = {
   skipTxValidation?: boolean;
 };
 
+/**
+ * The result of a profile() call.
+ */
+export type ProfileResult = PrivateKernelProverProfileResult & {
+  /** The result of the transaction as returned by the contract function. */
+  returnValues: any;
+};
+
 /**
  * This is the class that is returned when calling e.g. `contract.methods.myMethod(arg0, arg1)`.
  * It contains available interactions one can call on a method, including view.
@@ -110,4 +118,30 @@ export class ContractFunctionInteraction extends BaseContractInteraction {
 
     return rawReturnValues ? decodeFromAbi(this.functionDao.returnTypes, rawReturnValues) : [];
   }
+
+  /**
+   * Simulate a transaction and profile the gate count for each function in the transaction.
+   * @param options - Same options as `simulate`.
+   *
+   * @returns An object containing the function return value and profile result.
+   */
+  public async simulateWithProfile(options: SimulateMethodOptions = {}): Promise<ProfileResult> {
+    if (this.functionDao.functionType == FunctionType.UNCONSTRAINED) {
+      throw new Error("Can't profile an unconstrained function.");
+    }
+
+    const txRequest = await this.create();
+    const simulatedTx = await this.wallet.simulateTx(txRequest, true, options?.from, options?.skipTxValidation, true);
+
+    const rawReturnValues =
+      this.functionDao.functionType == FunctionType.PRIVATE
+        ? simulatedTx.getPrivateReturnValues().nested?.[0].values
+        : simulatedTx.getPublicReturnValues()?.[0].values;
+    const rawReturnValuesDecoded = rawReturnValues ? decodeFromAbi(this.functionDao.returnTypes, rawReturnValues) : [];
+
+    return {
+      returnValues: rawReturnValuesDecoded,
+      gateCounts: simulatedTx.profileResult!.gateCounts,
+    };
+  }
 }

yarn-project/aztec.js/src/index.ts

@@ -34,6 +34,7 @@ export {
   type DeployOptions,
   type SendMethodOptions,
   type WaitOpts,
+  type ProfileResult,
 } from './contract/index.js';
 
 export { ContractDeployer } from './deployment/index.js';

yarn-project/aztec.js/src/wallet/base_wallet.ts

@@ -116,8 +116,9 @@ export abstract class BaseWallet implements Wallet {
     simulatePublic: boolean,
     msgSender?: AztecAddress,
     skipTxValidation?: boolean,
+    profile?: boolean,
   ): Promise<TxSimulationResult> {
-    return this.pxe.simulateTx(txRequest, simulatePublic, msgSender, skipTxValidation, this.scopes);
+    return this.pxe.simulateTx(txRequest, simulatePublic, msgSender, skipTxValidation, profile, this.scopes);
   }
   sendTx(tx: Tx): Promise<TxHash> {
     return this.pxe.sendTx(tx);

yarn-project/bb-prover/src/bb/execute.ts

@@ -36,6 +36,8 @@ export type BBSuccess = {
   proofPath?: string;
   /** Full path of the contract. */
   contractPath?: string;
+  /** The number of gates in the circuit. */
+  circuitSize?: number;
 };
 
 export type BBFailure = {
@@ -872,6 +874,80 @@ export async function generateContractForCircuit(
   );
 }
 
+/**
+ * Compute bb gate count for a given circuit
+ * @param pathToBB - The full path to the bb binary
+ * @param workingDirectory - A temporary directory for writing the bytecode
+ * @param circuitName - The name of the circuit
+ * @param bytecode - The bytecode of the circuit
+ * @param flavor - The flavor of the backend - mega_honk or ultra_honk variants
+ * @returns An object containing the status, gate count, and time taken
+ */
+export async function computeGateCountForCircuit(
+  pathToBB: string,
+  workingDirectory: string,
+  circuitName: string,
+  bytecode: Buffer,
+  flavor: UltraHonkFlavor | 'mega_honk',
+): Promise<BBFailure | BBSuccess> {
+  // Check that the working directory exists
+  try {
+    await fs.access(workingDirectory);
+  } catch (error) {
+    return { status: BB_RESULT.FAILURE, reason: `Working directory ${workingDirectory} does not exist` };
+  }
+
+  // The bytecode is written to e.g. /workingDirectory/BaseParityArtifact-bytecode
+  const bytecodePath = `${workingDirectory}/${circuitName}-bytecode`;
+
+  const binaryPresent = await fs
+    .access(pathToBB, fs.constants.R_OK)
+    .then(_ => true)
+    .catch(_ => false);
+  if (!binaryPresent) {
+    return { status: BB_RESULT.FAILURE, reason: `Failed to find bb binary at ${pathToBB}` };
+  }
+
+  // Accumulate the stdout from bb
+  let stdout = '';
+  const logHandler = (message: string) => {
+    stdout += message;
+  };
+
+  try {
+    // Write the bytecode to the working directory
+    await fs.writeFile(bytecodePath, bytecode);
+    const timer = new Timer();
+
+    const result = await executeBB(
+      pathToBB,
+      flavor === 'mega_honk' ? `gates_mega_honk` : `gates`,
+      ['-b', bytecodePath, '-v'],
+      logHandler,
+    );
+    const duration = timer.ms();
+
+    if (result.status == BB_RESULT.SUCCESS) {
+      // Look for "circuit_size" in the stdout and parse the number
+      const circuitSizeMatch = stdout.match(/circuit_size": (\d+)/);
+      if (!circuitSizeMatch) {
+        return { status: BB_RESULT.FAILURE, reason: 'Failed to parse circuit_size from bb gates stdout.' };
+      }
+      const circuitSize = parseInt(circuitSizeMatch[1]);
+
+      return {
+        status: BB_RESULT.SUCCESS,
+        durationMs: duration,
+        circuitSize: circuitSize,
+      };
+    }
+
+    return { status: BB_RESULT.FAILURE, reason: 'Failed getting the gate count.' };
+  } catch (error) {
+    return { status: BB_RESULT.FAILURE, reason: `${error}` };
+  }
+}
+
 const CACHE_FILENAME = '.cache';
 async function fsCache<T>(
   dir: string,

yarn-project/bb-prover/src/prover/bb_private_kernel_prover.ts

@@ -52,6 +52,7 @@ import {
   BB_RESULT,
   PROOF_FIELDS_FILENAME,
   PROOF_FILENAME,
+  computeGateCountForCircuit,
   computeVerificationKey,
   executeBbClientIvcProof,
   verifyProof,
@@ -228,6 +229,21 @@ export class BBNativePrivateKernelProver implements PrivateKernelProver {
     this.log.info(`Successfully verified ${circuitType} proof in ${Math.ceil(result.durationMs)} ms`);
   }
 
+  public async computeGateCountForCircuit(bytecode: Buffer, circuitName: string): Promise<number> {
+    const result = await computeGateCountForCircuit(
+      this.bbBinaryPath,
+      this.bbWorkingDirectory,
+      circuitName,
+      bytecode,
+      'mega_honk',
+    );
+    if (result.status === BB_RESULT.FAILURE) {
+      throw new Error(result.reason);
+    }
+
+    return result.circuitSize as number;
+  }
+
   private async verifyProofFromKey(
     flavor: UltraHonkFlavor,
     verificationKey: Buffer,

yarn-project/circuit-types/src/interfaces/private_kernel_prover.ts

@@ -11,6 +11,10 @@ import {
 
 import { type WitnessMap } from '@noir-lang/acvm_js';
 
+export type PrivateKernelProverProfileResult = {
+  gateCounts: { circuitName: string; gateCount: number }[];
+};
+
 /**
  * Represents the output of the proof creation process for init and inner private kernel circuit.
  * Contains the public inputs required for the init and inner private kernel circuit and the generated proof.
@@ -28,6 +32,8 @@ export type PrivateKernelSimulateOutput<PublicInputsType> = {
   outputWitness: WitnessMap;
 
   bytecode: Buffer;
+
+  profileResult?: PrivateKernelProverProfileResult;
 };
 
 /**
@@ -97,4 +103,12 @@ export interface PrivateKernelProver {
    * @returns A Promise resolving to a Proof object
    */
   computeAppCircuitVerificationKey(bytecode: Buffer, appCircuitName?: string): Promise<AppCircuitSimulateOutput>;
+
+  /**
+   * Compute the gate count for a given circuit.
+   * @param bytecode - The circuit bytecode in gzipped bincode format
+   * @param circuitName - The name of the circuit
+   * @returns A Promise resolving to the gate count
+   */
+  computeGateCountForCircuit(bytecode: Buffer, circuitName: string): Promise<number>;
 }

yarn-project/circuit-types/src/interfaces/pxe.ts

@@ -161,6 +161,7 @@ export interface PXE {
    * @param simulatePublic - Whether to simulate the public part of the transaction.
    * @param msgSender - (Optional) The message sender to use for the simulation.
    * @param skipTxValidation - (Optional) If false, this function throws if the transaction is unable to be included in a block at the current state.
+   * @param profile - (Optional) If true, will run the private kernel prover with profiling enabled and include the result (gate count) in TxSimulationResult.
    * @param scopes - (Optional) The accounts whose notes we can access in this call. Currently optional and will default to all.
    * @returns A simulated transaction result object that includes public and private return values.
    * @throws If the code for the functions executed in this transaction has not been made available via `addContracts`.
@@ -171,6 +172,7 @@ export interface PXE {
     simulatePublic: boolean,
     msgSender?: AztecAddress,
     skipTxValidation?: boolean,
+    profile?: boolean,
     scopes?: AztecAddress[],
   ): Promise<TxSimulationResult>;
 

yarn-project/circuit-types/src/tx/simulated_tx.ts

@@ -1,6 +1,11 @@
 import { ClientIvcProof, PrivateKernelTailCircuitPublicInputs } from '@aztec/circuits.js';
 
-import { EncryptedNoteTxL2Logs, EncryptedTxL2Logs, UnencryptedTxL2Logs } from '../index.js';
+import {
+  EncryptedNoteTxL2Logs,
+  EncryptedTxL2Logs,
+  type PrivateKernelProverProfileResult,
+  UnencryptedTxL2Logs,
+} from '../index.js';
 import {
   PrivateExecutionResult,
   collectEnqueuedPublicFunctionCalls,
@@ -60,6 +65,7 @@ export class TxSimulationResult extends PrivateSimulationResult {
     privateExecutionResult: PrivateExecutionResult,
     publicInputs: PrivateKernelTailCircuitPublicInputs,
     public publicOutput?: PublicSimulationOutput,
+    public profileResult?: PrivateKernelProverProfileResult,
   ) {
     super(privateExecutionResult, publicInputs);
   }
@@ -71,11 +77,13 @@ export class TxSimulationResult extends PrivateSimulationResult {
   static fromPrivateSimulationResultAndPublicOutput(
     privateSimulationResult: PrivateSimulationResult,
     publicOutput?: PublicSimulationOutput,
+    profileResult?: PrivateKernelProverProfileResult,
   ) {
     return new TxSimulationResult(
       privateSimulationResult.privateExecutionResult,
       privateSimulationResult.publicInputs,
       publicOutput,
+      profileResult,
     );
   }
 
@@ -84,14 +92,16 @@ export class TxSimulationResult extends PrivateSimulationResult {
       privateExecutionResult: this.privateExecutionResult.toJSON(),
       publicInputs: this.publicInputs.toBuffer().toString('hex'),
       publicOutput: this.publicOutput ? this.publicOutput.toJSON() : undefined,
+      profileResult: this.profileResult,
     };
   }
 
   public static override fromJSON(obj: any) {
     const privateExecutionResult = PrivateExecutionResult.fromJSON(obj.privateExecutionResult);
     const publicInputs = PrivateKernelTailCircuitPublicInputs.fromBuffer(Buffer.from(obj.publicInputs, 'hex'));
     const publicOuput = obj.publicOutput ? PublicSimulationOutput.fromJSON(obj.publicOutput) : undefined;
-    return new TxSimulationResult(privateExecutionResult, publicInputs, publicOuput);
+    const profileResult = obj.profileResult;
+    return new TxSimulationResult(privateExecutionResult, publicInputs, publicOuput, profileResult);
   }
 }
 

yarn-project/cli-wallet/src/cmds/index.ts

@@ -34,6 +34,7 @@ import {
   createArgsOption,
   createArtifactOption,
   createContractAddressOption,
+  createProfileOption,
   createTypeOption,
   integerArgParser,
   parsePaymentMethod,
@@ -287,6 +288,7 @@ export function injectCommands(program: Command, log: LogFn, debugLogger: DebugL
     )
     .addOption(createAccountOption('Alias or address of the account to simulate from', !db, db))
     .addOption(createTypeOption(false))
+    .addOption(createProfileOption())
     .action(async (functionName, _options, command) => {
       const { simulate } = await import('./simulate.js');
       const options = command.optsWithGlobals();
@@ -299,13 +301,14 @@ export function injectCommands(program: Command, log: LogFn, debugLogger: DebugL
         type,
         secretKey,
         publicKey,
+        profile,
       } = options;
 
       const client = await createCompatibleClient(rpcUrl, debugLogger);
       const account = await createOrRetrieveAccount(client, parsedFromAddress, db, type, secretKey, Fr.ZERO, publicKey);
       const wallet = await getWalletWithScopes(account, db);
       const artifactPath = await artifactPathFromPromiseOrAlias(artifactPathPromise, contractAddress, db);
-      await simulate(wallet, functionName, args, artifactPath, contractAddress, log);
+      await simulate(wallet, functionName, args, artifactPath, contractAddress, profile, log);
     });
 
   program

yarn-project/cli-wallet/src/cmds/simulate.ts

@@ -1,22 +1,40 @@
-import { type AccountWalletWithSecretKey, type AztecAddress, Contract } from '@aztec/aztec.js';
+import { type AccountWalletWithSecretKey, type AztecAddress, Contract, type ProfileResult } from '@aztec/aztec.js';
 import { prepTx } from '@aztec/cli/utils';
 import { type LogFn } from '@aztec/foundation/log';
 
 import { format } from 'util';
 
+function printProfileResult(result: ProfileResult, log: LogFn) {
+  log(format('Simulation result:'));
+  log(format('Return value: ', JSON.stringify(result.returnValues, null, 2)));
+
+  log(format('Gate count: '));
+  let acc = 0;
+  result.gateCounts.forEach(r => {
+    acc += r.gateCount;
+    log(format('  ', r.circuitName.padEnd(30), 'Gates:', r.gateCount, '\tAcc:', acc));
+  });
+}
+
 export async function simulate(
   wallet: AccountWalletWithSecretKey,
   functionName: string,
   functionArgsIn: any[],
   contractArtifactPath: string,
   contractAddress: AztecAddress,
+  profile: boolean,
   log: LogFn,
 ) {
   const { functionArgs, contractArtifact } = await prepTx(contractArtifactPath, functionName, functionArgsIn, log);
 
   const contract = await Contract.at(contractAddress, contractArtifact, wallet);
   const call = contract.methods[functionName](...functionArgs);
 
-  const result = await call.simulate();
-  log(format('\nSimulation result: ', result, '\n'));
+  if (profile) {
+    const result = await call.simulateWithProfile();
+    printProfileResult(result, log);
+  } else {
+    const result = await call.simulate();
+    log(format('\nSimulation result: ', result, '\n'));
+  }
 }

yarn-project/cli-wallet/src/utils/options/index.ts

@@ -140,6 +140,13 @@ export function createArtifactOption(db?: WalletDB) {
     .makeOptionMandatory(false);
 }
 
+export function createProfileOption() {
+  return new Option(
+    '-p, --profile',
+    'Run the real prover and get the gate count for each function in the transaction.',
+  ).default(false);
+}
+
 async function contractArtifactFromWorkspace(pkg?: string, contractName?: string) {
   const cwd = process.cwd();
   try {