perf(sha256): compress blocks on ARMv8 SHA2 / x86 SHA-NI (#2441)

src/foundation/sha256.c

@@ -1,11 +1,49 @@
 /* SHA-256 per FIPS 180-4. Straightforward reference implementation; validated
- * against the NIST test vectors in tests/test_cli.c. */
+ * against the NIST test vectors in tests/test_cli.c.
+ *
+ * Block compression runs on the CPU's SHA-256 instructions when it has them
+ * (ARMv8 SHA2, x86 SHA-NI), chosen once at runtime; the portable transform
+ * below stays the fallback and the reference the tests compare against.
+ * #2441: every process start fingerprints its own ~300 MB executable, and the
+ * portable transform made that ~1.2 s of CPU — past the hook deadline. */
 
 #include "foundation/sha256.h"
 #include "foundation/secure_random.h"
 
+#include <stdatomic.h>
+#include <stdbool.h>
 #include <string.h>
 
+/* The ARM path needs a way to ask the OS whether the CPU has SHA2; on any
+ * other aarch64 OS the portable transform is used. */
+#if defined(__aarch64__) && (defined(__clang__) || defined(__GNUC__)) && \
+    (defined(__APPLE__) || defined(__linux__) || defined(_WIN32))
+#define SHA256_HW_ARM 1
+#include <arm_neon.h>
+#if defined(__APPLE__)
+#include <sys/sysctl.h>
+#elif defined(__linux__)
+#include <sys/auxv.h>
+#ifndef HWCAP_SHA2
+#define HWCAP_SHA2 (1UL << 6)
+#endif
+#else
+#include <windows.h>
+#endif
+#if defined(__clang__)
+#define SHA256_ARM_TARGET __attribute__((target("sha2")))
+#else
+#define SHA256_ARM_TARGET __attribute__((target("+sha2")))
+#endif
+#elif defined(__x86_64__) && (defined(__clang__) || defined(__GNUC__))
+#define SHA256_HW_X86 1
+#include <cpuid.h>
+#include <immintrin.h>
+#define SHA256_X86_TARGET __attribute__((target("sha,sse4.1,ssse3")))
+#endif
+
+enum { SHA256_BLOCK_BYTES = 64 };
+
 static const uint32_t K[64] = {
     0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
     0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
@@ -64,6 +102,181 @@ static void sha256_transform(cbm_sha256_ctx *c, const uint8_t *data) {
     c->state[7] += h;
 }
 
+static void sha256_blocks_portable(cbm_sha256_ctx *c, const uint8_t *data, size_t blocks) {
+    for (; blocks > 0; blocks--, data += SHA256_BLOCK_BYTES) {
+        sha256_transform(c, data);
+    }
+}
+
+#if defined(SHA256_HW_ARM)
+/* ARMv8 SHA2: the state stays {a,b,c,d} / {e,f,g,h}; each SHA256H/H2 pair runs
+ * four rounds, and SU0/SU1 extend the schedule four words at a time. */
+SHA256_ARM_TARGET static void sha256_blocks_arm(cbm_sha256_ctx *c, const uint8_t *data,
+                                                size_t blocks) {
+    uint32x4_t abcd = vld1q_u32(&c->state[0]);
+    uint32x4_t efgh = vld1q_u32(&c->state[4]);
+    for (; blocks > 0; blocks--, data += SHA256_BLOCK_BYTES) {
+        uint32x4_t abcd_in = abcd;
+        uint32x4_t efgh_in = efgh;
+        uint32x4_t w[4];
+        for (int i = 0; i < 4; i++) {
+            w[i] = vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(data + i * 16)));
+        }
+        for (int i = 0; i < 16; i++) {
+            uint32x4_t wk = vaddq_u32(w[i & 3], vld1q_u32(&K[i * 4]));
+            uint32x4_t abcd_prev = abcd;
+            abcd = vsha256hq_u32(abcd, efgh, wk);
+            efgh = vsha256h2q_u32(efgh, abcd_prev, wk);
+            if (i < 12) {
+                w[i & 3] = vsha256su1q_u32(vsha256su0q_u32(w[i & 3], w[(i + 1) & 3]),
+                                           w[(i + 2) & 3], w[(i + 3) & 3]);
+            }
+        }
+        abcd = vaddq_u32(abcd, abcd_in);
+        efgh = vaddq_u32(efgh, efgh_in);
+    }
+    vst1q_u32(&c->state[0], abcd);
+    vst1q_u32(&c->state[4], efgh);
+}
+
+static bool sha256_cpu_has_arm(void) {
+#if defined(__APPLE__)
+    /* Present on every Apple arm64 CPU; the key itself exists since macOS 12. */
+    int has = 0;
+    size_t size = sizeof(has);
+    return sysctlbyname("hw.optional.arm.FEAT_SHA256", &has, &size, NULL, 0) == 0 && has != 0;
+#elif defined(__linux__)
+    return (getauxval(AT_HWCAP) & HWCAP_SHA2) != 0;
+#else
+    return IsProcessorFeaturePresent(PF_ARM_V8_CRYPTO_INSTRUCTIONS_AVAILABLE) != 0;
+#endif
+}
+#endif
+
+#if defined(SHA256_HW_X86)
+/* x86 SHA-NI: SHA256RNDS2 wants the state as {a,b,e,f} / {c,d,g,h}, so it is
+ * shuffled in once per call and back out at the end. */
+SHA256_X86_TARGET static void sha256_blocks_x86(cbm_sha256_ctx *c, const uint8_t *data,
+                                                size_t blocks) {
+    const __m128i byteswap = _mm_set_epi64x(0x0c0d0e0f08090a0bULL, 0x0405060700010203ULL);
+    __m128i dcba = _mm_shuffle_epi32(_mm_loadu_si128((const __m128i *)&c->state[0]), 0xB1);
+    __m128i efgh = _mm_shuffle_epi32(_mm_loadu_si128((const __m128i *)&c->state[4]), 0x1B);
+    __m128i abef = _mm_alignr_epi8(dcba, efgh, 8);
+    __m128i cdgh = _mm_blend_epi16(efgh, dcba, 0xF0);
+    for (; blocks > 0; blocks--, data += SHA256_BLOCK_BYTES) {
+        __m128i abef_in = abef;
+        __m128i cdgh_in = cdgh;
+        __m128i w[4];
+        for (int i = 0; i < 4; i++) {
+            w[i] = _mm_shuffle_epi8(_mm_loadu_si128((const __m128i *)(data + i * 16)), byteswap);
+        }
+        for (int i = 0; i < 16; i++) {
+            __m128i wk = _mm_add_epi32(w[i & 3], _mm_loadu_si128((const __m128i *)&K[i * 4]));
+            cdgh = _mm_sha256rnds2_epu32(cdgh, abef, wk);
+            abef = _mm_sha256rnds2_epu32(abef, cdgh, _mm_shuffle_epi32(wk, 0x0E));
+            if (i < 12) {
+                __m128i next = _mm_add_epi32(_mm_sha256msg1_epu32(w[i & 3], w[(i + 1) & 3]),
+                                             _mm_alignr_epi8(w[(i + 3) & 3], w[(i + 2) & 3], 4));
+                w[i & 3] = _mm_sha256msg2_epu32(next, w[(i + 3) & 3]);
+            }
+        }
+        abef = _mm_add_epi32(abef, abef_in);
+        cdgh = _mm_add_epi32(cdgh, cdgh_in);
+    }
+    __m128i feba = _mm_shuffle_epi32(abef, 0x1B);
+    __m128i dchg = _mm_shuffle_epi32(cdgh, 0xB1);
+    _mm_storeu_si128((__m128i *)&c->state[0], _mm_blend_epi16(feba, dchg, 0xF0));
+    _mm_storeu_si128((__m128i *)&c->state[4], _mm_alignr_epi8(dchg, feba, 8));
+}
+
+static bool sha256_cpu_has_x86(void) {
+    unsigned int eax = 0;
+    unsigned int ebx = 0;
+    unsigned int ecx = 0;
+    unsigned int edx = 0;
+    if (!__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
+        return false;
+    }
+    bool ssse3 = (ecx & (1U << 9)) != 0;
+    bool sse41 = (ecx & (1U << 19)) != 0;
+    if (!ssse3 || !sse41 || !__get_cpuid_count(7, 0, &eax, &ebx, &ecx, &edx)) {
+        return false;
+    }
+    return (ebx & (1U << 29)) != 0; /* CPUID.(7,0):EBX.SHA */
+}
+#endif
+
+typedef enum {
+    SHA256_BACKEND_UNPROBED = 0,
+    SHA256_BACKEND_PORTABLE,
+    SHA256_BACKEND_ARM,
+    SHA256_BACKEND_X86,
+} sha256_backend_t;
+
+/* Probed once; every thread computes the same answer, so relaxed is enough. */
+static _Atomic int g_sha256_backend = SHA256_BACKEND_UNPROBED;
+static _Atomic bool g_sha256_force_portable = false;
+
+static sha256_backend_t sha256_probe_backend(void) {
+#if defined(SHA256_HW_ARM)
+    if (sha256_cpu_has_arm()) {
+        return SHA256_BACKEND_ARM;
+    }
+#elif defined(SHA256_HW_X86)
+    if (sha256_cpu_has_x86()) {
+        return SHA256_BACKEND_X86;
+    }
+#endif
+    return SHA256_BACKEND_PORTABLE;
+}
+
+static sha256_backend_t sha256_backend(void) {
+    if (atomic_load_explicit(&g_sha256_force_portable, memory_order_relaxed)) {
+        return SHA256_BACKEND_PORTABLE;
+    }
+    int backend = atomic_load_explicit(&g_sha256_backend, memory_order_relaxed);
+    if (backend == SHA256_BACKEND_UNPROBED) {
+        backend = (int)sha256_probe_backend();
+        atomic_store_explicit(&g_sha256_backend, backend, memory_order_relaxed);
+    }
+    return (sha256_backend_t)backend;
+}
+
+static void sha256_blocks(cbm_sha256_ctx *c, const uint8_t *data, size_t blocks) {
+    switch (sha256_backend()) {
+#if defined(SHA256_HW_ARM)
+    case SHA256_BACKEND_ARM:
+        sha256_blocks_arm(c, data, blocks);
+        return;
+#endif
+#if defined(SHA256_HW_X86)
+    case SHA256_BACKEND_X86:
+        sha256_blocks_x86(c, data, blocks);
+        return;
+#endif
+    default:
+        sha256_blocks_portable(c, data, blocks);
+        return;
+    }
+}
+
+#if defined(CBM_ENABLE_TEST_SEAMS) && CBM_ENABLE_TEST_SEAMS
+const char *cbm_sha256_backend_name_for_testing(void) {
+    switch (sha256_backend()) {
+    case SHA256_BACKEND_ARM:
+        return "arm-sha2";
+    case SHA256_BACKEND_X86:
+        return "x86-sha-ni";
+    default:
+        return "portable";
+    }
+}
+
+void cbm_sha256_force_portable_for_testing(bool force) {
+    atomic_store_explicit(&g_sha256_force_portable, force, memory_order_relaxed);
+}
+#endif
+
 void cbm_sha256_init(cbm_sha256_ctx *c) {
     c->bitlen = 0;
     c->buflen = 0;
@@ -78,14 +291,37 @@ void cbm_sha256_init(cbm_sha256_ctx *c) {
 }
 
 void cbm_sha256_update(cbm_sha256_ctx *c, const void *data, size_t len) {
+    if (len == 0) {
+        return; /* data may be NULL when len is 0 */
+    }
     const uint8_t *p = (const uint8_t *)data;
-    for (size_t i = 0; i < len; i++) {
-        c->buf[c->buflen++] = p[i];
-        if (c->buflen == 64) {
-            sha256_transform(c, c->buf);
-            c->bitlen += 512;
-            c->buflen = 0;
+    if (c->buflen > 0) {
+        size_t take = SHA256_BLOCK_BYTES - c->buflen;
+        if (take > len) {
+            take = len;
+        }
+        memcpy(c->buf + c->buflen, p, take);
+        c->buflen += take;
+        p += take;
+        len -= take;
+        if (c->buflen < SHA256_BLOCK_BYTES) {
+            return;
         }
+        sha256_blocks(c, c->buf, 1);
+        c->bitlen += 512;
+        c->buflen = 0;
+    }
+    /* Whole blocks straight from the input, without staging them in buf. */
+    size_t blocks = len / SHA256_BLOCK_BYTES;
+    if (blocks > 0) {
+        sha256_blocks(c, p, blocks);
+        c->bitlen += (uint64_t)blocks * 512;
+        p += blocks * SHA256_BLOCK_BYTES;
+        len -= blocks * SHA256_BLOCK_BYTES;
+    }
+    if (len > 0) {
+        memcpy(c->buf, p, len);
+        c->buflen = len;
     }
 }
 
@@ -98,7 +334,7 @@ void cbm_sha256_final(cbm_sha256_ctx *c, uint8_t out[CBM_SHA256_DIGEST_LEN]) {
         while (i < 64) {
             c->buf[i++] = 0;
         }
-        sha256_transform(c, c->buf);
+        sha256_blocks(c, c->buf, 1);
         i = 0;
     }
     while (i < 56) {
@@ -108,7 +344,7 @@ void cbm_sha256_final(cbm_sha256_ctx *c, uint8_t out[CBM_SHA256_DIGEST_LEN]) {
     for (int j = 0; j < 8; j++) {
         c->buf[56 + j] = (uint8_t)(c->bitlen >> (56 - 8 * j));
     }
-    sha256_transform(c, c->buf);
+    sha256_blocks(c, c->buf, 1);
 
     for (int j = 0; j < 8; j++) {
         out[j * 4] = (uint8_t)(c->state[j] >> 24);

src/foundation/sha256.h

@@ -41,4 +41,13 @@ void cbm_sha256_hex(const void *data, size_t len, char out[CBM_SHA256_HEX_LEN +
 void cbm_hmac_sha256(const void *key, size_t key_len, const void *data, size_t data_len,
                      uint8_t out[CBM_SHA256_DIGEST_LEN]);
 
+#if defined(CBM_ENABLE_TEST_SEAMS) && CBM_ENABLE_TEST_SEAMS
+#include <stdbool.h>
+/* The block backend in use: "arm-sha2", "x86-sha-ni" or "portable". */
+const char *cbm_sha256_backend_name_for_testing(void);
+/* Force the portable transform (true) or return to the probed backend (false),
+ * so a test can hash the same input both ways. */
+void cbm_sha256_force_portable_for_testing(bool force);
+#endif
+
 #endif /* CBM_SHA256_H */

tests/test_cli.c

@@ -17174,6 +17174,71 @@ TEST(cli_sha256_file_matches_known_vector) {
     PASS();
 }
 
+/* Hash `len` bytes fed in `chunk`-sized pieces (0 = one update) to hex. */
+static void sha256_chunked_hex(const uint8_t *data, size_t len, size_t chunk,
+                               char out[CBM_SHA256_HEX_LEN + 1]) {
+    static const char hex[] = "0123456789abcdef";
+    cbm_sha256_ctx ctx;
+    cbm_sha256_init(&ctx);
+    size_t step = chunk == 0 ? len : chunk;
+    for (size_t off = 0; off < len; off += step) {
+        cbm_sha256_update(&ctx, data + off, len - off < step ? len - off : step);
+    }
+    uint8_t digest[CBM_SHA256_DIGEST_LEN];
+    cbm_sha256_final(&ctx, digest);
+    for (int i = 0; i < CBM_SHA256_DIGEST_LEN; i++) {
+        out[i * 2] = hex[digest[i] >> 4];
+        out[i * 2 + 1] = hex[digest[i] & 0x0f];
+    }
+    out[CBM_SHA256_HEX_LEN] = '\0';
+}
+
+/* #2441: block compression runs on ARMv8 SHA2 / x86 SHA-NI when the CPU has
+ * them. The hardware path must agree with the portable transform bit for bit
+ * at every length around the block and padding boundaries, for every way the
+ * input is split across update calls, and on the NIST million-'a' vector that
+ * runs 15,625 blocks through it. */
+TEST(cli_sha256_hardware_matches_portable_issue2441) {
+    enum { MAX_LEN = (1 << 20) + 13 };
+    uint8_t *data = malloc(MAX_LEN);
+    ASSERT_NOT_NULL(data);
+    uint32_t x = 0x2441u;
+    for (size_t i = 0; i < MAX_LEN; i++) {
+        x ^= x << 13;
+        x ^= x >> 17;
+        x ^= x << 5;
+        data[i] = (uint8_t)x;
+    }
+    static const size_t big[] = {4096, 65535, 65536, 65537, MAX_LEN};
+    static const size_t chunks[] = {0, 1, 3, 55, 63, 64, 65, 1000};
+    char hw[CBM_SHA256_HEX_LEN + 1];
+    char portable[CBM_SHA256_HEX_LEN + 1];
+    for (size_t n = 0; n < 300 + sizeof(big) / sizeof(big[0]); n++) {
+        size_t len = n < 300 ? n : big[n - 300];
+        for (size_t c = 0; c < sizeof(chunks) / sizeof(chunks[0]); c++) {
+            if (len > 70000 && chunks[c] != 0 && chunks[c] < 64) {
+                continue; /* byte-sized updates over 1 MiB add time, not coverage */
+            }
+            cbm_sha256_force_portable_for_testing(false);
+            sha256_chunked_hex(data, len, chunks[c], hw);
+            cbm_sha256_force_portable_for_testing(true);
+            sha256_chunked_hex(data, len, chunks[c], portable);
+            cbm_sha256_force_portable_for_testing(false);
+            ASSERT_STR_EQ(hw, portable);
+        }
+    }
+    memset(data, 'a', 1000000);
+    sha256_chunked_hex(data, 1000000, 0, hw);
+    ASSERT_STR_EQ(hw, "cdc76e5c9914fb9281a1c7e284d73e67f1809a48a497200e046d39ccc7112cd0");
+    free(data);
+#if defined(__APPLE__) && defined(__aarch64__)
+    /* Every Apple arm64 CPU has the SHA2 extension: a portable answer here
+     * means the probe or dispatch broke, and startup is slow again. */
+    ASSERT_STR_EQ(cbm_sha256_backend_name_for_testing(), "arm-sha2");
+#endif
+    PASS();
+}
+
 /* #1544: v0.10.2 deleted the ui/standard chooser AND the flags that drove it,
  * so `update --ui` — a command people had in scripts and aliases — started
  * failing with "unknown update option". Retiring a choice is fine; breaking the
@@ -17669,6 +17734,7 @@ SUITE(cli) {
     RUN_TEST(cli_progress_sink_preserves_explicit_verbose_diagnostics);
     RUN_TEST(cli_progress_sink_serializes_concurrent_callbacks);
     RUN_TEST(cli_sha256_file_matches_known_vector);
+    RUN_TEST(cli_sha256_hardware_matches_portable_issue2441);
     RUN_TEST(cli_checksum_manifest_requires_exact_filename_and_accepts_star);
     RUN_TEST(cli_checksum_manifest_rejects_invalid_missing_and_conflicting_digest);
     RUN_TEST(cli_checksum_manifest_rejects_oversized_input);