从软件实现到硬件加速的数学算子演进:深入解析 ops-math 如何释放昇腾NPU的数学计算潜力
2026/6/11 17:57:00 网站建设 项目流程

前言

写了一段 PyTorch 代码,里面有几个 tensor 的转置、类型转换、还有几个随机数生成。代码跑起来没问题,结果也对。但当你把同样的逻辑搬到昇腾NPU上跑的时候,速度并没有你期待的那种"硬件加速"的感觉。

问题出在哪?很多开发者的第一反应是去查昇腾 CANN 的算子支持列表,看某个函数有没有对应的 NPU 实现。

很多开发者的第一反应是去查算子支持列表,看某个函数有没有对应的 NPU 实现。但这个思路忽略了一个更根本的问题:数学类的基础操作(类型转换、逐元素数学函数、随机数生成)在 CPU 上和在 NPU 上的执行方式完全不同。CPU 上你随手写的一行代码,到底层可能调用了不同的指令集、不同的内存布局、甚至不同的计算范式。

昇腾NPU的架构设计里,数学计算单元和矩阵计算单元是分开的。Cube 单元专门处理矩阵乘法这类大规模并行计算,Vector 单元专门处理逐元素的数学运算。如果你把应该放在 Vector 单元上跑的数学操作错误地用 Cube 单元去处理,或者反过来,性能差距可能达到一个数量级。

ops-math 这个仓库做的事情,就是把那些你在 CPU 上习以为常的数学操作,重新实现一遍,让它们能在昇腾NPU的 Vector 单元上高效执行。不是简单的"把代码移植到 NPU",而是针对 NPU 的硬件特性重新设计计算流程。

理解 ops-math 的价值,需要先从"CPU 上的数学计算是怎么执行的"说起,看 NPU 的硬件架构对这类计算有什么不同的要求,再看 ops-math 是怎么填补这个 gap 的。


数学类算子在 AI 框架中的真实位置

把 AI 框架(PyTorch、MindSpore、PaddlePaddle)想象成一个厨房。你作为厨师,关心的是"我要做哪些菜"(模型结构)、“食材怎么搭配”(张量形状和数据类型)、“火候怎么控制”(超参数)。但厨房里真正把食材变成菜肴的,是那些你平时不会特意关注的工具:刀具、砧板、炉灶、烤箱。

数学类算子就是这些"不会特意关注的工具"。

你写一个 Transformer 模型,核心逻辑是 self-attention 的计算:矩阵乘法、softmax、LayerNorm。这些是大菜。但在实现这些大菜的过程中,你会频繁地做以下几件事:

把 float32 的张量转换成 float16(或者 bfloat16),因为 NPU 上低精度计算更快。这个"转换"动作,就是 conversion 类算子。

对张量中的每个元素做数学变换,比如取指数、取对数、计算 sigmoid。这些操作看起来简单,但当一个张量有几十万、几百万个元素的时候,逐个处理也需要可观的计算资源。这些就是 math 类算子。

初始化模型参数时需要生成随机数,dropout 时需要生成随机掩码,这些数据也需要"看起来足够随机"。这些就是 random 类算子。

在 CPU 上,这些操作通常由底层库(比如 NumPy 的 C 实现、PyTorch 的 C++ 后端)高效完成。但当你把计算搬到 NPU 上的时候,这些底层实现不再可用。你需要一套新的实现,专门针对 NPU 的 Vector 计算单元优化。

这就是 ops-math 的定位:它是昇腾NPU上数学类基础算子的"底层工具库",为上层框架(PyTorch、MindSpore 等)提供高效的数学计算能力。

从 CANN 的五层架构来看,ops-math 属于第二层"昇腾计算服务层"中的 AOL 算子库的一部分,但具体实现会调用第一层的 Ascend C 编程语言来编写算子内核。它的输出会被 PyTorch 这样的框架通过 Framework Adaptor 来调用,也会直接被开发者通过 AscendCL 接口调用。

理解这个位置很重要。它决定了 ops-math 的设计约束:它不能假设上层框架的存在(因为可能直接被 AscendCL 调用),也不能忽略底层硬件特性(因为要在 Vector 单元上高效执行)。


Vector 计算单元的编程范式:为什么不能直接用 CPU 代码

昇腾NPU的达芬奇架构里,最核心的两个计算单元是 Cube 和 Vector。

Cube 单元专门做矩阵乘法。它的设计目标是"一次操作处理一大块数据"。假设你有两个 128×128 的矩阵要做乘法,Cube 单元可以把它拆成很多个小块,每个小块内部并行计算,再合并结果。这种"分块并行"的范式,非常适合矩阵乘法、卷积这类"输出每个元素都依赖输入很多元素"的操作。

Vector 单元做的事情完全不同。它处理的是"逐元素操作":对张量中的每个元素独立地做同样的数学变换。取指数、取对数、类型转换、生成随机数,这些都是逐元素操作。Vector 单元的设计目标是"同时处理很多个独立的计算任务"。

这两种计算范式的差异,决定了它们的编程模型完全不同。

在 CPU 上写逐元素操作,你通常会写一个 for 循环,遍历张量的每个元素,对每个元素做同样的变换。CPU 的编译器会自动做向量化(用 SIMD 指令),但你写代码的时候不需要关心这个。

在 NPU 的 Vector 单元上,你没有"for 循环"这个概念。你需要把数据切成很多个"分块"(tile),每个分块分配给一个计算核心(AI Core)去处理,每个核心内部再用 Vector 指令并行处理分块内的所有元素。

这个"切分数据→分配核心→并行执行"的过程,就是 Ascend C 编程语言要帮你管理的事情。但你写 Ascend C 代码的时候,仍然需要显式地指定:数据怎么切分、每个分块多大、分块之间有没有依赖关系、中间结果存在哪。

CPU 上你随手写的一行y = torch.exp(x),到底层可能只是一段调用了优化的 exp 函数的 C 代码。但在 NPU 上,同样的逻辑需要写成一个完整的 Ascend C kernel:定义输入张量的内存布局、定义分块大小、定义每个分块的计算逻辑、定义输出张量的内存布局、处理边界情况(当张量大小不是分块大小的整数倍时)、处理数据类型转换(如果输入是 float32 而输出是 float16)。

这就是 ops-math 存在的第一个理由:它把上面这一整套流程,针对常见的数学操作(exp、log、sigmoid、类型转换、随机数生成等),各实现了一遍。你不需要自己写 Ascend C kernel,只需要调用ops_math.exp(x)就行。

但 ops-math 的价值不止是"提供了现成的实现"。它更重要的是做了很多针对 NPU 架构的性能优化,这些优化在 CPU 代码里要么不需要、要么完全不一样。


conversion 类算子:类型转换背后的内存布局博弈

类型转换听起来是最简单的操作之一。你把 float32 的张量转换成 float16,不就是每个元素占的位数从 32 位变成 16 位吗?

在 CPU 上,确实是这样。CPU 的内存模型是"一段连续的字节",类型转换就是遍历这段内存,把每 4 个字节重新解释成 2 个字节(或者反过来)。

但在 NPU 上,内存布局不是简单的"一段连续的字节"。

昇腾NPU为了计算效率,会使用多种内存布局格式。最常见的两种是 ND 格式和 FRACTAL_Z 格式。

ND 格式就是你在 CPU 上熟悉的多维数组的线性存储:一个二维矩阵,先存第一行的所有元素,再存第二行的所有元素,以此类推。这种格式对人类友好,对逐元素操作也友好(因为相邻元素在内存中也相邻)。

但 FRACTAL_Z 格式完全不同。它会把矩阵切成很多个 16×16 的小块(这个大小跟 Cube 单元的计算粒度匹配),每个小块的 256 个元素在内存中连续存储,但小块和小块之间的顺序不是按行主序或列主序排列的,而是按一种对 Cube 单元友好的方式重新排列。

当你的张量在内存中是 FRACTAL_Z 格式的时候,你不能简单地"遍历每个元素做类型转换"。因为"每个元素"在内存中不是连续存储的,你需要先知道张量的内存布局,再按照正确的访问模式去读取和修改每个元素。

ops-math 的 conversion 类算子处理的核心问题之一,就是"在不知道输入张量内存布局的情况下,正确地做类型转换"。

它不能直接假设输入是 ND 格式,也不能直接假设输入是 FRACTAL_Z 格式。它需要做一次布局检测(通过 Ascend C 提供的接口查询张量的内存描述符),根据检测结果选择对应的转换路径。

如果输入是 ND 格式,转换逻辑相对简单:直接用 Vector 单元的 vconv 指令(专门做数据类型转换的指令),一次处理一个分块的所有元素。

如果输入是 FRACTAL_Z 格式,转换逻辑就复杂了:需要先按 FRACTAL_Z 的访问模式读取数据,转换成目标类型之后,再按 ND 格式写回(因为逐元素操作通常用 ND 格式更高效),或者继续保持 FRACTAL_Z 格式(如果后续操作是矩阵乘法)。

这个"布局检测→路径选择"的逻辑,在 CPU 上完全不存在。但在 NPU 上,它是类型转换算子必须处理的问题。

看一段简化版的类型转换代码(这是 Ascend C 的编程模式,不是你平时写的 PyTorch 代码):

// 这个 kernel 把一个 float32 的 ND 格式张量转换成 float16__global__voidcast_fp32_to_fp16(__gm__float*x,__gm__ half*y,intn){// 每个 AI Core 处理 256 个元素__aicore__int32_tblock_idx=GetBlockIdx();__aicore__int32_tblock_size=256;__aicore__int32_tstart=block_idx*block_size;__aicore__int32_tend=min(start+block_size,n);// 用 LocalTensor 在片上内存(UB)中分配空间__aicore__ LocalTensor<float>x_local=AllocTensor<float>(block_size);__aicore__ LocalTensor<half>y_local=AllocTensor<half>(block_size);// DMA 搬运:从全局内存(GM)读到片上内存(UB)DataCopy(x_local,x+start,end-start);// Vector 指令:类型转换(vconv 是 Ascend C 内置的向量化转换指令)// WHY: 这里用 vconv 而不是手写转换循环,因为 vconv 会调用 Vector 单元的硬件// 转换指令,一次可以处理 256 个元素,比逐元素循环快 10 倍以上vconv(y_local,x_local,end-start);// DMA 搬运:从片上内存(UB)写回全局内存(GM)DataCopy(y+start,y_local,end-start);}

这段代码的复杂之处不在于"类型转换"这个动作本身(就是那个vconv调用),而在于它显式地管理了数据的流动:全局内存 → 片上内存 → 计算 → 片上内存 → 全局内存。

在 CPU 上,这个过程是隐式的。你写y[i] = (half)x[i],编译器和硬件会自动处理缓存、向量化、内存对齐等细节。

在 NPU 上,这些细节你必须自己处理。DMA 搬运的时机、分块的大小、片上内存的分配策略,都会影响性能。ops-math 的价值之一就是把这些细节封装好,让上层调用者不需要关心。


math 类算子:逐元素数学函数的精度与性能权衡

逐元素数学函数(exp、log、sigmoid、tanh 等)的实现,涉及一个在 CPU 上也存在、但在 NPU 上更突出的问题:精度和性能的权衡。

数学函数(尤其是超越函数,比如指数、对数、三角函数)在硬件层面通常是不能直接计算的。没有哪条指令能"直接算出一个浮点数的自然对数"。实际的做法是:用多项式逼近(通常是切比雪夫逼近或最小二乘逼近)来计算函数的近似值。

这个多项式逼近的精度,取决于多项式的阶数。阶数越高,精度越高,但计算量也越大。

在 CPU 上,这个权衡通常由数学库(比如 Intel MKL、AMD AOCL)来帮你做。你调用exp(x),底层会根据 x 的取值范围和精度要求,选择合适的逼近多项式。

在 NPU 的 Vector 单元上,这个权衡需要算子系统来做的更多。

原因有几个。第一,NPU 的 Vector 单元的计算能力跟 CPU 的 SIMD 单元不一样。它一次能处理的元素数量更多(通常是 256 个 float16 元素),但每个元素的指令延迟可能更高。第二,NPU 上的内存层次结构不同。片上内存(UB)的大小有限(通常是几百 KB),你不能假设所有中间结果都能存在片上。第三,NPU 上的计算通常是"分块执行"的,你需要确保逼近多项式的计算能够很好地跟分块大小对齐。

ops-math 的 math 类算子针对这些问题做了优化。

以 exp 函数为例。标准的逼近方法是用泰勒展开或者 Padé 逼近。但这两个方法在 NPU 上都不是最优的。泰勒展开需要在很大的取值范围上保持高精度,导致多项式阶数很高。Padé 逼近虽然有更好的收敛性,但涉及除法操作,在 NPU 的 Vector 单元上除法比乘法慢很多。

ops-math 用的是一种叫做"分段线性逼近 + 查表修正"的方法。

基本思路是:把指数函数的输入范围分成很多个小区间(比如每个区间宽度是 0.1),在每个小区间内用线性函数逼近指数函数。线性函数的系数存在一张查找表里。计算exp(x)的时候,先根据 x 的值查表找到对应的线性系数,再做一次乘法和一次加法就能得到结果。

这种方法的精度比高阶多项式逼近低一些(通常是 10^-4 级别的相对误差,而高阶多项式可以达到 10^-7 级别),但速度快很多。对于深度学习中的大部分场景(比如 softmax 中的指数计算、激活函数中的指数计算),10^-4 的精度已经足够,而速度的提升(2-3 倍)带来的收益更大。

看一段简化版的 exp 实现代码:

// 用分段线性逼近计算 exp(x)__global__voidfast_exp(__gm__float*x,__gm__float*y,intn){__aicore__int32_tblock_idx=GetBlockIdx();__aicore__int32_tblock_size=256;__aicore__int32_tstart=block_idx*block_size;__aicore__int32_tend=min(start+block_size,n);__aicore__ LocalTensor<float>x_local=AllocTensor<float>(block_size);__aicore__ LocalTensor<float>y_local=AllocTensor<float>(block_size);DataCopy(x_local,x+start,end-start);// 对每个元素做 exp 计算for(inti=0;i<end-start;i++){floatval=x_local(i);// 把 val 映射到查找表的索引// WHY: 这里用查表而不是计算多项式,因为查表+一次乘加比计算// 5 阶多项式快 2-3 倍,而精度损失对深度学习应用通常可以接受intidx=(int)((val-EXP_TABLE_MIN)/EXP_TABLE_STEP);idx=max(0,min(idx,EXP_TABLE_SIZE-2));floata=exp_table_a[idx];// 线性系数 afloatb=exp_table_b[idx];// 线性系数 by_local(i)=a*val+b;// 线性逼近:exp(val) ≈ a*val + b}DataCopy(y+start,y_local,end-start);}

这段代码的关键设计决策是:用查表+线性逼近,而不是高阶多项式。这个决策背后的考量,就是前面说的"精度 vs 性能"的权衡。在 CPU 上,你可能会选择更高阶的多项式(因为 CPU 的 SIMD 单元很擅长做连续的乘加运算)。但在 NPU 的 Vector 单元上,查表+一次乘加的模式更友好,因为它减少了指令数量和寄存器压力。

另一个优化是"向量化查表"。上面的代码里,for循环是逐元素执行的。在实际的 ops-math 实现中,会用 Vector 单元的vgather指令(一次从查找表中读取多个不连续的元素)来批量查表,进一步减少指令数量。


random 类算子:随机数生成的质量与并行化矛盾

随机数生成在 CPU 上有一个已经很成熟的解法:用伪随机数生成器(PRNG),比如 Mersenne Twister、XORShift、PCG 等。这些算法能生成统计性质很好的伪随机数序列,执行速度也很快。

但这些算法都有一个假设:随机数是"顺序生成"的。你先生成第一个数,用第一个数做种子生成第二个数,再用第二个数做种子生成第三个数,以此类推。

这个"顺序生成"的范式,在 NPU 上遇到了根本性的困难。

NPU 的 Vector 单元是做并行计算的。你希望一次生成 256 个随机数(因为一个分块的大小通常是 256 个元素),而不是一个一个地生成。

但如果你用传统的 PRNG 算法,生成 256 个随机数需要 256 步顺序计算。这完全浪费了 Vector 单元的并行能力。

解决这个问题的方法是"跳过"(skip-ahead)技术。

基本思路是:不给 PRNG 算法一个一个地生成随机数,而是直接计算"第 k 个随机数的值是多少",而不需要先计算前 k-1 个随机数。

数学上,很多 PRNG 算法可以用线性递推关系来描述。比如,一个简单的线性同余生成器:x_{n+1} = (a * x_n + c) mod m。如果你想知道x_k是多少,你不需要从x_0开始一步一步算 k 次。你可以直接用矩阵快速幂(或者更复杂的代数技巧)算出x_k的闭式表达式。

一旦你能直接计算"第 k 个随机数",你就可以让 NPU 的 256 个计算核心各自去计算"第 k+1 个、第 k+2 个、…、第 k+256 个"随机数,这些计算是完全独立的,可以并行执行。

ops-math 的 random 类算子使用的就是这种"跳过"技术。

具体实现中,它用的是 PCG(Permuted Congruential Generator)算法的一个变种,这个算法的好处是"跳过"操作可以用很快的速度完成(有专门的数学技巧,不需要真的做矩阵快速幂)。

看一段简化版的随机数生成代码:

// 用 PCG 算法并行生成随机数__global__voidrandom_uniform(__gm__uint32_t*state,__gm__float*out,intn,intstride){__aicore__int32_tblock_idx=GetBlockIdx();__aicore__int32_tlane_idx=GetLaneIdx();// 当前核心内的线程编号(0-255)__aicore__int32_tglobal_idx=block_idx*256+lane_idx;// WHY: 这里每个线程独立计算自己的随机数,不需要跟其他线程同步。// 关键在于 pcg_advance 函数——它直接计算"跳过 k 步之后的状态",// 而不需要真的执行 k 次递推。这让 256 个线程可以完全并行地生成随机数。if(global_idx<n){// 从全局状态中读取当前种子的基础值uint32_tseed=state[0];// 计算"跳过 global_idx 步之后的 PCG 状态"uint32_ts=pcg_advance(seed,global_idx*stride);// 生成随机数(PCG 的输出函数,做一次置换)uint32_trand=pcg_output(s);// 转换成 [0, 1) 范围的 floatout[global_idx]=(float)rand/(float)UINT32_MAX;}}

这段代码最核心的设计是pcg_advance(seed, k)函数。它计算的是"从 seed 开始,跳过 k 步之后的 PCG 内部状态是什么"。这个计算可以在 O(log k) 的时间内完成(用类似快速幂的方法),而不是 O(k) 的时间。

当你有 256 个线程并行执行的时候,第 i 个线程调用pcg_advance(seed, i),就能得到它应该生成的那个随机数,所有线程的计算完全独立,不需要任何同步。

这种"并行 PRNG"的设计,在 CPU 上不是不需要,而是需求没那么迫切(因为 CPU 的随机数生成通常不是性能瓶颈)。但在 NPU 上,如果你要生成一个很大的随机掩码(比如 dropout 的掩码张量,可能有几百万个元素),随机数生成的速度就会直接影响训练性能。

ops-math 的 random 类算子处理的就是这个问题。它让你可以高效地生成大规模随机数张量,随机数的统计性质(均匀性、独立性)有足够保证,不会因为在并行生成的过程中引入了相关性而导致下游的机器学习模型出现问题。


使用前 vs 使用后的效率对比

场景使用前(CPU 实现或 naive NPU 实现)使用后(ops-math 优化实现)
类型转换性能数据需要在主机和设备之间搬运,转换操作无法利用 Vector 单元并行度原生 NPU 实现,数据不离开设备内存,Vector 单元并行处理,显著降低延迟
逐元素数学函数直接用 PyTorch 的 CPU 实现或用未优化的 kernel,指令数量和内存访问模式不是最优针对 Vector 单元指令集优化,用查表+向量化逼近替代高阶多项式,吞吐显著提升
随机数生成在 CPU 上生成后拷贝到 NPU,或者用未优化的逐次生成方法,无法并行化用跳过技术的并行 PRNG,256 个线程完全独立生成,大规模随机张量生成速度大幅提升
内存占用中间结果需要额外存储在全局内存,数据搬运频繁融合算子和片上内存管理减少全局内存访问,降低内存占用
代码可维护性手写 Ascend C kernel 需要管理内存布局、分块大小、DMA 搬运等底层细节调用现成的算子接口,底层细节被封装,代码更简洁易维护

ops-math 在 CANN 生态中的依赖与复用关系

ops-math 不是孤立存在的。它跟 CANN 生态中的其他仓库有密切的依赖和复用关系。

从依赖关系来看,几乎所有其他的 ops-* 仓库(ops-nn、ops-blas、ops-cv、ops-fft、ops-tensor)都会依赖 ops-math。原因是:这些仓库实现的算子,在计算过程中经常需要做类型转换(比如把输入转换成内部计算用的精度),或者需要做逐元素的后处理(比如对卷积输出做批量归一化的数学变换)。这些基础操作如果每个仓库都自己实现一遍,会造成大量的代码重复,质量也参差不齐。

ops-math 把这些基础操作实现好,让其他仓库直接调用。这种"基础库"的定位,跟 CPU 上的 NumPy 或 Intel MKL 的定位类似。

但 ops-math 跟 NumPy 有一个重要的区别:NumPy 是纯软件的实现,而 ops-math 是针对特定硬件(昇腾NPU)优化的实现。这意味着 ops-math 的实现细节(分块大小、指令选择、内存布局)是跟硬件特性强绑定的。如果未来昇腾NPU的架构发生变化(比如 Vector 单元的处理宽度从 256 变成 512),ops-math 需要相应地调整实现。

从复用关系来看,ops-math 会被上层的框架适配层(Framework Adaptor)调用。当你用 PyTorch 写一个模型,PyTorch 的底层会自动把某些操作(比如x.float()类型转换、torch.exp(x)数学函数)路由到 ops-math 的对应实现。这个路由过程是自动的,你不需要显式地调用 ops-math 的接口。

但如果你在做性能调优,或者你在写一个自定义的算子(用 Ascend C),你可能需要显式地调用 ops-math 的接口。比如,你的自定义算子在中间步骤需要做一次类型转换,你可以直接调用 ops-math 的 conversion 算子,而不是自己实现一个。

这种"显式调用"的使用方式,在算子融合(operator fusion)的场景中特别有用。假设你有一个自定义算子,它的计算逻辑是"先做矩阵乘法,对结果做 batch normalization,再做类型转换"。如果你把这三个操作写成三个独立的算子,中间结果需要写回全局内存,造成很大的内存带宽浪费。如果你把这三个操作融合成一个算子,类型转换那部分的逻辑可以直接调用 ops-math 的实现(因为它已经针对 Vector 单元优化好了),你不需要自己重新实现一遍。

ops-math 跟 opbase 仓库的关系也值得说明。opbase 是"算子基础组件/通用库",它提供的是算子开发的"基础设施":如何管理内存、如何处理错误、如何做日志记录、如何跟框架适配层对接。ops-math 在实现过程中会调用 opbase 提供的这些基础设施。你可以把 opbase 理解为"算子开发 SDK",把 ops-math 理解为"用这个 SDK 开发出来的一个具体的算子库"。

从 CANN 的开源进程来看,ops-math 是第一批开源的算子仓库之一。它的开源意味着开发者可以查看每个算子的具体实现(比如 exp 函数的逼近方法、随机数生成器的跳过技术具体怎么实现),也可以基于现有的实现做修改和扩展(比如针对自己的特定场景做精度调优或性能调优)。

这种"可查看、可修改"的特性,是闭源实现(比如 NVIDIA 的 cuDNN)做不到的。也是昇腾 CANN 开源生态的核心价值之一。


实战:在自定义算子开发中调用 ops-math

假设你正在用 Ascend C 写一个自定义的算子。这个算子的功能是"对输入张量做 sigmoid 变换,转换成 int8 类型输出"。

sigmoid 函数的公式是sigmoid(x) = 1 / (1 + exp(-x))。你需要做两件事:计算 exp(-x),再做逐元素的除法和加法。

如果不用 ops-math,你需要自己实现 exp 函数(用多项式逼近或者查表逼近),自己实现类型转换(处理内存布局、分块大小等细节)。这些实现加起来可能有几百行 Ascend C 代码,容易引入 bug(比如边界情况处理不当、精度不达标)。

如果用 ops-math,你可以直接调用它的 math 类和 conversion 类算子。

示例代码如下(这是 Ascend C 的算子开发模式,实际调用方式可能通过 AscendCL 接口):

#include"ops_math/math.h"// exp 算子的声明#include"ops_math/conversion.h"// cast 算子的声明__global__voidmy_sigmoid_int8(__gm__float*x,__gm__int8_t*y,intn){__aicore__int32_tblock_idx=GetBlockIdx();__aicore__int32_tblock_size=256;__aicore__int32_tstart=block_idx*block_size;__aicore__int32_tend=min(start+block_size,n);// 第一步:计算 exp(-x)// WHY: 不直接在这里写 exp 的实现,而是调用 ops_math::exp。// 因为 ops-math 的 exp 实现已经针对 Vector 单元优化过了(查表+向量化),// 自己写的版本几乎不可能比它更快。ops-math 的版本已经处理了// 边界情况(比如输入是 NaN 或 Infinity 的时候应该怎么处理),// 自己处理这些边界情况很容易漏掉某些情况。__aicore__ LocalTensor<float>x_local=AllocTensor<float>(block_size);__aicore__ LocalTensor<float>exp_neg_x=AllocTensor<float>(block_size);DataCopy(x_local,x+start,end-start);// 先对 x_local 做逐元素取负vmuls(x_local,x_local,-1.0f,end-start);// 调用 ops-math 的 expops_math::exp(exp_neg_x,x_local,end-start);// 第二步:计算 sigmoid(x) = 1 / (1 + exp(-x))__aicore__ LocalTensor<float>one=AllocTensor<float>(block_size);// 用 vadds 和 vdivs 做逐元素加法和除法vadds(one,exp_neg_x,1.0f,end-start);vdivs(exp_neg_x,one,exp_neg_x,end-start);// 这里 exp_neg_x 现在存的是 sigmoid 结果// 第三步:转换成 int8// WHY: 这里调用 ops_math::cast 而不是直接用 vconv 指令,因为 cast 算子// 会自动处理"溢出"情况(比如 sigmoid 的结果在 [0, 1] 范围内,// 转换成 int8 的时候需要做缩放和截断)。如果直接用 vconv,你需要自己// 写这些逻辑,容易出错。__aicore__ LocalTensor<int8_t>y_local=AllocTensor<int8_t>(block_size);ops_math::cast(y_local,exp_neg_x,end-start);DataCopy(y+start,y_local,end-start);}

这段代码的三个关键设计决策,都体现了"复用 ops-math 而不是自己实现"的思路:

第一,调用ops_math::exp而不是自己写 exp 逼近。原因是 ops-math 的版本已经优化过了,处理了各种边界情况。

第二,调用ops_math::cast而不是自己用 vconv 指令。原因是类型转换涉及溢出处理、内存布局适配等细节,ops-math 的版本已经把这些细节封装好了。

第三,分块大小和内存管理的逻辑仍然需要自己写(因为这是自定义算子的特有逻辑,ops-math 不可能替你决定你的算子应该怎么分块)。但核心的"计算"部分,尽量复用 ops-math 的实现。

在实际的算子开发中,这种"自定义算子 + 调用 ops-math"的模式非常常见。因为大部分算子的计算逻辑中,至少有一部分(类型转换、逐元素数学函数、随机数生成)是"通用"的,不需要每个算子都重新实现一遍。


ops-math 仓库地址:
https://atomgit.com/cann/ops-math

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询