Notes From The Montgomery Underground

Putting this faster code into Jolt

Ignoring speed for a minute, we plugged in the c-mul and smul multiplication algorithms in their native form² into Jolt. Then we ran cargo test -p jolt-core -- sha3_e2e_hyperkzg, and the unit tests fail for both s-mult and h-mult.

thread '<unnamed>' panicked at /Users/aribiswas3/Work-With-A16z/arkworks-algebra/ff/src/fields/models/fp/montgomery_backend.rs:1360:9:
High limb expected to be 0 if S >= 1 before 2p comparison


failures:
    jolt::vm::rv32i_vm::tests::sha3_e2e_hyperkzg

test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 219 filtered out; finished in 4.79s

Digging deeper, the part of the Jolt codebase that fails is code for subtraction. The Conditional Subtraction For Barrett Reduction throws an error claiming that the number we passed in was too large. If you look at line 1360, it specifically needs the Big integer listed as r_tmp not to have its high bit set to 1. This error goes away when using Arkworks-CIOS as the multiplication algorithm.

Why does this happen?

Scanning the lines of code for h-mult, we observe that the final answer is not reduced \(\mod p\). When you multiply two numbers that are both less than \(p\), the result may still cross over the modulus \(p\).

In Arkworks-CIOS, the last line of multiplication involves subtracting \(p\) from the final answer if the answer is greater than \(p\). The line __subtract_modulus(a)simply outputs \(p-a\) if the final output \(a \ge p\), otherwise leaves \(a\) unchanged.

In the multiplication code for s-mul, notice they have a reduce_ct function. In this function, we find that they subtract \(2p\) from \(a\) only if the most significant bit of \(a\) is 1. Thus, we have three algorithms, and although the final answer for each algorithm is the same (modulo p), the exact numerical value is different for all 3 algorithms. They are, however, not always different. Sometimes they are the same, sometimes two out of three of them are the same, and at times all three are different.

To summarise:

• c-mult: that passes unit tests and subtracts \(p\) from \(a\) when \(a \ge p\).
• h-mult: that fails unit tests and does not do any reduction.
• s-mult: that fails unit tests and subtracts \(2p\) from \(a\) when the msb of \(a\) is 1.

First attempt: We stick a __subtract_modulus(a) at the end of h-mul matching the behaviour in c-mul.

The tests still panic with the same error message

A few hours of debugging follow, we observe that even after subtracting \(p\) from the answer, the final answer is still sometimes greater than \(p\), i.e h-mul often outputs something outputs something that is sometimes greater than \(2p\).

Ok, what happens if we put in two instances of __subtract_modulus(a)?

Jolt passes!

Did we get lucky? One way to check is to try a lot of random inputs and try to break the test. \(100000\) randomly seeded inputs, with \(10\) chained multiplications in each trial, the numerical value of c-mult and h-mult are EXACTLY the same.

However, adding two conditional subtractions to s-mul – Jolt tests fail again.
But adding 3 conditional subtractions to s-mul makes Jolt pass with s-mul, and now all 3 algorithms are the same.

What is going on? Why do we need one subtraction for one algorithm, two and three for the others? How do we know we don’t need 3 or 4 or 5? \(100000\) trials are minuscule when compared to the input space of 256-bit numbers. What if there are edge cases that we can’t catch via fudging?

In what follows, we derive a formal guarantee on how large the answer of these algorithms can be, therefore establishing the maximum number of conditional subtracts we must employ. In the process, we establish that there is not just 1 new algorithm for Montgomery multiplication, but a family of algorithms that are all related to each other, each demonstrating speed for different input distributions.

Important Sub-Routine

Observe that \(c_0r^{-1} \equiv (c_0 + mp)r^{-1} \mod p\), where \(m\) such that \(c_0 + mp \equiv 0 \mod r\). Solving for \(m = -p^{-1}c_0 \equiv \mu c_0 \mod r\). Also as we are working modulo \(r\), it suffices to use \(m = \mu_{0}c_0\) discarding any carries, where \(\mu_{0}\) is the least significant word of \(\mu\).

Once we compute \(mp\), and add this to \(c\) to get \(c^{(1)}\), this should clear the least significant digit based on how we picked \(m\). See below

\[ \begin{align} c &\equiv (c_{2n-1} r^{2n-1} + c_{2n-2} r^{2n-2} + \ldots + c_1r + c_0) \mod p \\[10pt] cr^{-1} &\equiv ( c_{2n-1}r^{2n-1} + c_{2n-2} r^{2n-2} + \ldots + c_1r)r^{-1} + c_0 r^{-1} \mod p \\[10pt] &\equiv ( c_{2n-1}r^{2n-1} + c_{2n-2} r^{2n-2} + \ldots + c_1r)r^{-1} + (c_0+ mp)r^{-1} \mod p \\[10pt] &\equiv (c + mp)r^{-1} \mod p \\[10pt] &\equiv ( c^{(1)}_{2n-1}r^{2n-1} + \ldots c^{(1)}_1r + 0)r^{-1} \mod p \\[10pt] &\equiv ( c^{(1)}_{2n-1}r^{2n-2} + \ldots c^{(1)}_1) \mod p \\[10pt] \end{align} \]

Repeat The Subroutine \(n\) times

Now we can play the subroutine again with \(c^{(1)}_1\) replacing the role of \(c_0\), and the updated value of \(m = \mu_0c_1^{(1)}\) (NOTE the same \(\mu_0\) is always used)

\[ \begin{align} cr^{-1} &\equiv ( c^{(1)}_{2n-1}r^{2n-2} + \ldots c^{(1)}_1 ) \mod p \\ cr^{-2} &\equiv ( c^{(1)}_{2n-1}r^{2n-3} + \ldots + c_1^{(1)}+ c_1^{(1)}r^{-1} \mod p \\ &\equiv ( c^{(2)}_{2n-1}r^{2n-3} + \ldots + c_1^{(2)} ) \mod p \end{align} \]

In fact we play this game \(n\) times

\[cR^{-1} \equiv cr^{-n} \equiv c^{(n)}_{2n-1}r^{n-1} + c^{(n)}_{n+1}r +c^{(n)}_{n} \mod p\]

To get the final answer.

Why Can’t This Answer Be Greater Than \(2p\) ?

What did we actually do in the above process ? Over \(n\) iterations, in each iteration, we computed \(m\), added \(mp\) to \(c\) and divided by \(r\). We denote the \(m\) in iteration \(i\) with \(m_i\) here. More explicitly,

\[ \begin{align} cR^{-1} &\equiv (\ldots((c + m_0p)r^{-1} + m_1p)r^{-1} + \dots + m_{n-1}p)r^{-1} \\[10pt] &= \frac{c + p\sum_{i=0}^{n-1} m_ir^{i}}{r^n} \\ &= \frac{c + m'p }{R} \\ &< \frac{Rp + pR}{R}\\ &= 2p \end{align} \]

Here \(m'\) is some integer expressed using \(n\) limbs, therefore strictly less than \(R\). This is exactly why, if the answer is greater than \(p\), one subtraction suffices for c-mul. We NEVER need more than one subtraction!

A hybrid algorithm

More specifically, h-mul is a hybrid between the completely expanded s-mul algorithm and the c-mul algorithm. We explain this difference in the context of \(n=4\).

The h-mul algorithm can be described as an iteration of \(\log_2 n\) division operations. Contrasting this with s-mul, where we do \(n-1\) division operations by \(r^{-1}\) before adding things, or in c-mul where we NEVER divide before adding.

In h-mul, at each iteration, we reduce half of the remaining digits. If \(n=8\), we divide by \(r^{-4}, r^{-2}, r^{-1}\). If \(n=4\), we first divide by \(r^{-2} \equiv \rho_2 \mod p\), and then we divide by \(r^{-1}\equiv \rho_1 \mod p\).

So, for our bn-254 specific implementation, where \(n=4\), we first reduce by two digits \[ \begin{align} cr^{-2} &\equiv \Big(c_{7}r^{5} + \ldots +c_{3}r + c_2 \Big) + \rho_2(c_1r + c_0) \end{align} \]

In the next step, we reduce by 1 digit

\[ \begin{align} cr^{-3} &\equiv \Big(c_{7}r^{4} + \ldots +c_{3} \Big) + \rho_1(c_2 + \rho_2(c_1r + c_0)) \\ &\equiv Q + \text{Remainder} \end{align} \]

From here on with just do CIOS.

For each division operation, we blow up the current sum by at most \(rp\). See below \(n=4\).

\[ \begin{align} cr^{-2} &< \frac{p^2}{r^2} + \rho_2(c_1r + c_0) \\[10pt] &< \frac{p^2}{r^2} + \rho_2(r^2) \\[10pt] cr^{-3} &< \frac{p^2}{r^3} + \rho_1r + \rho_2r \end{align} \]

In h-mul we only divide twice to get \(Q\), whereas in s-mul we divide three times before adding. This is why the number of subtractions in h-mul will always be 1 less than the one needed for s-mul.

Note that this strategy gives rise to a family of algorithms – we can employ another division strategy. We just divide once by \(r^{-2}\), and then play CIOS.

\[ \begin{align} cr^{-2} &\equiv \Big(c_{7}r^{5} + \ldots +c_{3}r + c_2 \Big) + \rho_2(c_1r + c_0) \\[10pt] &\equiv Q +R \end{align} \]

In this case, the number of subtractions is always 1 less than that of h-mul but one more than c-mul. Thus, s-mul can be thought of as maximising parallel division operations before adding, and c-mul can be thought of as minimising division operations before an add. They are the two extremes of this family of algorithms. Using the above analysis, one can easily show the following:

Without Subtractions

The first step is to remove all subtractions from all algorithms and observe how fast things are.

warning: `minimal_mult` (example "ark-multiplication") generated 1 warning
    Finished `release` profile [optimized] target(s) in 3.76s
     Running `target/release/examples/ark-multiplication`
100000 Trials of 100 chained multiplications
Benchmarking done
Average Advantage: 11.905494513610297
C-mult    1779.19279
H-mult    1567.37109
dtype: float64

We are 11% faster. Note is very difficult to pinpoint if it’s due to fewer multiplications or how we wrote the code without using contiguous memory, but with 100,000 random inputs and 100 chained multiplications, we appear to be consistently faster. Although fast, this code will not pass – we need the answer to be reduced \(\mod p\).

Putting In The Subtractions

Now, for starters, ignoring the fact that h-mul actually needs two conditional subtractions, we stick one conditional subtraction for both c-mul and h-mul. We should still have a speed advantage, right? After all, we had one algorithm faster than the second algorithm, and we added the same line of code to both.

    Finished `release` profile [optimized] target(s) in 0.80s
     Running `target/release/examples/ark-multiplication`
100000 Trials of 100 chained multiplications
Benchmarking done
Average Advantage: -5.434536875745394
C-mult    1791.88866
H-mult    1889.26951
dtype: float64

We are 5% slower. With a little thought and listening Matt Godbolt, this observation is not surprising. As we are writing CPU code, the answer is branch prediction. The CPU gets it wrong a lot for h-mul. Why? It’s because the line of code we added was

if a >= p{
    a = p - a
}

Based on our theory so far, we know that c-mul overflows much less frequently than h-mul – With random inputs \(\mod p\), we see that CIOS only overflows above \(p\) 5% of the time, whereas h-mul overflows 68% of the time. If things were distributed uniformly⁶, it would imply that 95% of the time, the if statement evaluates to false for c-mul. So the CPU can pretty much always get rid of this line of code, and be wrong only 5% of the time.

If we wish to double down on this theory. We just put the if check, but remove the actual subtraction from the reduction (this will break correctness, but we’re only checking if our theory is correct). All our speed is recovered. The branch prediction is the same for both code bases, as the compiler will just get rid of the check in the first place.

100000 Trials of 100 chained multiplications
Benchmarking done
Average Advantage: 11.633622123165752
C-mult         1804.97929
H-mult         1594.99482
 C-overflow       0.05428
 H-overflow       0.68088
dtype: float64

Aside: Same Code Two Conclusions

As a brief distraction, in this section, we give an example of how finicky performance benchmarking can be.

One of the reasons we embarked on this journey was that preliminary experimentation showed 10-15% speedups⁷. Without the subtractions, we have a clear understanding of why this was the case (fewer multiplications and no big read/write chunks). However, even if we add the requisite subtractions to the Ingo skyscraper code code and run cargo bench based on Ingo skyscraper beches, we see something different from what’s discussed above.

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

     Running benches/my_benchmark.rs (target/release/deps/my_benchmark-8968c072d5fb1ae5)
Arkworks CIOS (10 chained c-mul, looped)
                        time:   [1.9783 µs 1.9847 µs 1.9929 µs]
Found 5 outliers among 100 measurements (5.00%)
  3 (3.00%) high mild
  2 (2.00%) high severe

h-mul (10 chained h-mul, looped)
                        time:   [1.8386 µs 1.8411 µs 1.8442 µs]

h-mul is FASTER even with the subtractions!!!! How can this be? This contradicts everything we have discussed so far. In terms of the body of the algorithms – this bench and our tests run the EXACT same code, and yet we have totally different outcomes. A lazy conclusion might be variance – but that’s not true. We can randomise, and bootstrap, build stratified experiments, h-mul is still marginally faster.

The real answer can be found by thinking about the CPU’s branch prediction algorithm and then staring at the inputs to the algorithm.

All of Montgomery’s theory relies on the fact that the initial inputs \(a\) and \(b\) are at most \(p-1\). In the Ingo skyscraper code, the inputs are random \(\mod r^4\). This strips away the overflow prevention property of CIOS – therefore, we lose the gains of the branch prediction, which makes CIOS fast in practice. Now the CPU cannot predict the branch, and this is much closer to the no-subtraction situation. Of course, we still lose a bit as h-mul needs 2 subtractions and 2 checks as opposed to 1 – hence we don’t see the 11% speedup, but still see some speedup.

In conclusion, something as innocuous as random inputs can lead to completely different conclusions about the real-world performance of code.