👏
A64FX SVE での bitnet b1.58 実装
なんで 1.58 って中途半端なビット数なんや...
-1.0, 0, 1.0 の重みで 3 パターンで log2(3) ~= 1.58 でした.
重みが 1.58bit で, activation(input) も 1.58 というわけではありませんでした
(ただ基本はこちらも量子化前提の模様)
sve だと predicate で 0 を扱えるので便利そ?
A64FX 向けの BitNet b1.58 ternary GEMV ベンチマークカーネルを書きます。A64FX 向けの BitNet b1.58 ternary GEMV ベンチマークを作成しました。
核心部分の解説
FMA 方式のカーネル
// 重みを2つの predicate に展開
svbool_t p_nonzero; // 0以外の位置
svbool_t p_neg; // -1 の位置
// sign vector: -1.0 or +1.0
svfloat32_t sign_vec = svsel(p_neg, neg_one, one);
// 1回の masked FMA で完結！
acc = svmla_f32_m(p_nonzero, acc, x, sign_vec);
パッキング形式
2bit/weight: 00=0, 01=+1, 11=-1
A64FX 512bit = 16 lanes → 4 bytes で 16 weights
比較用に Add/Sub 方式も実装
svbool_t p_pos = svbic_z(pg, p_nonzero, p_neg); // +1 の位置
acc = svadd_f32_m(p_pos, acc, x); // +1 → 加算
acc = svsub_f32_m(p_neg, acc, x); // -1 → 減算
使い方
tar xzf bitnet_sve_benchmark.tar.gz
cd bitnet_sve
# Fugaku で
make fugaku
./bitnet_benchmark
# または QEMU で試す
make
qemu-aarch64 -cpu max,sve512=on ./bitnet_benchmark
期待される結果
- FP16 比でメモリ帯域 8x 削減 (16bit → 2bit)
- FMA パイプラインを活用できれば演算律速になる可能性
- 展開オーバーヘッドがどの程度かは実測次第
Fugaku で試してみてね！
#include "ternary_gemv.h"
#include <arm_sve.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
// ============================================
// Packing/Unpacking utilities
// ============================================
void pack_ternary_weights(const int8_t *weights, uint8_t *packed, size_t n) {
// Pack 4 ternary weights into 1 byte (2 bits each)
// Encoding: 00 = 0, 01 = +1, 11 = -1
size_t packed_idx = 0;
for (size_t i = 0; i < n; i += 4) {
uint8_t byte = 0;
for (int j = 0; j < 4 && (i + j) < n; j++) {
int8_t w = weights[i + j];
uint8_t encoded;
if (w == 0) {
encoded = 0b00;
} else if (w == 1) {
encoded = 0b01;
} else { // w == -1
encoded = 0b11;
}
byte |= (encoded << (j * 2));
}
packed[packed_idx++] = byte;
}
}
void unpack_ternary_weights(const uint8_t *packed, int8_t *weights, size_t n) {
size_t packed_idx = 0;
for (size_t i = 0; i < n; i += 4) {
uint8_t byte = packed[packed_idx++];
for (int j = 0; j < 4 && (i + j) < n; j++) {
uint8_t encoded = (byte >> (j * 2)) & 0b11;
if (encoded == 0b00) {
weights[i + j] = 0;
} else if (encoded == 0b01) {
weights[i + j] = 1;
} else { // 0b11
weights[i + j] = -1;
}
}
}
}
// ============================================
// Scalar baseline
// ============================================
void ternary_gemv_scalar(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N) {
int8_t *w_unpacked = (int8_t *)malloc(N * sizeof(int8_t));
for (int i = 0; i < M; i++) {
unpack_ternary_weights(&W_packed[(i * N) / 4], w_unpacked, N);
float acc = 0.0f;
for (int j = 0; j < N; j++) {
acc += (float)w_unpacked[j] * X[j];
}
Y[i] = acc;
}
free(w_unpacked);
}
// ============================================
// SVE helper: unpack 2-bit weights to predicates
// Returns predicates for 16 weights (A64FX: 512-bit = 16 x float32)
// ============================================
static inline void unpack_to_predicates_16(const uint8_t *packed_4bytes,
svbool_t pg,
svbool_t *p_nonzero,
svbool_t *p_neg) {
// packed_4bytes contains 16 weights (4 bytes = 32 bits = 16 x 2-bit)
uint32_t bits = *(const uint32_t *)packed_4bytes;
// Create index vector: 0, 1, 2, ..., 15
svuint32_t indices = svindex_u32(0, 1);
// Extract 2-bit encoding for each position
// shift = indices * 2
svuint32_t shift = svlsl_x(pg, indices, 1);
svuint32_t bits_vec = svdup_u32(bits);
// shifted = bits >> (indices * 2)
svuint32_t shifted = svlsr_x(pg, bits_vec, shift);
// Extract 2-bit value: masked = shifted & 0b11
svuint32_t masked = svand_x(pg, shifted, 0x3);
// non_zero: bit 0 is set (values 01 or 11)
*p_nonzero = svcmpne(pg, svand_x(pg, masked, 0x1), 0);
// negative: both bits set (value 11)
*p_neg = svcmpeq(pg, masked, 0x3);
}
// ============================================
// SVE FMA-based implementation
// Key idea: acc = fma(acc, x, sign) where sign = ±1.0
// Use predicate to mask zero weights
// ============================================
void ternary_gemv_sve_fma(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N) {
const svfloat32_t one = svdup_f32(1.0f);
const svfloat32_t neg_one = svdup_f32(-1.0f);
// A64FX: svcntw() = 16 (512-bit / 32-bit = 16 lanes)
const int vlen = svcntw();
for (int i = 0; i < M; i++) {
svfloat32_t acc = svdup_f32(0.0f);
const uint8_t *row_packed = &W_packed[(i * N) / 4];
int j = 0;
for (; j + vlen <= N; j += vlen) {
svbool_t pg = svptrue_b32();
// Unpack 16 weights (4 bytes) to predicates
svbool_t p_nonzero, p_neg;
unpack_to_predicates_16(&row_packed[j / 4], pg, &p_nonzero, &p_neg);
// Load activation vector
svfloat32_t x = svld1(pg, &X[j]);
// Build sign vector: -1.0 where p_neg, else +1.0
svfloat32_t sign_vec = svsel(p_neg, neg_one, one);
// FMA: acc += x * sign (masked by p_nonzero)
// svmla_f32_m: acc = p_nonzero ? (acc + x * sign) : acc
acc = svmla_f32_m(p_nonzero, acc, x, sign_vec);
}
// Handle remainder
if (j < N) {
svbool_t pg = svwhilelt_b32(j, N);
svbool_t p_nonzero, p_neg;
unpack_to_predicates_16(&row_packed[j / 4], pg, &p_nonzero, &p_neg);
// Combine with tail predicate
p_nonzero = svand_z(pg, p_nonzero, pg);
p_neg = svand_z(pg, p_neg, pg);
svfloat32_t x = svld1(pg, &X[j]);
svfloat32_t sign_vec = svsel(p_neg, neg_one, one);
acc = svmla_f32_m(p_nonzero, acc, x, sign_vec);
}
// Horizontal reduction
Y[i] = svaddv(svptrue_b32(), acc);
}
}
// ============================================
// SVE Add/Sub implementation (no multiply)
// Use svadd_m / svsub_m with predicates
// ============================================
void ternary_gemv_sve_addsub(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N) {
const int vlen = svcntw();
for (int i = 0; i < M; i++) {
svfloat32_t acc = svdup_f32(0.0f);
const uint8_t *row_packed = &W_packed[(i * N) / 4];
int j = 0;
for (; j + vlen <= N; j += vlen) {
svbool_t pg = svptrue_b32();
svbool_t p_nonzero, p_neg;
unpack_to_predicates_16(&row_packed[j / 4], pg, &p_nonzero, &p_neg);
// p_pos = nonzero AND NOT neg
svbool_t p_pos = svbic_z(pg, p_nonzero, p_neg);
svfloat32_t x = svld1(pg, &X[j]);
// acc += x where positive
acc = svadd_f32_m(p_pos, acc, x);
// acc -= x where negative
acc = svsub_f32_m(p_neg, acc, x);
}
// Handle remainder
if (j < N) {
svbool_t pg = svwhilelt_b32(j, N);
svbool_t p_nonzero, p_neg;
unpack_to_predicates_16(&row_packed[j / 4], pg, &p_nonzero, &p_neg);
p_nonzero = svand_z(pg, p_nonzero, pg);
p_neg = svand_z(pg, p_neg, pg);
svbool_t p_pos = svbic_z(pg, p_nonzero, p_neg);
svfloat32_t x = svld1(pg, &X[j]);
acc = svadd_f32_m(p_pos, acc, x);
acc = svsub_f32_m(p_neg, acc, x);
}
Y[i] = svaddv(svptrue_b32(), acc);
}
}
// ============================================
// FP16 baseline using SVE
// ============================================
void fp16_gemv_sve(const uint16_t *W_fp16,
const float *X,
float *Y,
int M, int N) {
const int vlen = svcntw();
for (int i = 0; i < M; i++) {
svfloat32_t acc = svdup_f32(0.0f);
const uint16_t *row = &W_fp16[i * N];
int j = 0;
for (; j + vlen <= N; j += vlen) {
svbool_t pg = svptrue_b32();
// Load FP16 weights and convert to FP32
svfloat16_t w_f16 = svld1_f16(svptrue_b16(),
(const float16_t *)&row[j]);
// Convert lower half to FP32
svfloat32_t w = svcvt_f32_f16_x(pg, w_f16);
svfloat32_t x = svld1(pg, &X[j]);
acc = svmla_f32_x(pg, acc, x, w);
}
// Handle remainder
if (j < N) {
svbool_t pg = svwhilelt_b32(j, N);
svbool_t pg16 = svwhilelt_b16(j, N);
svfloat16_t w_f16 = svld1_f16(pg16, (const float16_t *)&row[j]);
svfloat32_t w = svcvt_f32_f16_x(pg, w_f16);
svfloat32_t x = svld1(pg, &X[j]);
acc = svmla_f32_m(pg, acc, x, w);
}
Y[i] = svaddv(svptrue_b32(), acc);
}
}
// ============================================
// Utility functions
// ============================================
void init_random_ternary(int8_t *weights, size_t n, unsigned int seed) {
srand(seed);
for (size_t i = 0; i < n; i++) {
int r = rand() % 3;
weights[i] = (int8_t)(r - 1); // -1, 0, or +1
}
}
void init_random_float(float *vec, size_t n, unsigned int seed) {
srand(seed);
for (size_t i = 0; i < n; i++) {
vec[i] = ((float)rand() / RAND_MAX) * 2.0f - 1.0f; // [-1, 1]
}
}
void ternary_to_fp16(const int8_t *ternary, uint16_t *fp16, size_t n) {
// Simple conversion using union trick for FP16
for (size_t i = 0; i < n; i++) {
float f = (float)ternary[i];
// Use ARM FP16 conversion
__fp16 h = (__fp16)f;
fp16[i] = *(uint16_t *)&h;
}
}
float verify_results(const float *ref, const float *test, size_t n) {
float max_err = 0.0f;
for (size_t i = 0; i < n; i++) {
float err = fabsf(ref[i] - test[i]);
if (err > max_err) {
max_err = err;
}
}
return max_err;
}
double get_time_sec(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return ts.tv_sec + ts.tv_nsec * 1e-9;
}
#ifndef TERNARY_GEMV_H
#define TERNARY_GEMV_H
#include <stdint.h>
#include <stddef.h>
#ifdef __cplusplus
extern "C" {
#endif
// Weight encoding (2 bits per weight):
// bit 0: non_zero flag (1 = weight is ±1, 0 = weight is 0)
// bit 1: sign flag (0 = +1, 1 = -1, only valid if non_zero=1)
//
// Encoding: 00 = 0, 01 = +1, 11 = -1 (10 is unused)
// Pack ternary weights into 2-bit format
// weights: array of int8_t with values {-1, 0, +1}
// packed: output buffer (N/4 bytes for N weights)
void pack_ternary_weights(const int8_t *weights, uint8_t *packed, size_t n);
// Unpack for verification
void unpack_ternary_weights(const uint8_t *packed, int8_t *weights, size_t n);
// ============================================
// GEMV implementations: Y = W * X
// W: M x N ternary weights (packed)
// X: N-dim float32 vector
// Y: M-dim float32 output
// ============================================
// Baseline: scalar implementation
void ternary_gemv_scalar(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N);
// SVE FMA-based implementation (main focus)
void ternary_gemv_sve_fma(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N);
// SVE with predicate-only (add/sub, no multiply)
void ternary_gemv_sve_addsub(const uint8_t *W_packed,
const float *X,
float *Y,
int M, int N);
// FP16 baseline for comparison
void fp16_gemv_sve(const uint16_t *W_fp16,
const float *X,
float *Y,
int M, int N);
// ============================================
// Utility functions
// ============================================
// Initialize random ternary weights
void init_random_ternary(int8_t *weights, size_t n, unsigned int seed);
// Initialize random float vector
void init_random_float(float *vec, size_t n, unsigned int seed);
// Convert ternary to FP16 for baseline comparison
void ternary_to_fp16(const int8_t *ternary, uint16_t *fp16, size_t n);
// Verify results (returns max absolute error)
float verify_results(const float *ref, const float *test, size_t n);
// Get current time in seconds (for benchmarking)
double get_time_sec(void);
#ifdef __cplusplus
}
#endif
#endif // TERNARY_GEMV_H
Discussion