Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud（背景信息）

使用Ascend C对TBE实现的MseLossGrad算子进行重构。
实现算子适配Atlas A2、Atlas 800I A2系列硬件。

Origin（信息来源）

社区任务

Benefit / Necessity （价值/作用）

• 算子功能：MseLossGrad 算子计算均方误差函数MseLoss的反向传播。

• 计算公式：

  当reduction为mean时：

  $$MselossBackward(grad, x, y) = grad * (x - y) * 2 / x.numel()$$

  其中x.numel()表示x中的元素个数。如果reduction不是mean, 那么：

  $$MselossBackward(grad, x, y) = grad * (x - y) * 2$$

• 支持设备：
  Atlas 800I/T A2

Design（设计方案）

TBE源码分析
对MseLossGrad算子TBE版本的功能分析如下：
1）功能概述
MSE Loss Grad 算子用于计算均方误差（Mean Squared Error）损失函数的梯度。它接收上游梯度 dout、前向算子的预测值 predict和真实标签值 label，根据指定的规约（reduction）模式计算并输出损失函数对 predict的梯度 grad。
2）数据类型支持
输入 predict、label和 dout支持 float16、float32 及 bfloat16 三种数据格式，且三者的数据类型必须保持一致。输出 grad的数据类型与输入数据类型相同。
3）计算精度与处理流程
为保障计算精度，当输入数据类型为 float16 时，算子内部会在核心计算过程中将其提升转换（cast）为 float32 进行计算。待计算完成后，再将结果转换回输入的原始数据类型。对于 float32 和 bfloat16 输入，则直接在原精度下进行计算。
4）规约模式与数学原理
算子的计算行为由 reduction属性控制，支持三种规约模式：
均值规约模式（reduction="mean"）：默认模式此模式计算所有元素的平均损失梯度，核心数学公式为：grad = (2.0 / n) × (predict - label) × dout，其中 n为 predict张量中的元素总数（reduce_elts），dout为上游梯度标量。
求和规约模式（reduction="sum"）：此模式计算所有元素的损失梯度总和，核心数学公式为：grad = 2.0 × (predict - label) × dout，其中 dout为上游梯度标量。
无规约模式（reduction="none"）：此模式保持逐元素的损失梯度，不进行求和或平均，核心数学公式为：grad = 2.0 × (predict - label) × dout，其中 dout的形状与 predict相同，表示逐元素的上游梯度。
5）形状与广播
算子通过以下步骤处理输入张量的形状：
首先调用 shape_util.unify_broadcast_shapes函数，将 predict、label和 dout的形状统一为可广播的最大公共形状。
随后将三个输入张量广播（broadcast）到该统一形状，确保逐元素计算的一致性。
输出张量 grad的形状与广播后的 predict形状完全相同。
6）计算流程
算子的核心计算流程如下：
形状统一：确定可广播的最大公共形状，将输入张量统一至此形状。
数据类型转换：如输入为 float16，转换为 float32 精度进行计算。
系数计算：根据 reduction模式计算系数：
"mean" 模式：cof = 2.0 / reduce_elts（其中 reduce_elts为 predict的元素总数）
"sum" 或 "none" 模式：cof = 2.0
梯度计算：执行三步计算：
计算差值：sub_res = predict - label
归一化梯度：norm_grad = sub_res × cof
最终梯度：grad_res = norm_grad × dout
数据类型恢复：如输入为 float16，将结果转换回 float16。
7）实现机制
算子实现采用 TVM 计算图表示，通过 tbe接口调用底层硬件优化指令：
通过 tbe.vsub计算张量差值
通过 tbe.vmuls执行标量乘法
通过 tbe.vmul执行逐元素乘法
自动调度（auto_schedule）生成优化计算图
最终编译为硬件可执行内核
8）类型安全
通过严格的数据类型转换流程确保类型安全。所有输入首先被检查数据类型一致性，在计算过程中根据需要进行精度提升，最终结果精确转换回与输入一致的数据类型，确保输入输出类型的一致性。
9）错误处理
算子对输入参数进行严格验证：
检查 reduction参数必须为 "mean"、"sum" 或 "none"
验证所有输入张量的数据类型一致性
检查 kernel 名称的有效性
对不支持的参数组合会抛出明确的错误信息
TBE版本流程图

算子原型

算子支持型号
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 =0 : dout的shape与predict相同；
   tilingkey =1 : dout的shape是[1]；

kernel侧设计方案
进行Init和Process两个阶段，其中Process包括数据搬入（CopyIn）、计算（Compute）、搬出（CopyOut）三个阶段。实现上，将非fp32数据精度转换为fp32计算后，再转换为原始精度类型。

AscendC实现流程图如下：