wehub-resource-sync
276 lines
8.9 KiB
C++
276 lines
8.9 KiB
C++
// Copyright (c) 2025 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/median_kernel.h"
|
|
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/nanmedian_utils.h"
|
|
#include "paddle/phi/kernels/top_k_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void CalcMedianFunc(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const std::vector<int64_t>& nan_counts,
|
|
const std::vector<int64_t>& nan_indice,
|
|
bool ignore_nan,
|
|
int64_t sort_k,
|
|
int64_t stride,
|
|
int64_t pre_dim,
|
|
T* o_ptr,
|
|
int64_t* m_ptr,
|
|
const std::string& mode) {
|
|
DenseTensor sort_out;
|
|
DenseTensor sort_indices;
|
|
auto sort_dim = x.dims();
|
|
int64_t rank = sort_dim.size();
|
|
sort_dim[static_cast<int>(rank - 1)] = sort_k;
|
|
sort_out.Resize(sort_dim);
|
|
sort_indices.Resize(sort_dim);
|
|
|
|
dev_ctx.template Alloc<T>(&sort_out);
|
|
T* sort_out_ptr = sort_out.data<T>();
|
|
dev_ctx.template Alloc<int64_t>(&sort_indices);
|
|
int64_t* sort_indices_ptr = sort_indices.data<int64_t>();
|
|
|
|
TopkKernel<T, Context>(
|
|
dev_ctx, x, Scalar(sort_k), -1, false, true, &sort_out, &sort_indices);
|
|
|
|
T div_factor = static_cast<T>(2.0);
|
|
int64_t offset = 0;
|
|
int64_t i = 0;
|
|
bool is_ori_odd = stride & 1;
|
|
if (ignore_nan) { // ignore_nan - has nan value; sort_k = max_valid_num
|
|
for (i = 0; i < pre_dim; i++) {
|
|
offset = i * sort_k;
|
|
if (nan_counts[i] == stride) {
|
|
if (mode == "avg") {
|
|
m_ptr[i * 2] = -1;
|
|
m_ptr[i * 2 + 1] = -1; // index is -1
|
|
} else {
|
|
m_ptr[i] = -1;
|
|
}
|
|
o_ptr[i] = sort_out_ptr[offset];
|
|
} else {
|
|
int64_t nan_k = nan_counts[i] > 0
|
|
? static_cast<int64_t>(stride - nan_counts[i])
|
|
: sort_k;
|
|
int64_t row_pos = static_cast<int64_t>(nan_k >> 1);
|
|
int64_t pos = offset + row_pos;
|
|
if (nan_k & 1) {
|
|
if (mode == "avg") {
|
|
m_ptr[2 * i] = sort_indices_ptr[pos];
|
|
m_ptr[2 * i + 1] = sort_indices_ptr[pos];
|
|
} else {
|
|
m_ptr[i] = sort_indices_ptr[pos];
|
|
}
|
|
o_ptr[i] = sort_out_ptr[pos];
|
|
} else {
|
|
// nan_k is even
|
|
T m_val_left =
|
|
row_pos > 0 ? sort_out_ptr[pos - 1] : sort_out_ptr[pos];
|
|
T m_val_right = sort_out_ptr[pos];
|
|
if (mode == "avg") {
|
|
m_ptr[2 * i] =
|
|
row_pos > 0 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
|
|
m_ptr[2 * i + 1] = sort_indices_ptr[pos];
|
|
o_ptr[i] = (m_val_left + m_val_right) / div_factor;
|
|
} else {
|
|
// mode == "min": output median value should be the left val since
|
|
// the sort_out is in ascending order
|
|
m_ptr[i] =
|
|
row_pos > 0 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
|
|
o_ptr[i] = m_val_left;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} else { // not ignore_nan - no nan value; sort_k = stride/2 + 1
|
|
if (is_ori_odd) {
|
|
for (i = 0; i < pre_dim; i++) {
|
|
if (nan_counts[i] > 0) {
|
|
o_ptr[i] = std::numeric_limits<T>::quiet_NaN();
|
|
m_ptr[i] = nan_indice[i];
|
|
continue;
|
|
}
|
|
offset = i * sort_k;
|
|
int64_t pos = offset + sort_k - 1;
|
|
o_ptr[i] = sort_out_ptr[pos];
|
|
if (mode == "avg") {
|
|
m_ptr[2 * i] = sort_indices_ptr[pos];
|
|
m_ptr[2 * i + 1] = sort_indices_ptr[pos];
|
|
} else {
|
|
m_ptr[i] = sort_indices_ptr[pos];
|
|
}
|
|
}
|
|
} else {
|
|
for (i = 0; i < pre_dim; i++) {
|
|
if (nan_counts[i] > 0) {
|
|
o_ptr[i] = std::numeric_limits<T>::quiet_NaN();
|
|
m_ptr[i] = nan_indice[i];
|
|
continue;
|
|
}
|
|
offset = i * sort_k;
|
|
int64_t pos = offset + sort_k - 1;
|
|
T m_val_left = sort_k > 1 ? sort_out_ptr[pos - 1] : sort_out_ptr[pos];
|
|
T m_val_right = sort_out_ptr[pos];
|
|
if (mode == "avg") {
|
|
m_ptr[2 * i] =
|
|
sort_k > 1 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
|
|
m_ptr[2 * i + 1] = sort_indices_ptr[pos];
|
|
o_ptr[i] = (m_val_left + m_val_right) / div_factor;
|
|
} else {
|
|
// mode == "min": output median value should be the left val since the
|
|
// sort_out is in ascending order
|
|
m_ptr[i] =
|
|
sort_k > 1 ? sort_indices_ptr[pos - 1] : sort_indices_ptr[pos];
|
|
o_ptr[i] = m_val_left;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void ProcessMedianKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const std::string& mode,
|
|
bool ignore_nan,
|
|
DenseTensor* out,
|
|
DenseTensor* median_index) {
|
|
const T* x_data = x.data<T>();
|
|
T* out_data = dev_ctx.template Alloc<T>(out);
|
|
int64_t* m_data = dev_ctx.template Alloc<int64_t>(median_index);
|
|
|
|
int64_t numel = x.numel();
|
|
auto x_dim = x.dims();
|
|
int64_t x_rank = x_dim.size();
|
|
int64_t stride = x_dim[static_cast<int>(x_rank - 1)];
|
|
|
|
PADDLE_ENFORCE_NE(stride,
|
|
0,
|
|
common::errors::InvalidArgument(
|
|
"The input Tensor x's shape[-1] should not "
|
|
"be 0, but shape is %s now.",
|
|
x_dim));
|
|
|
|
int64_t pre_dim = numel / stride;
|
|
int64_t i = 0;
|
|
|
|
int64_t max_valid_num = 0;
|
|
std::vector<int64_t> nan_counts;
|
|
std::vector<int64_t> nan_indice;
|
|
|
|
int64_t total_nan_num = 0;
|
|
std::vector<T> col_vec;
|
|
col_vec.reserve(stride);
|
|
col_vec.resize(stride);
|
|
nan_counts.clear();
|
|
nan_counts.reserve(pre_dim);
|
|
nan_counts.resize(pre_dim);
|
|
nan_indice.clear();
|
|
nan_indice.reserve(pre_dim);
|
|
nan_indice.resize(pre_dim);
|
|
for (int64_t i = 0; i < pre_dim; i++) {
|
|
col_vec.clear();
|
|
col_vec.insert(
|
|
col_vec.begin(), x_data + i * stride, x_data + (i + 1) * stride);
|
|
|
|
int64_t first_nan_idx = -1;
|
|
int64_t nan_count = 0;
|
|
|
|
for (int64_t j = 0; j < stride; ++j) {
|
|
if (std::isnan(static_cast<float>(col_vec[j]))) {
|
|
++nan_count;
|
|
if (first_nan_idx == -1) {
|
|
first_nan_idx = j;
|
|
}
|
|
}
|
|
}
|
|
|
|
nan_counts[i] = nan_count;
|
|
nan_indice[i] = first_nan_idx;
|
|
|
|
total_nan_num += nan_count;
|
|
if (stride - nan_count > max_valid_num) {
|
|
max_valid_num = stride - nan_count;
|
|
}
|
|
}
|
|
if (total_nan_num == numel) {
|
|
for (i = 0; i < pre_dim; i++) {
|
|
out_data[i] = std::numeric_limits<T>::quiet_NaN();
|
|
if (mode == "avg") {
|
|
m_data[2 * i] = 0;
|
|
m_data[2 * i + 1] = 1;
|
|
} else {
|
|
m_data[i] = 0;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
int64_t sort_k = ignore_nan ? max_valid_num : ((stride >> 1) + 1);
|
|
CalcMedianFunc<T, Context>(dev_ctx,
|
|
x,
|
|
nan_counts,
|
|
nan_indice,
|
|
ignore_nan,
|
|
sort_k,
|
|
stride,
|
|
pre_dim,
|
|
out_data,
|
|
m_data,
|
|
mode);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void MedianKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const IntArray& axes,
|
|
bool keepdim UNUSED,
|
|
const std::string& mode,
|
|
DenseTensor* out,
|
|
DenseTensor* median_index) {
|
|
if (x.numel() == 0) {
|
|
Full<T, Context>(dev_ctx, out->dims(), NAN, out);
|
|
Full<int64_t, Context>(dev_ctx, median_index->dims(), 0, median_index);
|
|
return;
|
|
}
|
|
DenseTensor tmp_x;
|
|
auto rank = x.dims().size();
|
|
if ((axes.size() == 0) || rank <= 1) {
|
|
tmp_x = x;
|
|
tmp_x.Resize({x.numel()}); // flatten
|
|
} else {
|
|
funcs::PreprocessMedianKernel<T, Context>(
|
|
dev_ctx,
|
|
x,
|
|
axes,
|
|
&tmp_x); // resize to 2D so as to compute median on last axis
|
|
}
|
|
|
|
ProcessMedianKernel<T, Context>(
|
|
dev_ctx, tmp_x, mode, false, out, median_index);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
median, CPU, ALL_LAYOUT, phi::MedianKernel, float, double, int, int64_t) {
|
|
kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
|
|
}
|