Use unified cpp_extensions.cc to register operators with pybind #4435
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
Thanks for your contribution! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
pybind注册的自定义算子对比于直接使用paddle注册的算子在动态图场景具有更高的调度效率，因此推荐同时使用pybind和paddle进行算子注册。而当前只有GPU和XPU算子使用pybind注册自定义算子。
由于pybind注册算子时是硬件无关的，所以这里统一编写了cpp_extensions.cc实现不同硬件的pybind注册（除了XPU，XPU的编译脚本在custom_ops/xpu_ops/build.sh）。
通过给不同硬件添加
-DCUSTOM_XXX的TAG（如-DCUSTOM_GPU）区分不同设备下要注册那些算子，并在cpp_extensions.cc中根据TAG进行算子注册与方法声明。风险：缺少对应的硬件设备，验证GPU以外的算子正确性。