8355216: Accelerate P-256 arithmetic on aarch64 - #30941
Conversation
|
👋 Welcome back ferakocz! A progress list of the required criteria for merging this PR into |
|
@ferakocz This change now passes all automated pre-integration checks. ℹ️ This project also has non-automated pre-integration requirements. Please see the file CONTRIBUTING.md for details. After integration, the commit message for the final commit will be: You can use pull request commands such as /summary, /contributor and /issue to adjust it as needed. At the time when this comment was updated there had been 11 new commits pushed to the
As there are no conflicts, your changes will automatically be rebased on top of these commits when integrating. If you prefer to avoid this automatic rebasing, please check the documentation for the /integrate command for further details. As you do not have Committer status in this project an existing Committer must agree to sponsor your change. Possible candidates are the reviewers of this PR (@adinn, @shipilev, @theRealAph) but any other Committer may sponsor as well. ➡️ To flag this PR as ready for integration with the above commit message, type |
|
@ferakocz this pull request can not be integrated into git checkout p256-aarch64
git fetch https://git.openjdk.org/jdk.git master
git merge FETCH_HEAD
# resolve conflicts and follow the instructions given by git merge
git commit -m "Merge master"
git push |
openjdk
Bot
added
merge-conflict
|
The total number of required reviews for this PR has been set to 2 based on the presence of this label: |
ferakocz
changed the title
|
@ferakocz |
Webrevs
|
openjdk
Bot
removed
the
merge-conflict
Can you provide more information about this? What hardware did you use? What benchmarks did you run? |
| __ umulh(high, a_i, b_0); | ||
| __ mul(low, a_i, b_0); | ||
| __ lsl(high, high, shift1); | ||
| __ lsr(tmp, low, shift2); | ||
| __ orr(high, high, tmp); | ||
| __ andr(low, low, limb_mask); |
There was a problem hiding this comment.
This is a recurring pattern:
__ umulh(hi, a, b);
__ mul(lo, a, b)
__ lsl(hi, hi, SHIFT1)
__ lsr(tmp, lo, SHIFT2)
__ orr(hi, hi, tmp)
__ andr(lo, lo, mask)
which could be abstracted as a macro method (picking an arbitrary name that you probably want to improve on):
p256_partial_mul(Register a, Register b, Register hi, Register lo, Register tmp, Register mask)
You can then simplify the code that processes this limb (likewise in each each subsequent limb) to make it clearer what is being done to combine the results of these macro computations:
__ ldr(a_i, __ post(a, 8));
p256_partial_mul(a_i, b_0, high, low, tmp, limb_mask);
__ andr(n, low, limb_mask);
neon_partial_mult_64(B, b_highs, a_vals, 0);
p256_partial_mul(n, mod_0, mod_high, mod_low, tmp, limb_mask)
__ add(low, low, mod_low);
__ add(high, high, mod_high);
__ lsr(c_i, low, shift2);
__ add(c_i, c_i, high);
Also, note that the function consumes SHIFT1 and SHIFT2 which should be defined as final int constants and would be better defined at file scope rather than being declared and initialized as local variables.
There was a problem hiding this comment.
Very good idea! Thanks a lot!
| __ st1(A[0], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(D[0], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(A[1], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(D[1], __ T2D, __ post(mul_ptr, 16)); | ||
|
|
||
| __ st1(A[2], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(D[2], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(A[3], __ T2D, __ post(mul_ptr, 16)); | ||
| __ st1(D[3], __ T2D, mul_ptr); |
There was a problem hiding this comment.
You could usefully abstract this as a VSeq template function
vs_st1_interleaved(VSeq<N> A, VSeq<N> B, Register dest) {
for (int i = 0; i < N; i++) {
__ st1(A[i], __ T2D, __ post(dest, 16));
__ st1(B[i], __ T2D, __ post(dest, 16));
}
}
There was a problem hiding this comment.
Good idea. Thanks!
| __ str(c_i, Address(c_ptr, 8)); | ||
| __ mov(c_i, high); | ||
|
|
||
| vs_shl(D, __ T2D, D, 12); |
There was a problem hiding this comment.
| vs_shl(D, __ T2D, D, 12); | |
| vs_shl(D, __ T2D, D, montMulP256Shift1); |
| // generate_intpoly_montgomeryMult_P256(). | ||
| // This function computes partial results of eight 52 x 52 bit multiplications, | ||
| // where the multiplicands are stored as 64-bit values, specifically | ||
| // (b_0, b_1, b_2, b_3) * (a_3, a_4). |
There was a problem hiding this comment.
Should this not be (a_0, a_1)?
There was a problem hiding this comment.
n.b. I see now why you originally wrote (a_3, a_4). My confusion arose because at this point it was not clear tome that the as being input here were actually the 4th and fifth limb of a sequence of 5 52-bit quantities taken from input array a.
I still think it is clearer to keep a_0 and a_1 here. The comments I suggested you add to the calling routine make it clear that in the context of that method the pair of 52-bit values in register as are the 4th and 5th limbs (a3, a4). Meanwhile in this routine where we accept two full a values and four half b values it makes more sense to label the as as a_0 and a_1.
There was a problem hiding this comment.
You are right. I added some more comments to explain why it is a_3 and a_4.
| // In a call to this function, either the high or low 32 bits of the b_i | ||
| // values are multiplied by either the high or low 32 bits of the a_j values, | ||
| // so four calls with the appropriate parameters will produce the 64-bit | ||
| // low32 * low32, low32 * high32, high32 * low32, high32 * high32 | ||
| // values in the output register sequences. |
There was a problem hiding this comment.
A little more detail would make it easier to understand this method and helpt to clarify what is happening in code where it is called
| // In a call to this function, either the high or low 32 bits of the b_i | |
| // values are multiplied by either the high or low 32 bits of the a_j values, | |
| // so four calls with the appropriate parameters will produce the 64-bit | |
| // low32 * low32, low32 * high32, high32 * low32, high32 * high32 | |
| // values in the output register sequences. | |
| // Calls to this function accept either the low 32 bis or high 20 bits | |
| // of each b_i packed into bs in ascending order. a_0 and a_1 are packed | |
| // into successive 64 bit elements of as. lane selects the low 32 or high | |
| // 20 bits of each a_j value. So four calls with the appropriate parameters | |
| // will produce the 64-bit low32 * low32, low32 * high20, high20 * low32, | |
| // high20 * high20 values in the output register sequences vs. The | |
| // 64-bit partial products are returned in vs in ascending order: | |
| // vs[0] = (b_0*a_0, b_1*a_0) . . . vs[3] = (b_2*a_1, b_3*a_1) |
There was a problem hiding this comment.
Accepted with minor changes.
| __ sub(sp, sp, 48); | ||
| __ mov(c_ptr, sp); | ||
|
|
||
| // Calculate limb mask |
There was a problem hiding this comment.
| // Calculate limb mask | |
| // Calculate (52-bit) limb masks for both gpr and vector registers |
| __ ldr(mod_3, __ post(mod_ptr, 8)); | ||
| __ ldr(mod_4, mod_ptr); | ||
| __ ld1(a_vals, __ T2D, a_ptr); | ||
| __ ld2(b_lows, b_highs, __ T4S, b); |
There was a problem hiding this comment.
| __ ld2(b_lows, b_highs, __ T4S, b); | |
| // use an interleaved load to group low 32 bits and high 20 bits | |
| // of 4 successive b values into two vector registers | |
| // n.b. these are the same inputs as the ones in b_0 ... b4 | |
| __ ld2(b_lows, b_highs, __ T4S, b); |
| tmp = *common_regs++, | ||
| n = *common_regs++; | ||
|
|
||
| VSeq<4> A(16); |
There was a problem hiding this comment.
| VSeq<4> A(16); | |
| // vector sequences used to compute and combine partial products of | |
| // b_i * a_j for i = {0,1,2,3} j = {3,4} | |
| VSeq<4> A(16); |
|
|
||
| //Load input arrays and modulus | ||
| Register a_ptr = *common_regs++, mod_ptr = *common_regs++; | ||
| __ add(a_ptr, a, 24); |
There was a problem hiding this comment.
| __ add(a_ptr, a, 24); | |
| // skip 3 limbs so a_ptr addresses trailing pair {a3, a4} | |
| __ add(a_ptr, a, 24); |
|
/integrate |
|
Thanks @adinn and @theRealAph! Could one of you also sponsor it? |
|
Hi @ferakocz, did you have a chance to run test/micro/org/openjdk/bench/javax/crypto/full/PolynomialP256Bench.java on this patch?
I got these numbers on my local laptop macOS on m4:
Baseline:
|
|
I am also seeing a slowdown for this specific micro-benchhmark on a fedora M2 Mac: Baseline:
Patched:
The benchMultiply and benchSquare micro-benchmarks both show an improvement Baseline:
Patched:
I'm not sure we should automatically trust this benchmark run in isolation -- it is most important to gauge what effect the use of the multiply and assign intrinsics has when exercising the P256 API. The micro-benchmark result does suggest that the intrinsification of conditionalAssign may not always help on AArch64. However, it might still be the case that when employed in combination with the multiply intrinsic it is of benefit - possibly also depending on what hardware we are running on. Your API/method level testing showed an improvement of 9% at the method level and 5% at the API level. Have you also run these tests on your M1 machine with the intrinsic for conditionalAssign omitted? If so what was the effect? If not then could you do so and let us know what difference it makes. If you provide details of the tests run and how to exercise them I will happily check what the effect is on my M2 box if I disable generation of the conditionalAssign intrinsic. Perhaps @mrserb can do the same on his M4 Mac. Depending on the outcome might also want to check this on other AArch64 CPUs. |
I agree with that, but it's surely worth a look at the generated code so see why hand-coded |
Sorry, wrong method. But the point remains... |
|
One more thought: it might just be unrolling and inlining. |
Apparently, it is C2 generating optimal code: I haven't paid much attention to this because it is not critical for the performance of the elliptic curve computation, but it is definitely better if it is optimal. |
Ok, so it seems like we don't need to generate the assign intrinsic on AArch64. @ferakocz Can you provide numbers to confirm that omitting generation of the assign intrinsic makes no difference to the crypto computation? An interesting follow-up question is whether we need the intrinsic on x86 and if so why? It's a little more complicated since we have slightly different code depending on whether UseAVX > 2 but C2 also knows about and responds to UseAVX so one might expect it to also generate good-to-optimal code for x86. I ran the micro benchmark on my x86 box (AMD Ryzen 9 7900 12 cores) first with the assign intrinsic generation enabled and then with it disabled and got the following results Intrinsic enabled:
Intrinsic disabled:
I'm not sure why the variance is so high in the disabled case but for this HW (where we have AVX2 and AVX512 support) it looks like a similar picture. So, . . .
|
|
Are both of you absolutely sure you're allowing adequate time for warmup? Sorry if this is a bit rude, but that variance looks really suspicious |
Yeah, it does look wrong . . . Looking deeper into the test it seems that the problem here is that running the benchmark without an intrinsic is actually just inviting an apples to pears comparison. However, I did try with a longer warmup. The benchmark test is provided in The run times for the case where there was no intrinsic were as follows When I ran it again with warmups set to 20 and got this
i.e. there is still great variability even after a lot more warmup. The test code actually does this where and the intrinsic candidate With an intrinsic in place the loop body cannot really be optimized. The callouts to the intrinsic are opaque to the compiler so it just has to make 4 calls per iteration. When we disable the intrinsic and compile with Java bytecode the compiler can immediately simplify inlined code for the 4 conditionalSet calls based on the That doesn't account for the high variance but it does suggest that the comparison is not valuable as it allows the compiler to win on a special case (set is known in advance). So, @ferakocz it would still be useful to know whether disabling the intrinsic changes the results of the crypto tests you have run. That would be a good reason to consider dropping the intrinsic on aarch64 (likewise, possibly on x86). The micro-benchmark itself does not really offer any reason to do so. |
|
I think it is rather unfortunate that this method was added to this microbenchmark suite as its contribution to the run time of any real crypto operation is minimal, so it makes almost no difference if it runs twice as fast. However, it is important that it runs in constant time (i.e. its running time is independent of the values in its input arrays and, more importantly, whether the value of the "set" argument is 0 or 1). The java code was written in such a way, but there is no guarantee that the compiler will not change it back to using a branch instead of the xors if it can figure out that only those 2 values are possible for "set". So the intrinsic here is more for guaranteeing "set" value independent execution than for any performance gains. |
Yes, the assign benchmark is definitely of no use for measuring performance of the intrinsic relative to the Java code given that it hard-wires set in each call. I'm not sure it is even of much interest when we do have an intrinsic for measuring a difference between cases where Your point that the intrinsic runs in constant time is the best argument for keeping it. So, I'm happy to push this as is. The same consideration would apply for x86. |
|
@theRealAph Do you agree? |
|
Then maybe an existing benchmark can be update, to show some improvement or the new one added? |
As @ferakocz said the intrinsic is useful to ensure constant time execution in all cases. We don't actually need it to improve performance. We would only need to worry about it if it was responsible for a noticeable degradation in performance. @ferakocz says we don't see that in the tests he has performed and I think that means we are likely to face that situation with any other kernels that rely on conditionalAssign. I'm reassured by the fact that the Java version runs roughly twice as fast as the intrinsic. That matches the opportunity the provision of bytecode offers the compiler, as noted above, to halve the number of EOR instructions. Since this would not be an option in cases where set is not hard-wired I don't think we need to care about the result of this test. @mrserb please feel free to create a new benchmark that reliably tests the Java code vs the intrinsic for the general case (i.e. where set cannot be derived or predicted) and if there is a true degradation on either x86 or aarch64 then we can consider whether or not to keep the intrinsic. |
|
I take the point about constant-time implementation, but I'm still curious why an intrinsic can't match C2. Never mind, it's not important. |
|
/sponsor |
|
Going to push as commit f1cd7f6.
Your commit was automatically rebased without conflicts. |
openjdk
Bot
removed
ready
|
@ferakocz Thanks for your patience! |
8 participants
An aarch64 implementation of the MontgomeryIntegerPolynomial256.mult() method and IntegerPolynomial.conditionalAssign(). Since 64-bit multiplication is not supported on Neon and manually performing this operation with 32-bit limbs is slower than with GPRs, a hybrid neon/gpr approach is used. Neon instructions are used to compute intermediate values used in the last two iterations of the main "loop", while the GPRs compute the first few iterations. At the method level this improves performance by ~9% and at the API level roughly 5%.
Progress
Issue
Reviewers
Reviewing
Using
gitCheckout this PR locally:
$ git fetch https://git.openjdk.org/jdk.git pull/30941/head:pull/30941$ git checkout pull/30941Update a local copy of the PR:
$ git checkout pull/30941$ git pull https://git.openjdk.org/jdk.git pull/30941/headUsing Skara CLI tools
Checkout this PR locally:
$ git pr checkout 30941View PR using the GUI difftool:
$ git pr show -t 30941Using diff file
Download this PR as a diff file:
https://git.openjdk.org/jdk/pull/30941.diff
Using Webrev
Link to Webrev Comment