-
Notifications
You must be signed in to change notification settings - Fork 5.6k
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
【Hackathon No.60】refactor unary sparse ops and add sparse sqrt, tanh, sin #41356
Changes from 2 commits
97ac270
26f4662
a7f3410
d4310af
7e5f102
71864fd
a99a5ba
95aa0b3
f706dea
d898df7
b770f41
f92e8cd
394ce5e
c577f46
3ad6fba
56fc5da
1f18c59
c606825
5dd4507
f59fa26
dea61c7
60c7359
ad8ceda
178dd27
1ace46f
7bb41d7
790cb0d
c44ac74
d35e923
b04ab6c
fa93d7d
a6d2cd0
67d14b4
268ac34
39c9750
06787c0
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 |
---|---|---|
@@ -0,0 +1,170 @@ | ||
#include "paddle/phi/backends/cpu/cpu_context.h" | ||
#include "paddle/phi/backends/gpu/gpu_context.h" | ||
#include "paddle/phi/core/kernel_registry.h" | ||
#include "paddle/phi/kernels/copy_kernel.h" | ||
#include "paddle/phi/kernels/empty_kernel.h" | ||
|
||
#define DEFINE_SPARSE_UNARY_KERNEL(dense_kernel_func) \ | ||
namespace phi { \ | ||
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. PD_REGISTER_KERNEL 也有必须写在顶层命名空间的要求,这个宏自身包含命名空间会让使用体验更一致一些 |
||
namespace sparse { \ | ||
\ | ||
template <typename T, typename Context> \ | ||
void SparseCoo##dense_kernel_func(const Context& dev_ctx, \ | ||
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. 已修改 |
||
const SparseCooTensor& x, \ | ||
SparseCooTensor* out) { \ | ||
DenseTensor non_zero_indices = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_indices()); \ | ||
DenseTensor non_zero_elements = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements()); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_indices(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_indices); \ | ||
phi::dense_kernel_func<T, Context>( \ | ||
dev_ctx, x.non_zero_elements(), &non_zero_elements); \ | ||
out->SetMember(non_zero_indices, non_zero_elements, x.dims(), true); \ | ||
} \ | ||
\ | ||
template <typename T, typename Context> \ | ||
void SparseCsr##dense_kernel_func(const Context& dev_ctx, \ | ||
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. 同时实现了 coo 和 csr 版本,python api 名也以 coo 和 csr 区分 |
||
const SparseCsrTensor& x, \ | ||
SparseCsrTensor* out) { \ | ||
DenseTensor non_zero_crows = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_crows()); \ | ||
DenseTensor non_zero_cols = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_cols()); \ | ||
DenseTensor non_zero_elements = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements()); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_crows(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_crows); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_cols(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_cols); \ | ||
phi::dense_kernel_func<T, Context>( \ | ||
dev_ctx, x.non_zero_elements(), &non_zero_elements); \ | ||
out->SetMember( \ | ||
non_zero_crows, non_zero_cols, non_zero_elements, x.dims()); \ | ||
} \ | ||
} \ | ||
} | ||
|
||
#define DEFINE_SPARSE_UNARY_GRAD_KERNEL(dense_kernel_func) \ | ||
namespace phi { \ | ||
namespace sparse { \ | ||
\ | ||
template <typename T, typename Context> \ | ||
void SparseCoo##dense_kernel_func(const Context& dev_ctx, \ | ||
const SparseCooTensor& x, \ | ||
const SparseCooTensor& out_grad, \ | ||
SparseCooTensor* x_grad) { \ | ||
DenseTensor non_zero_indices = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_indices()); \ | ||
DenseTensor non_zero_elements = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements()); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_indices(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_indices); \ | ||
phi::dense_kernel_func<T, Context>(dev_ctx, \ | ||
x.non_zero_elements(), \ | ||
out_grad.non_zero_elements(), \ | ||
&non_zero_elements); \ | ||
x_grad->SetMember(non_zero_indices, non_zero_elements, x.dims(), true); \ | ||
} \ | ||
\ | ||
template <typename T, typename Context> \ | ||
void SparseCsr##dense_kernel_func(const Context& dev_ctx, \ | ||
const SparseCsrTensor& x, \ | ||
const SparseCsrTensor& out_grad, \ | ||
SparseCsrTensor* out) { \ | ||
DenseTensor non_zero_crows = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_crows()); \ | ||
DenseTensor non_zero_cols = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_cols()); \ | ||
DenseTensor non_zero_elements = \ | ||
phi::EmptyLike<T, Context>(dev_ctx, x.non_zero_elements()); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_crows(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_crows); \ | ||
phi::Copy(dev_ctx, \ | ||
x.non_zero_cols(), \ | ||
dev_ctx.GetPlace(), \ | ||
false, \ | ||
&non_zero_cols); \ | ||
phi::dense_kernel_func<T, Context>(dev_ctx, \ | ||
x.non_zero_elements(), \ | ||
out_grad.non_zero_elements(), \ | ||
&non_zero_elements); \ | ||
out->SetMember( \ | ||
non_zero_crows, non_zero_cols, non_zero_elements, x.dims()); \ | ||
} \ | ||
} \ | ||
} | ||
|
||
#define REGISTER_CPU_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) \ | ||
PD_REGISTER_KERNEL(sparse_coo_##kernel_name, \ | ||
CPU, \ | ||
ALL_LAYOUT, \ | ||
phi::sparse::SparseCoo##dense_kernel_func, \ | ||
float, \ | ||
double) { \ | ||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ | ||
} \ | ||
PD_REGISTER_KERNEL(sparse_csr_##kernel_name, \ | ||
CPU, \ | ||
ALL_LAYOUT, \ | ||
phi::sparse::SparseCsr##dense_kernel_func, \ | ||
float, \ | ||
double) { \ | ||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ | ||
} | ||
|
||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) | ||
#define REGISTER_GPU_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) \ | ||
PD_REGISTER_KERNEL(sparse_coo_##kernel_name, \ | ||
GPU, \ | ||
ALL_LAYOUT, \ | ||
phi::sparse::SparseCoo##dense_kernel_func, \ | ||
float, \ | ||
double, \ | ||
phi::dtype::float16) { \ | ||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ | ||
} \ | ||
\ | ||
PD_REGISTER_KERNEL(sparse_csr_##kernel_name, \ | ||
GPU, \ | ||
ALL_LAYOUT, \ | ||
phi::sparse::SparseCsr##dense_kernel_func, \ | ||
float, \ | ||
double, \ | ||
phi::dtype::float16) { \ | ||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ | ||
} | ||
#else | ||
// This macro definition is empty when GPU is disabled | ||
#define REGISTER_GPU_SPARSE_UNARY_KERNEL(sparse_kernel_name, dense_kernel_func) | ||
#endif | ||
|
||
#define REGISTER_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) \ | ||
REGISTER_CPU_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) \ | ||
REGISTER_GPU_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) | ||
|
||
#define DEFINE_AND_REGISTER_SPARSE_UNARY_KERNEL(kernel_name, \ | ||
dense_kernel_func) \ | ||
DEFINE_SPARSE_UNARY_KERNEL(dense_kernel_func) \ | ||
REGISTER_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) | ||
|
||
#define DEFINE_AND_REGISTER_SPARSE_UNARY_GRAD_KERNEL(kernel_name, \ | ||
dense_kernel_func) \ | ||
DEFINE_SPARSE_UNARY_GRAD_KERNEL(dense_kernel_func) \ | ||
REGISTER_SPARSE_UNARY_KERNEL(kernel_name, dense_kernel_func) | ||
|
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.
之前 dense tensor 的 SqrtGrad kernel 没有在头文件中声明