New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[geometric]Add paddle.geometric.send_u_recv API #44580
Changes from 5 commits
07005a6
00e8cdd
2c8d2cb
2131886
a47b281
ad1692c
f8fac9a
978f139
bd2da52
b11a60c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -58,6 +58,10 @@ class GraphSendRecvOpMaker : public framework::OpProtoAndCheckerMaker { | |
"The input tensor with data type float32, float64, int32, int64."); | ||
AddInput("Src_index", "The source index tensor."); | ||
AddInput("Dst_index", "The destination index tensor."); | ||
AddInput("OutSizeTensor", | ||
"(Tensor<int>, optional). The 0th dimension of the output." | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Tensor, optional 为什么名字要加Tensor,是不是Out_size更好 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Tensor, optional的写法是参考其他类似的 OP,这里不改动;另外已经改成Out_size了。 |
||
"It has a higher priority than Attr(out_size).") | ||
.AsDispensable(); | ||
AddOutput("Out", "Output tensor of graph_send_recv op."); | ||
AddOutput("Dst_count", | ||
"Count tensor of Dst_index, mainly for MEAN pool_type.") | ||
|
@@ -68,12 +72,12 @@ class GraphSendRecvOpMaker : public framework::OpProtoAndCheckerMaker { | |
"tensors of Dst_index.") | ||
.SetDefault("SUM") | ||
.InEnum({"SUM", "MEAN", "MIN", "MAX"}); | ||
AddAttr<int64_t>( | ||
AddAttr<std::vector<int64_t>>( | ||
"out_size", | ||
"(int64_t, default 0)" | ||
"(vector<int64_t>, default {0})" | ||
"Define the first dimension of Output tensor." | ||
"If set default 0, then the shape of Out is the same with X.") | ||
.SetDefault(0); | ||
"If set default {0}, then the shape of Out is the same with X.") | ||
.SetDefault({0}); | ||
AddComment(R"DOC( | ||
Graph Learning Send_Recv combine operator. | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -886,7 +886,7 @@ | |
backward : gelu_grad | ||
|
||
- api : graph_send_recv | ||
args : (Tensor x, Tensor src_index, Tensor dst_index, str pool_type = "SUM", int64_t out_size = 0) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 这里求解释一下,为什么不用设置默认值了 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 发现可以加默认值,不会报错,改好了。 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. out_size 默认 -1? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 个人觉得还是默认0比较好。 |
||
args : (Tensor x, Tensor src_index, Tensor dst_index, str pool_type, IntArray out_size) | ||
output : Tensor(out), Tensor(dst_count) | ||
infer_meta : | ||
func : GraphSendRecvInferMeta | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -302,7 +302,7 @@ void GraphSendRecvInferMeta(const MetaTensor& x, | |
const MetaTensor& src_index, | ||
const MetaTensor& dst_index, | ||
const std::string& pool_type, | ||
int64_t out_size, | ||
const IntArray& out_size, | ||
MetaTensor* out, | ||
MetaTensor* dst_count) { | ||
auto src_index_dims = src_index.dims(); | ||
|
@@ -345,23 +345,11 @@ void GraphSendRecvInferMeta(const MetaTensor& x, | |
"Src_index and Dst_index should have the same shape.")); | ||
|
||
auto dims = x.dims(); | ||
if (out_size <= 0) { | ||
out->set_dims(dims); | ||
} else { | ||
std::vector<int64_t> dims_ = phi::vectorize(dims); | ||
if (dims_.size() > 0) { | ||
dims_[0] = out_size; | ||
} | ||
out->set_dims(phi::make_ddim(dims_)); | ||
} | ||
out->set_dims(dims); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 这里把不定长那个维度设置成-1 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
out->set_dtype(x.dtype()); | ||
|
||
if (pool_type == "MEAN") { | ||
if (out_size <= 0) { | ||
dst_count->set_dims({dims[0]}); | ||
} else { | ||
dst_count->set_dims({out_size}); | ||
} | ||
dst_count->set_dims({dims[0]}); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 同上 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
dst_count->set_dtype(DataType::INT32); | ||
} | ||
} | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -88,27 +88,34 @@ void GraphSendRecvOpKernelLaunchHelper(const Context& ctx, | |
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; | ||
if (out_size <= 0) { | ||
for (int i = 0; i < src_dims.size(); ++i) { | ||
memset_size *= src_dims[i]; | ||
} | ||
} else { | ||
// set out dim following out_size. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. set大写 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
std::vector<int64_t> dims_ = phi::vectorize(src_dims); | ||
if (dims_.size() > 0) { | ||
dims_[0] = out_size; | ||
} | ||
out->Resize(phi::make_ddim(dims_)); | ||
memset_size = out_size; | ||
for (int i = 1; i < src_dims.size(); ++i) { | ||
memset_size *= src_dims[i]; | ||
} | ||
} | ||
|
||
ctx.template Alloc<T>(out); | ||
T* p_output = out->data<T>(); | ||
const size_t& memset_bytes = memset_size * sizeof(T); | ||
memset(p_output, 0, memset_bytes); | ||
|
||
if (index_size == 0) return; | ||
|
||
const IndexT* s_index = src_index.data<IndexT>(); | ||
const IndexT* d_index = dst_index.data<IndexT>(); | ||
|
||
if (pool_type == "SUM") { | ||
GraphSendRecvCpuLoop<T, IndexT, GraphSendRecvSumFunctor<T>>( | ||
src_dims[0], index_size, s_index, d_index, x, out, pool_type); | ||
|
@@ -119,10 +126,12 @@ void GraphSendRecvOpKernelLaunchHelper(const Context& ctx, | |
GraphSendRecvCpuLoop<T, IndexT, GraphSendRecvMaxFunctor<T>>( | ||
src_dims[0], index_size, s_index, d_index, x, out, pool_type); | ||
} else if (pool_type == "MEAN") { | ||
int64_t input_size = out_size <= 0 ? src_dims[0] : out_size; | ||
dst_count->Resize({input_size}); | ||
ctx.template Alloc<int>(dst_count); | ||
int* p_dst_count = dst_count->data<int>(); | ||
memset(p_dst_count, 0, src_dims[0] * sizeof(int)); | ||
GraphSendRecvCpuLoop<T, IndexT, GraphSendRecvSumFunctor<T>>(src_dims[0], | ||
memset(p_dst_count, 0, input_size * sizeof(int)); | ||
GraphSendRecvCpuLoop<T, IndexT, GraphSendRecvSumFunctor<T>>(input_size, | ||
index_size, | ||
s_index, | ||
d_index, | ||
|
@@ -139,16 +148,29 @@ void GraphSendRecvKernel(const Context& ctx, | |
const DenseTensor& src_index, | ||
const DenseTensor& dst_index, | ||
const std::string& pool_type, | ||
int64_t out_size, | ||
const IntArray& out_size, | ||
DenseTensor* out, | ||
DenseTensor* dst_count) { | ||
auto index_type = src_index.dtype(); | ||
auto& out_size_data = out_size.GetData(); | ||
if (index_type == phi::DataType::INT32) { | ||
GraphSendRecvOpKernelLaunchHelper<Context, T, int32_t>( | ||
ctx, x, src_index, dst_index, pool_type, out_size, out, dst_count); | ||
GraphSendRecvOpKernelLaunchHelper<Context, T, int32_t>(ctx, | ||
x, | ||
src_index, | ||
dst_index, | ||
pool_type, | ||
out_size_data[0], | ||
out, | ||
dst_count); | ||
} else if (index_type == phi::DataType::INT64) { | ||
GraphSendRecvOpKernelLaunchHelper<Context, T, int64_t>( | ||
ctx, x, src_index, dst_index, pool_type, out_size, out, dst_count); | ||
GraphSendRecvOpKernelLaunchHelper<Context, T, int64_t>(ctx, | ||
x, | ||
src_index, | ||
dst_index, | ||
pool_type, | ||
out_size_data[0], | ||
out, | ||
dst_count); | ||
} | ||
} | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -37,20 +37,26 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
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; | ||
if (out_size <= 0) { | ||
for (int i = 0; i < src_dims.size(); ++i) { | ||
memset_size *= src_dims[i]; | ||
} | ||
} else { | ||
// set out dim following out_size. | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. set -> Set There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. done |
||
std::vector<int64_t> dims_ = phi::vectorize(src_dims); | ||
if (dims_.size() > 0) { | ||
dims_[0] = out_size; | ||
} | ||
out->Resize(phi::make_ddim(dims_)); | ||
memset_size = out_size; | ||
for (int i = 1; i < src_dims.size(); ++i) { | ||
memset_size *= src_dims[i]; | ||
} | ||
} | ||
ctx.template Alloc<T>(out); | ||
T* p_output = out->data<T>(); | ||
const size_t& memset_bytes = memset_size * sizeof(T); | ||
if (pool_type == "SUM" || pool_type == "MEAN") { | ||
#ifdef PADDLE_WITH_HIP | ||
|
@@ -63,7 +69,7 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
thrust::fill(thrust::device, | ||
p_output_ptr, | ||
p_output_ptr + memset_size, | ||
std::numeric_limits<T>::min()); | ||
std::numeric_limits<T>::lowest()); | ||
} else if (pool_type == "MIN") { | ||
thrust::device_ptr<T> p_output_ptr(p_output); | ||
thrust::fill(thrust::device, | ||
|
@@ -91,7 +97,7 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
int64_t max_grid_dimx = ctx.GetCUDAMaxGridDimSize()[0]; | ||
int64_t grid_tmp = (n + block - 1) / block; | ||
int64_t grid = grid_tmp < max_grid_dimx ? grid_tmp : max_grid_dimx; | ||
int64_t input_size = src_dims[0]; | ||
int64_t input_size = out_size <= 0 ? src_dims[0] : out_size; | ||
if (pool_type == "SUM") { | ||
GraphSendRecvSumCUDAFunctor<T, IndexT> functor; | ||
GraphSendRecvCUDAKernel<T, IndexT, GraphSendRecvSumCUDAFunctor<T, IndexT>> | ||
|
@@ -103,9 +109,6 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
<<<grid, block, 0, ctx.stream()>>>( | ||
p_src, s_index, d_index, p_output, index_size, slice_size, functor); | ||
|
||
if (out_size > 0) { | ||
input_size = out_size; | ||
} | ||
int64_t grid_max_tmp = (input_size * slice_size + block - 1) / block; | ||
int64_t grid_max = | ||
grid_max_tmp < max_grid_dimx ? grid_max_tmp : max_grid_dimx; | ||
|
@@ -117,9 +120,6 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
<<<grid, block, 0, ctx.stream()>>>( | ||
p_src, s_index, d_index, p_output, index_size, slice_size, functor); | ||
|
||
if (out_size > 0) { | ||
input_size = out_size; | ||
} | ||
int64_t grid_min_tmp = (input_size * slice_size + block - 1) / block; | ||
int64_t grid_min = | ||
grid_min_tmp < max_grid_dimx ? grid_min_tmp : max_grid_dimx; | ||
|
@@ -130,12 +130,9 @@ void GraphSendRecvOpCUDAKernelLaunchHelper(const Context& ctx, | |
GraphSendRecvCUDAKernel<T, IndexT, GraphSendRecvSumCUDAFunctor<T, IndexT>> | ||
<<<grid, block, 0, ctx.stream()>>>( | ||
p_src, s_index, d_index, p_output, index_size, slice_size, functor); | ||
|
||
dst_count->Resize({input_size}); | ||
ctx.template Alloc<int32_t>(dst_count); | ||
int32_t* p_dst_count = dst_count->data<int32_t>(); | ||
if (out_size > 0) { | ||
input_size = out_size; | ||
} | ||
int* p_dst_count = dst_count->data<int>(); | ||
|
||
#ifdef PADDLE_WITH_HIP | ||
hipMemset(p_dst_count, 0, input_size * sizeof(int)); | ||
|
@@ -161,16 +158,29 @@ void GraphSendRecvKernel(const Context& ctx, | |
const DenseTensor& src_index, | ||
const DenseTensor& dst_index, | ||
const std::string& pool_type, | ||
int64_t out_size, | ||
const IntArray& out_size, | ||
DenseTensor* out, | ||
DenseTensor* dst_count) { | ||
auto index_type = src_index.dtype(); | ||
auto& out_size_data = out_size.GetData(); | ||
if (index_type == phi::DataType::INT32) { | ||
GraphSendRecvOpCUDAKernelLaunchHelper<Context, T, int32_t>( | ||
ctx, x, src_index, dst_index, pool_type, out_size, out, dst_count); | ||
GraphSendRecvOpCUDAKernelLaunchHelper<Context, T, int32_t>(ctx, | ||
x, | ||
src_index, | ||
dst_index, | ||
pool_type, | ||
out_size_data[0], | ||
out, | ||
dst_count); | ||
} else if (index_type == phi::DataType::INT64) { | ||
GraphSendRecvOpCUDAKernelLaunchHelper<Context, T, int64_t>( | ||
ctx, x, src_index, dst_index, pool_type, out_size, out, dst_count); | ||
GraphSendRecvOpCUDAKernelLaunchHelper<Context, T, int64_t>(ctx, | ||
x, | ||
src_index, | ||
dst_index, | ||
pool_type, | ||
out_size_data[0], | ||
out, | ||
dst_count); | ||
} | ||
} | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
这里的命名为什么要强调是Tensor,看看是否命名是Out_size,和上面的变量命名一致
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
改成Out_size了