Fast way to compute a^d mod n

Question

For a Miller-Rabin primality test I need a fast way to compute a^d mod n.

Where a is one of 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37 with decreasing probability and increasing n. E.g. for n < 2047 only a = 2 is used. And tests are run in order until one says n is not a prime. Most non-primes are detected early.

The other bits of information I have about the parameter are n = 2^s * d + 1. n is odd, d is odd and 2 * d < n.

Below is the code I have to compute a^d mod n for uint32 values and the % operations result in __aeabi_uldivmod calls on ARM, which isn't exactly fast. Extending this to work for uint64_t requires the use of __int128, which gcc doesn't have on ARM. So I have to implement my own sqr(), mul() and mod() functions. Before I do that I wonder if the code could be improved:

#include <cstdint>
#include <bit>
#include <cassert>

// compute a^d mod n, d > 0    
uint32_t pow_mod_n(uint32_t a, uint32_t d, uint32_t n) {
    if (d == 0) __builtin_unreachable();
    unsigned shift = std::countl_zero(d) + 1;
    uint32_t t = a;
    int32_t m = d << shift;
    for (unsigned i = 32 - shift; i > 0; --i) {
      t = ((uint64_t)t * t) % n;
      if (m < 0) t = ((uint64_t)t * a) % n;
      m <<= 1;
    }
    return t;
}

Given the above limits for a, d and n I'm also looking for specializations of the code. E.g. when n < 65536 all the operations can be done with uint32_t, which speeds up the code on ARM a lot. Or code optimized for a * n < UINT32_MAX. I figure there must be a whole bunch of limits the code can be specialized for. I figure anything that lets me use a smaller type for some operation would be helpful. Has someone worked out all the thresholds that could be used to choose different functions?

Answer 1

Write readable code

Using single-letter variable names and making the source code compact doesn't have any effect on performance, but it does reduce readability and maintainability. So try to name variables more clearly, write only a single statement per line, and adding some whitespace might help as well.

The name of the function is a bit unclear, pow_n() doesn't make it clear it is going to do the power modulo some value. Why not name it pow_mod_n() instead?

Use fast integer types

You are already using fixed-width integer types, but you are using unsigned for shift and i. Here you only need something that holds a small value, and should be fast. Consider using uint_fast8_t, making your requirements explicit.

Shifting negative numbers

You already require C++20 since you use std::countl_zero(), and then shifting negative numbers is actually well-defined. However, not everyone might know that, as it was undefined behavior prior to C++20. An explanation in a comment might be helpful, or you could have just used an unsigned number and checked the top bit.

Possibility to avoid modulo in some cases

If you know something about the inputs, you might be able to avoid expensive modulo operations. For example:

E.g. for n < 2047 only a = 2 is used.

If you know a == 2, then you know that if you have a value which is already smaller than n, multiplying it by two can never bring it above 2 * n, so instead of doing a modulo you can just do an if plus a subtraction:

if (m < 0) {
    t = static_cast<uint64_t>(t) * a;
    if (t >= n) {
        t -= n;
    }
}

As you mentioned, non-primes are detected early, so this might be a worthwhile optimization to implement. It might even pay off for some more values of a.