如何在Go语言中使用BF16

在大语言模型流行的当下，低精度浮点数对于开发者们不再陌生，而其中BF16是其中支持范围最广的低精度浮点格式。本文将介绍如何在Go语言中使用BF16。

BF16介绍

+---------+--------+--------+
| 1符号位 | 8指数位 | 7尾数位 |
+---------+--------+--------+

BF16由1位符号位、 8位指数位和7位尾数位构成。和FP32相比，指数位相同，但尾数位大幅减少。因此BF16动态范围与FP32完全一致（约$±3.4×10³⁸$），牺牲部分精度（约$2^{-7}$）换取更大的指数范围。大指数范围可有效避免梯度溢出，适合深度学习中的反向传播。同时与FP32兼容性强，硬件设计更简单。BF16用于在模型训练和推理场景替代FP32，降低内存占用和带宽需求，提升SIMD指令一次处理数据的数量。

Go语言中的BF16

显然Go语言不会原生支持BF16，我们无法像使用单精度浮点或者双精度浮点一样使用BF16。在Go语言中使用BF16会面临两个问题：

1. 如何在不支持BF16的CPU上计算BF16？ 可以使用软件模拟BF16计算，例如github.com/chenxingqiang/go-floatx这个项目。软件模拟BF16计算效率低下，完美融合了BF16精度低缺点的和FP32浮点计算并行度低的缺点。如果CPU不支持硬件级别的BF16计算，还是老老实实使用FP32。
2. 如何在支持BF16的CPU上计算BF16？ 如果检测到CPU支持BF16或者确保程序运行的CPU支持BF16，那么就可以考虑使用BF16进行计算。本文剩余部分将介绍如何分别使用Cgo和汇编两种方法调用BF16计算指令。由于我们最容易获取的支持BF16的消费级CPU是来自于苹果的M3和之后的处理器，因此本文将以ARM下的BF16计算为例。

通过Cgo使用BF16

Cgo可以借助C语言生成计算BF16的指令，BF16通常需要批量计算，以下为BF16批量相加和转换的Cgo实现：

package cgo

/*
#cgo CFLAGS: -march=armv8.2-a+bf16 -O3
#include <arm_bf16.h>
#include <stdint.h>

void convert_float32_to_bf16(float *a, uint16_t *b, int n) {
	bfloat16_t *bf16_b = (bfloat16_t *)b;
	for (int i = 0; i < n; i++) {
		bf16_b[i] = (bfloat16_t)a[i];
	}
}

void convert_bf16_to_float32(uint16_t *a, float *b, int n) {
	bfloat16_t *bf16_a = (bfloat16_t *)a;
	for (int i = 0; i < n; i++) {
		b[i] = (float)bf16_a[i];
	}
}

void add_bf16(uint16_t *a, uint16_t *b, uint16_t *c, int n) {
	bfloat16_t *bf16_a = (bfloat16_t *)a;
	bfloat16_t *bf16_b = (bfloat16_t *)b;
	bfloat16_t *bf16_c = (bfloat16_t *)c;
	for (int i = 0; i < n; i++) {
		bf16_c[i] = bf16_a[i] + bf16_b[i];
	}
}
*/
import "C"

func ConvertFloat32ToBF16(a []float32) []uint16 {
	b := make([]uint16, len(a))
	C.convert_float32_to_bf16((*C.float)(&a[0]), (*C.uint16_t)(&b[0]), C.int(len(a)))
	return b
}

func ConvertBF16ToFloat32(a []uint16) []float32 {
	b := make([]float32, len(a))
	C.convert_bf16_to_float32((*C.uint16_t)(&a[0]), (*C.float)(&b[0]), C.int(len(a)))
	return b
}

func AddBF16(a, b, c []uint16) {
	if len(a) != len(b) || len(a) != len(c) {
		panic("slices must have the same length")
	}
	C.add_bf16((*C.uint16_t)(&a[0]), (*C.uint16_t)(&b[0]), (*C.uint16_t)(&c[0]), C.int(len(a)))
}

• convert_float32_to_bf16将FP32转换为BF16
• convert_bf16_to_float32将BF16转换为FP32
• add_bf16执行BF16相加

C和GO之间无法传递bfloat16_t类型的指针，因此需要通过uint16_t指针传递。

使用GoAT生成BF16汇编

Cgo给予Go语言更多扩展性的同时存在很多问题，其中包括存在性能问题（参见《cgo is not Go》一文），如果希望避免Cgo调用的开销，那么可以使用GoAT直接将C语言转换为Go汇编进行调用。

首先，将BF16向量相加和转换函数实现保存到src/bfloat16.c文件中：

#include <arm_bf16.h>

void convert_float32_to_bf16(float *a, void *b, long n) {
	bfloat16_t *bf16_b = (bfloat16_t *)b;
	for (long i = 0; i < n; i++) {
		bf16_b[i] = (bfloat16_t)a[i];
	}
}

void convert_bf16_to_float32(void *a, float *b, long n) {
	bfloat16_t *bf16_a = (bfloat16_t *)a;
	for (long i = 0; i < n; i++) {
		b[i] = (float)bf16_a[i];
	}
}

void add_bf16(void *a, void *b, void *result, long n) {
	bfloat16_t *bf16_a = (bfloat16_t *)a;
	bfloat16_t *bf16_b = (bfloat16_t *)b;
	bfloat16_t *bf16_c = (bfloat16_t *)result;
	for (int i = 0; i < n; i++) {
		bf16_c[i] = bf16_a[i] + bf16_b[i];
	}
}

然后，安装好GoAT之后，使用GoAT将C语言编译成Go汇编：

goat src/bfloat16.c -o . -O3 -march=armv8.2-a+bf16

编译得到两个文件，bfloat16.s和bfloat16.go：

• bfloat16.s为Go汇编文件
• bfloat16.go为函数声明

//go:build !noasm && arm64
// Code generated by GoAT. DO NOT EDIT.
// versions:
// 	clang   18.1.3 (1ubuntu1)
// 	objdump 2.42
// flags: -march=armv8.2-a+bf16 -O3
// source: bf16/asm_bf16/src/bfloat16.c

package asm_bf16

import "unsafe"

//go:noescape
func convert_float32_to_bf16(a, b unsafe.Pointer, n int64)

//go:noescape
func convert_bf16_to_float32(a, b unsafe.Pointer, n int64)

//go:noescape
func add_bf16(a, b, result unsafe.Pointer, n int64)

调用函数的方式就是将数组转换为unsafe.Pointer以及传入数组长度。

func AddBF16(a, b, c []uint16) {
	if len(a) != len(b) || len(a) != len(c) {
		panic("slices must have the same length")
	}
	add_bf16(unsafe.Pointer(&a[0]), unsafe.Pointer(&b[0]), unsafe.Pointer(&c[0]), int64(len(a)))
}

性能对比

本文在阿里云自研倚天710 ARM架构CPU上比较了Cgo和汇编实现的BF16向量加法的性能，同时也测试了FP32向量加法以作参考，测试代码已上传至GitHub。Cgo和GoAT使用了相同的clang版本，可以认为它们实际上调用的汇编代码是完全相同的。对8至1024维向量的相加的耗时对比如下：

BF16向量加法性能对比

图中可以找到以下两点有意思的发现：

1. BF16加法的Cgo实现和汇编实现的耗时之间始终存在固定时间的差距。 由于Cgo调用存在开销，Cgo实现始终慢于汇编实现。
2. BF16加法比FP32加法快一倍。 倚天710 CPU的SIMD指令一次能处理8个BF16，但是如果是FP32则只能处理4个，因此BF16加法的性能是FP32加法的两倍。

总结

Go语言并不擅长和计算机底层打交道，在Go语言中使用BF16并不是一个容易完成的任务。如果在使用Go语言的同时又不得不使用BF16等扩展指令集，相比使用Cgo，直接编写汇编代码或者使用GoAT编译生成汇编代码会有更好的性能。