CANN/PyPTO：kernel出参未写回问题

kernel函数出参未写回导致计算不生效【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto问题现象描述当前PyPTO框架用pypto.frontend.jit装饰的kernel函数不支持有返回值输出需要通过参数的形式传入并使用[:]等进行写回操作。如果直接使用等号赋值无法将数据写入输出Tensor中。示例代码pypto.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创建的未经初始化的随机值输出数据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]])原因分析在add_kernel函数内部执行y x 1时这里的y是函数的局部变量相当于创建了一个新的变量y它会覆盖传入参数y的引用。也就是说这行代码只是让函数内的y指向了x 1的新Tensor并不会修改外部传入的Tensory的内容。解决措施通过全切片操作符[:]将计算结果写入函数参数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也可以替换为y.move(x 1)或者y.assemble(x 1, [0, 0])。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考