nim-dagger/benchmarks/create_circuits.nim

182 lines
6.0 KiB
Nim
Raw Normal View History

2024-05-01 21:12:24 +00:00
import std/[hashes, json, strutils, strformat, os, osproc]
2024-04-24 10:00:05 +00:00
2024-05-02 12:56:05 +00:00
import ./utils
2024-05-01 21:12:24 +00:00
type
CircuitEnv* = object
2024-05-02 12:38:31 +00:00
nimCircuitCli*: string
circuitDirIncludes*: string
ptauPath*: string
ptauUrl*: string
codexProjDir*: string
2024-05-01 21:12:24 +00:00
2024-05-02 12:47:12 +00:00
CircuitArgs* = object
2024-05-01 21:12:24 +00:00
depth*: int
maxslots*: int
cellsize*: int
blocksize*: int
nsamples*: int
entropy*: int
seed*: int
nslots*: int
ncells*: int
index*: int
2024-05-02 12:45:25 +00:00
proc findCodexProjectDir(): string =
## find codex proj dir -- assumes this script is in codex/benchmarks
result = currentSourcePath().parentDir.parentDir
2024-05-02 12:38:31 +00:00
func default*(tp: typedesc[CircuitEnv]): CircuitEnv =
2024-05-02 12:45:25 +00:00
let codexDir = findCodexProjectDir()
result.nimCircuitCli = codexDir / "vendor" / "codex-storage-proofs-circuits" / "reference" / "nim" / "proof_input" / "cli"
result.circuitDirIncludes = codexDir / "vendor" / "codex-storage-proofs-circuits" / "circuit"
result.ptauPath = codexDir / "benchmarks" / "ceremony" / "powersOfTau28_hez_final_23.ptau"
2024-05-02 12:38:31 +00:00
result.ptauUrl = "https://storage.googleapis.com/zkevm/ptau/"
2024-05-02 12:45:25 +00:00
result.codexProjDir = codexDir
2024-05-02 12:38:31 +00:00
2024-05-02 12:48:49 +00:00
proc check*(env: var CircuitEnv) =
2024-05-01 21:21:45 +00:00
## check that the CWD of script is in the codex parent
2024-05-02 12:38:31 +00:00
let codexProjDir = findCodexProjectDir()
echo "\n\nFound project dir: ", codexProjDir
let snarkjs = findExe("snarkjs")
if snarkjs == "":
echo dedent"""
ERROR: must install snarkjs first
npm install -g snarkjs@latest
"""
let circom = findExe("circom")
if circom == "":
echo dedent"""
ERROR: must install circom first
git clone https://github.com/iden3/circom.git
cargo install --path circom
"""
if snarkjs == "" or circom == "":
quit 2
echo "Found SnarkJS: ", snarkjs
echo "Found Circom: ", circom
2024-05-02 12:38:31 +00:00
if not env.nimCircuitCli.fileExists:
echo "Nim Circuit reference cli not found: ", env.nimCircuitCli
echo "Building Circuit reference cli...\n"
2024-05-02 12:38:31 +00:00
withDir env.nimCircuitCli.parentDir:
2024-05-01 21:31:15 +00:00
runit "nimble build -d:release --styleCheck:off cli"
echo "CWD: ", getCurrentDir()
2024-05-02 12:38:31 +00:00
assert env.nimCircuitCli.fileExists()
2024-05-02 12:38:31 +00:00
echo "Found NimCircuitCli: ", env.nimCircuitCli
echo "Found Circuit Path: ", env.circuitDirIncludes
echo "Found PTAU file: ", env.ptauPath
proc downloadPtau*(ptauPath, ptauUrl: string) =
2024-05-01 21:22:36 +00:00
## download ptau file using curl if needed
if not ptauPath.fileExists:
echo "Ceremony file not found, downloading..."
2024-05-01 21:33:30 +00:00
createDir ptauPath.parentDir
withDir ptauPath.parentDir:
2024-05-01 21:31:15 +00:00
runit fmt"curl --output '{ptauPath}' '{ptauUrl}'"
2024-04-29 10:55:08 +00:00
else:
echo "Found PTAU file at: ", ptauPath
2024-05-02 13:04:58 +00:00
proc getCircuitBenchPath*(args: CircuitArgs, env: CircuitEnv): string =
## generate folder name for unique circuit args
2024-04-24 14:28:53 +00:00
var an = ""
2024-05-02 13:04:58 +00:00
for f, v in fieldPairs(args): an &= "_" & f & $v
env.codexProjDir / "benchmarks/circuit_bench" & an
2024-05-02 12:47:12 +00:00
proc generateCircomAndSamples*(args: CircuitArgs, env: CircuitEnv, name: string) =
2024-04-24 12:24:35 +00:00
## run nim circuit and sample generator
2024-05-02 12:38:31 +00:00
var cliCmd = env.nimCircuitCli
2024-04-24 12:24:35 +00:00
for f, v in fieldPairs(args):
cliCmd &= " --" & f & "=" & $v
if not "input.json".fileExists:
echo "Generating Circom Files..."
2024-05-01 21:31:15 +00:00
runit fmt"{cliCmd} -v --circom={name}.circom --output=input.json"
2024-04-24 12:24:35 +00:00
proc createCircuit*(
2024-05-02 12:47:12 +00:00
args: CircuitArgs,
2024-05-02 12:38:31 +00:00
env: CircuitEnv,
name = "proof_main",
2024-05-02 13:04:58 +00:00
circBenchDir = getCircuitBenchPath(args, env),
someEntropy = "some_entropy_75289v3b7rcawcsyiur",
2024-05-01 21:12:24 +00:00
doGenerateWitness = false,
2024-04-24 12:11:55 +00:00
): tuple[dir: string, name: string] =
## Generates all the files needed for to run a proof circuit. Downloads the PTAU file if needed.
##
2024-05-02 13:04:58 +00:00
## All needed circuit files will be generated in `circBenchDir` which
## defaults to a folder like:
## `nim-codex/benchmarks/circuit_bench_depth32_maxslots256_cellsize2048_blocksize65536_nsamples9_entropy1234567_seed12345_nslots11_ncells512_index3`
## with all the given CircuitArgs.
##
let circdir = circBenchDir
2024-05-02 12:38:31 +00:00
downloadPtau env.ptauPath, env.ptauUrl
echo "Creating circuit dir: ", circdir
2024-05-01 21:33:30 +00:00
createDir circdir
withDir circdir:
writeFile("circuit_params.json", pretty(%*args))
let
2024-05-01 21:12:24 +00:00
inputs = circdir / "input.json"
zkey = circdir / fmt"{name}.zkey"
wasm = circdir / fmt"{name}.wasm"
r1cs = circdir / fmt"{name}.r1cs"
2024-05-01 21:12:24 +00:00
wtns = circdir / fmt"{name}.wtns"
2024-05-02 12:38:31 +00:00
generateCircomAndSamples(args, env, name)
if not wasm.fileExists or not r1cs.fileExists:
2024-05-02 12:38:31 +00:00
runit fmt"circom --r1cs --wasm --O2 -l{env.circuitDirIncludes} {name}.circom"
2024-05-01 21:33:30 +00:00
moveFile fmt"{name}_js" / fmt"{name}.wasm", fmt"{name}.wasm"
echo "Found wasm: ", wasm
echo "Found r1cs: ", r1cs
if not zkey.fileExists:
echo "Zkey not found, generating..."
2024-05-01 21:33:30 +00:00
putEnv "NODE_OPTIONS", "--max-old-space-size=8192"
if not fmt"{name}_0000.zkey".fileExists:
2024-05-02 12:38:31 +00:00
runit fmt"snarkjs groth16 setup {r1cs} {env.ptauPath} {name}_0000.zkey"
echo fmt"Generated {name}_0000.zkey"
let cmd =
fmt"snarkjs zkey contribute {name}_0000.zkey {name}_0001.zkey --name='1st Contributor Name'"
echo "CMD: ", cmd
let cmdRes = execCmdEx(cmd, options = {}, input = someEntropy & "\n")
assert cmdRes.exitCode == 0
2024-04-24 10:00:05 +00:00
2024-05-01 21:33:30 +00:00
moveFile fmt"{name}_0001.zkey", fmt"{name}.zkey"
removeFile fmt"{name}_0000.zkey"
2024-04-24 10:00:05 +00:00
2024-05-01 21:12:24 +00:00
if not wtns.fileExists and doGenerateWitness:
2024-05-01 21:31:15 +00:00
runit fmt"node generate_witness.js {wtns} ../input.json ../witness.wtns"
2024-05-01 21:12:24 +00:00
2024-04-24 12:11:55 +00:00
return (circdir, name)
when isMainModule:
2024-05-02 12:38:31 +00:00
echo "findCodexProjectDir: ", findCodexProjectDir()
2024-05-01 21:21:45 +00:00
## test run creating a circuit
2024-05-02 12:38:31 +00:00
var env = CircuitEnv.default()
2024-05-02 12:48:49 +00:00
env.check()
2024-04-24 10:00:05 +00:00
2024-05-02 12:47:12 +00:00
let args = CircuitArgs(
depth: 32, # maximum depth of the slot tree
2024-05-01 21:35:05 +00:00
maxslots: 256, # maximum number of slots
cellsize: 2048, # cell size in bytes
blocksize: 65536, # block size in bytes
nsamples: 5, # number of samples to prove
entropy: 1234567, # external randomness
seed: 12345, # seed for creating fake data
nslots: 11, # number of slots in the dataset
index: 3, # which slot we prove (0..NSLOTS-1)
ncells: 512, # number of cells in this slot
)
2024-05-02 12:38:31 +00:00
let benchenv = createCircuit(args, env)
2024-05-01 21:12:24 +00:00
echo "\nBench dir:\n", benchenv