bfloat16に公式対応していないx86/64CPUでbfloat16を使う方法

はじめに

OpenCV Advent Calendar 2023 15日目の記事として投稿します．
他の記事は目次からご覧ください．

@tomoaki_teshima(Aki Teshima) さんのOpenCV bfloat16のサポート始めるってよの付録的な内容です．
こっちはOpenCV専用というよりもx86/64 CPU専用な話です．ユニバーサルな話になってません．

fp16やbfloat16は，16ビットで浮動小数点を表現可能なある形式です．
通常32ビットあるfloat型の精度を半分のビットで表し，符号部，指数部、仮数部の位置がそれぞれ異なります．詳細はリンク先のウィキペディアに任せます．

これらの半浮動小数点のメリットは２つあります．

1. データ帯域の半減
2. 演算速度の向上

1については単純で，データが半分になるため，データＩＯの量が半減します．
2については専用回路がないと意味がありません．

fp16についてはfloat型からの型変換だけならAVX（not AVX2）が使えるCPUからサポートしています．

fp16からfp32の変換_mm256_cvtph_ps（vcvtph2ps）: レイテンシ７スループット１
fp32からfp16の変換_mm256_cvtps_ph（vcvtps2ph）: レイテンシ７スループット１

2の演算速度の向上では，fp16の半浮動小数点用計算演算はXeon第4世代のSapphire Rapidsのみ対応しており，float型の倍の速度で動作しますが，普通の人々（）はまだ使うことはありません．

bfloat16はIntel Cooperlake(第3世代のXeon（2020年ごろ）やAMD Zen4で型変換と，内積演算（_mm_dpbf16_ps）だけサポートしていますが，汎用の計算命令はサポートしていません．
このbfloat16という型は，ラウンディングを気にしなければ，floatの下位16ビットを切り捨てて詰めるだけで型変換可能です．
この記事では，この型変換で作ったbfloat16の性能検証していきましょう．

floatからのbfloat16への変換

まず，float型からbfloat型へ変換するintrinsicを示します．
下位ビットだけ切り捨てて詰めればいいだけなのですが，AVX2だと適切な命令があまりないので少し冗長です．
基本的には16ビットずらして，ずらさないものと偶数奇数でインターリーブした後にshuffle命令で並べ返してもとに戻しています．
(AVX512なら直接１命令で書けます)．
また，floatを2つ同時に処理したらレジスタ幅がちょうどAVX2の幅になるため若干効率化します．

//__m128i cvtpspbhmask128 = _mm_setr_epi8(0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15);
static const __m128i cvtpspbhmask128 = _mm_setr_epi64x(940136352262127872, 1084816697938281218);
static inline __m128i _mm256_cvtps_bph(__m256 a)
{
	__m128i b = _mm_srli_epi32(_mm_castps_si128(_mm256_castps256_ps128(a)), 16);
	b = _mm_blend_epi16(b, _mm_castps_si128(_mm256_extractf128_ps(a, 1)), 0b10101010);
	return  _mm_shuffle_epi8(b, cvtpspbhmask128);
}

static inline __m256i _mm256_cvtpsx2_bph(__m256 a, __m256 b)
{
	__m256i c = _mm256_srli_epi32(_mm256_castps_si256(_mm256_insertf128_ps(a, _mm256_castps256_ps128(b), 3)), 16);
	__m256i d = _mm256_castps_si256(_mm256_permute2f128_ps(a, b, 0x31));
	return  _mm256_shuffle_epi8(_mm256_blend_epi16(c, d, 0b1010101010101010), cvtpspbhmask256);
}

bfloat16からfloatのへの変換

bfloat型からfloat型に戻すのは非常に簡単です．
bfloat型をshort型だと思ってまずintにアップキャストします．そうすると上位ビットほうがが０になっているので，ビットシフトでそっちを詰めて，下位ビットに０を充填します．
型変換がレイテンシ３，スループット１，ビットシフトはレイテンシ１スループット0.5のため，理論的には，fp16よりもこっちのほうが動作が早くなってほしいです．

併せて２つ出力する命令は書いていますが，こちらのほうは大差ないはずです．

static inline __m256 _mm256_cvtbph_ps(const __m128i& a)
{
	return _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepi16_epi32(a), 16));
}

static inline void _mm256_cvtbph_psx2(const __m256i& a, __m256& d0, __m256& d1)
{
	d0 = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepi16_epi32(_mm256_castsi256_si128(a)), 16));
	d1 = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepi16_epi32(_mm256_extractf128_si256(a, 1)), 16));
}

検証

1024x1024の画像を，二乗する処理を１００００回繰り返した計算時間を示します．
使用した計算機はIntel Core i9 10980HX（ゲーミングノートPC）です．

下記結果は，入力の型，出力の型で，計算はfp32で行っています．
例えば，fp16fp16の計算なら，入力としてfp16を入れ，fp32にキャストしたのちに演算してそれをfp16に戻しています．

void func_fp16fp16(const Mat& src, Mat& dest)
{
	const int size = src.size().area();
	const short* s = src.ptr<short>();
	short* d = dest.ptr<short>();
	for (int i = 0; i < size; i += 8)
	{
		__m256 a = _mm256_cvtph_ps(_mm_load_si128((__m128i*)(s + i + 0)));
		_mm_store_si128((__m128i*)(d + i + 0), _mm256_cvtps_ph(_mm256_mul_ps(a, a), 0));
	}
}

なお，OpenCVのMatの中身をFP16にするには，下記のように引数CV_16Fを指定すると変換できます．
（ここが唯一のOpenCV要素．．・）

	src.convertTo(src16, CV_16F, 1 / 255.0);

また，bfloat16に変換するメソッドはOpenCVにはないため，下記のような関数で変換してあげる必要があります．8で割れない端数処理はしていないため，実際はmaskstoreで対応する必要がありますが割愛します．

void cvt_fp32_bfloat16(const Mat& src, Mat& dest)
{
	const int size = src.size().area();
	const float* s = src.ptr<float>();
	short* d = dest.ptr<short>();
	for (int i = 0; i < size; i += 8)
	{
		_mm_store_si128((__m128i*)(d + i), _mm256_cvtps_bph(_mm256_load_ps(s + i)));
	}
}

結果

方法	時間[sec]
fp32fp32	1.577
fp16fp32	1.180
bf16fp32	1.211
bf16fp32x2	1.201
fp16fp16	0.845
bf16bf16	1.283
bf16bf16x2	1.115

入出力ともにfp16になると実行時間がほぼ半減するという結果になりました．
出力がfp32のものは，計算後に型変換をせずそのままfloatを書き出しています．
またx2がついているものは，２つまとめて処理している方法です．

肝心なbfloat16は，fp32よりは速いものの，実行速度は専用回路のあるfp16よりも遅くなりました．
理論的には，bf16をfp32に戻すのはfp16よりも速いはずですが，ベンチとったものが悪いんでしょうか．．．

おわりに

AVX2のマシンだとレジスタ数が少ないからか？性能でないという結果でした．
ちょっと計算機を変えてやってみたいですが，流行りのあれに，今更なかんじで自宅療養中なため，復帰したらということで，いったんリリース．

明日は@dandelion1124さんの予定で，
dnnモジュール（Inference Engine backend）がONNXモデル読み込みサポートしたってばよ！