kernel函数出参未写回导致计算不生效

问题现象描述

当前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, dtype=torch.float32)
y = torch.empty(4, 4, dtype=torch.float32)
add_kernel(pypto.from_torch(x), pypto.from_torch(y))
print(y)  # 输出torch.empty创建的未经初始化的随机值

输出数据：

tensor([[2.0703e-19, 7.1833e+22, 1.8502e+28, 6.8608e+22],
        [4.8011e+30, 1.2123e+25, 4.7418e+30, 1.8465e+25],
        [1.2122e+25, 4.6114e+24, 1.7836e+31, 1.7591e+22],
        [1.1306e+24, 4.2245e-39, 6.8664e-44, 0.0000e+00]])

原因分析

在add_kernel函数内部执行y = x + 1时，这里的y是函数的局部变量（相当于创建了一个新的变量y），它会覆盖传入参数y的引用。也就是说，这行代码只是让函数内的y指向了x + 1的新Tensor，并不会修改外部传入的Tensor y的内容。

解决措施

通过全切片操作符[:]，将计算结果写入函数参数y的原有内存空间。

示例代码：

@pypto.frontend.jit
def add_kernel(x, y):
    pypto.set_vec_tile_shapes(4, 4)
    y[:] = x + 1  # 将x+1的结果写入函数参数y的原有内存空间

torch.npu.set_device(0)
x = torch.ones(4, 4, dtype=torch.float32)
y = torch.empty(4, 4, dtype=torch.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])。