bandersnatch/fp: use optimized sqrt #38
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>
Signed-off-by: Ignacio Hagopian <jsign.uy@gmail.com>
jsign
commented
Mar 23, 2023
Comment on lines
+5
to
+17
| // The following code is _almost_ the original code from: | ||
| // https://github.com/GottfriedHerold/Bandersnatch/blob/f665f90b64892b9c4c89cff3219e70456bb431e5/bandersnatch/fieldElements/field_element_square_root.go | ||
| // | ||
| // We had to do some changes to make it work with gnark: | ||
| // - The type `feType_SquareRoot` was aliased to `Element` so everything looks the same. These types didn't have the exact | ||
| // same underlying representation, so it leaded to some minor adjustements. (e.g: accessing the limbs) | ||
| // - Original APIs regarding finite-field multiplications (e.g: MulEq) were adjusted to use gnark Mul APIs. | ||
| // - The original code had to explicitely do `Normalize()` after field element operations, but this isn't needed in gnark. | ||
| // - The primitive 2^32-root-of unity value (see init()) was pulled from gnark FFT domain code. | ||
| // - The original code used anonymous functions to define global vars, but we changed to use a init() function. | ||
| // This was required since we have other init() in the package that configure other globals (e.g: _modulus). | ||
| // By the way init() functions execution order works, we'll have these configured before the sqrt init() is called, | ||
| // compared with the original anonymous function global calls. |
There was a problem hiding this comment.
I wrote a summary here of the work I did to adjust things to work for go-ipa.
Sorry if sounds a bit cryptic. LMK any questions.
| @@ -0,0 +1,293 @@ | |||
| package fp | |||
There was a problem hiding this comment.
Note this is a new file. We should try to avoid as much as possible "pollute" generated code to avoid headaches in the future.
Comment on lines
+7
to
+44
| func TestCustomSqrt(t *testing.T) { | ||
| for i := 0; i < 10_000; i++ { | ||
| // Take a random fp. | ||
| var a Element | ||
| a.SetUint64(uint64(i)) | ||
|
|
||
| // Calculate the square root with thew optimized algorithm. | ||
| var sqrtNew Element | ||
| if sqrtNew.SqrtPrecomp(&a) == nil { | ||
| continue // Doesn't exist? Skip. | ||
| } | ||
|
|
||
| // Recalculate the original element using the calculated sqrt. | ||
| var regenNew Element | ||
| regenNew.Mul(&sqrtNew, &sqrtNew) | ||
|
|
||
| // Check the obvious: regenNew should be equal to the original element. | ||
| if !regenNew.Equal(&a) { | ||
| t.Fatalf("regenNew != a for %d", i) | ||
| } | ||
|
|
||
| // Calculate the sqrt with the original gnark code. | ||
| var sqrtOld Element | ||
| sqrtOld.Sqrt(&a) | ||
| var regenOld Element | ||
| regenOld.Mul(&sqrtOld, &sqrtOld) | ||
| if !regenOld.Equal(&a) { // Somewhat obvious, but still. | ||
| t.Fatalf("regenOld != a for %d", i) | ||
| } | ||
|
|
||
| // Check that both sqrt's are equal, *considering* the case that they have opposite signs. | ||
| // We need to do that since both algorithm can return either the positive or negative sqrt, | ||
| // which is fine. | ||
| if !sqrtNew.Equal(&sqrtOld) && !sqrtNew.Neg(&sqrtNew).Equal(&sqrtOld) { | ||
| t.Fatalf("sqrtNew != sqrtOld for %d", i) | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Wrote this extra test while I was integrating the work and double checking it was working correctly with the original implementation (that still is available).
| is_nil := y.Sqrt(&y) | ||
| is_nil := y.SqrtPrecomp(&y) |
There was a problem hiding this comment.
For computeY use the faster implementation now.
This is the only line change in generated files code.
There was a problem hiding this comment.
Open the file diff and see a comment! (Mentioning it since can be hidden since is marked as generated)
kevaundray
approved these changes
Apr 5, 2023
|
@kevaundray, is this PR mergeable? |
This was referenced Apr 10, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR includes optimization for sqrt.
My approach here is to make the minimal changes from the Bandersnatch repo implementation so we can take advantage of the speedup as soon as possible.
Upstreaming this to gnark is a separate beast, in the sense that from what I saw it might require some refactors in their
fftpackage to extract roots of unity and might not be entirely trivial. Still is a good idea to do it, but I don't want that to block our progress in the VKT speedups.My strategy was to include this separate "faster" sqrt so that doesn't mess up with generated code from gnark, since that would be painful for us in the near future when we try to (attempt) to not depend on generated code and try using gnark directly. (Again, that's a different topic!)
For more context or info: