#include "stats.h"
#include <cstddef>
#include <numeric>
#include "../collective/aggregator.h"
#include "linalg_op.h"
#include "optional_weight.h"
#include "threading_utils.h"
#include "transform_iterator.h"
#include "xgboost/context.h"
#include "xgboost/host_device_vector.h"
#include "xgboost/linalg.h"
#include "xgboost/logging.h"
namespace xgboost::common {
void Median(Context const* ctx, linalg::Matrix<float> const& t,
HostDeviceVector<float> const& weights, linalg::Tensor<float, 1>* out) {
if (ctx->IsCUDA()) {
weights.SetDevice(ctx->Device());
auto opt_weights = OptionalWeights(weights.ConstDeviceSpan());
auto t_v = t.View(ctx->Device());
cuda_impl::Median(ctx, t_v, opt_weights, out);
}
auto opt_weights = OptionalWeights(weights.ConstHostSpan());
auto t_v = t.HostView();
out->Reshape(t.Shape(1));
auto h_out = out->HostView();
for (std::size_t i{0}; i < t.Shape(1); ++i) {
auto ti_v = t_v.Slice(linalg::All(), i);
auto iter = linalg::cbegin(ti_v);
float q{0};
if (opt_weights.Empty()) {
q = common::Quantile(ctx, 0.5, iter, iter + ti_v.Size());
} else {
CHECK_NE(t_v.Shape(1), 0);
auto w_it = common::MakeIndexTransformIter([&](std::size_t i) { return opt_weights[i]; });
q = common::WeightedQuantile(ctx, 0.5, iter, iter + ti_v.Size(), w_it);
}
h_out(i) = q;
}
}
void Mean(Context const* ctx, linalg::Vector<float> const& v, linalg::Vector<float>* out) {
v.SetDevice(ctx->Device());
out->SetDevice(ctx->Device());
out->Reshape(1);
if (ctx->IsCUDA()) {
cuda_impl::Mean(ctx, v.View(ctx->Device()), out->View(ctx->Device()));
} else {
auto h_v = v.HostView();
float n = v.Size();
MemStackAllocator<float, DefaultMaxThreads()> tloc(ctx->Threads(), 0.0f);
ParallelFor(v.Size(), ctx->Threads(),
[&](auto i) { tloc[omp_get_thread_num()] += h_v(i) / n; });
auto ret = std::accumulate(tloc.cbegin(), tloc.cend(), .0f);
out->HostView()(0) = ret;
}
}
void SampleMean(Context const* ctx, bool is_column_split, linalg::Matrix<float> const& v,
linalg::Vector<float>* out) {
*out = linalg::Zeros<float>(ctx, std::max(v.Shape(1), decltype(v.Shape(1)){1}));
if (!ctx->IsCUDA()) {
auto h_v = v.HostView();
CHECK(h_v.CContiguous());
std::int64_t n_samples = v.Shape(0);
SafeColl(collective::GlobalSum(ctx, is_column_split, linalg::MakeVec(&n_samples, 1)));
auto n_columns = v.Shape(1);
auto h_out = out->HostView();
auto n_rows_f64 = static_cast<double>(n_samples);
for (std::size_t j = 0; j < n_columns; ++j) {
MemStackAllocator<double, DefaultMaxThreads()> mean_tloc(ctx->Threads(), 0.0);
ParallelFor(v.Shape(0), ctx->Threads(),
[&](auto i) { mean_tloc[omp_get_thread_num()] += (h_v(i, j) / n_rows_f64); });
auto mean = std::accumulate(mean_tloc.cbegin(), mean_tloc.cend(), 0.0);
h_out(j) = mean;
}
SafeColl(collective::GlobalSum(ctx, is_column_split, h_out));
} else {
auto d_v = v.View(ctx->Device());
auto d_out = out->View(ctx->Device());
cuda_impl::SampleMean(ctx, is_column_split, d_v, d_out);
}
}
void WeightedSampleMean(Context const* ctx, bool is_column_split, linalg::Matrix<float> const& v,
HostDeviceVector<float> const& w, linalg::Vector<float>* out) {
*out = linalg::Zeros<float>(ctx, std::max(v.Shape(1), decltype(v.Shape(1)){1}));
CHECK_EQ(v.Shape(0), w.Size());
if (!ctx->IsCUDA()) {
auto h_v = v.HostView();
auto h_w = w.ConstHostSpan();
auto sum_w = std::accumulate(h_w.data(), h_w.data() + h_w.size(), 0.0);
SafeColl(collective::GlobalSum(ctx, is_column_split, linalg::MakeVec(&sum_w, 1)));
auto h_out = out->HostView();
for (std::size_t j = 0; j < v.Shape(1); ++j) {
MemStackAllocator<double, DefaultMaxThreads()> mean_tloc(ctx->Threads(), 0.0);
ParallelFor(v.Shape(0), ctx->Threads(),
[&](auto i) { mean_tloc[omp_get_thread_num()] += (h_v(i, j) / sum_w * h_w(i)); });
auto mean = std::accumulate(mean_tloc.cbegin(), mean_tloc.cend(), 0.0);
h_out(j) = mean;
}
SafeColl(collective::GlobalSum(ctx, is_column_split, h_out));
} else {
auto d_v = v.View(ctx->Device());
w.SetDevice(ctx->Device());
auto d_w = w.ConstDeviceSpan();
auto d_out = out->View(ctx->Device());
cuda_impl::WeightedSampleMean(ctx, is_column_split, d_v, d_w, d_out);
}
}
}