chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,223 @@
|
||||
// Copyright (c) 2026 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/std_var_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/kernel_utils.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
|
||||
#include "paddle/phi/kernels/funcs/dense_tensor_iterator.h"
|
||||
#include "paddle/phi/kernels/gpu/reduce.h"
|
||||
|
||||
#define C10_HOST_DEVICE __host__ __device__
|
||||
#define C10_DEVICE __device__
|
||||
#define C10_HOST __host__
|
||||
|
||||
#if defined(__CUDACC__) || defined(__HIPCC__)
|
||||
#include <thrust/pair.h>
|
||||
#else
|
||||
#include <cmath>
|
||||
#define device_sqrt std::sqrt
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
#if defined(USE_ROCM)
|
||||
#include <math.h>
|
||||
template <typename scalar_t>
|
||||
static __forceinline__ __device__ scalar_t device_sqrt(scalar_t val);
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ float device_sqrt(float val) {
|
||||
return ::sqrtf(val);
|
||||
}
|
||||
|
||||
template <>
|
||||
__forceinline__ __device__ double device_sqrt(double val) {
|
||||
return ::sqrt(val);
|
||||
}
|
||||
#else
|
||||
template <typename scalar_t>
|
||||
__forceinline__ __device__ double device_sqrt(scalar_t val) {
|
||||
return std::sqrt(val);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
C10_DEVICE __forceinline__ T WARP_SHFL_DOWN(T value,
|
||||
unsigned int delta,
|
||||
int width = warpSize,
|
||||
unsigned int mask = 0xffffffff) {
|
||||
#ifndef __HIPCC__
|
||||
return __shfl_down_sync(mask, value, delta, width);
|
||||
#else
|
||||
return __shfl_down(value, delta, width);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename index_t>
|
||||
struct WelfordData {
|
||||
scalar_t mean;
|
||||
scalar_t m2;
|
||||
index_t n;
|
||||
scalar_t nf;
|
||||
|
||||
C10_HOST_DEVICE WelfordData() : mean(0), m2(0), n(0), nf(0) {}
|
||||
|
||||
C10_HOST_DEVICE WelfordData(scalar_t mean,
|
||||
scalar_t m2,
|
||||
index_t n,
|
||||
scalar_t nf)
|
||||
: mean(mean), m2(m2), n(n), nf(nf) {}
|
||||
};
|
||||
|
||||
template <typename scalar_t,
|
||||
typename acc_scalar_t,
|
||||
typename index_t,
|
||||
typename res_t>
|
||||
struct WelfordOps {
|
||||
acc_scalar_t correction;
|
||||
bool take_sqrt;
|
||||
|
||||
public:
|
||||
using acc_t = WelfordData<acc_scalar_t, index_t>;
|
||||
inline C10_DEVICE acc_t compute(acc_t acc, scalar_t data) const {
|
||||
index_t new_n = acc.n + 1;
|
||||
acc_scalar_t new_nf = static_cast<acc_scalar_t>(new_n);
|
||||
acc_scalar_t delta = static_cast<acc_scalar_t>(data) - acc.mean;
|
||||
acc_scalar_t new_mean = acc.mean + delta / new_nf;
|
||||
acc_scalar_t new_delta = static_cast<acc_scalar_t>(data) - new_mean;
|
||||
return {
|
||||
new_mean,
|
||||
acc.m2 + delta * new_delta,
|
||||
new_n,
|
||||
new_nf,
|
||||
};
|
||||
}
|
||||
inline C10_DEVICE acc_t reduce(acc_t a, acc_t b) const {
|
||||
if (a.nf == 0) {
|
||||
return b;
|
||||
}
|
||||
if (b.nf == 0) {
|
||||
return a;
|
||||
}
|
||||
acc_scalar_t delta = b.mean - a.mean;
|
||||
acc_scalar_t new_count = a.nf + b.nf;
|
||||
acc_scalar_t nb_over_n = b.nf / new_count;
|
||||
return {a.mean + delta * nb_over_n,
|
||||
a.m2 + b.m2 + delta * delta * a.nf * nb_over_n,
|
||||
-1,
|
||||
new_count};
|
||||
}
|
||||
inline C10_DEVICE res_t post_process(acc_t acc) const {
|
||||
const auto mean = static_cast<scalar_t>(acc.mean);
|
||||
const auto divisor = acc.nf > correction ? acc.nf - correction : 0;
|
||||
const auto var = acc.m2 / divisor;
|
||||
res_t results(take_sqrt ? static_cast<scalar_t>(device_sqrt(var))
|
||||
: static_cast<scalar_t>(var),
|
||||
mean);
|
||||
return results;
|
||||
}
|
||||
|
||||
#if defined(__CUDACC__) || defined(__HIPCC__)
|
||||
inline __device__ acc_t shfl_sync(unsigned int mask,
|
||||
acc_t acc,
|
||||
int offset) const {
|
||||
return {WARP_SHFL_DOWN(acc.mean, offset),
|
||||
WARP_SHFL_DOWN(acc.m2, offset),
|
||||
WARP_SHFL_DOWN(acc.n, offset),
|
||||
WARP_SHFL_DOWN(acc.nf, offset)};
|
||||
}
|
||||
#endif
|
||||
C10_HOST_DEVICE WelfordOps(acc_scalar_t correction, bool take_sqrt)
|
||||
: correction(correction), take_sqrt(take_sqrt) {}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
void Std_VarKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const std::vector<int64_t>& axis,
|
||||
bool keepdim,
|
||||
double correction,
|
||||
bool take_sqrt,
|
||||
DenseTensor* out) {
|
||||
if (x.numel() == 0) {
|
||||
Full<T, Context>(dev_ctx, out->dims(), static_cast<T>(NAN), out);
|
||||
return;
|
||||
}
|
||||
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
|
||||
int64_t ndim = x.dims().size();
|
||||
std::vector<int32_t> axis32(axis.begin(), axis.end());
|
||||
auto positive_reduce_dims = ConvertToPositiveDims(axis32, ndim);
|
||||
auto mask = MakeDimMask(positive_reduce_dims, ndim);
|
||||
auto viewed_result = ReviewReduceResult(x, *(out), ndim, mask);
|
||||
|
||||
DenseTensorIteratorConfig dense_iter_config;
|
||||
dense_iter_config.is_reduction(true);
|
||||
dense_iter_config.add_output(viewed_result);
|
||||
dense_iter_config.add_const_input(x);
|
||||
DenseTensorIterator iter = dense_iter_config.build();
|
||||
|
||||
using AccT = typename MPTypeTrait<T>::Type;
|
||||
using ops_t = WelfordOps<T, AccT, int32_t, thrust::pair<T, T>>;
|
||||
ops_t ops(static_cast<AccT>(correction), take_sqrt);
|
||||
|
||||
GPUReduceScheduler<T, T, 2>(dev_ctx, iter, ops, typename ops_t::acc_t{});
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void VarKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const std::vector<int64_t>& axis,
|
||||
bool keepdim,
|
||||
bool unbiased,
|
||||
double correction,
|
||||
DenseTensor* out) {
|
||||
Std_VarKernel<T, Context>(dev_ctx, x, axis, keepdim, correction, false, out);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void StdKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const std::vector<int64_t>& axis,
|
||||
bool keepdim,
|
||||
bool unbiased,
|
||||
double correction,
|
||||
DenseTensor* out) {
|
||||
Std_VarKernel<T, Context>(dev_ctx, x, axis, keepdim, correction, true, out);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(var,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::VarKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(std,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::StdKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
Reference in New Issue
Block a user