Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
使用AscendC对TBE实现的SoftMarginLossGrad算子进行重构,实现了AscendC实现的SoftMarginLossGrad算子对Atlas A2 训练系列产品/Atlas 800I A2推理产品硬件的适配。
外部贡献
SoftMarginLossGrad算子的主要功能是计算Soft Margin Loss的反向传播梯度。该算子基于多分类场景下的软间隔损失函数,通过计算预测值与目标标签之间的差异,生成用于参数更新的梯度信号。其核心机制采用对数求和指数(Log-Sum-Exp)技巧处理多类别概率分布,将输入的未归一化分数(logits)与目标类别索引进行比对,生成每个样本对应所有类别的梯度贡献。这种计算方式在保证数值稳定性的同时,有效实现了间隔最大化原则,使得正确类别的分数相对于其他类别具有更大的决策边界。SoftMarginLossGrad支持批量张量计算,可高效处理浮点型数据,其计算结果直接服务于全连接层输出端的梯度回传,为神经网络分类器的端到端训练提供关键梯度信息。实现了SoftMarginLossGrad算子的AscendC实现,替代原有TBE算子在昇腾硬件上的适配。
host侧设计方案
由于算子支持三输入广播,计算过程涉及输入形状的维度对齐与广播推导,故在Host侧需先完成广播形状的归一化计算,再将广播后的统一数据视为一维向量进行任务调度。任务均分策略根据广播后的总元素数和UB块大小动态调整,确保每个核心处理的数据量均匀。 维度补齐:对input_predict、input_label、input_dout三个输入的形状进行维度补齐,将较低维度的输入形状高位补1,使其与最高维度对齐。 广播推导:推导三输入广播后的统一形状shape_max,计算总元素数作为后续任务调度的依据
分核策略
优先使用满核的原则。 如果核间能均分,可视作无大小核区分,大核小核数据块一致; 如果核间不能均分,需要将余出的数据块分配到前几个核上。 输入数据大小计算:通过GetInputShape和GetDataTypeLength函数获取输入数据的大小和类型长度,计算出输入数据的总字节数。 UB内存大小和核心数量获取:通过平台信息获取UB内存大小和核心数量,并根据这些信息调整核心数量。
数据分块和内存优化策略
充分使用UB空间的原则。 需要考虑不同硬件的UB大小不同、是否开启double buffer、kernel侧API实现过程中是否需要临时数据的储存,综合考虑单核内切分的大小。 UB内存大小获取:通过GetCoreMemSize函数获取UB内存的大小,用于后续的数据切分计算。 Tile块计算:根据UB内存大小和预定义的BLOCK_SIZE及BUFFER_NUM和不同类型下的ubDataNum,计算出每个Tile块的数据数量。 数据切分:将输入数据按照计算出的Tile块大小进行切分,计算出每个core需要处理的数据块数量和最后一个block的剩余数据量。 设置切分参数:将计算出的切分参数(如每个core的数据量、Tile块大小等)设置到MishTilingData对象中。 这些策略确保了数据在多个核心之间的均匀分布,并且在单个核心内进行了合理的切分,以提高并行处理的效率。
tilingkey规划策略
不进行tilingkey划分,在kernel侧利用输入数据的类型来走不同的分支。
kernel侧设计方案
进行Init和Process两个阶段,其中Process包括数据搬入(CopyIn)、计算(Compute)、搬出(CopyOut)三个阶段。 对于 bfloat16,float16 数据类型,统一转换为 float32 进行计算; 对于 float32 数据类型,直接计算。 无广播填充需求,直接在Compute阶段执行元素级运算。 根据tilingkey(固定)执行核函数。 把原TBE算子中的tbe.vmuls、tbe.vmul、tbe.cast_to指令替换为AscendC的Muls、Mul、Cast等矢量指令。 按照TBE算子的计算逻辑完成SoftMarginLossGrad计算。
/assign
请@fullt 组织需求评审分析 有需要上nn sig 例会评审讨论
Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
Backgroud(背景信息)
使用AscendC对TBE实现的SoftMarginLossGrad算子进行重构,实现了AscendC实现的SoftMarginLossGrad算子对Atlas A2 训练系列产品/Atlas 800I A2推理产品硬件的适配。
Origin(信息来源)
外部贡献
Benefit / Necessity (价值/作用)
SoftMarginLossGrad算子的主要功能是计算Soft Margin Loss的反向传播梯度。该算子基于多分类场景下的软间隔损失函数,通过计算预测值与目标标签之间的差异,生成用于参数更新的梯度信号。其核心机制采用对数求和指数(Log-Sum-Exp)技巧处理多类别概率分布,将输入的未归一化分数(logits)与目标类别索引进行比对,生成每个样本对应所有类别的梯度贡献。这种计算方式在保证数值稳定性的同时,有效实现了间隔最大化原则,使得正确类别的分数相对于其他类别具有更大的决策边界。SoftMarginLossGrad支持批量张量计算,可高效处理浮点型数据,其计算结果直接服务于全连接层输出端的梯度回传,为神经网络分类器的端到端训练提供关键梯度信息。实现了SoftMarginLossGrad算子的AscendC实现,替代原有TBE算子在昇腾硬件上的适配。
Design(设计方案)
host侧设计方案
由于算子支持三输入广播,计算过程涉及输入形状的维度对齐与广播推导,故在Host侧需先完成广播形状的归一化计算,再将广播后的统一数据视为一维向量进行任务调度。任务均分策略根据广播后的总元素数和UB块大小动态调整,确保每个核心处理的数据量均匀。
维度补齐:对input_predict、input_label、input_dout三个输入的形状进行维度补齐,将较低维度的输入形状高位补1,使其与最高维度对齐。
广播推导:推导三输入广播后的统一形状shape_max,计算总元素数作为后续任务调度的依据
分核策略
优先使用满核的原则。
如果核间能均分,可视作无大小核区分,大核小核数据块一致;
如果核间不能均分,需要将余出的数据块分配到前几个核上。
输入数据大小计算:通过GetInputShape和GetDataTypeLength函数获取输入数据的大小和类型长度,计算出输入数据的总字节数。
UB内存大小和核心数量获取:通过平台信息获取UB内存大小和核心数量,并根据这些信息调整核心数量。
数据分块和内存优化策略
充分使用UB空间的原则。
需要考虑不同硬件的UB大小不同、是否开启double buffer、kernel侧API实现过程中是否需要临时数据的储存,综合考虑单核内切分的大小。
UB内存大小获取:通过GetCoreMemSize函数获取UB内存的大小,用于后续的数据切分计算。
Tile块计算:根据UB内存大小和预定义的BLOCK_SIZE及BUFFER_NUM和不同类型下的ubDataNum,计算出每个Tile块的数据数量。
数据切分:将输入数据按照计算出的Tile块大小进行切分,计算出每个core需要处理的数据块数量和最后一个block的剩余数据量。
设置切分参数:将计算出的切分参数(如每个core的数据量、Tile块大小等)设置到MishTilingData对象中。
这些策略确保了数据在多个核心之间的均匀分布,并且在单个核心内进行了合理的切分,以提高并行处理的效率。
tilingkey规划策略
不进行tilingkey划分,在kernel侧利用输入数据的类型来走不同的分支。
kernel侧设计方案
进行Init和Process两个阶段,其中Process包括数据搬入(CopyIn)、计算(Compute)、搬出(CopyOut)三个阶段。
对于 bfloat16,float16 数据类型,统一转换为 float32 进行计算;
对于 float32 数据类型,直接计算。
无广播填充需求,直接在Compute阶段执行元素级运算。
根据tilingkey(固定)执行核函数。
把原TBE算子中的tbe.vmuls、tbe.vmul、tbe.cast_to指令替换为AscendC的Muls、Mul、Cast等矢量指令。
按照TBE算子的计算逻辑完成SoftMarginLossGrad计算。