Custom PyTorch Operators on Ascend 910B

Environment The hardware environment this article is based on is the Ascend 910B3, and the software environment includes CANN 7.0-RC1, PyTorch 1.11.0, and Ascend PyTorch Adapter v5.0.rc3-pytorch1.11.0. The situation on other CANN and PyTorch versions may differ slightly. Registration Process Adding a Custom Operator in the Ascend PyTorch Adapter References: https://www.hiascend.com/document/detail/zh/canncommercial/70RC1/operatordev/Ascendcopdevg/atlas_ascendc_10_0045.html https://gitee.com/ascend/samples/tree/master/operator/AddCustomSample/FrameworkLaunch/PytorchInvocation Add the npu_add_custom function in torch_npu/csrc/aten/npu_native_functions.yaml: 1 2 custom: - func: npu_add_custom(Tensor x, Tensor y) -> Tensor # 添加的函数 Add the file AddCustomKernelNpu.cpp in torch_npu/csrc/aten/ops/op_api: ...

2023-11-14 · 2 min · Monsoon