Skip to content

Commit

Permalink
feat: add wasm mode for profiler (#12407)
Browse files Browse the repository at this point in the history 
+fix the profiler
  • Loading branch information
@saleel
saleel authored Mar 3, 2025
1 parent 3bdb886 commit 035bdb6
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 33 deletions.
13 changes: 13 additions & 0 deletions yarn-project/bb-prover/src/wasm/bb_wasm_private_kernel_prover.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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];
}
}
50 changes: 22 additions & 28 deletions yarn-project/pxe/src/kernel_prover/kernel_prover.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ export type ProvingConfig = {
simulate: boolean;
skipFeeEnforcement: boolean;
profile: boolean;
dryRun: boolean;
};

/**
Expand All @@ -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 skipProofGeneration = this.fakeProofs || simulate;
const generateWitnesses = !skipProofGeneration || profile;

const timer = new Timer();

Expand All @@ -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);
Expand All @@ -175,9 +169,9 @@ 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 = generateWitnesses
? await this.proofCreator.generateResetOutput(privateInputs)
: await this.proofCreator.simulateReset(privateInputs);
// TODO(#7368) consider refactoring this redundant bytecode pushing
acirs.push(output.bytecode);
witnessStack.push(output.outputWitness);
Expand Down Expand Up @@ -228,9 +222,9 @@ export class KernelProver {

pushTestData('private-kernel-inputs-init', proofInput);

output = simulate
? await this.proofCreator.simulateInit(proofInput)
: await this.proofCreator.generateInitOutput(proofInput);
output = generateWitnesses
? await this.proofCreator.generateInitOutput(proofInput)
: await this.proofCreator.simulateInit(proofInput);

acirs.push(output.bytecode);
witnessStack.push(output.outputWitness);
Expand All @@ -249,9 +243,9 @@ export class KernelProver {

pushTestData('private-kernel-inputs-inner', proofInput);

output = simulate
? await this.proofCreator.simulateInner(proofInput)
: await this.proofCreator.generateInnerOutput(proofInput);
output = generateWitnesses
? await this.proofCreator.generateInnerOutput(proofInput)
: await this.proofCreator.simulateInner(proofInput);

acirs.push(output.bytecode);
witnessStack.push(output.outputWitness);
Expand All @@ -271,9 +265,9 @@ 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 = generateWitnesses
? await this.proofCreator.generateResetOutput(privateInputs)
: await this.proofCreator.simulateReset(privateInputs);

acirs.push(output.bytecode);
witnessStack.push(output.outputWitness);
Expand All @@ -290,7 +284,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);
Expand All @@ -312,9 +306,9 @@ export class KernelProver {

pushTestData('private-kernel-inputs-ordering', privateInputs);

const tailOutput = simulate
? await this.proofCreator.simulateTail(privateInputs)
: await this.proofCreator.generateTailOutput(privateInputs);
const tailOutput = generateWitnesses
? await this.proofCreator.generateTailOutput(privateInputs)
: await this.proofCreator.simulateTail(privateInputs);
if (tailOutput.publicInputs.forPublic) {
const privateLogs = privateInputs.previousKernel.publicInputs.end.privateLogs;
const nonRevertiblePrivateLogs = tailOutput.publicInputs.forPublic.nonRevertibleAccumulatedData.privateLogs;
Expand All @@ -329,12 +323,12 @@ export class KernelProver {
tailOutput.profileResult = { gateCounts };
}

if (!simulate) {
if (generateWitnesses) {
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 {
Expand Down
7 changes: 2 additions & 5 deletions yarn-project/pxe/src/pxe_service/pxe_service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -430,7 +429,6 @@ export class PXEService implements PXE {
simulate: !profile,
skipFeeEnforcement,
profile,
dryRun: true,
});

const privateSimulationResult = new PrivateSimulationResult(privateExecutionResult, publicInputs);
Expand Down Expand Up @@ -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,
});
}

Expand Down

1 comment on commit 035bdb6

@AztecBot
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'C++ Benchmark'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.05.

Benchmark suite Current: 035bdb6 Previous: 3bdb886 Ratio
wasmconstruct_proof_ultrahonk_power_of_2/20 10419.563221 ms/iter 9793.011663 ms/iter 1.06

This comment was automatically generated by workflow using github-action-benchmark.

CC: @ludamad @codygunton

Please sign in to comment.