Skip to content

Commit

Permalink
Merge pull request #240 from mooori/inst-helper
Browse files Browse the repository at this point in the history
feat(bench): add `InstructionTracer.js` helper
  • Loading branch information
mooori authored Mar 8, 2024
2 parents c96b25f + af25607 commit 0c0d969
Show file tree
Hide file tree
Showing 6 changed files with 225 additions and 20 deletions.
5 changes: 4 additions & 1 deletion cranelift/codegen/src/isa/zkasm/inst/emit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -369,13 +369,16 @@ impl Inst {
}

/// Emits zkASM to trace executed instructions.
///
/// Calls a method implemented by the `InstructionTracer.js` helper. An instance of this helper
/// must be provided to `zkevm-proverjs` for the successful execution of the generated zkASM.
fn emit_inst_instrumentation(&self, sink: &mut MachBuffer<Inst>) {
match self {
// Labels are handled separately since benchmarking will provide a separate command to
// analyze labels.
&MInst::Label { .. } => {}
_ => put_string(
";; TODO(mooori) call the helper to trace `MInst` execution\n",
&format!("$${{traceInstruction({})}}\n", self.print_name()),
sink,
),
}
Expand Down
61 changes: 61 additions & 0 deletions cranelift/codegen/src/isa/zkasm/inst/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -752,6 +752,8 @@ pub fn reg_name(reg: Reg) -> String {
}

impl Inst {
// TODO(#218): remove Risc-V info from `print_with_state`
// TODO(#217): print consistent names by reusing `print_name`
fn print_with_state(
&self,
_state: &mut EmitState,
Expand Down Expand Up @@ -1258,6 +1260,65 @@ impl Inst {
&Inst::ECall {} => String::from("ecall"),
}
}

/// Prints the enum variant identifier.
fn print_name(&self) -> String {
let name = match self {
&Inst::Nop => "Nop",
&Inst::Label { .. } => "Label",
&Inst::LoadConst32 { .. } => "LoadConst32",
&Inst::LoadConst64 { .. } => "LoadConst64",
&Inst::AluRRR { .. } => "AluRRR",
&Inst::MulArith { .. } => "MulArith",
&Inst::Shl64 { .. } => "Shl64",
&Inst::Shru64 { .. } => "Shru64",
&Inst::DivArith { .. } => "DivArith",
&Inst::UDivArith { .. } => "UDivArith",
&Inst::RemArith { .. } => "RemArith",
&Inst::URemArith { .. } => "URemArith",
&Inst::Ineg { .. } => "Ineg",
&Inst::Bnot { .. } => "Bnot",
&Inst::Load { .. } => "Load",
&Inst::Store { .. } => "Store",
&Inst::Args { .. } => "Args",
&Inst::Ret { .. } => "Ret",
&Inst::Extend { .. } => "Extend",
&Inst::ReserveSp { .. } => "ReserveSp",
&Inst::ReleaseSp { .. } => "ReleaseSp",
&Inst::Call { .. } => "Call",
&Inst::CallInd { .. } => "CallInd",
&Inst::ReturnCall { .. } => "ReturnCall",
&Inst::ReturnCallInd { .. } => "ReturnCallInd",
&Inst::TrapIf { .. } => "TrapIf",
&Inst::TrapIfC { .. } => "TrapIfC",
&Inst::Jal { .. } => "Jal",
&Inst::CondBr { .. } => "CondBr",
&Inst::LoadExtName { .. } => "LoadExtName",
&Inst::LoadAddr { .. } => "LoadAddr",
&Inst::VirtualSPOffsetAdj { .. } => "VirtualSPOffsetAdj",
&Inst::Mov { .. } => "Mov",
&Inst::MovFromPReg { .. } => "MovFromPReg",
&Inst::ECall => "ECall",
&Inst::EBreak => "EBreak",
&Inst::Udf { .. } => "Udf",
&Inst::Jalr { .. } => "Jalr",
&Inst::Select { .. } => "Select",
&Inst::BrTable { .. } => "BrTable",
&Inst::IntSelect { .. } => "IntSelect",
&Inst::Icmp { .. } => "Icmp",
&Inst::SelectReg { .. } => "SelectReg",
&Inst::RawData { .. } => "RawData",
&Inst::Unwind { .. } => "Unwind",
&Inst::DummyUse { .. } => "DummyUse",
&Inst::Popcnt { .. } => "Popcnt",
&Inst::Cltz { .. } => "Cltz",
&Inst::Rev8 { .. } => "Rev8",
&Inst::Brev8 { .. } => "Brev8",
&Inst::StackProbeLoop { .. } => "StackProbeLoop",
&Inst::AddImm32 { .. } => "AddImm32",
};
String::from(name)
}
}

/// Different forms of label references for different instruction formats.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,16 +4,16 @@ start:
:JMP(function_0)
:JMP(finalizeExecution)
function_0:
;; TODO(mooori) call the helper to trace `MInst` execution
$${traceInstruction(ReserveSp)}
SP - 1 => SP
;; TODO(mooori) call the helper to trace `MInst` execution
$${traceInstruction(Store)}
RR :MSTORE(SP)
;; TODO(mooori) call the helper to trace `MInst` execution
;; TODO(mooori) call the helper to trace `MInst` execution
$${traceInstruction(Jal)}
$${traceInstruction(Load)}
$ => RR :MLOAD(SP)
;; TODO(mooori) call the helper to trace `MInst` execution
$${traceInstruction(ReleaseSp)}
SP + 1 => SP
;; TODO(mooori) call the helper to trace `MInst` execution
$${traceInstruction(Ret)}
:JMP(RR)
finalizeExecution:
${beforeLast()} :JMPN(finalizeExecution)
Expand Down
51 changes: 51 additions & 0 deletions tests/zkasm/helpers/InstructionTracer.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
const fs = require("fs");

const dataDump = require("./data-dump");

/**
* Handles the generation of traces of instructions executed at runtime.
*/
class InstructionTracer {
constructor() {
// Contains executed instructions in order of execution.
this.rawTrace = [];
}

setup() {
// `zkevm-proverjs` requires a `setup` function on helper objects.
}

/**
* @param {Object} ctx - The context passed by `zkevm-proverjs` to helper calls
* @param {Object} tag - Helper call specific information passed by `zkevm-proverjs`
* @param {Object} tag.params[0] - expects the identifier as first parameter of
* `$${instrumentInst(...)}`
*
* @example
* // To trace the execution of an `AluRRR` instruction, call
* // $${instrumentInst(AluRRR)}
*/
eval_traceInstruction(ctx, tag) {
const instruction = tag.params[0].varName;
this.rawTrace.push(instruction);
}

/**
* Writes the raw trace to `path`.
*
* @param {string} path
*/
writeRawTrace(path) {
if (typeof path !== "string") {
// Writing to a descriptor will append instead of replace content,
// which might result in invalid traces, see
// https://nodejs.org/api/fs.html#using-fswritefile-with-file-descriptors
throw new Error("provide a file name (not descriptor) to write to");
}
// Writing in chunks of 1000 instructions to limit memory usage,
// as raw traces might grow big.
dataDump.writeToFileInChunks(this.rawTrace, path, 1000);
}
}

module.exports = InstructionTracer;
25 changes: 25 additions & 0 deletions tests/zkasm/helpers/data-dump.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
const fs = require('fs');

/**
* Writes `lines` to a file in chunks to limit memory usage.
*
* @param {string[]} lines - each string in the array is written on its own line
* @param {string} filePath
* @param {number} chunkSize
*/
function writeToFileInChunks(lines, filePath, chunkSize) {
// If the file already exists, clear it by writing empty string.
if (fs.existsSync(filePath)) {
fs.writeFileSync(filePath, "", "utf-8");
}

// Append chunks to the file (will be created if it doesn't exist).
for (let i = 0; i < lines.length; i += chunkSize) {
let linesToWrite = lines.slice(i, i + chunkSize).join("\n") + "\n";
fs.appendFileSync(filePath, linesToWrite, "utf-8");
}
}

module.exports = {
writeToFileInChunks
}
91 changes: 78 additions & 13 deletions tests/zkasm/run-tests-zkasm.js
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ const {
} = require('pilcom');
const buildPoseidon = require('@0xpolygonhermez/zkevm-commonjs').getPoseidon;
const AssertHelper = require('./assert_helper');
const InstructionTracer = require('./helpers/InstructionTracer');

const emptyInput = require('@0xpolygonhermez/zkevm-proverjs/test/inputs/empty_input.json');

Expand All @@ -35,21 +36,38 @@ function value_to_json(key, value) {
* Run this script with `--help` to print docs.
*/
async function main() {
const argv = require("yargs/yargs")(process.argv.slice(2))
.command("$0 <path> [outfile]", "the default command runs zkASM tests", (yargs) => {
yargs.positional("path", {
describe: "The zkASM file to run or a directory to search for zkASM files.",
type: "string"
})
yargs.positional("outfile", {
describe: "If provided, results are written to this file. Otherwise they are printed to stdout.",
type: "string"
})
require("yargs/yargs")(process.argv.slice(2))
.command({
command: "$0 <path> [outfile]",
desc: "the default command runs zkASM tests",
builder: (yargs) => {
yargs.positional("path", {
describe: "The zkASM file to run or a directory to search for zkASM files.",
type: "string"
})
yargs.positional("outfile", {
describe: "If provided, results are written to this file. Otherwise they are printed to stdout.",
type: "string"
})
},
handler: (argv) => runTestsCmd(argv.path, argv.outfile)
})
.command({
command: "profile-instructions <path> [outfile]",
desc: "profiles instructions executed at runtime, assuming zkASM was instrumented",
builder: (yargs) => {
yargs.positional("path", {
describe: "The instrumented zkASM file to execute.",
type: "string"
})
yargs.positional("outfile", {
describe: "If provided, results are written to this file. Otherwise they are printed to stdout.",
type: "string"
})
},
handler: (argv) => profileInstructions(argv.path, argv.outfile)
})
.parse();

// Run the default command.
runTestsCmd(argv.path, argv.outfile);
}

/**
Expand Down Expand Up @@ -176,4 +194,51 @@ async function runTest(pathTest, cmPols) {
}
}

/**
* Executes a zkASM file instrumented for instruction tracing and produces the
* trace of executed instructions.
*
* @param {string} zkasmFile - Path to the zkASM file.
* @param {string} [outfile] - Path to a file where output is written. If not
* given, the trace is written to `stdout`.
*/
async function profileInstructions(zkasmFile, outfile) {
const configZkasm = {
defines: [],
allowUndefinedLabels: true,
allowOverwriteLabels: true,
};

// Construct helper classes.
const instructionTracer = new InstructionTracer();

// Compile rom.
const config = {
debug: true,
stepsN: 8388608,
assertOutputs: false,
helpers: [
instructionTracer
]
};
const cmPols = await compilePil();
const rom = await zkasm.compile(zkasmFile, null, configZkasm);

const result = await smMain.execute(cmPols.Main, emptyInput, rom, config);

console.log("Execution finished");
if (result.output) {
console.log(`Output: ${result.output}`)
}
if (result.logs && Object.keys(result.logs).length > 0) {
console.log(result.logs);
}

if (outfile) {
instructionTracer.writeRawTrace(outfile);
} else {
console.log(instructionTracer.rawTrace);
}
}

main();

0 comments on commit 0c0d969

Please sign in to comment.