Speed up the u64 mining kernel ~1.8x on Apple Metal (#87)

Apple M5 Pro, CLI benchmark: 24.0 -> 43.2 MH/s. Kernel output is
unchanged and bit-exact against the CPU reference.

Field arithmetic:
- Deferred-carry additions (Acc): the external layer and the internal
  row sum accumulate unreduced with a carry counter and fold once per
  output instead of two epsilon corrections per add.
- gf64_mul_add folds the row sum into the 128-bit product before a
  single plain reduction.
- mul_wide / gf64_sqr assemble the 128-bit product from 32-bit-valued
  partial sums with no carry compares.

Kernel structure (code size dominates on Apple's compiler; every
unrolled variant measured slower):
- One loop drives all 30 rounds so each layer is emitted once, with
  round constants added by the linear layer preceding each S-box
  (RC_EXT table, RC_INTERNAL padded with a trailing zero).
- One S-box site with a runtime lane count (12 external, 1 internal).
- mining_main runs pad / first squeeze / second squeeze through a
  single inlined permute64 call in a phase loop.

Measured and rejected: uniform buffers and unhoisted loads, workgroup
sizes 64/128/512, 8-add MDS, one-compare reduce, fully unrolled
layers, two nonces per thread.
This commit is contained in:
Nikolaus Heger
2026-08-31 22:50:27 +08:00
committed by GitHub
parent 696eb51125
commit ea5e5e742e

View File

@@ -15,8 +15,10 @@ const P64: u64 = 0xFFFFFFFF00000001lu;
// EPS64 = 2^32 - 1 = 2^64 mod P
const EPS64: u64 = 0xFFFFFFFFlu;
const RC_INTERNAL: array<u64, 22> = array<u64, 22>(
0x97f7798a784ad863lu, 0xd1d2bf082f60d4f0lu, 0x69a377a79f9ad206lu, 0xa9d06906a3858e24lu, 0x295275001eede5b5lu, 0x5874e441117bd746lu, 0x8a084bbba8ed86cclu, 0x3defd7645cde6425lu, 0x3998cfe6871cc137lu, 0x3e52ef8bca48314alu, 0x964a209f85dc9ecclu, 0x3fcc9ee82cc4577elu, 0x8e79b4a5d0096d6dlu, 0x8492362ad2392556lu, 0xee72f470262574d6lu, 0x1e0e18496da2444alu, 0x0f3a74bf215eaac6lu, 0x1b061b76a1c0ded3lu, 0x192c42d86803d7a6lu, 0xf6d49ff997ae0260lu, 0x3ec372e7a0fa3786lu, 0x5538cdf4f23445d3lu
// Padded with a trailing zero so the internal layer can add the next round's
// constant unconditionally.
const RC_INTERNAL: array<u64, 23> = array<u64, 23>(
0x97f7798a784ad863lu, 0xd1d2bf082f60d4f0lu, 0x69a377a79f9ad206lu, 0xa9d06906a3858e24lu, 0x295275001eede5b5lu, 0x5874e441117bd746lu, 0x8a084bbba8ed86cclu, 0x3defd7645cde6425lu, 0x3998cfe6871cc137lu, 0x3e52ef8bca48314alu, 0x964a209f85dc9ecclu, 0x3fcc9ee82cc4577elu, 0x8e79b4a5d0096d6dlu, 0x8492362ad2392556lu, 0xee72f470262574d6lu, 0x1e0e18496da2444alu, 0x0f3a74bf215eaac6lu, 0x1b061b76a1c0ded3lu, 0x192c42d86803d7a6lu, 0xf6d49ff997ae0260lu, 0x3ec372e7a0fa3786lu, 0x5538cdf4f23445d3lu, 0lu
);
const RC_INITIAL: array<array<u64, 12>, 4> = array<array<u64, 12>, 4>(
array<u64, 12>(0xc002e770975b1607lu, 0xbca51a8dfe14593alu, 0x72938dfbe774f7f9lu, 0xe4f2fe29e03234aclu, 0xd5e0ba2f541b6449lu, 0xec33b868f3cc46c1lu, 0x486dcb55419d475alu, 0x6c1cb2a358cc24f1lu, 0xe3f30d509a1436bblu, 0xd9a64f068dca7c29lu, 0xe59b3f57aabba1aelu, 0x2a3dd4505b478fdclu),
@@ -43,19 +45,46 @@ fn gf64_add(a: u64, b: u64) -> u64 {
return s1 + select(0lu, EPS64, c2);
}
// Sum accumulator: value = lo + carries * 2^64, folded once via 2^64 ≡ EPS64 (mod P).
struct Acc {
lo: u64,
carries: u32,
}
fn acc_add(a: Acc, b: u64) -> Acc {
let s = a.lo + b;
return Acc(s, a.carries + select(0u, 1u, s < a.lo));
}
fn acc_add2(a: Acc, b: Acc) -> Acc {
let s = a.lo + b.lo;
return Acc(s, a.carries + b.carries + select(0u, 1u, s < a.lo));
}
fn acc_fold(a: Acc) -> u64 {
let c = u64(a.carries);
let t = a.lo + ((c << 32u) - c);
return t + select(0lu, EPS64, t < a.lo);
}
struct U128 {
lo: u64,
hi: u64,
}
// Reduce a 128-bit value (lo + hi*2^64) mod P using
// 2^64 ≡ EPS64 and 2^96 ≡ -1 (mod P).
fn gf64_reduce(lo: u64, hi: u64) -> u64 {
let hi_hi = hi >> 32u;
let hi_lo = hi & EPS64;
var t0 = lo - hi_hi;
t0 = t0 - select(0lu, EPS64, lo < hi_hi);
fn gf64_reduce(v: U128) -> u64 {
let hi_hi = v.hi >> 32u;
let hi_lo = v.hi & EPS64;
var t0 = v.lo - hi_hi;
t0 = t0 - select(0lu, EPS64, v.lo < hi_hi);
let t1 = hi_lo * EPS64;
let t2 = t0 + t1;
return t2 + select(0lu, EPS64, t2 < t0);
}
fn gf64_mul(a: u64, b: u64) -> u64 {
fn mul_wide(a: u64, b: u64) -> U128 {
let a_lo = a & EPS64;
let a_hi = a >> 32u;
let b_lo = b & EPS64;
@@ -64,12 +93,20 @@ fn gf64_mul(a: u64, b: u64) -> u64 {
let lh = a_lo * b_hi;
let hl = a_hi * b_lo;
let hh = a_hi * b_hi;
let mid = lh + hl;
let mid_c = select(0lu, 1lu, mid < lh);
let lo = ll + (mid << 32u);
let lo_c = select(0lu, 1lu, lo < ll);
let hi = hh + (mid >> 32u) + (mid_c << 32u) + lo_c;
return gf64_reduce(lo, hi);
let mid = (ll >> 32u) + (lh & EPS64) + (hl & EPS64);
return U128((mid << 32u) | (ll & EPS64), hh + (lh >> 32u) + (hl >> 32u) + (mid >> 32u));
}
// (a*b + addend) mod P for b <= 2^64 - 2^32 (all MDS_DIAG entries): the addend's
// value and carries are folded into the 128-bit product before reduction.
fn gf64_mul_add(a: u64, b: u64, addend: Acc) -> u64 {
let v = mul_wide(a, b);
let lo = v.lo + addend.lo;
return gf64_reduce(U128(lo, v.hi + u64(addend.carries) + select(0lu, 1lu, lo < v.lo)));
}
fn gf64_mul(a: u64, b: u64) -> u64 {
return gf64_reduce(mul_wide(a, b));
}
fn gf64_sqr(a: u64) -> u64 {
@@ -78,12 +115,8 @@ fn gf64_sqr(a: u64) -> u64 {
let ll = a_lo * a_lo;
let lh = a_lo * a_hi;
let hh = a_hi * a_hi;
let mid = lh << 1u;
let mid_c = lh >> 63u;
let lo = ll + (mid << 32u);
let lo_c = select(0lu, 1lu, lo < ll);
let hi = hh + (mid >> 32u) + (mid_c << 32u) + lo_c;
return gf64_reduce(lo, hi);
let mid = (ll >> 32u) + ((lh & EPS64) << 1u);
return gf64_reduce(U128((mid << 32u) | (ll & EPS64), hh + ((lh >> 32u) << 1u) + (mid >> 32u)));
}
fn gf64_sbox(x: u64) -> u64 {
@@ -97,67 +130,109 @@ fn gf64_canon(a: u64) -> u64 {
return a - select(0lu, P64, a >= P64);
}
// External linear layer: 4x4 MDS on each chunk, then circulant sums.
fn ext_layer64(state: ptr<function, array<u64, 12>>) {
// 4x4 MDS circ(2, 3, 1, 1) on one chunk, results left unreduced.
fn mds4(x0: u64, x1: u64, x2: u64, x3: u64) -> array<Acc, 4> {
let t01 = acc_add(Acc(x0, 0u), x1);
let t23 = acc_add(Acc(x2, 0u), x3);
let t0123 = acc_add2(t01, t23);
let t01123 = acc_add(t0123, x1);
let t01233 = acc_add(t0123, x3);
return array<Acc, 4>(
acc_add2(t01123, t01),
acc_add2(t01123, acc_add(Acc(x2, 0u), x2)),
acc_add2(t01233, t23),
acc_add2(t01233, acc_add(Acc(x0, 0u), x0))
);
}
// External linear layer: 4x4 MDS on each chunk, then circulant sums, plus the
// next round's constants. Additions are accumulated unreduced (at most 27
// carries) and folded once per output.
fn ext_layer64(state: ptr<function, array<u64, 12>>, rc: array<u64, 12>) {
var y: array<Acc, 12>;
for (var chunk = 0u; chunk < 3u; chunk++) {
let o = chunk * 4u;
let x0 = (*state)[o];
let x1 = (*state)[o + 1u];
let x2 = (*state)[o + 2u];
let x3 = (*state)[o + 3u];
let t01 = gf64_add(x0, x1);
let t23 = gf64_add(x2, x3);
let t0123 = gf64_add(t01, t23);
let t01123 = gf64_add(t0123, x1);
let t01233 = gf64_add(t0123, x3);
(*state)[o + 3u] = gf64_add(t01233, gf64_add(x0, x0));
(*state)[o + 1u] = gf64_add(t01123, gf64_add(x2, x2));
(*state)[o] = gf64_add(t01123, t01);
(*state)[o + 2u] = gf64_add(t01233, t23);
let m = mds4((*state)[o], (*state)[o + 1u], (*state)[o + 2u], (*state)[o + 3u]);
y[o] = m[0];
y[o + 1u] = m[1];
y[o + 2u] = m[2];
y[o + 3u] = m[3];
}
var sums: array<u64, 4>;
for (var k = 0u; k < 4u; k++) {
sums[k] = gf64_add(gf64_add((*state)[k], (*state)[k + 4u]), (*state)[k + 8u]);
}
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_add((*state)[i], sums[i % 4u]);
let s = acc_add2(acc_add2(y[k], y[k + 4u]), y[k + 8u]);
(*state)[k] = acc_fold(acc_add(acc_add2(y[k], s), rc[k]));
(*state)[k + 4u] = acc_fold(acc_add(acc_add2(y[k + 4u], s), rc[k + 4u]));
(*state)[k + 8u] = acc_fold(acc_add(acc_add2(y[k + 8u], s), rc[k + 8u]));
}
}
// Internal linear layer: diagonal matrix plus full sum.
fn int_layer64(state: ptr<function, array<u64, 12>>) {
var sum = (*state)[0];
for (var i = 1u; i < 12u; i++) {
sum = gf64_add(sum, (*state)[i]);
}
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_add(gf64_mul((*state)[i], MDS_DIAG[i]), sum);
// Internal linear layer: diagonal matrix plus full sum, plus the next round's
// constant on element 0 (the S-box output, which is summed last).
fn int_layer64(state: ptr<function, array<u64, 12>>, rc0: u64) {
let s12 = acc_add(Acc((*state)[1], 0u), (*state)[2]);
let s34 = acc_add(Acc((*state)[3], 0u), (*state)[4]);
let s56 = acc_add(Acc((*state)[5], 0u), (*state)[6]);
let s78 = acc_add(Acc((*state)[7], 0u), (*state)[8]);
let s910 = acc_add(Acc((*state)[9], 0u), (*state)[10]);
let s1234 = acc_add2(s12, s34);
let s5678 = acc_add2(s56, s78);
let s91011 = acc_add(s910, (*state)[11]);
let sum = acc_add(acc_add2(acc_add2(s1234, s5678), s91011), (*state)[0]);
(*state)[0] = gf64_mul_add((*state)[0], MDS_DIAG[0], acc_add(sum, rc0));
(*state)[1] = gf64_mul_add((*state)[1], MDS_DIAG[1], sum);
(*state)[2] = gf64_mul_add((*state)[2], MDS_DIAG[2], sum);
(*state)[3] = gf64_mul_add((*state)[3], MDS_DIAG[3], sum);
(*state)[4] = gf64_mul_add((*state)[4], MDS_DIAG[4], sum);
(*state)[5] = gf64_mul_add((*state)[5], MDS_DIAG[5], sum);
(*state)[6] = gf64_mul_add((*state)[6], MDS_DIAG[6], sum);
(*state)[7] = gf64_mul_add((*state)[7], MDS_DIAG[7], sum);
(*state)[8] = gf64_mul_add((*state)[8], MDS_DIAG[8], sum);
(*state)[9] = gf64_mul_add((*state)[9], MDS_DIAG[9], sum);
(*state)[10] = gf64_mul_add((*state)[10], MDS_DIAG[10], sum);
(*state)[11] = gf64_mul_add((*state)[11], MDS_DIAG[11], sum);
}
fn sbox_lanes(state: ptr<function, array<u64, 12>>, lanes: u32) {
for (var i = 0u; i < lanes; i++) {
(*state)[i] = gf64_sbox((*state)[i]);
}
}
fn add_rc(state: ptr<function, array<u64, 12>>, rc: array<u64, 12>) {
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_add((*state)[i], rc[i]);
}
}
const RC_ZERO: array<u64, 12> = array<u64, 12>(0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu);
// Constants added by each external layer, which feeds the S-box that follows it:
// entries 0..3 precede the initial rounds, entry 4 carries the first internal
// round's constant on element 0, entries 5..7 precede terminal rounds 1..3.
const RC_EXT: array<array<u64, 12>, 9> = array<array<u64, 12>, 9>(
RC_INITIAL[0], RC_INITIAL[1], RC_INITIAL[2], RC_INITIAL[3],
array<u64, 12>(RC_INTERNAL[0], 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu, 0lu),
RC_TERMINAL[1], RC_TERMINAL[2], RC_TERMINAL[3], RC_ZERO
);
// One loop drives all 30 rounds so each layer is emitted once: step 0 is the
// initial external layer, steps 1..4 and 27..30 are external rounds, steps
// 5..26 are internal rounds. Round constants are added by the linear layer
// preceding each S-box (the last internal layer is followed by the first
// terminal round's constants).
fn permute64(state: ptr<function, array<u64, 12>>) {
ext_layer64(state);
for (var r = 0u; r < 4u; r++) {
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_add((*state)[i], RC_INITIAL[r][i]);
for (var k = 0u; k < 31u; k++) {
let is_ext = k < 5u || k > 26u;
if (k > 0u) {
sbox_lanes(state, select(1u, 12u, is_ext));
}
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_sbox((*state)[i]);
if (is_ext) {
ext_layer64(state, RC_EXT[select(k, k - 22u, k > 26u)]);
} else {
int_layer64(state, RC_INTERNAL[k - 4u]);
if (k == 26u) {
add_rc(state, RC_TERMINAL[0]);
}
}
ext_layer64(state);
}
for (var r = 0u; r < 22u; r++) {
(*state)[0] = gf64_sbox(gf64_add((*state)[0], RC_INTERNAL[r]));
int_layer64(state);
}
for (var r = 0u; r < 4u; r++) {
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_add((*state)[i], RC_TERMINAL[r][i]);
}
for (var i = 0u; i < 12u; i++) {
(*state)[i] = gf64_sbox((*state)[i]);
}
ext_layer64(state);
}
}
@@ -221,63 +296,65 @@ fn mining_main(@builtin(global_invocation_id) global_id: vec3<u32>) {
// Resume the sponge from the precomputed midstate: absorb the low
// nonce half, pad, squeeze twice (3 permutations instead of 5).
var st: array<u64, 12>;
for (var i = 0u; i < 12u; i++) {
st[i] = mid[i];
}
var st = mid;
for (var i = 0u; i < 8u; i++) {
st[i] = gf64_add(st[i], u64(bswap32(current_nonce[7u - i])));
}
permute64(&st);
st[0] = gf64_add(st[0], 1lu);
st[1] = gf64_add(st[1], 1lu);
permute64(&st);
// First squeeze yields the most significant 256 bits of the hash, which
// decide hash-vs-target on their own unless they exactly equal the
// target's high half. Only candidates pay for the second squeeze, and
// byte-swapped hash words are produced on demand during the compare.
var first: array<u32, 8>;
for (var i = 0u; i < 4u; i++) {
let c = gf64_canon(st[i]);
first[2u * i] = u32(c & EPS64);
first[2u * i + 1u] = u32(c >> 32u);
}
// Squeeze-and-compare phases share one inlined permutation: phase 0
// pads after absorbing, phase 1 yields the most significant 256 bits of
// the hash, which decide hash-vs-target on their own unless they exactly
// equal the target's high half, and only candidates run phase 2 for the
// low half. Byte-swapped hash words are produced on demand.
var hash_le: array<u32, 16>;
var cmp = 0u;
for (var i = 0u; i < 8u; i++) {
let h = bswap32(first[i]);
let t = tgt[15u - i];
if (h != t) {
cmp = select(2u, 1u, h > t);
break;
var below = false;
for (var phase = 0u; phase < 3u; phase++) {
permute64(&st);
if (phase == 0u) {
st[0] = gf64_add(st[0], 1lu);
st[1] = gf64_add(st[1], 1lu);
continue;
}
var words: array<u32, 8>;
for (var i = 0u; i < 4u; i++) {
let c = gf64_canon(st[i]);
words[2u * i] = bswap32(u32(c & EPS64));
words[2u * i + 1u] = bswap32(u32(c >> 32u));
}
let base = select(15u, 7u, phase == 2u);
for (var i = 0u; i < 8u; i++) {
hash_le[base - i] = words[i];
}
if (phase == 1u) {
for (var i = 0u; i < 8u; i++) {
let h = words[i];
let t = tgt[15u - i];
if (h != t) {
cmp = select(2u, 1u, h > t);
break;
}
}
if (cmp == 1u) {
break;
}
} else {
below = cmp == 2u;
if (!below) {
for (var i = 0u; i < 8u; i++) {
let h = words[i];
let t = tgt[7u - i];
if (h != t) {
below = h < t;
break;
}
}
}
}
}
if (cmp == 1u) {
continue;
}
var hash_le: array<u32, 16>;
for (var i = 0u; i < 8u; i++) {
hash_le[15u - i] = bswap32(first[i]);
}
permute64(&st);
for (var i = 0u; i < 4u; i++) {
let c = gf64_canon(st[i]);
hash_le[7u - 2u * i] = bswap32(u32(c & EPS64));
hash_le[6u - 2u * i] = bswap32(u32(c >> 32u));
}
var below = cmp == 2u;
if (!below) {
for (var i = 0u; i < 8u; i++) {
let h = hash_le[7u - i];
let t = tgt[7u - i];
if (h != t) {
below = h < t;
break;
}
}
}
if (below) {
if (atomicExchange(&results[0], 1u) == 0u) {
for (var i = 0u; i < 16u; i++) {
@@ -504,14 +581,14 @@ fn state_unpack(v: ptr<function, array<u64, 12>>, state: ptr<function, array<Gol
fn external_linear_layer(state: ptr<function, array<GoldilocksField, 12>>) {
var st: array<u64, 12>;
state_pack(state, &st);
ext_layer64(&st);
ext_layer64(&st, RC_ZERO);
state_unpack(&st, state);
}
fn internal_linear_layer(state: ptr<function, array<GoldilocksField, 12>>) {
var st: array<u64, 12>;
state_pack(state, &st);
int_layer64(&st);
int_layer64(&st, 0lu);
state_unpack(&st, state);
}