标量和矢量编程简介 - 2023.2 简体中文

AI 引擎内核与计算图编程指南 (UG1079)

Document ID
UG1079
Release Date
2023-12-04
Version
2023.2 简体中文

本节提供了有关标量和矢量处理元素的内核编程关键要素的概述。如需了解有关每个元素和最优化技巧的详细信息,请参阅后续章节。

以下示例仅使用标量引擎。其中展示了迭代穿过 512 个 int32 元素的 for 循环。每次循环迭代都执行 int32 a 和 int32 b 的单一乘法,将结果存储在 c 中并将其写入输出缓冲器。scalar_mul 内核会对 input_buffer<int32> 数据的两个输出块(缓冲器)进行运算,并生成 output_buffer<int32> 数据输出缓冲器。

缓冲器是通过标量和矢量迭代器来访问的。如需了解有关缓冲器 API 的更多详细信息,请参阅 串流数据 API

void scalar_mul(input_buffer<int32> & data1,
			input_buffer<int32> & data2,
			output_buffer<int32> & out){
        auto pin1 = aie::begin(data1);
        auto pin2 = aie::begin(data2);
        auto pout = aie::begin(out);
	for(int i=0;i<512;i++)
	{
		int32 a=*pin1++;
		int32 b=*pin2++;
		int32 c=a*b;
		*pout++ = c;
	}
}

以下示例是相同内核的矢量化版本。

void vect_mul(input_buffer<int32> & __restrict data1,
			input_buffer<int32> & __restrict data2,
			output_buffer<int32> & __restrict out){
        auto pin1 = aie::begin_vector<8>(data1);
        auto pin2 = aie::begin_vector<8>(data2);
        auto pout = aie::begin_vector<8>(out);
	for(int i=0;i<64;i++)
	chess_prepare_for_pipelining
	{
		aie::vector<int32,8> va=*pin1++;
		aie::vector<int32,8> vb=*pin2++;
		aie::accum<acc80,8> vt=mul(va,vb);
		aie::vector<int32,8> vc=srs(vt,0);
		*pout++ = vc;
	}
}

请注意,先前内核代码中使用的数据类型是 vector<int32,8>accum<acc80,8>。缓冲器 API begin_vector<8> 返回的矢量将基于 8 个 int32 的矢量进行迭代,并将其存储在名为 vavb 的变量中。这两个变量均为矢量类型的变量,并传递至内部函数 mul,该函数会输出 accum<acc80,8> 数据类型的 vt移位 - 舍入 - 饱和函数 srs 会对 accum<acc80,8> 类型进行缩减,该函数允许返回 vector<int32,8> 类型的 vc 变量,随后将其写入输出缓冲器。如需了解有关 AI 引擎支持的数据类型的更多详细信息,请参阅后续章节。

vect_mul 函数的输入和输出参数上使用的 __restrict 关键字允许通过显式声明数据之间的独立性来执行更激进的编译器最优化。

chess_prepare_for_pipelining 是编译器的编译指示,用于指令内核编译器为循环达成最优化的流水线。

此示例函数的标量版本耗时 1055 个周期,而矢量化版本仅耗时 99 个周期。如您所见,内核的矢量化版本速度提升超过 10 倍。矢量处理本身能对 int32 乘法提供 8 倍的吞吐量,但时延更高,总体无法达到 8 倍的吞吐量。但完成循环最优化后,可达到接近 10 倍。下一节详细描述了可使用的各种数据类型、可用的寄存器、以及在 AI 引擎上可达成的各种最优化,要达成这些最优化,需要使用诸如循环中的软件流水打拍和 __restrict 之类的关键字等概念。