【Paddle Tensor No.26】Svdvals new branch

paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc

@@ -44,6 +44,14 @@ void OperatorDialect::initialize() {
       >();
   RegisterOps<
 #define GET_OP_LIST2
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
+      >();
+  RegisterOps<
+#define GET_OP_LIST3
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
+      >();
+  RegisterOps<
+#define GET_OP_LIST4
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
       >();
 #else

paddle/fluid/pir/dialect/op_generator/op_gen.py

@@ -245,6 +245,12 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 #elif defined(GET_OP_LIST2)
 #undef GET_OP_LIST2
 {op_declare_second_part}
+#elif defined(GET_OP_LIST3)
+#undef GET_OP_LIST3
+{op_declare_third_part}
+#elif defined(GET_OP_LIST4)
+#undef GET_OP_LIST4
+{op_declare_fourth_part}
 """
 
 CC_OP_INFO_FILE_TEMPLATE_PART2 = """
@@ -2390,9 +2396,11 @@ def OpGenerator(
 
     if op_info_file is not None:
         if sys.platform == "win32":
-            n = len(op_list_strs) // 2
+            n = len(op_list_strs) // 4
             first_part_op_info = op_list_strs[:n]
-            second_part_op_info = op_list_strs[n:]
+            second_part_op_info = op_list_strs[n : 2 * n]
+            third_part_op_info = op_list_strs[2 * n : 3 * n]
+            fourth_part_op_info = op_list_strs[3 * n :]
             CC_OP_INFO_FILE_TEMPLATE = (
                 CC_OP_INFO_FILE_TEMPLATE_WIN_PART1
                 + CC_OP_INFO_FILE_TEMPLATE_PART2
@@ -2404,6 +2412,12 @@ def OpGenerator(
                 op_declare_second_part=",".join(second_part_op_info).replace(
                     "\n", ""
                 ),
+                op_declare_third_part=",".join(third_part_op_info).replace(
+                    "\n", ""
+                ),
+                op_declare_fourth_part=",".join(fourth_part_op_info).replace(
+                    "\n", ""
+                ),
                 other_info=other_info_str,
                 h_file=op_def_h_file[:-4],
             )

paddle/fluid/pir/dialect/operator/ir/op_dialect.cc

@@ -369,6 +369,17 @@ void OperatorDialect::initialize() {
 #define GET_OP_LIST2
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
       >();
+
+  RegisterOps<
+#define GET_OP_LIST3
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
+      >();
+
+  RegisterOps<
+#define GET_OP_LIST4
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
+      >();
+
 #else
   RegisterOps<
 #define GET_OP_LIST

paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.cc

@@ -57,6 +57,15 @@ void OneDNNOperatorDialect::initialize() {
 #define GET_OP_LIST2
 #include "paddle/fluid/pir/dialect/operator/ir/onednn_op_info.cc"  // NOLINT
       >();
+  RegisterOps<
+#define GET_OP_LIST3
+#include "paddle/fluid/pir/dialect/operator/ir/onednn_op_info.cc"  // NOLINT
+      >();
+  RegisterOps<
+#define GET_OP_LIST4
+#include "paddle/fluid/pir/dialect/operator/ir/onednn_op_info.cc"  // NOLINT
+      >();
+
 #else
   RegisterOps<
 #define GET_OP_LIST

paddle/phi/infermeta/unary.cc

@@ -4902,6 +4902,31 @@ void PartialConcatInferMeta(const std::vector<const MetaTensor*>& xs,
   out->set_dtype(xs[0]->dtype());
 }
 
+void SvdvalsInferMeta(const MetaTensor& x, MetaTensor* s) {
+  auto SDDim = [](const DDim& x_dim, int k) {
+    auto x_vec = common::vectorize(x_dim);
+    x_vec.erase(x_vec.end() - 2, x_vec.end());
+    x_vec.push_back(k);
+    return common::make_ddim(x_vec);
+  };
+
+  auto in_dims = x.dims();
+  int64_t x_rank = in_dims.size();
+
+  PADDLE_ENFORCE_GE(
+      x_rank,
+      2,
+      common::errors::InvalidArgument("The rank of input tensor must be >= 2"));
+
+  int64_t m = in_dims[x_rank - 2];
+  int64_t n = in_dims[x_rank - 1];
+
+  int64_t k = std::min(m, n);
+  s->set_dims(SDDim(in_dims, k));
+  s->share_lod(x);
+  s->set_dtype(x.dtype());
+}
+
 void SvdInferMeta(const MetaTensor& x,
                   bool full_matrices,
                   MetaTensor* u,

paddle/phi/infermeta/unary.h

@@ -810,6 +810,8 @@ void PartialSumInferMeta(const std::vector<const MetaTensor*>& xs,
                          MetaTensor* out,
                          MetaConfig config = MetaConfig());
 
+void SvdvalsInferMeta(const MetaTensor& x, MetaTensor* s);
+
 void SvdInferMeta(const MetaTensor& x,
                   bool full_matrices,
                   MetaTensor* u,

paddle/phi/kernels/cpu/svdvals_grad_kernel.cc

@@ -0,0 +1,20 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/svdvals_grad_kernel.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/svdvals_grad_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    svdvals_grad, CPU, ALL_LAYOUT, phi::SvdvalsGradKernel, float, double) {}

paddle/phi/kernels/cpu/svdvals_kernel.cc

@@ -0,0 +1,130 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/svdvals_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/complex_functors.h"
+#include "paddle/phi/kernels/funcs/lapack/lapack_function.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi {
+
+template <typename T>
+void LapackSvdvals(const T* X, T* S, int rows, int cols) {
+  // Using N to neglect computing U、VH
+  char jobz = 'N';
+  T* a = const_cast<T*>(X);
+  int lda = rows;
+  int lwork = -1;
+  std::vector<T> work(1);
+  int info = 0;
+  // Get the best lwork
+  phi::funcs::lapackSvd<T>(jobz,
+                           rows,
+                           cols,
+                           a,
+                           lda,
+                           S,
+                           nullptr,  // U is not needed
+                           1,        // dummy dimension for U
+                           nullptr,  // VH is not needed
+                           1,        // dummy dimension for VH
+                           work.data(),
+                           lwork,
+                           nullptr,  // iwork is not needed
+                           &info);
+  if (info != 0) {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "Error during LAPACK lwork query. Invalid matrix or arguments."));
+  }
+  lwork = static_cast<int>(work[0]);
+  work.resize(lwork);
+  phi::funcs::lapackSvd<T>(jobz,
+                           rows,
+                           cols,
+                           a,
+                           lda,
+                           S,
+                           nullptr,  // U is not needed
+                           1,        // dummy dimension for U
+                           nullptr,  // VH is not needed
+                           1,        // dummy dimension for VH
+                           work.data(),
+                           lwork,
+                           nullptr,  // iwork is not needed
+                           &info);
+  if (info < 0) {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "This %s-th argument has an illegal value.", info));
+  }
+  if (info > 0) {
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "SVD computation did not converge. Input matrix may be invalid."));
+  }
+}
+
+template <typename T>
+void BatchSvdvals(const T* X, T* S, int rows, int cols, int batches) {
+  int stride = rows * cols;
+  int stride_s = std::min(rows, cols);
+  for (int i = 0; i < batches; i++) {
+    LapackSvdvals<T>(X + i * stride, S + i * stride_s, rows, cols);
+  }
+}
+
+template <typename T, typename Context>
+void SvdvalsKernel(const Context& dev_ctx,
+                   const DenseTensor& X,
+                   DenseTensor* S) {
+  auto x_dims = X.dims();
+  int rows = static_cast<int>(x_dims[x_dims.size() - 2]);
+  int cols = static_cast<int>(x_dims[x_dims.size() - 1]);
+  // Validate dimensions
+  PADDLE_ENFORCE_GT(
+      rows,
+      0,
+      phi::errors::InvalidArgument("The row of Input(X) must be > 0."));
+  PADDLE_ENFORCE_GT(
+      cols,
+      0,
+      phi::errors::InvalidArgument("The column of Input(X) must be > 0."));
+  int k = std::min(rows, cols);
+  int batches = static_cast<int>(X.numel() / (rows * cols));
+  PADDLE_ENFORCE_GT(
+      batches,
+      0,
+      phi::errors::InvalidArgument("The batch size of Input(X) must be > 0."));
+  DDim s_dims;
+  if (batches == 1) {
+    s_dims = {k};
+  } else {
+    s_dims = {batches, k};
+  }
+  S->Resize(s_dims);
+  // Allocate memory for output
+  auto* S_out = dev_ctx.template Alloc<phi::dtype::Real<T>>(S);
+
+  // Transpose the last two dimensions for LAPACK compatibility
+  DenseTensor trans_x = ::phi::TransposeLast2Dim<T>(dev_ctx, X);
+  auto* x_data = trans_x.data<T>();
+  // Perform batch SVD computation for singular values
+  BatchSvdvals<T>(x_data, S_out, rows, cols, batches);
+}
+
+}  // namespace phi
+
+// Register the kernel for CPU
+PD_REGISTER_KERNEL(
+    svdvals, CPU, ALL_LAYOUT, phi::SvdvalsKernel, float, double) {}

paddle/phi/kernels/impl/svdvals_grad_kernel_impl.h

@@ -0,0 +1,62 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/infermeta/unary.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+#include "paddle/phi/kernels/diag_kernel.h"
+#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
+#include "paddle/phi/kernels/funcs/lapack/lapack_function.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/impl/diag_embed_impl.h"
+#include "paddle/phi/kernels/matmul_kernel.h"
+#include "paddle/phi/kernels/slice_kernel.h"
+#include "paddle/phi/kernels/svd_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SvdvalsGradKernel(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& s_grad,
+                       DenseTensor* x_grad) {
+  auto x_dims = x.dims();
+  int rows = static_cast<int>(x_dims[x_dims.size() - 2]);
+  int cols = static_cast<int>(x_dims[x_dims.size() - 1]);
+  int batches = static_cast<int>(x.numel() / (rows * cols));
+  DenseTensor dX_term;
+  if (batches == 1) {
+    dX_term = Diag<T, Context>(dev_ctx, s_grad, 0, 0);
+  } else {
+    MetaTensor meta_dX(&dX_term);
+    DiagEmbedInferMeta(s_grad, 0, -1, -2, &meta_dX);
+    phi::DiagEmbedKernel<T, Context>(dev_ctx, s_grad, 0, -1, -2, &dX_term);
+  }
+
+  DenseTensor U, VH, S_recomputed;
+  MetaTensor meta_u(&U), meta_s(&S_recomputed), meta_vh(&VH);
+  SvdInferMeta(x, false, &meta_u, &meta_s, &meta_vh);
+  phi::SvdKernel<T, Context>(dev_ctx,
+                             x,
+                             false,
+                             &U,
+                             &S_recomputed,
+                             &VH);  // Crucial: recomputing SVD
+  *x_grad =
+      Matmul<T, Context>(dev_ctx, Matmul<T, Context>(dev_ctx, U, dX_term), VH);
+}
+}  // namespace phi

paddle/phi/kernels/svdvals_grad_kernel.h

@@ -0,0 +1,26 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SvdvalsGradKernel(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& s_grad,
+                       DenseTensor* x_grad);
+}  // namespace phi

paddle/phi/kernels/svdvals_kernel.h

@@ -0,0 +1,27 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SvdvalsKernel(const Context& dev_ctx,
+                   const DenseTensor& X,
+                   DenseTensor* S);
+
+}  // namespace phi