数据搬移操作【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. 概述HFusion 方言提供丰富的数据搬移操作包括广播、转置、拼接、填充、交织、翻转和聚集等。这些操作用于在不同形状和布局的张量间进行数据重排。源码参考HFusionOps.td、HFusionStructuredOps.td2. interleave2.1 功能将 n 个输入张量沿最后一维交织构造一个输出张量。当前仅支持 n2。2.2 操作签名操作数/结果类型说明inputVariadicAnyRankedTensor输入张量2 个形状相同outputAnyRankedTensor输出张量2.3 TraitsPure,CommutativeSameOperandsAndResultRankReifyRankedShapedTypeOpInterface2.4 MLIR 示例%output hfusion.interleave %a, %b : tensor128x256xf32, tensor128x256xf32 - tensor128x512xf323. deinterleave3.1 功能将一个输入张量沿最后一维反交织为两个张量。偶数索引元素进入第一个输出奇数索引元素进入第二个输出。3.2 操作签名操作数/结果类型说明inputAnyRankedTensor输入张量channelIndexI64Attr通道选择-1全部, 0偶数, 1奇数outputVariadicAnyRankedTensor输出张量3.3 channelIndex 行为值行为-1输出两个张量偶数索引 奇数索引0仅输出偶数索引通道1仅输出奇数索引通道3.4 约束输入张量最后一维大小必须是 2 的倍数3.5 MLIR 示例%even, %odd hfusion.deinterleave %input channel_index -1 : tensor128x512xf32 - tensor128x256xf32, tensor128x256xf32 %even_only hfusion.deinterleave %input channel_index 0 : tensor128x512xf32 - tensor128x256xf324. flip4.1 功能沿指定维度翻转张量。当前仅支持最后一维。4.2 操作签名操作数/结果类型说明inputAnyRankedTensor输入张量flip_axisI64Attr翻转轴outputAnyRankedTensor输出张量4.3 MLIR 示例%output hfusion.flip %input flip_axis 1 : tensor128x256xf32 - tensor128x256xf325. gather5.1 功能沿指定轴从源张量中聚集元素。对应triton.language.gather语义。5.2 操作签名操作数/结果类型说明srcAnyShaped源张量indexAnyShaped索引张量initAnyShaped初始值axisI64Attr聚集轴resultVariadicAnyTensor输出5.3 TraitsAllRanksMatch[src, index, init]AllShapesMatch[index, init]AllElementTypesMatch[src, init]BiShengIRAggregatedOpInterface支持 decomposeOperation5.4 语义给定 src:tensor16x16 和 index:tensor16x4axis1for i in 0 to 16: for j in 0 to 4: for k in 0 to 16: output[i][j] (index[i][j] k) ? src[i][k] : output[i][j]5.5 MLIR 示例%result hfusion.gather ins(%src, %index, %init : tensor16x16xf32, tensor16x4xi32, tensor16x4xf32) axis 1 - tensor16x4xf326. arange6.1 功能生成等差数列张量支持偏移和多维步长。6.2 操作签名操作数/结果类型说明offsetOptionalIndex偏移量默认 0stridesVariadicIndex各维步长initAnyShaped初始张量决定输出形状result_tensorOptionalAnyShaped输出张量6.3 语义3D arangearange[i, j, k] offset stride[0] * i stride[1] * j stride[2] * k7. 其他数据搬移操作以下操作由 Linalg 命名操作或 HFusion 结构化操作覆盖通过 OpDSL 或 Linalg 通用操作实现操作说明实现方式broadcast广播张量到目标形状Linalg 广播语义transpose转置张量维度Linalg 转置语义concat沿指定维度拼接张量Linalg 拼接语义pad填充张量边界Linalg 填充语义【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考