[phi] Update graph_send_recv OP #40509
Merged
Conversation
This file contains 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
|
✅ This PR's description meets the template requirements! |
|
Thanks for your contribution! |
DesmonDay
changed the title
wawltor
reviewed
Mar 15, 2022
| {"pool_type", "out_size"}, | ||
| {"Out", "Dst_count"}); | ||
| } | ||
|
|
| @@ -32,15 +32,23 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |||
| const DenseTensor& src_index, | |||
| const DenseTensor& dst_index, | |||
| const std::string& pool_type, | |||
| int64_t out_size, | |||
| DenseTensor* out, | |||
| DenseTensor* dst_count = nullptr) { | |||
| const int& index_size = src_index.dims()[0]; | |||
| ctx.template Alloc<T>(out); | |||
| T* p_output = out->data<T>(); | |||
| const auto& src_dims = x.dims(); | |||
| int64_t memset_size = 1; | |||
There was a problem hiding this comment.
这里有个疑问，这里预置shape的目的是什么了？是否要对设置的out_size与实际的size进行检查了？
There was a problem hiding this comment.
预置shape的目的是希望可以动态减小输出的实际形状。比如我输入X维度为(100, 128)，而实际要进行消息传递的结果只有节点编号最大为60，那我可以把输出shape设置为(61, 128)，其他结果则动态去掉。
这块设置检查我是想在python端处理。或者也可以在python端函数输入一个flag来确定是否要动态压缩，然后out_size我们直接自己设置为max(dst_index) + 1.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
PR types
Function optimization
PR changes
OPs
Describe
Add out_size attribute for graph_send_recv. After adding this attribute, we can set the shape of output according to actual need, which can bring us less computation cost. Take PGL GraphSage model on Reddit dataset as an example. After using this
out_sizeattribute, the model computation cost decreased. We can see the following picture for a look.