In-place算子插件化入图样例

假设您已经实现了In-place算子开发和入图样例中的PyTorch算子torch.ops.npu.my_inplace，并能在Eager模式下正常调用。

- func: my_inplace(Tensor(a!) x, Tensor(b!) y) -> ()

插件化适配时的Meta推导函数、函数化转换及Converter与示例实现完全一致，唯一的区别是可以将实现写在单独的py文件中，并在模型执行前加载。

完整的代码实现如下所示：

import torch
import torch_npu
import torchair

m = torch.library.Library("npu", "FRAGMENT")
m.define("my_inplace_functional(Tensor x, Tensor y) -> (Tensor, Tensor)")
# 实现my_inplace_functional的NPU实现（注：NPU通过PrivateUse1设备扩展接入PyTorch）
@torch.library.impl(m, "my_inplace_functional", "PrivateUse1")
def my_inplace_functional_npu(x, y):
    x_clone = x.clone()
    y_clone = y.clone()
    torch.ops.npu.my_inplace(x_clone, y_clone)
    return x_clone, y_clone

# 实现Inplace算子的Meta推导
@torch.library.impl(m, "my_inplace", "Meta")
def my_inplace_functional_npu(x, y):
    pass

# 实现对应非Inplace算子的Meta推导
@torch.library.impl(m, "my_inplace_functional", "Meta")
def my_inplace_functional_npu(x, y):
    return torch.empty_like(x), torch.empty_like(y)

# 实现Inplace算子的函数化实现
@torch.library.impl(m, "my_inplace", "Functionalize")
def my_inplace_functional_npu(x, y):
    x_out, y_out = torch.ops.npu.my_inplace_functional(x, y)
    x.copy_(x_out)
    y.copy_(y_out)

# 仅当需要使用max-autotune模式，且对应Ascend C算子基于算子工程开发时需要
from torchair._ge_concrete_graph import ge_apis as ge
@torchair.register_fx_node_ge_converter(torch.ops.npu.my_inplace_functional.default)
def convert_npu_my_inplace(x, y):
     x = torchair.ge.Clone(x)
     y = torchair.ge.Clone(y)
     return torchair.ge.custom_op("MyInplace", x, y)