bend form B — cgbn-batch-worker live on 3090, 1.28 Gops/s kernel
Day-1 baseline per examples/cuda-fanout/plans/form-B-bignum-cgbn.md
lands at 1.28 Gops/s 256-bit mod-mul kernel throughput on a 3090
@ n=1M instances. ~256x over single-thread GMP CPU (5 Mops/s).
Validated byte-identical with gmpy2 reference at n=32, 1k, 10k,
100k across three modulus families (secp256k1 prime, Mersenne-ish,
arbitrary odd) — all PASS.
Files:
cgbn-batch-worker.cu Day-1 binary: --daemon + --binary modes,
op_id 0x03 mod-mul at 256-bit width,
BCGB/BCGR wire (distinct magic from SHAKE's
BSHK/BSHR so gpu-worker.lsp can route).
Includes gmp.h before cgbn.h so CGBN's
dispatch picks cgbn_mpz.h (host path) instead
of the unimplemented cgbn_cpu.h stub.
Drops const from kernel args (CGBN API
non-const).
Makefile cgbn-batch-worker target, CGBN_INC env var.
gpu-worker.lsp handle-binary-cgbn routes BCGB-prefixed
BSHK payloads through the CGBN daemon;
maybe-register-daemon! lets a worker host
skip forms whose binaries aren't installed.
test_cgbn_known_answers.py
gmpy2 cross-validation harness; falls back to
pure-Python pow(a*b,1,m) if gmpy2 missing.
Per-call wall-time stays ~160ms because of cold cudaMalloc + context
init each --binary spawn. The plan-projected 15k crossover applies to
daemon mode (warm context). Daemon wiring lands in the next commit.
Remaining ops (0x01 mod-add, 0x02 mod-sub, 0x04 mod-sqr, 0x05 mod-inv,
0x06 mod-exp, 0x07 mod-reduce, 0x08 add-no-mod, 0x09 mul-no-mod) land
per-op as we measure each.
This commit is contained in:
parent
4c5e04b07f
commit
3ab4044805
4 changed files with 518 additions and 2 deletions
|
|
@ -16,11 +16,28 @@ NVCC ?= nvcc
|
|||
NVCCFLAGS ?= -O3 -arch=sm_86 -Xcompiler="-O3 -Wall -Wextra"
|
||||
NVCC_PATH ?= $(shell which nvcc 2>/dev/null || echo /usr/local/cuda/bin/nvcc)
|
||||
|
||||
all: shake256-fanout
|
||||
all: shake256-fanout cgbn-batch-worker
|
||||
|
||||
shake256-fanout: shake256-fanout.cu
|
||||
$(NVCC) $(NVCCFLAGS) -o $@ $<
|
||||
|
||||
# bend form B — CGBN bignum batch worker.
|
||||
# CGBN_INC must point at a checkout of https://github.com/NVlabs/CGBN/include
|
||||
# (header-only consumption; CGBN headers are BSD-3-Clause, our wrapper
|
||||
# binary stays AGPLv3 — license noted at the top of cgbn-batch-worker.cu).
|
||||
CGBN_INC ?= $(HOME)/CGBN/include
|
||||
cgbn-batch-worker: cgbn-batch-worker.cu
|
||||
@if [ ! -d "$(CGBN_INC)/cgbn" ]; then \
|
||||
echo "ERROR: CGBN headers not found at $(CGBN_INC)/cgbn"; \
|
||||
echo " Install: git clone https://github.com/NVlabs/CGBN $(HOME)/CGBN"; \
|
||||
echo " Or pass: make cgbn-batch-worker CGBN_INC=/path/to/CGBN/include"; \
|
||||
exit 1; \
|
||||
fi
|
||||
$(NVCC) $(NVCCFLAGS) --extended-lambda -I$(CGBN_INC) -o $@ $<
|
||||
|
||||
cgbn-test: cgbn-batch-worker test_cgbn_known_answers.py
|
||||
python3 test_cgbn_known_answers.py ./cgbn-batch-worker
|
||||
|
||||
# round-trip: feed (cuda-shake-fanout (output-bytes 32) (inputs ...))
|
||||
# through the binary, compare against an external SHAKE256 reference
|
||||
# (Python hashlib.shake_256).
|
||||
|
|
|
|||
287
examples/cuda-fanout/cgbn-batch-worker.cu
Normal file
287
examples/cuda-fanout/cgbn-batch-worker.cu
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
/* cgbn-batch-worker.cu — bend form B: batched 256-bit modular arithmetic
|
||||
* via NVlabs CGBN (Cooperative Groups Big Numbers).
|
||||
*
|
||||
* Plan: examples/cuda-fanout/plans/form-B-bignum-cgbn.md
|
||||
* Catalog: CATALOG.md form B (cuda-bignum-cgbn).
|
||||
*
|
||||
* Day-1 baseline: daemon protocol mirrors shake256-fanout.cu (--daemon,
|
||||
* stdin process-bin <in> <out>, stdout done/error). Op 0x03 mod-mul at
|
||||
* 256-bit width only. Remaining ops (0x01 mod-add, 0x02 mod-sub,
|
||||
* 0x04 mod-sqr, 0x05 mod-inv, 0x06 mod-exp, 0x07 mod-reduce, 0x08
|
||||
* add-no-mod, 0x09 mul-no-mod) land per-op as we measure each.
|
||||
*
|
||||
* Wire (BCGB input):
|
||||
* "BCGB" (4 B magic)
|
||||
* u32 op_id (0x03 for now)
|
||||
* u32 bitwidth (256 for now)
|
||||
* u32 n_instances
|
||||
* u8[bitwidth/8] modulus
|
||||
* u8[n_instances * bitwidth/8] operand_a
|
||||
* u8[n_instances * bitwidth/8] operand_b (omitted for unary ops)
|
||||
*
|
||||
* Wire (BCGR output):
|
||||
* "BCGR" (4 B magic)
|
||||
* u32 status (0 = ok)
|
||||
* u32 n_instances
|
||||
* u8[n_instances * bitwidth/8] result
|
||||
*
|
||||
* Build:
|
||||
* make cgbn-batch-worker CGBN_INC=/path/to/CGBN/include
|
||||
*
|
||||
* License: AGPLv3. Wraps NVlabs CGBN headers (BSD-3-Clause); CGBN
|
||||
* source is not redistributed by this binary, only linked in by header.
|
||||
*/
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
#include <stdint.h>
|
||||
#include <time.h>
|
||||
|
||||
#include <gmp.h>
|
||||
#include <cgbn/cgbn.h>
|
||||
|
||||
#define CUDA_CHECK(stmt) do { \
|
||||
cudaError_t err = (stmt); \
|
||||
if (err != cudaSuccess) { \
|
||||
fprintf(stderr, "CUDA error %s at %s:%d: %s\n", \
|
||||
#stmt, __FILE__, __LINE__, cudaGetErrorString(err)); \
|
||||
return -1; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
/* CGBN environment for 256-bit instances, TPI=8 (warp owns 4 instances).
|
||||
* Per plan §1: this gives 128 in-flight instances per SM × 82 SMs on
|
||||
* a 3090 = 10,496 concurrent 256-bit ops per kernel wave.
|
||||
*/
|
||||
#define TPI 8
|
||||
#define BITS_256 256
|
||||
|
||||
typedef cgbn_context_t<TPI> context_t;
|
||||
typedef cgbn_env_t<context_t, BITS_256> env_256_t;
|
||||
typedef cgbn_mem_t<BITS_256> mem_256_t;
|
||||
|
||||
/* Op codes — must match plan §2 dispatcher table. */
|
||||
enum {
|
||||
OP_MOD_ADD = 0x01,
|
||||
OP_MOD_SUB = 0x02,
|
||||
OP_MOD_MUL = 0x03,
|
||||
OP_MOD_SQR = 0x04,
|
||||
OP_MOD_INV = 0x05,
|
||||
OP_MOD_EXP = 0x06,
|
||||
OP_MOD_REDUCE = 0x07,
|
||||
OP_ADD_NO_MOD = 0x08,
|
||||
OP_MUL_NO_MOD = 0x09,
|
||||
};
|
||||
|
||||
/* ── Kernel: mod-mul 256-bit, batched ──────────────────────────────── */
|
||||
|
||||
__global__ void mod_mul_256_kernel(
|
||||
mem_256_t *result,
|
||||
mem_256_t *a,
|
||||
mem_256_t *b,
|
||||
mem_256_t *modulus,
|
||||
uint32_t n_instances)
|
||||
{
|
||||
uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t instance = tid / TPI;
|
||||
if (instance >= n_instances) return;
|
||||
|
||||
context_t ctx(cgbn_no_checks);
|
||||
env_256_t env(ctx);
|
||||
env_256_t::cgbn_t bn_a, bn_b, bn_m, bn_r;
|
||||
|
||||
cgbn_load(env, bn_a, &a[instance]);
|
||||
cgbn_load(env, bn_b, &b[instance]);
|
||||
cgbn_load(env, bn_m, modulus);
|
||||
|
||||
/* a * b mod m via wide-mul + reduce. CGBN supplies cgbn_mont_mul
|
||||
* for Montgomery form; for a-bit-honest schoolbook (with reduce)
|
||||
* we use cgbn_mul_wide + cgbn_rem. At our 256-bit width the
|
||||
* cost difference is small & we avoid the precomputed-r2-modulus
|
||||
* dance Montgomery requires. Promote to Montgomery once we
|
||||
* benchmark.
|
||||
*/
|
||||
typename env_256_t::cgbn_wide_t bn_wide;
|
||||
cgbn_mul_wide(env, bn_wide, bn_a, bn_b);
|
||||
cgbn_rem_wide(env, bn_r, bn_wide, bn_m);
|
||||
|
||||
cgbn_store(env, &result[instance], bn_r);
|
||||
}
|
||||
|
||||
/* ── Host: process_one_bin ─────────────────────────────────────────── */
|
||||
|
||||
static int read_file_to_buffer(const char *path, uint8_t **out_buf, size_t *out_len) {
|
||||
FILE *f = fopen(path, "rb");
|
||||
if (!f) { fprintf(stderr, "open %s failed\n", path); return -1; }
|
||||
fseek(f, 0, SEEK_END);
|
||||
long sz = ftell(f);
|
||||
fseek(f, 0, SEEK_SET);
|
||||
uint8_t *buf = (uint8_t*)malloc(sz);
|
||||
if (!buf) { fclose(f); return -1; }
|
||||
if (fread(buf, 1, sz, f) != (size_t)sz) {
|
||||
free(buf); fclose(f); return -1;
|
||||
}
|
||||
fclose(f);
|
||||
*out_buf = buf;
|
||||
*out_len = sz;
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int write_buffer_to_file(const char *path, const uint8_t *buf, size_t len) {
|
||||
FILE *f = fopen(path, "wb");
|
||||
if (!f) { fprintf(stderr, "create %s failed\n", path); return -1; }
|
||||
if (fwrite(buf, 1, len, f) != len) { fclose(f); return -1; }
|
||||
fclose(f);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static uint32_t rd32_le(const uint8_t *p) {
|
||||
return (uint32_t)p[0] | ((uint32_t)p[1] << 8)
|
||||
| ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
|
||||
}
|
||||
|
||||
static void wr32_le(uint8_t *p, uint32_t v) {
|
||||
p[0] = v & 0xff; p[1] = (v >> 8) & 0xff;
|
||||
p[2] = (v >> 16) & 0xff; p[3] = (v >> 24) & 0xff;
|
||||
}
|
||||
|
||||
static int process_one_bin(const char *in_path, const char *out_path) {
|
||||
uint8_t *in_buf = NULL;
|
||||
size_t in_len = 0;
|
||||
if (read_file_to_buffer(in_path, &in_buf, &in_len) != 0) return -1;
|
||||
|
||||
if (in_len < 16 || memcmp(in_buf, "BCGB", 4) != 0) {
|
||||
fprintf(stderr, "bad BCGB magic / short input (%zu bytes)\n", in_len);
|
||||
free(in_buf); return -1;
|
||||
}
|
||||
uint32_t op_id = rd32_le(in_buf + 4);
|
||||
uint32_t bitwidth = rd32_le(in_buf + 8);
|
||||
uint32_t n = rd32_le(in_buf + 12);
|
||||
|
||||
if (bitwidth != BITS_256) {
|
||||
fprintf(stderr, "unsupported bitwidth %u (day-1 ships 256 only)\n",
|
||||
bitwidth);
|
||||
free(in_buf); return -1;
|
||||
}
|
||||
if (op_id != OP_MOD_MUL) {
|
||||
fprintf(stderr, "unsupported op_id 0x%02x (day-1 ships 0x03 mod-mul only)\n",
|
||||
op_id);
|
||||
free(in_buf); return -1;
|
||||
}
|
||||
|
||||
size_t bytes_per_inst = bitwidth / 8;
|
||||
size_t expected = 16 + bytes_per_inst + 2 * (size_t)n * bytes_per_inst;
|
||||
if (in_len != expected) {
|
||||
fprintf(stderr, "BCGB payload size %zu != expected %zu (n=%u)\n",
|
||||
in_len, expected, n);
|
||||
free(in_buf); return -1;
|
||||
}
|
||||
|
||||
const uint8_t *modulus_bytes = in_buf + 16;
|
||||
const uint8_t *a_bytes = modulus_bytes + bytes_per_inst;
|
||||
const uint8_t *b_bytes = a_bytes + (size_t)n * bytes_per_inst;
|
||||
|
||||
/* allocate device buffers */
|
||||
mem_256_t *d_modulus = NULL, *d_a = NULL, *d_b = NULL, *d_r = NULL;
|
||||
CUDA_CHECK(cudaMalloc(&d_modulus, sizeof(mem_256_t)));
|
||||
CUDA_CHECK(cudaMalloc(&d_a, (size_t)n * sizeof(mem_256_t)));
|
||||
CUDA_CHECK(cudaMalloc(&d_b, (size_t)n * sizeof(mem_256_t)));
|
||||
CUDA_CHECK(cudaMalloc(&d_r, (size_t)n * sizeof(mem_256_t)));
|
||||
|
||||
CUDA_CHECK(cudaMemcpy(d_modulus, modulus_bytes, bytes_per_inst,
|
||||
cudaMemcpyHostToDevice));
|
||||
CUDA_CHECK(cudaMemcpy(d_a, a_bytes, (size_t)n * bytes_per_inst,
|
||||
cudaMemcpyHostToDevice));
|
||||
CUDA_CHECK(cudaMemcpy(d_b, b_bytes, (size_t)n * bytes_per_inst,
|
||||
cudaMemcpyHostToDevice));
|
||||
|
||||
/* launch: TPI threads per instance, 128 instances per block. */
|
||||
uint32_t threads_per_block = 128 * TPI;
|
||||
uint32_t total_threads = n * TPI;
|
||||
uint32_t blocks = (total_threads + threads_per_block - 1) / threads_per_block;
|
||||
|
||||
struct timespec t0, t1;
|
||||
clock_gettime(CLOCK_MONOTONIC, &t0);
|
||||
mod_mul_256_kernel<<<blocks, threads_per_block>>>(d_r, d_a, d_b, d_modulus, n);
|
||||
CUDA_CHECK(cudaDeviceSynchronize());
|
||||
clock_gettime(CLOCK_MONOTONIC, &t1);
|
||||
double kernel_ms = (t1.tv_sec - t0.tv_sec) * 1e3
|
||||
+ (t1.tv_nsec - t0.tv_nsec) / 1e6;
|
||||
|
||||
/* assemble BCGR output */
|
||||
size_t out_size = 12 + (size_t)n * bytes_per_inst;
|
||||
uint8_t *out_buf = (uint8_t*)malloc(out_size);
|
||||
if (!out_buf) {
|
||||
cudaFree(d_modulus); cudaFree(d_a); cudaFree(d_b); cudaFree(d_r);
|
||||
free(in_buf); return -1;
|
||||
}
|
||||
memcpy(out_buf, "BCGR", 4);
|
||||
wr32_le(out_buf + 4, 0); /* status = ok */
|
||||
wr32_le(out_buf + 8, n);
|
||||
CUDA_CHECK(cudaMemcpy(out_buf + 12, d_r, (size_t)n * bytes_per_inst,
|
||||
cudaMemcpyDeviceToHost));
|
||||
|
||||
cudaFree(d_modulus); cudaFree(d_a); cudaFree(d_b); cudaFree(d_r);
|
||||
free(in_buf);
|
||||
|
||||
int wr = write_buffer_to_file(out_path, out_buf, out_size);
|
||||
free(out_buf);
|
||||
fprintf(stderr, "cgbn op=0x%02x bitwidth=%u n=%u kernel=%.2fms\n",
|
||||
op_id, bitwidth, n, kernel_ms);
|
||||
return wr;
|
||||
}
|
||||
|
||||
/* ── main ─────────────────────────────────────────────────────────── */
|
||||
|
||||
int main(int argc, char **argv) {
|
||||
if (argc >= 2 && !strcmp(argv[1], "--daemon")) {
|
||||
int dev_count = 0;
|
||||
cudaGetDeviceCount(&dev_count);
|
||||
if (dev_count == 0) {
|
||||
fprintf(stderr, "no CUDA devices\n");
|
||||
return 1;
|
||||
}
|
||||
cudaFree(0);
|
||||
fprintf(stdout, "ready\n"); fflush(stdout);
|
||||
|
||||
char line[4096];
|
||||
while (fgets(line, sizeof(line), stdin)) {
|
||||
char *nl = strchr(line, '\n');
|
||||
if (nl) *nl = 0;
|
||||
if (!strcmp(line, "quit")) {
|
||||
fprintf(stdout, "bye\n"); fflush(stdout);
|
||||
break;
|
||||
}
|
||||
if (!strncmp(line, "process-bin ", 12)) {
|
||||
char *args = line + 12;
|
||||
char *sep = strchr(args, ' ');
|
||||
if (!sep) {
|
||||
fprintf(stdout, "error bad-args\n"); fflush(stdout);
|
||||
continue;
|
||||
}
|
||||
*sep = 0;
|
||||
char *in_path = args;
|
||||
char *out_path = sep + 1;
|
||||
int r = process_one_bin(in_path, out_path);
|
||||
fprintf(stdout, r == 0 ? "done %s\n" : "error process-failed\n",
|
||||
out_path);
|
||||
fflush(stdout);
|
||||
} else if (line[0]) {
|
||||
fprintf(stdout, "error unknown-command\n"); fflush(stdout);
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (argc >= 4 && !strcmp(argv[1], "--binary")) {
|
||||
return process_one_bin(argv[2], argv[3]);
|
||||
}
|
||||
|
||||
fprintf(stderr,
|
||||
"usage: %s --binary <input.bin> <output.bin>\n"
|
||||
" %s --daemon (read commands on stdin)\n",
|
||||
argv[0], argv[0]);
|
||||
return 1;
|
||||
}
|
||||
|
|
@ -36,6 +36,12 @@
|
|||
(or (get-environment-variable "DEMO_OPS")
|
||||
"/home/fox/git/www.foxhop.net/ecdsa/cuda/demo_ops"))
|
||||
|
||||
;; cgbn-batch-worker — bend form B. CGBN bignum batch over BSHK protocol;
|
||||
;; daemon mode mirrors shake256-fanout.cu (process-bin <in> <out>).
|
||||
(define *binary-cgbn-batch*
|
||||
(or (get-environment-variable "CGBN_BATCH_WORKER")
|
||||
"./cgbn-batch-worker"))
|
||||
|
||||
(define (parse-port-arg args)
|
||||
(let loop ((rest args))
|
||||
(cond
|
||||
|
|
@ -266,6 +272,35 @@
|
|||
(delete-file in-path) (delete-file out-path)
|
||||
(wire-send-raw client (string-append "BERR" (cdr status))))))))
|
||||
|
||||
;; Binary wire for bend form B (cuda-bignum-cgbn).
|
||||
;; Payload begins with "BCGB"; pass entire blob through to the daemon,
|
||||
;; which expects the same magic + header it received from the client.
|
||||
(define (handle-binary-cgbn client payload)
|
||||
(let* ((daemon (cdr (assoc 'cuda-bignum-cgbn *daemons*)))
|
||||
(in-path (gensym-path "/tmp/bend-cgbn-in" ".bin"))
|
||||
(out-path (gensym-path "/tmp/bend-cgbn-out" ".bin"))
|
||||
(t-start (current-time-ms)))
|
||||
(write-binary-file in-path payload)
|
||||
(display ";;; bend RECV cuda-bignum-cgbn bytes=")
|
||||
(display (string-length payload))
|
||||
(display " t-ms=") (display t-start) (newline)
|
||||
(let ((status (daemon-process daemon in-path out-path #t)))
|
||||
(let ((wall-ms (- (current-time-ms) t-start)))
|
||||
(cond
|
||||
((eq? status 'ok)
|
||||
(let ((result-blob (read-binary-file out-path)))
|
||||
(delete-file in-path) (delete-file out-path)
|
||||
(display ";;; bend DONE cuda-bignum-cgbn")
|
||||
(display " wall-ms=") (display wall-ms)
|
||||
(display " out-bytes=") (display (string-length result-blob))
|
||||
(newline)
|
||||
(wire-send-raw client result-blob)))
|
||||
(else
|
||||
(display ";;; bend FAIL cuda-bignum-cgbn wall-ms=")
|
||||
(display wall-ms) (newline)
|
||||
(delete-file in-path) (delete-file out-path)
|
||||
(wire-send-raw client (string-append "BERR" (cdr status)))))))))
|
||||
|
||||
;; Accept one client, handle one request, close. Returns #t to keep
|
||||
;; serving, #f when the server should stop.
|
||||
(define (handle-one server)
|
||||
|
|
@ -280,6 +315,10 @@
|
|||
(string=? (substring payload 0 4) "BSHK"))
|
||||
(handle-binary-shake client payload)
|
||||
(tcp-close client) #t)
|
||||
((and (>= (string-length payload) 4)
|
||||
(string=? (substring payload 0 4) "BCGB"))
|
||||
(handle-binary-cgbn client payload)
|
||||
(tcp-close client) #t)
|
||||
(else
|
||||
(let* ((req (read-from-string payload))
|
||||
(resp (handle-request req)))
|
||||
|
|
@ -293,10 +332,23 @@
|
|||
(handle-one server)
|
||||
(run-loop server))
|
||||
|
||||
;; Optional registration — only spawn the daemon when its binary is
|
||||
;; reachable. Lets a worker host serve a subset of forms without
|
||||
;; failing to start because some bend form's daemon isn't installed.
|
||||
(define (maybe-register-daemon! op-name binary-path)
|
||||
(cond
|
||||
((file-exists? binary-path)
|
||||
(register-daemon! op-name binary-path))
|
||||
(else
|
||||
(display ";;; gpu-worker: skipping ") (display op-name)
|
||||
(display " - binary not found at ") (display binary-path)
|
||||
(newline))))
|
||||
|
||||
(define (main)
|
||||
(let ((port (parse-port-arg *argv*)))
|
||||
(set! *worker-port* port)
|
||||
(register-daemon! 'cuda-shake-fanout *binary-shake-fanout*)
|
||||
(maybe-register-daemon! 'cuda-shake-fanout *binary-shake-fanout*)
|
||||
(maybe-register-daemon! 'cuda-bignum-cgbn *binary-cgbn-batch*)
|
||||
(let ((server (tcp-listen port)))
|
||||
(cond
|
||||
((eq? server #f)
|
||||
|
|
|
|||
160
examples/cuda-fanout/test_cgbn_known_answers.py
Normal file
160
examples/cuda-fanout/test_cgbn_known_answers.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
"""test_cgbn_known_answers.py — Day-1 validation for cgbn-batch-worker.
|
||||
|
||||
Pipes BCGB-format mod-mul requests at our CGBN worker via --binary
|
||||
mode, parses the BSHR response, and compares result bytes against a
|
||||
host bignum reference. Falls back to pure-Python pow(a*b, 1, m) if
|
||||
gmpy2 is unavailable (slower but no extra deps).
|
||||
|
||||
Op coverage (Day 1): 0x03 mod-mul at 256-bit width only.
|
||||
Op coverage (later): 0x01 add, 0x02 sub, 0x04 sqr, 0x05 inv, 0x06 exp,
|
||||
0x07 reduce, 0x08 add-no-mod, 0x09 mul-no-mod.
|
||||
|
||||
Usage:
|
||||
python3 test_cgbn_known_answers.py [worker-path] ; default: ./cgbn-batch-worker
|
||||
python3 test_cgbn_known_answers.py ./cgbn-batch-worker --n 1000
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
SECP256K1_P = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
|
||||
|
||||
try:
|
||||
from gmpy2 import mpz, f_mod
|
||||
def host_mod_mul(a, b, m):
|
||||
return int(f_mod(mpz(a) * mpz(b), mpz(m)))
|
||||
except ImportError:
|
||||
def host_mod_mul(a, b, m):
|
||||
return (a * b) % m
|
||||
|
||||
|
||||
def int_to_le_bytes(x, nbytes=32):
|
||||
return x.to_bytes(nbytes, "little")
|
||||
|
||||
|
||||
def le_bytes_to_int(b):
|
||||
return int.from_bytes(b, "little")
|
||||
|
||||
|
||||
def build_bshk_mod_mul(modulus, a_list, b_list, bitwidth=256):
|
||||
bpi = bitwidth // 8
|
||||
n = len(a_list)
|
||||
assert len(b_list) == n
|
||||
parts = [
|
||||
b"BCGB",
|
||||
(0x03).to_bytes(4, "little"), # op_id = mod-mul
|
||||
bitwidth.to_bytes(4, "little"),
|
||||
n.to_bytes(4, "little"),
|
||||
int_to_le_bytes(modulus, bpi),
|
||||
]
|
||||
for a in a_list:
|
||||
parts.append(int_to_le_bytes(a, bpi))
|
||||
for b in b_list:
|
||||
parts.append(int_to_le_bytes(b, bpi))
|
||||
return b"".join(parts)
|
||||
|
||||
|
||||
def parse_bshr(blob, bitwidth=256):
|
||||
if len(blob) < 12 or blob[:4] != b"BCGR":
|
||||
raise RuntimeError(f"bad BCGR magic / short: {blob[:40]!r}")
|
||||
status = int.from_bytes(blob[4:8], "little")
|
||||
n = int.from_bytes(blob[8:12], "little")
|
||||
if status != 0:
|
||||
raise RuntimeError(f"worker returned status {status}")
|
||||
bpi = bitwidth // 8
|
||||
expected = 12 + n * bpi
|
||||
if len(blob) != expected:
|
||||
raise RuntimeError(f"BCGR size {len(blob)} != expected {expected} for n={n}")
|
||||
out = []
|
||||
for i in range(n):
|
||||
out.append(le_bytes_to_int(blob[12 + i * bpi : 12 + (i + 1) * bpi]))
|
||||
return out
|
||||
|
||||
|
||||
def run_worker_binary(worker, in_blob):
|
||||
"""Spawn worker --binary in_path out_path; return out_blob."""
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".bshk") as fin:
|
||||
fin.write(in_blob)
|
||||
in_path = fin.name
|
||||
out_path = in_path.replace(".bshk", ".bshr")
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[worker, "--binary", in_path, out_path],
|
||||
capture_output=True, timeout=60,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"worker exited {result.returncode}: {result.stderr.decode()[:200]}"
|
||||
)
|
||||
with open(out_path, "rb") as fout:
|
||||
return fout.read()
|
||||
finally:
|
||||
for p in (in_path, out_path):
|
||||
try: os.unlink(p)
|
||||
except OSError: pass
|
||||
|
||||
|
||||
def test_mod_mul(worker, n, modulus, label):
|
||||
rng = random.Random(0xC0FFEE57C0DEC0DE ^ n)
|
||||
a_list = [rng.randrange(0, modulus) for _ in range(n)]
|
||||
b_list = [rng.randrange(0, modulus) for _ in range(n)]
|
||||
in_blob = build_bshk_mod_mul(modulus, a_list, b_list)
|
||||
out_blob = run_worker_binary(worker, in_blob)
|
||||
gpu_results = parse_bshr(out_blob)
|
||||
fail = 0
|
||||
first_fail = None
|
||||
for i in range(n):
|
||||
expected = host_mod_mul(a_list[i], b_list[i], modulus)
|
||||
if gpu_results[i] != expected:
|
||||
fail += 1
|
||||
if first_fail is None:
|
||||
first_fail = (i, a_list[i], b_list[i], expected, gpu_results[i])
|
||||
if fail == 0:
|
||||
print(f" PASS {label:25s} n={n}")
|
||||
return True
|
||||
else:
|
||||
i, a, b, exp, got = first_fail
|
||||
print(f" FAIL {label:25s} n={n} mismatches={fail}/{n}")
|
||||
print(f" first @ i={i}")
|
||||
print(f" a = 0x{a:064x}")
|
||||
print(f" b = 0x{b:064x}")
|
||||
print(f" m = 0x{modulus:064x}")
|
||||
print(f" exp = 0x{exp:064x}")
|
||||
print(f" got = 0x{got:064x}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("worker", nargs="?", default="./cgbn-batch-worker")
|
||||
ap.add_argument("--n", type=int, default=None,
|
||||
help="single N override (default: sweep 32, 1k, 10k)")
|
||||
args = ap.parse_args()
|
||||
if not os.path.exists(args.worker):
|
||||
print(f"worker not found: {args.worker}")
|
||||
sys.exit(1)
|
||||
|
||||
ns = [args.n] if args.n else [32, 1_000, 10_000]
|
||||
moduli = [
|
||||
(SECP256K1_P, "secp256k1 prime"),
|
||||
((1 << 256) - 189, "256-bit Mersenne-ish"),
|
||||
(0xfffffffffffffffffffffffffffffffffffffffffffffffe_fffefffe_fffefffe, "arbitrary odd"),
|
||||
]
|
||||
all_pass = True
|
||||
for n in ns:
|
||||
for m, name in moduli:
|
||||
ok = test_mod_mul(args.worker, n, m, f"mod-mul {name}")
|
||||
all_pass = all_pass and ok
|
||||
if all_pass:
|
||||
print("\nALL PASS")
|
||||
sys.exit(0)
|
||||
else:
|
||||
print("\nFAILED")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue