梯度计算方法、设备、存储介质及程序产品与流程

本公开涉及芯片，更具体地，涉及一种梯度计算方法、设备、存储介质及程序产品。

背景技术：

1、在基于图形处理器（graphics processing unit，gpu）等芯片进行深度学习等人工智能（artificial intelligence，ai）模型训练或推理的过程中，反向传播过程的计算效率对人工智能模型的训练或推理效率影响较大。例如，在深度学习算子的反向传播中，大量标量梯度（如注意力权重梯度、采样坐标梯度等）需要对特定维度（例如通道维度）进行归约运算后写出。

2、在相关技术中，可以通过gpu的多线程并行计算各线程的通道部分和，随后通过原子加法（atomicadd）等原子操作在全局显存中完成归约（reduce）。这会导致每个梯度计算都引发多次原子操作。由于原子操作必须串行通过gpu高速缓存控制器的仲裁队列，会产生访问延迟与存储带宽争用，影响梯度计算效率。

技术实现思路

1、本公开实施例的一个目的是提供一种关于梯度计算的新的技术方案。

2、根据本公开实施例的第一方面，提供了一种梯度计算方法，应用于目标芯片，所述目标芯片包括多个线程束，每一线程束包括多个线程、每一线程所对应的寄存器、以及用于多个寄存器互联访问的寄存器互联网络；所述方法包括：

3、从待计算的候选梯度张量中确定满足预设梯度输出条件的目标梯度张量；其中，所述预设梯度输出条件包括：所述目标梯度张量的多个输出元素中每一所述输出元素均基于设定归约维度进行归约运算生成，并且不同的输出元素分别写入所述目标芯片的存储器中不同存储地址；

4、将所述目标梯度张量中同一输出元素对应的多个目标计算任务分配至同一线程束内的不同线程执行，以使得所述线程束基于所述寄存器以及所述寄存器互联网络执行多个目标计算任务以实现所述设定归约维度的归约运算后得到所述输出元素，由所述线程束内的指定线程将所述输出元素写入所述目标芯片的存储器。

5、可选地，所述输出元素对应的元素表示维度为多个；所述将所述目标梯度张量中同一输出元素对应的多个目标计算任务分配至同一线程束内的不同线程执行，包括：

6、将所述输出元素的多个元素表示维度转换为对应于所述输出元素的一维元素索引；

7、根据所述一维元素索引确定所述输出元素对应的线程束标识；

8、针对任一输出元素，将所述输出元素对应的多个目标计算任务分配至所述线程束标识对应的线程束内的不同线程执行。

9、可选地，将所述输出元素的多个元素表示维度转换为对应于所述输出元素的一维元素索引，包括：

10、确定所述输出元素在各个所述元素表示维度上的维度索引值；

11、获取所述目标梯度张量的多个输出元素在各个所述元素表示维度上的维度索引数量；

12、基于所述维度索引值和所述维度索引数量，计算得到所述输出元素对应的一维元素索引。

13、可选地，同一输出元素所对应多个目标计算任务中各个目标计算任务与所述输出元素所对应设定归约维度的归约维度索引值一一对应；所述将所述目标梯度张量中同一输出元素对应的多个目标计算任务分配至同一线程束内的不同线程执行，包括：

14、确定所述线程束内的线程数量以及所述设定归约维度的归约维度索引数量；

15、根据所述线程数量和所述归约维度索引数量，确定所述线程束内各个线程的线程索引值与所述归约维度索引值的线程分配索引对应关系；

16、根据所述线程分配索引对应关系，将不同归约维度索引值对应的目标计算任务分配至对应于所述归约维度索引值的目标线程。

17、可选地，所述候选梯度张量为目标算子内的梯度张量，所述从待计算的候选梯度张量中确定满足预设梯度输出条件的目标梯度张量，包括：

18、在算子开发阶段，确定所述目标算子中待计算的候选梯度张量；

19、从待计算的所述候选梯度张量中确定满足所述预设梯度输出条件的候选梯度张量作为所述目标梯度张量；

20、所述方法还包括：将不满足所述预设梯度输出条件的候选梯度张量作为其他梯度张量；其中，所述其他梯度张量的梯度计算采用与所述目标梯度张量不同的任务分配方式。

21、可选地，所述预设梯度输出条件还包括：所述输出元素的输出维度不包括所述设定归约维度。

22、可选地，所述线程束实现多个通道的归约运算的方式包括：

23、所述线程束的每一线程基于所述线程对应的寄存器执行所分配的目标计算任务以计算得到至少一个通道的线程级归约结果，将所述线程级归约结果输出至所述寄存器中；

24、通过寄存器互联网络对所述线程束内的多个线程级归约结果进行归约运算，得到所述输出元素的元素归约结果。

25、可选地，所述通过寄存器互联网络对所述线程束内的多个线程级归约结果进行归约运算，得到所述目标梯度张量的输出元素，包括：

26、通过寄存器互联网络对所述线程束内执行线程束洗牌指令，以对多个线程级归约结果进行树形归约，得到所述输出元素的元素归约结果。

27、可选地，所述目标梯度张量包括权重梯度和/或坐标梯度。

28、可选地，所述设定归约维度包括通道维度，所述通道维度的归约维度索引数量为所述线程束内的线程数量的整数倍。

29、根据本公开实施例的第二方面，提供了一种电子设备，包括存储器和处理器，所述存储器用于存储计算机指令，所述处理器用于从所述存储器中调用所述计算机指令，以执行第一方面所述的方法。

30、根据本公开实施例的第三方面，提供了一种计算机可读存储介质，其上存储有计算机程序，所述计算机程序在被处理器执行时实现第一方面所述的方法。

31、根据本公开实施例的第四方面，提供了一种计算机程序产品，所述计算机程序产品包括计算机程序，所述计算机程序在被处理器执行时实现第一方面所述的方法。

32、基于本公开实施例提供的梯度计算方法，将同一输出元素的归约运算分配至一个线程束，从而将目标梯度张量的归约运算从传统的基于全局存储层次的原子累加路径转移至线程束内的寄存器互联网络中完成，能够减少或消除相关原子操作对缓存控制器cas队列的占用，降低存储访问延迟、缓解存储带宽争用，并利用极低延迟的寄存器以及寄存器互联网络实现了归约运算的硬件加速，进而提高了梯度计算效率。

33、通过以下参照附图对本公开的示例性实施例的详细描述，本公开的其他特征及其优点将会变得清楚。

1.一种梯度计算方法，其特征在于，应用于目标芯片，所述目标芯片包括多个线程束，每一线程束包括多个线程、每一线程所对应的寄存器、以及用于多个寄存器互联访问的寄存器互联网络；所述方法包括：

2.根据权利要求1所述的方法，其特征在于，所述输出元素对应的元素表示维度为多个；所述将所述目标梯度张量中同一输出元素对应的多个目标计算任务分配至同一线程束内的不同线程执行，包括：

3.根据权利要求2所述的方法，其特征在于，将所述输出元素的多个元素表示维度转换为对应于所述输出元素的一维元素索引，包括：

4.根据权利要求1所述的方法，其特征在于，同一输出元素所对应多个目标计算任务中各个目标计算任务与所述输出元素所对应设定归约维度的归约维度索引值一一对应；所述将所述目标梯度张量中同一输出元素对应的多个目标计算任务分配至同一线程束内的不同线程执行，包括：

5.根据权利要求1所述的方法，其特征在于，所述候选梯度张量为目标算子内的梯度张量，所述从待计算的候选梯度张量中确定满足预设梯度输出条件的目标梯度张量，包括：

6.根据权利要求1所述的方法，其特征在于，所述预设梯度输出条件还包括：所述输出元素的输出维度不包括所述设定归约维度。

7.根据权利要求1所述的方法，其特征在于，所述线程束实现多个通道的归约运算的方式包括：

8.根据权利要求7所述的方法，其特征在于，所述通过寄存器互联网络对所述线程束内的多个线程级归约结果进行归约运算，得到所述目标梯度张量的输出元素，包括：

9.根据权利要求1所述的方法，其特征在于，

10.一种电子设备，其特征在于，包括存储器和处理器，所述存储器用于存储计算机指令，所述处理器用于从所述存储器中调用所述计算机指令，以执行权利要求1至9中任一项所述的方法。

11.一种计算机可读存储介质，其特征在于，其上存储有计算机程序，所述计算机程序在被处理器执行时实现权利要求1至9中任一项所述的方法。

12.一种计算机程序产品，其特征在于，包括计算机程序，所述计算机程序在被处理器执行时实现权利要求1至9中任一项所述的方法。