/** * mlkem_test - ML-KEM hardware accelerator example (CIC-IEC 2026) * * Exercises the mlkem_axi_slave at 0xbf400000 by running a full * KeyGen -> Encaps -> Decaps round trip on ML-KEM-512 (k=2) and verifying * that the decapsulated shared secret matches the encapsulated one. This is a * self-consistency check: it needs no external golden vectors, yet touches * every register path (seeds, start/poll, byte preload, byte readback, ss). * * KeyGen : d,z -> ek (800 B), dk (1632 B) * Encaps : ek + msg -> ss, ct (768 B) * Decaps : dk + ct -> ss' (must equal ss) * * The core runs at sys_clk/2 (25 MHz); the wrapper hides the clock crossing, * so software just does normal MMIO loads/stores (preload writes are * automatically wait-stated on the bus). */ #include #include #include #include #include "dma.h" #include "mlkem_native.h" /* ─── BSP globals ─────────────────────────────────────────────────── */ unsigned long UART_BASE = 0xbf000000; unsigned long CONFREG_TIMER_BASE = 0xbf20f100; unsigned long CONFREG_CLOCKS_PER_SEC = 50000000L; unsigned long CORE_CLOCKS_PER_SEC = 33000000L; /* ─── ML-KEM IP register map (base 0xbf400000) ────────────────────── */ #define MLKEM_BASE 0xbf400000 #define MLKEM_CTRL (MLKEM_BASE + 0x000) /* W bit0 = start */ #define MLKEM_STATUS (MLKEM_BASE + 0x004) /* R busy/done/preload_busy */ #define MLKEM_CONFIG (MLKEM_BASE + 0x008) /* RW [2:0]=k [5:4]=op */ #define MLKEM_D (MLKEM_BASE + 0x010) /* RW seed d (8 words) */ #define MLKEM_Z (MLKEM_BASE + 0x030) /* RW seed z (8 words) */ #define MLKEM_MSG (MLKEM_BASE + 0x050) /* RW message (8 words) */ #define MLKEM_SS (MLKEM_BASE + 0x070) /* R shared secret (8 words) */ #define MLKEM_PRE_SEL (MLKEM_BASE + 0x100) /* RW 0=ek 1=dk 2=c */ #define MLKEM_PRE_ADR (MLKEM_BASE + 0x104) /* RW byte addr (auto-incr) */ #define MLKEM_PRE_DAT (MLKEM_BASE + 0x108) /* W byte -> selected port */ #define MLKEM_RDB_SEL (MLKEM_BASE + 0x10C) /* RW 0=ek 1=dk_pke 2=dk 3=ct */ #define MLKEM_RDB_ADR (MLKEM_BASE + 0x110) /* RW output byte index */ #define MLKEM_RDB_DAT (MLKEM_BASE + 0x114) /* R selected byte (low 8b) */ #define MLKEM_EK_WINDOW (MLKEM_BASE + 0x1000) /* W preload ek, R read ek */ #define MLKEM_DK_WINDOW (MLKEM_BASE + 0x2000) /* W preload dk, R read dk */ #define MLKEM_CT_WINDOW (MLKEM_BASE + 0x3000) /* W preload c, R read ct */ /* PRNG sub-window in the shared ML-KEM/FFT/PRNG AXI slave. */ #define PRNG_BASE (MLKEM_BASE + 0xD000) #define PRNG_CTRL (PRNG_BASE + 0x000) /* W bit0=mix seed bit1=clr count */ #define PRNG_STATUS (PRNG_BASE + 0x004) /* R bit0=ready */ #define PRNG_DATA (PRNG_BASE + 0x008) /* R next pseudo-random word */ #define PRNG_COUNT (PRNG_BASE + 0x00C) /* R data read count */ #define PRNG_SEED (PRNG_BASE + 0x010) /* RW software seed mix value */ #define CTRL_START 0x1 #define STATUS_BUSY 0x1 #define STATUS_DONE 0x2 #define PRNG_READY 0x1 #define PRNG_CTRL_MIX 0x1 #define OP_KEYGEN 0 #define OP_ENCAPS 1 #define OP_DECAPS 2 /* preload select */ #define PRE_EK 0 #define PRE_DK 1 #define PRE_C 2 /* readback select */ #define RDB_EK 0 #define RDB_DKP 1 #define RDB_DK 2 #define RDB_CT 3 /* ─── ML-KEM-512 (k=2) byte sizes ─────────────────────────────────── */ #define KVAL 2 #define EK_BYTES (384 * KVAL + 32) /* 800 */ #define DKP_BYTES (384 * KVAL) /* 768 */ #define DK_BYTES (DKP_BYTES + EK_BYTES + 32 + 32) /* 1632 */ #define CT_BYTES 768 /* ─── Convenience macros ──────────────────────────────────────────── */ #define REG_WR(addr, val) (*(volatile uint32_t *)(addr) = (uint32_t)(val)) #define REG_RD(addr) (*(volatile uint32_t *)(addr)) #define PHYS_ADDR(addr) ((uint32_t)(addr) & 0x1FFFFFFFu) #define UNCACHED_ADDR(addr) ((PHYS_ADDR(addr)) | 0xA0000000u) /* PLACEHOLDER */ /* ─── Wait until the core finishes the current op ─────────────────── */ static void wait_done(void) { while (!(REG_RD(MLKEM_STATUS) & STATUS_DONE)) ; } /* Clear the latched done flag (W1C) before launching a new op. */ static void clear_done(void) { REG_WR(MLKEM_STATUS, STATUS_DONE); } /* Launch the op currently configured in CONFIG and block until done. */ static void run_op(uint32_t op) { REG_WR(MLKEM_CONFIG, ((op & 0x3) << 4) | (KVAL & 0x7)); clear_done(); REG_WR(MLKEM_CTRL, CTRL_START); wait_done(); } /* Write a 256-bit value (32 bytes) into one of the seed/msg windows. * Byte i of the logical value lands in word[i>>2] byte (i&3) — i.e. the * core sees byte 0 in D[0][7:0], matching mlkem_top's d_i[7:0] convention. */ static void write_block32(uint32_t base, const uint8_t *bytes) { int w; for (w = 0; w < 8; w++) { uint32_t v = (uint32_t)bytes[w*4 + 0] | ((uint32_t)bytes[w*4 + 1] << 8) | ((uint32_t)bytes[w*4 + 2] << 16) | ((uint32_t)bytes[w*4 + 3] << 24); REG_WR(base + w*4, v); } } /* Read a 256-bit value (32 bytes) back from a window (used for SS). */ static void read_block32(uint32_t base, uint8_t *bytes) { int w; for (w = 0; w < 8; w++) { uint32_t v = REG_RD(base + w*4); bytes[w*4 + 0] = (uint8_t)(v & 0xFF); bytes[w*4 + 1] = (uint8_t)((v >> 8) & 0xFF); bytes[w*4 + 2] = (uint8_t)((v >> 16) & 0xFF); bytes[w*4 + 3] = (uint8_t)((v >> 24) & 0xFF); } } /* Filled from the PRNG endpoint at runtime. */ static uint8_t SEED_D[32]; static uint8_t SEED_Z[32]; static uint8_t MSG_M[32]; static uint32_t prng_read_word(void) { while ((REG_RD(PRNG_STATUS) & PRNG_READY) == 0) ; return REG_RD(PRNG_DATA); } static void prng_mix_seed(uint32_t seed) { REG_WR(PRNG_SEED, seed); REG_WR(PRNG_CTRL, PRNG_CTRL_MIX); } static void prng_fill(uint8_t *bytes, int n) { int i; for (i = 0; i < n; i += 4) { uint32_t v = prng_read_word(); bytes[i + 0] = (uint8_t)(v & 0xFF); if (i + 1 < n) bytes[i + 1] = (uint8_t)((v >> 8) & 0xFF); if (i + 2 < n) bytes[i + 2] = (uint8_t)((v >> 16) & 0xFF); if (i + 3 < n) bytes[i + 3] = (uint8_t)((v >> 24) & 0xFF); } } static void generate_prng_inputs(void) { prng_mix_seed((uint32_t)get_clock_count()); prng_fill(SEED_D, sizeof(SEED_D)); prng_fill(SEED_Z, sizeof(SEED_Z)); prng_fill(MSG_M, sizeof(MSG_M)); } /* Stream `n` bytes into a preload port. Each PRE_DAT write auto-increments * the preload address; the bus wait-states each write until the cross-domain * byte handshake completes, so a plain loop is correct. */ static void preload(uint32_t sel, const uint8_t *buf, int n) { int i; REG_WR(MLKEM_PRE_SEL, sel); REG_WR(MLKEM_PRE_ADR, 0); for (i = 0; i < n; i++) REG_WR(MLKEM_PRE_DAT, buf[i]); } /* Read `n` bytes out of a readback tap. RDB_ADR has no auto-increment, so set * the index for each byte (write addr, then read data). */ static void readback(uint32_t sel, uint8_t *buf, int n) { int i; REG_WR(MLKEM_RDB_SEL, sel); for (i = 0; i < n; i++) { REG_WR(MLKEM_RDB_ADR, i); buf[i] = (uint8_t)(REG_RD(MLKEM_RDB_DAT) & 0xFF); } } /* PLACEHOLDER2 */ static volatile uint8_t *uncached_u8(uint8_t *buf) { return (volatile uint8_t *)UNCACHED_ADDR(buf); } static void dma_to_mlkem(uint32_t dst_window, volatile uint8_t *src, int n) { dma_start_transfer(0, PHYS_ADDR(src), PHYS_ADDR(dst_window), (uint32_t)n, 0); dma_wait_polling(0); } static void dma_from_mlkem(volatile uint8_t *dst, uint32_t src_window, int n) { dma_start_transfer(0, PHYS_ADDR(src_window), PHYS_ADDR(dst), (uint32_t)n, 0); dma_wait_polling(0); } static void print_hex_uc(const char *label, volatile uint8_t *b, int n) { int i; printf("%s", label); for (i = 0; i < n; i++) printf("%02x", b[i]); printf("\n"); } static void clear_uc(volatile uint8_t *b, int n) { int i; for (i = 0; i < n; i++) b[i] = 0; } static void copy_from_uc(uint8_t *dst, volatile uint8_t *src, int n) { int i; for (i = 0; i < n; i++) dst[i] = src[i]; } typedef struct { uint8_t ek[EK_BYTES] __attribute__((aligned(64))); uint8_t dk[DK_BYTES] __attribute__((aligned(64))); uint8_t ct[CT_BYTES] __attribute__((aligned(64))); uint8_t ss_enc[32]; uint8_t ss_dec[32]; } mlkem_result_t; static mlkem_result_t cpu_res; static mlkem_result_t dma_res; static mlkem_result_t sw_res; static void print_hex(const char *label, const uint8_t *b, int n) { int i; printf("%s", label); for (i = 0; i < n; i++) printf("%02x", b[i]); printf("\n"); } static int verify_result(const mlkem_result_t *res) { int i, mismatch; mismatch = 0; for (i = 0; i < 32; i++) if (res->ss_enc[i] != res->ss_dec[i]) mismatch++; return mismatch; } static void clear_result_cpu(mlkem_result_t *res) { memset(res, 0, sizeof(*res)); } static void clear_result_dma(mlkem_result_t *res) { volatile uint8_t *p = (volatile uint8_t *)UNCACHED_ADDR(res); clear_uc(p, sizeof(*res)); } static unsigned long run_roundtrip_sw(const uint8_t *seed_d, const uint8_t *seed_z, const uint8_t *msg, mlkem_result_t *res, int *status) { unsigned long t0, t1; uint8_t coins[2 * MLKEM_SYMBYTES]; int rc; clear_result_cpu(res); memcpy(coins, seed_d, MLKEM_SYMBYTES); memcpy(coins + MLKEM_SYMBYTES, seed_z, MLKEM_SYMBYTES); t0 = get_clock_count(); rc = crypto_kem_keypair_derand(res->ek, res->dk, coins); if (rc == 0) rc = crypto_kem_enc_derand(res->ct, res->ss_enc, res->ek, msg); if (rc == 0) rc = crypto_kem_dec(res->ss_dec, res->ct, res->dk); t1 = get_clock_count(); *status = rc; return t1 - t0; } static unsigned long run_roundtrip_cpu(const uint8_t *seed_d, const uint8_t *seed_z, const uint8_t *msg, mlkem_result_t *res) { unsigned long t0, t1; clear_result_cpu(res); t0 = get_clock_count(); write_block32(MLKEM_D, seed_d); write_block32(MLKEM_Z, seed_z); run_op(OP_KEYGEN); readback(RDB_EK, res->ek, EK_BYTES); readback(RDB_DK, res->dk, DK_BYTES); preload(PRE_EK, res->ek, EK_BYTES); write_block32(MLKEM_MSG, msg); run_op(OP_ENCAPS); read_block32(MLKEM_SS, res->ss_enc); readback(RDB_CT, res->ct, CT_BYTES); preload(PRE_DK, res->dk, DK_BYTES); preload(PRE_C, res->ct, CT_BYTES); run_op(OP_DECAPS); read_block32(MLKEM_SS, res->ss_dec); t1 = get_clock_count(); return t1 - t0; } static unsigned long run_roundtrip_dma(const uint8_t *seed_d, const uint8_t *seed_z, const uint8_t *msg, mlkem_result_t *res) { unsigned long t0, t1; volatile uint8_t *ek_uc = uncached_u8(res->ek); volatile uint8_t *dk_uc = uncached_u8(res->dk); volatile uint8_t *ct_uc = uncached_u8(res->ct); clear_result_dma(res); t0 = get_clock_count(); write_block32(MLKEM_D, seed_d); write_block32(MLKEM_Z, seed_z); run_op(OP_KEYGEN); dma_from_mlkem(ek_uc, MLKEM_EK_WINDOW, EK_BYTES); dma_from_mlkem(dk_uc, MLKEM_DK_WINDOW, DK_BYTES); dma_to_mlkem(MLKEM_EK_WINDOW, ek_uc, EK_BYTES); write_block32(MLKEM_MSG, msg); run_op(OP_ENCAPS); read_block32(MLKEM_SS, res->ss_enc); dma_from_mlkem(ct_uc, MLKEM_CT_WINDOW, CT_BYTES); dma_to_mlkem(MLKEM_DK_WINDOW, dk_uc, DK_BYTES); dma_to_mlkem(MLKEM_CT_WINDOW, ct_uc, CT_BYTES); run_op(OP_DECAPS); read_block32(MLKEM_SS, res->ss_dec); t1 = get_clock_count(); copy_from_uc(res->ek, ek_uc, EK_BYTES); copy_from_uc(res->dk, dk_uc, DK_BYTES); copy_from_uc(res->ct, ct_uc, CT_BYTES); return t1 - t0; } static int compare_results(const mlkem_result_t *a, const mlkem_result_t *b) { return memcmp(a->ek, b->ek, EK_BYTES) || memcmp(a->dk, b->dk, DK_BYTES) || memcmp(a->ct, b->ct, CT_BYTES) || memcmp(a->ss_enc, b->ss_enc, sizeof(a->ss_enc)) || memcmp(a->ss_dec, b->ss_dec, sizeof(a->ss_dec)); } static unsigned long cycles_to_us(unsigned long cycles) { return cycles / (CORE_CLOCKS_PER_SEC / USEC_PER_SEC); } int main(void) { int mismatch_sw, mismatch_cpu, mismatch_dma, sw_status, cpu_vs_dma_diff; int sw_vs_cpu_diff, sw_vs_dma_diff; unsigned long sw_cycles, cpu_cycles, dma_cycles; printf("MLKEM_START\n"); generate_prng_inputs(); sw_cycles = run_roundtrip_sw(SEED_D, SEED_Z, MSG_M, &sw_res, &sw_status); cpu_cycles = run_roundtrip_cpu(SEED_D, SEED_Z, MSG_M, &cpu_res); dma_cycles = run_roundtrip_dma(SEED_D, SEED_Z, MSG_M, &dma_res); mismatch_sw = verify_result(&sw_res); mismatch_cpu = verify_result(&cpu_res); mismatch_dma = verify_result(&dma_res); cpu_vs_dma_diff = compare_results(&cpu_res, &dma_res); sw_vs_cpu_diff = compare_results(&sw_res, &cpu_res); sw_vs_dma_diff = compare_results(&sw_res, &dma_res); print_hex("SW_EK32=", sw_res.ek, 32); print_hex("SW_SS_ENC=", sw_res.ss_enc, 32); print_hex("SW_SS_DEC=", sw_res.ss_dec, 32); print_hex("CPU_EK32=", cpu_res.ek, 32); print_hex("CPU_SS_ENC=", cpu_res.ss_enc, 32); print_hex("CPU_SS_DEC=", cpu_res.ss_dec, 32); print_hex("DMA_EK32=", dma_res.ek, 32); print_hex("DMA_SS_ENC=", dma_res.ss_enc, 32); print_hex("DMA_SS_DEC=", dma_res.ss_dec, 32); printf("MLKEM_SW_STATUS=%d\n", sw_status); printf("MLKEM_SW_CYCLES=%lu\n", sw_cycles); printf("MLKEM_CPU_CYCLES=%lu\n", cpu_cycles); printf("MLKEM_DMA_CYCLES=%lu\n", dma_cycles); printf("MLKEM_SW_US=%lu\n", cycles_to_us(sw_cycles)); printf("MLKEM_CPU_US=%lu\n", cycles_to_us(cpu_cycles)); printf("MLKEM_DMA_US=%lu\n", cycles_to_us(dma_cycles)); if (cpu_cycles != 0) { unsigned long hw_speedup_x100 = (sw_cycles * 100u) / cpu_cycles; printf("MLKEM_HW_VS_SW_SPEEDUP_X100=%lu\n", hw_speedup_x100); } if (dma_cycles != 0) { unsigned long hw_dma_speedup_x100 = (sw_cycles * 100u) / dma_cycles; printf("MLKEM_HW_DMA_VS_SW_SPEEDUP_X100=%lu\n", hw_dma_speedup_x100); } if (dma_cycles != 0) { unsigned long speedup_x100 = (cpu_cycles * 100u) / dma_cycles; printf("MLKEM_DMA_SPEEDUP_X100=%lu\n", speedup_x100); } if (sw_status == 0 && mismatch_sw == 0 && mismatch_cpu == 0 && mismatch_dma == 0 && cpu_vs_dma_diff == 0 && sw_vs_cpu_diff == 0 && sw_vs_dma_diff == 0) printf("MLKEM_RESULT=PASS\n"); else { printf("MLKEM_RESULT=FAIL sw_status=%d sw_mismatch=%d cpu_mismatch=%d dma_mismatch=%d cpu_dma_diff=%d sw_cpu_diff=%d sw_dma_diff=%d\n", sw_status, mismatch_sw, mismatch_cpu, mismatch_dma, cpu_vs_dma_diff, sw_vs_cpu_diff, sw_vs_dma_diff); } printf("MLKEM_DONE\n"); return 0; }