feat: Multi threading support for circom prover

codex/codex.nim

@@ -293,8 +293,9 @@ proc new*(
     store = NetworkStore.new(engine, repoStore)
     prover =
       if config.prover:
-        let backend =
-          config.initializeBackend().expect("Unable to create prover backend.")
+        let backend = config.initializeBackend(taskpool = taskpool).expect(
+            "Unable to create prover backend."
+          )
         some Prover.new(store, backend, config.numProofSamples)
       else:
         none Prover

codex/slots/proofs/backendfactory.nim

@@ -2,6 +2,7 @@ import os
 import strutils
 import pkg/chronos
 import pkg/chronicles
+import pkg/taskpools
 import pkg/questionable
 import pkg/confutils/defs
 import pkg/stew/io2
@@ -11,7 +12,9 @@ import ../../conf
 import ./backends
 import ./backendutils
 
-proc initializeFromConfig(config: CodexConf, utils: BackendUtils): ?!AnyBackend =
+proc initializeFromConfig(
+    config: CodexConf, utils: BackendUtils, taskpool: Taskpool
+): ?!AnyBackend =
   if not fileAccessible($config.circomR1cs, {AccessFlags.Read}) or
       not endsWith($config.circomR1cs, ".r1cs"):
     return failure("Circom R1CS file not accessible")
@@ -27,7 +30,7 @@ proc initializeFromConfig(config: CodexConf, utils: BackendUtils): ?!AnyBackend
   trace "Initialized prover backend from cli config"
   success(
     utils.initializeCircomBackend(
-      $config.circomR1cs, $config.circomWasm, $config.circomZkey
+      $config.circomR1cs, $config.circomWasm, $config.circomZkey, taskpool
     )
   )
 
@@ -41,14 +44,14 @@ proc zkeyFilePath(config: CodexConf): string =
   config.circuitDir / "proof_main.zkey"
 
 proc initializeFromCircuitDirFiles(
-    config: CodexConf, utils: BackendUtils
+    config: CodexConf, utils: BackendUtils, taskpool: Taskpool
 ): ?!AnyBackend {.gcsafe.} =
   if fileExists(config.r1csFilePath) and fileExists(config.wasmFilePath) and
       fileExists(config.zkeyFilePath):
     trace "Initialized prover backend from local files"
     return success(
       utils.initializeCircomBackend(
-        config.r1csFilePath, config.wasmFilePath, config.zkeyFilePath
+        config.r1csFilePath, config.wasmFilePath, config.zkeyFilePath, taskpool
       )
     )
 
@@ -68,11 +71,11 @@ proc suggestDownloadTool(config: CodexConf) =
     instructions
 
 proc initializeBackend*(
-    config: CodexConf, utils: BackendUtils = BackendUtils()
+    config: CodexConf, utils: BackendUtils = BackendUtils(), taskpool: Taskpool
 ): ?!AnyBackend =
-  without backend =? initializeFromConfig(config, utils), cliErr:
+  without backend =? initializeFromConfig(config, utils, taskpool), cliErr:
     info "Could not initialize prover backend from CLI options...", msg = cliErr.msg
-    without backend =? initializeFromCircuitDirFiles(config, utils), localErr:
+    without backend =? initializeFromCircuitDirFiles(config, utils, taskpool), localErr:
       info "Could not initialize prover backend from circuit dir files...",
         msg = localErr.msg
       suggestDownloadTool(config)

codex/slots/proofs/backends/circomcompat.nim

@@ -9,9 +9,11 @@
 
 {.push raises: [].}
 
-import std/sugar
+import std/[sugar, atomics, locks]
 
 import pkg/chronos
+import pkg/taskpools
+import pkg/chronos/threadsync
 import pkg/questionable/results
 import pkg/circomcompat
 
@@ -22,6 +24,7 @@
 import ./converters
 
 export circomcompat, converters
+export taskpools
 
 type
   CircomCompat* = object
@@ -35,9 +38,25 @@
     zkeyPath: string # path to the zkey file
     backendCfg: ptr CircomBn254Cfg
     vkp*: ptr CircomKey
+    taskpool: Taskpool
+    lock: ptr Lock
 
   NormalizedProofInputs*[H] {.borrow: `.`.} = distinct ProofInputs[H]
 
+  ProveTask = object
+    circom: ptr CircomCompat
+    ctx: ptr CircomCompatCtx
+    proof: ptr Proof
+    success: Atomic[bool]
+    signal: ThreadSignalPtr
+
+  VerifyTask = object
+    proof: ptr CircomProof
+    vkp: ptr CircomKey
+    inputs: ptr CircomInputs
+    success: VerifyResult
+    signal: ThreadSignalPtr
+
 func normalizeInput*[H](
     self: CircomCompat, input: ProofInputs[H]
 ): NormalizedProofInputs[H] =
@@ -79,7 +98,33 @@
   if not isNil(self.vkp):
     self.vkp.unsafeAddr.release_key()
 
-proc prove[H](self: CircomCompat, input: NormalizedProofInputs[H]): ?!CircomProof =
+  if not isNil(self.lock):
+    deinitLock(self.lock[]) # Cleanup the lock
+    dealloc(self.lock) # Free the memory
+
+proc circomProveTask(task: ptr ProveTask) {.gcsafe.} =
+  withLock task[].circom.lock[]:
+    defer:
+      discard task[].signal.fireSync()
+
+    var proofPtr: ptr Proof = nil
+    try:
+      if (
+        let res = task.circom.backendCfg.prove_circuit(task.ctx, proofPtr.addr)
+        res != ERR_OK
+      ) or proofPtr == nil:
+        task.success.store(false)
+        return
+
+      copyProof(task.proof, proofPtr[])
+      task.success.store(true)
+    finally:
+      if proofPtr != nil:
+        proofPtr.addr.release_proof()
+
+proc asyncProve*[H](
+    self: CircomCompat, input: NormalizedProofInputs[H], proof: ptr Proof
+): Future[?!void] {.async.} =
   doAssert input.samples.len == self.numSamples, "Number of samples does not match"
 
   doAssert input.slotProof.len <= self.datasetDepth,
@@ -143,58 +188,130 @@
 
   for s in input.samples:
     var
-      merklePaths = s.merklePaths.mapIt(it.toBytes)
+      merklePaths = s.merklePaths.mapIt(@(it.toBytes)).concat
       data = s.cellData.mapIt(@(it.toBytes)).concat
 
     if ctx.push_input_u256_array(
-      "merklePaths".cstring,
-      merklePaths[0].addr,
-      uint (merklePaths[0].len * merklePaths.len),
+      "merklePaths".cstring, merklePaths[0].addr, uint (merklePaths.len)
     ) != ERR_OK:
       return failure("Failed to push merkle paths")
 
     if ctx.push_input_u256_array("cellData".cstring, data[0].addr, data.len.uint) !=
         ERR_OK:
       return failure("Failed to push cell data")
 
-  var proofPtr: ptr Proof = nil
+  without threadPtr =? ThreadSignalPtr.new():
+    return failure("Unable to create thread signal")
 
-  let proof =
-    try:
-      if (let res = self.backendCfg.prove_circuit(ctx, proofPtr.addr); res != ERR_OK) or
-          proofPtr == nil:
-        return failure("Failed to prove - err code: " & $res)
+  defer:
+    threadPtr.close().expect("closing once works")
 
-      proofPtr[]
-    finally:
-      if proofPtr != nil:
-        proofPtr.addr.release_proof()
+  var task = ProveTask(circom: addr self, ctx: ctx, proof: proof, signal: threadPtr)
 
-  success proof
+  let taskPtr = addr task
+
+  doAssert task.circom.taskpool.numThreads > 1,
+    "Must have at least one separate thread or signal will never be fired"
+  task.circom.taskpool.spawn circomProveTask(taskPtr)
+  let threadFut = threadPtr.wait()
+
+  if joinErr =? catch(await threadFut.join()).errorOption:
+    if err =? catch(await noCancel threadFut).errorOption:
+      return failure(err)
+    if joinErr of CancelledError:
+      raise joinErr
+    else:
+      return failure(joinErr)
 
-proc prove*[H](self: CircomCompat, input: ProofInputs[H]): ?!CircomProof =
-  self.prove(self.normalizeInput(input))
+  if not task.success.load():
+    return failure("Failed to prove circuit")
+
+  success()
+
+proc prove*[H](
+    self: CircomCompat, input: ProofInputs[H]
+): Future[?!CircomProof] {.async, raises: [CancelledError].} =
+  var proof = ProofPtr.new()
+  defer:
+    destroyProof(proof)
+
+  try:
+    if error =? (await self.asyncProve(self.normalizeInput(input), proof)).errorOption:
+      return failure(error)
+    return success(deepCopy(proof)[])
+  except CancelledError as exc:
+    raise exc
+
+proc circomVerifyTask(task: ptr VerifyTask) {.gcsafe.} =
+  defer:
+    task[].inputs[].releaseCircomInputs()
+    discard task[].signal.fireSync()
+
+  let res = verify_circuit(task[].proof, task[].inputs, task[].vkp)
+  if res == ERR_OK:
+    task[].success[] = true
+  elif res == ERR_FAILED_TO_VERIFY_PROOF:
+    task[].success[] = false
+  else:
+    task[].success[] = false
+    error "Failed to verify proof", errorCode = res
+
+proc asyncVerify*[H](
+    self: CircomCompat,
+    proof: CircomProof,
+    inputs: ProofInputs[H],
+    success: VerifyResult,
+): Future[?!void] {.async.} =
+  var proofPtr = unsafeAddr proof
+  var inputs = inputs.toCircomInputs()
+
+  without threadPtr =? ThreadSignalPtr.new():
+    return failure("Unable to create thread signal")
+
+  defer:
+    threadPtr.close().expect("closing once works")
+
+  var task = VerifyTask(
+    proof: proofPtr,
+    vkp: self.vkp,
+    inputs: addr inputs,
+    success: success,
+    signal: threadPtr,
+  )
+
+  let taskPtr = addr task
+
+  doAssert self.taskpool.numThreads > 1,
+    "Must have at least one separate thread or signal will never be fired"
+
+  self.taskpool.spawn circomVerifyTask(taskPtr)
+
+  let threadFut = threadPtr.wait()
+
+  if joinErr =? catch(await threadFut.join()).errorOption:
+    if err =? catch(await noCancel threadFut).errorOption:
+      return failure(err)
+    if joinErr of CancelledError:
+      raise joinErr
+    else:
+      return failure(joinErr)
+
+  success()
 
 proc verify*[H](
     self: CircomCompat, proof: CircomProof, inputs: ProofInputs[H]
-): ?!bool =
+): Future[?!bool] {.async, raises: [CancelledError].} =
   ## Verify a proof using a ctx
   ##
-
-  var
-    proofPtr = unsafeAddr proof
-    inputs = inputs.toCircomInputs()
-
+  var res = VerifyResult.new()
+  defer:
+    destroyVerifyResult(res)
   try:
-    let res = verify_circuit(proofPtr, inputs.addr, self.vkp)
-    if res == ERR_OK:
-      success true
-    elif res == ERR_FAILED_TO_VERIFY_PROOF:
-      success false
-    else:
-      failure("Failed to verify proof - err code: " & $res)
-  finally:
-    inputs.releaseCircomInputs()
+    if error =? (await self.asyncVerify(proof, inputs, res)).errorOption:
+      return failure(error)
+    return success(res[])
+  except CancelledError as exc:
+    raise exc
 
 proc init*(
     _: type CircomCompat,
@@ -206,10 +323,13 @@
     blkDepth = DefaultBlockDepth,
     cellElms = DefaultCellElms,
     numSamples = DefaultSamplesNum,
+    taskpool: Taskpool,
 ): CircomCompat =
-  ## Create a new ctx
-  ##
+  # Allocate and initialize the lock
+  var lockPtr = create(Lock) # Allocate memory for the lock
+  initLock(lockPtr[]) # Initialize the lock
 
+  ## Create a new ctx
   var cfg: ptr CircomBn254Cfg
   var zkey = if zkeyPath.len > 0: zkeyPath.cstring else: nil
 
@@ -237,4 +357,6 @@
     numSamples: numSamples,
     backendCfg: cfg,
     vkp: vkpPtr,
+    taskpool: taskpool,
+    lock: lockPtr,
   )

codex/slots/proofs/backends/converters.nim

@@ -10,6 +10,7 @@
 {.push raises: [].}
 
 import pkg/circomcompat
+import std/atomics
 
 import ../../../contracts
 import ../../types
@@ -22,6 +23,14 @@
   CircomProof* = Proof
   CircomKey* = VerifyingKey
   CircomInputs* = Inputs
+  VerifyResult* = ptr bool
+  ProofPtr* = ptr Proof
+
+proc new*(_: type ProofPtr): ProofPtr =
+  cast[ptr Proof](allocShared0(sizeof(Proof)))
+
+proc new*(_: type VerifyResult): VerifyResult =
+  cast[ptr bool](allocShared0(sizeof(bool)))
 
 proc toCircomInputs*(inputs: ProofInputs[Poseidon2Hash]): CircomInputs =
   var
@@ -52,3 +61,26 @@
 
 func toGroth16Proof*(proof: CircomProof): Groth16Proof =
   Groth16Proof(a: proof.a.toG1, b: proof.b.toG2, c: proof.c.toG1)
+
+proc destroyVerifyResult*(result: VerifyResult) =
+  if result != nil:
+    deallocShared(result)
+
+proc destroyProof*(proof: ProofPtr) =
+  if proof != nil:
+    deallocShared(proof)
+
+proc copyInto*(dest: var G1, src: G1) =
+  copyMem(addr dest.x[0], addr src.x[0], 32)
+  copyMem(addr dest.y[0], addr src.y[0], 32)
+
+proc copyInto*(dest: var G2, src: G2) =
+  for i in 0 .. 1:
+    copyMem(addr dest.x[i][0], addr src.x[i][0], 32)
+    copyMem(addr dest.y[i][0], addr src.y[i][0], 32)
+
+proc copyProof*(dest: ptr Proof, src: Proof) =
+  if not isNil(dest):
+    copyInto(dest.a, src.a)
+    copyInto(dest.b, src.b)
+    copyInto(dest.c, src.c)