CANN pyasc 算子开发：asc.language.basic.shift_right 右移标量指令详解-北京尧图网络

CANN pyasc 算子开发asc.language.basic.shift_right 右移标量指令详解【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc导读本文围绕 CANN / pyasc 开源仓库中asc.language.basic.shift_right接口展开系统讲解如何在昇腾 AI 处理器上用 Python 原生语法对矢量Vector数据执行按标量右移运算。文章完整覆盖该接口的三种重载形式、参数语义、掩码与重复迭代repeat机制、底层 Ascend C 原型与 IR 指令映射并结合仓库单元测试给出可运行的算子样例。读完本文你将掌握shift_right的全部调用方式理解mask、repeat_times、UnaryRepeatParams对高维数据切分计算的控制原理并能在自己的算子中正确使用这一接口。一、接口定位pyasc 中的标量右移指令shift_right属于 pyasc 提供的矢量-标量二元运算vec-binary-scalar指令族与adds、muls、maxs、mins、shift_left等接口同属一族。其语义为对源操作数src中的每个元素执行右移操作右移的位数由输入参数scalar决定。在 pyasc 的源码中该接口定义于 vec_binary_scalar.py并通过 basic/init.py 以asc.shift_right的形式对外导出。调用该接口时会经由统一的vec_binary_scalar_op_impl分发逻辑见 basic/utils.py按参数形式分别映射到底层的ShiftRightL0Op、ShiftRightL1Op、ShiftRightL2Op三类 IR 指令。接口的函数签名如下asc.shift_right(dst: LocalTensor, src: LocalTensor, scalar: int | float, count: int, is_set_mask: bool True) - None asc.shift_right(dst: LocalTensor, src: LocalTensor, scalar: int | float, mask: int, repeat_times: int, repeat_params: UnaryRepeatParams, is_set_mask: bool True) - None asc.shift_right(dst: LocalTensor, src: LocalTensor, scalar: int | float, mask: List[int], repeat_times: int, repeat_params: UnaryRepeatParams, is_set_mask: bool True) - None三种重载形式分别对应三类底层指令适用于不同的计算场景重载形式mask 类型对应底层 IR 指令典型场景count 形式无ShiftRightL2Op对 Tensor 前 n 个元素整体计算mask 连续模式intShiftRightL0Op每次迭代按固定元素个数计算mask 连续模式mask 逐 bit 模式List[int]ShiftRightL1Op每次迭代按 bit 精确控制参与计算的元素从源码结构看shift_left与shift_right共享完全相同的实现骨架vec_binary_scalar.py仅底层指令不同ShiftLeftL0/L1/L2OpvsShiftRightL0/L1/L2Op因此本文关于参数机制的分析对shift_left同样适用。二、对应的 Ascend C 函数原型pyasc 的接口与 Ascend C 算子开发接口一一对应。shift_right对应的 Ascend C 函数原型如下template typename T, bool isSetMask true __aicore__ inline void ShiftRight(const LocalTensorT dstLocal, const LocalTensorT srcLocal, const T scalarValue, const int32_t calCount) template typename T, bool isSetMask true __aicore__ inline void ShiftRight(const LocalTensorT dstLocal, const LocalTensorT srcLocal, const T scalarValue, uint64_t mask[], const uint8_t repeatTimes, const UnaryRepeatParams repeatParams) template typename T, bool isSetMask true __aicore__ inline void ShiftRight(const LocalTensorT dstLocal, const LocalTensorT srcLocal, const T scalarValue, uint64_t mask, const uint8_t repeatTimes, const UnaryRepeatParams repeatParams)可以清楚看到三者的对应关系Python 的count参数对应 C 的calCountint32_t类型Python 的mask: int对应 C 的uint64_t mask单值Python 的mask: List[int]对应 C 的uint64_t mask[]数组Python 的repeat_times对应 C 的repeatTimesuint8_t类型repeat_params直接对应 C 的UnaryRepeatParams结构体。模板参数isSetMask对应 Python 的is_set_mask布尔参数用于控制是否在接口内部设置 mask 模式和 mask 值。三、参数说明参数类型说明is_set_maskbool是否在接口内部设置 mask 模式和 mask 值。默认True。dstLocalTensor目的操作数。类型为 LocalTensor支持的 TPosition 为 VECIN / VECCALC / VECOUT。srcLocalTensor源操作数。类型为 LocalTensor支持的 TPosition 为 VECIN / VECCALC / VECOUT。scalarint / float源操作数右移的位数。数据类型需要与目的操作数中的元素类型保持一致。countint参与计算的元素个数。maskint / List[int]用于控制每次迭代内参与计算的元素。repeat_timesint重复迭代次数。repeat_paramsUnaryRepeatParams元素操作控制结构信息。3.1 关于数据类型的说明从源码中可以看到右移操作本质上是位运算因此shift_right常用于整数类型数据。仓库单元测试 test_vector_scalar.py 中即使用asc.int16类型构造 src/dst Tensorx_local asc.LocalTensor(dtypeasc.int16, posasc.TPosition.VECIN, addr0, tile_size512) z_local asc.LocalTensor(dtypeasc.int16, posasc.TPosition.VECOUT, addr0, tile_size512) asc.shift_right(z_local, x_local, 1, count512)在 pyasc 的实现中scalar会被转换为与src相同的数据类型后再参与构建 IR见 basic/utils.py 中的scalar _mat(scalar, src.dtype).to_ir()这正是文档要求scalar数据类型需要与目的操作数中的元素类型保持一致的原因所在。3.2 关于UnaryRepeatParamsUnaryRepeatParams是控制重复迭代时数据地址步长的结构体定义于 core/types.py。其构造函数签名如下UnaryRepeatParams(dst_blk_stride: int 1, src_blk_stride: int 1, dst_rep_stride: int 8, src_rep_stride: int 8)四个参数的含义dst_blk_stride/src_blk_stride单次迭代repeat内相邻数据块block之间目的/源地址的步长以 block 为单位dst_rep_stride/src_rep_stride相邻两次迭代repeat之间目的/源地址的步长以 block 为单位。从源码结构看该结构体在构建 IR 时前两个字段以ui16类型、后两个字段以ui8类型写入UnaryRepeatParamsType属性因此步长取值应分别不超过 65535 与 255。四、调用示例详解4.1 Tensor 前 n 个数据计算count 模式asc.shift_right(dst, src, scalar, count512)该形式直接指定参与计算的元素个数count512即对src的前 512 个元素逐一右移scalar位并写入dst。此形式参数最少适合从头开始连续计算一段数据的场景对应底层ShiftRightL2Op指令见 basic/utils.py。4.2 Tensor 高维切分计算mask 连续模式mask 128 scalar 2 # repeat_times 4一次迭代计算128个数共计算512个数 # dst_blk_stride, src_blk_stride 1单次迭代内数据连续读取和写入 # dst_rep_stride, src_rep_stride 8相邻迭代间数据连续读取和写入 params asc.UnaryRepeatParams(1, 1, 8, 8) asc.shift_right(dst, src, scalar, maskmask, repeat_times4, repeat_paramsparams)该形式中mask 128表示每次迭代计算 128 个元素repeat_times 4表示重复 4 次迭代因此一共计算128 × 4 512个数。UnaryRepeatParams(1, 1, 8, 8)表示单次迭代内数据连续读取/写入blk_stride 1相邻迭代之间地址步长为 8 个 blockrep_stride 8对 int16 类型而言即连续排布。这种mask repeat模式适合高维 Tensor 的分块tile计算把一次大的数据搬运拆分为多次小粒度迭代从而充分利用矢量单元的流水线。4.3 Tensor 高维切分计算mask 逐 bit 模式mask [uint64_max, uint64_max] scalar 2 # repeat_times 4一次迭代计算128个数共计算512个数 # dst_blk_stride, src_blk_stride 1单次迭代内数据连续读取和写入 # dst_rep_stride, src_rep_stride 8相邻迭代间数据连续读取和写入 params asc.UnaryRepeatParams(1, 1, 8, 8) asc.shift_right(dst, src, scalar, maskmask, repeat_times4, repeat_paramsparams)当需要精确控制每次迭代中每个元素是否参与计算时使用List[int]形式的 mask。每个int元素对应一个 64 位掩码uint64_tbit 为 1 表示该元素参与计算。示例中mask [uint64_max, uint64_max]即两个2^64 - 1表示 128 个 bit 全部置 1与mask 128的效果等价——一次迭代同样计算 128 个元素。在 pyasc 实现中List[int]形式的 mask 会被逐个转换为uint64类型的 IR 值见 basic/utils.py对应底层ShiftRightL1Op指令。五、底层实现从 Python 调用到 IR 指令shift_right的完整调用链如下Python 入口asc.shift_right(dst, src, scalar, ...)进入 vec_binary_scalar.py参数分发调用vec_binary_scalar_op_implbasic/utils.py内部通过OverloadDispatcher依据mask/repeat_times/count等参数形态选择重载分支IR 构建三种重载分别调用builder.create_asc_ShiftRightL0Opmask 单值、create_asc_ShiftRightL1Opmask 列表、create_asc_ShiftRightL2Opcount 模式生成对应 IR 指令类型转换scalar按src.dtype做类型物化mask转为int64、repeat_times转为int8、count转为int32与 Ascend C 原型中的参数类型严格对应。在 IR 定义侧ShiftRight指令族定义于 include/ascir/Dialect/Asc/IR/Basic/OpVecBinaryScalar.tddefm ShiftRight : VecScalarL012Opshift_right, ShiftRight;其中L012后缀正表明该算子族同时拥有 L0mask 单值、L1mask 列表、L2count三种 IR 变体与 Python 侧三种重载一一对应。此外在标量寄存器指令Reg体系中还存在shift_rightsOpReg.td面向标量寄存器操作与本文的矢量shift_right属不同指令类别使用时注意区分。六、约束与注意事项使用shift_right时需关注以下约束操作数地址对齐dst 与 src 的地址对齐要求请参见《Ascend C 算子开发接口》中的通用说明和约束—通用地址对齐约束操作数地址重叠dst 与 src 的地址重叠约束请参考《Ascend C 算子开发接口》中的通用说明和约束—通用地址重叠约束数据类型一致性scalar的类型需与 dst 元素类型保持一致源码实现中会将其按src.dtype物化类型不一致可能导致类型转换异常TPosition 限制dst 与 src 仅支持 VECIN / VECCALC / VECOUT 三种存储位置步长取值范围UnaryRepeatParams的 blk_stride 字段为 ui16、rep_stride 字段为 ui8见 core/types.py取值需在对应类型范围内。七、结合单元测试验证接口行为仓库单元测试 test_vector_scalar.py 中test_shift_right_kernel完整验证了三种调用形式def test_shift_right_kernel(mock_launcher_run): asc.jit def shift_right_kernel(): x_local asc.LocalTensor(dtypeasc.int16, posasc.TPosition.VECIN, addr0, tile_size512) z_local asc.LocalTensor(dtypeasc.int16, posasc.TPosition.VECOUT, addr0, tile_size512) asc.shift_right(z_local, x_local, 1, count512) params asc.UnaryRepeatParams(1, 1, 8, 8) asc.shift_right(z_local, x_local, 1, mask512, repeat_times1, repeat_paramsparams) uint64_max 2**64 - 1 mask [uint64_max, uint64_max] asc.shift_right(z_local, x_local, 1, maskmask, repeat_times1, repeat_paramsparams) shift_right_kernel[1]() assert mock_launcher_run.call_count 1该测试展示了编写 pyasc 算子 Kernel 的标准姿势用asc.jit装饰 Kernel 函数通过asc.LocalTensor(dtype..., pos..., addr..., tile_size...)在指定存储位置声明局部 Tensor在 Kernel 内依次调用shift_right的 count、mask 单值、mask 列表三种形式通过kernel[1]()以 1 个核启动 Kernel。test_shift_left_kerneltest_vector_scalar.py结构与右移完全对称可作为对照参考。对于完整的算子开发流程含 GlobalTensor 输入、数据搬运、结果回写等可进一步参考 examples/01_add/add.py 与 examples/06_gelu/gelu.py 等示例。八、小结asc.language.basic.shift_right是 pyasc 矢量标量指令族中实现按标量右移的核心接口其价值在于三种重载适配不同计算粒度count适合整段连续计算mask: int适合定长分块迭代mask: List[int]适合 bit 级精确控制与 Ascend C 一一对应参数类型、模板语义、UnaryRepeatParams结构均与原 C 接口严格对齐Pyasc 开发者无需感知底层硬件细节实现链路清晰从 Python 入口到OverloadDispatcher分发、再到ShiftRightL0/L1/L2OpIR 指令每一层都有可查证的源码依据vec_binary_scalar.py、basic/utils.py、OpVecBinaryScalar.td。在实现量化缩放、定点数运算、位宽压缩等需要按位操作的算子时shift_right以及与之对称的shift_left是不可或缺的基础指令。【免费下载链接】pyasc本项目为Python用户提供算子编程接口支持在昇腾AI处理器上加速计算接口与Ascend C一一对应并遵守Python原生语法。项目地址: https://gitcode.com/cann/pyasc创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

CANN pyasc 算子开发：asc.language.basic.shift_right 右移标量指令详解

相关推荐

AgentMesh：为多 Agent 系统构建企业级信任与治理层的实战指南（Agent Governance Toolkit）

Python安装失败0x8007007E：MSI错误原理与修复指南

Node.js 18.14.2 (LTS) 发布公告深度解读：npm 9.5.0 升级、安全变更与全平台下载校验指南

更多资讯

3 步开卖数字产品：Gumroad 零门槛销售平台完整上手路径

STM32CubeIDE从下载到建工程：安装配置与高频问题排查全指南

租房数据分析毕设：从Python爬虫到Hadoop+Vue可视化

uni-app 团队 AI 协作指南：uni-agent 订阅套餐（企业版）席位、共享资源包与成员授权完全解析

ObserverList引用方式与设计取舍：从静态属性到接口抽象

3ds Max 2025建模精髓：从选择修改范围到数学建模的工程化实践

今日更新

AReaL 调试指南：从 Agent Workflow 验证到分布式训练死锁诊断

MATLAB实现GPS L1 C/A信号仿真与二维捕获验证

彻底搞懂ASCII、Unicode与UTF-8：从乱码根源到编码实战

本周精选

AI驱动的自动化专利分析平台技术解析

LeetCode-Go 题解：1252. Cells with Odd Values in a Matrix（奇数值单元格计数）双解法剖析

在线答疑系统Java毕设实战：Spring Boot前后端分离与状态流转

本月热门

基于YOLO的焊缝缺陷检测：从数据标注到工业部署全流程实践

分布式服务的决策记录

创业产品的首版范围

读完文章，想聊聊您的网站？