paddle.linalg.matrix_rank add attributes atol and rtol 易用性提升 (#66929)

NKNaN · web-flow · commit 3372fe75ac35 · 2024-09-03T17:32:19.000+08:00
* update matrix rank op

* udpate

* update kernel

* update

* update kernel
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -21,6 +21,7 @@ limitations under the License. */
 #include "paddle/common/layout.h"
 #include "paddle/phi/core/ddim.h"
 #include "paddle/phi/core/enforce.h"
+#include "paddle/phi/infermeta/binary.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
 #include "paddle/phi/kernels/impl/box_coder.h"
 
@@ -1305,6 +1306,14 @@ void MatchMatrixTensorInferMeta(const MetaTensor& x,
   }
 }
 
+void MatrixRankAtolRtolInferMeta(const MetaTensor& x,
+                                 const MetaTensor& atol,
+                                 const MetaTensor& rtol,
+                                 bool hermitian,
+                                 MetaTensor* out) {
+  MatrixRankTolInferMeta(x, atol, true, hermitian, out);
+}
+
 void MultiClassNMSInferMeta(const MetaTensor& bboxes,
                             const MetaTensor& scores,
                             const MetaTensor& rois_num,
diff --git a/paddle/phi/infermeta/ternary.h b/paddle/phi/infermeta/ternary.h
@@ -231,6 +231,12 @@ void MatchMatrixTensorInferMeta(const MetaTensor& x,
                                 MetaTensor* tmp,
                                 MetaConfig config = MetaConfig());
 
+void MatrixRankAtolRtolInferMeta(const MetaTensor& x,
+                                 const MetaTensor& atol,
+                                 const MetaTensor& rtol,
+                                 bool hermitian,
+                                 MetaTensor* out);
+
 void MovingAverageAbsMaxScaleInferMeta(const MetaTensor& x,
                                        const MetaTensor& in_accum,
                                        const MetaTensor& in_state,
diff --git a/paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc b/paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/abs_kernel.h"
+#include "paddle/phi/kernels/compare_kernel.h"
 #include "paddle/phi/kernels/elementwise_multiply_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/compare_functors.h"
@@ -26,6 +27,7 @@
 #include "paddle/phi/kernels/reduce_max_kernel.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
 #include "paddle/phi/kernels/transpose_kernel.h"
+#include "paddle/phi/kernels/where_kernel.h"
 
 namespace phi {
 
@@ -169,9 +171,135 @@ void MatrixRankTolKernel(const Context& dev_ctx,
                           false,
                           out);
 }
+
+template <typename T, typename Context>
+void MatrixRankAtolRtolKernel(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const DenseTensor& atol,
+                              const paddle::optional<DenseTensor>& rtol,
+                              bool hermitian,
+                              DenseTensor* out) {
+  dev_ctx.template Alloc<int64_t>(out);
+  auto dim_x = x.dims();
+  auto dim_out = out->dims();
+  int rows = static_cast<int>(dim_x[dim_x.size() - 2]);
+  int cols = static_cast<int>(dim_x[dim_x.size() - 1]);
+  int k = std::min(rows, cols);
+  int batches = static_cast<int>(x.numel() / (rows * cols));
+
+  DenseTensor eigenvalue_tensor;
+  eigenvalue_tensor.Resize(detail::GetEigenvalueDim(dim_x, k));
+  auto* eigenvalue_data = dev_ctx.template Alloc<T>(&eigenvalue_tensor);
+
+  if (hermitian) {
+    phi::funcs::MatrixEighFunctor<Context, T> functor;
+    functor(dev_ctx, x, &eigenvalue_tensor, nullptr, true, false);
+    phi::AbsKernel<T, Context>(dev_ctx, eigenvalue_tensor, &eigenvalue_tensor);
+  } else {
+    DenseTensor trans_x = phi::TransposeLast2Dim<T>(dev_ctx, x);
+    auto* x_data = trans_x.data<T>();
+    BatchSVD<T>(x_data, eigenvalue_data, batches, rows, cols);
+  }
+
+  DenseTensor max_eigenvalue_tensor;
+  max_eigenvalue_tensor.Resize(detail::RemoveLastDim(eigenvalue_tensor.dims()));
+  dev_ctx.template Alloc<T>(&max_eigenvalue_tensor);
+  phi::MaxKernel<T, Context>(dev_ctx,
+                             eigenvalue_tensor,
+                             phi::IntArray({-1}),
+                             false,
+                             &max_eigenvalue_tensor);
+
+  DenseTensor tol_tensor;
+  tol_tensor.Resize(dim_out);
+  dev_ctx.template Alloc<T>(&tol_tensor);
+
+  if (rtol) {
+    DenseTensor tmp_rtol_tensor;
+    tmp_rtol_tensor = phi::Multiply<T>(dev_ctx, *rtol, max_eigenvalue_tensor);
+    funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
+        dev_ctx,
+        atol,
+        tmp_rtol_tensor,
+        GreaterElementFunctor<T>(),
+        &tol_tensor);
+  } else {
+    // when `rtol` is specified to be None in py api
+    // use rtol=eps*max(m, n) only if `atol` is passed with value 0.0, else use
+    // rtol=0.0
+    DenseTensor zero_tensor;
+    zero_tensor = phi::FullLike<T, Context>(dev_ctx, atol, static_cast<T>(0.0));
+
+    T rtol_T = std::numeric_limits<T>::epsilon() * std::max(rows, cols);
+    DenseTensor default_rtol_tensor;
+    default_rtol_tensor =
+        phi::Full<T, Context>(dev_ctx, {1}, static_cast<T>(rtol_T));
+    default_rtol_tensor =
+        phi::Multiply<T>(dev_ctx, default_rtol_tensor, max_eigenvalue_tensor);
+
+    DenseTensor atol_compare_result;
+    atol_compare_result.Resize(atol.dims());
+    phi::EqualKernel<T, Context>(
+        dev_ctx, atol, zero_tensor, &atol_compare_result);
+
+    DenseTensor selected_rtol_tensor;
+    selected_rtol_tensor.Resize(atol.dims());
+    phi::WhereKernel<T, Context>(dev_ctx,
+                                 atol_compare_result,
+                                 default_rtol_tensor,
+                                 zero_tensor,
+                                 &selected_rtol_tensor);
+    funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
+        dev_ctx,
+        atol,
+        selected_rtol_tensor,
+        GreaterElementFunctor<T>(),
+        &tol_tensor);
+  }
+
+  tol_tensor.Resize(detail::NewAxisDim(tol_tensor.dims(), 1));
+
+  DenseTensor compare_result;
+  compare_result.Resize(detail::NewAxisDim(dim_out, k));
+  dev_ctx.template Alloc<int64_t>(&compare_result);
+  int axis = -1;
+  if (eigenvalue_tensor.dims().size() >= tol_tensor.dims().size()) {
+    funcs::ElementwiseCompute<funcs::GreaterThanFunctor<T, int64_t>, T, int>(
+        dev_ctx,
+        eigenvalue_tensor,
+        tol_tensor,
+        funcs::GreaterThanFunctor<T, int64_t>(),
+        &compare_result,
+        axis);
+  } else {
+    funcs::ElementwiseCompute<funcs::LessThanFunctor<T, int64_t>, T, int>(
+        dev_ctx,
+        eigenvalue_tensor,
+        tol_tensor,
+        funcs::LessThanFunctor<T, int64_t>(),
+        &compare_result,
+        axis);
+  }
+
+  phi::SumKernel<int64_t>(dev_ctx,
+                          compare_result,
+                          std::vector<int64_t>{-1},
+                          compare_result.dtype(),
+                          false,
+                          out);
+}
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
     matrix_rank_tol, CPU, ALL_LAYOUT, phi::MatrixRankTolKernel, float, double) {
   kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
 }
+
+PD_REGISTER_KERNEL(matrix_rank_atol_rtol,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::MatrixRankAtolRtolKernel,
+                   float,
+                   double) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
diff --git a/paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu b/paddle/phi/kernels/gpu/matrix_rank_tol_kernel.cu
@@ -24,13 +24,15 @@
 #include "paddle/phi/common/memory_utils.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/abs_kernel.h"
+#include "paddle/phi/kernels/compare_kernel.h"
 #include "paddle/phi/kernels/elementwise_multiply_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/broadcast_function.h"
 #include "paddle/phi/kernels/funcs/compare_functors.h"
 #include "paddle/phi/kernels/impl/matrix_rank_kernel_impl.h"
 #include "paddle/phi/kernels/reduce_max_kernel.h"
 #include "paddle/phi/kernels/reduce_sum_kernel.h"
+#include "paddle/phi/kernels/where_kernel.h"
 
 namespace phi {
 
@@ -447,6 +449,151 @@ void MatrixRankTolKernel(const Context& dev_ctx,
                           out);
 }
 
+template <typename T, typename Context>
+void MatrixRankAtolRtolKernel(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const DenseTensor& atol,
+                              const paddle::optional<DenseTensor>& rtol,
+                              bool hermitian,
+                              DenseTensor* out) {
+  auto* x_data = x.data<T>();
+  dev_ctx.template Alloc<int64_t>(out);
+
+  auto dim_x = x.dims();
+  auto dim_out = out->dims();
+  int rows = dim_x[dim_x.size() - 2];
+  int cols = dim_x[dim_x.size() - 1];
+  PADDLE_ENFORCE_NE(
+      rows,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-2] should not "
+                                   "be 0, but shape is %s now.",
+                                   dim_x));
+  PADDLE_ENFORCE_NE(
+      cols,
+      0,
+      phi::errors::InvalidArgument("The input Tensor x's shape[-1] should not "
+                                   "be 0, but shape is %s now.",
+                                   dim_x));
+  int k = std::min(rows, cols);
+  auto numel = x.numel();
+  int batches = numel / (rows * cols);
+
+  // Must Copy X once, because the gesvdj will destroy the content when exit.
+  DenseTensor x_tmp;
+  phi::Copy(dev_ctx, x, dev_ctx.GetPlace(), false, &x_tmp);
+  auto info = phi::memory_utils::Alloc(
+      dev_ctx.GetPlace(),
+      sizeof(int) * batches,
+      phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
+  int* info_ptr = reinterpret_cast<int*>(info->ptr());
+
+  DenseTensor eigenvalue_tensor;
+  eigenvalue_tensor.Resize(detail::GetEigenvalueDim(dim_x, k));
+  auto* eigenvalue_data = dev_ctx.template Alloc<T>(&eigenvalue_tensor);
+
+  if (hermitian) {
+    SyevjBatched<T>(
+        dev_ctx, batches, rows, x_tmp.data<T>(), eigenvalue_data, info_ptr);
+
+    phi::AbsKernel<T, Context>(dev_ctx, eigenvalue_tensor, &eigenvalue_tensor);
+
+  } else {
+    DenseTensor U, VH;
+    U.Resize(detail::GetUDDim(dim_x, k));
+    VH.Resize(detail::GetVHDDim(dim_x, k));
+    auto* u_data = dev_ctx.template Alloc<T>(&U);
+    auto* vh_data = dev_ctx.template Alloc<T>(&VH);
+    GesvdjBatched<T>(dev_ctx,
+                     batches,
+                     cols,
+                     rows,
+                     k,
+                     x_tmp.data<T>(),
+                     vh_data,
+                     u_data,
+                     eigenvalue_data,
+                     info_ptr,
+                     1);
+  }
+
+  DenseTensor max_eigenvalue_tensor;
+  dev_ctx.template Alloc<T>(&max_eigenvalue_tensor);
+  max_eigenvalue_tensor.Resize(detail::RemoveLastDim(eigenvalue_tensor.dims()));
+
+  phi::MaxKernel<T, Context>(dev_ctx,
+                             eigenvalue_tensor,
+                             phi::IntArray({-1}),
+                             false,
+                             &max_eigenvalue_tensor);
+
+  DenseTensor tol_tensor;
+  tol_tensor.Resize(dim_out);
+  dev_ctx.template Alloc<T>(&tol_tensor);
+
+  if (rtol) {
+    DenseTensor tmp_rtol_tensor;
+    tmp_rtol_tensor = phi::Multiply<T>(dev_ctx, *rtol, max_eigenvalue_tensor);
+    funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
+        dev_ctx,
+        atol,
+        tmp_rtol_tensor,
+        GreaterElementFunctor<T>(),
+        &tol_tensor);
+  } else {
+    // when `rtol` is specified to be None in py api
+    // use rtol=eps*max(m, n) only if `atol` is passed with value 0.0, else use
+    // rtol=0.0
+    DenseTensor zero_tensor;
+    zero_tensor = phi::FullLike<T, Context>(dev_ctx, atol, static_cast<T>(0.0));
+
+    T rtol_T = std::numeric_limits<T>::epsilon() * std::max(rows, cols);
+    DenseTensor default_rtol_tensor;
+    default_rtol_tensor =
+        phi::Full<T, Context>(dev_ctx, {1}, static_cast<T>(rtol_T));
+    default_rtol_tensor =
+        phi::Multiply<T>(dev_ctx, default_rtol_tensor, max_eigenvalue_tensor);
+
+    DenseTensor atol_compare_result;
+    atol_compare_result.Resize(atol.dims());
+    phi::EqualKernel<T, Context>(
+        dev_ctx, atol, zero_tensor, &atol_compare_result);
+
+    DenseTensor selected_rtol_tensor;
+    selected_rtol_tensor.Resize(atol.dims());
+    phi::WhereKernel<T, Context>(dev_ctx,
+                                 atol_compare_result,
+                                 default_rtol_tensor,
+                                 zero_tensor,
+                                 &selected_rtol_tensor);
+    funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
+        dev_ctx,
+        atol,
+        selected_rtol_tensor,
+        GreaterElementFunctor<T>(),
+        &tol_tensor);
+  }
+
+  tol_tensor.Resize(detail::NewAxisDim(tol_tensor.dims(), 1));
+
+  DenseTensor compare_result;
+  compare_result.Resize(detail::NewAxisDim(dim_out, k));
+  dev_ctx.template Alloc<int64_t>(&compare_result);
+
+  funcs::ElementwiseCompute<funcs::GreaterThanFunctor<T, int64_t>, T, int64_t>(
+      dev_ctx,
+      eigenvalue_tensor,
+      tol_tensor,
+      funcs::GreaterThanFunctor<T, int64_t>(),
+      &compare_result);
+
+  phi::SumKernel<int64_t>(dev_ctx,
+                          compare_result,
+                          std::vector<int64_t>{-1},
+                          compare_result.dtype(),
+                          false,
+                          out);
+}
 }  // namespace phi
 
 PD_REGISTER_KERNEL(matrix_rank_tol,  // cuda_only
@@ -458,4 +605,13 @@ PD_REGISTER_KERNEL(matrix_rank_tol,  // cuda_only
   kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
 }
 
+PD_REGISTER_KERNEL(matrix_rank_atol_rtol,  // cuda_only
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::MatrixRankAtolRtolKernel,
+                   float,
+                   double) {
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
+
 #endif  // not PADDLE_WITH_HIP
diff --git a/paddle/phi/kernels/matrix_rank_tol_kernel.h b/paddle/phi/kernels/matrix_rank_tol_kernel.h
@@ -26,4 +26,12 @@ void MatrixRankTolKernel(const Context& dev_ctx,
                          bool hermitian,
                          DenseTensor* out);
 
+template <typename T, typename Context>
+void MatrixRankAtolRtolKernel(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const DenseTensor& atol,
+                              const paddle::optional<DenseTensor>& rtol,
+                              bool hermitian,
+                              DenseTensor* out);
+
 }  // namespace phi
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -3141,6 +3141,15 @@
   kernel :
     func : matrix_rank
 
+- op : matrix_rank_atol_rtol
+  args : (Tensor x, Tensor atol, Tensor rtol, bool hermitian=false)
+  output : Tensor(out)
+  infer_meta :
+    func : MatrixRankAtolRtolInferMeta
+  kernel :
+    func : matrix_rank_atol_rtol
+  optional : rtol
+
 - op : matrix_rank_tol
   args : (Tensor x, Tensor atol_tensor, bool use_default_tol=true, bool hermitian=false)
   output : Tensor(out)
diff --git a/python/paddle/tensor/linalg.py b/python/paddle/tensor/linalg.py
diff --git a/test/legacy_test/test_matrix_rank_atol_rtol_op.py b/test/legacy_test/test_matrix_rank_atol_rtol_op.py
