diff --git a/paddle/fluid/eager/grad_node_info.cc b/paddle/fluid/eager/grad_node_info.cc index ce7f7caf1f44c..6ec5a17b9361d 100644 --- a/paddle/fluid/eager/grad_node_info.cc +++ b/paddle/fluid/eager/grad_node_info.cc @@ -435,6 +435,32 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in, meta.SetDistAttr(dist_attr); meta.SetDistTensorGlobalDims(dist_tensor->dims()); SetIsRunAutoParallel(true); + } else if (phi::SparseCsrTensor::classof(fwd_in.impl().get())) { + phi::SparseCsrTensor* sparse_tensor = + static_cast(fwd_in.impl().get()); + const phi::DenseTensor dense_tensor = + static_cast(sparse_tensor->values()); + PADDLE_ENFORCE_NE( + dense_tensor.dtype(), + phi::DataType::UNDEFINED, + paddle::platform::errors::Fatal("Attempting to copy DenseTensorMeta " + "with phi::DataType::UNDEFINED," + "which is illegal.")); + meta.SetTensorMeta(dense_tensor.meta()); + meta.SetPlace(fwd_in.place()); + } else if (phi::SparseCooTensor::classof(fwd_in.impl().get())) { + phi::SparseCooTensor* sparse_tensor = + static_cast(fwd_in.impl().get()); + const phi::DenseTensor dense_tensor = + static_cast(sparse_tensor->values()); + PADDLE_ENFORCE_NE( + dense_tensor.dtype(), + phi::DataType::UNDEFINED, + paddle::platform::errors::Fatal("Attempting to copy DenseTensorMeta " + "with phi::DataType::UNDEFINED," + "which is illegal.")); + meta.SetTensorMeta(dense_tensor.meta()); + meta.SetPlace(fwd_in.place()); } else { VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with " diff --git a/paddle/phi/api/yaml/sparse_ops.yaml b/paddle/phi/api/yaml/sparse_ops.yaml index 82c1548bf35da..ac230be485c09 100644 --- a/paddle/phi/api/yaml/sparse_ops.yaml +++ b/paddle/phi/api/yaml/sparse_ops.yaml @@ -2,7 +2,7 @@ args : (Tensor x) output : Tensor(out) infer_meta : - func : UnchangedInferMeta + func : RealAndImagInferMeta kernel : func : abs_coo{sparse_coo -> sparse_coo}, abs_csr{sparse_csr -> sparse_csr} diff --git a/paddle/phi/infermeta/sparse/unary.cc b/paddle/phi/infermeta/sparse/unary.cc index 01da3ae08eb74..f624df0d8c55a 100644 --- a/paddle/phi/infermeta/sparse/unary.cc +++ b/paddle/phi/infermeta/sparse/unary.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/phi/infermeta/sparse/unary.h" - #include "paddle/phi/core/infermeta_utils.h" namespace phi { diff --git a/paddle/phi/kernels/sparse/cpu/mask_kernel.cc b/paddle/phi/kernels/sparse/cpu/mask_kernel.cc index b92ebccbefbc8..5213dd44a4c07 100644 --- a/paddle/phi/kernels/sparse/cpu/mask_kernel.cc +++ b/paddle/phi/kernels/sparse/cpu/mask_kernel.cc @@ -121,9 +121,11 @@ void MaskHelperCooCPUKernel(const CPUContext& dev_ctx, for (uint64_t i = 0; i < x_indexs.size(); i++) { x_indexs_map[x_indexs[i]] = i; } + *out = phi::EmptyLike(dev_ctx, x.values()); + phi::funcs::SetConstant set_zero; + set_zero(dev_ctx, out, static_cast(0)); T* out_ptr = out->data(); - memset(out_ptr, static_cast(0), out->numel() * sizeof(T)); const int64_t stride = x.dims().size() == sparse_dim ? 1 : x.values().dims()[1]; const T* in_ptr = x.values().data(); @@ -166,7 +168,9 @@ PD_REGISTER_KERNEL(mask_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -176,9 +180,12 @@ PD_REGISTER_KERNEL(mask_helper_coo, phi::sparse::MaskHelperCooKernel, float, double, + phi::dtype::float16, uint8_t, int16_t, int, - int64_t) { + int64_t, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } diff --git a/paddle/phi/kernels/sparse/cpu/sparse_utils_kernel.cc b/paddle/phi/kernels/sparse/cpu/sparse_utils_kernel.cc index 6ec3f747dd381..0c5e6857de24c 100644 --- a/paddle/phi/kernels/sparse/cpu/sparse_utils_kernel.cc +++ b/paddle/phi/kernels/sparse/cpu/sparse_utils_kernel.cc @@ -328,7 +328,9 @@ PD_REGISTER_KERNEL(dense_to_coo, int8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(csr_to_coo, CPU, @@ -342,7 +344,9 @@ PD_REGISTER_KERNEL(csr_to_coo, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(coo_to_csr, CPU, @@ -356,7 +360,9 @@ PD_REGISTER_KERNEL(coo_to_csr, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(dense_to_csr, CPU, @@ -369,7 +375,9 @@ PD_REGISTER_KERNEL(dense_to_csr, int8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(coo_to_dense, CPU, @@ -383,7 +391,9 @@ PD_REGISTER_KERNEL(coo_to_dense, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(csr_to_dense, CPU, @@ -397,7 +407,9 @@ PD_REGISTER_KERNEL(csr_to_dense, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(values_coo, CPU, @@ -411,7 +423,9 @@ PD_REGISTER_KERNEL(values_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -442,7 +456,9 @@ PD_REGISTER_KERNEL(values_csr, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); } @@ -456,4 +472,6 @@ PD_REGISTER_KERNEL(sparse_coo_tensor, uint8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/sparse/cpu/unary_grad_kernel.cc b/paddle/phi/kernels/sparse/cpu/unary_grad_kernel.cc index 4c993e3a27425..5d65412a5dc05 100644 --- a/paddle/phi/kernels/sparse/cpu/unary_grad_kernel.cc +++ b/paddle/phi/kernels/sparse/cpu/unary_grad_kernel.cc @@ -37,6 +37,29 @@ kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ } +#define PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(name, prefix) \ + PD_REGISTER_KERNEL(name##_coo_grad, \ + CPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CooGradKernel, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ + } \ + \ + PD_REGISTER_KERNEL(name##_csr_grad, \ + CPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CsrGradKernel, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ + } + PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(sin, Sin) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(tan, Tan) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(asin, Asin) @@ -49,12 +72,13 @@ PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(sqrt, Sqrt) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(square, Square) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(log1p, Log1p) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(relu, Relu) -PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(abs, Abs) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(pow, Pow) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(expm1, Expm1) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(relu6, Relu6) PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL(leaky_relu, LeakyRelu) +PD_REGISTER_SPARSE_UNARY_CPU_GRAD_KERNEL_WITH_COMPLEX(abs, Abs) + PD_REGISTER_KERNEL(cast_coo_grad, CPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/sparse/cpu/unary_kernel.cc b/paddle/phi/kernels/sparse/cpu/unary_kernel.cc index 53956174044fa..6066d79bc26be 100644 --- a/paddle/phi/kernels/sparse/cpu/unary_kernel.cc +++ b/paddle/phi/kernels/sparse/cpu/unary_kernel.cc @@ -78,6 +78,29 @@ void DivScalarCsrKernel(const Context& dev_ctx, kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ } +#define PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(name, prefix) \ + PD_REGISTER_KERNEL(name##_coo, \ + CPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CooKernel, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ + } \ + \ + PD_REGISTER_KERNEL(name##_csr, \ + CPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CsrKernel, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ + } + PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(sin, Sin) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(tan, Tan) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(asin, Asin) @@ -90,13 +113,14 @@ PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(sqrt, Sqrt) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(square, Square) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(log1p, Log1p) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(relu, Relu) -PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(abs, Abs) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(pow, Pow) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(scale, Scale) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(expm1, Expm1) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(relu6, Relu6) PD_REGISTER_SPARSE_UNARY_CPU_KERNEL(leaky_relu, LeakyRelu) +PD_REGISTER_SPARSE_UNARY_CPU_KERNEL_WITH_COMPLEX(abs, Abs) + PD_REGISTER_KERNEL(divide_scalar_coo, CPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/sparse/empty_kernel.cc b/paddle/phi/kernels/sparse/empty_kernel.cc index 49a377ca70f67..2fb11e7a66f2e 100644 --- a/paddle/phi/kernels/sparse/empty_kernel.cc +++ b/paddle/phi/kernels/sparse/empty_kernel.cc @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and limitations under the License. */ #include "paddle/phi/kernels/sparse/empty_kernel.h" - #include "paddle/phi/backends/cpu/cpu_context.h" #include "paddle/phi/backends/gpu/gpu_context.h" #include "paddle/phi/core/kernel_registry.h" @@ -48,7 +47,6 @@ void EmptyLikeCsrKernel(const Context& dev_ctx, out->set_meta(x.meta()); dev_ctx.template Alloc(out_values); } - } // namespace sparse } // namespace phi @@ -63,7 +61,9 @@ PD_REGISTER_KERNEL(empty_like_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -78,7 +78,9 @@ PD_REGISTER_KERNEL(empty_like_csr, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); } @@ -95,7 +97,9 @@ PD_REGISTER_KERNEL(empty_like_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -111,7 +115,10 @@ PD_REGISTER_KERNEL(empty_like_csr, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); } + #endif diff --git a/paddle/phi/kernels/sparse/gpu/mask_kernel.cu b/paddle/phi/kernels/sparse/gpu/mask_kernel.cu index ab367efb11fd6..0941ad69b0dd2 100644 --- a/paddle/phi/kernels/sparse/gpu/mask_kernel.cu +++ b/paddle/phi/kernels/sparse/gpu/mask_kernel.cu @@ -308,7 +308,9 @@ PD_REGISTER_KERNEL(mask_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -322,6 +324,8 @@ PD_REGISTER_KERNEL(mask_helper_coo, uint8_t, int16_t, int, - int64_t) { + int64_t, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } diff --git a/paddle/phi/kernels/sparse/gpu/sparse_utils_kernel.cu b/paddle/phi/kernels/sparse/gpu/sparse_utils_kernel.cu index e9001d2c506b8..6f080e4df90cd 100644 --- a/paddle/phi/kernels/sparse/gpu/sparse_utils_kernel.cu +++ b/paddle/phi/kernels/sparse/gpu/sparse_utils_kernel.cu @@ -589,7 +589,9 @@ PD_REGISTER_KERNEL(dense_to_coo, int8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(csr_to_coo, GPU, @@ -603,7 +605,9 @@ PD_REGISTER_KERNEL(csr_to_coo, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(coo_to_csr, GPU, @@ -617,7 +621,9 @@ PD_REGISTER_KERNEL(coo_to_csr, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(dense_to_csr, GPU, @@ -630,7 +636,9 @@ PD_REGISTER_KERNEL(dense_to_csr, int8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(coo_to_dense, GPU, @@ -644,7 +652,9 @@ PD_REGISTER_KERNEL(coo_to_dense, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(csr_to_dense, GPU, @@ -658,7 +668,9 @@ PD_REGISTER_KERNEL(csr_to_dense, int16_t, int, int64_t, - bool) {} + bool, + phi::dtype::complex, + phi::dtype::complex) {} PD_REGISTER_KERNEL(values_coo, GPU, @@ -672,7 +684,9 @@ PD_REGISTER_KERNEL(values_coo, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -688,7 +702,9 @@ PD_REGISTER_KERNEL(values_csr, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); } @@ -717,4 +733,6 @@ PD_REGISTER_KERNEL(sparse_coo_tensor, uint8_t, int16_t, int, - int64_t) {} + int64_t, + phi::dtype::complex, + phi::dtype::complex) {} diff --git a/paddle/phi/kernels/sparse/gpu/unary_grad_kernel.cu b/paddle/phi/kernels/sparse/gpu/unary_grad_kernel.cu index ef66e91364c02..248556c1d6b53 100644 --- a/paddle/phi/kernels/sparse/gpu/unary_grad_kernel.cu +++ b/paddle/phi/kernels/sparse/gpu/unary_grad_kernel.cu @@ -39,6 +39,31 @@ kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ } +#define PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(name, prefix) \ + PD_REGISTER_KERNEL(name##_coo_grad, \ + GPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CooGradKernel, \ + phi::dtype::float16, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ + } \ + \ + PD_REGISTER_KERNEL(name##_csr_grad, \ + GPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CsrGradKernel, \ + phi::dtype::float16, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ + } + PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(sin, Sin) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(tan, Tan) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(asin, Asin) @@ -51,12 +76,13 @@ PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(sqrt, Sqrt) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(square, Square) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(log1p, Log1p) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(relu, Relu) -PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(abs, Abs) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(pow, Pow) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(expm1, Expm1) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(relu6, Relu6) PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL(leaky_relu, LeakyRelu) +PD_REGISTER_SPARSE_UNARY_GPU_GRAD_KERNEL_WITH_COMPLEX(abs, Abs) + PD_REGISTER_KERNEL(cast_coo_grad, GPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/sparse/gpu/unary_kernel.cu b/paddle/phi/kernels/sparse/gpu/unary_kernel.cu index 22392a9fea5d5..8d24b3011d9ba 100644 --- a/paddle/phi/kernels/sparse/gpu/unary_kernel.cu +++ b/paddle/phi/kernels/sparse/gpu/unary_kernel.cu @@ -69,6 +69,31 @@ void DivScalarCsrKernel(const Context& dev_ctx, kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ } +#define PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(name, prefix) \ + PD_REGISTER_KERNEL(name##_coo, \ + GPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CooKernel, \ + phi::dtype::float16, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); \ + } \ + \ + PD_REGISTER_KERNEL(name##_csr, \ + GPU, \ + ALL_LAYOUT, \ + phi::sparse::prefix##CsrKernel, \ + phi::dtype::float16, \ + float, \ + double, \ + phi::dtype::complex, \ + phi::dtype::complex) { \ + kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR); \ + } + PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(sin, Sin) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(tan, Tan) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(asin, Asin) @@ -81,13 +106,14 @@ PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(sqrt, Sqrt) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(square, Square) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(log1p, Log1p) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(relu, Relu) -PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(abs, Abs) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(pow, Pow) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(scale, Scale) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(expm1, Expm1) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(relu6, Relu6) PD_REGISTER_SPARSE_UNARY_GPU_KERNEL(leaky_relu, LeakyRelu) +PD_REGISTER_SPARSE_UNARY_GPU_KERNEL_WITH_COMPLEX(abs, Abs) + PD_REGISTER_KERNEL(divide_scalar_coo, GPU, ALL_LAYOUT, diff --git a/paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h b/paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h index 6a71a92452673..763293de836e5 100644 --- a/paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h +++ b/paddle/phi/kernels/sparse/impl/unary_grad_kernel_impl.h @@ -91,8 +91,8 @@ DEFINE_SPARSE_UNARY_GRAD_KERNEL(Atanh) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Sqrt) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Square) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Log1p) -DEFINE_SPARSE_UNARY_GRAD_KERNEL(Relu) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Abs) +DEFINE_SPARSE_UNARY_GRAD_KERNEL(Relu) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Expm1) DEFINE_SPARSE_UNARY_GRAD_KERNEL(Relu6) DEFINE_SPARSE_UNARY_GRAD_KERNEL_WITH_ONE_ATTR(Pow, factor) diff --git a/paddle/phi/kernels/sparse/impl/unary_kernel_impl.h b/paddle/phi/kernels/sparse/impl/unary_kernel_impl.h index 84cd885f862f0..dd2461253ad6f 100644 --- a/paddle/phi/kernels/sparse/impl/unary_kernel_impl.h +++ b/paddle/phi/kernels/sparse/impl/unary_kernel_impl.h @@ -13,7 +13,6 @@ // limitations under the License. #pragma once - #include "paddle/phi/core/meta_tensor.h" #include "paddle/phi/core/sparse_coo_tensor.h" #include "paddle/phi/core/sparse_csr_tensor.h" @@ -87,12 +86,44 @@ DEFINE_SPARSE_UNARY_KERNEL(Sqrt) DEFINE_SPARSE_UNARY_KERNEL(Square) DEFINE_SPARSE_UNARY_KERNEL(Log1p) DEFINE_SPARSE_UNARY_KERNEL(Relu) -DEFINE_SPARSE_UNARY_KERNEL(Abs) DEFINE_SPARSE_UNARY_KERNEL(Expm1) DEFINE_SPARSE_UNARY_KERNEL(Relu6) DEFINE_SPARSE_UNARY_KERNEL_WITH_ONE_ATTR(Pow, factor) DEFINE_SPARSE_UNARY_KERNEL_WITH_ONE_ATTR(LeakyRelu, alpha) +template +void AbsCooKernel(const Context& dev_ctx, + const SparseCooTensor& x, + SparseCooTensor* out) { + *(out->mutable_indices()) = x.indices(); + + DenseTensor* out_values = out->mutable_values(); + const DenseTensor& x_values = x.values(); + out_values->Resize(x_values.dims()); + dev_ctx.template Alloc(out_values); + + phi::AbsKernel( + dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements()); + + out->SetIndicesDict(x.GetIndicesDict()); +} + +template +void AbsCsrKernel(const Context& dev_ctx, + const SparseCsrTensor& x, + SparseCsrTensor* out) { + *(out->mutable_crows()) = x.crows(); + *(out->mutable_cols()) = x.cols(); + + DenseTensor* out_values = out->mutable_values(); + const DenseTensor& x_values = x.values(); + out_values->Resize(x_values.dims()); + dev_ctx.template Alloc(out_values); + + phi::AbsKernel( + dev_ctx, x.non_zero_elements(), out->mutable_non_zero_elements()); +} + template void ScaleCooKernel(const Context& dev_ctx, const SparseCooTensor& x, diff --git a/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc b/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc index 064867610d719..f5915c7acb84c 100644 --- a/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc +++ b/paddle/phi/kernels/sparse/sparse_utils_grad_kernel.cc @@ -49,7 +49,9 @@ PD_REGISTER_KERNEL(values_coo_grad, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -64,7 +66,9 @@ PD_REGISTER_KERNEL(coo_to_dense_grad, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -74,10 +78,13 @@ PD_REGISTER_KERNEL(sparse_coo_tensor_grad, phi::sparse::SparseCooTensorGradKernel, float, double, + paddle::float16, uint8_t, int16_t, int, - int64_t) { + int64_t, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); } @@ -94,7 +101,9 @@ PD_REGISTER_KERNEL(values_coo_grad, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } PD_REGISTER_KERNEL(coo_to_dense_grad, @@ -109,7 +118,9 @@ PD_REGISTER_KERNEL(coo_to_dense_grad, int16_t, int, int64_t, - bool) { + bool, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); } PD_REGISTER_KERNEL(sparse_coo_tensor_grad, @@ -121,7 +132,9 @@ PD_REGISTER_KERNEL(sparse_coo_tensor_grad, uint8_t, int16_t, int, - int64_t) { + int64_t, + phi::dtype::complex, + phi::dtype::complex) { kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO); } #endif diff --git a/python/paddle/sparse/unary.py b/python/paddle/sparse/unary.py index c4f54631deee5..3f378e242e346 100644 --- a/python/paddle/sparse/unary.py +++ b/python/paddle/sparse/unary.py @@ -648,7 +648,7 @@ def abs(x, name=None): out = |x| Parameters: - x (Tensor): The input Sparse Tensor with data type float32, float64. + x (Tensor): The input Sparse Tensor with data type float32, float64, complex64, complex128. name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`. diff --git a/test/legacy_test/test_sparse_unary_op.py b/test/legacy_test/test_sparse_unary_op.py index 1f04694747e71..70a29ce0c62e3 100644 --- a/test/legacy_test/test_sparse_unary_op.py +++ b/test/legacy_test/test_sparse_unary_op.py @@ -19,6 +19,8 @@ import paddle from paddle.base.framework import convert_np_dtype_to_dtype_ +devices = ['cpu', 'gpu'] + class TestSparseUnary(unittest.TestCase): def to_sparse(self, x, format): @@ -27,16 +29,39 @@ def to_sparse(self, x, format): elif format == 'csr': return x.detach().to_sparse_csr() - def check_result(self, dense_func, sparse_func, format, *args): - origin_x = paddle.rand([8, 16, 32], dtype='float32') - mask = paddle.randint(0, 2, [8, 16, 32]).astype('float32') - while paddle.sum(mask) == 0: + def check_result( + self, + dense_func, + sparse_func, + format, + device='cpu', + dtype='float32', + *args + ): + if dtype == 'complex64': + origin_x_real = paddle.rand([8, 16, 32], 'float32') + origin_x_com = paddle.rand([8, 16, 32], 'float32') + origin_x = (origin_x_real + 1j * origin_x_com).astype('complex64') mask = paddle.randint(0, 2, [8, 16, 32]).astype("float32") + while paddle.sum(mask) == 0: + mask = paddle.randint(0, 2, [8, 16, 32]).astype("float32") + elif dtype == 'complex128': + origin_x_real = paddle.rand([8, 16, 32], 'float64') + origin_x_com = paddle.rand([8, 16, 32], 'float64') + origin_x = (origin_x_real + 1j * origin_x_com).astype('complex128') + mask = paddle.randint(0, 2, [8, 16, 32]).astype("float64") + while paddle.sum(mask) == 0: + mask = paddle.randint(0, 2, [8, 16, 32]).astype("float64") + else: + origin_x = paddle.rand([8, 16, 32], dtype) + mask = paddle.randint(0, 2, [8, 16, 32]).astype(dtype) + while paddle.sum(mask) == 0: + mask = paddle.randint(0, 2, [8, 16, 32]).astype(dtype) # --- check sparse coo with dense --- # dense_x = origin_x * mask + dense_x.to(device) sp_x = self.to_sparse(dense_x, format) - sp_x.stop_gradient = False if len(args) == 0: sp_out = sparse_func(sp_x) @@ -79,19 +104,59 @@ def check_result(self, dense_func, sparse_func, format, *args): sp_x.grad.to_dense().numpy(), expect_grad, rtol=1e-05 ) - def compare_with_dense(self, dense_func, sparse_func): - self.check_result(dense_func, sparse_func, 'coo') - self.check_result(dense_func, sparse_func, 'csr') + def compare_with_dense(self, dense_func, sparse_func, dtype='float32'): + for device in devices: + # The sparse unary op is only compatible with float16 on the CUDA. + if (device == 'cpu' and dtype != 'float16') or ( + device == 'gpu' and paddle.is_compiled_with_cuda() + ): + self.check_result(dense_func, sparse_func, 'coo', device, dtype) + self.check_result(dense_func, sparse_func, 'csr', device, dtype) def compare_with_dense_one_attr(self, dense_func, sparse_func, attr1): - self.check_result(dense_func, sparse_func, 'coo', attr1) - self.check_result(dense_func, sparse_func, 'csr', attr1) + for device in devices: + if device == 'cpu' or ( + device == 'gpu' and paddle.is_compiled_with_cuda() + ): + self.check_result( + dense_func, sparse_func, 'coo', device, 'float32', attr1 + ) + self.check_result( + dense_func, sparse_func, 'csr', device, 'float32', attr1 + ) def compare_with_dense_two_attr( self, dense_func, sparse_func, attr1, attr2 ): - self.check_result(dense_func, sparse_func, 'coo', attr1, attr2) - self.check_result(dense_func, sparse_func, 'csr', attr1, attr2) + for device in devices: + if device == 'cpu' or ( + device == 'gpu' and paddle.is_compiled_with_cuda() + ): + self.check_result( + dense_func, + sparse_func, + 'coo', + device, + 'float32', + attr1, + attr2, + ) + self.check_result( + dense_func, + sparse_func, + 'csr', + device, + 'float32', + attr1, + attr2, + ) + + def test_sparse_abs(self): + self.compare_with_dense(paddle.abs, paddle.sparse.abs, 'float16') + self.compare_with_dense(paddle.abs, paddle.sparse.abs, 'float32') + self.compare_with_dense(paddle.abs, paddle.sparse.abs, 'float64') + self.compare_with_dense(paddle.abs, paddle.sparse.abs, 'complex64') + self.compare_with_dense(paddle.abs, paddle.sparse.abs, 'complex128') def test_sparse_sin(self): self.compare_with_dense(paddle.sin, paddle.sparse.sin) @@ -105,9 +170,6 @@ def test_sparse_asin(self): def test_sparse_atan(self): self.compare_with_dense(paddle.atan, paddle.sparse.atan) - def test_sparse_sinh(self): - self.compare_with_dense(paddle.sinh, paddle.sparse.sinh) - def test_sparse_tanh(self): self.compare_with_dense(paddle.tanh, paddle.sparse.tanh) @@ -137,8 +199,8 @@ def test_sparse_leaky_relu(self): paddle.nn.LeakyReLU(0.1), paddle.sparse.nn.LeakyReLU(0.1) ) - def test_sparse_abs(self): - self.compare_with_dense(paddle.abs, paddle.sparse.abs) + def test_sparse_sinh(self): + self.compare_with_dense(paddle.sinh, paddle.sparse.sinh) def test_sparse_expm1(self): self.compare_with_dense(paddle.expm1, paddle.sparse.expm1)