CANN/asc-devkit张量广播API文档

发布时间:2026/7/17 12:31:40

CANN/asc-devkit张量广播API文档 Broadcast【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit产品支持情况Ascend 950PR/Ascend 950DT支持Atlas A3 训练系列产品/Atlas A3 推理系列产品支持Atlas A2 训练系列产品/Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品AI Core支持Atlas 推理系列产品Vector Core不支持Atlas 训练系列产品不支持Kirin X90支持Kirin 9030支持功能说明将输入按照输出shape进行广播。比如A的shape为(2,1)广播的目标shape为(2,16)则会将原来的一列扩展为相同的16列。输入数据 [[ 1] [ 2]] 输出数据 [[ 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1] [ 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2]]实现原理以float类型ND格式[m, 1]广播到[m, k]为例描述Broadcast高阶API内部算法框图如下图所示。图1Broadcast算法框图![](https://raw.gitcode.com/cann/asc-devkit/raw/1cd1d1192e01ae4a8684cd07552c35e13e5c4bdf/docs/zh/api/figures/Broadcast算法框图.png Broadcast算法框图?utm_sourcegitcode_repo_files)计算过程分为如下几步均在Vector上进行brcb步骤将每个元素广播为一个datablockCopy步骤将每个datablock均复制为多个datablockk对齐场景下即为结果y对于k非对齐的场景再使用GatherMask截取[m, k]个元素其中k表示k向上对齐32B的大小。函数原型通过sharedTmpBuffer入参传入临时空间template typename T, int32_t dim, int32_t axis, bool isReuseSource false __aicore__ inline void Broadcast(const LocalTensorT dstLocal, const LocalTensorT srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim], LocalTensoruint8_t sharedTmpBuffer)接口框架申请临时空间template typename T, int32_t dim, int32_t axis, bool isReuseSource false __aicore__ inline void Broadcast(const LocalTensorT dstLocal, const LocalTensorT srcLocal, const uint32_t dstShape[dim], const uint32_t srcShape[dim])支持动态shape仅在Ascend 950PR/Ascend 950DT上支持。templateclass T, int constRank-1, uint32_t* constDstShape nullptr, uint32_t* constSrcShape nullptr, bool constSrcInnerPad false __aicore__ inline void Broadcast(const LocalTensorT dst, const LocalTensorT src, const uint32_t* dstShape, const uint32_t* srcShape, BroadcastTiling* tiling)该接口需要额外的临时空间来存储计算过程中的中间变量。临时空间支持开发者通过sharedTmpBuffer入参传入和接口框架申请两种方式。通过sharedTmpBuffer入参传入使用该tensor作为临时空间进行处理接口框架不再申请。该方式开发者可以自行管理sharedTmpBuffer内存空间并在接口调用完成后复用该部分内存内存不会反复申请释放灵活性较高内存利用率也较高。接口框架申请临时空间开发者无需申请但是需要预留临时空间的大小。通过sharedTmpBuffer传入的情况开发者需要为tensor申请空间接口框架申请的方式开发者需要预留临时空间。临时空间大小BufferSize的获取方式如下通过GetBroadCastMaxMinTmpSize中提供的接口获取需要预留空间范围的大小。另外提供了一个Kernel侧计算Tiling的接口针对Broadcast的实现计算Tiling获取Tiling结果。该接口的模板参数功能与支持动态shape的Broadcast接口模板参数相同其余参数说明请参见表5。kernel侧tiling计算接口仅在Ascend 950PR/Ascend 950DT上支持。templateclass T, int constRank-1, uint32_t* constDstShape nullptr, uint32_t* constSrcShape nullptr __aicore__ inline void GetBroadcastTilingInfo(uint32_t rank, const uint32_t* dstShape, const uint32_t* srcShape, bool srcInnerPad, BroadcastTiling tiling)参数说明表1模板参数说明参数名称功能T操作数的数据类型。不同型号支持的数据类型请参考支持的数据类型。dim输入/输出tensor的维度目前仅支持1维和2维。axis要广播的维度目前仅支持0和1。参数取值为0代表要广播第1维取值为1代表要广播第2维。isReuseSource是否允许修改源操作数。该参数预留传入默认值false即可。表2支持动态shape接口的模板参数说明参数名称功能T操作数的数据类型目前支持int8_t、uint8_t、int16_t、uint16_t、half、bfloat16_t、int32_t、uint32_t、float、int64_t、uint64_t。constRank输入/输出tensor的维度数目。默认值-1为动态shape场景内部按照GetBroadcastTilingInfo接口中的参数rank计算constRank大于0时必须与GetBroadcastTilingInfo接口中的参数rank取值相同。rank当前支持的范围为[1, 9]。constDstShape输出tensor的shape。uint32_t类型的数组。该数组中任一维度取值为0表示该维度为动态场景该维度实际取值由参数dstShape对应维度取值决定。该数组中任一维度取值大于0表示该维度为静态场景该维度取值与参数dstShape中对应维度取值相同。该参数预留传入默认值nullptr即可。constSrcShape输入tensor的shape。uint32_t类型的数组。该数组中任一维度取值为0表示该维度为动态场景实际shape由参数srcShape决定。该数组中任一维度取值大于0表示该维度为静态场景与参数srcShape中对应维度取值相同。该参数预留传入默认值nullptr即可。constSrcInnerPad表示输入的最后一维srcShape[rank-1]是否32B对齐其中rank为输入/输出tensor的维度数目。该参数预留传入默认值false即可。表3接口参数说明参数名称输入/输出描述dstLocal输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。srcLocal输入源操作数。源操作数的数据类型需要与目的操作数保持一致。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。dstShape输入输出tensor的shapeuint32_t类型的数组长度为1或者2输入/输出的shape维度数目必须一致。srcShape输入输入tensor的shapeuint32_t类型的数组长度为1或者2输入/输出的shape维度数目必须一致。sharedTmpBuffer输入临时缓存。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。用于Broadcast内部复杂计算时存储中间变量由开发者提供。临时空间大小BufferSize的获取方式请参考GetBroadCastMaxMinTmpSize。表4支持动态shape接口的参数说明参数名称输入/输出描述dst输出目的操作数。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。src输入源操作数。源操作数的数据类型需要与目的操作数保持一致。类型为LocalTensor支持的TPosition为VECIN/VECCALC/VECOUT。dstShape输入输出tensor的shapeuint32_t类型的数组长度取值范围为[1, 9]。输入/输出的shape维度数目必须一致且满足条件dstShape[i] srcShape[i]。srcShape输入输入tensor的shapeuint32_t类型的数组长度取值范围为[1, 9]。输入/输出的shape维度数目必须一致且满足条件dstShape[i] srcShape[i]。当srcShape[i]的值为1且dstShape[i]不等于srcShape[i]时表示i轴为广播轴。tiling输入Broadcast接口所需的Tiling信息。BroadcastTiling*类型通过调用Kernel侧的tiling计算接口GetBroadcastTilingInfo获取。表5kernel侧tiling计算接口参数说明| 参数名称 | 输入/输出 | 功能 | | --- | --- | --- | | rank | 输入 | 输入/输出tensor的维度数目目前支持的取值为[1, 9]。 | | dstShape | 输入 | 输出tensor的shapeuint32_t类型的数组长度取值范围为[1, 9]。输入/输出的shape维度数目必须一致且满足条件dstShape[i] srcShape[i]。 | | srcShape | 输入 | 输入tensor的shapeuint32_t类型的数组长度取值范围为[1, 9]。输入/输出的shape维度数目必须一致且满足条件dstShape[i] srcShape[i]。当srcShape[i]的值为1且dstShape[i]不等于srcShape[i]时表示i轴为广播轴。 | | srcInnerPad | 输入 | 表示输入的最后一维srcShape[rank-1]是否32B对齐。当前仅支持取值为false。 | | tiling | 输出 | 计算返回的Tiling信息。BroadcastTiling类型。 |返回值说明无约束说明操作数地址对齐要求请参见通用地址对齐约束。不支持源操作数与目的操作数地址重叠。当前仅支持ND格式的输入不支持其他格式。dim目前仅支持1或者2 axis目前仅支持0或者1。对于Atlas 推理系列产品AI Core在dim2且axis1时srcShape[0]必须为32B对齐即当输入/输出Tensor的维度为2维且广播的维度为1时输入Tensor的第0维数据需为32字节的倍数。在dim2axis0时要求srcShape[1]必须32B对齐。对于Ascend 950PR/Ascend 950DT输入/输出tensor支持的维度数目即rank支持的取值范围为[1, 9]。支持的数据类型Ascend 950PR/Ascend 950DT支持的数据类型为b8、b16、b32、b64位宽对应的数据类型具体数据类型请参考不同位宽对应的数据类型。Atlas A3 训练系列产品/Atlas A3 推理系列产品支持的数据类型为int8_t、uint8_t、half、float。Atlas A2 训练系列产品/Atlas A2 推理系列产品支持的数据类型为int8_t、uint8_t、half、float。Atlas 推理系列产品AI Core支持的数据类型为int8_t、uint8_t、half、float。Kirin X90支持的数据类型为int8_t、uint8_t、half、float。Kirin 9030支持的数据类型为int8_t、uint8_t、half、float。调用示例// dstLocal输出Tensor // srcLocal输入Tensor const uint32_t srcShape[2] {1, 16}; // 源数据shape const uint32_t dstShape[2] {16, 16}; // broadcast数据shape AscendC::Broadcastfloat, 2, 1( dstLocal, srcLocal, dstShape, srcShape); // broadcast数据类型为float源数据是2维 broadcast第1维度结果示例如下输入数据srcLocal: [[ 1] [ 2] [ 3] [ 4] [ 5] [ 6] [ 7] [ 8] [ 9] [10] [11] [12] [13] [14] [15] [16]] dim2 axis1 输出数据dstLocal: [[ 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1] [ 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2] [ 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3] [ 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4] [ 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5 5] [ 6 6 6 6 6 6 6 6 6 6 6 6 6 6 6 6] [ 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7 7] [ 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8] [ 9 9 9 9 9 9 9 9 9 9 9 9 9 9 9 9] [10 10 10 10 10 10 10 10 10 10 10 10 10 10 10 10] [11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11] [12 12 12 12 12 12 12 12 12 12 12 12 12 12 12 12] [13 13 13 13 13 13 13 13 13 13 13 13 13 13 13 13] [14 14 14 14 14 14 14 14 14 14 14 14 14 14 14 14] [15 15 15 15 15 15 15 15 15 15 15 15 15 15 15 15] [16 16 16 16 16 16 16 16 16 16 16 16 16 16 16 16]]【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址: https://gitcode.com/cann/asc-devkit创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

相关新闻