企业官网建设流程全解析

• 人工智能
• 编译器
• 模型编译
• 高性能计算
• 深度学习
• CANN

导读

vf.unsqueeze是 PyPTO（Parallel Tensor/Tile Operation 编程范式）SIMD 矢量函数（VF）体系中的掩码解压算子：它以 mask_reg 为操作对象，将掩码中每个有效位扩展为对应的寄存器 lane 值（有效为 1、无效为 0），从而把"哪些元素参与运算"的掩码信息物化为可直接存储、参与后续计算的 reg_tensor 数据。本文围绕该算子展开：先介绍其在 VF 计算体系中的定位与产品支持情况，再给出算法语义与掩码粒度原理，随后完整继承函数原型、参数与返回值说明，并基于当前仓库给出可直接运行的调用示例和源码级实现证据，帮助读者在 Ascend 950 系列产品上正确使用该能力。

产品支持情况

vf.unsqueeze的支持范围与当前仓库中其他 VF 掩码类算子（vf.create_mask、vf.update_mask、vf.mask_reg）保持一致，按昇腾产品型号划分如下：

说明：以上支持矩阵取自 unsqueeze.md 原文档，编写 VF 内核前请先确认目标设备的型号归属。

功能说明：把掩码"位"展开为寄存器"lane"

vf.unsqueeze的功能定义是：以 dst 为操作对象，根据 preg（mask_reg）进行解压缩。具体算法为：dst 的首位固定为 0；其后，当 mask[i] 对应掩码值为 1 时，dst[i] 的值为 dst[i-1] + 1；当 mask[i] 对应掩码值为 0 时，dst[i] 的值为 dst[i-1]。mask 的最高位被忽略，不参与统计。

其形式化表达为：

$$dstReg_i = \begin{cases} 1 & \text{if } mask_i = 1 \ 0 & \text{if } mask_i = 0 \end{cases}$$

从仓库源码看，该算子在 python/pypto_pro/language/_vf_api.py 中的声明为：

@staticmethod @_api_decl def unsqueeze(preg): """Unsqueeze mask bits into a register (vusqz instruction). Expands each mask bit into the corresponding register lane (1 for active, 0 for inactive). Args: preg: Mask register to unsqueeze Returns: Destination register (``RegTensor``) holding one lane per mask bit: 1 for active bits, 0 for inactive bits. """

可以确认：该算子底层对应vusqz指令，其作用是把掩码中每个有效位（1）展开为对应寄存器 lane 上的数值 1、无效位（0）展开为数值 0。示意图如下：

图中展示了 mask 中连续的 1 与 0 如何被依次映射到 dst 寄存器各 lane 的过程，最高位被忽略、首位输出为 0 的细节均体现其中。

掩码来源：create_mask 与 update_mask

preg参数必须是由 vf.create_mask 或 vf.update_mask 产生的 mask_reg，它作为 mask_reg 类型的参数直接传递给矢量计算 API，控制哪些元素参与运算：

• vf.create_mask(pattern, dtype)：按MaskPattern模式创建掩码。常用模式包括ALL（全有效）、ALLF（全无效）、VL1~VL128（最低 N 个元素有效，用于尾块处理）、H（最低一半有效）、Q（最低四分之一有效）、M3/M4（按倍数位置有效）等；
• vf.update_mask(scalar, dtype)：根据标量值scalar的比特位生成对应长度的有效位掩码，适合动态确定有效元素范围的场景。

两者生成的均为 mask_reg 类型，可直接作为vf.unsqueeze的输入。典型组合用法是：先用vf.ge、vf.eq等比较算子得到比较掩码，再用vf.unsqueeze把比较结果展开成 0/1 数据存出。

掩码粒度原理

mask_reg 总位宽固定为 256 bit，其粒度由创建时指定的 dtype 决定：每个数据元素对应的掩码位数随元素位宽变化，具体对应关系如下：

注意：dtype 决定的是掩码粒度（每多少 bit 对应一个数据元素），而非 mask_reg 本身的类型——mask_reg 类型始终不变。vf.unsqueeze展开时按"每个元素对应位"逐 lane 输出，因此掩码粒度直接影响输出寄存器中 1/0 的排布密度。

函数原型

unsqueeze(preg) -> dst

该接口为@pl.vector_function内部的 VF 算子，返回目标 reg_tensor。需要注意的是，仓库测试中还存在携带dtype的变体调用（如 python/tests/st/pypto_pro/frontend/vf_api/test_vf_basic_ops.py 中的vf.unsqueeze(cmp_mask, dtype=pl.DT_UINT32)），用于指定输出寄存器的数据类型，使用时以实际编译通过的签名与文档版本为准。

参数说明

约束说明

无。该算子无额外使用约束，但需遵守其输入输出类型的前提：

• preg必须为 mask_reg 类型，不能直接传入普通 reg_tensor；
• mask_reg 在使用前需被vf.create_mask/vf.update_mask（或比较算子）初始化，未初始化的掩码内容未定义；
• 调用时需在@pl.vector_function装饰的矢量函数体内进行，mask_reg 在函数结束后自动释放。

返回值说明

返回 dst 目标操作数，类型为 reg_tensor，存放掩码解压结果，支持的数据类型为：DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32。

reg_tensor 是 VF 计算的基本数据容器（寄存器总大小固定为 256 字节），不同 dtype 对应不同元素个数：8 bit 类型 256 个元素、16 bit 类型 128 个元素、32 bit 类型 64 个元素、64 bit 类型 32 个元素。vf.unsqueeze的输出通常配合vf.store_align写回 UB Tile，用于把掩码信息持久化或供后续标量/主机侧逻辑使用。

调用示例

原文档给出了一个完整的端到端示例，将vf.unsqueeze嵌入 kernel 流程：创建全有效掩码 → 从 Tile 对齐加载数据 → 对掩码执行 unsqueeze → 带掩码对齐存储。完整代码如下：

import os import pypto_pro.language as pl import torch import torch_npu @pl.vector_function def example_vf(src_tile, dst_tile): preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_UINT32) src = vf.load_align(src_tile, 0) dst = vf.unsqueeze(preg) vf.store_align(dst_tile, dst, preg) @pl.jit() def example_kernel( a: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_UINT32], out: pl.Tensor[[pl.DYNAMIC, pl.DYNAMIC], pl.DT_UINT32], ): tf = pl.TileType(shape=[1, 64], dtype=pl.DT_UINT32, target_memory=pl.MemorySpace.Vec) in_a_grp = pl.make_tile_group(type=tf, addrs=0x0, mutex_ids=[0]) in_a = in_a_grp.current() t_out_grp = pl.make_tile_group(type=tf, addrs=0x100, mutex_ids=[1]) t_out = t_out_grp.current() with pl.section_vector(): pl.load(in_a, a, [0, 0]) example_vf(in_a, t_out) pl.store(out, t_out, [0, 0]) def test_example(): device_id = int(os.environ.get("TILE_FWK_DEVICE_ID", 0)) device = f"npu:{device_id}" core_nums = 1 torch.npu.set_device(device) a = torch.randint(0, 100, [1, 64], device=device, dtype=torch.int32) out = torch.empty([1, 64], device=device, dtype=torch.int32) example_kernelNone, core_nums torch.npu.synchronize() assert out.shape == torch.Size([1, 64]) if __name__ == "__main__": test_example() print("PASSED")

示例解读

1. 矢量函数定义：@pl.vector_function装饰的example_vf内以隐式的vf命名空间调用 VF 算子；vf.unsqueeze(preg)在preg = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_UINT32)之后执行，此时掩码全为有效位，unsqueeze 结果对应位置均为 1；
2. kernel 组装：@pl.jit()装饰的example_kernel通过pl.make_tile_group在 Vector 内存空间（pl.MemorySpace.Vec）上声明输入/输出 Tile（地址 0x0 与 0x100 需不重叠），在pl.section_vector()段内依次完成 load → vector_function → store；
3. host 侧驱动：test_example通过TILE_FWK_DEVICE_ID环境变量选择 NPU 设备，example_kernel[None, core_nums]以单核启动 kernel，torch.npu.synchronize()等待执行完成后校验输出 shape。

进阶用法：比较掩码 → unsqueeze 物化

仓库测试 python/tests/st/pypto_pro/frontend/vf_api/test_vf_basic_ops.py 展示了更贴近实际业务的模式——先用比较算子产生掩码，再 unsqueeze 成可存储的 0/1 数据：

@pl.vector_function def _vf_kernel_15_cmp_unsqueeze_0(in_a, in_b, t_u0): preg_f32 = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_FP32) preg_u32 = vf.create_mask(pattern=pl.MaskPattern.ALL, dtype=pl.DT_UINT32) reg_a = vf.load_align(in_a, 0) reg_b = vf.load_align(in_b, 0) cmp_mask = vf.ge(reg_a, 0.0, preg_f32) # 比较得到 mask_reg：reg_a >= 0 的元素有效 reg_dst_u32 = vf.unsqueeze(cmp_mask, dtype=pl.DT_UINT32) # 解压为 0/1 数据 vf.store_align(t_u0, reg_dst_u32, preg_u32) # 写回 UB reg_i0, reg_i1 = vf.interleave(reg_a, reg_b)

这段代码验证了两个要点：vf.unsqueeze的输入可以是比较算子（如vf.ge）动态生成的 mask_reg；展开后的 0/1 结果可作为普通 reg_tensor 通过vf.store_align存出。这种"条件筛选结果物化"能力常用于统计、直方图、条件计数等需要把谓词结果落盘的场景。

实现原理补充：与 vusqz 指令的对应关系

从 python/pypto_pro/language/_vf_api.py 的接口声明可以看到，vf.unsqueeze的文档注释明确标注其对应vusqz 指令。该指令在硬件层面读取掩码寄存器（由SetVectorMask/SPR 体系维护的 256 bit 掩码），逐位展开为 lane 值：有效位写 1、无效位写 0，写入目的寄存器。结合 vf.mask_reg 的语义（比特位为 1 表示元素有效、比特位为 0 表示元素无效且目的位置零），可以推断：unsqueeze 本质上把掩码从"控制面"（决定运算是否发生）翻转为"数据面"（以 0/1 形式参与后续运算），这也是它与 vf.squeeze（掩码压缩）互为反向操作的原因。

典型使用场景总结

• 掩码可视化/落盘：将动态掩码（如尾块 VL 掩码、比较掩码）转换为 0/1 数据写回全局内存，便于调试与主机侧分析；
• 条件计数与统计：对 unsqueeze 结果继续做归约（如 vf.reduce_sum），即可获得有效元素个数；
• 谓词结果物化：配合vf.ge/vf.eq/vf.gt等比较算子，把逐元素比较结论编码为整型数据流，供后续分支或数据混合逻辑使用。

结合产品支持矩阵，该能力目前仅在 Ascend 950 系列（950PR/950DT）上可用；在 Atlas A2/A3 系列上如需实现类似语义，需评估使用标量侧逻辑或其他平台可用算子替代。

参考文档

• vf.mask_reg（掩码寄存器语义与粒度）
• vf.reg_tensor（目标寄存器类型）
• vf.create_mask（掩码创建）
• vf.update_mask（掩码更新）
• 高级计算章节索引
• vf.unsqueeze 接口声明源码
• vf.unsqueeze 比较掩码组合测试用例

• 人工智能
• 编译器
• 模型编译
• 高性能计算
• 深度学习
• CANN