PyPTO Kernel 出参写回指南解决 kernel 函数返回值不生效的问题【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读在 CANN PyPTO 框架中使用pypto.frontend.jit装饰的 kernel 函数不支持返回值所有计算结果都必须通过传入参数写回。很多开发者会在 kernel 内部直接使用等号赋值如y x 1导致计算结果无法写入输出 Tensor最终打印出torch.empty产生的未初始化随机值。本文从问题现象出发结合仓库源码剖析根因并给出[:]全切片写回、y.move()与y.assemble()三种等价的正确写法帮助你在编写 PyPTO kernel 时避免计算不生效的陷阱。问题现象kernel 计算悄悄失效输出全是随机值在 PyPTO 中一个典型的 kernel 写法如下用pypto.frontend.jit装饰函数通过pypto.set_vec_tile_shapes(4, 4)声明 vector 计算的 tile 形状再对输入执行逐元素运算。下面的示例代码试图实现逐元素加 1pypto.frontend.jit def add_kernel(x, y): pypto.set_vec_tile_shapes(4, 4) y x 1 # 此处会创建新的Tensor y torch.npu.set_device(0) x torch.ones(4, 4, dtypetorch.float32) y torch.empty(4, 4, dtypetorch.float32) add_kernel(pypto.from_torch(x), pypto.from_torch(y)) print(y) # 输出torch.empty创建的未经初始化的随机值运行后y中并不是期望的全 2 矩阵而是一堆随机值tensor([[2.0703e-19, 7.1833e22, 1.8502e28, 6.8608e22], [4.8011e30, 1.2123e25, 4.7418e30, 1.8465e25], [1.2122e25, 4.6114e24, 1.7836e31, 1.7591e22], [1.1306e24, 4.2245e-39, 6.8664e-44, 0.0000e00]])这些数值正是torch.empty(4, 4, dtypetorch.float32)创建时未初始化的内存内容。这说明 kernel 里的y x 1这行代码根本没有把结果写回外部传入的yTensor。原因分析等号赋值只是换引用不是写内存从 Python 语义看函数内的y是局部变量在add_kernel函数内部执行y x 1时这里的y是函数的局部变量相当于新建了一个变量y它会覆盖传入参数y的引用。也就是说这行代码只是让函数内的y指向了x 1计算得到的新 Tensor并不会修改外部传入 Tensory的内容。这一点与普通 Python 函数的传参行为一致x 1会构造一个新对象赋值语句将名字y重新绑定到这个新对象上而调用方持有的原对象不受任何影响。PyPTO kernel 同样遵循这一语义因此不能用等号赋值来写出计算结果。从 PyPTO 编译原理看kernel 的输出必须绑定到参数内存PyPTO 的pypto.frontend.jit装饰器会将 Python 函数解析并编译为 PTO 中间表示IR。从 python/pypto/frontend/parser/entry.py 的实现可以看到jit()装饰器支持jit与jit()两种用法其编译产物是一个JitCallableWrapper被装饰函数最终替换为编译后的 PTO 函数def jit( func: Optional[Callable] None, *, host_options: Optional[dict[str, Any]] None, codegen_options: Optional[dict[str, Any]] None, pass_options: Optional[dict[str, Any]] None, runtime_options: Optional[dict[str, Any]] None, verify_options: Optional[dict[str, Any]] None, debug_options: Optional[dict[str, Any]] None, new_ir: bool True, create_new_logical_tensor: bool False, ) - Union[Callable, Callable[[Callable], JitCallableWrapper]]: JIT decorator for compiling Python functions to PTO IR.在 PTO 的数据流图中kernel 的入参包括输出参数作为图的输入节点被记录而y x 1产生的是一个全新的中间 Tensor它没有被接线到参数y对应的存储上因此编译生成的 kernel 不会把结果搬运回y的内存地址。从源码结构看这正是 PyPTO 选择输出以参数形式传入、以写回方式生效这一设计的原因。一个小提示正确设置 device示例中通过torch.npu.set_device(0)指定 NPU 设备这是因为pypto.from_torch()转换的输入来自torch.npu上的 Tensor。若未设置 device 而直接运行通常会遇到设备相关的报错可参考 常见问题-未设置设备 一文排查。解决措施用写回语义替代等号赋值方式一全切片操作符[:]推荐通过全切片操作符[:]将计算结果写入函数参数y的原有内存空间pypto.frontend.jit def add_kernel(x, y): pypto.set_vec_tile_shapes(4, 4) y[:] x 1 # 将x1的结果写入函数参数y的原有内存空间 torch.npu.set_device(0) x torch.ones(4, 4, dtypetorch.float32) y torch.empty(4, 4, dtypetorch.float32) add_kernel(pypto.from_torch(x), pypto.from_torch(y)) print(y) # 输出x 1的结果此时输出符合预期tensor([[2., 2., 2., 2.], [2., 2., 2., 2.], [2., 2., 2., 2.], [2., 2., 2., 2.]])y[:] x 1表达的是把右值x 1的结果 Tensor整体写入左值y所指向的存储区域而不是重新绑定y这个名字。这与 PyTorch 中out[:] ...的 in-place 写回习惯完全一致也是 PyPTO kernel 中最通用、最直观的写回方式。方式二y.move(x 1)y[:] x 1也可以替换为y.move(x 1)。move是Tensor对象的内置方法其语义就是把另一个 Tensor 的数据搬移到当前 Tensor 的存储中。从 python/pypto/tensor.py 的源码可以看到其实现def move(self, other: Tensor) - None: if isinstance(other, Tensor): self._base.Move(other._base) else: raise FeError(TypeError(f{type(other).__name__} type cannot be moved to Tensor))move要求入参必须是Tensor否则会抛出FeError。在 PyPTO 的 IR 层Move操作会为当前 Tensor 绑定新的逻辑存储logical tensor使x 1的计算结果直接落到y的内存上。因此y.move(x 1)与y[:] x 1在效果上等价。方式三y.assemble(x 1, [0, 0])第三种等价写法是y.assemble(x 1, [0, 0])。assemble的语义是将一个小 Tensor 按指定的 offsets 拼装进一个大 Tensor特别适合分块计算后汇聚结果的场景。从 python/pypto/tensor.py 可以看到source_location def assemble(self, input: Tensor, offsets: List[Union[int, SymbolicScalar]]) - None: Assemble a small Tensor into a larger Tensor based on specified offsets. Args: input (Tensor): The small input tensor to be assembled into the larger tensor. offsets (Union[List[int], List[SymbolicScalar]]): Offset for placing the input tensor. example: s pypto.tensor((16, 16), pypto.DT_FP32) a pypto.tensor((2, 2), pypto.DT_FP32) s.assemble(a, [0, 0]) pypto.assemble(input, offsets, self)assemble方法的第一个参数是小 Tensorx 1的结果第二个参数[0, 0]表示放置的偏移量。当结果大小与输出y完全一致时偏移量取[0, 0]即为整块写回当 kernel 采用分块tiling计算时每一块结果可以通过不同的偏移量逐块写回大 Tensor 的对应区域。三种写回方式的底层汇聚点从 python/pypto/operation.py 可以看到assemble的函数形式最终都汇聚到 IR 层的Assemble调用source_location def assemble(*args, parallel: bool False) - None: if len(args) 3: src, offsets, dst args check_type(src, Tensor, assemble(): src) check_type(offsets, Sequence[Union[int, SymbolicScalar]], assemble(): offsets) check_type(dst, Tensor, assemble(): dst) pypto_impl.Assemble(src.base(), to_syms(offsets), dst.base(), parallel) ...assemble支持三种调用形态assemble(src, offsets, dst)、assemble([(src1, off1), (src2, off2), ...], dst)多块一次拼装以及dst.assemble(src, offsets)方法形式。多块拼装时还可通过parallelTrue允许并行执行。偏移量支持int与SymbolicScalar符号标量因此可以配合pypto.loop在动态循环中使用。三种写回方式的对比与选型建议写回方式写法适用场景是否支持偏移量全切片写回y[:] x 1整体结果直接写回输出参数最通用、最直观否整块move 搬移y.move(x 1)结果整体搬移到目标 Tensor 存储否整块assemble 拼装y.assemble(x 1, [0, 0])分块计算结果按偏移写回大 Tensor是支持int/SymbolicScalar偏移选型建议单次计算、结果与输出形状一致时优先使用y[:] x 1可读性最好语义上强调把结果搬进目标存储时可用y.move(x 1)涉及 tiling 分块、需要在循环中按[row_off, col_off]逐块写回时使用y.assemble(result, offsets)偏移量可以是符号标量以支持动态形状。深入理解几个相关的核心 APIpypto.set_vec_tile_shapes声明 vector 计算的 tile 形状示例中 kernel 第一行调用的pypto.set_vec_tile_shapes(4, 4)用于设置 vector 计算中每个维度的 tile 形状。其实现位于 python/pypto/_controller.py本质上是把 shape 写入当前编译 scopedef set_vec_tile_shapes(*shapes: int): set the tile shapes in vector computation concrete_shapes [it.concrete() if isinstance(it, SymbolicScalar) else it for it in shapes] pypto_impl.SetScope({vec_tile_shapes: concrete_shapes})set_vec_tile_shapes也接受SymbolicScalar并会先求值.concrete()说明 tile 形状同样支持符号化表达。对应的查询接口是pypto.get_vec_tile_shapes()。此外还有用于 cube矩阵乘计算的pypto.set_cube_tile_shapes(m, k, n, enable_split_k)和卷积相关的pypto.set_conv_tile_shapes(...)、pypto.set_convbp_input_tile_shapes(...)它们共同构成了 PyPTO 的 tiling 配置体系。pypto.from_torchTorch Tensor 与 PyPTO Tensor 的桥接示例中用pypto.from_torch(x)将torch.npu上的 Tensor 转换为 PyPTO Tensor 后再传入 kernel。该函数位于 python/pypto/converter.py支持name命名、dynamic_axis标记动态维度、tensor_format如TILEOP_ND/TILEOP_NZ和dtype等可选参数_count_calls def from_torch( tensor, name: str , dynamic_axis: Optional[List[int]] None, tensor_format: Optional[TileOpFormat] None, dtype: Optional[DataType] None, ): convert the input into a PyPTO Tensor转换得到的 PyPTO Tensor 保留了 shape、data_ptr内存地址、format、dtype 等属性。kernel 计算完成后外部torchTensor 即通过这份共享内存拿到结果这也是写回能生效的前提——PyPTO 与 Torch 共享底层存储。常见误区与最佳实践不要依赖 kernel 的返回值PyPTO 的 JIT kernel 设计上不返回结果返回的 Tensor 不会自动写回调用方。所有输出都应通过参数 写回操作完成。不要用等号重置输出参数y ...只改变函数内局部名字的绑定任何 PyTorch / PyPTO 场景下都不应以此作为写回手段。在循环内写回时注意写回位置tiling 循环中应使用y.assemble(result, offsets)将每一块结果按偏移写入偏移量支持SymbolicScalar可配合pypto.loop动态计算。区分新 Tensor与写回中间计算结果如tmp x 1是合法的局部 Tensor但最终必须通过[:]/move/assemble之一落到输出参数上kernel 才算真正产生外部可见的效果。验证输出是否生效运行后直接print(y)检查数值是否符合预期若输出仍是未初始化的随机值优先检查 kernel 内部是否使用了等号赋值写输出。延伸阅读本问题对应的原始 FAQ 文档kernel函数出参未写回导致计算不生效其他常见问题合集FAQ 索引包括未初始化 Tensor、view/assemble 循环依赖、tile shape 维度不匹配等问题Tensor.move与Tensor.assemble实现python/pypto/tensor.pypypto.assemble三种调用形态python/pypto/operation.pyjit装饰器与编译管线python/pypto/frontend/parser/entry.pyTensor / Torch 互转python/pypto/converter.pytile shape 设置接口python/pypto/_controller.py【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考