| 添加aclgraph+<<<>>> 示例代码 Co-authored-by: wangkai<wangkai579@huawei.com> # message auto-generated for no-merge-commit merge: !4019 merge master into master 添加aclgraph+<<<>>> 示例代码 Created-by: mihudan Commit-by: wangkai Merged-by: ascend-robot Description: <!-- Thanks for sending a pull request! --> **What type of PR is this?** > Uncomment only one /kind <> line, hit enter to put that in a new line, and remove leading whitespaces from that line: > > /kind bug /kind task > /kind feature **What does this PR do / why do we need it**: 添加aclgraph 与<<<>>>联调的demo测试代码,进行验证并供用户参考。 展示了如何使用PyTorch的torch.librar以及pybind两张方式注册自定义算子,通过<<<>>>内核调用符调用核函数,并适配aclgraph使用该自定义算子,以简单的Add算子和三角函数计算的原地算子为例,实现aclgraph下自定义算子的调用。 展示了3种aclgraph的使能方式,通过对比NPU输出与CPU标准加法结果来验证自定义算子的数值正确性。 1. torch.npu.NPUGraph() 2. torch.npu.make_graphed_callables 3. backend="npugraph_ex" ## 算子描述 ### Add算子 - 算子功能: Add算子实现了两个数据相加,返回相加结果的功能。对应的算子原型为: ascendc_add(Tensor x, Tensor y) -> Tensor - 算子规格: <table> <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">add_custom</td></tr> </tr> <tr><td rowspan="3" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr> <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr> <tr><td align="center">y</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr> </tr> </tr> <tr><td rowspan="1" align="center">算子输出</td><td align="center">z</td><td align="center">8 * 2048</td><td align="center">int</td><td align="center">ND</td></tr> </tr> </table> ### 原地三角函数算子 - 算子功能: 该算子入参为x, out_sin ,out_cos, 算子调用后,out_sin会被原地修改为sin(x)计算结果,out_cos会被原地修改为cos(x)计算结果,返回值tan(x)计算结果。对应的算子原型为: ascendc_trig(Tensor x, Tensor(a!) out_sin, Tensor(b!) out_cos) -> Tensor - 算子规格: <table> <tr><td rowspan="1" align="center">核函数名</td><td colspan="4" align="center">trig_inplace_custom</td></tr> </tr> <tr><td rowspan="4" align="center">算子输入</td><td align="center">name</td><td align="center">shape</td><td align="center">data type</td><td align="center">format</td></tr> <tr><td align="center">x</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> <tr><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> </tr> </tr> <tr><td rowspan="3" align="center">算子输出</td><td align="center">out_sin</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> <tr><td align="center">out_cos</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> <tr><td align="center">out_tan</td><td align="center">8 * 2048</td><td align="center">float</td><td align="center">ND</td></tr> </tr> </table> ## 代码实现介绍 - 以Add算子为例,样例在*.asc文件中定义了一个名为ascendc_ops的命名空间,并在其中注册了ascendc_add函数。在ascendc_add函数中通过c10_npu::getCurrentNPUStream()函数获取当前NPU上的流,并通过内核调用符<<<>>>调用自定义的Kernel函数add_custom,在NPU上执行算子。 c++ add_custom<<<blockDim, nullptr, aclStream>>>(xGm, yGm, zGm, totalLength); - PyTorch提供TORCH_LIBRARY_FRAGMENT宏作为自定义算子注册的核心接口,用于创建并初始化自定义算子库,注册后在Python侧可以通过torch.ops.namespace.op_name方式进行调用,例如: c++ TORCH_LIBRARY_FRAGMENT(ascendc_ops, m) { m.def(ascendc_add"(Tensor x, Tensor y) -> Tensor"); } - TORCH_LIBRARY_IMPL用于将算子逻辑绑定到特定的DispatchKey(PyTorch设备调度标识)。针对NPU设备,需要将算子实现注册到PrivateUse1这一专属的DispatchKey上,例如: c++ TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m) { m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add)); } - 注册Meta函数: 注册Meta函数使faketensor流程正常工作,在使用fx, compile等功能涉及,注册代码如下: c++ TORCH_LIBRARY_IMPL(ascendc_ops, Meta, m) { m.impl("ascendc_add", &add_impl_meta); } **Special notes for your reviewers**: See merge request: Ascend/op-plugin!4019 | 5 个月前 |