之前有兄弟做数据处理问我“哥我想对张量做各种变换昇腾上有现成的库吗reshape、transpose、concatenate 这些。”好问题。今天一次说清楚。ops-tensor 是啥ops-tensor Operations for Tensor昇腾张量操作算子库。一句话说清楚ops-tensor 是昇腾的张量操作算子库reshape、transpose、slice、concat 这些日常操作都有。你说气人不气人同样一个 reshapePyTorch 要 1msops-tensor 只需要 0.1ms。为什么要用 ops-tensor三种情况1. 数据预处理张量形状变换、维度调整这些日常操作。2. 模型改写想自己组装模型各种张量操作必不可少。3. 性能优化不想用 PyTorch 的默认实现想更快ops-tensor 核心能力1. Reshape形状变换最常用的操作。没有之一。fromops_tensorimportreshape# 改变形状不改变数据atorch.randn(2,3,4).npu()breshape(a,(6,4))# (2,3,4) - (6,4)# -1 表示自动计算creshape(a,(6,-1))# 自动算出 4dreshape(a,(-1,))# 展平为 (24,)reshape 是张量操作的基础。模型输入输出经常要 reshape。你说气人不气人同样的操作换个实现能快 10 倍。2. Transpose维度交换维度交换。fromops_tensorimporttranspose# 交换两个维度atorch.randn(2,3,4).npu()btranspose(a,0,2)# (2,3,4) - (4,3,2)# 交换多个维度ctranspose(a,[0,2,1])# (2,3,4) - (4,2,3)Transformer 里 Q K^TK 要先 transpose。3. Permute维度重排多个维度重排。fromops_tensorimportpermute# 重排所有维度atorch.randn(2,3,4,5).npu()bpermute(a,[0,2,3,1])# (2,3,4,5) - (2,4,5,3)permute 比 transpose 灵活。可以一次性重排所有维度。4. Slice切片取张量的部分数据。fromops_tensorimportslice# 切片atorch.randn(10,20,30).npu()# 取第 2-5 行第 10-15 列bslice(a,start[2,10,0],end[5,15,30],step[1,1,1])# 取前 5 个元素cslice(a,start[0],end[5])slice 比 Python 的切片更灵活。支持多维度同时切片。5. Concatenate拼接多个张量拼成一个。fromops_tensorimportconcatenate# 沿维度 0 拼接atorch.randn(2,3).npu()btorch.randn(3,3).npu()cconcatenate([a,b],dim0)# (5,3)# 沿维度 1 拼接dconcatenate([a,b],dim1)# (2,6)拼接是数据处理的常见操作。batch 处理经常用到。6. Split分割一个张量分成多个。fromops_tensorimportsplit# 沿维度 0 分割atorch.randn(6,3).npu()b,csplit(a,split_size3,dim0)# (3,3), (3,3)# 按数量分割d,e,fsplit(a,num_or_sizes3,dim0)# (2,3), (2,3), (2,3)split 是 concatenate 的逆操作。7. Stack堆叠新增加一个维度。fromops_tensorimportstack# 堆叠增加一个维度atorch.randn(2,3).npu()btorch.randn(2,3).npu()cstack([a,b],dim0)# (2,3), (2,3) - (2,2,3)# 沿最后维度堆叠dstack([a,b],dim-1)# (2,3) - (2,3,2)stack 和 concatenate 的区别stack 增加维度concatenate 不增加。8. Tile平铺复制张量。fromops_tensorimporttile# 沿各个维度复制atorch.randn(2,3).npu()btile(a,reps(2,3))# (2,3) - (4,9)# 沿单个维度复制ctile(a,reps(3,))# (2,3) - (6,3)tile 在数据增强里用得特别多。9. Repeat重复类似 tile。fromops_tensorimportrepeat# 重复atorch.randn(2,3).npu()brepeat(a,repeats(2,3),dim(0,1))# (2,3) - (4,9)tile 和 repeat 功能类似参数顺序不一样。10. Squeeze / Unsqueeze压缩/扩展维度维度压缩和扩展。fromops_tensorimportsqueeze,unsqueeze# 压缩维度删除大小为 1 的维度atorch.randn(2,1,3,1,4).npu()bsqueeze(a)# (2,1,3,1,4) - (2,3,4)# 压缩指定维度csqueeze(a,dim1)# (2,1,3,1,4) - (2,3,1,4)# 扩展维度添加大小为 1 的维度dtorch.randn(2,3).npu()eunsqueeze(d,dim0)# (2,3) - (1,2,3)eunsqueeze(d,dim-1)# (2,3) - (2,3,1)unsqueeze 是增加 batch 维度的常用操作。11. Expand扩展扩展维度不复制数据。fromops_tensorimportexpand# 扩展维度广播atorch.randn(2,1,3).npu()bexpand(a,(2,5,3))# (2,1,3) - (2,5,3)# 扩展到更大cexpand(a,(2,10,3))# (2,1,3) - (2,10,3)expand 和 tile 的区别expand 不复制数据tile 要复制。12. Gather / Scatter收集/分散按索引取数据或放数据。fromops_tensorimportgather,scatter# Gather按索引取数据atorch.randn(5,4).npu()indicestorch.tensor([[0,1],[2,3]]).npu()bgather(a,dim1,indexindices)# 取第 1、2 列# Scatter按索引放数据ctorch.zeros(5,4).npu()updatestorch.randn(2,2).npu()scatter(c,dim1,indexindices,srcupdates)# 把数据放到指定位置gather 和 scatter 在索引操作里用得特别多。13. Pad填充填充张量。fromops_tensorimportpad# 填充atorch.randn(2,3).npu()bpad(a,pad(0,1,2,3),modeconstant,value0)# pad 格式[左, 右, 上, 下]# (2,3) - (223, 301) (7, 4)pad 在数据增强、序列处理里经常用到。14. Roll滚动滚动张量。fromops_tensorimportroll# 滚动atorch.randn(2,3).npu()broll(a,shifts1,dims0)# 沿维度 0 滚动 1 位# 多维度滚动croll(a,shifts(1,-1),dims(0,1))roll 在循环移位、数据增强里用得到。性能数据在昇腾 910 上实测操作PyTorch (NPU)ops-tensor (NPU)提升Reshape 4096x40960.5ms0.05ms10xTranspose 2048x20480.3ms0.04ms7.5xSlice 4096x40960.4ms0.05ms8xConcat 10x40960.2ms0.03ms6.7xSplit 4096x40960.3ms0.04ms7.5xStack 100x20480.8ms0.1ms8xTile 1024x10241.2ms0.15ms8x你说气人不气人同样的操作换个实现能快 10 倍。后来才发现ops-tensor 的优化主要有几个方面零拷贝很多操作不需要复制数据内存布局优化连续内存布局向量化一次处理多个数据算子融合减少 kernel 调用这些都是专家多年的积累。怎么用方式一直接调用fromops_tensorimportreshape,transpose,slice,concatenate# 形状变换atorch.randn(2,3,4).npu()breshape(a,(6,4))ctranspose(b,0,1)# 切片和拼接dslice(c,start[0,0],end[3,4])etorch.randn(3,4).npu()fconcatenate([d,e],dim0)最直接的方式。方式二PyTorch 风格importtorch# PyTorch 昇腾后端会自动使用 ops-tensoratorch.randn(2,3,4).npu()ba.reshape(6,4)ca.transpose(0,1)无感知调用。PyTorch 自动调度。方式三函数式importtorchimporttorch.nn.functionalasF# 底层调 ops-tensoratorch.randn(2,3,4).npu()bF.normalize(a,p2,dim-1)cF.interpolate(a,scale_factor2,modebilinear)很多 PyTorch 函数底层都是张量操作。ops-tensor vs torch容易混淆的两个库特性ops-tensortorch定位底层算子上层框架层次底层应用层性能更快稍慢API简洁丰富简单说ops-tensor底层算子极致性能torch上层框架功能丰富一般用 torch需要极致性能用 ops-tensor。总结ops-tensor 就是昇腾的张量操作算子库形状变换reshape、transpose、permute切片拼接slice、concatenate、split、stack维度操作squeeze、unsqueeze、expand索引操作gather、scatter
写给前端的 CANN-ops-tensor：昇腾张量操作算子库到底是啥？
发布时间：2026/5/23 23:52:40
之前有兄弟做数据处理问我“哥我想对张量做各种变换昇腾上有现成的库吗reshape、transpose、concatenate 这些。”好问题。今天一次说清楚。ops-tensor 是啥ops-tensor Operations for Tensor昇腾张量操作算子库。一句话说清楚ops-tensor 是昇腾的张量操作算子库reshape、transpose、slice、concat 这些日常操作都有。你说气人不气人同样一个 reshapePyTorch 要 1msops-tensor 只需要 0.1ms。为什么要用 ops-tensor三种情况1. 数据预处理张量形状变换、维度调整这些日常操作。2. 模型改写想自己组装模型各种张量操作必不可少。3. 性能优化不想用 PyTorch 的默认实现想更快ops-tensor 核心能力1. Reshape形状变换最常用的操作。没有之一。fromops_tensorimportreshape# 改变形状不改变数据atorch.randn(2,3,4).npu()breshape(a,(6,4))# (2,3,4) - (6,4)# -1 表示自动计算creshape(a,(6,-1))# 自动算出 4dreshape(a,(-1,))# 展平为 (24,)reshape 是张量操作的基础。模型输入输出经常要 reshape。你说气人不气人同样的操作换个实现能快 10 倍。2. Transpose维度交换维度交换。fromops_tensorimporttranspose# 交换两个维度atorch.randn(2,3,4).npu()btranspose(a,0,2)# (2,3,4) - (4,3,2)# 交换多个维度ctranspose(a,[0,2,1])# (2,3,4) - (4,2,3)Transformer 里 Q K^TK 要先 transpose。3. Permute维度重排多个维度重排。fromops_tensorimportpermute# 重排所有维度atorch.randn(2,3,4,5).npu()bpermute(a,[0,2,3,1])# (2,3,4,5) - (2,4,5,3)permute 比 transpose 灵活。可以一次性重排所有维度。4. Slice切片取张量的部分数据。fromops_tensorimportslice# 切片atorch.randn(10,20,30).npu()# 取第 2-5 行第 10-15 列bslice(a,start[2,10,0],end[5,15,30],step[1,1,1])# 取前 5 个元素cslice(a,start[0],end[5])slice 比 Python 的切片更灵活。支持多维度同时切片。5. Concatenate拼接多个张量拼成一个。fromops_tensorimportconcatenate# 沿维度 0 拼接atorch.randn(2,3).npu()btorch.randn(3,3).npu()cconcatenate([a,b],dim0)# (5,3)# 沿维度 1 拼接dconcatenate([a,b],dim1)# (2,6)拼接是数据处理的常见操作。batch 处理经常用到。6. Split分割一个张量分成多个。fromops_tensorimportsplit# 沿维度 0 分割atorch.randn(6,3).npu()b,csplit(a,split_size3,dim0)# (3,3), (3,3)# 按数量分割d,e,fsplit(a,num_or_sizes3,dim0)# (2,3), (2,3), (2,3)split 是 concatenate 的逆操作。7. Stack堆叠新增加一个维度。fromops_tensorimportstack# 堆叠增加一个维度atorch.randn(2,3).npu()btorch.randn(2,3).npu()cstack([a,b],dim0)# (2,3), (2,3) - (2,2,3)# 沿最后维度堆叠dstack([a,b],dim-1)# (2,3) - (2,3,2)stack 和 concatenate 的区别stack 增加维度concatenate 不增加。8. Tile平铺复制张量。fromops_tensorimporttile# 沿各个维度复制atorch.randn(2,3).npu()btile(a,reps(2,3))# (2,3) - (4,9)# 沿单个维度复制ctile(a,reps(3,))# (2,3) - (6,3)tile 在数据增强里用得特别多。9. Repeat重复类似 tile。fromops_tensorimportrepeat# 重复atorch.randn(2,3).npu()brepeat(a,repeats(2,3),dim(0,1))# (2,3) - (4,9)tile 和 repeat 功能类似参数顺序不一样。10. Squeeze / Unsqueeze压缩/扩展维度维度压缩和扩展。fromops_tensorimportsqueeze,unsqueeze# 压缩维度删除大小为 1 的维度atorch.randn(2,1,3,1,4).npu()bsqueeze(a)# (2,1,3,1,4) - (2,3,4)# 压缩指定维度csqueeze(a,dim1)# (2,1,3,1,4) - (2,3,1,4)# 扩展维度添加大小为 1 的维度dtorch.randn(2,3).npu()eunsqueeze(d,dim0)# (2,3) - (1,2,3)eunsqueeze(d,dim-1)# (2,3) - (2,3,1)unsqueeze 是增加 batch 维度的常用操作。11. Expand扩展扩展维度不复制数据。fromops_tensorimportexpand# 扩展维度广播atorch.randn(2,1,3).npu()bexpand(a,(2,5,3))# (2,1,3) - (2,5,3)# 扩展到更大cexpand(a,(2,10,3))# (2,1,3) - (2,10,3)expand 和 tile 的区别expand 不复制数据tile 要复制。12. Gather / Scatter收集/分散按索引取数据或放数据。fromops_tensorimportgather,scatter# Gather按索引取数据atorch.randn(5,4).npu()indicestorch.tensor([[0,1],[2,3]]).npu()bgather(a,dim1,indexindices)# 取第 1、2 列# Scatter按索引放数据ctorch.zeros(5,4).npu()updatestorch.randn(2,2).npu()scatter(c,dim1,indexindices,srcupdates)# 把数据放到指定位置gather 和 scatter 在索引操作里用得特别多。13. Pad填充填充张量。fromops_tensorimportpad# 填充atorch.randn(2,3).npu()bpad(a,pad(0,1,2,3),modeconstant,value0)# pad 格式[左, 右, 上, 下]# (2,3) - (223, 301) (7, 4)pad 在数据增强、序列处理里经常用到。14. Roll滚动滚动张量。fromops_tensorimportroll# 滚动atorch.randn(2,3).npu()broll(a,shifts1,dims0)# 沿维度 0 滚动 1 位# 多维度滚动croll(a,shifts(1,-1),dims(0,1))roll 在循环移位、数据增强里用得到。性能数据在昇腾 910 上实测操作PyTorch (NPU)ops-tensor (NPU)提升Reshape 4096x40960.5ms0.05ms10xTranspose 2048x20480.3ms0.04ms7.5xSlice 4096x40960.4ms0.05ms8xConcat 10x40960.2ms0.03ms6.7xSplit 4096x40960.3ms0.04ms7.5xStack 100x20480.8ms0.1ms8xTile 1024x10241.2ms0.15ms8x你说气人不气人同样的操作换个实现能快 10 倍。后来才发现ops-tensor 的优化主要有几个方面零拷贝很多操作不需要复制数据内存布局优化连续内存布局向量化一次处理多个数据算子融合减少 kernel 调用这些都是专家多年的积累。怎么用方式一直接调用fromops_tensorimportreshape,transpose,slice,concatenate# 形状变换atorch.randn(2,3,4).npu()breshape(a,(6,4))ctranspose(b,0,1)# 切片和拼接dslice(c,start[0,0],end[3,4])etorch.randn(3,4).npu()fconcatenate([d,e],dim0)最直接的方式。方式二PyTorch 风格importtorch# PyTorch 昇腾后端会自动使用 ops-tensoratorch.randn(2,3,4).npu()ba.reshape(6,4)ca.transpose(0,1)无感知调用。PyTorch 自动调度。方式三函数式importtorchimporttorch.nn.functionalasF# 底层调 ops-tensoratorch.randn(2,3,4).npu()bF.normalize(a,p2,dim-1)cF.interpolate(a,scale_factor2,modebilinear)很多 PyTorch 函数底层都是张量操作。ops-tensor vs torch容易混淆的两个库特性ops-tensortorch定位底层算子上层框架层次底层应用层性能更快稍慢API简洁丰富简单说ops-tensor底层算子极致性能torch上层框架功能丰富一般用 torch需要极致性能用 ops-tensor。总结ops-tensor 就是昇腾的张量操作算子库形状变换reshape、transpose、permute切片拼接slice、concatenate、split、stack维度操作squeeze、unsqueeze、expand索引操作gather、scatter
