Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Pickles] Avoid recomputation of cached verification keys #14412

Merged
merged 3 commits into from
Oct 26, 2023
Merged

[Pickles] Avoid recomputation of cached verification keys #14412

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6377db9
avoid wrap vk creation when caching
mitschabaude Oct 23, 2023
796262b
avoid step vk creation when caching
mitschabaude Oct 23, 2023
83e7f97
Merge branch 'o1js-main' into feature/avoid-cached-vk-recomputation
mitschabaude Oct 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 7 additions & 9 deletions src/lib/pickles/cache.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,8 +94,7 @@ module Step = struct
Key_cache.Sync.read cache s_p (Lazy.force k_p) )
with
| Ok (pk, dirty) ->
Common.time "step keypair create" (fun () ->
(Keypair.create ~pk ~vk:(Backend.Tick.Keypair.vk pk), dirty) )
Common.time "step keypair create" (fun () -> (pk, dirty))
| Error _e ->
let r =
Common.time "stepkeygen" (fun () ->
Expand All @@ -106,7 +105,7 @@ module Step = struct
ignore
( Key_cache.Sync.write cache s_p (Lazy.force k_p) (Keypair.pk r)
: unit Or_error.t ) ;
(r, `Generated_something) )
(Keypair.pk r, `Generated_something) )
in
let vk =
lazy
Expand All @@ -119,7 +118,7 @@ module Step = struct
(vk, `Cache_hit)
| Error _e ->
let pk, c = Lazy.force pk in
let vk = Keypair.vk pk in
let vk = Backend.Tick.Keypair.vk pk in
ignore (Key_cache.Sync.write cache s_v k_v vk : unit Or_error.t) ;
(vk, c) )
in
Expand Down Expand Up @@ -194,7 +193,7 @@ module Wrap = struct
Key_cache.Sync.read cache s_p k )
with
| Ok (pk, d) ->
(Keypair.create ~pk ~vk:(Backend.Tock.Keypair.vk pk), d)
(pk, d)
| Error _e ->
let r =
Common.time "wrapkeygen" (fun () ->
Expand All @@ -204,7 +203,7 @@ module Wrap = struct
ignore
( Key_cache.Sync.write cache s_p k (Keypair.pk r)
: unit Or_error.t ) ;
(r, `Generated_something) )
(Keypair.pk r, `Generated_something) )
in
let vk =
lazy
Expand Down Expand Up @@ -247,9 +246,8 @@ module Wrap = struct
| Ok (vk, d) ->
(vk, d)
| Error _e ->
let kp, _dirty = Lazy.force pk in
let vk = Keypair.vk kp in
let pk = Keypair.pk kp in
let pk, _dirty = Lazy.force pk in
let vk = Backend.Tock.Keypair.vk pk in
let vk : Vk.t =
{ index = vk
; commitments =
Expand Down
4 changes: 2 additions & 2 deletions src/lib/pickles/cache.mli
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ module Step : sig
-> ('a, 'b) Impls.Step.Typ.t
-> ('c, 'd) Impls.Step.Typ.t
-> ('a -> unit -> 'c)
-> ( Impls.Step.Keypair.t
-> ( Impls.Step.Proving_key.t
* [> `Cache_hit | `Generated_something | `Locally_generated ] )
lazy_t
* ( Kimchi_bindings.Protocol.VerifierIndex.Fp.t
Expand Down Expand Up @@ -67,7 +67,7 @@ module Wrap : sig
-> ('a, 'b) Impls.Wrap.Typ.t
-> ('c, 'd) Impls.Wrap.Typ.t
-> ('a -> unit -> 'c)
-> ( Impls.Wrap.Keypair.t
-> ( Impls.Wrap.Proving_key.t
* [> `Cache_hit | `Generated_something | `Locally_generated ] )
lazy_t
* ( Verification_key.Stable.V2.t
Expand Down
6 changes: 3 additions & 3 deletions src/lib/pickles/compile.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,7 +277,7 @@ struct

module Lazy_keys = struct
type t =
(Impls.Step.Keypair.t * Dirty.t) Lazy.t
(Impls.Step.Proving_key.t * Dirty.t) Lazy.t
* (Kimchi_bindings.Protocol.VerifierIndex.Fp.t * Dirty.t) Lazy.t

(* TODO Think this is right.. *)
Expand Down Expand Up @@ -733,7 +733,7 @@ struct
~f:(fun x -> [| x |])
wrap_vk.commitments )
~public_input ~auxiliary_typ ~feature_flags
(Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
(fst (Lazy.force step_pk))
wrap_vk.index
in
let step_vk = fst (Lazy.force step_vk) in
Expand Down Expand Up @@ -781,7 +781,7 @@ struct
wrap_vk.commitments )
wrap_main ~typ ~step_vk ~step_plonk_indices:(Lazy.force step_vks)
~actual_wrap_domains
(Impls.Wrap.Keypair.pk (fst (Lazy.force wrap_pk)))
(fst (Lazy.force wrap_pk))
proof
in
( return_value
Expand Down
6 changes: 3 additions & 3 deletions src/lib/pickles/pickles.ml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1076,7 +1076,7 @@ module Make_str (_ : Wire_types.Concrete) = struct

module Lazy_keys = struct
type t =
(Impls.Step.Keypair.t * Dirty.t) Lazy.t
(Impls.Step.Proving_key.t * Dirty.t) Lazy.t
* (Kimchi_bindings.Protocol.VerifierIndex.Fp.t * Dirty.t) Lazy.t

(* TODO Think this is right.. *)
Expand Down Expand Up @@ -1297,7 +1297,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
((* TODO *) Plonk_verification_key_evals.map
~f:(fun x -> [| x |])
wrap_vk.commitments )
(Impls.Step.Keypair.pk (fst (Lazy.force step_pk)))
(fst (Lazy.force step_pk))
wrap_vk.index
in
let pairing_vk = fst (Lazy.force step_vk) in
Expand Down Expand Up @@ -1785,7 +1785,7 @@ module Make_str (_ : Wire_types.Concrete) = struct
wrap_main A_value.to_field_elements ~pairing_vk
~step_domains:b.domains
~pairing_plonk_indices:(Lazy.force step_vks) ~wrap_domains
(Impls.Wrap.Keypair.pk (fst (Lazy.force wrap_pk)))
(fst (Lazy.force wrap_pk))
proof
in
Proof.T
Expand Down