PyPTO 数学函数 expands 详解：Tile 标量填充（splat）的原理与实战用法

人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载导读pypto_pro.language.expands是 CANN PyPTO 编程范式中用于将整块 Tile 的所有元素填充为同一个标量值的数学函数其语义等价于out[i] scalar在 Kernel 编程中常用于初始化负无穷 Tile构造因果掩码或零 Tile累加器清零。本文以 expands.md 为骨架完整梳理该函数的函数原型、参数约束、支持的数据类型与运行平台并结合仓库源码Python API 声明、IR 构建、类型与取值范围校验、ST/UT 测试深入讲解其底层实现原理最终给出可直接运行的基本用法与典型业务场景示例帮助读者在 PyPTO 矢量编程中正确、高效地使用标量填充操作。功能说明expands用于将目的 Tile 填充为指定标量值其计算语义为out[i] scalar 对 out 中的每一个元素 i在 PyPTO 的 Tile 编程模型中expands属于**矢量计算Vector/SIMD**操作需要在section_vector()矢量流水段中执行。该函数最常见的两个用途是初始化负无穷 Tile为 Softmax、注意力Attention等算子构造因果掩码时需要将掩码区域填充为-inf保证指数运算后该区域权重归零初始化零 Tile为归约、累加等算子预先清零或为不同数据类型的 Tile 提供统一的零值初始化入口。从源码注释看其本质是一次标量广播splat操作API 声明 中明确写为Fill Tile with a scalar (splat): out[i] scalar。产品支持情况expands的硬件支持情况与 PyPTO 的矢量流水线支持范围一致具体如下产品支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持说明以上支持矩阵来自 expands.md 的产品支持情况章节。在调用前请确认目标运行环境属于 Ascend 950 系列否则编译或运行阶段可能报错。函数原型pypto_pro.language.expands( out: Tile, scalar: Scalar, ) - None函数位于pypto_pro.language命名空间下通过import pypto_pro.language as pl后以pl.expands(...)方式调用。它是一个无返回值的原地写操作结果直接写入out指向的 Tile函数返回None。参数说明参数输入/输出说明out输出目的操作数Tile类型全部元素被填充为scalar值。数据类型支持DT_UINT8、DT_INT8、DT_UINT16、DT_INT16、DT_UINT32、DT_INT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32。位于 UBUnified Buffer或 L1 Buffer。scalar输入填充值。为整型或浮点型常量或运行时整型或浮点型标量表达式类型须与out元素类型兼容。关键点解读out的数据类型覆盖整型UINT8 至 UINT64与浮点型FP16/BF16/FP32共 11 种未包含 FP64 与 FP8 系列这与 IR 层的 dtype 白名单校验 中_EXPANDS_DTYPES的约束一一对应out必须位于 UB 或 L1 Bufferexpands是矢量单元执行的写操作操作数必须落在矢量单元可直接访问的存储空间。在示例中通过TileType(target_memorypl.MemorySpace.Vec)显式指定scalar的兼容性既可以是编译期字面量常量也可以是运行期的标量表达式如循环变量、symbolic_scalar计算结果。若字面量超出out元素类型的可表示范围编译期会抛出范围校验错误见下文范围校验。约束说明官方文档声明该函数无额外约束。需要补充说明的是无约束指函数本身没有形状、步长等特殊限制但实际使用中仍需满足 PyPTO 通用编程约束expands属于矢量操作必须在pl.section_vector()代码块内调用填充的 Tile 需通过pl.make_tile_group创建并保证地址、互斥锁mutex分配不与同流水段其他 Tile 冲突scalar的值必须能被out的元素类型精确表示否则在编译期触发范围检查错误。返回值说明无返回值None。填充结果直接写入out参数对应的目的 Tile。调用示例以下示例均可在 Ascend 950 系列设备上运行并与仓库 ST 测试 test_reduce_expand_fp32.py 中的expands_kernel保持一致。基本用法import pypto_pro.language as pl K_VALUE 2.0 pl.jit(auto_mutexTrue) def expands_kernel(dummy: pl.Tensor[[64, 64], pl.DT_FP32], out: pl.Tensor[[64, 64], pl.DT_FP32]): tt pl.TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_out pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) with pl.section_vector(): cur_out tile_out.current() pl.expands(cur_out, K_VALUE) pl.store(out, cur_out, [0, 0])运行结果文档示例输出输入数据K_VALUE2 输出数据out[[2 2 2 2 2 2 2 2 ...], [2 2 2 2 2 2 2 2 ...], [2 2 2 2 2 2 2 2 ...], [2 2 2 2 2 2 2 2 ...], ...]示例关键步骤拆解创建 Tile 类型TileType(shape[64, 64], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec)定义了一块 64×64 的 FP32 矢量 Tile绑定 Tile 组make_tile_group(typett, addrs0x0000, mutex_ids[0])将 Tile 绑定到 UB 起始地址0x0000并申请互斥锁0auto_mutexTrue会自动管理锁的申请与释放进入矢量流水段with pl.section_vector():声明后续指令在矢量单元执行取当前 Tile 并填充cur_out tile_out.current()获取当前实例pl.expands(cur_out, K_VALUE)将整块填充为2.0写回全局内存pl.store(out, cur_out, [0, 0])将 UB 中的 Tile 存回 Tensorout。对应的 ST 测试在 test_reduce_expand_fp32.py 中验证构造torch.full((64, 64), K_VALUE)作为参考结果与 Kernel 输出逐元素比对rtol1e-2, atol1e-2确认整块 Tile 均被填充为2.0。初始化负无穷 Tile# 初始化负无穷 Tile因果掩码 pl.expands(neg_inf_vec, NEG_INF)这是expands最具代表性的应用场景。在实现 Softmax 或 Attention 的因果掩码causal mask时需要把上三角区域填充为-inf使exp(-inf) 0从而在后续 softmax 归一化中屏蔽非法位置的注意力权重。与逐元素赋值相比expands一条指令即可完成整块 Tile 的填充避免了循环与地址计算开销。初始化零 Tile# 初始化零 Tile pl.expands(score_u16_row, 0)当需要对 Tile 清零时expands同样是最直接的方式。示例中score_u16_row为DT_UINT16类型的 Tile填充值0与无符号整型兼容。清零操作常用于累加器初始化、输出缓冲区预填充等场景此时expands可替代逐元素store显著减少指令数。源码级原理剖析1. Python API 声明层expands在 python/pypto_pro/language/_api.py 中声明并通过 python/pypto_pro/language/init.py 的__all__导出最终以pl.expands形式对用户可见。它是一个使用_api_decl修饰的声明式 API函数体只承载文档字符串真正的编译逻辑由解析器与 IR 构建器接管。2. 前端解析与操作数角色在 PyPTO 的前端解析流水线中expands被注册在 python/pypto_pro/language/parser/_op_pipeline.py 的操作数角色表中expands: [W, None],该表声明了expands的两个参数中第一个参数out是写操作数W第二个参数scalar是非 Tile 操作数。get_op_tile_rolespython/pypto_pro/language/parser/_op_pipeline.py在跨核同步分析等场景中依赖该角色信息确保写操作数的跨核依赖被正确追踪。此外python/pypto_pro/language/parser/_call_parser.py 在调用解析时对expands做了专门分支处理。3. IR 构建与双重校验expands的 IR 构建逻辑位于 python/pypto_pro/ir/op/block_ops.py 的_ir_expands并注册到 block 算子表block_ops.py 中的expands: OpSpec(builder_ir_expands)。其核心流程包含两道校验def _ir_expands(out: Expr, scalar: Expr, *, span: Span | None None) - Expr: out_dtype getattr(out.type, dtype, None) _check_dtype(expands, out_dtype, _EXPANDS_DTYPES) # The splat value lands in the out tile, so it has to be representable there. check_const_expr_fits_dtype(scalar, out_dtype, spanspan, apipl.expands) return _ir_core.create_op_call(block_ir_op(expands), [out, scalar], {}, span or _span())第一道dtype 白名单校验。_check_dtype(expands, out_dtype, _EXPANDS_DTYPES)检查out的 dtype 是否在文档列出的 11 种类型之内非法类型在编译期直接报错第二道标量取值范围校验。check_const_expr_fits_dtype(scalar, out_dtype, ...)检查字面量scalar是否能被out的元素类型表示。代码注释明确指出splat 值最终要落进 out Tile因此必须在其 dtype 可表示范围内——如果不做此检查标量只受 IR 存储带宽storage band限制而 IR 存储带宽远宽于窄位宽的 Tile dtype如 UINT8可能导致溢出后静默回绕。4. 取值范围校验的测试佐证取值范围校验的边界行为在 UT 与 ST 测试中被完整覆盖test_scalar_range_validation.py向DT_INT64的 Tile 填充INT64_MAX 1即9223372036854775808该值虽能落入 IR 的 uint64 存储带宽但超出 INT64 元素类型可表示范围编译期抛出OutOfRange错误信息为pl.expands: scalar operand must be representable in int64, i.e. in [-9223372036854775808, 9223372036854775807], got ...test_scalar_range_validation.py填充INT64_MAX合法边界值的 Kernel 可在 950 设备上真实运行输出与torch.full(..., INT64_MAX, dtypetorch.int64)完全一致rtol0, atol0test_scalar_range.pyUT 注释明确pl.expands: the splat value must fit the out tiles dtype从单测层面对该约束进行了系统验证。这些测试共同印证了编译期拦截非法填充值、边界合法值可正常下板执行的实现事实。与相近操作的对比在 PyPTO 算子库中expands与以下操作易混淆使用时需注意区分操作语义关键差异pl.expands将整块 Tile 填充为同一标量标量广播无源 Tile仅 2 个参数pl.full以标量值填充见 _op_pipeline.py同属填充类算子expands为矢量流水段内的 splat 写操作pl.fillpad填充 Tile 的 padding 区域仅填充边缘/填充区域非全量填充python/pypto_pro/language/_api.pyrow_expand/col_expand将归约结果沿行/列方向扩展回原形状源是归约输出 Tile目标是形状扩展非标量填充_op_pipeline.py需要特别说明的是row_expand/col_expand属于**归约扩展Reductions / expands**类别python/pypto_pro/language/_api.py 的 B8 分组注释其语义是维度广播与expands的标量广播是两类不同操作不应混淆。小结pypto_pro.language.expands是 PyPTO 矢量编程中最基础也最高效的标量填充原语用法上一条指令完成整块 Tile 的标量填充是初始化负无穷掩码、零 Tile 的首选方案约束上out支持 11 种常见整型/浮点 dtype必须位于 UB 或 L1 Bufferscalar须与out元素类型兼容实现上从 API 声明 到 IR 构建 再到 操作数角色表形成声明—解析—校验—建图的完整链路其中 dtype 白名单与标量可表示范围双重校验在编译期即拦截非法用法验证上ST 测试 test_reduce_expand_fp32.py 与范围校验测试 test_scalar_range_validation.py 共同保证了功能正确性与边界行为。在编写需要掩码、清零或常量预填充的 PyPTO Kernel 时优先考虑pl.expands它能让代码更简洁也让矢量流水线获得更高的指令效率。赞分享人工智能编译器模型编译高性能计算深度学习CANN【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PTO-ISA TEXPANDS 指令详解标量广播填充 Tile 的语义、约束与实战用法PTO ISA TEXPANDS 指令详解标量广播填充 Tile 的语义、约束与实战用法 TEXPANDS 是 CANN PTO ISA 虚拟指令集中的一条人工智能指令集算子库CANNAscendes-toolkit 数组批量填充指南fill 函数用法、负索引与原地修改原理es toolkit 数组批量填充指南fill 函数用法、负索引与原地修改原理 fill 是 es toolkit 提供的数组实用函数用于将数组全部或指定范前端后端CANN PyPTO SIMT 标量计算pl.simt.sin 正弦函数 API 详解与实战CANN PyPTO SIMT 标量计算pl.simt.sin 正弦函数 API 详解与实战 导读 pypto_pro.language.simt.sin 人工智能编译器模型编译高性能计算深度学习CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考