update kernel

Author: NKNaN

paddle/phi/infermeta/multiary.cc

@@ -3845,20 +3845,6 @@ void LogspaceInferMeta(const MetaTensor& start,
   out->set_dtype(dtype);
 }
 
-void MatrixRankAtolRtolInferMeta(const MetaTensor& x,
-                                 const MetaTensor& tol,
-                                 const MetaTensor& atol,
-                                 const MetaTensor& rtol,
-                                 bool use_default_tol,
-                                 bool hermitian,
-                                 MetaTensor* out) {
-  if (tol) {
-    MatrixRankTolInferMeta(x, tol, use_default_tol, hermitian, out);
-  } else {
-    MatrixRankTolInferMeta(x, atol, use_default_tol, hermitian, out);
-  }
-}
-
 void MergedAdamInferMeta(
     const std::vector<const MetaTensor*>& param,
     const std::vector<const MetaTensor*>& grad,

paddle/phi/infermeta/multiary.h

@@ -704,14 +704,6 @@ void LogspaceInferMeta(const MetaTensor& start,
                        DataType dtype,
                        MetaTensor* out);
 
-void MatrixRankAtolRtolInferMeta(const MetaTensor& x,
-                                 const MetaTensor& tol,
-                                 const MetaTensor& atol,
-                                 const MetaTensor& rtol,
-                                 bool use_default_tol,
-                                 bool hermitian,
-                                 MetaTensor* out);
-
 void MergedAdamInferMeta(
     const std::vector<const MetaTensor*>& param,
     const std::vector<const MetaTensor*>& grad,

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,

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,

paddle/phi/kernels/cpu/matrix_rank_tol_kernel.cc

@@ -175,128 +175,118 @@ void MatrixRankTolKernel(const Context& dev_ctx,
 template <typename T, typename Context>
 void MatrixRankAtolRtolKernel(const Context& dev_ctx,
                               const DenseTensor& x,
-                              const paddle::optional<DenseTensor>& tol,
-                              const paddle::optional<DenseTensor>& atol,
+                              const DenseTensor& atol,
                               const paddle::optional<DenseTensor>& rtol,
-                              bool use_default_tol,
                               bool hermitian,
                               DenseTensor* out) {
-  if (tol) {
-    MatrixRankTolKernel<T, Context>(
-        dev_ctx, x, *tol, use_default_tol, hermitian, 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 {
-    DenseTensor atol_tensor = *atol;
-    DenseTensor rtol_tensor = *rtol;
-    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 (use_default_tol) {
-      // when `rtol` is specified to be None in py api
-      // use tol=eps*max(m, n)*sigma_1 only if `atol` is specified as 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 tmp_zeros;
-      tmp_zeros = phi::Full<T, Context>(dev_ctx, {1}, static_cast<T>(0.0));
-      DenseTensor atol_compare_result;
-      atol_compare_result.Resize(atol_tensor.dims());
-      phi::EqualKernel<T, Context>(
-          dev_ctx, atol_tensor, tmp_zeros, &atol_compare_result);
-
-      DenseTensor selected_rtol_tensor;
-      selected_rtol_tensor.Resize(rtol_tensor.dims());
-      phi::WhereKernel<T, Context>(dev_ctx,
-                                   atol_compare_result,
-                                   default_rtol_tensor,
-                                   rtol_tensor,
-                                   &selected_rtol_tensor);
-      funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
-          dev_ctx,
-          atol_tensor,
-          selected_rtol_tensor,
-          GreaterElementFunctor<T>(),
-          &tol_tensor);
-    } else {
-      DenseTensor tmp_rtol_tensor;
-      tmp_rtol_tensor =
-          phi::Multiply<T>(dev_ctx, rtol_tensor, max_eigenvalue_tensor);
-      funcs::ElementwiseCompute<GreaterElementFunctor<T>, T>(
-          dev_ctx,
-          atol_tensor,
-          tmp_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);
+    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