feat: add wasm mode for profiler

AztecProtocol · Mar 3, 2025 · 629b3c5 · 629b3c5
1 parent 443d615
commit 629b3c5
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 34 deletions.
diff --git a/yarn-project/bb-prover/src/wasm/bb_wasm_private_kernel_prover.ts b/yarn-project/bb-prover/src/wasm/bb_wasm_private_kernel_prover.ts
@@ -39,4 +39,17 @@ export abstract class BBWASMPrivateKernelProver extends BBPrivateKernelProver {
     });
     return new ClientIvcProof(Buffer.from(proof), Buffer.from(vk));
   }
+
+  public override async computeGateCountForCircuit(_bytecode: Buffer, _circuitName: string): Promise<number> {
+    const backend = new AztecClientBackend([ungzip(_bytecode)], {
+      threads: this.threads,
+      logger: this.log.verbose,
+      wasmPath: process.env.BB_WASM_PATH,
+    });
+
+    const gateCount = await backend.gates();
+    await backend.destroy();
+
+    return gateCount[0];
+  }
 }
diff --git a/yarn-project/pxe/src/kernel_prover/kernel_prover.ts b/yarn-project/pxe/src/kernel_prover/kernel_prover.ts
@@ -91,7 +91,6 @@ export type ProvingConfig = {
   simulate: boolean;
   skipFeeEnforcement: boolean;
   profile: boolean;
-  dryRun: boolean;
 };
 
 /**
@@ -118,25 +117,20 @@ export class KernelProver {
    * @param txRequest - The authenticated transaction request object.
    * @param executionResult - The execution result object containing nested executions and preimages.
    * @param profile - Set true to profile the gate count for each circuit
-   * @param dryRun - Set true to skip the IVC proof generation (only simulation is run). Useful for profiling gate count without proof gen.
    * @returns A Promise that resolves to a KernelProverOutput object containing proof, public inputs, and output notes.
    * TODO(#7368) this should be refactored to not recreate the ACIR bytecode now that it operates on a program stack
    */
   async prove(
     txRequest: TxRequest,
     executionResult: PrivateExecutionResult,
-    { simulate, skipFeeEnforcement, profile, dryRun }: ProvingConfig = {
+    { simulate, skipFeeEnforcement, profile }: ProvingConfig = {
       simulate: false,
       skipFeeEnforcement: false,
       profile: false,
-      dryRun: false,
     },
   ): Promise<PrivateKernelSimulateOutput<PrivateKernelTailCircuitPublicInputs>> {
-    if (simulate && profile) {
-      throw new Error('Cannot simulate and profile at the same time');
-    }
-
-    simulate = simulate || this.fakeProofs;
+    const skipWitnessGeneration = simulate && !profile;
+    const skipProofGeneration = this.fakeProofs || simulate;
 
     const timer = new Timer();
 
@@ -152,7 +146,7 @@ export class KernelProver {
       const gateCount = (await this.proofCreator.computeGateCountForCircuit(bytecode, circuitName)) as number;
       gateCounts.push({ circuitName, gateCount });
 
-      this.log.info(`Tx ${txRequest.hash()}: bb gates for ${circuitName} - ${gateCount}`);
+      this.log.debug(`Gate count for ${circuitName} - ${gateCount}`);
     };
 
     const noteHashLeafIndexMap = collectNoteHashLeafIndexMap(executionResult);
@@ -175,9 +169,10 @@ export class KernelProver {
         );
         while (resetBuilder.needsReset()) {
           const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
-          output = simulate
-            ? await this.proofCreator.simulateReset(privateInputs)
-            : await this.proofCreator.generateResetOutput(privateInputs);
+          output =
+            skipWitnessGeneration
+              ? await this.proofCreator.simulateReset(privateInputs)
+              : await this.proofCreator.generateResetOutput(privateInputs);
           // TODO(#7368) consider refactoring this redundant bytecode pushing
           acirs.push(output.bytecode);
           witnessStack.push(output.outputWitness);
@@ -228,9 +223,10 @@ export class KernelProver {
 
         pushTestData('private-kernel-inputs-init', proofInput);
 
-        output = simulate
-          ? await this.proofCreator.simulateInit(proofInput)
-          : await this.proofCreator.generateInitOutput(proofInput);
+        output =
+          skipWitnessGeneration
+            ? await this.proofCreator.simulateInit(proofInput)
+            : await this.proofCreator.generateInitOutput(proofInput);
 
         acirs.push(output.bytecode);
         witnessStack.push(output.outputWitness);
@@ -249,9 +245,10 @@ export class KernelProver {
 
         pushTestData('private-kernel-inputs-inner', proofInput);
 
-        output = simulate
-          ? await this.proofCreator.simulateInner(proofInput)
-          : await this.proofCreator.generateInnerOutput(proofInput);
+        output =
+          skipWitnessGeneration
+            ? await this.proofCreator.simulateInner(proofInput)
+            : await this.proofCreator.generateInnerOutput(proofInput);
 
         acirs.push(output.bytecode);
         witnessStack.push(output.outputWitness);
@@ -271,9 +268,10 @@ export class KernelProver {
     );
     while (resetBuilder.needsReset()) {
       const privateInputs = await resetBuilder.build(this.oracle, noteHashLeafIndexMap);
-      output = simulate
-        ? await this.proofCreator.simulateReset(privateInputs)
-        : await this.proofCreator.generateResetOutput(privateInputs);
+      output =
+        skipWitnessGeneration
+          ? await this.proofCreator.simulateReset(privateInputs)
+          : await this.proofCreator.generateResetOutput(privateInputs);
 
       acirs.push(output.bytecode);
       witnessStack.push(output.outputWitness);
@@ -290,7 +288,7 @@ export class KernelProver {
     }
 
     if (output.publicInputs.feePayer.isZero() && skipFeeEnforcement) {
-      if (!dryRun && !simulate) {
+      if (!skipProofGeneration) {
         throw new Error('Fee payment must be enforced when creating real proof.');
       }
       output.publicInputs.feePayer = new AztecAddress(Fr.MAX_FIELD_VALUE);
@@ -312,9 +310,10 @@ export class KernelProver {
 
     pushTestData('private-kernel-inputs-ordering', privateInputs);
 
-    const tailOutput = simulate
-      ? await this.proofCreator.simulateTail(privateInputs)
-      : await this.proofCreator.generateTailOutput(privateInputs);
+    const tailOutput =
+      skipWitnessGeneration
+        ? await this.proofCreator.simulateTail(privateInputs)
+        : await this.proofCreator.generateTailOutput(privateInputs);
     if (tailOutput.publicInputs.forPublic) {
       const privateLogs = privateInputs.previousKernel.publicInputs.end.privateLogs;
       const nonRevertiblePrivateLogs = tailOutput.publicInputs.forPublic.nonRevertibleAccumulatedData.privateLogs;
@@ -329,12 +328,12 @@ export class KernelProver {
       tailOutput.profileResult = { gateCounts };
     }
 
-    if (!simulate) {
+    if (!skipWitnessGeneration) {
       this.log.info(`Private kernel witness generation took ${timer.ms()}ms`);
     }
 
     // TODO(#7368) how do we 'bincode' encode these inputs?
-    if (!dryRun && !simulate) {
+    if (!skipProofGeneration) {
       const ivcProof = await this.proofCreator.createClientIvcProof(acirs, witnessStack);
       tailOutput.clientIvcProof = ivcProof;
     } else {

diff --git a/yarn-project/pxe/src/pxe_service/pxe_service.ts b/yarn-project/pxe/src/pxe_service/pxe_service.ts
@@ -390,7 +390,6 @@ export class PXEService implements PXE {
         simulate: false,
         skipFeeEnforcement: false,
         profile: false,
-        dryRun: false,
       });
       return new TxProvingResult(privateExecutionResult, publicInputs, clientIvcProof!);
     } catch (err: any) {
@@ -427,10 +426,9 @@ export class PXEService implements PXE {
       const privateExecutionResult = await this.#executePrivate(txRequest, msgSender, scopes);
 
       const { publicInputs, profileResult } = await this.#prove(txRequest, this.proofCreator, privateExecutionResult, {
-        simulate: !profile,
+        simulate: true,
         skipFeeEnforcement,
         profile,
-        dryRun: true,
       });
 
       const privateSimulationResult = new PrivateSimulationResult(privateExecutionResult, publicInputs);
@@ -727,19 +725,18 @@ export class PXEService implements PXE {
     txExecutionRequest: TxExecutionRequest,
     proofCreator: PrivateKernelProver,
     privateExecutionResult: PrivateExecutionResult,
-    { simulate, skipFeeEnforcement, profile, dryRun }: ProvingConfig,
+    { simulate, skipFeeEnforcement, profile }: ProvingConfig,
   ): Promise<PrivateKernelSimulateOutput<PrivateKernelTailCircuitPublicInputs>> {
     // use the block the tx was simulated against
     const block =
       privateExecutionResult.entrypoint.publicInputs.historicalHeader.globalVariables.blockNumber.toNumber();
     const kernelOracle = new KernelOracle(this.contractDataOracle, this.keyStore, this.node, block);
     const kernelProver = new KernelProver(kernelOracle, proofCreator, !this.proverEnabled);
-    this.log.debug(`Executing kernel prover (simulate: ${simulate}, profile: ${profile}, dryRun: ${dryRun})...`);
+    this.log.debug(`Executing kernel prover (simulate: ${simulate}, profile: ${profile})...`);
     return await kernelProver.prove(txExecutionRequest.toTxRequest(), privateExecutionResult, {
       simulate,
       skipFeeEnforcement,
       profile,
-      dryRun,
     });
   }