[SelectionDAG] General case expansion for CLMUL is bad
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
Our `TargetLowering::expandCLMUL` is not as good as it could be. It essentially does a naive bit-by-bit multiplication with the results being combined with XOR instead of added together. There are some better approaches described in https://inria.hal.science/inria-00188261v4/document and implemented in https://github.com/libntl/ntl/blob/be43be3554f366f3710d4121323ba67a5a256c96/src/GF2X.cpp#L555
I have brought this up previously on https://github.com/llvm/llvm-project/pull/140301
I have since generalized the approach from NTL, and the results for 64-to-128-bit CLMUL are:
https://quick-bench.com/q/oC7TMzhFe6LrTcyEdLa_08RJzjo
For 32-to-64-bit CLMUL (though I think the `Fast` approach can actually be greatly improved by not computing the high bits and doing everything with 64-bit operands:
https://quick-bench.com/q/CDnkvj11bAFyAdiSGLmWK7kMC7k
In any case, it seems like we can make the current software implementation 12x faster with the code shown in the benchmarks. We can port this generic `_BitInt` code over to `expandCLMUL`, which shouldn't be too hard.
For full disclosure, I let Claude figure out the generalization of the NTL code for any bit size. I verified correctness by fuzz testing against the x86 instruction with 50 million random inputs.
The approach also massively benefits from one of the operands being constant because the computation of the `A[8]` table in the code disappears completely.
CC @artagnon
```
#include
#include
#include
#include
#include
#include
#include
#include // _mm_clmulepi64_si128
#include // _mm_extract_epi64
template
using uintN_t = unsigned _BitInt(N);
// =============================================================================
// native_clmul: hardware-backed reference.
// =============================================================================
//
// PCLMULQDQ multiplies two 64-bit values to produce a 128-bit result. We
// place each operand zero-extended in the low quadword of a `__m128i`,
// invoke the intrinsic with `imm8 = 0` (selecting low * low), and read the
// lo/hi quadwords back out of the resulting vector. Because the inputs are
// zero-extended from N to 64 bits, the meaningful output occupies at most
// 2*N bits and the truncation to `_BitInt(2*N)` is exact.
template
__attribute__((target("pclmul")))
inline uintN_t<2 * N> native_clmul(uintN_t a, uintN_t b) {
static_assert(N >= 1 && N <= 64,
"native_clmul currently supports widths 1..64 (PCLMULQDQ "
"operates on 64-bit lanes)");
const unsigned long long a64 = static_cast(a);
const unsigned long long b64 = static_cast(b);
const __m128i va = _mm_set_epi64x(0, static_cast(a64));
const __m128i vb = _mm_set_epi64x(0, static_cast(b64));
const __m128i vr = _mm_clmulepi64_si128(va, vb, 0);
const unsigned long long lo =
static_cast(_mm_extract_epi64(vr, 0));
const unsigned long long hi =
static_cast(_mm_extract_epi64(vr, 1));
// Recombine into 128 bits then narrow to the public 2*N-bit result.
const unsigned _BitInt(128) full =
(static_cast(hi) << 64) |
static_cast(lo);
return static_cast>(full);
}
// =============================================================================
// naive_clmul: textbook bit-by-bit reference.
// =============================================================================
//
// For each bit i ∈ [0, N) of `b` that is set, XOR `a << i` (zero-extended to
// 2*N bits) into the running result. The operation is mathematically the
// definition of carry-less multiplication and is straightforward to audit.
// We keep it `constexpr` so the compiler can constant-fold it for fixed
// inputs in tests.
template
constexpr uintN_t<2 * N> naive_clmul(uintN_t a, uintN_t b) {
static_assert(N >= 1 && N <= 64);
using R = uintN_t<2 * N>;
R result = 0;
R aw = static_cast(a); // a zero-extended to 2*N so high shifts work
for (unsigned i = 0; i < N; ++i) {
if (static_cast((b >> i) & uintN_t{1})) {
result ^= aw << i;
}
}
return result;
}
// =============================================================================
// fast_clmul: width-generic version of NTL's 3-bit-chunked algorithm.
// =============================================================================
//
// Theory (mirrors the NTL implementation that is hard-coded for M=32):
//
// 1. Build a small table A[0..7] with A[i] = low M bits of clmul(a, i),
// where i is treated as a 3-bit unsigned. Because i has at most 3 bits,
// every entry can be obtained from at most one left-shift of `a` and
// possibly one XOR with `a`, so the table costs O(1) operations and is
// independent of M.
//
// 2. View the second operand `b` as a stream of 3-bit groups
// g_j = (b >> 3j) & 7 for j = 0, 1, …. Then
//
// clmul(a, b) = ⊕_j (clmul(a, g_j) << 3j).
//
// We accumulate that sum in two M-bit halves `lo` and `hi` representing
// the 2M-bit running result. For each group:
// lo ^= A[g_j] << 3j
// hi ^= A[g_j] >> (M - 3j) (skipped when 3j == 0 to avoid UB)
//
// 3. The table only stores the LOW M bits of clmul(a, g_j). The two top
// bits (positions M and M+1 of the true clmul) are functions of a's
// high bits and g_j's bits 1 and 2 — they are restored by two final
// correction terms. For each chunk position 3j those missing bits land
// at hi[3j] and hi[3j+1]; summed over all j this becomes
// if (a[M-1]) hi ^= (b & mask12) >> 1
// if (a[M-2]) hi ^= (b & mask2) >> 2
// where `mask12` is the union of bits {3j+1, 3j+2 : 3j+k < M} and
// `mask2` is the set of bits {3j+2 : 3j+2 < M}. The masks reduce to
// 0xb6db6db6 and 0x24924924 when M = 32, exactly what NTL uses.
//
// Generalising to widths that are not multiples of 3 just means the final
// loop iteration may have fewer than 3 leftover bits in `b` (e.g. M=8 has
// chunks at positions 0 and 3, then 2 remaining bits at positions 6,7), and
// the mask helpers need to skip out-of-range bit positions; both are handled
// uniformly by the loop bounds and the constexpr mask builders.
//
// For non-power-of-two N we round up to M = bit_ceil(N). The upper M-N bits
// of `a` are zero after zero-extension, so a[M-1] and a[M-2] are zero
// whenever they fall in those upper bits, neatly disabling the corrections
// that they control — which is exactly what we need: at width N, no bits of
// the partial product live above position N + chunk_size, so the bits the
// corrections would re-inject simply don't exist.
namespace fast_clmul_detail {
// Bits 1 and 2 of every 3-bit group of `b`, but only for positions inside
// [0, M). NTL hard-codes this as 0xb6db6db6 for M = 32; this generalisation
// derives the same value (and the analogous values for other M) at compile
// time, so the algorithm scales to any working width without touching the
// implementation.
template
constexpr uintN_t compute_mask_12() {
uintN_t mask = 0;
for (unsigned j = 0; 3u * j < M; ++j) {
if (3u * j + 1u < M) mask |= static_cast>(1) << (3u * j + 1u);
if (3u * j + 2u < M) mask |= static_cast>(1) << (3u * j + 2u);
}
return mask;
}
// Bit 2 of every 3-bit group, restricted to [0, M). NTL value: 0x24924924
// for M = 32.
template
constexpr uintN_t compute_mask_2() {
uintN_t mask = 0;
for (unsigned j = 0; 3u * j < M; ++j) {
if (3u * j + 2u < M) mask |= static_cast>(1) << (3u * j + 2u);
}
return mask;
}
} // namespace fast_clmul_detail
template
uintN_t<2 * N> fast_clmul(uintN_t a_in, uintN_t b_in) {
static_assert(N >= 1 && N <= 64,
"fast_clmul currently supports widths 1..64");
// Working width: smallest power of two ≥ N. The loop bounds and mask
// helpers all scale on M, so the only assumption being baked into the
// implementation is that the algorithm operates on a power-of-two
// storage type (which is what makes the M-bit truncation in `lo`/`hi`
// match the table semantics). N itself is unconstrained beyond the
// M ≤ 64 limit imposed by `_BitInt` storage on x86_64.
constexpr unsigned M = std::bit_ceil(N);
using U = uintN_t; // working unsigned type for lo/hi/A[*]
using R = uintN_t<2 * N>; // public return type
const U a = static_cast(a_in);
const U b = static_cast(b_in);
// --- Step 1: precomputed table A[i] = low M bits of clmul(a, i) -------
//
// Entries are derived from each other so the whole table costs at most a
// handful of XORs and 1-bit shifts regardless of M.
U A[8];
A[0] = 0;
A[1] = a;
A[2] = static_cast(A[1] << 1);
A[3] = static_cast(A[2] ^ A[1]);
A[4] = static_cast(A[2] << 1);
A[5] = static_cast(A[4] ^ A[1]);
A[6] = static_cast(A[3] << 1);
A[7] = static_cast(A[6] ^ A[1]);
// --- Step 2: process b in 3-bit chunks --------------------------------
//
// For every chunk position j with shift = 3*j inside [0, M):
// * read `chunk` = up to 3 bits of `b` starting at `shift`
// * fold A[chunk] into the running (hi, lo) pair, with the high bits
// of `t << shift` rolled into `hi` via a right shift.
//
// The final iteration may have fewer than 3 bits available when M is
// not a multiple of 3 (M ∈ {8, 16, 32, 64} here gives leftovers of
// {2, 1, 2, 1} respectively). The mask `last_mask` reduces the chunk to
// the still-valid portion of `b`; A[chunk] is then a partial-product
// entry that ignores the never-set high bit(s) of the would-be 3-bit
// index, which is exactly what we want.
U lo = 0;
U hi = 0;
for (unsigned j = 0; 3u * j < M; ++j) {
const unsigned shift = 3u * j;
const unsigned remaining = M - shift;
unsigned chunk;
if (remaining >= 3u) {
chunk = static_cast((b >> shift) & static_cast(7));
} else {
// remaining ∈ {1, 2}: build a width-correct mask for the tail.
const U last_mask =
static_cast((static_cast(1) << remaining) -
static_cast(1));
chunk = static_cast((b >> shift) & last_mask);
}
const U t = A[chunk];
lo ^= static_cast(t << shift);
// Skip the high-half update at j=0 because `t >> M` would be UB;
// also there's nothing meaningful to roll into `hi` at zero shift.
if (shift > 0u) {
hi ^= static_cast(t >> (M - shift));
}
}
// --- Step 3: high-bit corrections -------------------------------------
//
// Each A[g] discards bits ≥ M of the true clmul(a, g). Those discarded
// bits are bit M (= a[M-1]·g[1] ⊕ a[M-2]·g[2]) and bit M+1
// (= a[M-1]·g[2]). Across all chunk positions these missing bits form
// an affine function of `b` selected by the top two bits of `a`; the
// compile-time constants `mask12` and `mask2` capture exactly which
// bits of `b` need to be folded back into `hi`.
constexpr U mask12 = fast_clmul_detail::compute_mask_12();
constexpr U mask2 = fast_clmul_detail::compute_mask_2();
if ((a >> (M - 1u)) & static_cast(1)) {
hi ^= static_cast((b & mask12) >> 1u);
}
if constexpr (M >= 2u) {
if ((a >> (M - 2u)) & static_cast(1)) {
hi ^= static_cast((b & mask2) >> 2u);
}
}
// --- Step 4: recombine into a 2*M-bit value and truncate to 2*N -------
//
// When N < M the upper 2*M − 2*N bits of (hi, lo) are zero (because the
// input `a_in` had no set bits above position N-1), so the final
// narrowing cast is a value-preserving truncation.
using F = uintN_t<2u * M>;
const F full = (static_cast(hi) << M) | static_cast(lo);
return static_cast(full);
}
std::default_random_engine rng;
std::uniform_int_distribution d;
constexpr std::size_t samples = 1024 * 1024;
using T = uintN_t<32>;
const auto xs = [] {
std::array result;
for (T& e : result)
e = d(rng);
return result;
}();
const auto ys = [] {
std::array result;
for (T& e : result)
e = d(rng);
return result;
}();
// needed because DoNotOptimize does not take _BitInt
template
std::uint64_t hash(uintN_t x) {
if constexpr (N > 64) { return std::uint64_t(x >> 64) ^ std::uint64_t(x); }
else { return x; }
}
static void Native(benchmark::State& state) {
for (auto _ : state) {
for (std::size_t i = 0; i < samples; ++i) {
benchmark::DoNotOptimize(hash(native_clmul(xs[i], ys[i])));
}
}
}
BENCHMARK(Native);
static void Naive(benchmark::State& state) {
for (auto _ : state) {
for (std::size_t i = 0; i < samples; ++i) {
benchmark::DoNotOptimize(hash(naive_clmul(xs[i], ys[i])));
}
}
}
BENCHMARK(Naive);
static void Fast(benchmark::State& state) {
for (auto _ : state) {
for (std::size_t i = 0; i < samples; ++i) {
benchmark::DoNotOptimize(hash(fast_clmul(xs[i], ys[i])));
}
}
}
BENCHMARK(Fast);
```
Contributor guide
Assessment
This issue has not been assessed yet.