【complex op No.33】abs_coo/abs_csr(sparse)

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<phi::SparseCsrTensor*>(fwd_in.impl().get());
+      const phi::DenseTensor dense_tensor =
+          static_cast<const phi::DenseTensor>(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<phi::SparseCooTensor*>(fwd_in.impl().get());
+      const phi::DenseTensor dense_tensor =
+          static_cast<const phi::DenseTensor>(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 "

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}

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 {

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<T>(dev_ctx, x.values());
+  phi::funcs::SetConstant<CPUContext, T> set_zero;
+  set_zero(dev_ctx, out, static_cast<T>(0));
   T* out_ptr = out->data<T>();
-  memset(out_ptr, static_cast<T>(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<T>();
@@ -166,7 +168,9 @@ PD_REGISTER_KERNEL(mask_coo,
                    int16_t,
                    int,
                    int64_t,
-                   bool) {
+                   bool,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }

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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {}
 
 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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {}

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<float>,                            \
+                     phi::dtype::complex<double>) {                         \
+    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<float>,                            \
+                     phi::dtype::complex<double>) {                         \
+    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,

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<float>,                       \
+                     phi::dtype::complex<double>) {                    \
+    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<float>,                       \
+                     phi::dtype::complex<double>) {                    \
+    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,

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<T>(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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_CSR);
 }
+
 #endif

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<float>,
+                   phi::dtype::complex<double>) {
   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<float>,
+                   phi::dtype::complex<double>) {
   kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
 }