| 【社区任务】AxpyV2算子开发-算子提交
Co-authored-by: Nice-try-zzw<1849404617@qq.com>
# message auto-generated for no-merge-commit merge:
!479 merge submit-axpy_v2 into master
【社区任务】AxpyV2算子开发-算子提交
Created-by: Nice_try
Commit-by: Nice-try-zzw
Merged-by: cann-robot
Description: ## 描述
### 背景信息
基于AxpyV2算子历史TBE版本使用Ascend C编程语言进行优化。
### TBE源码分析
通过对AxpyV2算子TBE版本的功能分析，当前支持的能力如下：
① 算子支持float16，float32，bfloat16，int32四种格式的输入输出。
② AxpyV2算子涉及到对输入数据进行广播（输入数据的shape调整到相同大小），任务书说明：暂可不支持广播场景
③ 在对输入数据进行广播操作后，调用vmul、vadd、vmla和cast接口实现相应的表达式：
- 算子功能：源操作数2(src2Tensor)中每个元素与标量(alphaScalar)对应元素求积后和源操作数1(src1Tensor)中的对应元素相加。
> 
> 函数说明中，可以确定alpha是Scalar
- 计算公式：
$$
dstTensor_i = src1Tensor_i + alphaScalar * src2Tensor_i
$$
> 对于数据类型为float16和bfloat16，需要类型转换为float32进行计算，另外支持alphaTensor与src1Tensor数据类型不一致。
AxpyV2算子TBE版本的整体流程图如下图所示：
> 在对tbe算子进行st测试时，三个输入以及一个输出的数据类型仅需满足是四个数据类型的其中一个即可，故共有$4^4=256$种数据类型组合情况。
### 算子原型
<table style="undefined;table-layout: fixed; width: 1494px"><colgroup>
<col style="width: 146px">
<col style="width: 110px">
<col style="width: 301px">
<col style="width: 219px">
<col style="width: 328px">
<col style="width: 101px">
<col style="width: 143px">
<col style="width: 146px">
</colgroup>
<thead>
<tr>
<th>参数名</th>
<th>输入/输出</th>
<th>描述</th>
<th>使用说明</th>
<th>数据类型</th>
<th>数据格式</th>
<th>维度(shape)</th>
<th>非连续Tensor</th>
</tr></thead>
<tbody>
<tr>
<td>x1</td>
<td>输入</td>
<td>待进行axpy_v2计算的入参，公式中的src1Tensor。</td>
<td>无</td>
<td>FLOAT、FLOAT16、BFLOAT16、INT32</td>
<td>ND</td>
<td>0-8</td>
<td>√</td>
</tr>
<tr>
<td>x2</td>
<td>输入</td>
<td>待进行axpy_v2计算的入参，公式中的src2Tensor。</td>
<td>shape与x1相同。</td>
<td>FLOAT、FLOAT16、BFLOAT16、INT32</td>
<td>ND</td>
<td>0-8</td>
<td>√</td>
</tr>
<tr>
<td>alpha</td>
<td>输入</td>
<td>待进行axpy_v2计算的入参，公式中的alphaScalar。</td>
<td>shape为[]</td>
<td>FLOAT、FLOAT16、BFLOAT16、INT32</td>
<td>ND</td>
<td>0-8</td>
<td>√</td>
</tr>
<tr>
<td>out</td>
<td>输出</td>
<td>待进行axpy_v2计算的出参，公式中的dstTensor。</td>
<td>shape与x1相同。</td>
<td>FLOAT、FLOAT16、BFLOAT16、INT32</td>
<td>ND</td>
<td>0-8</td>
<td>√</td>
</tr>
</tbody>
</table>
### 算子支持型号
Atlas A2 训练系列产品/Atlas 800I A2推理产品
### host侧设计方案
算子计算过程不涉及数据的维度信息，故在host侧将数据视为一维向量，仅考虑数据个数，不考虑数据维度信息。
任务均分：coreNum 根据输入长度和块大小动态调整，确保每个核心处理的数据块数均匀。
批量搬运：tileBlockNum 和 tileDataNum 计算单次搬运的数据量，通过 finalSmallTileNum 和 finalBigTileNum 确定小核/大核的搬运次数，将多次搬运合并为批量操作，减少冗余开销。尾块的处理逻辑确保不完整块也能被合并到计算流程中，避免数据碎片。
#### 1) 分核策略
优先使用满核的原则。
如果核间能均分，可视作无大小核区分，大核小核数据块一致；
如果核间不能均分，需要将余出的数据块分配到前几个核上。
输入数据大小计算：通过GetInputShape和GetDataTypeLength函数获取输入数据的大小和类型长度，计算出输入数据的总字节数。
UB内存大小和核心数量获取：通过平台信息获取UB内存大小和核心数量，并根据这些信息调整核心数量。
#### 2) 数据分块和内存优化策略
充分使用UB空间的原则。
需要考虑不同硬件的UB大小不同、是否开启double buffer、kernel侧API实现过程中是否需要临时数据的储存，综合考虑单核内切分的大小。
UB内存大小获取：通过GetCoreMemSize函数获取UB内存的大小，用于后续的数据切分计算。
Tile块计算：根据UB内存大小和预定义的BLOCK_SIZE及BUFFER_NUM和不同类型下的ubDataNum，计算出每个Tile块的数据数量。
数据切分：将输入数据按照计算出的Tile块大小进行切分，计算出每个core需要处理的数据块数量和最后一个block的剩余数据量。
设置切分参数：将计算出的切分参数（如每个core的数据量、Tile块大小等）设置到RealDivTilingData对象中。
这些策略确保了数据在多个核心之间的均匀分布，并且在单个核心内进行了合理的切分，以提高并行处理的效率。
#### 3) tilingkey规划策略
不进行tilingkey划分，在kernel侧利用输入数据的类型来走不同的分支。
### kernel侧设计方案
进行Init和Process两个阶段，其中Process包括数据搬入（CopyIn）、计算（Compute）、搬出（CopyOut）三个阶段。
1) 依照TBE实现，对于三个输入均为int32的情况下，使用Mul和Add来实现，并且无需转换数据类型，其余数据类型组合均转换为fp32进行计算；计算完成后，转换为输出的数据类型进行搬出。
2) Ascend C的AxpyV2算子流程见下图。
## 关联的Issue
[#281](https://gitcode.com/cann/ops-math/issues/281)
[#231](https://gitcode.com/cann/ops-math/issues/231)
## 测试
<!--描述进行了哪些测试来验证你的改动。包括但不限于二级冒烟、算子泛化等。-->
不涉及
## 文档更新
<!--如果这个PR包含文档的更新，请在这里指出。例如：更新了README.md文件。-->
## 类型标签
<!-- [x] 表示选中 -->
- [ ] Bug修复
- [ ] 新特性
- [ ] 性能优化
- [ ] 文档更新
- [x] 其他，请描述：社区任务算子设计文档
See merge request: cann/ops-math!479 | 4 个月前 |
