chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) 2022 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/abs_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/common/type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/abs_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(abs_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AbsGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(1).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
PD_REGISTER_KERNEL(abs_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AbsDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(1).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
// Copyright (c) 2022 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/abs_kernel.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/complex_functors.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Enable = void>
|
||||
struct CudaAbsFunctor;
|
||||
|
||||
template <typename T>
|
||||
struct CudaAbsFunctor<T, funcs::Complex<T, dtype::Real<T>>> {
|
||||
__device__ __forceinline__ dtype::Real<T> operator()(const T x) const {
|
||||
return abs(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct CudaAbsFunctor<T,
|
||||
std::enable_if_t<std::is_same<T, dtype::Real<T>>::value &&
|
||||
std::is_same<T, bfloat16>::value>> {
|
||||
__device__ __forceinline__ T operator()(const T x) const { return abs(x); }
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct CudaAbsFunctor<T,
|
||||
std::enable_if_t<std::is_same<T, dtype::Real<T>>::value &&
|
||||
!std::is_same<T, bfloat16>::value>> {
|
||||
__device__ __forceinline__ T operator()(const T x) const {
|
||||
return std::abs(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
PADDLE_API void AbsKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
dev_ctx.template Alloc<dtype::Real<T>>(out);
|
||||
std::vector<const DenseTensor*> ins = {&x};
|
||||
std::vector<DenseTensor*> outs = {out};
|
||||
auto functor = CudaAbsFunctor<T>();
|
||||
|
||||
funcs::ElementwiseKernel<dtype::Real<T>>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(abs,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AbsKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->OutputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
// Copyright (c) 2024 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/accuracy_check_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/data_type.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/accuracy_check_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(accuracy_check,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AccuracyCheckKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
bool,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,146 @@
|
||||
// Copyright (c) 2022 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/accuracy_kernel.h"
|
||||
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/reduce.h>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <int BlockSize, typename T>
|
||||
__global__ void AccuracyCudaKernel(const int64_t N,
|
||||
const int D,
|
||||
const int64_t* Xdata,
|
||||
const int64_t* labeldata,
|
||||
int* correct_data,
|
||||
T* accuracy,
|
||||
int* total_data) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
int count = 0;
|
||||
__shared__ int total[BlockSize];
|
||||
|
||||
// support only 1 block
|
||||
for (int64_t i = threadIdx.x; i < (N); i += BlockSize) {
|
||||
for (int j = 0; j < D; ++j) {
|
||||
if (Xdata[i * D + j] == labeldata[i]) {
|
||||
++count;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
total[threadIdx.x] = count;
|
||||
__syncthreads();
|
||||
|
||||
// reduce the count with init value 0, and output accuracy.
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
int result = thrust::reduce(thrust::device, total, total + BlockSize, 0);
|
||||
#else
|
||||
// HIP thrust::reduce not support __device__
|
||||
for (int s = BlockSize / 2; s > 0; s >>= 1) {
|
||||
if (threadIdx.x < s) {
|
||||
total[threadIdx.x] += total[threadIdx.x + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
int result = total[0];
|
||||
#endif
|
||||
if (threadIdx.x == 0) {
|
||||
*correct_data = result;
|
||||
*accuracy = static_cast<T>(static_cast<MT>(result) / static_cast<MT>(N));
|
||||
*total_data = N;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AccuracyKernel(const Context& dev_ctx,
|
||||
const DenseTensor& inference,
|
||||
const DenseTensor& indices,
|
||||
const DenseTensor& label,
|
||||
DenseTensor* accuracy,
|
||||
DenseTensor* correct,
|
||||
DenseTensor* total) {
|
||||
// FIXME(typhoonzero): only support indices currently
|
||||
// if add support for output values, how to detect the data type?
|
||||
const int64_t* indices_data = indices.data<int64_t>();
|
||||
const int64_t* label_data = label.data<int64_t>();
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
inference.dims().size(),
|
||||
2,
|
||||
common::errors::InvalidArgument(
|
||||
"Rank(Input) of AccuracyOp must be 2, with shape "
|
||||
"[sample_number, class_dim], But received rank(Input) is %d",
|
||||
inference.dims().size()));
|
||||
|
||||
int* correct_data = dev_ctx.template Alloc<int>(correct);
|
||||
int* total_data = dev_ctx.template Alloc<int>(total);
|
||||
T* accuracy_data = dev_ctx.template Alloc<T>(accuracy);
|
||||
|
||||
int64_t num_samples = inference.dims()[0];
|
||||
size_t infer_width = inference.dims()[1];
|
||||
auto stream = dev_ctx.stream();
|
||||
backends::gpu::GpuMemsetAsync(accuracy_data, 0, sizeof(T), stream);
|
||||
|
||||
PADDLE_ENFORCE_GT(label.dims().size(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"Rank(Label) of AccuracyOp must greater than 0, "
|
||||
"But received rank(Label) is %d",
|
||||
label.dims().size()));
|
||||
|
||||
PADDLE_ENFORCE_GE(label.dims()[0],
|
||||
inference.dims()[0],
|
||||
common::errors::InvalidArgument(
|
||||
"num_samples(%d) of Label should less than "
|
||||
"or equal to num_samples(%d) of Input",
|
||||
label.dims()[0],
|
||||
num_samples));
|
||||
|
||||
if (num_samples == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
AccuracyCudaKernel<PADDLE_CUDA_NUM_THREADS, T>
|
||||
<<<1, PADDLE_CUDA_NUM_THREADS, 0, stream>>>(num_samples,
|
||||
infer_width,
|
||||
indices_data,
|
||||
label_data,
|
||||
correct_data,
|
||||
accuracy_data,
|
||||
total_data);
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
// FIXME(typhoonzero): types of T is for inference data.
|
||||
// label data is always int64
|
||||
PD_REGISTER_KERNEL(accuracy,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AccuracyKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double) {
|
||||
kernel->InputAt(1).SetDataType(phi::DataType::INT64);
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::INT64);
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
|
||||
}
|
||||
@@ -0,0 +1,759 @@
|
||||
/* Copyright (c) 2022 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/activation_grad_kernel.h"
|
||||
#include "paddle/common/flags.h"
|
||||
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
#include "paddle/phi/kernels/impl/activation_grad_impl.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context, typename Functor>
|
||||
void ActivationGradGPUImpl(const Context& dev_ctx,
|
||||
const DenseTensor* x,
|
||||
const DenseTensor* out,
|
||||
const DenseTensor* d_out,
|
||||
DenseTensor* d_x,
|
||||
const Functor& functor) {
|
||||
if (static_cast<int>(Functor::FwdDeps()) &
|
||||
static_cast<int>(funcs::ActBwdOpFwdDeps::kDepOut)) {
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
out, errors::NotFound("The input DenseTensor Out can not be nullptr"));
|
||||
}
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
d_out, errors::NotFound("The input DenseTensor dOut can not be nullptr"));
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
d_x, errors::NotFound("The output DenseTensor dX can not be nullptr"));
|
||||
|
||||
if (!out) {
|
||||
out = d_out; // fake out
|
||||
}
|
||||
if (static_cast<int>(Functor::FwdDeps()) &
|
||||
static_cast<int>(funcs::ActBwdOpFwdDeps::kDepX)) {
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
x, errors::NotFound("The input DenseTensor X can not be nullptr"));
|
||||
} else {
|
||||
VLOG(10) << "Inplace activation of Op Functor: " << typeid(Functor).name();
|
||||
x = d_x;
|
||||
}
|
||||
|
||||
dev_ctx.template Alloc<T>(d_x);
|
||||
if (d_x->numel() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<const DenseTensor*> ins = {d_out};
|
||||
std::vector<DenseTensor*> outs = {d_x};
|
||||
|
||||
if (static_cast<int>(Functor::FwdDeps()) ==
|
||||
static_cast<int>(funcs::ActBwdOpFwdDeps::kDepOut)) {
|
||||
// Only need forward output Out
|
||||
ins.push_back(out);
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
} else if (static_cast<int>(Functor::FwdDeps()) ==
|
||||
static_cast<int>(funcs::ActBwdOpFwdDeps::kDepX)) {
|
||||
// Only need forward input X
|
||||
ins.push_back(x);
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
} else {
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(name, functor_class) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
const DenseTensor& dout, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX( \
|
||||
name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
const DenseTensor& dout, \
|
||||
float attr, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_DOUBLE_ATTRS_DEPX( \
|
||||
name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
const DenseTensor& dout, \
|
||||
double attr, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPX( \
|
||||
name, functor_class, attr1, attr2) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
const DenseTensor& dout, \
|
||||
float attr1, \
|
||||
float attr2, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr1; \
|
||||
*(attrs[1].second) = attr2; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_DOUBLE_ATTRS_DEPX( \
|
||||
name, functor_class, attr1, attr2) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
const DenseTensor& dout, \
|
||||
double attr1, \
|
||||
double attr2, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr1; \
|
||||
*(attrs[1].second) = attr2; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(name, functor_class) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& out, \
|
||||
const DenseTensor& dout, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPOUT( \
|
||||
name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& out, \
|
||||
const DenseTensor& dout, \
|
||||
float attr, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_DOUBLE_ATTRS_DEPOUT( \
|
||||
name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& out, \
|
||||
const DenseTensor& dout, \
|
||||
double attr, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPOUT( \
|
||||
name, functor_class, attr1, attr2) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel(const Context& dev_ctx, \
|
||||
const DenseTensor& out, \
|
||||
const DenseTensor& dout, \
|
||||
float attr1, \
|
||||
float attr2, \
|
||||
DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr1; \
|
||||
*(attrs[1].second) = attr2; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(name, functor_class) \
|
||||
template <typename T, typename Context> \
|
||||
void name##GradKernel( \
|
||||
const Context& dev_ctx, const DenseTensor& dout, DenseTensor* dx) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
ActivationGradGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, nullptr, nullptr, &dout, dx, functor); \
|
||||
}
|
||||
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu, CudaReluGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Tanh, CudaTanhGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid, CudaSigmoidGradFunctor);
|
||||
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Rint, CudaZeroGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Round, CudaZeroGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Floor, CudaZeroGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_NODEP(Ceil, CudaZeroGradFunctor);
|
||||
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Cos, CudaCosGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Tan, CudaTanGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Acos, CudaAcosGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Sin, CudaSinGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Asin, CudaAsinGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Atan, CudaAtanGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Sinh, CudaSinhGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Cosh, CudaCoshGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Asinh, CudaAsinhGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Acosh, CudaAcoshGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Atanh, CudaAtanhGradFunctor);
|
||||
// TanhShrinkGrad: custom kernel to support FLAGS_use_accuracy_compatible_kernel
|
||||
template <typename T, typename Context>
|
||||
void TanhShrinkGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& dout,
|
||||
DenseTensor* dx) {
|
||||
funcs::CudaTanhShrinkGradFunctor<T> functor;
|
||||
functor.compatible = FLAGS_use_accuracy_compatible_kernel;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaTanhShrinkGradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
}
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Square, CudaSquareGradFunctor);
|
||||
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Exp, CudaExpGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Expm1, CudaExpm1GradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Reciprocal, CudaReciprocalGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Sqrt, CudaSqrtGradFunctor);
|
||||
// RsqrtGrad: custom kernel to support FLAGS_use_accuracy_compatible_kernel
|
||||
template <typename T, typename Context>
|
||||
void RsqrtGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& out,
|
||||
const DenseTensor& dout,
|
||||
DenseTensor* dx) {
|
||||
if (FLAGS_use_accuracy_compatible_kernel) {
|
||||
funcs::CudaRsqrtGradFunctor<T, true> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaRsqrtGradFunctor<T, true>>(
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor);
|
||||
} else {
|
||||
funcs::CudaRsqrtGradFunctor<T, false> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaRsqrtGradFunctor<T, false>>(
|
||||
dev_ctx, nullptr, &out, &dout, dx, functor);
|
||||
}
|
||||
}
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu6, CudaRelu6GradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Softsign, CudaSoftsignGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(LogSigmoid, CudaLogSigmoidGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Log, CudaLogGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Log2, CudaLog2GradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Log10, CudaLog10GradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Log1p, CudaLog1pGradFunctor);
|
||||
DEFINE_GPU_ACTIVATION_GRAD_KERNEL_DEPX(Swish, CudaSwishGradFunctor);
|
||||
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_DOUBLE_ATTRS_DEPX(LeakyRelu,
|
||||
CudaLeakyReluGradFunctor,
|
||||
alpha);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX(SoftShrink,
|
||||
CudaSoftShrinkGradFunctor,
|
||||
lambda);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX(HardShrink,
|
||||
CudaHardShrinkGradFunctor,
|
||||
threshold);
|
||||
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX(Mish,
|
||||
CudaMishGradFunctor,
|
||||
threshold);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX(Celu,
|
||||
CudaCELUGradFunctor,
|
||||
alpha);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_ONE_DOUBLE_ATTRS_DEPOUT(LogitCUDA,
|
||||
CudaLogitGradFunctor,
|
||||
eps);
|
||||
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPX(HardTanh,
|
||||
CudaHardTanhGradFunctor,
|
||||
t_min,
|
||||
t_max);
|
||||
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPX(STanh,
|
||||
CudaSTanhGradFunctor,
|
||||
scale_a,
|
||||
scale_b);
|
||||
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_DOUBLE_ATTRS_DEPX(Softplus,
|
||||
CudaSoftplusGradFunctor,
|
||||
beta,
|
||||
threshold);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPOUT(HardSigmoid,
|
||||
CudaHardSigmoidGradFunctor,
|
||||
slope,
|
||||
offset);
|
||||
DEFINE_GPU_ACT_GRAD_KERNEL_WITH_TWO_ATTRS_DEPX(ThresholdedRelu,
|
||||
CudaThresholdedReluGradFunctor,
|
||||
threshold,
|
||||
value);
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SiluGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& out,
|
||||
const DenseTensor& dout,
|
||||
DenseTensor* dx) {
|
||||
funcs::CudaSiluGradFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaSiluGradFunctor<T>>(
|
||||
dev_ctx, &x, &out, &dout, dx, functor);
|
||||
}
|
||||
template <typename T, typename Context>
|
||||
void EluGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& out,
|
||||
const DenseTensor& dout,
|
||||
float alpha,
|
||||
DenseTensor* dx) {
|
||||
dev_ctx.template Alloc<T>(dx);
|
||||
if (dx->numel() == 0) {
|
||||
return;
|
||||
}
|
||||
std::vector<const DenseTensor*> ins = {&dout, &out};
|
||||
std::vector<DenseTensor*> outs = {dx};
|
||||
if (alpha > 0) {
|
||||
funcs::CudaELUGradFunctor<T> functor;
|
||||
functor.alpha = alpha;
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
} else {
|
||||
funcs::CudaELUGradNegativeAlphaFunctor<T> functor;
|
||||
functor.alpha = alpha;
|
||||
ins.push_back(&x);
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void HardSwishGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& dout,
|
||||
DenseTensor* dx) {
|
||||
funcs::CudaHardSwishGradFunctor<T> functor;
|
||||
float threshold = 6;
|
||||
float scale = 6;
|
||||
float offset = 3;
|
||||
auto attrs = functor.GetAttrs();
|
||||
*(attrs[0].second) = threshold;
|
||||
*(attrs[1].second) = scale;
|
||||
*(attrs[2].second) = offset;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaHardSwishGradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void PowGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& dout,
|
||||
const Scalar& factor,
|
||||
DenseTensor* dx) {
|
||||
if (factor.to<double>() == 0) {
|
||||
std::vector<int64_t> vec_dims = vectorize(dx->dims());
|
||||
Full<T, Context>(dev_ctx, IntArray(vec_dims), static_cast<T>(0), dx);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 1) {
|
||||
std::vector<int64_t> vec_dims = vectorize(dx->dims());
|
||||
Copy<Context>(dev_ctx, dout, dev_ctx.GetPlace(), false, dx);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 2) {
|
||||
funcs::CudaSquareGradFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaSquareGradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 3) {
|
||||
funcs::CudaCubeGradFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaCubeGradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 4) {
|
||||
funcs::CudaPow4GradFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaPow4GradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
if constexpr (!std::is_integral<T>::value) {
|
||||
if (factor.to<double>() == 1.5) {
|
||||
funcs::CudaPow1p5GradFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaPow1p5GradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 0.5) {
|
||||
funcs::CudaSqrtGradDepXFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaSqrtGradDepXFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == -1) {
|
||||
funcs::CudaReciprocalGradDepXFunctor<T> functor;
|
||||
ActivationGradGPUImpl<T,
|
||||
Context,
|
||||
funcs::CudaReciprocalGradDepXFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
return;
|
||||
}
|
||||
}
|
||||
funcs::CudaPowGradFunctor<T> functor;
|
||||
functor.SetFactor(factor.to<double>());
|
||||
ActivationGradGPUImpl<T, Context, funcs::CudaPowGradFunctor<T>>(
|
||||
dev_ctx, &x, nullptr, &dout, dx, functor);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PD_REGISTER_KERNEL(relu_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ReluGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
PD_REGISTER_KERNEL(relu_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ReluDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
#else
|
||||
PD_REGISTER_KERNEL(relu_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ReluGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(relu_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ReluDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
#define PD_REGISTER_ACTIVATION_GRAD_KERNEL(name, func) \
|
||||
PD_REGISTER_KERNEL(name, \
|
||||
GPU, \
|
||||
ALL_LAYOUT, \
|
||||
phi::func, \
|
||||
float, \
|
||||
double, \
|
||||
phi::float16, \
|
||||
phi::bfloat16) {}
|
||||
|
||||
#define PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(name, func) \
|
||||
PD_REGISTER_KERNEL(name, \
|
||||
GPU, \
|
||||
ALL_LAYOUT, \
|
||||
phi::func, \
|
||||
float, \
|
||||
double, \
|
||||
phi::float16, \
|
||||
phi::bfloat16, \
|
||||
phi::complex64, \
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sin_grad, SinGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(cos_grad, CosGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(tan_grad, TanGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(acos_grad, AcosGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(asin_grad, AsinGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(atan_grad, AtanGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sinh_grad, SinhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(cosh_grad, CoshGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(asinh_grad, AsinhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(acosh_grad, AcoshGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(atanh_grad, AtanhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(tanh_grad, TanhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(tanh_double_grad,
|
||||
TanhDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(tanh_triple_grad,
|
||||
TanhTripleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(hardtanh_grad, HardTanhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(leaky_relu_grad, LeakyReluGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(leaky_relu_double_grad,
|
||||
LeakyReluDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(thresholded_relu_grad,
|
||||
ThresholdedReluGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(relu6_grad, Relu6GradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(mish_grad, MishGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(stanh_grad, STanhGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(reciprocal_grad,
|
||||
ReciprocalGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(softplus_grad,
|
||||
SoftplusGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(softplus_double_grad,
|
||||
SoftplusDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sqrt_grad, SqrtGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(sqrt_double_grad, SqrtDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(rsqrt_grad, RsqrtGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(rsqrt_double_grad, RsqrtDoubleGradKernel)
|
||||
|
||||
PD_REGISTER_KERNEL(exp_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ExpGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(softshrink_grad, SoftShrinkGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(hard_shrink_grad, HardShrinkGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(tanh_shrink_grad, TanhShrinkGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(silu_grad, SiluGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(elu_grad, EluGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(elu_double_grad, EluDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(logit_grad, LogitCUDAGradKernel)
|
||||
|
||||
PD_REGISTER_KERNEL(expm1_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Expm1GradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(square_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SquareGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(square_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SquareDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(sin_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SinDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(sin_triple_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SinTripleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(cos_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CosDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(cos_triple_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CosTripleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(softsign_grad,
|
||||
SoftsignGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sigmoid_grad, SigmoidGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sigmoid_double_grad,
|
||||
SigmoidDoubleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(sigmoid_triple_grad,
|
||||
SigmoidTripleGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(hardsigmoid_grad, HardSigmoidGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(logsigmoid_grad,
|
||||
LogSigmoidGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(log_grad, LogGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(log2_grad, Log2GradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(log10_grad, Log10GradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(log1p_grad, Log1pGradKernel)
|
||||
PD_REGISTER_KERNEL(log_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::LogDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL_WITH_COMPLEX(hardswish_grad,
|
||||
HardSwishGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(swish_grad, SwishGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(celu_grad, CeluGradKernel)
|
||||
PD_REGISTER_ACTIVATION_GRAD_KERNEL(celu_double_grad, CeluDoubleGradKernel)
|
||||
|
||||
PD_REGISTER_KERNEL(rint_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::RintGradKernel,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(round_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::RoundGradKernel,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(pow_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::PowGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(pow_double_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::PowDoubleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(pow_triple_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::PowTripleGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(ceil_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CeilGradKernel,
|
||||
float,
|
||||
double,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(floor_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::FloorGradKernel,
|
||||
float,
|
||||
double,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,532 @@
|
||||
/* Copyright (c) 2022 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/activation_kernel.h"
|
||||
#include "paddle/common/flags.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
#include "paddle/phi/kernels/impl/activation_grad_impl.h"
|
||||
#include "paddle/phi/kernels/impl/activation_impl.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context, typename Functor>
|
||||
void ActivationGPUImpl(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out,
|
||||
const Functor& functor) {
|
||||
PADDLE_ENFORCE_NOT_NULL(out,
|
||||
errors::NotFound("Output Out should not be nullptr"));
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
if (out->numel() == 0) {
|
||||
return;
|
||||
}
|
||||
std::vector<const DenseTensor*> ins = {&x};
|
||||
std::vector<DenseTensor*> outs = {out};
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACTIVATION_KERNEL(name, functor_class) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel( \
|
||||
const Context& dev_ctx, const DenseTensor& x, DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
ActivationGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(name, \
|
||||
functor_class) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel( \
|
||||
const Context& dev_ctx, const DenseTensor& x, DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
using U = \
|
||||
typename std::conditional_t<std::is_integral<T>::value, float, T>; \
|
||||
ActivationGPUImpl<U, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
float attr, \
|
||||
DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_KERNEL_WITH_ONE_DOUBLE_ATTRS(name, functor_class, attr) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
double attr, \
|
||||
DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr; \
|
||||
ActivationGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS( \
|
||||
name, functor_class, attr1, attr2) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
float attr1, \
|
||||
float attr2, \
|
||||
DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr1; \
|
||||
*(attrs[1].second) = attr2; \
|
||||
ActivationGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
#define DEFINE_GPU_ACT_KERNEL_WITH_TWO_DOUBLE_ATTRS( \
|
||||
name, functor_class, attr1, attr2) \
|
||||
template <typename T, typename Context> \
|
||||
void name##Kernel(const Context& dev_ctx, \
|
||||
const DenseTensor& x, \
|
||||
double attr1, \
|
||||
double attr2, \
|
||||
DenseTensor* out) { \
|
||||
funcs::functor_class<T> functor; \
|
||||
auto attrs = functor.GetAttrs(); \
|
||||
*(attrs[0].second) = attr1; \
|
||||
*(attrs[1].second) = attr2; \
|
||||
ActivationGPUImpl<T, Context, funcs::functor_class<T>>( \
|
||||
dev_ctx, x, out, functor); \
|
||||
}
|
||||
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Cos, CudaCosFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Tan, CudaTanFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Acos, CudaAcosFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Sin, CudaSinFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Asin, CudaAsinFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Atan, CudaAtanFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Sinh, CudaSinhFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Cosh, CudaCoshFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Asinh, CudaAsinhFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Acosh, CudaAcoshFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Atanh, CudaAtanhFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Relu, CudaReluFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Tanh, CudaTanhFunctor)
|
||||
// TanhShrink: custom kernel to support FLAGS_use_accuracy_compatible_kernel
|
||||
template <typename T, typename Context>
|
||||
void TanhShrinkKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
funcs::CudaTanhShrinkFunctor<T> functor;
|
||||
functor.compatible = FLAGS_use_accuracy_compatible_kernel;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaTanhShrinkFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Silu, CudaSiluFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Reciprocal, CudaReciprocalFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Square, CudaSquareFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Sqrt, CudaSqrtFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Rsqrt, CudaRsqrtFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Softsign, CudaSoftsignFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Sigmoid, CudaSigmoidFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(LogSigmoid, CudaLogSigmoidFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Floor, CudaFloorFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Ceil, CudaCeilFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL(Rint, CudaRintFunctor)
|
||||
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log, CudaLogFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log2, CudaLog2Functor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log10, CudaLog10Functor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Log1p, CudaLog1pFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Exp, CudaExpFunctor)
|
||||
DEFINE_GPU_ACTIVATION_KERNEL_WITH_INT_IN_FLOAT_OUT(Expm1, CudaExpm1Functor)
|
||||
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_DOUBLE_ATTRS(LeakyRelu,
|
||||
CudaLeakyReluFunctor,
|
||||
alpha)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_DOUBLE_ATTRS(LogitCUDA, CudaLogitFunctor, eps)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(HardShrink,
|
||||
CudaHardShrinkFunctor,
|
||||
threshold)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(SoftShrink, CudaSoftShrinkFunctor, lambda)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(Elu, CudaELUFunctor, alpha)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(Mish, CudaMishFunctor, threshold)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_ONE_ATTRS(Celu, CudaCELUFunctor, alpha)
|
||||
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS(HardTanh,
|
||||
CudaHardTanhFunctor,
|
||||
t_min,
|
||||
t_max)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS(Stanh, CudaSTanhFunctor, scale_a, scale_b)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_DOUBLE_ATTRS(Softplus,
|
||||
CudaSoftplusFunctor,
|
||||
beta,
|
||||
threshold)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS(HardSigmoid,
|
||||
CudaHardSigmoidFunctor,
|
||||
slope,
|
||||
offset)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS(Selu, CudaSeluFunctor, scale, alpha)
|
||||
DEFINE_GPU_ACT_KERNEL_WITH_TWO_ATTRS(ThresholdedRelu,
|
||||
CudaThresholdedReluFunctor,
|
||||
threshold,
|
||||
value)
|
||||
|
||||
template <typename T, typename Context>
|
||||
void HardSwishKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
funcs::CudaHardSwishFunctor<T> functor;
|
||||
float threshold = 6;
|
||||
float scale = 6;
|
||||
float offset = 3;
|
||||
auto attrs = functor.GetAttrs();
|
||||
*(attrs[0].second) = threshold;
|
||||
*(attrs[1].second) = scale;
|
||||
*(attrs[2].second) = offset;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaHardSwishFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SwishKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
funcs::CudaSwishFunctor<T> functor;
|
||||
auto attrs = functor.GetAttrs();
|
||||
*(attrs[0].second) = 1.0;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaSwishFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void Relu6Kernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
funcs::CudaRelu6Functor<T> functor;
|
||||
auto attrs = functor.GetAttrs();
|
||||
*(attrs[0].second) = 6.0;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaRelu6Functor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void RoundKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const int decimals,
|
||||
DenseTensor* out) {
|
||||
funcs::CudaRoundFunctor<T> functor;
|
||||
auto attrs = functor.GetAttrs();
|
||||
*(attrs[0].second) = decimals;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaRoundFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void PowKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const Scalar& factor,
|
||||
DenseTensor* out) {
|
||||
if constexpr (std::is_integral<T>::value) {
|
||||
PADDLE_ENFORCE_GE(
|
||||
factor.to<double>(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"Integers to negative integer powers are not allowed."));
|
||||
} else {
|
||||
if (factor.to<double>() == 0.5) {
|
||||
funcs::CudaSqrtFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaSqrtFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == -0.5) {
|
||||
funcs::CudaRsqrtFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaRsqrtFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == -1) {
|
||||
funcs::CudaReciprocalFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaReciprocalFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == -2) {
|
||||
funcs::CudaRsquareFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaRsquareFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (factor.to<double>() == 0) {
|
||||
std::vector<int64_t> vec_dims = vectorize(out->dims());
|
||||
Full<T, Context>(dev_ctx, IntArray(vec_dims), static_cast<T>(1), out);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 1) {
|
||||
Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 2) {
|
||||
funcs::CudaSquareFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaSquareFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
if (factor.to<double>() == 3) {
|
||||
funcs::CudaCubeFunctor<T> functor;
|
||||
ActivationGPUImpl<T, Context, funcs::CudaCubeFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
return;
|
||||
}
|
||||
|
||||
funcs::CudaPowFunctor<T> functor;
|
||||
functor.SetFactor(factor.to<double>());
|
||||
ActivationGPUImpl<T, Context, funcs::CudaPowFunctor<T>>(
|
||||
dev_ctx, x, out, functor);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PD_REGISTER_KERNEL(
|
||||
relu, GPU, ALL_LAYOUT, phi::ReluKernel, float, double, phi::float16) {}
|
||||
#else
|
||||
PD_REGISTER_KERNEL(relu,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ReluKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
|
||||
#define PD_REGISTER_ACTIVATION_KERNEL(name, func) \
|
||||
PD_REGISTER_KERNEL(name, \
|
||||
GPU, \
|
||||
ALL_LAYOUT, \
|
||||
phi::func, \
|
||||
float, \
|
||||
double, \
|
||||
phi::float16, \
|
||||
phi::bfloat16) {}
|
||||
|
||||
#define PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(name, func) \
|
||||
PD_REGISTER_KERNEL(name, \
|
||||
GPU, \
|
||||
ALL_LAYOUT, \
|
||||
phi::func, \
|
||||
float, \
|
||||
double, \
|
||||
phi::float16, \
|
||||
phi::bfloat16, \
|
||||
phi::complex64, \
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(sin, SinKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(cos, CosKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(tan, TanKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(acos, AcosKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(asin, AsinKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(atan, AtanKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(sinh, SinhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(cosh, CoshKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(asinh, AsinhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(acosh, AcoshKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(atanh, AtanhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(tanh, TanhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(hardtanh, HardTanhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(thresholded_relu, ThresholdedReluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(relu6, Relu6Kernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(leaky_relu, LeakyReluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(mish, MishKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(stanh, StanhKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(reciprocal, ReciprocalKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(sqrt, SqrtKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(rsqrt, RsqrtKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(softplus, SoftplusKernel)
|
||||
|
||||
PD_REGISTER_KERNEL(exp,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ExpKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(expm1,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Expm1Kernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(square,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SquareKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_ACTIVATION_KERNEL(hard_shrink, HardShrinkKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(softshrink, SoftShrinkKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(tanh_shrink, TanhShrinkKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(elu, EluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(silu, SiluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(softsign, SoftsignKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(sigmoid, SigmoidKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(logsigmoid, LogSigmoidKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(hardsigmoid, HardSigmoidKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL_WITH_COMPLEX(hardswish, HardSwishKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(swish, SwishKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(celu, CeluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(selu, SeluKernel)
|
||||
PD_REGISTER_ACTIVATION_KERNEL(logit, LogitCUDAKernel)
|
||||
|
||||
PD_REGISTER_KERNEL(rint,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::RintKernel,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(round,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::RoundKernel,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(log,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::LogKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(log2,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Log2Kernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(log10,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Log10Kernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(log1p,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Log1pKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(pow,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::PowKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
PD_REGISTER_KERNEL(ceil,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CeilKernel,
|
||||
float,
|
||||
double,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
PD_REGISTER_KERNEL(floor,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::FloorKernel,
|
||||
float,
|
||||
double,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright (c) 2022 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/adadelta_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/adadelta_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(adadelta,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AdadeltaKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
// Copyright (c) 2022 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/adagrad_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
|
||||
#include "paddle/phi/kernels/impl/adagrad_kernel_impl.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void AdagradGPUKernel(const T* param,
|
||||
const T* grad,
|
||||
const MT* moment,
|
||||
const MT* lr,
|
||||
const MT* master_param,
|
||||
MT epsilon,
|
||||
T* param_out,
|
||||
MT* moment_out,
|
||||
MT* master_param_out,
|
||||
int64_t num) {
|
||||
int64_t idx =
|
||||
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(blockIdx.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
MT lr_data = static_cast<MT>(lr[0]);
|
||||
|
||||
for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) {
|
||||
MT grad_data = static_cast<MT>(grad[i]);
|
||||
MT moment_out_data = static_cast<MT>(moment[i]) + grad_data * grad_data;
|
||||
moment_out[i] = static_cast<MT>(moment_out_data);
|
||||
auto in = master_param_out ? master_param[i] : static_cast<MT>(param[i]);
|
||||
MT param_out_data =
|
||||
in - (lr_data * grad_data) / (sqrt(moment_out_data) + epsilon);
|
||||
|
||||
param_out[i] = static_cast<T>(param_out_data);
|
||||
|
||||
if (master_param_out) {
|
||||
master_param_out[i] = param_out_data;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct DenseAdagradFunctor<GPUContext, T> {
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const DenseTensor& param_t,
|
||||
const DenseTensor& grad_t,
|
||||
const DenseTensor& moment_t,
|
||||
const DenseTensor& learning_rate,
|
||||
const optional<DenseTensor>& master_param,
|
||||
float epsilon_t,
|
||||
bool multi_precision,
|
||||
DenseTensor* param_out_tensor,
|
||||
DenseTensor* moment_out_tensor,
|
||||
DenseTensor* master_param_outs) {
|
||||
using MT = typename dtype::template MPTypeTrait<T>::Type;
|
||||
T* param_out_data = dev_ctx.template Alloc<T>(param_out_tensor);
|
||||
MT* moment_out_data = dev_ctx.template Alloc<MT>(moment_out_tensor);
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data = multi_precision
|
||||
? dev_ctx.template Alloc<MT>(master_param_outs)
|
||||
: nullptr;
|
||||
|
||||
MT epsilon = static_cast<MT>(epsilon_t);
|
||||
|
||||
int64_t numel = param_t.numel();
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 1);
|
||||
int grid = config.block_per_grid.x;
|
||||
int block = config.thread_per_block.x;
|
||||
auto stream = dev_ctx.stream();
|
||||
AdagradGPUKernel<T, MT>
|
||||
<<<block, grid, 0, stream>>>(param_t.data<T>(),
|
||||
grad_t.data<T>(),
|
||||
moment_t.data<MT>(),
|
||||
learning_rate.data<MT>(),
|
||||
master_in_data,
|
||||
epsilon,
|
||||
param_out_data,
|
||||
moment_out_data,
|
||||
master_out_data,
|
||||
numel);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, int block_size>
|
||||
__global__ void MergeGradKernel(const T* grad,
|
||||
const int64_t* grad_rows,
|
||||
T* grad_merge,
|
||||
const int64_t* grad_merge_rows,
|
||||
size_t grad_merge_rows_size,
|
||||
int64_t row_numel) {
|
||||
const int ty = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
__shared__ size_t grad_merge_idx;
|
||||
|
||||
if (tid == 0) {
|
||||
for (size_t i = 0; i < grad_merge_rows_size; i++) {
|
||||
if (grad_rows[ty] == grad_merge_rows[i]) {
|
||||
grad_merge_idx = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
grad += ty * row_numel;
|
||||
grad_merge += grad_merge_idx * row_numel;
|
||||
for (int64_t index = tid; index < row_numel; index += block_size) {
|
||||
CudaAtomicAdd(grad_merge + index, grad[index]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, int block_size>
|
||||
__global__ void SparseAdagradFunctorKernel(const T* grad,
|
||||
const int64_t* rows,
|
||||
const T* learning_rate,
|
||||
T* param,
|
||||
T* moment,
|
||||
int64_t row_numel,
|
||||
T epsilon) {
|
||||
const int ty = blockIdx.y;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
grad += ty * row_numel;
|
||||
param += rows[ty] * row_numel;
|
||||
moment += rows[ty] * row_numel;
|
||||
|
||||
for (int64_t index = tid; index < row_numel; index += block_size) {
|
||||
// Since index in rows of SelectedRows can be duplicate, we have to use
|
||||
// Atomic Operation to avoid concurrent write error.
|
||||
CudaAtomicAdd(param + index,
|
||||
-1.0 * learning_rate[0] * grad[index] /
|
||||
(sqrt(moment[index]) + epsilon));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct SparseAdagradFunctor<GPUContext, T> {
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const SelectedRows& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
T epsilon,
|
||||
DenseTensor* moment,
|
||||
DenseTensor* param) {
|
||||
// 1. g_m.rows = set(g.rows)
|
||||
auto grad_width = grad.value().dims()[1];
|
||||
funcs::scatter::MergeAdd<GPUContext, T> merge_func;
|
||||
auto grad_merge = merge_func(dev_ctx, grad);
|
||||
auto* grad_merge_data = grad_merge.mutable_value()->template data<T>();
|
||||
Vector<int64_t> merge_rows(grad_merge.rows());
|
||||
// 2. m += g_m * g_m
|
||||
auto grad_square = SquareSelectedRows<GPUContext, T>(dev_ctx, grad_merge);
|
||||
|
||||
funcs::SelectedRowsAddToTensor<GPUContext, T> functor;
|
||||
functor(dev_ctx, grad_square, moment);
|
||||
|
||||
// 3. update parameter
|
||||
auto* lr = learning_rate.data<T>();
|
||||
auto* param_data = param->data<T>();
|
||||
auto* moment_data = moment->data<T>();
|
||||
|
||||
const int block_size = 256;
|
||||
dim3 threads(block_size, 1);
|
||||
dim3 grid2(1, merge_rows.size());
|
||||
MixVector<int64_t> mixv_merge_rows(&merge_rows);
|
||||
SparseAdagradFunctorKernel<T, 256>
|
||||
<<<grid2,
|
||||
threads,
|
||||
0,
|
||||
reinterpret_cast<const GPUContext&>(dev_ctx).stream()>>>(
|
||||
grad_merge_data,
|
||||
mixv_merge_rows.CUDAMutableData(dev_ctx.GetPlace()),
|
||||
lr,
|
||||
param_data,
|
||||
moment_data,
|
||||
grad_width,
|
||||
epsilon);
|
||||
mixv_merge_rows.CopyToCPU();
|
||||
}
|
||||
};
|
||||
|
||||
template struct SparseAdagradFunctor<GPUContext, float>;
|
||||
template struct SparseAdagradFunctor<GPUContext, double>;
|
||||
template struct DenseAdagradFunctor<GPUContext, float>;
|
||||
template struct DenseAdagradFunctor<GPUContext, double>;
|
||||
template struct DenseAdagradFunctor<GPUContext, float16>;
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(adagrad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AdagradDenseKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(adagrad_dense_param_sparse_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AdagradSparseKernel,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,991 @@
|
||||
// Copyright (c) 2022 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/adam_kernel.h"
|
||||
|
||||
#include <math.h> // for sqrt in CPU and CUDA
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/common/flags.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/adam_functors.h"
|
||||
#include "paddle/phi/kernels/funcs/for_range.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename TG, typename MT>
|
||||
__global__ void AdamKernelREG(MT beta1,
|
||||
MT beta2,
|
||||
MT epsilon,
|
||||
MT beta1_pow_,
|
||||
MT beta2_pow_,
|
||||
const MT* moment1,
|
||||
MT* moment1_out,
|
||||
const MT* moment2,
|
||||
MT* moment2_out,
|
||||
const MT* moment2_max,
|
||||
MT* moment2_max_out,
|
||||
const double* lr_,
|
||||
const TG* grad,
|
||||
const T* param,
|
||||
T* param_out,
|
||||
const MT* master_param,
|
||||
MT* master_param_out,
|
||||
int64_t ndim,
|
||||
bool amsgrad) {
|
||||
MT lr = static_cast<MT>(*lr_);
|
||||
MT beta1_pow = beta1_pow_;
|
||||
MT beta2_pow = beta2_pow_;
|
||||
|
||||
int64_t id =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
|
||||
for (; id < ndim; id += gridDim.x * blockDim.x) {
|
||||
MT p = master_param ? master_param[id] : static_cast<MT>(param[id]);
|
||||
MT g = static_cast<MT>(grad[id]);
|
||||
MT mom1 = static_cast<MT>(moment1[id]);
|
||||
MT mom2 = static_cast<MT>(moment2[id]);
|
||||
|
||||
mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
|
||||
mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
|
||||
|
||||
MT denom;
|
||||
if (amsgrad) {
|
||||
MT mom2_max = static_cast<MT>(moment2_max[id]);
|
||||
MT mom2_max_ = std::max(mom2, mom2_max);
|
||||
moment2_max_out[id] = mom2_max_;
|
||||
|
||||
denom =
|
||||
(sqrt(mom2_max_) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
} else {
|
||||
denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
}
|
||||
|
||||
p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
|
||||
|
||||
moment1_out[id] = mom1;
|
||||
moment2_out[id] = mom2;
|
||||
param_out[id] = static_cast<T>(p);
|
||||
if (master_param_out) {
|
||||
master_param_out[id] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename TG, typename MT>
|
||||
__global__ void AdamKernelMEM(MT beta1,
|
||||
MT beta2,
|
||||
MT epsilon,
|
||||
const MT* beta1_pow_,
|
||||
const MT* beta2_pow_,
|
||||
const MT* moment1,
|
||||
MT* moment1_out,
|
||||
const MT* moment2,
|
||||
MT* moment2_out,
|
||||
const MT* moment2_max,
|
||||
MT* moment2_max_out,
|
||||
const double* lr_,
|
||||
const TG* grad,
|
||||
const T* param,
|
||||
T* param_out,
|
||||
const MT* master_param,
|
||||
MT* master_param_out,
|
||||
int64_t ndim,
|
||||
bool amsgrad) {
|
||||
MT lr = static_cast<MT>(*lr_);
|
||||
MT beta1_pow = *beta1_pow_;
|
||||
MT beta2_pow = *beta2_pow_;
|
||||
|
||||
int64_t id =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
|
||||
for (; id < ndim; id += gridDim.x * blockDim.x) {
|
||||
MT p = master_param ? master_param[id] : static_cast<MT>(param[id]);
|
||||
MT g = static_cast<MT>(grad[id]);
|
||||
MT mom1 = static_cast<MT>(moment1[id]);
|
||||
MT mom2 = static_cast<MT>(moment2[id]);
|
||||
|
||||
mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
|
||||
mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
|
||||
|
||||
MT denom;
|
||||
if (amsgrad) {
|
||||
MT mom2_max = static_cast<MT>(moment2_max[id]);
|
||||
MT mom2_max_ = std::max(mom2, mom2_max);
|
||||
moment2_max_out[id] = mom2_max_;
|
||||
|
||||
denom =
|
||||
(sqrt(mom2_max_) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
} else {
|
||||
denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
}
|
||||
|
||||
p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
|
||||
|
||||
moment1_out[id] = mom1;
|
||||
moment2_out[id] = mom2;
|
||||
param_out[id] = static_cast<T>(p);
|
||||
if (master_param_out) {
|
||||
master_param_out[id] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void UpdateBetaPow(T beta1,
|
||||
T beta2,
|
||||
const T* beta1_pow_,
|
||||
const T* beta2_pow_,
|
||||
T* beta1_pow_out,
|
||||
T* beta2_pow_out) {
|
||||
*beta1_pow_out = beta1 * beta1_pow_[0];
|
||||
*beta2_pow_out = beta2 * beta2_pow_[0];
|
||||
}
|
||||
|
||||
// ---- Torch-compatible Adam infrastructure ----
|
||||
|
||||
// LrAccessor for Adam: lr tensor is MT (float), upcast to double on device.
|
||||
template <typename MT, bool IsCpu>
|
||||
struct AdamLrAccessor;
|
||||
|
||||
template <typename MT>
|
||||
struct AdamLrAccessor<MT, true> {
|
||||
const double lr_double;
|
||||
explicit AdamLrAccessor(double lr) : lr_double(lr) {}
|
||||
__device__ __forceinline__ double GetLrDouble() const { return lr_double; }
|
||||
};
|
||||
|
||||
template <typename MT>
|
||||
struct AdamLrAccessor<MT, false> {
|
||||
const double* lr;
|
||||
explicit AdamLrAccessor(const double* lr) : lr(lr) {}
|
||||
__device__ __forceinline__ double GetLrDouble() const { return *lr; }
|
||||
};
|
||||
|
||||
// Bias correction accessors matching torch's step_count-based computation.
|
||||
template <typename MT, bool IsCpu>
|
||||
struct AdamBiasCorrAccessorCompat;
|
||||
|
||||
template <typename MT>
|
||||
struct AdamBiasCorrAccessorCompat<MT, true> {
|
||||
const double beta1;
|
||||
const double beta2;
|
||||
const float step_count;
|
||||
|
||||
AdamBiasCorrAccessorCompat(double b1, double b2, float sc)
|
||||
: beta1(b1), beta2(b2), step_count(sc) {}
|
||||
|
||||
__device__ __forceinline__ double GetBc1() const {
|
||||
return 1.0 - ::pow(beta1, step_count);
|
||||
}
|
||||
__device__ __forceinline__ double GetBc2() const {
|
||||
return 1.0 - ::pow(beta2, step_count);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename MT>
|
||||
struct AdamBiasCorrAccessorCompat<MT, false> {
|
||||
const double beta1;
|
||||
const double beta2;
|
||||
const MT* beta1_pow;
|
||||
|
||||
AdamBiasCorrAccessorCompat(double b1, double b2, const MT* bp1)
|
||||
: beta1(b1), beta2(b2), beta1_pow(bp1) {}
|
||||
|
||||
__device__ __forceinline__ double GetBc1() const {
|
||||
const float sc = static_cast<float>(
|
||||
::round(::log(static_cast<double>(*beta1_pow)) / ::log(beta1)));
|
||||
return 1.0 - ::pow(beta1, sc);
|
||||
}
|
||||
__device__ __forceinline__ double GetBc2() const {
|
||||
const float sc = static_cast<float>(
|
||||
::round(::log(static_cast<double>(*beta1_pow)) / ::log(beta1)));
|
||||
return 1.0 - ::pow(beta2, sc);
|
||||
}
|
||||
};
|
||||
|
||||
// Torch-compatible Adam kernel: no weight decay, float lr upcast to double,
|
||||
// FMA-based moment updates matching torch's fused adam math.
|
||||
template <typename T,
|
||||
typename TG,
|
||||
typename MT,
|
||||
typename LrAccessor,
|
||||
typename BiasCorrAccessor>
|
||||
__global__ void AdamStyleKernel(const double beta1,
|
||||
const double beta2,
|
||||
const double epsilon,
|
||||
LrAccessor lr_accessor,
|
||||
BiasCorrAccessor bias_corr_accessor,
|
||||
const TG* __restrict__ grad,
|
||||
const T* __restrict__ param,
|
||||
T* __restrict__ param_out,
|
||||
const MT* __restrict__ master_param,
|
||||
MT* __restrict__ master_param_out,
|
||||
const MT* __restrict__ moment1,
|
||||
MT* __restrict__ moment1_out,
|
||||
const MT* __restrict__ moment2,
|
||||
MT* __restrict__ moment2_out,
|
||||
const MT* __restrict__ moment2_max,
|
||||
MT* __restrict__ moment2_max_out,
|
||||
int64_t ndim,
|
||||
bool amsgrad) {
|
||||
__shared__ double one_minus_beta1_shared;
|
||||
__shared__ double one_minus_beta2_shared;
|
||||
__shared__ MT bias_correction2_sqrt_shared;
|
||||
__shared__ MT step_size_shared;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
const double lr_double = lr_accessor.GetLrDouble();
|
||||
const double bc1_dbl = bias_corr_accessor.GetBc1();
|
||||
const double bc2_dbl = bias_corr_accessor.GetBc2();
|
||||
const double bc2_sqrt_dbl = ::sqrt(bc2_dbl);
|
||||
|
||||
one_minus_beta1_shared = 1.0 - beta1;
|
||||
one_minus_beta2_shared = 1.0 - beta2;
|
||||
const MT bias_correction1 = static_cast<MT>(bc1_dbl);
|
||||
bias_correction2_sqrt_shared = static_cast<MT>(bc2_sqrt_dbl);
|
||||
step_size_shared = lr_double / bias_correction1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const double one_minus_beta1 = one_minus_beta1_shared;
|
||||
const double one_minus_beta2 = one_minus_beta2_shared;
|
||||
const MT bias_correction2_sqrt = bias_correction2_sqrt_shared;
|
||||
const MT step_size = step_size_shared;
|
||||
|
||||
int64_t id =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
|
||||
for (; id < ndim; id += static_cast<int64_t>(gridDim.x) *
|
||||
static_cast<int64_t>(blockDim.x)) {
|
||||
MT p = master_param ? master_param[id] : static_cast<MT>(param[id]);
|
||||
MT g = static_cast<MT>(grad[id]);
|
||||
MT exp_avg = static_cast<MT>(moment1[id]);
|
||||
MT exp_avg_sq = static_cast<MT>(moment2[id]);
|
||||
const double g_d = static_cast<double>(g);
|
||||
|
||||
// exp_avg = beta1 * exp_avg + (1 - beta1) * grad
|
||||
// FMA variant A: compute (1-beta1)*grad first, then fma(beta1, exp_avg,
|
||||
// t2). This matches NVCC's fused adam: beta1*exp_avg computed exactly in
|
||||
// FMA.
|
||||
{
|
||||
const double exp_avg_d = static_cast<double>(exp_avg);
|
||||
const double t2 = __dmul_rn(one_minus_beta1, g_d);
|
||||
exp_avg = static_cast<MT>(__fma_rn(beta1, exp_avg_d, t2));
|
||||
}
|
||||
|
||||
// exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * grad^2
|
||||
// Left-to-right: ((1-beta2)*g)*g, then separate add, matching PyTorch.
|
||||
{
|
||||
const double exp_avg_sq_d = static_cast<double>(exp_avg_sq);
|
||||
const double t1 = __dmul_rn(beta2, exp_avg_sq_d);
|
||||
const double t2 = __dmul_rn(one_minus_beta2, g_d);
|
||||
const double t3 = __dmul_rn(t2, g_d);
|
||||
exp_avg_sq = static_cast<MT>(__dadd_rn(t1, t3));
|
||||
}
|
||||
|
||||
MT denom;
|
||||
if (amsgrad) {
|
||||
MT max_exp_avg_sq = static_cast<MT>(moment2_max[id]);
|
||||
max_exp_avg_sq =
|
||||
max_exp_avg_sq > exp_avg_sq ? max_exp_avg_sq : exp_avg_sq;
|
||||
moment2_max_out[id] = max_exp_avg_sq;
|
||||
denom = (sqrt(max_exp_avg_sq) / bias_correction2_sqrt) + epsilon;
|
||||
} else {
|
||||
denom = (sqrt(exp_avg_sq) / bias_correction2_sqrt) + epsilon;
|
||||
}
|
||||
|
||||
p -= step_size * exp_avg / denom;
|
||||
|
||||
moment1_out[id] = exp_avg;
|
||||
moment2_out[id] = exp_avg_sq;
|
||||
param_out[id] = static_cast<T>(p);
|
||||
if (master_param_out) {
|
||||
master_param_out[id] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AdamDenseKernel_compatible(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment1,
|
||||
const DenseTensor& moment2,
|
||||
const optional<DenseTensor>& moment2_max,
|
||||
const DenseTensor& beta1_pow,
|
||||
const DenseTensor& beta2_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
const optional<DenseTensor>& skip_update,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
bool lazy_mode,
|
||||
int64_t min_row_size_to_use_multithread,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment1_out,
|
||||
DenseTensor* moment2_out,
|
||||
DenseTensor* moment2_max_out,
|
||||
DenseTensor* beta1_pow_out,
|
||||
DenseTensor* beta2_pow_out,
|
||||
DenseTensor* master_param_outs) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
|
||||
bool skip_update_ = false;
|
||||
if (skip_update.is_initialized()) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
skip_update->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
||||
skip_update->numel()));
|
||||
std::vector<bool> skip_update_vec;
|
||||
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
||||
skip_update_ = skip_update_vec[0];
|
||||
}
|
||||
|
||||
if (skip_update_) {
|
||||
VLOG(4) << "Adam skip update";
|
||||
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
||||
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
||||
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
||||
if (amsgrad) {
|
||||
Copy(dev_ctx,
|
||||
moment2_max.get(),
|
||||
dev_ctx.GetPlace(),
|
||||
false,
|
||||
moment2_max_out);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
||||
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
double beta1_ = beta1.to<double>();
|
||||
double beta2_ = beta2.to<double>();
|
||||
double epsilon_ = epsilon.to<double>();
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta1_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta1 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta1_pow_out->numel()));
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta2_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta2 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta2_pow_out->numel()));
|
||||
|
||||
const bool beta_pow_on_cpu =
|
||||
beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace();
|
||||
const bool lr_on_cpu = learning_rate.place() == CPUPlace();
|
||||
|
||||
// Read float lr on host if available; on device it's read in kernel.
|
||||
double lr_host_double = 0.0;
|
||||
if (lr_on_cpu) {
|
||||
lr_host_double = learning_rate.data<double>()[0];
|
||||
}
|
||||
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
||||
|
||||
const MT* moment2_max_in_data =
|
||||
amsgrad ? moment2_max.get().data<MT>() : nullptr;
|
||||
MT* moment2_max_out_data =
|
||||
amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out) : nullptr;
|
||||
|
||||
int threads = 512;
|
||||
int64_t blocks_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int blocks = std::min((param.numel() + threads - 1) / threads, blocks_max);
|
||||
|
||||
const bool use_float32_grad = grad.dtype() == DataType::FLOAT32;
|
||||
|
||||
// Use decltype so template args are inferred from the local accessor variables.
|
||||
#define LAUNCH_ADAM_STYLE_KERNEL() \
|
||||
if (use_float32_grad) { \
|
||||
AdamStyleKernel<T, \
|
||||
float, \
|
||||
MT, \
|
||||
decltype(lr_accessor), \
|
||||
decltype(bias_corr_accessor)> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
lr_accessor, \
|
||||
bias_corr_accessor, \
|
||||
grad.data<float>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MT>(), \
|
||||
dev_ctx.template Alloc<MT>(moment1_out), \
|
||||
moment2.data<MT>(), \
|
||||
dev_ctx.template Alloc<MT>(moment2_out), \
|
||||
moment2_max_in_data, \
|
||||
moment2_max_out_data, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} else { \
|
||||
AdamStyleKernel<T, \
|
||||
T, \
|
||||
MT, \
|
||||
decltype(lr_accessor), \
|
||||
decltype(bias_corr_accessor)> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
lr_accessor, \
|
||||
bias_corr_accessor, \
|
||||
grad.data<T>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MT>(), \
|
||||
dev_ctx.template Alloc<MT>(moment1_out), \
|
||||
moment2.data<MT>(), \
|
||||
dev_ctx.template Alloc<MT>(moment2_out), \
|
||||
moment2_max_in_data, \
|
||||
moment2_max_out_data, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
}
|
||||
|
||||
if (lr_on_cpu) {
|
||||
AdamLrAccessor<MT, true> lr_accessor(lr_host_double);
|
||||
if (beta_pow_on_cpu) {
|
||||
const float sc = static_cast<float>(
|
||||
std::round(std::log(static_cast<double>(beta1_pow.data<MT>()[0])) /
|
||||
std::log(beta1_)));
|
||||
AdamBiasCorrAccessorCompat<MT, true> bias_corr_accessor(
|
||||
beta1_, beta2_, sc);
|
||||
LAUNCH_ADAM_STYLE_KERNEL()
|
||||
} else {
|
||||
AdamBiasCorrAccessorCompat<MT, false> bias_corr_accessor(
|
||||
beta1_, beta2_, beta1_pow.data<MT>());
|
||||
LAUNCH_ADAM_STYLE_KERNEL()
|
||||
}
|
||||
} else {
|
||||
AdamLrAccessor<MT, false> lr_accessor(learning_rate.data<double>());
|
||||
if (beta_pow_on_cpu) {
|
||||
const float sc = static_cast<float>(
|
||||
std::round(std::log(static_cast<double>(beta1_pow.data<MT>()[0])) /
|
||||
std::log(beta1_)));
|
||||
AdamBiasCorrAccessorCompat<MT, true> bias_corr_accessor(
|
||||
beta1_, beta2_, sc);
|
||||
LAUNCH_ADAM_STYLE_KERNEL()
|
||||
} else {
|
||||
AdamBiasCorrAccessorCompat<MT, false> bias_corr_accessor(
|
||||
beta1_, beta2_, beta1_pow.data<MT>());
|
||||
LAUNCH_ADAM_STYLE_KERNEL()
|
||||
}
|
||||
}
|
||||
#undef LAUNCH_ADAM_STYLE_KERNEL
|
||||
|
||||
if (!use_global_beta_pow) {
|
||||
if (beta_pow_on_cpu) {
|
||||
dev_ctx.template HostAlloc<MT>(beta1_pow_out)[0] =
|
||||
static_cast<MT>(beta1_) * beta1_pow.data<MT>()[0];
|
||||
dev_ctx.template HostAlloc<MT>(beta2_pow_out)[0] =
|
||||
static_cast<MT>(beta2_) * beta2_pow.data<MT>()[0];
|
||||
} else {
|
||||
UpdateBetaPow<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
static_cast<MT>(beta1_),
|
||||
static_cast<MT>(beta2_),
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(beta1_pow_out),
|
||||
dev_ctx.template Alloc<MT>(beta2_pow_out));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
PADDLE_API void AdamDenseKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment1,
|
||||
const DenseTensor& moment2,
|
||||
const optional<DenseTensor>& moment2_max,
|
||||
const DenseTensor& beta1_pow,
|
||||
const DenseTensor& beta2_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
const optional<DenseTensor>& skip_update,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
bool lazy_mode,
|
||||
int64_t min_row_size_to_use_multithread,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment1_out,
|
||||
DenseTensor* moment2_out,
|
||||
DenseTensor* moment2_max_out,
|
||||
DenseTensor* beta1_pow_out,
|
||||
DenseTensor* beta2_pow_out,
|
||||
DenseTensor* master_param_outs) {
|
||||
if (FLAGS_use_accuracy_compatible_kernel) {
|
||||
AdamDenseKernel_compatible<T, Context>(dev_ctx,
|
||||
param,
|
||||
grad,
|
||||
learning_rate,
|
||||
moment1,
|
||||
moment2,
|
||||
moment2_max,
|
||||
beta1_pow,
|
||||
beta2_pow,
|
||||
master_param,
|
||||
skip_update,
|
||||
beta1,
|
||||
beta2,
|
||||
epsilon,
|
||||
lazy_mode,
|
||||
min_row_size_to_use_multithread,
|
||||
multi_precision,
|
||||
use_global_beta_pow,
|
||||
amsgrad,
|
||||
param_out,
|
||||
moment1_out,
|
||||
moment2_out,
|
||||
moment2_max_out,
|
||||
beta1_pow_out,
|
||||
beta2_pow_out,
|
||||
master_param_outs);
|
||||
return;
|
||||
}
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
const auto grad_type = grad.dtype();
|
||||
|
||||
VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
|
||||
VLOG(4) << "amsgrad: " << amsgrad;
|
||||
|
||||
bool skip_update_ = false;
|
||||
if (skip_update.is_initialized()) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
skip_update->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
||||
skip_update->numel()));
|
||||
std::vector<bool> skip_update_vec;
|
||||
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
||||
skip_update_ = skip_update_vec[0];
|
||||
}
|
||||
// skip_update=true, just copy input to output, and TensorCopy will call
|
||||
// mutable_data
|
||||
if (skip_update_) {
|
||||
VLOG(4) << "Adam skip update";
|
||||
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
||||
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
||||
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
||||
if (amsgrad) {
|
||||
Copy(dev_ctx,
|
||||
moment2_max.get(),
|
||||
dev_ctx.GetPlace(),
|
||||
false,
|
||||
moment2_max_out);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
||||
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
MT beta1_ = beta1.to<MT>();
|
||||
MT beta2_ = beta2.to<MT>();
|
||||
MT epsilon_ = epsilon.to<MT>();
|
||||
VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel()
|
||||
<< "beta2_pow.numel() : " << beta2_pow.numel();
|
||||
VLOG(3) << "param.numel(): " << param.numel();
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta1_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta1 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta1_pow_out->numel()));
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta2_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta2 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta2_pow_out->numel()));
|
||||
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
||||
|
||||
const MT* moment2_max_in_data =
|
||||
amsgrad ? moment2_max.get().data<MT>() : nullptr;
|
||||
MT* moment2_max_out_data =
|
||||
amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out) : nullptr;
|
||||
|
||||
// update param and moment
|
||||
int threads = 512;
|
||||
int64_t blocks_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int blocks = std::min((param.numel() + threads - 1) / threads, blocks_max);
|
||||
|
||||
if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
|
||||
// Compute with betapow in REG
|
||||
if (grad_type == DataType::FLOAT32) {
|
||||
AdamKernelREG<T, float, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
*beta1_pow.data<MT>(),
|
||||
*beta2_pow.data<MT>(),
|
||||
moment1.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out),
|
||||
moment2.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate.data<double>(),
|
||||
grad.data<float>(),
|
||||
param.data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param.numel(),
|
||||
amsgrad);
|
||||
} else {
|
||||
AdamKernelREG<T, T, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
*beta1_pow.data<MT>(),
|
||||
*beta2_pow.data<MT>(),
|
||||
moment1.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out),
|
||||
moment2.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate.data<double>(),
|
||||
grad.data<T>(),
|
||||
param.data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param.numel(),
|
||||
amsgrad);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
// Cpu update
|
||||
dev_ctx.template HostAlloc<MT>(beta1_pow_out)[0] =
|
||||
beta1_ * beta1_pow.data<MT>()[0];
|
||||
dev_ctx.template HostAlloc<MT>(beta2_pow_out)[0] =
|
||||
beta2_ * beta2_pow.data<MT>()[0];
|
||||
}
|
||||
} else {
|
||||
if (grad_type == DataType::FLOAT32) {
|
||||
AdamKernelMEM<T, float, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
moment1.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out),
|
||||
moment2.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate.data<double>(),
|
||||
grad.data<float>(),
|
||||
param.data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param.numel(),
|
||||
amsgrad);
|
||||
} else {
|
||||
AdamKernelMEM<T, T, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
moment1.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out),
|
||||
moment2.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate.data<double>(),
|
||||
grad.data<T>(),
|
||||
param.data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param.numel(),
|
||||
amsgrad);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
// Update with gpu
|
||||
UpdateBetaPow<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(beta1_pow_out),
|
||||
dev_ctx.template Alloc<MT>(beta2_pow_out));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void MergedAdamKernel(
|
||||
const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& param,
|
||||
const std::vector<const DenseTensor*>& grad,
|
||||
const std::vector<const DenseTensor*>& learning_rate,
|
||||
const std::vector<const DenseTensor*>& moment1,
|
||||
const std::vector<const DenseTensor*>& moment2,
|
||||
const optional<std::vector<const DenseTensor*>>& moment2_max,
|
||||
const std::vector<const DenseTensor*>& beta1_pow,
|
||||
const std::vector<const DenseTensor*>& beta2_pow,
|
||||
const optional<std::vector<const DenseTensor*>>& master_param,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
std::vector<DenseTensor*> param_out,
|
||||
std::vector<DenseTensor*> moment1_out,
|
||||
std::vector<DenseTensor*> moment2_out,
|
||||
std::vector<DenseTensor*> moment2_max_out,
|
||||
std::vector<DenseTensor*> beta1_pow_out,
|
||||
std::vector<DenseTensor*> beta2_pow_out,
|
||||
std::vector<DenseTensor*> master_param_out) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
|
||||
MT beta1_ = beta1.to<MT>();
|
||||
MT beta2_ = beta2.to<MT>();
|
||||
MT epsilon_ = epsilon.to<MT>();
|
||||
|
||||
size_t param_num = param.size();
|
||||
|
||||
for (size_t idx = 0; idx < param_num; idx++) {
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param.get()[idx]->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_out[idx])
|
||||
: nullptr;
|
||||
|
||||
const MT* moment2_max_in_data =
|
||||
amsgrad ? moment2_max.get()[idx]->data<MT>() : nullptr;
|
||||
MT* moment2_max_out_data =
|
||||
amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out[idx]) : nullptr;
|
||||
|
||||
// update param and moment
|
||||
int threads = 512;
|
||||
int64_t blocks_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int blocks =
|
||||
std::min((param[idx]->numel() + threads - 1) / threads, blocks_max);
|
||||
|
||||
const auto grad_type = grad[idx]->dtype();
|
||||
if (beta1_pow[idx]->place() == CPUPlace() &&
|
||||
beta2_pow[idx]->place() == CPUPlace()) {
|
||||
// Compute with betapow in REG
|
||||
if (grad_type == DataType::FLOAT32) {
|
||||
AdamKernelREG<T, float, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
*beta1_pow[idx]->data<MT>(),
|
||||
*beta2_pow[idx]->data<MT>(),
|
||||
moment1[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out[idx]),
|
||||
moment2[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out[idx]),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate[idx]->data<double>(),
|
||||
grad[idx]->data<float>(),
|
||||
param[idx]->data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out[idx]),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param[idx]->numel(),
|
||||
amsgrad);
|
||||
} else {
|
||||
AdamKernelREG<T, T, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
*beta1_pow[idx]->data<MT>(),
|
||||
*beta2_pow[idx]->data<MT>(),
|
||||
moment1[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out[idx]),
|
||||
moment2[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out[idx]),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate[idx]->data<double>(),
|
||||
grad[idx]->data<T>(),
|
||||
param[idx]->data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out[idx]),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param[idx]->numel(),
|
||||
amsgrad);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
// Cpu update
|
||||
dev_ctx.template HostAlloc<MT>(beta1_pow_out[idx])[0] =
|
||||
beta1_ * beta1_pow[idx]->data<MT>()[0];
|
||||
dev_ctx.template HostAlloc<MT>(beta2_pow_out[idx])[0] =
|
||||
beta2_ * beta2_pow[idx]->data<MT>()[0];
|
||||
}
|
||||
} else {
|
||||
if (grad_type == DataType::FLOAT32) {
|
||||
AdamKernelMEM<T, float, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
beta1_pow[idx]->data<MT>(),
|
||||
beta2_pow[idx]->data<MT>(),
|
||||
moment1[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out[idx]),
|
||||
moment2[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out[idx]),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate[idx]->data<double>(),
|
||||
grad[idx]->data<float>(),
|
||||
param[idx]->data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out[idx]),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param[idx]->numel(),
|
||||
amsgrad);
|
||||
} else {
|
||||
AdamKernelMEM<T, T, MT><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
beta1_pow[idx]->data<MT>(),
|
||||
beta2_pow[idx]->data<MT>(),
|
||||
moment1[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment1_out[idx]),
|
||||
moment2[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(moment2_out[idx]),
|
||||
moment2_max_in_data,
|
||||
moment2_max_out_data,
|
||||
learning_rate[idx]->data<double>(),
|
||||
grad[idx]->data<T>(),
|
||||
param[idx]->data<T>(),
|
||||
dev_ctx.template Alloc<T>(param_out[idx]),
|
||||
master_in_data,
|
||||
master_out_data,
|
||||
param[idx]->numel(),
|
||||
amsgrad);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
// Update with gpu
|
||||
UpdateBetaPow<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
beta1_pow[idx]->data<MT>(),
|
||||
beta2_pow[idx]->data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(beta1_pow_out[idx]),
|
||||
dev_ctx.template Alloc<MT>(beta2_pow_out[idx]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(adam,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AdamDenseKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT64);
|
||||
// Skip beta1_pow, beta2_pow, skip_update data transform
|
||||
kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
||||
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
|
||||
kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(merged_adam,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::MergedAdamKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT64);
|
||||
// Skip beta1_pow, beta2_pow data transform
|
||||
kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
||||
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
|
||||
kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
// Copyright (c) 2022 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/adamax_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
namespace phi {
|
||||
template <typename T, typename MT>
|
||||
__global__ void AdamaxGPUKernel(const T* param,
|
||||
const T* grad,
|
||||
const MT* learning_rate,
|
||||
const MT* moment,
|
||||
const MT* inf_norm,
|
||||
const MT* beta1_pow,
|
||||
const MT* master_param,
|
||||
MT d_beta1,
|
||||
MT d_beta2,
|
||||
MT d_epsilon,
|
||||
int64_t num,
|
||||
T* param_out,
|
||||
MT* moment_out,
|
||||
MT* inf_norm_out,
|
||||
MT* master_param_out) {
|
||||
int64_t idx =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
|
||||
MT lr = static_cast<MT>(learning_rate[0]);
|
||||
MT d_pow = static_cast<MT>(beta1_pow[0]);
|
||||
MT one = static_cast<MT>(1.0f);
|
||||
auto l_r = lr / (one - d_pow);
|
||||
|
||||
for (int64_t index = idx; index < num; index += gridDim.x * blockDim.x) {
|
||||
// load and cast input to MT
|
||||
MT d_param =
|
||||
master_param ? master_param[index] : static_cast<MT>(param[index]);
|
||||
MT d_grad = static_cast<MT>(grad[index]);
|
||||
MT d_moment = static_cast<MT>(moment[index]);
|
||||
MT d_inf = static_cast<MT>(inf_norm[index]);
|
||||
// compute
|
||||
auto mom_out = d_beta1 * d_moment + (one - d_beta1) * d_grad;
|
||||
auto norm_out = std::max(std::abs(d_grad), d_beta2 * d_inf + d_epsilon);
|
||||
auto out_data = d_param - l_r * (mom_out / norm_out);
|
||||
// store
|
||||
param_out[index] = static_cast<T>(out_data);
|
||||
moment_out[index] = static_cast<T>(mom_out);
|
||||
inf_norm_out[index] = static_cast<T>(norm_out);
|
||||
|
||||
if (master_param_out) {
|
||||
master_param_out[index] = out_data;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AdamaxKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment,
|
||||
const DenseTensor& inf_norm,
|
||||
const DenseTensor& beta1_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
float beta1,
|
||||
float beta2,
|
||||
float epsilon,
|
||||
bool multi_precision,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment_out,
|
||||
DenseTensor* inf_norm_out,
|
||||
DenseTensor* master_param_outs) {
|
||||
using MT = typename dtype::template MPTypeTrait<T>::Type;
|
||||
T* param_out_data = dev_ctx.template Alloc<T>(param_out);
|
||||
MT* moment_out_data = dev_ctx.template Alloc<MT>(moment_out);
|
||||
MT* inf_norm_out_data = dev_ctx.template Alloc<MT>(inf_norm_out);
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta1_pow.numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta1 pow's size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta1_pow.numel()));
|
||||
|
||||
MT beta1_ = static_cast<MT>(beta1);
|
||||
MT beta2_ = static_cast<MT>(beta2);
|
||||
MT epsilon_ = static_cast<MT>(epsilon);
|
||||
|
||||
int64_t numel = param.numel();
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 1);
|
||||
int grid = config.block_per_grid.x;
|
||||
int block = config.thread_per_block.x;
|
||||
auto stream = dev_ctx.stream();
|
||||
|
||||
AdamaxGPUKernel<T, MT><<<block, grid, 0, stream>>>(param.data<T>(),
|
||||
grad.data<T>(),
|
||||
learning_rate.data<MT>(),
|
||||
moment.data<MT>(),
|
||||
inf_norm.data<MT>(),
|
||||
beta1_pow.data<MT>(),
|
||||
master_in_data,
|
||||
beta1_,
|
||||
beta2_,
|
||||
epsilon_,
|
||||
numel,
|
||||
param_out_data,
|
||||
moment_out_data,
|
||||
inf_norm_out_data,
|
||||
master_out_data);
|
||||
}
|
||||
} // namespace phi
|
||||
PD_REGISTER_KERNEL(
|
||||
adamax, GPU, ALL_LAYOUT, phi::AdamaxKernel, float, double, phi::float16) {
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,911 @@
|
||||
// Copyright (c) 2022 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/adamw_kernel.h"
|
||||
|
||||
#include <math.h> // for sqrt in CPU and CUDA
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/common/flags.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/adam_functors.h"
|
||||
#include "paddle/phi/kernels/funcs/for_range.h"
|
||||
#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
|
||||
|
||||
COMMON_DECLARE_bool(use_accuracy_compatible_kernel);
|
||||
|
||||
namespace phi {
|
||||
|
||||
// Template accessor design
|
||||
template <typename MT, bool IsCpu>
|
||||
struct BetaPowAccessor;
|
||||
|
||||
template <typename MT>
|
||||
struct BetaPowAccessor<MT, true> { // CPU accessor
|
||||
const MT beta1;
|
||||
const MT beta2;
|
||||
|
||||
BetaPowAccessor(const MT* beta1_pow, const MT* beta2_pow)
|
||||
: beta1(*beta1_pow), beta2(*beta2_pow) {}
|
||||
|
||||
__device__ MT GetBeta1() const { return beta1; }
|
||||
__device__ MT GetBeta2() const { return beta2; }
|
||||
};
|
||||
|
||||
template <typename MT>
|
||||
struct BetaPowAccessor<MT, false> { // GPU pointer
|
||||
const MT* beta1_pow;
|
||||
const MT* beta2_pow;
|
||||
|
||||
BetaPowAccessor(const MT* beta1, const MT* beta2)
|
||||
: beta1_pow(beta1), beta2_pow(beta2) {}
|
||||
|
||||
__device__ MT GetBeta1() const { return *beta1_pow; }
|
||||
__device__ MT GetBeta2() const { return *beta2_pow; }
|
||||
};
|
||||
|
||||
// Unified kernel template
|
||||
template <typename T, // Parameter type
|
||||
typename TG, // Gradient type
|
||||
typename MT, // Multi-precision type
|
||||
typename TM, // Moment estimation type
|
||||
typename BetaAccessor>
|
||||
__global__ void AdamWKernel(MT beta1,
|
||||
MT beta2,
|
||||
MT epsilon,
|
||||
MT coeff,
|
||||
MT lr_ratio,
|
||||
const double* lr_,
|
||||
const TG* grad,
|
||||
const T* param,
|
||||
T* param_out,
|
||||
const MT* master_param,
|
||||
MT* master_param_out,
|
||||
const TM* moment1,
|
||||
TM* moment1_out,
|
||||
const TM* moment2,
|
||||
TM* moment2_out,
|
||||
const TM* moment2_max,
|
||||
TM* moment2_max_out,
|
||||
BetaAccessor beta_accessor,
|
||||
int64_t ndim,
|
||||
bool amsgrad) {
|
||||
int64_t id =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
MT lr = *lr_ * lr_ratio;
|
||||
// Get beta powers
|
||||
MT beta1_pow = beta_accessor.GetBeta1();
|
||||
MT beta2_pow = beta_accessor.GetBeta2();
|
||||
|
||||
for (; id < ndim; id += gridDim.x * blockDim.x) {
|
||||
MT p = master_param ? master_param[id] : static_cast<MT>(param[id]);
|
||||
MT g = static_cast<MT>(grad[id]);
|
||||
MT mom1 = static_cast<MT>(moment1[id]);
|
||||
MT mom2 = static_cast<MT>(moment2[id]);
|
||||
|
||||
p *= (static_cast<MT>(1.0) - lr * coeff);
|
||||
|
||||
mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
|
||||
mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
|
||||
|
||||
MT denom;
|
||||
if (amsgrad) {
|
||||
MT mom2_max = static_cast<MT>(moment2_max[id]);
|
||||
MT mom2_max_ = std::max(mom2, mom2_max);
|
||||
moment2_max_out[id] = mom2_max_;
|
||||
|
||||
denom =
|
||||
(sqrt(mom2_max_) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
} else {
|
||||
denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
|
||||
}
|
||||
|
||||
p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
|
||||
|
||||
moment1_out[id] = mom1;
|
||||
moment2_out[id] = mom2;
|
||||
param_out[id] = static_cast<T>(p);
|
||||
if (master_param_out) {
|
||||
master_param_out[id] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Beta power update kernel
|
||||
template <typename MT>
|
||||
__global__ void UpdateBetaPowKernel(MT beta1,
|
||||
MT beta2,
|
||||
const MT* beta1_pow,
|
||||
const MT* beta2_pow,
|
||||
MT* beta1_pow_out,
|
||||
MT* beta2_pow_out) {
|
||||
beta1_pow_out[0] = beta1 * beta1_pow[0];
|
||||
beta2_pow_out[0] = beta2 * beta2_pow[0];
|
||||
}
|
||||
|
||||
// Forward declaration
|
||||
template <typename T, typename Context>
|
||||
PADDLE_API void AdamwDenseKernel_compatible(
|
||||
const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment1,
|
||||
const DenseTensor& moment2,
|
||||
const optional<DenseTensor>& moment2_max,
|
||||
const DenseTensor& beta1_pow,
|
||||
const DenseTensor& beta2_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
const optional<DenseTensor>& skip_update,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
double lr_ratio,
|
||||
double coeff,
|
||||
bool with_decay,
|
||||
bool lazy_mode,
|
||||
int64_t min_row_size_to_use_multithread,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment1_out,
|
||||
DenseTensor* moment2_out,
|
||||
DenseTensor* moment2_max_out,
|
||||
DenseTensor* beta1_pow_out,
|
||||
DenseTensor* beta2_pow_out,
|
||||
DenseTensor* master_param_outs);
|
||||
|
||||
template <typename T, typename Context>
|
||||
PADDLE_API void AdamwDenseKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment1,
|
||||
const DenseTensor& moment2,
|
||||
const optional<DenseTensor>& moment2_max,
|
||||
const DenseTensor& beta1_pow,
|
||||
const DenseTensor& beta2_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
const optional<DenseTensor>& skip_update,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
double lr_ratio,
|
||||
double coeff,
|
||||
bool with_decay,
|
||||
bool lazy_mode,
|
||||
int64_t min_row_size_to_use_multithread,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment1_out,
|
||||
DenseTensor* moment2_out,
|
||||
DenseTensor* moment2_max_out,
|
||||
DenseTensor* beta1_pow_out,
|
||||
DenseTensor* beta2_pow_out,
|
||||
DenseTensor* master_param_outs) {
|
||||
if (FLAGS_use_accuracy_compatible_kernel) {
|
||||
AdamwDenseKernel_compatible<T, Context>(dev_ctx,
|
||||
param,
|
||||
grad,
|
||||
learning_rate,
|
||||
moment1,
|
||||
moment2,
|
||||
moment2_max,
|
||||
beta1_pow,
|
||||
beta2_pow,
|
||||
master_param,
|
||||
skip_update,
|
||||
beta1,
|
||||
beta2,
|
||||
epsilon,
|
||||
lr_ratio,
|
||||
coeff,
|
||||
with_decay,
|
||||
lazy_mode,
|
||||
min_row_size_to_use_multithread,
|
||||
multi_precision,
|
||||
use_global_beta_pow,
|
||||
amsgrad,
|
||||
param_out,
|
||||
moment1_out,
|
||||
moment2_out,
|
||||
moment2_max_out,
|
||||
beta1_pow_out,
|
||||
beta2_pow_out,
|
||||
master_param_outs);
|
||||
return;
|
||||
}
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
MT coeff_ = static_cast<MT>(coeff);
|
||||
MT lr_ratio_ = static_cast<MT>(lr_ratio);
|
||||
|
||||
bool skip_update_ = false;
|
||||
if (skip_update.is_initialized()) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
skip_update->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
||||
skip_update->numel()));
|
||||
std::vector<bool> skip_update_vec;
|
||||
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
||||
skip_update_ = skip_update_vec[0];
|
||||
}
|
||||
|
||||
// skip_update=true, just copy input to output, and TensorCopy will call
|
||||
// mutable_data
|
||||
if (skip_update_) {
|
||||
VLOG(4) << "Adamw skip update";
|
||||
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
||||
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
||||
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
||||
if (amsgrad) {
|
||||
Copy(dev_ctx,
|
||||
moment2_max.get(),
|
||||
dev_ctx.GetPlace(),
|
||||
false,
|
||||
moment2_max_out);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
||||
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// if with_decay = false, coeff = 0
|
||||
if (!with_decay) {
|
||||
coeff_ = static_cast<MT>(0.0);
|
||||
}
|
||||
|
||||
MT beta1_ = beta1.to<MT>();
|
||||
MT beta2_ = beta2.to<MT>();
|
||||
MT epsilon_ = epsilon.to<MT>();
|
||||
VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel()
|
||||
<< "beta2_pow.numel() : " << beta2_pow.numel();
|
||||
VLOG(3) << "param.numel(): " << param.numel();
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta1_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta1 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta1_pow_out->numel()));
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta2_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta2 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta2_pow_out->numel()));
|
||||
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
||||
|
||||
const MT* moment2_max_in_data =
|
||||
amsgrad ? moment2_max.get().data<MT>() : nullptr;
|
||||
MT* moment2_max_out_data =
|
||||
amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out) : nullptr;
|
||||
|
||||
// update param and moment
|
||||
int threads = 512;
|
||||
int64_t blocks_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int blocks = std::min((param.numel() + threads - 1) / threads, blocks_max);
|
||||
|
||||
// Determine BetaPow location
|
||||
const bool beta_pow_on_cpu =
|
||||
beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace();
|
||||
|
||||
// Determine gradient type
|
||||
const bool use_bfloat32_grad = grad.dtype() == DataType::FLOAT32;
|
||||
// Determine moment type
|
||||
const bool use_bfloat16_moments = moment1.dtype() == DataType::BFLOAT16 &&
|
||||
moment2.dtype() == DataType::BFLOAT16;
|
||||
|
||||
#define LAUNCH_ADAMW_KERNEL(MOMENT_T) \
|
||||
if (beta_pow_on_cpu) { \
|
||||
BetaPowAccessor<MT, true> accessor(beta1_pow.data<MT>(), \
|
||||
beta2_pow.data<MT>()); \
|
||||
if (use_bfloat32_grad) { \
|
||||
AdamWKernel<T, float, MT, MOMENT_T, BetaPowAccessor<MT, true>> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
coeff_, \
|
||||
lr_ratio_, \
|
||||
learning_rate.data<double>(), \
|
||||
grad.data<float>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
moment2_max ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
accessor, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} else { \
|
||||
AdamWKernel<T, T, MT, MOMENT_T, BetaPowAccessor<MT, true>> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
coeff_, \
|
||||
lr_ratio_, \
|
||||
learning_rate.data<double>(), \
|
||||
grad.data<T>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
moment2_max ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
accessor, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} \
|
||||
} else { \
|
||||
BetaPowAccessor<MT, false> accessor(beta1_pow.data<MT>(), \
|
||||
beta2_pow.data<MT>()); \
|
||||
if (use_bfloat32_grad) { \
|
||||
AdamWKernel<T, float, MT, MOMENT_T, BetaPowAccessor<MT, false>> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
coeff_, \
|
||||
lr_ratio_, \
|
||||
learning_rate.data<double>(), \
|
||||
grad.data<float>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
moment2_max ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
accessor, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} else { \
|
||||
AdamWKernel<T, T, MT, MOMENT_T, BetaPowAccessor<MT, false>> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
coeff_, \
|
||||
lr_ratio_, \
|
||||
learning_rate.data<double>(), \
|
||||
grad.data<T>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
moment2_max ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
accessor, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} \
|
||||
}
|
||||
|
||||
// Select template instantiation based on moment type
|
||||
if (use_bfloat16_moments) {
|
||||
LAUNCH_ADAMW_KERNEL(bfloat16)
|
||||
} else {
|
||||
LAUNCH_ADAMW_KERNEL(MT)
|
||||
}
|
||||
#undef LAUNCH_ADAMW_KERNEL
|
||||
|
||||
// Update beta_pow
|
||||
if (!use_global_beta_pow) {
|
||||
if (beta_pow_on_cpu) {
|
||||
auto* beta1_pow_out_data = dev_ctx.template HostAlloc<MT>(beta1_pow_out);
|
||||
auto* beta2_pow_out_data = dev_ctx.template HostAlloc<MT>(beta2_pow_out);
|
||||
beta1_pow_out_data[0] = beta1_ * beta1_pow.data<MT>()[0];
|
||||
beta2_pow_out_data[0] = beta2_ * beta2_pow.data<MT>()[0];
|
||||
} else {
|
||||
UpdateBetaPowKernel<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(beta1_pow_out),
|
||||
dev_ctx.template Alloc<MT>(beta2_pow_out));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
template <bool IsCpu>
|
||||
struct AdamWLrAccessor;
|
||||
|
||||
// cpu
|
||||
template <>
|
||||
struct AdamWLrAccessor<true> {
|
||||
const double lr_double;
|
||||
|
||||
explicit AdamWLrAccessor(double lr) : lr_double(lr) {}
|
||||
|
||||
__device__ __forceinline__ double GetLrDouble() const { return lr_double; }
|
||||
};
|
||||
|
||||
// gpu
|
||||
template <>
|
||||
struct AdamWLrAccessor<false> {
|
||||
const double* lr;
|
||||
const double lr_ratio;
|
||||
|
||||
AdamWLrAccessor(const double* lr, double lr_ratio)
|
||||
: lr(lr), lr_ratio(lr_ratio) {}
|
||||
|
||||
__device__ __forceinline__ double GetLrDouble() const {
|
||||
return *lr * lr_ratio;
|
||||
}
|
||||
};
|
||||
|
||||
// Device-side pow matching torch's at::native::pow_ (promotes float exp to
|
||||
// double, then calls ::pow(double, double))
|
||||
template <typename Base_type, typename Exp_type>
|
||||
static __device__ __forceinline__ Base_type torch_pow_(Base_type base,
|
||||
Exp_type exp) {
|
||||
return ::pow(base, exp);
|
||||
}
|
||||
|
||||
// Accuracy-compatible bias correction: computes 1-beta^step_count on device,
|
||||
// matching torch's FusedAdamMathFunctor. After torch's "Use opmath_t and not
|
||||
// double compute" change, the bias correction is computed in opmath_t (float
|
||||
// for fp32/fp16/bf16, double for fp64), with beta and step_count both cast to
|
||||
// opmath_t before pow_. We therefore compute everything in MT (= opmath_t).
|
||||
template <typename MT, bool IsCpu>
|
||||
struct AdamWBiasCorrAccessorCompat;
|
||||
|
||||
// CPU specialization: step_count pre-computed on host
|
||||
template <typename MT>
|
||||
struct AdamWBiasCorrAccessorCompat<MT, true> {
|
||||
const double beta1;
|
||||
const double beta2;
|
||||
const float step_count;
|
||||
|
||||
AdamWBiasCorrAccessorCompat(double b1, double b2, float sc)
|
||||
: beta1(b1), beta2(b2), step_count(sc) {}
|
||||
|
||||
__device__ __forceinline__ MT GetBc1() const {
|
||||
return static_cast<MT>(1) -
|
||||
torch_pow_(static_cast<MT>(beta1), static_cast<MT>(step_count));
|
||||
}
|
||||
__device__ __forceinline__ MT GetBc2() const {
|
||||
return static_cast<MT>(1) -
|
||||
torch_pow_(static_cast<MT>(beta2), static_cast<MT>(step_count));
|
||||
}
|
||||
};
|
||||
|
||||
// GPU specialization: recover step_count from beta1_pow pointer on device
|
||||
template <typename MT>
|
||||
struct AdamWBiasCorrAccessorCompat<MT, false> {
|
||||
const double beta1;
|
||||
const double beta2;
|
||||
const MT* beta1_pow;
|
||||
|
||||
AdamWBiasCorrAccessorCompat(double b1, double b2, const MT* bp1)
|
||||
: beta1(b1), beta2(b2), beta1_pow(bp1) {}
|
||||
|
||||
__device__ __forceinline__ MT GetStepCount() const {
|
||||
return static_cast<MT>(
|
||||
::round(::log(static_cast<double>(*beta1_pow)) / ::log(beta1)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ MT GetBc1() const {
|
||||
return static_cast<MT>(1) -
|
||||
torch_pow_(static_cast<MT>(beta1), GetStepCount());
|
||||
}
|
||||
__device__ __forceinline__ MT GetBc2() const {
|
||||
return static_cast<MT>(1) -
|
||||
torch_pow_(static_cast<MT>(beta2), GetStepCount());
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, // Parameter type (may be fp16/bf16)
|
||||
typename TG, // Gradient type
|
||||
typename MT, // Master precision type (= opmath_t, float for
|
||||
// float/fp16/bf16)
|
||||
typename TM, // Moment estimation type (can be bfloat16)
|
||||
typename LrAccessor,
|
||||
typename BiasCorrAccessor>
|
||||
__global__ void AdamWStyleKernel(const double beta1,
|
||||
const double beta2,
|
||||
const double epsilon,
|
||||
const double weight_decay,
|
||||
LrAccessor lr_accessor,
|
||||
BiasCorrAccessor bias_corr_accessor,
|
||||
const TG* __restrict__ grad,
|
||||
const T* __restrict__ param,
|
||||
T* __restrict__ param_out,
|
||||
const MT* __restrict__ master_param,
|
||||
MT* __restrict__ master_param_out,
|
||||
const TM* __restrict__ moment1,
|
||||
TM* __restrict__ moment1_out,
|
||||
const TM* __restrict__ moment2,
|
||||
TM* __restrict__ moment2_out,
|
||||
const TM* __restrict__ moment2_max,
|
||||
TM* __restrict__ moment2_max_out,
|
||||
int64_t ndim,
|
||||
bool amsgrad) {
|
||||
// Matches torch >= 2.12's fused Adam(W) math after PR#173224 (use fma) and
|
||||
// PR#173227 (use opmath_t, not double): every scalar lives in opmath_t
|
||||
// (== MT, float for fp32/fp16/bf16, double for fp64) and the moment updates
|
||||
// use nested fma in opmath_t.
|
||||
__shared__ MT lr_weight_decay_shared;
|
||||
__shared__ MT bias_correction2_sqrt_shared;
|
||||
__shared__ MT step_size_shared;
|
||||
|
||||
const MT beta1_o = static_cast<MT>(beta1);
|
||||
const MT beta2_o = static_cast<MT>(beta2);
|
||||
const MT eps_o = static_cast<MT>(epsilon);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// lr cast to opmath_t (matches lr_opmath = static_cast<opmath_t>(lr)).
|
||||
const MT lr_o = static_cast<MT>(lr_accessor.GetLrDouble());
|
||||
// bias_correction{1,2} are computed in opmath_t inside the accessor.
|
||||
const MT bias_correction1 = bias_corr_accessor.GetBc1();
|
||||
const MT bias_correction2 = bias_corr_accessor.GetBc2();
|
||||
bias_correction2_sqrt_shared = static_cast<MT>(sqrt(bias_correction2));
|
||||
// step_size = lr / bias_correction1 (opmath_t / opmath_t).
|
||||
step_size_shared = lr_o / bias_correction1;
|
||||
// weight decay: param -= lr * weight_decay * param (all opmath_t).
|
||||
lr_weight_decay_shared = lr_o * static_cast<MT>(weight_decay);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const MT lr_weight_decay = lr_weight_decay_shared;
|
||||
const MT bias_correction2_sqrt = bias_correction2_sqrt_shared;
|
||||
const MT step_size = step_size_shared;
|
||||
|
||||
int64_t id =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
|
||||
for (; id < ndim; id += static_cast<int64_t>(gridDim.x) *
|
||||
static_cast<int64_t>(blockDim.x)) {
|
||||
MT p = master_param ? master_param[id] : static_cast<MT>(param[id]);
|
||||
MT g = static_cast<MT>(grad[id]);
|
||||
MT exp_avg = static_cast<MT>(moment1[id]);
|
||||
MT exp_avg_sq = static_cast<MT>(moment2[id]);
|
||||
|
||||
// Weight decay: param -= lr * weight_decay * param
|
||||
if (weight_decay != 0) {
|
||||
p -= lr_weight_decay * p;
|
||||
}
|
||||
|
||||
// exp_avg = fma(beta1, exp_avg, fma(-beta1, grad, grad))
|
||||
exp_avg = fma(beta1_o, exp_avg, fma(-beta1_o, g, g));
|
||||
// exp_avg_sq = fma(beta2, exp_avg_sq, fma(-beta2, grad*grad, grad*grad))
|
||||
const MT g_sq = g * g;
|
||||
exp_avg_sq = fma(beta2_o, exp_avg_sq, fma(-beta2_o, g_sq, g_sq));
|
||||
|
||||
MT denom;
|
||||
if (amsgrad) {
|
||||
MT max_exp_avg_sq_val = static_cast<MT>(moment2_max[id]);
|
||||
max_exp_avg_sq_val =
|
||||
max_exp_avg_sq_val > exp_avg_sq ? max_exp_avg_sq_val : exp_avg_sq;
|
||||
moment2_max_out[id] = static_cast<TM>(max_exp_avg_sq_val);
|
||||
denom = (sqrt(max_exp_avg_sq_val) / bias_correction2_sqrt) + eps_o;
|
||||
} else {
|
||||
denom = (sqrt(exp_avg_sq) / bias_correction2_sqrt) + eps_o;
|
||||
}
|
||||
|
||||
// param -= step_size * exp_avg / denom
|
||||
p -= step_size * exp_avg / denom;
|
||||
|
||||
moment1_out[id] = static_cast<TM>(exp_avg);
|
||||
moment2_out[id] = static_cast<TM>(exp_avg_sq);
|
||||
param_out[id] = static_cast<T>(p);
|
||||
if (master_param_out) {
|
||||
master_param_out[id] = p;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
PADDLE_API void AdamwDenseKernel_compatible(
|
||||
const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& moment1,
|
||||
const DenseTensor& moment2,
|
||||
const optional<DenseTensor>& moment2_max,
|
||||
const DenseTensor& beta1_pow,
|
||||
const DenseTensor& beta2_pow,
|
||||
const optional<DenseTensor>& master_param,
|
||||
const optional<DenseTensor>& skip_update,
|
||||
const Scalar& beta1,
|
||||
const Scalar& beta2,
|
||||
const Scalar& epsilon,
|
||||
double lr_ratio,
|
||||
double coeff,
|
||||
bool with_decay,
|
||||
bool lazy_mode,
|
||||
int64_t min_row_size_to_use_multithread,
|
||||
bool multi_precision,
|
||||
bool use_global_beta_pow,
|
||||
bool amsgrad,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* moment1_out,
|
||||
DenseTensor* moment2_out,
|
||||
DenseTensor* moment2_max_out,
|
||||
DenseTensor* beta1_pow_out,
|
||||
DenseTensor* beta2_pow_out,
|
||||
DenseTensor* master_param_outs) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
|
||||
bool skip_update_ = false;
|
||||
if (skip_update.is_initialized()) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
skip_update->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
|
||||
skip_update->numel()));
|
||||
std::vector<bool> skip_update_vec;
|
||||
TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
|
||||
skip_update_ = skip_update_vec[0];
|
||||
}
|
||||
|
||||
if (skip_update_) {
|
||||
VLOG(4) << "Adamw skip update";
|
||||
Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
|
||||
Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
|
||||
Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
|
||||
if (amsgrad) {
|
||||
Copy(dev_ctx,
|
||||
moment2_max.get(),
|
||||
dev_ctx.GetPlace(),
|
||||
false,
|
||||
moment2_max_out);
|
||||
}
|
||||
if (!use_global_beta_pow) {
|
||||
Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
|
||||
Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// weight_decay: if with_decay is false, set to 0 (matching torch behavior)
|
||||
double weight_decay = with_decay ? coeff : 0.0;
|
||||
|
||||
double beta1_ = beta1.to<double>();
|
||||
double beta2_ = beta2.to<double>();
|
||||
double epsilon_ = epsilon.to<double>();
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta1_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta1 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta1_pow_out->numel()));
|
||||
PADDLE_ENFORCE_EQ(
|
||||
beta2_pow_out->numel(),
|
||||
1,
|
||||
errors::InvalidArgument("beta2 pow output size should be 1, but received "
|
||||
"value is:%d.",
|
||||
beta2_pow_out->numel()));
|
||||
|
||||
const bool beta_pow_on_cpu =
|
||||
beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace();
|
||||
|
||||
// Get learning rate as double. For GPU learning_rate, load it in the CUDA
|
||||
// kernel to avoid a host copy and synchronization.
|
||||
const bool lr_on_cpu = learning_rate.place() == CPUPlace();
|
||||
double lr_double = 0.0;
|
||||
if (lr_on_cpu) {
|
||||
lr_double = learning_rate.data<double>()[0] * lr_ratio;
|
||||
}
|
||||
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_outs) : nullptr;
|
||||
|
||||
// Launch kernel
|
||||
int threads = 512;
|
||||
int64_t blocks_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int blocks = std::min((param.numel() + threads - 1) / threads, blocks_max);
|
||||
|
||||
// Determine gradient type
|
||||
const bool use_bfloat32_grad = grad.dtype() == DataType::FLOAT32;
|
||||
// Determine moment type
|
||||
const bool use_bfloat16_moments = moment1.dtype() == DataType::BFLOAT16 &&
|
||||
moment2.dtype() == DataType::BFLOAT16;
|
||||
|
||||
#define LAUNCH_ADAMW_STYLE_KERNEL(MOMENT_T) \
|
||||
if (use_bfloat32_grad) { \
|
||||
AdamWStyleKernel<T, \
|
||||
float, \
|
||||
MT, \
|
||||
MOMENT_T, \
|
||||
decltype(lr_accessor), \
|
||||
decltype(bias_corr_accessor)> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
weight_decay, \
|
||||
lr_accessor, \
|
||||
bias_corr_accessor, \
|
||||
grad.data<float>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
amsgrad ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
} else { \
|
||||
AdamWStyleKernel<T, \
|
||||
T, \
|
||||
MT, \
|
||||
MOMENT_T, \
|
||||
decltype(lr_accessor), \
|
||||
decltype(bias_corr_accessor)> \
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>( \
|
||||
beta1_, \
|
||||
beta2_, \
|
||||
epsilon_, \
|
||||
weight_decay, \
|
||||
lr_accessor, \
|
||||
bias_corr_accessor, \
|
||||
grad.data<T>(), \
|
||||
param.data<T>(), \
|
||||
dev_ctx.template Alloc<T>(param_out), \
|
||||
master_in_data, \
|
||||
master_out_data, \
|
||||
moment1.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment1_out), \
|
||||
moment2.data<MOMENT_T>(), \
|
||||
dev_ctx.template Alloc<MOMENT_T>(moment2_out), \
|
||||
amsgrad ? moment2_max->data<MOMENT_T>() : nullptr, \
|
||||
amsgrad ? dev_ctx.template Alloc<MOMENT_T>(moment2_max_out) \
|
||||
: nullptr, \
|
||||
param.numel(), \
|
||||
amsgrad); \
|
||||
}
|
||||
|
||||
#define DISPATCH_ADAMW_STYLE_COMPAT_KERNEL(MOMENT_T) \
|
||||
if (lr_on_cpu) { \
|
||||
AdamWLrAccessor<true> lr_accessor(lr_double); \
|
||||
if (beta_pow_on_cpu) { \
|
||||
const float sc = static_cast<float>( \
|
||||
std::round(std::log(static_cast<double>(beta1_pow.data<MT>()[0])) / \
|
||||
std::log(beta1_))); \
|
||||
AdamWBiasCorrAccessorCompat<MT, true> bias_corr_accessor( \
|
||||
beta1_, beta2_, sc); \
|
||||
LAUNCH_ADAMW_STYLE_KERNEL(MOMENT_T) \
|
||||
} else { \
|
||||
AdamWBiasCorrAccessorCompat<MT, false> bias_corr_accessor( \
|
||||
beta1_, beta2_, beta1_pow.data<MT>()); \
|
||||
LAUNCH_ADAMW_STYLE_KERNEL(MOMENT_T) \
|
||||
} \
|
||||
} else { \
|
||||
AdamWLrAccessor<false> lr_accessor(learning_rate.data<double>(), \
|
||||
lr_ratio); \
|
||||
if (beta_pow_on_cpu) { \
|
||||
const float sc = static_cast<float>( \
|
||||
std::round(std::log(static_cast<double>(beta1_pow.data<MT>()[0])) / \
|
||||
std::log(beta1_))); \
|
||||
AdamWBiasCorrAccessorCompat<MT, true> bias_corr_accessor( \
|
||||
beta1_, beta2_, sc); \
|
||||
LAUNCH_ADAMW_STYLE_KERNEL(MOMENT_T) \
|
||||
} else { \
|
||||
AdamWBiasCorrAccessorCompat<MT, false> bias_corr_accessor( \
|
||||
beta1_, beta2_, beta1_pow.data<MT>()); \
|
||||
LAUNCH_ADAMW_STYLE_KERNEL(MOMENT_T) \
|
||||
} \
|
||||
}
|
||||
|
||||
// This function is only reached when FLAGS_use_accuracy_compatible_kernel is
|
||||
// true (see AdamwDenseKernel), so only the torch-compatible dispatch exists.
|
||||
if (use_bfloat16_moments) {
|
||||
DISPATCH_ADAMW_STYLE_COMPAT_KERNEL(bfloat16)
|
||||
} else {
|
||||
DISPATCH_ADAMW_STYLE_COMPAT_KERNEL(MT)
|
||||
}
|
||||
#undef DISPATCH_ADAMW_STYLE_COMPAT_KERNEL
|
||||
#undef LAUNCH_ADAMW_STYLE_KERNEL
|
||||
|
||||
// Update beta_pow (same as original)
|
||||
if (!use_global_beta_pow) {
|
||||
if (beta_pow_on_cpu) {
|
||||
auto* beta1_pow_out_data = dev_ctx.template HostAlloc<MT>(beta1_pow_out);
|
||||
auto* beta2_pow_out_data = dev_ctx.template HostAlloc<MT>(beta2_pow_out);
|
||||
beta1_pow_out_data[0] = beta1_ * beta1_pow.data<MT>()[0];
|
||||
beta2_pow_out_data[0] = beta2_ * beta2_pow.data<MT>()[0];
|
||||
} else {
|
||||
UpdateBetaPowKernel<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
beta1_,
|
||||
beta2_,
|
||||
beta1_pow.data<MT>(),
|
||||
beta2_pow.data<MT>(),
|
||||
dev_ctx.template Alloc<MT>(beta1_pow_out),
|
||||
dev_ctx.template Alloc<MT>(beta2_pow_out));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(adamw,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AdamwDenseKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::FLOAT64);
|
||||
kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
||||
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
|
||||
kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
|
||||
kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,356 @@
|
||||
// Copyright (c) 2022 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/add_n_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/kernels/impl/add_n_kernel_impl.h"
|
||||
namespace phi {
|
||||
|
||||
#define CEIL_DIV(x, y) (((x) + (y)-1) / (y))
|
||||
|
||||
template <class T>
|
||||
__global__ void SumArrayCUDAKernel(
|
||||
T **in, T *out, int64_t N, size_t in_size, bool read_dst) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
|
||||
MT total(read_dst ? static_cast<MT>(out[idx]) : static_cast<MT>(0));
|
||||
for (int i = 0; i < in_size; ++i) {
|
||||
const T *tmp = in[i];
|
||||
if (tmp) {
|
||||
total += static_cast<MT>(tmp[idx]);
|
||||
}
|
||||
}
|
||||
out[idx] = static_cast<T>(total);
|
||||
}
|
||||
}
|
||||
|
||||
template <class T, class HALF>
|
||||
__global__ void SumArrayMixedTypeCUDAKernel(const T *in_0,
|
||||
void **in_others,
|
||||
T *out,
|
||||
int64_t N,
|
||||
size_t in_others_size,
|
||||
bool read_dst) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
|
||||
MT total(read_dst ? static_cast<MT>(out[idx]) : static_cast<MT>(0));
|
||||
total += static_cast<MT>(in_0[idx]);
|
||||
for (int i = 0; i < in_others_size; ++i) {
|
||||
const HALF *tmp = static_cast<HALF *>(in_others[i]);
|
||||
if (tmp) {
|
||||
total += static_cast<MT>(tmp[idx]);
|
||||
}
|
||||
}
|
||||
out[idx] = static_cast<T>(total);
|
||||
}
|
||||
}
|
||||
|
||||
template <class T>
|
||||
__global__ void SumSelectedRowsCUDAKernel(T **sr_in_out,
|
||||
int64_t N,
|
||||
size_t rows) {
|
||||
CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
|
||||
for (int i = 0; i < 2 * rows; i += 2) {
|
||||
const T *tmp = sr_in_out[i];
|
||||
T *tmp_out = sr_in_out[i + 1];
|
||||
if (tmp && tmp_out) {
|
||||
tmp_out[idx] += tmp[idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AddNKernel(const Context &dev_ctx,
|
||||
const std::vector<const TensorBase *> &x,
|
||||
DenseTensor *out) {
|
||||
if (out && out->numel() == 0) {
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
return;
|
||||
}
|
||||
const size_t in_num = x.size();
|
||||
for (int i = 0; i < in_num; ++i) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
x[i]->valid() && x[i]->has_allocation(),
|
||||
true,
|
||||
common::errors::InvalidArgument(
|
||||
"This argument is invalid, %d-th tensor is uninitialized.", i));
|
||||
}
|
||||
|
||||
constexpr size_t theory_sm_threads = 1024;
|
||||
auto stream = dev_ctx.stream();
|
||||
|
||||
auto max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
auto sm_count = max_threads / theory_sm_threads;
|
||||
size_t tile_size = 0;
|
||||
dim3 grids;
|
||||
dim3 blocks;
|
||||
|
||||
auto ComputeKernelParameter = [&](size_t length) {
|
||||
if (length >= max_threads)
|
||||
tile_size = 1024;
|
||||
else if (length < max_threads && length > sm_count * 128)
|
||||
tile_size = 512;
|
||||
else if (length <= sm_count * 128)
|
||||
tile_size = 256;
|
||||
grids = dim3(CEIL_DIV(length, tile_size), 1, 1);
|
||||
blocks = dim3(tile_size, 1, 1);
|
||||
};
|
||||
auto *out_ptr = dev_ctx.template Alloc<T>(out);
|
||||
bool in_place = false;
|
||||
|
||||
if (x.size() > 0 && x[0]->initialized() && DenseTensor::classof(x[0])) {
|
||||
if ((static_cast<const DenseTensor *>(x[0]))->data() == out->data()) {
|
||||
in_place = true;
|
||||
}
|
||||
}
|
||||
|
||||
// Sum of two tensors
|
||||
if (in_num == 2 && DenseTensor::classof(x[0]) && DenseTensor::classof(x[1])) {
|
||||
auto &in_0 = *(static_cast<const DenseTensor *>(x[0]));
|
||||
auto &in_1 = *(static_cast<const DenseTensor *>(x[1]));
|
||||
int64_t length_0 = in_0.numel();
|
||||
int64_t length_1 = in_1.numel();
|
||||
if (length_0 && length_1 && in_0.IsInitialized() && in_1.IsInitialized()) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
auto result = EigenVector<T>::Flatten(*out);
|
||||
auto &place = *dev_ctx.eigen_device();
|
||||
auto in_0_e = EigenVector<T>::Flatten(in_0).template cast<MT>();
|
||||
auto in_1_e = EigenVector<T>::Flatten(in_1).template cast<MT>();
|
||||
result.device(place) = (in_0_e + in_1_e).template cast<T>();
|
||||
} else if (length_0 && in_0.IsInitialized()) {
|
||||
auto result = EigenVector<T>::Flatten(*out);
|
||||
auto &place = *dev_ctx.eigen_device();
|
||||
result.device(place) = EigenVector<T>::Flatten(in_0);
|
||||
} else if (length_1 && in_1.IsInitialized()) {
|
||||
auto result = EigenVector<T>::Flatten(*out);
|
||||
auto &place = *dev_ctx.eigen_device();
|
||||
result.device(place) = EigenVector<T>::Flatten(in_1);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
int start = in_place ? 1 : 0;
|
||||
if (!in_place) {
|
||||
funcs::SetConstant<GPUContext, T> constant_functor;
|
||||
constant_functor(dev_ctx, out, static_cast<T>(0));
|
||||
}
|
||||
|
||||
// Support mixed inputs for master grad accumulation
|
||||
// conditions:
|
||||
// 1. all inputs are DensorTensor and number >= 2
|
||||
// 2. the first tensor is fp32 type and the others are fp16/bf16 type
|
||||
if (in_num >= 2 && DenseTensor::classof(x[0]) &&
|
||||
x[0]->dtype() == DataType::FLOAT32 &&
|
||||
x[1]->dtype() != DataType::FLOAT32) {
|
||||
auto in_other_dtype = x[1]->dtype();
|
||||
int64_t numel = static_cast<const DenseTensor *>(x[0])->numel();
|
||||
bool all_dense_tensor = true;
|
||||
std::vector<const void *> in_data;
|
||||
const T *in_0 = static_cast<const DenseTensor *>(x[0])->data<T>();
|
||||
for (int i = 1; i < in_num; ++i) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
in_other_dtype,
|
||||
x[i]->dtype(),
|
||||
errors::InvalidArgument("The dtype of inputs should be the same, "
|
||||
"but received the dtype of input 1 is %s, "
|
||||
"input %d is %s",
|
||||
in_other_dtype,
|
||||
i,
|
||||
x[i]->dtype()));
|
||||
if (DenseTensor::classof(x[i])) {
|
||||
auto &in_i = *(static_cast<const DenseTensor *>(x[i]));
|
||||
if (in_i.IsInitialized()) {
|
||||
in_data.emplace_back(in_i.data());
|
||||
}
|
||||
} else {
|
||||
all_dense_tensor = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (all_dense_tensor && (in_other_dtype == DataType::BFLOAT16 ||
|
||||
in_other_dtype == DataType::FLOAT16)) {
|
||||
auto tmp_in_array = phi::memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(), in_data.size() * sizeof(void *));
|
||||
size_t nbytes_in = in_data.size() * sizeof(void *);
|
||||
const void *stable_in = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
reinterpret_cast<uint8_t *>(const_cast<void **>(in_data.data())),
|
||||
nbytes_in);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
tmp_in_array->ptr(),
|
||||
CPUPlace(),
|
||||
stable_in,
|
||||
nbytes_in,
|
||||
dev_ctx.stream());
|
||||
|
||||
void **in_array_data = reinterpret_cast<void **>(tmp_in_array->ptr());
|
||||
ComputeKernelParameter(numel);
|
||||
VLOG(4) << "Call SumArrayMixedTypeCUDAKernel";
|
||||
if (in_other_dtype == DataType::FLOAT16) {
|
||||
SumArrayMixedTypeCUDAKernel<T, phi::float16>
|
||||
<<<grids, blocks, 0, stream>>>(in_0,
|
||||
in_array_data,
|
||||
out->data<T>(),
|
||||
numel,
|
||||
in_data.size(),
|
||||
in_place);
|
||||
} else if (in_other_dtype == DataType::BFLOAT16) {
|
||||
SumArrayMixedTypeCUDAKernel<T, phi::bfloat16>
|
||||
<<<grids, blocks, 0, stream>>>(in_0,
|
||||
in_array_data,
|
||||
out->data<T>(),
|
||||
numel,
|
||||
in_data.size(),
|
||||
in_place);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<const T *> in_data;
|
||||
std::vector<int> selectrow_index;
|
||||
int64_t lod_length = 0;
|
||||
bool dst_write = false;
|
||||
for (int i = start; i < in_num; ++i) {
|
||||
if (DenseTensor::classof(x[i])) {
|
||||
auto &in_i = *(static_cast<const DenseTensor *>(x[i]));
|
||||
lod_length = in_i.numel();
|
||||
if (lod_length && in_i.IsInitialized()) {
|
||||
in_data.emplace_back(in_i.data<T>());
|
||||
}
|
||||
} else if (SelectedRows::classof(x[i])) {
|
||||
selectrow_index.push_back(i);
|
||||
}
|
||||
}
|
||||
|
||||
// compute select rows separately.
|
||||
if (!selectrow_index.empty()) {
|
||||
std::vector<const T *> sr_in_out_data;
|
||||
size_t rows = 0;
|
||||
int64_t length = 0;
|
||||
for (auto index : selectrow_index) {
|
||||
auto &sr = *(static_cast<const SelectedRows *>(x[index]));
|
||||
auto &sr_value = sr.value();
|
||||
auto &sr_rows = sr.rows();
|
||||
|
||||
auto row_numel = sr_value.numel() / sr_rows.size();
|
||||
auto out_dims = out->dims();
|
||||
|
||||
PADDLE_ENFORCE_EQ(sr.height(),
|
||||
out_dims[0],
|
||||
errors::InvalidArgument(
|
||||
"The table height of input must be same as output, "
|
||||
"but received input height is %d"
|
||||
", output height is %d",
|
||||
sr.height(),
|
||||
out_dims[0]));
|
||||
PADDLE_ENFORCE_EQ(row_numel,
|
||||
out->numel() / sr.height(),
|
||||
errors::InvalidArgument(
|
||||
"The table width of input must be same as output, "
|
||||
"but received input width is %d"
|
||||
", output width is %d",
|
||||
row_numel,
|
||||
out->numel() / sr.height()));
|
||||
|
||||
auto *sr_data = sr_value.data<T>();
|
||||
auto *sr_out_data = out->data<T>();
|
||||
rows += sr_rows.size();
|
||||
length = row_numel;
|
||||
|
||||
for (size_t i = 0; i < sr_rows.size(); ++i) {
|
||||
sr_in_out_data.emplace_back(&sr_data[i * row_numel]);
|
||||
sr_in_out_data.emplace_back(&sr_out_data[sr_rows[i] * row_numel]);
|
||||
}
|
||||
}
|
||||
if (!sr_in_out_data.empty()) {
|
||||
auto tmp_sr_in_out_array = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(), sr_in_out_data.size() * sizeof(T *));
|
||||
|
||||
size_t nbytes_sr = sr_in_out_data.size() * sizeof(T *);
|
||||
const void *stable_sr = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
reinterpret_cast<uint8_t *>(sr_in_out_data.data()), nbytes_sr);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
tmp_sr_in_out_array->ptr(),
|
||||
CPUPlace(),
|
||||
stable_sr,
|
||||
nbytes_sr,
|
||||
dev_ctx.stream());
|
||||
|
||||
T **sr_in_out_array_data =
|
||||
reinterpret_cast<T **>(tmp_sr_in_out_array->ptr());
|
||||
|
||||
ComputeKernelParameter(length);
|
||||
SumSelectedRowsCUDAKernel<T>
|
||||
<<<grids, blocks, 0, stream>>>(sr_in_out_array_data, length, rows);
|
||||
dst_write = true;
|
||||
}
|
||||
}
|
||||
// if indata not null, merge into one kernel call.
|
||||
if (!in_data.empty()) {
|
||||
auto tmp_in_array =
|
||||
memory_utils::Alloc(dev_ctx.GetPlace(), in_data.size() * sizeof(T *));
|
||||
|
||||
size_t nbytes_in2 = in_data.size() * sizeof(T *);
|
||||
const void *stable_in2 = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
reinterpret_cast<uint8_t *>(const_cast<T **>(in_data.data())),
|
||||
nbytes_in2);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
tmp_in_array->ptr(),
|
||||
CPUPlace(),
|
||||
stable_in2,
|
||||
nbytes_in2,
|
||||
dev_ctx.stream());
|
||||
|
||||
T **in_array_data = reinterpret_cast<T **>(tmp_in_array->ptr());
|
||||
ComputeKernelParameter(lod_length);
|
||||
SumArrayCUDAKernel<T><<<grids, blocks, 0, stream>>>(in_array_data,
|
||||
out->data<T>(),
|
||||
lod_length,
|
||||
in_data.size(),
|
||||
dst_write | in_place);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(add_n,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AddNKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(add_n_array,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AddNArrayKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
/* Copyright (c) 2022 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/addmm_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/addmm_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(addmm_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AddmmGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
/* Copyright (c) 2022 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/addmm_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/addmm_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(addmm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AddmmKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,174 @@
|
||||
// Copyright (c) 2024 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/affine_channel_grad_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, DataLayout layout, bool HasBias>
|
||||
__global__ static inline void KeAffineChannelCUDA(const T* x,
|
||||
const T* scale,
|
||||
const T* bias,
|
||||
const int C,
|
||||
const int64_t HxW,
|
||||
const int64_t num,
|
||||
T* y) {
|
||||
int64_t gid =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
for (int64_t i = gid; i < num; i += stride) {
|
||||
const int c = layout == DataLayout::NCHW ? i / HxW % C : i % C;
|
||||
if (HasBias) {
|
||||
y[i] = scale[c] * x[i] + bias[c];
|
||||
} else {
|
||||
y[i] = scale[c] * x[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, int BlockDim, DataLayout layout>
|
||||
__global__ void AffineChannelScaleBiasGradientCUDAKernel(const T* dy,
|
||||
const T* x,
|
||||
const int N,
|
||||
const int C,
|
||||
const int64_t HxW,
|
||||
T* dscale,
|
||||
T* dbias) {
|
||||
const int outer_size = C;
|
||||
const int64_t inner_size = HxW * N;
|
||||
typedef cub::BlockReduce<double, BlockDim> BlockReduce;
|
||||
__shared__ typename BlockReduce::TempStorage ds_storage;
|
||||
__shared__ typename BlockReduce::TempStorage db_storage;
|
||||
|
||||
for (int i = blockIdx.x; i < outer_size; i += gridDim.x) {
|
||||
T ds_sum = 0;
|
||||
T db_sum = 0;
|
||||
for (int64_t j = threadIdx.x; j < inner_size; j += blockDim.x) {
|
||||
const int64_t index = layout == DataLayout::NCHW
|
||||
? (j / HxW * C + i) * HxW + j % HxW
|
||||
: j * outer_size + i;
|
||||
ds_sum += dy[index] * x[index];
|
||||
db_sum += dy[index];
|
||||
}
|
||||
__syncthreads();
|
||||
auto ds_out =
|
||||
BlockReduce(ds_storage).Reduce(static_cast<double>(ds_sum), cub::Sum());
|
||||
auto db_out =
|
||||
BlockReduce(db_storage).Reduce(static_cast<double>(db_sum), cub::Sum());
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
dscale[i] = ds_out;
|
||||
dbias[i] = db_out;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineChannelGradCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x_in,
|
||||
const DenseTensor& scale_in,
|
||||
const DenseTensor& bias_in,
|
||||
const DenseTensor& out_grad,
|
||||
const std::string& data_layout,
|
||||
DenseTensor* x_grad,
|
||||
DenseTensor* scale_grad,
|
||||
DenseTensor* bias_grad) {
|
||||
auto* x = &x_in;
|
||||
auto* scale = &scale_in;
|
||||
auto* bias = &bias_in;
|
||||
auto* dy = &out_grad;
|
||||
|
||||
auto* dx = x_grad;
|
||||
auto* dscale = scale_grad;
|
||||
auto* dbias = bias_grad;
|
||||
|
||||
const DataLayout layout = StringToDataLayout(data_layout);
|
||||
|
||||
auto dims = dy->dims();
|
||||
const int64_t num = dy->numel();
|
||||
int64_t N = dims[0];
|
||||
int64_t C = layout == DataLayout::NCHW ? dims[1] : dims[dims.size() - 1];
|
||||
int64_t HxW = num / N / C;
|
||||
|
||||
const T* dy_d = dy->data<T>();
|
||||
const T* s_d = scale->data<T>();
|
||||
|
||||
T* dx_d = dx ? dev_ctx.template Alloc<T>(dx) : nullptr;
|
||||
T* ds_d = dscale ? dev_ctx.template Alloc<T>(dscale) : nullptr;
|
||||
T* db_d = dbias ? dev_ctx.template Alloc<T>(dbias) : nullptr;
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
const int block = 256;
|
||||
#else
|
||||
const int block = 1024;
|
||||
#endif // PADDLE_WITH_HIP
|
||||
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
const int max_blocks = std::max(max_threads / block, 1);
|
||||
int grid1 = (num + block - 1) / block;
|
||||
int grid2 = std::min(static_cast<int>(C), max_blocks);
|
||||
|
||||
// NOTE(large-tensor): Kernel functions expect int for N and C parameters
|
||||
PADDLE_ENFORCE_LE_INT_MAX(N, "N");
|
||||
PADDLE_ENFORCE_LE_INT_MAX(C, "C");
|
||||
if (layout == DataLayout::NCHW) {
|
||||
if (dscale && dbias) {
|
||||
const T* x_d = x->data<T>();
|
||||
AffineChannelScaleBiasGradientCUDAKernel<T, block, DataLayout::NCHW>
|
||||
<<<grid2, block, 0, dev_ctx.stream()>>>(dy_d,
|
||||
x_d,
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(C),
|
||||
HxW,
|
||||
ds_d,
|
||||
db_d);
|
||||
}
|
||||
if (dx) {
|
||||
KeAffineChannelCUDA<T, DataLayout::NCHW, false>
|
||||
<<<grid1, block, 0, dev_ctx.stream()>>>(
|
||||
dy_d, s_d, nullptr, static_cast<int>(C), HxW, num, dx_d);
|
||||
}
|
||||
} else {
|
||||
if (dscale && dbias) {
|
||||
const T* x_d = x->data<T>();
|
||||
AffineChannelScaleBiasGradientCUDAKernel<T, block, DataLayout::NHWC>
|
||||
<<<grid2, block, 0, dev_ctx.stream()>>>(dy_d,
|
||||
x_d,
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(C),
|
||||
HxW,
|
||||
ds_d,
|
||||
db_d);
|
||||
}
|
||||
|
||||
if (dx) {
|
||||
KeAffineChannelCUDA<T, DataLayout::NHWC, false>
|
||||
<<<grid1, block, 0, dev_ctx.stream()>>>(
|
||||
dy_d, s_d, nullptr, static_cast<int>(C), HxW, num, dx_d);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(affine_channel_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AffineChannelGradCUDAKernel,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,102 @@
|
||||
// Copyright (c) 2024 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/affine_channel_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, DataLayout layout, bool HasBias>
|
||||
__global__ static inline void KeAffineChannelCUDA(const T* x,
|
||||
const T* scale,
|
||||
const T* bias,
|
||||
const int C,
|
||||
const int64_t HxW,
|
||||
const int64_t num,
|
||||
T* y) {
|
||||
int64_t gid =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
for (int64_t i = gid; i < num; i += stride) {
|
||||
const int c = layout == DataLayout::NCHW ? i / HxW % C : i % C;
|
||||
if (HasBias) {
|
||||
y[i] = scale[c] * x[i] + bias[c];
|
||||
} else {
|
||||
y[i] = scale[c] * x[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineChannelCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x_in,
|
||||
const DenseTensor& scale_in,
|
||||
const DenseTensor& bias_in,
|
||||
const std::string& data_layout,
|
||||
DenseTensor* out) {
|
||||
auto* x = &x_in;
|
||||
auto* scale = &scale_in;
|
||||
auto* bias = &bias_in;
|
||||
|
||||
auto* y = out;
|
||||
dev_ctx.template Alloc<T>(y);
|
||||
|
||||
const DataLayout layout = StringToDataLayout(data_layout);
|
||||
|
||||
auto dims = x->dims();
|
||||
const int64_t num = x->numel();
|
||||
int64_t N = dims[0];
|
||||
int64_t C = layout == DataLayout::NCHW ? dims[1] : dims[dims.size() - 1];
|
||||
int64_t HxW = num / N / C;
|
||||
|
||||
const T* x_d = x->data<T>();
|
||||
const T* scale_d = scale->data<T>();
|
||||
const T* bias_d = bias->data<T>();
|
||||
T* y_d = y->data<T>();
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
int block = 256;
|
||||
#else
|
||||
int block = 1024;
|
||||
#endif // PADDLE_WITH_HIP
|
||||
int grid = (num + block - 1) / block;
|
||||
|
||||
int max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
grid = std::min(std::max(max_threads / block, 1), grid);
|
||||
|
||||
// NOTE(large-tensor): KeAffineChannelCUDA function signature uses int for C
|
||||
// parameter
|
||||
PADDLE_ENFORCE_LE_INT_MAX(C, "C");
|
||||
if (layout == DataLayout::NCHW) {
|
||||
KeAffineChannelCUDA<T, DataLayout::NCHW, true>
|
||||
<<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
x_d, scale_d, bias_d, static_cast<int>(C), HxW, num, y_d);
|
||||
} else {
|
||||
KeAffineChannelCUDA<T, DataLayout::NHWC, true>
|
||||
<<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
x_d, scale_d, bias_d, static_cast<int>(C), HxW, num, y_d);
|
||||
}
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(affine_channel,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AffineChannelCUDAKernel,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,178 @@
|
||||
// Copyright (c) 2022 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/phi/kernels/affine_grid_grad_kernel.h"
|
||||
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/int_array.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/bmm_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/affine_grid_utils.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGridGrad4DCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& output_grad,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* input_grad) {
|
||||
// The shape of the output grad is [N, H, W, 2]
|
||||
auto grad_grid_dims = output_grad.dims();
|
||||
int64_t n = grad_grid_dims[0];
|
||||
int64_t h = grad_grid_dims[1];
|
||||
int64_t w = grad_grid_dims[2];
|
||||
|
||||
// The shape of input_grad (theta gradient) should be [N, 2, 3]
|
||||
input_grad->Resize({n, 2, 3});
|
||||
T* grad_theta_data = dev_ctx.template Alloc<T>(input_grad);
|
||||
|
||||
if (output_grad.numel() == 0) {
|
||||
Full<T, Context>(dev_ctx, input_grad->dims(), 0, input_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// 1. Directly create the basic grid using the same kernel as the forward
|
||||
// direction
|
||||
DenseTensor base_grid;
|
||||
base_grid.Resize({n, h, w, 3});
|
||||
T* base_grid_data = dev_ctx.template Alloc<T>(&base_grid);
|
||||
|
||||
funcs::CreateBaseGridKernel_4D<T, Context>(
|
||||
dev_ctx, base_grid_data, n, h, w, align_corners);
|
||||
|
||||
// 2. Reshaping base_grid to [N, H * W, 3]
|
||||
DenseTensor base_grid_reshaped;
|
||||
base_grid_reshaped.ShareDataWith(base_grid);
|
||||
base_grid_reshaped.Resize({n, h * w, 3});
|
||||
|
||||
// 3. Transposition base_grid: [N, H * W, 3] ->[N, 3, H * W]
|
||||
DenseTensor base_grid_transposed;
|
||||
base_grid_transposed.Resize({n, 3, h * w});
|
||||
TransposeKernel<T, Context>(
|
||||
dev_ctx, base_grid_reshaped, {0, 2, 1}, &base_grid_transposed);
|
||||
|
||||
// 4. Reshaping Output_grad to [N, H * W, 2]
|
||||
DenseTensor grad_grid_reshaped;
|
||||
grad_grid_reshaped.ShareDataWith(output_grad);
|
||||
grad_grid_reshaped.Resize({n, h * w, 2});
|
||||
|
||||
// 5. Batch matrix multiplication: [N, 3, H * W] x [N, H * W, 2]=[N, 3, 2]
|
||||
DenseTensor grad_theta_temp;
|
||||
grad_theta_temp.Resize({n, 3, 2});
|
||||
|
||||
BmmKernel<T, Context>(
|
||||
dev_ctx, base_grid_transposed, grad_grid_reshaped, &grad_theta_temp);
|
||||
|
||||
// 6. Transposition yields the final result: [N, 3, 2] ->[N, 2, 3]
|
||||
TransposeKernel<T, Context>(dev_ctx, grad_theta_temp, {0, 2, 1}, input_grad);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGridGrad5DCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& output_grad,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* input_grad) {
|
||||
// The shape of the output grad is [N, D, H, W, 3]
|
||||
auto grad_grid_dims = output_grad.dims();
|
||||
int64_t n = grad_grid_dims[0];
|
||||
int64_t d = grad_grid_dims[1];
|
||||
int64_t h = grad_grid_dims[2];
|
||||
int64_t w = grad_grid_dims[3];
|
||||
|
||||
// The shape of input_grad (theta gradient) should be [N, 3, 4]
|
||||
input_grad->Resize({n, 3, 4});
|
||||
T* grad_theta_data = dev_ctx.template Alloc<T>(input_grad);
|
||||
|
||||
if (output_grad.numel() == 0) {
|
||||
Full<T, Context>(dev_ctx, input_grad->dims(), 0, input_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// 1. Directly create the basic grid using the same kernel as the forward
|
||||
// direction
|
||||
DenseTensor base_grid;
|
||||
base_grid.Resize({n, d, h, w, 4});
|
||||
T* base_grid_data = dev_ctx.template Alloc<T>(&base_grid);
|
||||
|
||||
funcs::CreateBaseGridKernel_5D<T, Context>(
|
||||
dev_ctx, base_grid_data, n, d, h, w, align_corners);
|
||||
|
||||
// 2. Reshaping base_grid to [N, D * H * W, 4]
|
||||
DenseTensor base_grid_reshaped;
|
||||
base_grid_reshaped.ShareDataWith(base_grid);
|
||||
base_grid_reshaped.Resize({n, d * h * w, 4});
|
||||
|
||||
// 3. Transpose base_grid:[N，D*H*W，4]->[N，4，D*H*W]
|
||||
DenseTensor base_grid_transposed;
|
||||
base_grid_transposed.Resize({n, 4, d * h * w});
|
||||
TransposeKernel<T, Context>(
|
||||
dev_ctx, base_grid_reshaped, {0, 2, 1}, &base_grid_transposed);
|
||||
|
||||
// 4. Reshaping Output_grad to [N, D * H * W, 3]
|
||||
DenseTensor grad_grid_reshaped;
|
||||
grad_grid_reshaped.ShareDataWith(output_grad);
|
||||
grad_grid_reshaped.Resize({n, d * h * w, 3});
|
||||
|
||||
// 5. Batch matrix multiplication: [N, 4, D * H * W] x [N, D * H * W, 3]=[N,
|
||||
// 4, 3]
|
||||
DenseTensor grad_theta_temp;
|
||||
grad_theta_temp.Resize({n, 4, 3});
|
||||
|
||||
BmmKernel<T, Context>(
|
||||
dev_ctx, base_grid_transposed, grad_grid_reshaped, &grad_theta_temp);
|
||||
|
||||
// 6. Transposition yields the final result: [N, 4, 3] ->[N, 3, 4]
|
||||
TransposeKernel<T, Context>(dev_ctx, grad_theta_temp, {0, 2, 1}, input_grad);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGridGradCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* output) {
|
||||
auto* theta = &input;
|
||||
auto theta_size = theta->dims().size();
|
||||
if (output->numel() == 0 || input.numel() == 0) {
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
funcs::SetConstant<GPUContext, T>()(dev_ctx, output, static_cast<T>(0));
|
||||
return;
|
||||
}
|
||||
if (theta_size == 4) {
|
||||
AffineGridGrad4DCUDAKernel<T, Context>(
|
||||
dev_ctx, input, outputShape, align_corners, output);
|
||||
} else {
|
||||
AffineGridGrad5DCUDAKernel<T, Context>(
|
||||
dev_ctx, input, outputShape, align_corners, output);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(affine_grid_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AffineGridGradCUDAKernel,
|
||||
float,
|
||||
double){};
|
||||
@@ -0,0 +1,144 @@
|
||||
// Copyright (c) 2022 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/phi/kernels/affine_grid_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/int_array.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/bmm_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/affine_grid_utils.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGrid4DCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* output) {
|
||||
auto* theta = &input;
|
||||
int64_t n = theta->dims()[0];
|
||||
auto& size_attr = outputShape.GetData();
|
||||
int64_t h = size_attr[2];
|
||||
int64_t w = size_attr[3];
|
||||
|
||||
if (input.numel() == 0) {
|
||||
output->Resize({n, h, w, 2});
|
||||
Full<T, Context>(dev_ctx, output->dims(), 0, output);
|
||||
return;
|
||||
}
|
||||
|
||||
// Directly create the base mesh
|
||||
DenseTensor base_grid;
|
||||
base_grid.Resize({n, h, w, 3});
|
||||
T* base_grid_data = dev_ctx.template Alloc<T>(&base_grid);
|
||||
|
||||
funcs::CreateBaseGridKernel_4D<T, Context>(
|
||||
dev_ctx, base_grid_data, n, h, w, align_corners);
|
||||
|
||||
// Apply affine transformation
|
||||
DenseTensor base_grid_new;
|
||||
base_grid_new.ShareDataWith(base_grid);
|
||||
base_grid_new.Resize({n, h * w, 3});
|
||||
|
||||
// Transpose theta: [N, 2, 3] -> [N, 3, 2]
|
||||
DenseTensor theta_transposed;
|
||||
theta_transposed.Resize({n, 3, 2});
|
||||
TransposeKernel<T, Context>(dev_ctx, input, {0, 2, 1}, &theta_transposed);
|
||||
|
||||
DenseTensor grid_flat;
|
||||
grid_flat.Resize({n, h * w, 2});
|
||||
BmmKernel<T, Context>(dev_ctx, base_grid_new, theta_transposed, &grid_flat);
|
||||
|
||||
// Reshaping Output
|
||||
output->ShareDataWith(grid_flat);
|
||||
output->Resize({n, h, w, 2});
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGrid5DCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* output) {
|
||||
auto* theta = &input;
|
||||
int64_t n = theta->dims()[0];
|
||||
auto& size_attr = outputShape.GetData();
|
||||
int64_t d = size_attr[2]; // depth
|
||||
int64_t h = size_attr[3]; // height
|
||||
int64_t w = size_attr[4]; // width
|
||||
|
||||
if (input.numel() == 0) {
|
||||
output->Resize({n, d, h, w, 3});
|
||||
Full<T, Context>(dev_ctx, output->dims(), 0, output);
|
||||
return;
|
||||
}
|
||||
|
||||
// Create a basic grid
|
||||
DenseTensor base_grid;
|
||||
base_grid.Resize({n, d, h, w, 4});
|
||||
T* base_grid_data = dev_ctx.template Alloc<T>(&base_grid);
|
||||
|
||||
funcs::CreateBaseGridKernel_5D<T, Context>(
|
||||
dev_ctx, base_grid_data, n, d, h, w, align_corners);
|
||||
|
||||
// Apply affine transformation
|
||||
DenseTensor base_grid_new;
|
||||
base_grid_new.ShareDataWith(base_grid);
|
||||
base_grid_new.Resize({n, d * h * w, 4});
|
||||
|
||||
// Transpose theta: [N, 3, 4] -> [N, 4, 3]
|
||||
DenseTensor theta_transposed;
|
||||
theta_transposed.Resize({n, 4, 3});
|
||||
TransposeKernel<T, Context>(dev_ctx, input, {0, 2, 1}, &theta_transposed);
|
||||
|
||||
// Perform batch matrix multiplication
|
||||
DenseTensor grid_flat;
|
||||
grid_flat.Resize({n, d * h * w, 3});
|
||||
BmmKernel<T, Context>(dev_ctx, base_grid_new, theta_transposed, &grid_flat);
|
||||
|
||||
// Reshaping Output
|
||||
output->ShareDataWith(grid_flat);
|
||||
output->Resize({n, d, h, w, 3});
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AffineGridCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const IntArray& outputShape,
|
||||
bool align_corners,
|
||||
DenseTensor* output) {
|
||||
auto* theta = &input;
|
||||
int64_t theta_h = theta->dims()[1];
|
||||
if (theta_h == 2) {
|
||||
AffineGrid4DCUDAKernel<T, Context>(
|
||||
dev_ctx, input, outputShape, align_corners, output);
|
||||
} else {
|
||||
AffineGrid5DCUDAKernel<T, Context>(
|
||||
dev_ctx, input, outputShape, align_corners, output);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
affine_grid, GPU, ALL_LAYOUT, phi::AffineGridCUDAKernel, float, double){};
|
||||
@@ -0,0 +1,75 @@
|
||||
// Copyright (c) 2023 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/all_gather_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AllGatherKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
int nranks,
|
||||
DenseTensor* out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
auto out_dims = x.dims();
|
||||
out_dims[0] *= nranks;
|
||||
out->Resize(out_dims);
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(
|
||||
comm_ctx,
|
||||
nullptr,
|
||||
errors::Unavailable("NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
PADDLE_ENFORCE_EQ(
|
||||
nranks,
|
||||
comm_ctx->GetSize(),
|
||||
errors::InvalidArgument(
|
||||
"nranks: %s should equal to %s", nranks, comm_ctx->GetSize()));
|
||||
|
||||
gpuStream_t stream = dev_ctx.stream();
|
||||
comm_ctx->AllGather(out, x, stream);
|
||||
#else
|
||||
PADDLE_THROW(
|
||||
errors::PreconditionNotMet("PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(all_gather,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AllGatherKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int64_t,
|
||||
bool,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,98 @@
|
||||
// Copyright (c) 2023 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/all_reduce_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AllReduceKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
int reduce_type,
|
||||
DenseTensor* out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
out->Resize(x.dims());
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(
|
||||
comm_ctx,
|
||||
nullptr,
|
||||
errors::Unavailable("NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
gpuStream_t stream = dev_ctx.stream();
|
||||
PADDLE_ENFORCE_NOT_NULL(stream,
|
||||
errors::NotFound("Should initialize NCCL firstly."));
|
||||
|
||||
ncclRedOp_t red_type = ncclSum;
|
||||
switch (static_cast<ReduceType>(reduce_type)) {
|
||||
case ReduceType::kRedSum:
|
||||
red_type = ncclSum;
|
||||
break;
|
||||
case ReduceType::kRedMax:
|
||||
red_type = ncclMax;
|
||||
break;
|
||||
case ReduceType::kRedMin:
|
||||
red_type = ncclMin;
|
||||
break;
|
||||
case ReduceType::kRedProd:
|
||||
red_type = ncclProd;
|
||||
break;
|
||||
#if NCCL_VERSION_CODE >= 21000
|
||||
case ReduceType::kRedAvg:
|
||||
red_type = ncclAvg;
|
||||
break;
|
||||
#endif
|
||||
case ReduceType::kRedAll:
|
||||
// NOTE(zhonghui): There is no reduce_all type of ncclRedOp_t, just use
|
||||
// min to replace
|
||||
red_type = ncclMin;
|
||||
break;
|
||||
case ReduceType::kRedAny:
|
||||
// NOTE(ooooo): There is no reduce_any type of ncclRedOp_t, just use
|
||||
// max to replace
|
||||
red_type = ncclMax;
|
||||
break;
|
||||
}
|
||||
comm_ctx->AllReduce(out, x, red_type, stream);
|
||||
#else
|
||||
PADDLE_THROW(
|
||||
errors::PreconditionNotMet("PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(all_reduce,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AllReduceKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
bool,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int64_t,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,94 @@
|
||||
// Copyright (c) 2023 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/all_to_all_kernel.h"
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#include "paddle/phi/core/distributed/utils.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AllToAllKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
auto x_dims = x.dims();
|
||||
out->Resize(x_dims);
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(
|
||||
comm_ctx,
|
||||
nullptr,
|
||||
errors::Unavailable("NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
gpuStream_t stream = dev_ctx.stream();
|
||||
PADDLE_ENFORCE_NOT_NULL(stream,
|
||||
errors::NotFound("Should initialize NCCL firstly."));
|
||||
|
||||
int nranks = comm_ctx->GetSize();
|
||||
int64_t send_numel = x.numel() / nranks;
|
||||
size_t offset = 0;
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
x_dims[0] % nranks,
|
||||
0,
|
||||
errors::InvalidArgument(
|
||||
"The first dimension size (%d) of the input tensor must be "
|
||||
"divisible by the number of ranks (%d).",
|
||||
x_dims[0],
|
||||
nranks));
|
||||
|
||||
comm_ctx->GroupStart();
|
||||
|
||||
const auto* send_buf = x.data<T>();
|
||||
auto* recv_buf = out->data<T>();
|
||||
for (auto i = 0; i < nranks; ++i) {
|
||||
auto send_buf = distributed::GetPartialTensor(x, offset, send_numel);
|
||||
comm_ctx->Send(send_buf, send_numel, i, stream);
|
||||
auto recv_buf = distributed::GetPartialTensor(*out, offset, send_numel);
|
||||
comm_ctx->Recv(&recv_buf, send_numel, i, stream);
|
||||
offset += send_numel;
|
||||
}
|
||||
comm_ctx->GroupEnd();
|
||||
#else
|
||||
PADDLE_THROW(
|
||||
errors::PreconditionNotMet("PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(all_to_all,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AllToAllKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int64_t,
|
||||
bool,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,131 @@
|
||||
// Copyright (c) 2022 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/allclose_kernel.h"
|
||||
|
||||
#include <type_traits>
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/data_type.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename IndexType>
|
||||
__global__ void AllcloseCUDAKernel(const T* in_data,
|
||||
const T* other_data,
|
||||
const double rtol,
|
||||
const double atol,
|
||||
bool equal_nan,
|
||||
IndexType num,
|
||||
bool* out_data) {
|
||||
unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
bool val;
|
||||
using BaseMT = typename MPTypeTrait<T>::Type;
|
||||
|
||||
using MT = typename std::conditional<std::is_same<T, int32_t>::value ||
|
||||
std::is_same<T, int64_t>::value ||
|
||||
std::is_same<T, bool>::value,
|
||||
double,
|
||||
BaseMT>::type;
|
||||
|
||||
for (IndexType i = idx; i < num; i += blockDim.x * gridDim.x) {
|
||||
const MT a = static_cast<MT>(in_data[i]);
|
||||
const MT b = static_cast<MT>(other_data[i]);
|
||||
if (isnan(a) || isnan(b)) {
|
||||
val = equal_nan && isnan(a) == isnan(b);
|
||||
} else {
|
||||
MT left = (a > b ? a - b : b - a);
|
||||
MT right = atol + (b > 0 ? rtol * b : (-rtol) * b);
|
||||
MT diff = (left > right ? left - right : right - left);
|
||||
val = a == b || left <= right || diff <= 1e-15;
|
||||
}
|
||||
if (!val) *out_data = false;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AllCloseKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& y,
|
||||
const Scalar& rtol,
|
||||
const Scalar& atol,
|
||||
bool equal_nan,
|
||||
DenseTensor* out) {
|
||||
if (x.numel() == 0 || y.numel() == 0) {
|
||||
Full<bool, Context>(dev_ctx, out->dims(), true, out);
|
||||
return;
|
||||
}
|
||||
double rtol_v, atol_v;
|
||||
if (rtol.dtype() == DataType::FLOAT64) {
|
||||
rtol_v = rtol.to<double>();
|
||||
} else if (rtol.dtype() == DataType::FLOAT32) {
|
||||
rtol_v = rtol.to<float>();
|
||||
} else {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"Input (Rtol) type must be double or float, but get %s.",
|
||||
rtol.dtype()));
|
||||
}
|
||||
if (atol.dtype() == DataType::FLOAT64) {
|
||||
atol_v = atol.to<double>();
|
||||
} else if (atol.dtype() == DataType::FLOAT32) {
|
||||
atol_v = atol.to<float>();
|
||||
} else {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"Input (Atol) type must be double or float, but get %s.",
|
||||
atol.dtype()));
|
||||
}
|
||||
VLOG(3) << "rtol and atol is : " << rtol_v << " " << atol_v;
|
||||
const T* in_data = x.data<T>();
|
||||
const T* other_data = y.data<T>();
|
||||
bool* out_data = dev_ctx.template Alloc<bool>(out);
|
||||
|
||||
int64_t num = x.numel();
|
||||
const int vec_size = 4;
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, num, vec_size);
|
||||
uint32_t grid = config.block_per_grid.x;
|
||||
uint32_t block = config.thread_per_block.x;
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hipMemset(out_data, true, sizeof(bool));
|
||||
#else
|
||||
cudaMemset(out_data, true, sizeof(bool));
|
||||
#endif
|
||||
if (num > std::numeric_limits<int32_t>::max()) {
|
||||
AllcloseCUDAKernel<T, int64_t><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
in_data, other_data, rtol_v, atol_v, equal_nan, num, out_data);
|
||||
} else {
|
||||
AllcloseCUDAKernel<T, int32_t><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
in_data, other_data, rtol_v, atol_v, equal_nan, num, out_data);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(allclose,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AllCloseKernel,
|
||||
float,
|
||||
double,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
// Copyright (c) 2022 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/amp_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/empty_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/amp_kernel_impl.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
// Utils
|
||||
|
||||
template <typename T>
|
||||
__global__ void InverseAndMemset(const T* s, T* o, bool* found_inf) {
|
||||
*o = 1.0 / *s;
|
||||
*found_inf = false;
|
||||
}
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void CheckFiniteAndUnscale(const T** xs,
|
||||
const MT* scale,
|
||||
int64_t size,
|
||||
int64_t* starts,
|
||||
bool* found_inf,
|
||||
T** outs) {
|
||||
const int64_t tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// copy starts array from global memory to shared memory
|
||||
extern __shared__ int64_t s_starts[];
|
||||
for (int64_t i = threadIdx.x; i <= size; i += blockDim.x) {
|
||||
s_starts[i] = starts[i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int64_t num = s_starts[size];
|
||||
int xs_index = 0;
|
||||
bool local_found_inf = false;
|
||||
const MT local_scale = *scale;
|
||||
for (int64_t idx = tid; idx < num; idx += gridDim.x * blockDim.x) {
|
||||
// get the "out" index of "id"
|
||||
// For example:
|
||||
// idx = 15, starts = [0, 10, 10, 20, 30]
|
||||
// because 10 <= idx < 20 ==>
|
||||
// the idx element locate in the 3rd tensor (notice the 2nd tensor size is
|
||||
// 0)
|
||||
int next_xs_index = xs_index;
|
||||
while (idx >= s_starts[next_xs_index]) next_xs_index++;
|
||||
xs_index = next_xs_index - 1;
|
||||
|
||||
// get in data and out data
|
||||
const T* in = xs[xs_index];
|
||||
T* out = outs[xs_index];
|
||||
int64_t in_idx = idx - s_starts[xs_index];
|
||||
|
||||
// Unscale
|
||||
MT val = static_cast<MT>(in[in_idx]) * local_scale;
|
||||
T narrow_val = static_cast<T>(val);
|
||||
out[in_idx] = narrow_val;
|
||||
|
||||
// CheckFinite
|
||||
if (!isfinite(narrow_val)) {
|
||||
local_found_inf = true;
|
||||
}
|
||||
}
|
||||
if (local_found_inf) {
|
||||
*found_inf = true;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename FoundNanInfFlagT>
|
||||
__global__ void GpuUpdateLossScaling(const FoundNanInfFlagT found_inf_data,
|
||||
const T* pre_loss_scaling_data,
|
||||
const int* good_in_data,
|
||||
const int* bad_in_data,
|
||||
const int incr_every_n_steps,
|
||||
const int decr_every_n_nan_or_inf,
|
||||
const float incr_ratio,
|
||||
const float decr_ratio,
|
||||
T* updated_loss_scaling_data,
|
||||
int* good_out_data,
|
||||
int* bad_out_data) {
|
||||
Update<T>(found_inf_data,
|
||||
pre_loss_scaling_data,
|
||||
good_in_data,
|
||||
bad_in_data,
|
||||
incr_every_n_steps,
|
||||
decr_every_n_nan_or_inf,
|
||||
incr_ratio,
|
||||
decr_ratio,
|
||||
updated_loss_scaling_data,
|
||||
good_out_data,
|
||||
bad_out_data);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void FusedFillIf(T** outs,
|
||||
const size_t xs_size,
|
||||
const int64_t* starts,
|
||||
const T value,
|
||||
const bool* has_inf) {
|
||||
if (!(*has_inf)) return;
|
||||
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// copy starts array from global memory to shared memory
|
||||
extern __shared__ int64_t s_starts[];
|
||||
for (size_t i = threadIdx.x; i <= xs_size; i += blockDim.x) {
|
||||
s_starts[i] = starts[i];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int64_t total_num = s_starts[xs_size];
|
||||
int out_index = 0;
|
||||
|
||||
for (int64_t id = tid; id < total_num; id += blockDim.x * gridDim.x) {
|
||||
// get the "out" index of "id"
|
||||
// For example:
|
||||
// id = 15, starts = [0, 10, 10, 20, 30]
|
||||
// because 10 <= id < 20 ==>
|
||||
// the id element locate in the 3rd tensor (notice the 2nd tensor size is 0)
|
||||
int next_out_index = out_index;
|
||||
while (id >= s_starts[next_out_index]) next_out_index++;
|
||||
out_index = next_out_index - 1;
|
||||
|
||||
// get data pointer and index
|
||||
T* out_data = outs[out_index];
|
||||
int64_t idx = id - s_starts[out_index];
|
||||
|
||||
// set value
|
||||
out_data[idx] = value;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
class LazyZeros<GPUContext, T> {
|
||||
public:
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const bool* found_inf_data,
|
||||
const std::vector<const DenseTensor*>& xs,
|
||||
const std::vector<DenseTensor*>& outs) {
|
||||
size_t xs_size = xs.size();
|
||||
if (xs_size == 0) return;
|
||||
|
||||
const auto& cpu_place = CPUPlace();
|
||||
// alloc each tensor's start index and copy to device
|
||||
auto h_in_starts_mem =
|
||||
memory_utils::Alloc(cpu_place, (xs_size + 1) * sizeof(int64_t));
|
||||
int64_t* h_starts = reinterpret_cast<int64_t*>(h_in_starts_mem->ptr());
|
||||
|
||||
auto d_in_starts_mem = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(),
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
int64_t* d_starts = reinterpret_cast<int64_t*>(d_in_starts_mem->ptr());
|
||||
|
||||
// the start index value of each tensor is
|
||||
// the sum of previous tensor's size. For example:
|
||||
// outs = [10, 0, 10, 10] ==> starts = [0, 10, 10, 20, 30]
|
||||
h_starts[0] = 0;
|
||||
for (int i = 0; i < xs_size; i++) {
|
||||
h_starts[i + 1] = h_starts[i] + outs[i]->numel();
|
||||
}
|
||||
auto* stable_h_starts = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
h_starts, xs_size + 1);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_starts,
|
||||
cpu_place,
|
||||
stable_h_starts,
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
dev_ctx.stream());
|
||||
|
||||
// copy each tensor of "outs" data address array to device
|
||||
auto h_out_addrs_mem = memory_utils::Alloc(cpu_place, xs_size * sizeof(T*));
|
||||
T** h_out_addrs = reinterpret_cast<T**>(h_out_addrs_mem->ptr());
|
||||
|
||||
auto d_out_addrs_mem = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(),
|
||||
xs_size * sizeof(T*),
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
T** d_out_addrs = reinterpret_cast<T**>(d_out_addrs_mem->ptr());
|
||||
|
||||
for (size_t i = 0; i < xs_size; ++i) {
|
||||
h_out_addrs[i] = dev_ctx.Alloc<T>(outs[i]);
|
||||
}
|
||||
auto* stable_h_out_addrs =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(h_out_addrs, xs_size);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_out_addrs,
|
||||
cpu_place,
|
||||
stable_h_out_addrs,
|
||||
xs_size * sizeof(T*),
|
||||
dev_ctx.stream());
|
||||
|
||||
// launch cuda kernel
|
||||
int64_t total_num = h_starts[xs_size];
|
||||
int64_t threads_per_block = std::min(static_cast<int64_t>(1024), total_num);
|
||||
int64_t elements_per_block =
|
||||
threads_per_block * 50; // each thread deal with 50 data
|
||||
int64_t blocks_per_grid =
|
||||
(total_num + elements_per_block - 1) / elements_per_block;
|
||||
FusedFillIf<T><<<blocks_per_grid,
|
||||
threads_per_block,
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
dev_ctx.stream()>>>(
|
||||
d_out_addrs, xs_size, d_starts, static_cast<T>(0), found_inf_data);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, bool IsFoundInfOnCPU>
|
||||
class UpdateLossScalingFunctor<GPUContext, T, IsFoundInfOnCPU> {
|
||||
public:
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const bool* found_inf_data,
|
||||
const T* pre_loss_scaling_data,
|
||||
const int* good_in_data,
|
||||
const int* bad_in_data,
|
||||
const int incr_every_n_steps,
|
||||
const int decr_every_n_nan_or_inf,
|
||||
const float incr_ratio,
|
||||
const float decr_ratio,
|
||||
T* updated_loss_scaling_data,
|
||||
int* good_out_data,
|
||||
int* bad_out_data) const {
|
||||
if (IsFoundInfOnCPU) {
|
||||
GpuUpdateLossScaling<T>
|
||||
<<<1, 1, 0, dev_ctx.stream()>>>(*found_inf_data,
|
||||
pre_loss_scaling_data,
|
||||
good_in_data,
|
||||
bad_in_data,
|
||||
incr_every_n_steps,
|
||||
decr_every_n_nan_or_inf,
|
||||
incr_ratio,
|
||||
decr_ratio,
|
||||
updated_loss_scaling_data,
|
||||
good_out_data,
|
||||
bad_out_data);
|
||||
} else {
|
||||
GpuUpdateLossScaling<T>
|
||||
<<<1, 1, 0, dev_ctx.stream()>>>(found_inf_data,
|
||||
pre_loss_scaling_data,
|
||||
good_in_data,
|
||||
bad_in_data,
|
||||
incr_every_n_steps,
|
||||
decr_every_n_nan_or_inf,
|
||||
incr_ratio,
|
||||
decr_ratio,
|
||||
updated_loss_scaling_data,
|
||||
good_out_data,
|
||||
bad_out_data);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Kernels
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CheckFiniteAndUnscaleKernel(const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& xs,
|
||||
const DenseTensor& scale,
|
||||
std::vector<DenseTensor*> outs,
|
||||
DenseTensor* found_infinite) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
|
||||
const MT* scale_data = scale.data<MT>();
|
||||
bool* found_inf_data = dev_ctx.template Alloc<bool>(found_infinite);
|
||||
|
||||
DenseTensor inverse_scale = Empty<MT>(dev_ctx, {1});
|
||||
MT* inverse_scale_v = inverse_scale.template data<MT>();
|
||||
|
||||
InverseAndMemset<MT><<<1, 1, 0, dev_ctx.stream()>>>(
|
||||
scale_data, inverse_scale_v, found_inf_data);
|
||||
|
||||
size_t xs_size = xs.size();
|
||||
if (xs_size == 0) return;
|
||||
|
||||
const auto& cpu_place = CPUPlace();
|
||||
// calculate each tensor's start index and copy to device
|
||||
auto h_starts_tensor =
|
||||
memory_utils::Alloc(cpu_place, (xs_size + 1) * sizeof(int64_t));
|
||||
int64_t* h_starts = reinterpret_cast<int64_t*>(h_starts_tensor->ptr());
|
||||
|
||||
auto d_starts_tensor =
|
||||
memory_utils::Alloc(dev_ctx.GetPlace(),
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
int64_t* d_starts = reinterpret_cast<int64_t*>(d_starts_tensor->ptr());
|
||||
|
||||
// the start index value of each tensor is
|
||||
// the sum of previous tensor's size. For example:
|
||||
// x = [10, 0, 10, 10] ==> starts = [0, 10, 10, 20, 30]
|
||||
h_starts[0] = 0;
|
||||
for (int i = 1; i <= xs_size; i++) {
|
||||
h_starts[i] = h_starts[i - 1] + xs[i - 1]->numel();
|
||||
}
|
||||
int64_t total_num = h_starts[xs_size];
|
||||
auto* stable_h_starts =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(h_starts, xs_size + 1);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_starts,
|
||||
cpu_place,
|
||||
stable_h_starts,
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
dev_ctx.stream());
|
||||
|
||||
// copy each tensor's data address to device
|
||||
auto h_mem = memory_utils::Alloc(cpu_place, 2 * xs_size * sizeof(T*));
|
||||
const T** h_xs = reinterpret_cast<const T**>(h_mem->ptr());
|
||||
T** h_outs = reinterpret_cast<T**>(h_mem->ptr()) + xs_size;
|
||||
|
||||
auto d_mem =
|
||||
memory_utils::Alloc(dev_ctx.GetPlace(),
|
||||
2 * xs_size * sizeof(T*),
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
const T** d_xs = reinterpret_cast<const T**>(d_mem->ptr());
|
||||
T** d_outs = reinterpret_cast<T**>(d_mem->ptr()) + xs_size;
|
||||
|
||||
for (size_t i = 0; i < xs_size; ++i) {
|
||||
h_xs[i] = xs[i]->data<T>();
|
||||
h_outs[i] = dev_ctx.template Alloc<T>(outs[i]);
|
||||
}
|
||||
auto* stable_h_xs =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(h_xs, 2 * xs_size);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
d_xs,
|
||||
cpu_place,
|
||||
stable_h_xs,
|
||||
2 * xs_size * sizeof(T*),
|
||||
dev_ctx.stream());
|
||||
|
||||
// Launch Kernel
|
||||
int threads_per_block = std::min(static_cast<int64_t>(1024), total_num);
|
||||
int elements_per_block =
|
||||
threads_per_block * 20; // each thread deal with 20 number
|
||||
int blocks_per_grid =
|
||||
(total_num + elements_per_block - 1) / elements_per_block;
|
||||
CheckFiniteAndUnscale<T, MT><<<blocks_per_grid,
|
||||
threads_per_block,
|
||||
(xs_size + 1) * sizeof(int64_t),
|
||||
dev_ctx.stream()>>>(
|
||||
d_xs, inverse_scale_v, xs_size, d_starts, found_inf_data, d_outs);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(check_finite_and_unscale,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CheckFiniteAndUnscaleKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::BOOL);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(update_loss_scaling,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::UpdateLossScalingKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
|
||||
kernel_key.dtype() == phi::DataType::BFLOAT16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
|
||||
}
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::INT32);
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::INT32);
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright (c) 2022 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/angle_grad_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/angle_grad_kernel_impl.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
PD_REGISTER_KERNEL(angle_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AngleGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(1).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
|
||||
// Copyright (c) 2022 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/angle_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/angle_kernel_impl.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
PD_REGISTER_KERNEL(angle,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AngleKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
// 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/gpu/ap_facade_kernel.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApFacadeKernel(const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
int64_t num_outputs,
|
||||
const std::string& custom_op_name,
|
||||
const std::string& infer_meta_func_name,
|
||||
const std::string& infer_symbolic_func_name,
|
||||
const std::string& serialized_attributes,
|
||||
std::vector<DenseTensor*> outs) {
|
||||
PADDLE_THROW(
|
||||
common::errors::Unimplemented("ap_facade has no kernel registered."));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(ap_facade,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApFacadeKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,33 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApFacadeKernel(const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
int64_t num_outputs,
|
||||
const std::string& custom_op_name,
|
||||
const std::string& infer_meta_func_name,
|
||||
const std::string& infer_symbolic_func_name,
|
||||
const std::string& serialized_attributes,
|
||||
std::vector<DenseTensor*> outs);
|
||||
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,45 @@
|
||||
// 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/gpu/ap_trivial_fusion_begin_kernel.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApTrivialFusionBeginKernel(
|
||||
const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
DenseTensor* out) {
|
||||
PADDLE_THROW(common::errors::Unimplemented(
|
||||
"pd_op.ap_trivial_fusion_begin has no kernel registered."));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(ap_trivial_fusion_begin,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApTrivialFusionBeginKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApTrivialFusionBeginKernel(
|
||||
const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
DenseTensor* out);
|
||||
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,45 @@
|
||||
// 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/gpu/ap_trivial_fusion_end_kernel.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApTrivialFusionEndKernel(
|
||||
const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
DenseTensor* out) {
|
||||
PADDLE_THROW(common::errors::Unimplemented(
|
||||
"pd_op.ap_trivial_fusion_end has no kernel registered."));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(ap_trivial_fusion_end,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApTrivialFusionEndKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApTrivialFusionEndKernel(
|
||||
const Context& dev_ctx,
|
||||
const optional<std::vector<const DenseTensor*>>& xs,
|
||||
DenseTensor* out);
|
||||
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,126 @@
|
||||
// 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/gpu/ap_variadic_kernel.h"
|
||||
#include "paddle/ap/include/axpr/data_type_util.h"
|
||||
#include "paddle/ap/include/kernel_dispatch/ap_variadic_kernel.h"
|
||||
#include "paddle/ap/include/paddle/phi/device_ctx.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename Context>
|
||||
void AllocateOutTensors(const Context& dev_ctx,
|
||||
const std::vector<DenseTensor*>& outs) {
|
||||
for (auto* out : outs) {
|
||||
auto out_dtype = ap::axpr::GetDataTypeFromPhiDataType(out->dtype());
|
||||
PADDLE_ENFORCE_EQ(out_dtype.HasOkValue(),
|
||||
true,
|
||||
common::errors::InvalidArgument(
|
||||
"GetDataTypeFromPhiDataType() failed !"));
|
||||
|
||||
out_dtype.GetOkValue().Match(
|
||||
[&](const ap::axpr::CppDataType<ap::adt::Undefined>&) {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"allocate not support undefined !"));
|
||||
},
|
||||
[&](const ap::axpr::CppDataType<uint64_t>&) {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"allocate not support uint64_t !"));
|
||||
},
|
||||
[&](const ap::axpr::CppDataType<uint32_t>&) {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"allocate not support uint32_t !"));
|
||||
},
|
||||
[&](const ap::axpr::CppDataType<uint16_t>&) {
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"allocate not support uint16_t !"));
|
||||
},
|
||||
[&](const auto& impl) {
|
||||
using tensor_type = typename std::decay_t<decltype(impl)>::type;
|
||||
dev_ctx.template Alloc<tensor_type>(out);
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApVariadicKernel(const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& xs,
|
||||
int num_outputs,
|
||||
const std::string& code_module_lambda,
|
||||
const std::string& infer_symbolic_lambda,
|
||||
const std::string& infer_meta_lambda,
|
||||
const std::string& kernel_dispatch_lambda,
|
||||
const std::string& kernel_dispatch_const_data_lambda,
|
||||
std::vector<DenseTensor*> outs) {
|
||||
PADDLE_ENFORCE_GT(
|
||||
xs.size(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"At least 1 input is required. current number out uts: // %d",
|
||||
xs.size()));
|
||||
PADDLE_ENFORCE_GT(
|
||||
outs.size(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"num_outputs must be greater than 1. current _outputs: // %d",
|
||||
outs.size()));
|
||||
|
||||
AllocateOutTensors<Context>(dev_ctx, outs);
|
||||
std::shared_ptr<ap::kernel_dispatch::DeviceCtxImpl> impl =
|
||||
std::make_shared<ap::paddle::DeviceCtx<Context>>(&dev_ctx);
|
||||
ap::kernel_dispatch::DeviceCtx ap_device_ctx{impl};
|
||||
const auto& ret =
|
||||
ap::kernel_dispatch::ApVariadicKernel(ap_device_ctx,
|
||||
xs,
|
||||
num_outputs,
|
||||
code_module_lambda,
|
||||
infer_meta_lambda,
|
||||
kernel_dispatch_lambda,
|
||||
kernel_dispatch_const_data_lambda,
|
||||
outs);
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ret.HasError(),
|
||||
false,
|
||||
common::errors::Fatal("ap_variadic failed. \nTraceback (most "
|
||||
"recent call last):\n%s\n%s: %s. ",
|
||||
ret.GetError().CallStackToString(),
|
||||
ret.GetError().class_name(),
|
||||
ret.GetError().msg()));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
PD_REGISTER_KERNEL(ap_variadic,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApVariadicKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
#else
|
||||
PD_REGISTER_KERNEL(ap_variadic,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApVariadicKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
#endif
|
||||
@@ -0,0 +1,34 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApVariadicKernel(const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& xs,
|
||||
int num_outputs,
|
||||
const std::string& code_module_lambda,
|
||||
const std::string& infer_symbolic_lambda,
|
||||
const std::string& infer_meta_lambda,
|
||||
const std::string& kernel_dispatch_lambda,
|
||||
const std::string& kernel_dispatch_const_data_lambda,
|
||||
std::vector<DenseTensor*> outs);
|
||||
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,208 @@
|
||||
// Copyright (c) 2023 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.
|
||||
|
||||
/*
|
||||
* Copyright (c) 2022-2023, NVIDIA CORPORATION. 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/apply_per_channel_scale_kernel.h"
|
||||
|
||||
#include <assert.h>
|
||||
#include <stdint.h>
|
||||
#include <cmath>
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/datatype_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
namespace {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
template <typename T>
|
||||
struct CUDA_HALF_2_TYPE_TARIS {};
|
||||
|
||||
template <>
|
||||
struct CUDA_HALF_2_TYPE_TARIS<half> {
|
||||
using type = half2;
|
||||
};
|
||||
|
||||
#ifdef PADDLE_CUDA_BF16
|
||||
template <>
|
||||
struct CUDA_HALF_2_TYPE_TARIS<__nv_bfloat16> {
|
||||
using type = __nv_bfloat162;
|
||||
};
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
struct HalfMul2 {};
|
||||
|
||||
template <>
|
||||
struct HalfMul2<half2> {
|
||||
static __device__ __forceinline__ half2 apply(const half2& x,
|
||||
const half2& y) {
|
||||
return __hmul2(x, y);
|
||||
}
|
||||
};
|
||||
|
||||
#ifdef PADDLE_CUDA_BF16
|
||||
template <>
|
||||
struct HalfMul2<__nv_bfloat162> {
|
||||
static __device__ __forceinline__ __nv_bfloat162
|
||||
apply(const __nv_bfloat162& x, const __nv_bfloat162& y) {
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
|
||||
return __hmul2(x, y);
|
||||
#else
|
||||
float fxl, fxh, fyl, fyh;
|
||||
fxl = __low2float(x);
|
||||
fxh = __high2float(x);
|
||||
fyl = __low2float(y);
|
||||
fyh = __high2float(y);
|
||||
return __floats2bfloat162_rn(fxl * fyl, fxh * fyh);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
template <typename T, int kProcessRows, typename AccessType>
|
||||
__global__ void apply_per_channel_scale(
|
||||
const T* act, const T* scales, int rows, int cols, T* out) {
|
||||
using HALF_2_TYPE = typename CUDA_HALF_2_TYPE_TARIS<T>::type;
|
||||
static constexpr int kElems = sizeof(AccessType) / sizeof(T);
|
||||
T scale[kElems], act_vec[kElems];
|
||||
int64_t col_offset =
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
|
||||
static_cast<int64_t>(threadIdx.x);
|
||||
int row_offset = blockIdx.y;
|
||||
if (col_offset * kElems >= cols || row_offset * kProcessRows >= rows) return;
|
||||
act += row_offset * kProcessRows * cols;
|
||||
out += row_offset * kProcessRows * cols;
|
||||
*reinterpret_cast<AccessType*>(scale) =
|
||||
reinterpret_cast<const AccessType*>(scales)[col_offset];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kProcessRows; ++i) {
|
||||
*reinterpret_cast<AccessType*>(act_vec) =
|
||||
reinterpret_cast<const AccessType*>(act + i * cols)[col_offset];
|
||||
if constexpr (kElems % 2 == 0 && (std::is_same<T, half>::value
|
||||
#ifdef PADDLE_CUDA_BF16
|
||||
|| std::is_same<T, __nv_bfloat16>::value
|
||||
#endif
|
||||
)) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kElems; j += 2) {
|
||||
*reinterpret_cast<HALF_2_TYPE*>(act_vec + j) =
|
||||
HalfMul2<HALF_2_TYPE>::apply(
|
||||
*reinterpret_cast<HALF_2_TYPE*>(act_vec + j),
|
||||
*reinterpret_cast<HALF_2_TYPE*>(scale + j));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kElems; ++j) {
|
||||
act_vec[j] *= scale[j];
|
||||
}
|
||||
}
|
||||
reinterpret_cast<AccessType*>(out + i * cols)[col_offset] =
|
||||
*reinterpret_cast<AccessType*>(act_vec);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, int kProcessRows, typename AccessType = float4>
|
||||
void apply_per_channel_scale_launcher(const T* act,
|
||||
const T* scales,
|
||||
int rows,
|
||||
int cols,
|
||||
T* out,
|
||||
cudaStream_t stream = 0) {
|
||||
static constexpr int kElems = sizeof(AccessType) / sizeof(T);
|
||||
dim3 block(128);
|
||||
dim3 grid((cols / kElems + block.x - 1) / block.x,
|
||||
(rows + kProcessRows - 1) / kProcessRows);
|
||||
apply_per_channel_scale<T, kProcessRows, AccessType>
|
||||
<<<grid, block, 0, stream>>>(act, scales, rows, cols, out);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
#endif
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ApplyPerChannelScaleKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& scales,
|
||||
DenseTensor* out) {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
using DataType = typename PDDataTypeTraits<T>::DataType;
|
||||
int64_t rows = x.dims()[0];
|
||||
int64_t cols = x.dims()[1];
|
||||
int64_t elems = rows * cols;
|
||||
const T* x_data = x.data<T>();
|
||||
const T* scales_data = scales.data<T>();
|
||||
T* out_data = dev_ctx.template Alloc<T>(out);
|
||||
if (elems < 2048 * 2048) {
|
||||
apply_per_channel_scale_launcher<DataType, 1, float4>(
|
||||
reinterpret_cast<const DataType*>(x_data),
|
||||
reinterpret_cast<const DataType*>(scales_data),
|
||||
rows,
|
||||
cols,
|
||||
reinterpret_cast<DataType*>(out_data),
|
||||
dev_ctx.stream());
|
||||
} else if (elems < 4096 * 4096) {
|
||||
apply_per_channel_scale_launcher<DataType, 4, float4>(
|
||||
reinterpret_cast<const DataType*>(x_data),
|
||||
reinterpret_cast<const DataType*>(scales_data),
|
||||
rows,
|
||||
cols,
|
||||
reinterpret_cast<DataType*>(out_data),
|
||||
dev_ctx.stream());
|
||||
} else if (elems < 8192 * 8192) {
|
||||
apply_per_channel_scale_launcher<DataType, 8, float4>(
|
||||
reinterpret_cast<const DataType*>(x_data),
|
||||
reinterpret_cast<const DataType*>(scales_data),
|
||||
rows,
|
||||
cols,
|
||||
reinterpret_cast<DataType*>(out_data),
|
||||
dev_ctx.stream());
|
||||
} else {
|
||||
PADDLE_ENFORCE_LE_INT_MAX(rows, "rows");
|
||||
PADDLE_ENFORCE_LE_INT_MAX(cols, "cols");
|
||||
apply_per_channel_scale_launcher<DataType, 16, float4>(
|
||||
reinterpret_cast<const DataType*>(x_data),
|
||||
reinterpret_cast<const DataType*>(scales_data),
|
||||
rows,
|
||||
cols,
|
||||
reinterpret_cast<DataType*>(out_data),
|
||||
dev_ctx.stream());
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(apply_per_channel_scale,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ApplyPerChannelScaleKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,171 @@
|
||||
// Copyright (c) 2022 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/arange_kernel.h"
|
||||
|
||||
#include "paddle/common/errors.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/core/visit_type.h"
|
||||
#include "paddle/phi/kernels/funcs/range_function.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename OUT_TYPE>
|
||||
__global__ void Range(T start, T step, int64_t size, OUT_TYPE* out) {
|
||||
CUDA_KERNEL_LOOP_TYPE(index, size, int64_t) {
|
||||
out[index] = static_cast<OUT_TYPE>(start + step * index);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArangeTensorKernel(const Context& dev_ctx,
|
||||
const DenseTensor& start,
|
||||
const DenseTensor& end,
|
||||
const DenseTensor& step,
|
||||
DenseTensor* out) {
|
||||
bool any_float = phi::IsFloatingType(start.dtype()) ||
|
||||
phi::IsFloatingType(end.dtype()) ||
|
||||
phi::IsFloatingType(step.dtype());
|
||||
int64_t size = 0;
|
||||
using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
|
||||
Scalar start_scalar(start);
|
||||
Scalar end_scalar(end);
|
||||
Scalar step_scalar(step);
|
||||
if (any_float) {
|
||||
double sv = start_scalar.to<double>();
|
||||
double ev = end_scalar.to<double>();
|
||||
double stv = step_scalar.to<double>();
|
||||
funcs::GetSize<double>(sv, ev, stv, &size);
|
||||
} else {
|
||||
int64_t sv = start_scalar.to<int64_t>();
|
||||
int64_t ev = end_scalar.to<int64_t>();
|
||||
int64_t stv = step_scalar.to<int64_t>();
|
||||
funcs::GetSize<int64_t>(sv, ev, stv, &size);
|
||||
}
|
||||
MPType start_value = start_scalar.to<MPType>();
|
||||
MPType step_value = step_scalar.to<MPType>();
|
||||
|
||||
out->Resize({size});
|
||||
T* out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto stream = dev_ctx.stream();
|
||||
int64_t block = std::min(size, static_cast<int64_t>(256));
|
||||
if (block == 0) {
|
||||
return;
|
||||
}
|
||||
int64_t grid = (size + block - 1) / block;
|
||||
Range<MPType, T>
|
||||
<<<grid, block, 0, stream>>>(start_value, step_value, size, out_data);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArangeNullaryKernel(const Context& dev_ctx,
|
||||
const T start_value,
|
||||
const T end_value,
|
||||
const T step_value,
|
||||
DenseTensor* out) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
MT start_value_mpt = static_cast<MT>(start_value);
|
||||
MT end_value_mpt = static_cast<MT>(end_value);
|
||||
MT step_value_mpt = static_cast<MT>(step_value);
|
||||
int64_t size = 0;
|
||||
funcs::GetSize(start_value_mpt, end_value_mpt, step_value_mpt, &size);
|
||||
out->Resize({size});
|
||||
T* out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto stream = dev_ctx.stream();
|
||||
int64_t block = std::min(size, static_cast<int64_t>(256));
|
||||
if (block == 0) {
|
||||
return;
|
||||
}
|
||||
int64_t grid = (size + block - 1) / block;
|
||||
Range<MT, T><<<grid, block, 0, stream>>>(
|
||||
start_value_mpt, step_value_mpt, size, out_data);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArangeKernel(const Context& dev_ctx,
|
||||
const Scalar& start,
|
||||
const Scalar& end,
|
||||
const Scalar& step,
|
||||
DenseTensor* out) {
|
||||
bool is_floating = phi::IsFloatingType(start.dtype()) ||
|
||||
phi::IsFloatingType(end.dtype()) ||
|
||||
phi::IsFloatingType(step.dtype());
|
||||
using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
|
||||
int64_t size = 0;
|
||||
if (is_floating) {
|
||||
double sv = start.to<double>();
|
||||
double ev = end.to<double>();
|
||||
double stv = step.to<double>();
|
||||
funcs::GetSize<double>(sv, ev, stv, &size);
|
||||
} else {
|
||||
int64_t sv = start.to<int64_t>();
|
||||
int64_t ev = end.to<int64_t>();
|
||||
int64_t stv = step.to<int64_t>();
|
||||
funcs::GetSize<int64_t>(sv, ev, stv, &size);
|
||||
}
|
||||
MPType start_value = start.to<MPType>();
|
||||
MPType step_value = step.to<MPType>();
|
||||
out->Resize({size});
|
||||
T* out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto stream = dev_ctx.stream();
|
||||
int64_t block = std::min(size, static_cast<int64_t>(256));
|
||||
if (block == 0) {
|
||||
return;
|
||||
}
|
||||
int64_t grid = (size + block - 1) / block;
|
||||
Range<MPType, T>
|
||||
<<<grid, block, 0, stream>>>(start_value, step_value, size, out_data);
|
||||
}
|
||||
|
||||
template void ArangeNullaryKernel<int64_t, GPUContext>(const GPUContext&,
|
||||
const int64_t,
|
||||
const int64_t,
|
||||
const int64_t,
|
||||
DenseTensor*);
|
||||
template void ArangeNullaryKernel<int, GPUContext>(
|
||||
const GPUContext&, const int, const int, const int, DenseTensor*);
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(arange_tensor,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArangeTensorKernel,
|
||||
float,
|
||||
double,
|
||||
int64_t,
|
||||
int,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
kernel->InputAt(2).SetBackend(phi::Backend::ALL_BACKEND);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(arange,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArangeKernel,
|
||||
float,
|
||||
double,
|
||||
int64_t,
|
||||
int,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,297 @@
|
||||
// Copyright (c) 2022 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/arg_min_max_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
#if defined(__NVCC__) || defined(__HIPCC__)
|
||||
|
||||
#include <limits>
|
||||
#include "paddle/common/ddim.h"
|
||||
#include "paddle/phi/core/utils/data_type.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
namespace phi {
|
||||
|
||||
namespace { // NOLINT
|
||||
template <typename K, typename V>
|
||||
using KeyValuePair = cub::KeyValuePair<K, V>;
|
||||
|
||||
} // namespace
|
||||
|
||||
#define FIXED_BLOCK_DIM_CASE_BASE(log2_block_dim, ...) \
|
||||
case (1 << (log2_block_dim)): { \
|
||||
constexpr auto kBlockDim = (1 << (log2_block_dim)); \
|
||||
__VA_ARGS__; \
|
||||
} break
|
||||
|
||||
#define FIXED_BLOCK_DIM_CASE(...) \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(10, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(9, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(8, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(7, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(6, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(5, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(4, ##__VA_ARGS__); \
|
||||
FIXED_BLOCK_DIM_CASE_BASE(3, ##__VA_ARGS__);
|
||||
|
||||
template <typename T,
|
||||
typename IndType,
|
||||
class Reducer,
|
||||
size_t BlockDim,
|
||||
typename IndexType>
|
||||
__global__ void ArgCUDAKernel(const int64_t height, // n * h
|
||||
const int64_t width, // c
|
||||
const int64_t post_size, // h
|
||||
const Reducer reducer,
|
||||
const T init,
|
||||
const T* in,
|
||||
IndType* out) {
|
||||
typedef cub::BlockReduce<KeyValuePair<IndexType, T>, BlockDim> BlockReduce;
|
||||
__shared__ typename BlockReduce::TempStorage temp_storage;
|
||||
|
||||
for (IndexType idx = blockIdx.x; idx < height; idx += gridDim.x) {
|
||||
KeyValuePair<IndexType, T> kv_pair = {-1, init};
|
||||
IndexType h = idx / post_size;
|
||||
IndexType w = idx % post_size;
|
||||
for (IndexType k = threadIdx.x; k < width; k += blockDim.x) {
|
||||
kv_pair =
|
||||
reducer({k, in[h * width * post_size + k * post_size + w]}, kv_pair);
|
||||
}
|
||||
kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, reducer);
|
||||
if (threadIdx.x == 0) {
|
||||
out[idx] = static_cast<IndType>(kv_pair.key);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndType, class Reducer, typename IndexType>
|
||||
void ComputeFullArg(const GPUContext& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
DenseTensor* indices,
|
||||
const int64_t pre,
|
||||
const int64_t post,
|
||||
const int64_t n) {
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
auto ComputeBlockSize = [](int64_t col) {
|
||||
auto block_size = 8;
|
||||
if (col > 512)
|
||||
block_size = 1024;
|
||||
else if (col > 256)
|
||||
block_size = 512;
|
||||
else if (col > 128)
|
||||
block_size = 256;
|
||||
else if (col > 64)
|
||||
block_size = 128;
|
||||
else if (col > 32)
|
||||
block_size = 64;
|
||||
else if (col > 16)
|
||||
block_size = 32;
|
||||
else if (col > 8)
|
||||
block_size = 16;
|
||||
return block_size;
|
||||
};
|
||||
|
||||
int64_t max_grid_dimx = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int64_t height = pre * post;
|
||||
int64_t width = n;
|
||||
int64_t grid_size = height < max_grid_dimx ? height : max_grid_dimx;
|
||||
|
||||
const T* in_data = input.data<T>();
|
||||
IndType* out_data = dev_ctx.template Alloc<IndType>(indices);
|
||||
|
||||
if (typeid(Reducer) == typeid(cub::ArgMax)) {
|
||||
switch (ComputeBlockSize(width)) {
|
||||
FIXED_BLOCK_DIM_CASE(
|
||||
ArgCUDAKernel<T, IndType, Reducer, kBlockDim, IndexType>
|
||||
<<<grid_size, kBlockDim, 0, cu_stream>>>(
|
||||
height,
|
||||
width,
|
||||
post,
|
||||
Reducer(),
|
||||
std::numeric_limits<T>::lowest(),
|
||||
in_data,
|
||||
out_data));
|
||||
}
|
||||
} else {
|
||||
switch (ComputeBlockSize(width)) {
|
||||
FIXED_BLOCK_DIM_CASE(
|
||||
ArgCUDAKernel<T, IndType, Reducer, kBlockDim, IndexType>
|
||||
<<<grid_size, kBlockDim, 0, cu_stream>>>(
|
||||
height,
|
||||
width,
|
||||
post,
|
||||
Reducer(),
|
||||
std::numeric_limits<T>::max(),
|
||||
in_data,
|
||||
out_data));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Context, typename T, class Reducer>
|
||||
struct VisitDataCudaArgMinMaxFunctor {
|
||||
const Context& dev_ctx;
|
||||
const DenseTensor& x;
|
||||
int64_t axis;
|
||||
bool keepdims;
|
||||
bool flatten;
|
||||
DenseTensor* out;
|
||||
|
||||
explicit VisitDataCudaArgMinMaxFunctor(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
int64_t axis,
|
||||
bool keepdims,
|
||||
bool flatten,
|
||||
DenseTensor* out)
|
||||
: dev_ctx(dev_ctx),
|
||||
x(x),
|
||||
axis(axis),
|
||||
keepdims(keepdims),
|
||||
flatten(flatten),
|
||||
out(out) {}
|
||||
|
||||
template <typename IndType>
|
||||
void apply() const {
|
||||
DDim x_dims;
|
||||
int new_axis = axis;
|
||||
if (flatten) {
|
||||
x_dims = make_ddim({x.numel()});
|
||||
// if flatten, the axis just as 0
|
||||
new_axis = 0;
|
||||
} else {
|
||||
x_dims = x.dims();
|
||||
if (axis < 0) new_axis = axis + x.dims().size();
|
||||
}
|
||||
if (x.numel() == 0) {
|
||||
dev_ctx.template Alloc<IndType>(out);
|
||||
return;
|
||||
}
|
||||
// For 0D Tensor
|
||||
if (x.dims().size() == 0) {
|
||||
dev_ctx.template Alloc<IndType>(out);
|
||||
funcs::set_constant(dev_ctx, out, static_cast<IndType>(0));
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t numel = x.numel();
|
||||
int64_t groups = numel / x_dims[new_axis];
|
||||
int64_t pre = 1;
|
||||
int64_t post = 1;
|
||||
int64_t n = x_dims[new_axis];
|
||||
|
||||
for (int i = 0; i < new_axis; i++) {
|
||||
pre *= x_dims[i];
|
||||
}
|
||||
|
||||
for (int i = new_axis + 1; i < x_dims.size(); i++) {
|
||||
post *= x_dims[i];
|
||||
}
|
||||
|
||||
if (numel > std::numeric_limits<int32_t>::max()) {
|
||||
ComputeFullArg<T, IndType, Reducer, int64_t>(
|
||||
dev_ctx, x, out, pre, post, n);
|
||||
} else {
|
||||
ComputeFullArg<T, IndType, Reducer, int32_t>(
|
||||
dev_ctx, x, out, pre, post, n);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Context, typename T, class Reducer>
|
||||
void ArgMinMaxOpCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const Scalar& axis,
|
||||
bool keepdims,
|
||||
bool flatten,
|
||||
DataType dtype,
|
||||
DenseTensor* out) {
|
||||
PADDLE_ENFORCE_GE(
|
||||
x.numel(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"argmin/argmax input numel must > 0, bug got %d", x.numel()));
|
||||
if (dtype == DataType::UNDEFINED) {
|
||||
phi::VisitDataTypeTiny(
|
||||
DataType::INT64,
|
||||
VisitDataCudaArgMinMaxFunctor<Context, T, Reducer>(
|
||||
dev_ctx, x, axis.to<int64_t>(), keepdims, flatten, out));
|
||||
return;
|
||||
}
|
||||
VisitDataTypeTiny(
|
||||
dtype,
|
||||
VisitDataCudaArgMinMaxFunctor<Context, T, Reducer>(
|
||||
dev_ctx, x, axis.to<int64_t>(), keepdims, flatten, out));
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArgMinKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const Scalar& axis,
|
||||
bool keepdims,
|
||||
bool flatten,
|
||||
DataType dtype,
|
||||
DenseTensor* out) {
|
||||
ArgMinMaxOpCUDAKernel<Context, T, cub::ArgMin>(
|
||||
dev_ctx, x, axis, keepdims, flatten, dtype, out);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArgMaxKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const Scalar& axis,
|
||||
bool keepdims,
|
||||
bool flatten,
|
||||
DataType dtype,
|
||||
DenseTensor* out) {
|
||||
ArgMinMaxOpCUDAKernel<Context, T, cub::ArgMax>(
|
||||
dev_ctx, x, axis, keepdims, flatten, dtype, out);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(argmin,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArgMinKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double,
|
||||
int32_t,
|
||||
int64_t,
|
||||
int16_t,
|
||||
uint8_t) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(argmax,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArgMaxKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double,
|
||||
int32_t,
|
||||
int64_t,
|
||||
int16_t,
|
||||
uint8_t) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
// Copyright (c) 2022 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 <thrust/copy.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/sequence.h>
|
||||
#include <thrust/sort.h>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/argsort_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
#ifdef __HIPCC__
|
||||
#include <rocprim/config.hpp>
|
||||
#if defined(ROCPRIM_VERSION) && ROCPRIM_VERSION >= 400000
|
||||
// rocPRIM 4.x (ROCm 7.0+) removed rocprim::detail::radix_key_codec_base.
|
||||
// This TU has no actual cub::*/thrust::*/rocprim::* sort calls, so no
|
||||
// replacement traits are needed; keep this arm empty.
|
||||
#else
|
||||
namespace rocprim {
|
||||
namespace detail {
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::float16>
|
||||
: radix_key_codec_integral<phi::float16, uint16_t> {};
|
||||
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::bfloat16>
|
||||
: radix_key_codec_integral<phi::bfloat16, uint16_t> {};
|
||||
} // namespace detail
|
||||
} // namespace rocprim
|
||||
#endif // ROCPRIM_VERSION
|
||||
#else
|
||||
// set cub base traits in order to handle float16
|
||||
namespace cub {
|
||||
template <>
|
||||
struct NumericTraits<phi::float16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::float16> {};
|
||||
|
||||
template <>
|
||||
struct NumericTraits<phi::bfloat16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::bfloat16> {};
|
||||
} // namespace cub
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename IndType>
|
||||
static __global__ void FillFlattenGrad(const T* dO,
|
||||
const IndType* indices,
|
||||
int64_t size,
|
||||
T* dX) {
|
||||
int64_t index =
|
||||
static_cast<int64_t>(threadIdx.x) +
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
for (int64_t i = index; i < size; i += stride) {
|
||||
dX[indices[i]] = dO[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndType>
|
||||
static __global__ void FillGrad(const T* dO,
|
||||
const IndType* indices,
|
||||
T* dX,
|
||||
IndType num_rows,
|
||||
IndType num_cols) {
|
||||
int col_id = threadIdx.x;
|
||||
int row_id = blockIdx.x;
|
||||
|
||||
for (IndType j = row_id; j < num_rows; j += gridDim.x) {
|
||||
for (IndType i = col_id; i < num_cols; i += blockDim.x) {
|
||||
dX[j * num_cols + indices[j * num_cols + i]] = dO[j * num_cols + i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndType>
|
||||
void ArgFullAssign(const GPUContext& dev_ctx,
|
||||
const DenseTensor* dO,
|
||||
const DenseTensor* indices,
|
||||
DenseTensor* dX,
|
||||
const IndType num_rows,
|
||||
const IndType num_cols) {
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
|
||||
auto ComputeBlockSize = [](IndType col) {
|
||||
if (col > 512)
|
||||
return 1024;
|
||||
else if (col > 256 && col <= 512)
|
||||
return 512;
|
||||
else if (col > 128 && col <= 256)
|
||||
return 256;
|
||||
else if (col > 64 && col <= 128)
|
||||
return 128;
|
||||
else
|
||||
return 64;
|
||||
};
|
||||
|
||||
int block_size = ComputeBlockSize(num_cols);
|
||||
|
||||
int maxGridDimX = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
// actually, int num_rows < max_grid_size
|
||||
int grid_size = num_rows < maxGridDimX ? num_rows : maxGridDimX;
|
||||
FillGrad<<<grid_size, block_size, 0, cu_stream>>>(dO->data<T>(),
|
||||
indices->data<IndType>(),
|
||||
dX->data<T>(),
|
||||
num_rows,
|
||||
num_cols);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void ArgFlattenAssign(const GPUContext& dev_ctx,
|
||||
const DenseTensor* dO,
|
||||
const DenseTensor* indices,
|
||||
int64_t size,
|
||||
DenseTensor* dX) {
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
|
||||
const int64_t block_size =
|
||||
std::min(size, static_cast<int64_t>(dev_ctx.GetMaxThreadsPerBlock()));
|
||||
int64_t max_threads = dev_ctx.GetMaxPhysicalThreadCount();
|
||||
const int64_t max_blocks =
|
||||
std::max(((max_threads - 1) / block_size + 1), static_cast<int64_t>(1));
|
||||
const int64_t grid_size =
|
||||
std::min(max_blocks, (size + block_size - 1) / block_size);
|
||||
|
||||
FillFlattenGrad<<<grid_size, block_size, 0, cu_stream>>>(
|
||||
dO->data<T>(), indices->data<int64_t>(), size, dX->data<T>());
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArgsortGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& indices,
|
||||
const DenseTensor& input,
|
||||
const DenseTensor& out_grad,
|
||||
int axis,
|
||||
bool descending,
|
||||
bool stable,
|
||||
DenseTensor* in_grad) {
|
||||
dev_ctx.template Alloc<T>(in_grad);
|
||||
funcs::set_constant(dev_ctx, in_grad, static_cast<T>(0.0));
|
||||
if (out_grad.numel() == 0) return;
|
||||
auto in_dims = in_grad->dims();
|
||||
auto rank = in_dims.size();
|
||||
axis = (axis < 0) ? (in_dims.size() + axis) : axis;
|
||||
int64_t size = in_grad->numel();
|
||||
|
||||
if (rank == 0) {
|
||||
Copy<Context>(dev_ctx, out_grad, dev_ctx.GetPlace(), false, in_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// Parallel acceleration when the input size is equal to the length of the
|
||||
// 'axis' dimension.
|
||||
// Compared to 'special case for full sort' below, the gradient calculation
|
||||
// is 10 times faster.
|
||||
if (size == in_dims[axis]) {
|
||||
ArgFlattenAssign<T>(dev_ctx, &out_grad, &indices, size, in_grad);
|
||||
return;
|
||||
}
|
||||
|
||||
// Special case for full sort, speedup ~190x.
|
||||
if (axis == -1 || axis + 1 == in_dims.size()) {
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(in_dims, 0, in_dims.size() - 1));
|
||||
const int64_t input_width = in_dims[in_dims.size() - 1];
|
||||
ArgFullAssign<T, int64_t>(
|
||||
dev_ctx, &out_grad, &indices, in_grad, input_height, input_width);
|
||||
} else {
|
||||
// if not full sort, do transpose first
|
||||
std::vector<int> trans;
|
||||
for (int i = 0; i < axis; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(in_dims.size() - 1);
|
||||
for (int i = axis + 1; i < in_dims.size() - 1; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(axis);
|
||||
DDim trans_dims(in_dims);
|
||||
for (int i = 0; i < trans.size(); i++) {
|
||||
trans_dims[i] = in_dims[trans[i]];
|
||||
}
|
||||
|
||||
DenseTensor trans_dO;
|
||||
trans_dO.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<T>(&trans_dO);
|
||||
DenseTensor trans_ind;
|
||||
trans_ind.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<int64_t>(&trans_ind);
|
||||
TransposeKernel<T, Context>(dev_ctx, out_grad, trans, &trans_dO);
|
||||
TransposeKernel<int64_t, Context>(dev_ctx, indices, trans, &trans_ind);
|
||||
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(trans_dims, 0, trans_dims.size() - 1));
|
||||
const int64_t input_width = trans_dims[trans_dims.size() - 1];
|
||||
|
||||
DenseTensor tmp_out;
|
||||
tmp_out.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<T>(&tmp_out);
|
||||
|
||||
ArgFullAssign<T, int64_t>(
|
||||
dev_ctx, &trans_dO, &trans_ind, &tmp_out, input_height, input_width);
|
||||
|
||||
// transpose back
|
||||
TransposeKernel<T, Context>(dev_ctx, tmp_out, trans, in_grad);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(argsort_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArgsortGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,512 @@
|
||||
// Copyright (c) 2022 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/argsort_kernel.h"
|
||||
|
||||
#include <thrust/copy.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/sequence.h>
|
||||
#include <thrust/sort.h>
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
#ifdef __HIPCC__
|
||||
#include <rocprim/config.hpp>
|
||||
#if defined(ROCPRIM_VERSION) && ROCPRIM_VERSION >= 400000
|
||||
// rocPRIM 4.x (ROCm 7.0+) replaces detail::radix_key_codec_base
|
||||
// with traits::define for non-builtin / wrapper types.
|
||||
namespace rocprim {
|
||||
namespace traits {
|
||||
template <>
|
||||
struct define<phi::float16> {
|
||||
using float_bit_mask =
|
||||
float_bit_mask::values<uint16_t, 0x8000, 0x7C00, 0x03FF>;
|
||||
};
|
||||
template <>
|
||||
struct define<phi::bfloat16> {
|
||||
using float_bit_mask =
|
||||
float_bit_mask::values<uint16_t, 0x8000, 0x7F80, 0x007F>;
|
||||
};
|
||||
} // namespace traits
|
||||
} // namespace rocprim
|
||||
#else
|
||||
namespace rocprim {
|
||||
namespace detail {
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::float16>
|
||||
: radix_key_codec_integral<phi::float16, uint16_t> {};
|
||||
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::bfloat16>
|
||||
: radix_key_codec_integral<phi::bfloat16, uint16_t> {};
|
||||
|
||||
#if HIP_VERSION >= 50400000
|
||||
template <>
|
||||
struct float_bit_mask<phi::float16> : float_bit_mask<rocprim::half> {};
|
||||
|
||||
template <>
|
||||
struct float_bit_mask<phi::bfloat16> : float_bit_mask<rocprim::bfloat16> {};
|
||||
#endif
|
||||
} // namespace detail
|
||||
} // namespace rocprim
|
||||
#endif // ROCPRIM_VERSION
|
||||
#else
|
||||
// set cub base traits in order to handle float16
|
||||
namespace cub {
|
||||
template <>
|
||||
struct NumericTraits<phi::float16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::float16> {};
|
||||
|
||||
template <>
|
||||
struct NumericTraits<phi::bfloat16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::bfloat16> {};
|
||||
} // namespace cub
|
||||
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
// Iter for move to next row
|
||||
struct SegmentOffsetIter {
|
||||
EIGEN_DEVICE_FUNC
|
||||
explicit SegmentOffsetIter(int num_cols) : num_cols_(num_cols) {}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(int idx) const {
|
||||
return idx * num_cols_;
|
||||
}
|
||||
|
||||
int num_cols_;
|
||||
};
|
||||
|
||||
template <typename T, typename IndType>
|
||||
__global__ void merge_kernel(const T* A,
|
||||
size_t sizeA,
|
||||
const T* B,
|
||||
size_t sizeB,
|
||||
const IndType* ids_A,
|
||||
const IndType* ids_B,
|
||||
T* out,
|
||||
IndType* out_ids,
|
||||
bool descending) {
|
||||
int64_t thread =
|
||||
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
||||
int64_t num_per_thread = (sizeA + sizeB + thread) / thread;
|
||||
for (int64_t offset = 0; offset < num_per_thread; offset++) {
|
||||
size_t idx =
|
||||
static_cast<size_t>(blockIdx.x) * static_cast<size_t>(blockDim.x) +
|
||||
static_cast<size_t>(threadIdx.x) + offset * thread;
|
||||
size_t total = sizeA + sizeB;
|
||||
if (idx >= total) return;
|
||||
size_t left = (idx > sizeB) ? idx - sizeB : 0;
|
||||
size_t right = (idx < sizeA) ? idx : sizeA;
|
||||
while (left < right) {
|
||||
size_t mid = (left + right) / 2;
|
||||
size_t b_idx = idx - mid;
|
||||
|
||||
T A_mid, B_bidx;
|
||||
if (descending) {
|
||||
A_mid = (mid >= sizeA) ? std::numeric_limits<T>::lowest() : A[mid];
|
||||
B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::lowest() : B[b_idx];
|
||||
} else {
|
||||
A_mid = (mid >= sizeA) ? std::numeric_limits<T>::max() : A[mid];
|
||||
B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::max() : B[b_idx];
|
||||
}
|
||||
|
||||
if (descending ? (A_mid >= B_bidx) : (A_mid <= B_bidx))
|
||||
left = mid + 1;
|
||||
else
|
||||
right = mid;
|
||||
}
|
||||
|
||||
size_t a_idx = left;
|
||||
size_t b_idx = idx - a_idx;
|
||||
if (a_idx >= sizeA) {
|
||||
if (descending ? (A[sizeA - 1] < B[b_idx]) : (A[sizeA - 1] > B[b_idx])) {
|
||||
out[idx] = A[sizeA - 1];
|
||||
out_ids[idx] = ids_A[sizeA - 1];
|
||||
} else {
|
||||
out[idx] = B[b_idx];
|
||||
out_ids[idx] = ids_B[b_idx];
|
||||
}
|
||||
} else if (b_idx >= sizeB) {
|
||||
out[idx] = A[a_idx];
|
||||
out_ids[idx] = ids_A[a_idx];
|
||||
} else {
|
||||
if (descending ? (A[a_idx] >= B[b_idx]) : (A[a_idx] <= B[b_idx])) {
|
||||
out[idx] = A[a_idx];
|
||||
out_ids[idx] = ids_A[a_idx];
|
||||
} else if (descending ? (a_idx > 0 && (A[a_idx - 1] < B[b_idx]))
|
||||
: (a_idx > 0 && (A[a_idx - 1] > B[b_idx]))) {
|
||||
out[idx] = A[a_idx - 1];
|
||||
out_ids[idx] = ids_A[a_idx - 1];
|
||||
} else {
|
||||
out[idx] = B[b_idx];
|
||||
out_ids[idx] = ids_B[b_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static __global__ void FillIndex(T* indices, T num_rows, T num_cols) {
|
||||
int col_id = threadIdx.x;
|
||||
int row_id = blockIdx.x;
|
||||
|
||||
for (T j = row_id; j < num_rows; j += gridDim.x) {
|
||||
for (T i = col_id; i < num_cols; i += blockDim.x) {
|
||||
indices[j * num_cols + i] = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define CUB_ARGSORT_WRAPPER(func, ...) \
|
||||
{ \
|
||||
size_t temp_storage_bytes = 0; \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
func(nullptr, temp_storage_bytes, __VA_ARGS__)); \
|
||||
DenseTensor temp_storage; \
|
||||
int64_t temp_size = static_cast<int64_t>(temp_storage_bytes); \
|
||||
PADDLE_ENFORCE_GT( \
|
||||
temp_size, \
|
||||
0, \
|
||||
common::errors::InvalidArgument( \
|
||||
"Argsort temp storage size is %d, but should be greater than 0.", \
|
||||
temp_size)); \
|
||||
temp_storage.Resize({temp_size}); \
|
||||
dev_ctx.template Alloc<uint8_t>(&temp_storage); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
func(temp_storage.data<uint8_t>(), temp_storage_bytes, __VA_ARGS__)); \
|
||||
}
|
||||
|
||||
#define PREDICATE_CUB_ARGSORT(predicate, if_func, else_func, ...) \
|
||||
if (predicate) \
|
||||
CUB_ARGSORT_WRAPPER(if_func, __VA_ARGS__) \
|
||||
else \
|
||||
CUB_ARGSORT_WRAPPER(else_func, __VA_ARGS__)
|
||||
|
||||
// Sort by flag descending, True: descending. False: Ascending.
|
||||
// Default is false.
|
||||
template <typename T, typename IndType>
|
||||
void ArgFullSort(const GPUContext& dev_ctx,
|
||||
const DenseTensor* input,
|
||||
DenseTensor* output,
|
||||
DenseTensor* indices,
|
||||
const int64_t num_rows,
|
||||
const int64_t num_cols,
|
||||
const bool descending) {
|
||||
PADDLE_ENFORCE_LE_INT_MAX(num_cols, "num_cols");
|
||||
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
auto ComputeBlockSize = [](IndType col) {
|
||||
if (col > 512)
|
||||
return 1024;
|
||||
else if (col > 256 && col <= 512)
|
||||
return 512;
|
||||
else if (col > 128 && col <= 256)
|
||||
return 256;
|
||||
else
|
||||
return 128;
|
||||
};
|
||||
const int block_size = ComputeBlockSize(num_cols);
|
||||
const int64_t maxGridDimX = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
|
||||
const T* inp = input->data<T>();
|
||||
IndType* sorted_indices_ptr = indices->data<IndType>();
|
||||
|
||||
// create iter for counting input
|
||||
cub::CountingInputIterator<IndType> counting_iter(0);
|
||||
// segment_offset is used for move to next row
|
||||
cub::TransformInputIterator<IndType,
|
||||
SegmentOffsetIter,
|
||||
cub::CountingInputIterator<IndType>>
|
||||
segment_offsets_t(counting_iter, SegmentOffsetIter(num_cols));
|
||||
|
||||
// num_rows is the total segments to be sorted
|
||||
constexpr int64_t max_elements = 1 << 30;
|
||||
const int64_t total_elements = num_cols * num_rows;
|
||||
const int64_t segment_size = num_cols;
|
||||
const int64_t element_per_call = std::min(max_elements, total_elements);
|
||||
|
||||
// make sure element_per_call >= segment_size
|
||||
const int64_t adjusted_elements_per_call =
|
||||
std::max(max_elements, segment_size);
|
||||
|
||||
// make sure batch size is the multiple of segment_size
|
||||
const int64_t batch_size =
|
||||
(adjusted_elements_per_call / segment_size) * segment_size;
|
||||
int64_t offset = 0;
|
||||
DenseTensor input_indices;
|
||||
|
||||
T* sorted_out_ptr = sorted_out_ptr = output->data<T>();
|
||||
IndType* ind_ptr = nullptr;
|
||||
|
||||
while (offset < total_elements) {
|
||||
const int64_t n_elements = std::min(batch_size, total_elements - offset);
|
||||
const int64_t n_segments = n_elements / segment_size;
|
||||
|
||||
// allocate a temporary storage for input indices, with shape:
|
||||
// [num_segments = n_elements / segment_size, segment_size]
|
||||
// will be de-allocated once the sort is done, to save memory and
|
||||
// avoid repeated allocation and deallocation
|
||||
if (input_indices.initialized()) {
|
||||
ind_ptr = input_indices.data<IndType>();
|
||||
} else {
|
||||
input_indices.Resize({n_segments, segment_size});
|
||||
ind_ptr = dev_ctx.template Alloc<IndType>(&input_indices);
|
||||
}
|
||||
const int64_t grid_size = std::min(n_segments, maxGridDimX);
|
||||
// Init a index array
|
||||
FillIndex<<<grid_size, block_size, 0, cu_stream>>>(
|
||||
ind_ptr, n_segments, segment_size);
|
||||
|
||||
PREDICATE_CUB_ARGSORT(descending,
|
||||
cub::DeviceSegmentedRadixSort::SortPairsDescending,
|
||||
cub::DeviceSegmentedRadixSort::SortPairs,
|
||||
inp + offset,
|
||||
sorted_out_ptr + offset,
|
||||
ind_ptr,
|
||||
sorted_indices_ptr + offset,
|
||||
n_elements,
|
||||
n_segments,
|
||||
segment_offsets_t,
|
||||
segment_offsets_t + 1,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
cu_stream);
|
||||
offset += n_elements;
|
||||
}
|
||||
}
|
||||
template <typename T, typename IndType>
|
||||
void PerSort(const GPUContext& dev_ctx,
|
||||
T* out_data,
|
||||
int64_t* ids_data,
|
||||
IndType start,
|
||||
IndType end,
|
||||
bool stable,
|
||||
bool descending) {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(dev_ctx.GetPlace(),
|
||||
dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
|
||||
#else
|
||||
phi::memory_utils::ThrustAllocator<hipStream_t> allocator(dev_ctx.GetPlace(),
|
||||
dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
|
||||
#endif
|
||||
if (stable) {
|
||||
if (descending) {
|
||||
thrust::stable_sort_by_key(exec_policy,
|
||||
out_data + start,
|
||||
out_data + end,
|
||||
ids_data + start,
|
||||
thrust::greater<T>());
|
||||
} else {
|
||||
thrust::stable_sort_by_key(
|
||||
exec_policy, out_data + start, out_data + end, ids_data + start);
|
||||
}
|
||||
return;
|
||||
} else {
|
||||
thrust::sort_by_key(
|
||||
exec_policy, out_data + start, out_data + end, ids_data + start);
|
||||
if (descending) {
|
||||
thrust::reverse(exec_policy, out_data + start, out_data + end);
|
||||
thrust::reverse(exec_policy, ids_data + start, ids_data + end);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArgsortKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
int axis,
|
||||
bool descending,
|
||||
bool stable,
|
||||
DenseTensor* output,
|
||||
DenseTensor* indices) {
|
||||
auto in_dims = input.dims();
|
||||
auto rank = in_dims.size();
|
||||
|
||||
if (input.numel() == 0) {
|
||||
output->Resize(in_dims);
|
||||
indices->Resize(in_dims);
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
return;
|
||||
}
|
||||
|
||||
axis = (axis < 0) ? (in_dims.size() + axis) : axis;
|
||||
const T* in_data = input.data<T>();
|
||||
auto size = input.numel();
|
||||
|
||||
if (rank == 0) {
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
Copy<Context>(dev_ctx, input, dev_ctx.GetPlace(), false, output);
|
||||
funcs::set_constant(dev_ctx, indices, static_cast<int64_t>(0));
|
||||
return;
|
||||
}
|
||||
|
||||
// Use thrust for parallel acceleration when the input size is equal to the
|
||||
// length of the 'axis' dimension.
|
||||
// Compared to the following 'Special case for full sort', ascending sort is
|
||||
// 34 times faster and descending sort is 31 times faster.
|
||||
if (size == in_dims[axis]) {
|
||||
T* out_data = dev_ctx.template Alloc<T>(output);
|
||||
int64_t* ids_data = dev_ctx.template Alloc<int64_t>(indices);
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(
|
||||
dev_ctx.GetPlace(), dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
|
||||
#else
|
||||
phi::memory_utils::ThrustAllocator<hipStream_t> allocator(
|
||||
dev_ctx.GetPlace(), dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
|
||||
#endif
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
thrust::sequence(exec_policy, ids_data, ids_data + size);
|
||||
thrust::copy(exec_policy, in_data, in_data + size, out_data);
|
||||
const int64_t per_number = (1LL << 31) - 1;
|
||||
int64_t start = 0;
|
||||
int64_t end = std::min(start + per_number, size);
|
||||
if (end == size) {
|
||||
PerSort<T, int64_t>(
|
||||
dev_ctx, out_data, ids_data, start, end, stable, descending);
|
||||
} else {
|
||||
// Sorting the segments and then merging them
|
||||
DenseTensor temp;
|
||||
DenseTensor ids;
|
||||
temp.Resize(in_dims);
|
||||
ids.Resize(in_dims);
|
||||
T* temp_data = dev_ctx.template Alloc<T>(&temp);
|
||||
int64_t* temp_ids = dev_ctx.template Alloc<int64_t>(&ids);
|
||||
|
||||
while (start != size) {
|
||||
PerSort<T, int64_t>(
|
||||
dev_ctx, out_data, ids_data, start, end, stable, descending);
|
||||
if (start != 0) {
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, end);
|
||||
merge_kernel<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
cu_stream>>>(out_data,
|
||||
start,
|
||||
out_data + start,
|
||||
end - start,
|
||||
ids_data,
|
||||
ids_data + start,
|
||||
temp_data,
|
||||
temp_ids,
|
||||
descending);
|
||||
thrust::copy(exec_policy, temp_ids, temp_ids + end, ids_data);
|
||||
thrust::copy(exec_policy, temp_data, temp_data + end, out_data);
|
||||
}
|
||||
start = end;
|
||||
end = std::min(start + per_number, size);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Special case for full sort, speedup ~190x.
|
||||
if (axis == -1 || axis + 1 == in_dims.size()) {
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(in_dims, 0, in_dims.size() - 1));
|
||||
const int64_t input_width = in_dims[in_dims.size() - 1];
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
ArgFullSort<T, int64_t>(dev_ctx,
|
||||
&input,
|
||||
output,
|
||||
indices,
|
||||
input_height,
|
||||
input_width,
|
||||
descending);
|
||||
} else {
|
||||
// if not full sort, do transpose first
|
||||
std::vector<int> trans;
|
||||
for (int i = 0; i < axis; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(in_dims.size() - 1);
|
||||
for (int i = axis + 1; i < in_dims.size() - 1; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(axis);
|
||||
DDim trans_dims(in_dims);
|
||||
for (int i = 0; i < trans.size(); i++) {
|
||||
trans_dims[i] = in_dims[trans[i]];
|
||||
}
|
||||
|
||||
DenseTensor trans_inp;
|
||||
trans_inp.Resize(trans_dims);
|
||||
T* trans_inp_data = dev_ctx.template Alloc<T>(&trans_inp);
|
||||
// Do transpose
|
||||
TransposeKernel<T, Context>(dev_ctx, input, trans, &trans_inp);
|
||||
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(trans_dims, 0, trans_dims.size() - 1));
|
||||
const int64_t input_width = trans_dims[trans_dims.size() - 1];
|
||||
|
||||
DenseTensor tmp_out;
|
||||
tmp_out.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<T>(&tmp_out);
|
||||
|
||||
DenseTensor tmp_indices;
|
||||
// temp indices for sorting
|
||||
tmp_indices.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<int64_t>(&tmp_indices);
|
||||
|
||||
ArgFullSort<T, int64_t>(dev_ctx,
|
||||
&trans_inp,
|
||||
&tmp_out,
|
||||
&tmp_indices,
|
||||
input_height,
|
||||
input_width,
|
||||
descending);
|
||||
// delay output allocation until after transpose, to avoid
|
||||
// allocating too much memory
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
// transpose back
|
||||
TransposeKernel<T, Context>(dev_ctx, tmp_out, trans, output);
|
||||
TransposeKernel<int64_t, Context>(dev_ctx, tmp_indices, trans, indices);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(argsort,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArgsortKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
// Copyright (c) 2022 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/as_complex_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/as_complex_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
as_complex, GPU, ALL_LAYOUT, phi::AsComplexKernel, float, double) {
|
||||
kernel->OutputAt(0).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2022 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/as_real_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/as_real_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(as_real,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AsRealKernel,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,106 @@
|
||||
// Copyright (c) 2022 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/asgd_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_helper.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/mixed_vector.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void ASGDKernelGPUImpl(const T* param,
|
||||
const T* grad,
|
||||
const T* learning_rate,
|
||||
const T* d,
|
||||
const T* y,
|
||||
const T* n,
|
||||
const MT* master_param,
|
||||
int num,
|
||||
T* param_out,
|
||||
T* d_out,
|
||||
T* y_out,
|
||||
MT* master_param_out) {
|
||||
MT learning_rate_MT = static_cast<MT>(learning_rate[0]);
|
||||
MT n_MT = static_cast<MT>(n[0]);
|
||||
CUDA_KERNEL_LOOP_TYPE(i, num, int64_t) {
|
||||
MT param_data = master_param ? master_param[i] : static_cast<MT>(param[i]);
|
||||
MT grad_data = static_cast<MT>(grad[i]);
|
||||
MT d_data = static_cast<MT>(d[i]);
|
||||
MT y_data = static_cast<MT>(y[i]);
|
||||
d_data = d_data - y_data + grad_data;
|
||||
y_data = grad_data;
|
||||
param_data = param_data - (learning_rate_MT / n_MT) * d_data;
|
||||
param_out[i] = static_cast<T>(param_data);
|
||||
d_out[i] = static_cast<T>(d_data);
|
||||
y_out[i] = static_cast<T>(y_data);
|
||||
if (master_param_out) {
|
||||
master_param_out[i] = param_data;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ASGDKernel(const Context& dev_ctx,
|
||||
const DenseTensor& param,
|
||||
const DenseTensor& grad,
|
||||
const DenseTensor& learning_rate,
|
||||
const DenseTensor& d,
|
||||
const DenseTensor& y,
|
||||
const DenseTensor& n,
|
||||
const optional<DenseTensor>& master_param,
|
||||
bool multi_precision,
|
||||
DenseTensor* param_out,
|
||||
DenseTensor* d_out,
|
||||
DenseTensor* y_out,
|
||||
DenseTensor* master_param_out) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
const MT* master_in_data =
|
||||
multi_precision ? master_param->data<MT>() : nullptr;
|
||||
MT* master_out_data =
|
||||
multi_precision ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;
|
||||
|
||||
int block = 512;
|
||||
int64_t grid_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int grid = std::min((param.numel() + block - 1) / block, grid_max);
|
||||
|
||||
ASGDKernelGPUImpl<T, MT><<<grid, block, 0, dev_ctx.stream()>>>(
|
||||
param.data<T>(),
|
||||
grad.data<T>(),
|
||||
learning_rate.data<T>(),
|
||||
d.data<T>(),
|
||||
y.data<T>(),
|
||||
n.data<T>(),
|
||||
master_in_data,
|
||||
param.numel(),
|
||||
dev_ctx.template Alloc<T>(param_out),
|
||||
dev_ctx.template Alloc<T>(d_out),
|
||||
dev_ctx.template Alloc<T>(y_out),
|
||||
master_out_data);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(asgd,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ASGDKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,89 @@
|
||||
// Copyright (c) 2023 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/assign_pos_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void AssignPos(T* cum_count,
|
||||
const T* numbers,
|
||||
T* out,
|
||||
int64_t limit) {
|
||||
CUDA_KERNEL_LOOP(i, limit) {
|
||||
int number_idx = numbers[i];
|
||||
if (number_idx > -1) {
|
||||
int p = CudaAtomicAdd(cum_count + number_idx, -1);
|
||||
out[p - 1] = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AssignPosKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const DenseTensor& cum_count,
|
||||
const DenseTensor& eff_num_len,
|
||||
DenseTensor* out) {
|
||||
// assign pos decides which tokens should be fetched belong to specially
|
||||
// counter orderingly.
|
||||
auto cum_count_ptr = &cum_count; // (counter number) int32 | int64
|
||||
auto numbers = &x; // (batch_size * seq_len, topk) int32
|
||||
auto eff_num_len_ptr = &eff_num_len; // (sum(cum_count))
|
||||
auto out_ptr = &out; // (cum_count) value ranges
|
||||
// from 0 to batch_size *
|
||||
// seq_len * topk
|
||||
auto numel = numbers->numel();
|
||||
T* cum_data = const_cast<T*>(cum_count_ptr->data<T>());
|
||||
auto cum_size = cum_count_ptr->numel();
|
||||
|
||||
DenseTensor cpu_eff_num_len;
|
||||
int64_t cpu_eff_num_len_data = 0;
|
||||
bool is_cpu_place = eff_num_len_ptr->place() == CPUPlace();
|
||||
if (is_cpu_place) {
|
||||
cpu_eff_num_len_data = eff_num_len_ptr->data<T>()[0];
|
||||
} else {
|
||||
Copy(dev_ctx, eff_num_len, CPUPlace(), false, &cpu_eff_num_len);
|
||||
cpu_eff_num_len_data = cpu_eff_num_len.data<T>()[0];
|
||||
}
|
||||
|
||||
DDim out_dims = make_ddim({cpu_eff_num_len_data});
|
||||
out->Resize(out_dims);
|
||||
auto out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
const T* num_data = numbers->data<T>();
|
||||
|
||||
int64_t blocks = NumBlocks(numel);
|
||||
int threads = kNumCUDAThreads;
|
||||
|
||||
AssignPos<T><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
cum_data, num_data, out_data, numel);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(assign_pos, GPU, ALL_LAYOUT, phi::AssignPosKernel, int64_t) {
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright (c) 2021 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/atan2_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(atan2_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Atan2GradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2021 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/atan2_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(atan2,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::Atan2Kernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
int,
|
||||
int64_t) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
// Copyright (c) 2022 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/auc_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
__global__ void ClearObsoleteDataKernel(int64_t *pos,
|
||||
int64_t *neg,
|
||||
const int bucket_length,
|
||||
const int slide_steps) {
|
||||
int cur_step_index =
|
||||
static_cast<int>(pos[(slide_steps + 1) * bucket_length]) % slide_steps;
|
||||
int64_t cur_step_begin = static_cast<int64_t>(cur_step_index) * bucket_length;
|
||||
int64_t sum_step_begin = static_cast<int64_t>(slide_steps) * bucket_length;
|
||||
CUDA_KERNEL_LOOP(i, bucket_length) {
|
||||
pos[sum_step_begin + i] -= pos[cur_step_begin + i];
|
||||
neg[sum_step_begin + i] -= neg[cur_step_begin + i];
|
||||
pos[cur_step_begin + i] = neg[cur_step_begin + i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void UpdateSumDataKernel(int64_t *pos,
|
||||
int64_t *neg,
|
||||
const int bucket_length,
|
||||
const int slide_steps) {
|
||||
int cur_step_index =
|
||||
static_cast<int>(pos[(slide_steps + 1) * bucket_length]) % slide_steps;
|
||||
int64_t cur_step_begin = static_cast<int64_t>(cur_step_index) * bucket_length;
|
||||
int64_t sum_step_begin = static_cast<int64_t>(slide_steps) * bucket_length;
|
||||
CUDA_KERNEL_LOOP(i, bucket_length) {
|
||||
pos[sum_step_begin + i] += pos[cur_step_begin + i];
|
||||
neg[sum_step_begin + i] += neg[cur_step_begin + i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void AddDataKernel(const int64_t *label_data,
|
||||
const T *pred_data,
|
||||
const int inference_width,
|
||||
const int num_thresholds,
|
||||
int64_t *pos,
|
||||
int64_t *neg,
|
||||
const int numel,
|
||||
const int slide_steps) {
|
||||
int cur_step_begin = 0;
|
||||
if (slide_steps > 0) {
|
||||
int cur_step_index =
|
||||
static_cast<int>(pos[(slide_steps + 1) * (1 + num_thresholds)]) %
|
||||
slide_steps;
|
||||
cur_step_begin = cur_step_index * (1 + num_thresholds);
|
||||
}
|
||||
CUDA_KERNEL_LOOP(i, numel) {
|
||||
auto predict_data = pred_data[i * inference_width + (inference_width - 1)];
|
||||
PADDLE_ENFORCE(predict_data <= 1, "The predict data must less or equal 1.");
|
||||
PADDLE_ENFORCE(predict_data >= 0,
|
||||
"The predict data must gather or equal 0.");
|
||||
uint32_t binIdx = static_cast<uint32_t>(predict_data * num_thresholds);
|
||||
if (label_data[i]) {
|
||||
CudaAtomicAdd(pos + cur_step_begin + binIdx, 1);
|
||||
} else {
|
||||
CudaAtomicAdd(neg + cur_step_begin + binIdx, 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void CalcAucKernel(int64_t *stat_pos,
|
||||
int64_t *stat_neg,
|
||||
int num_thresholds,
|
||||
double *auc,
|
||||
bool need_add_batch_num) {
|
||||
*auc = 0.0f;
|
||||
double totPos = 0.0;
|
||||
double totNeg = 0.0;
|
||||
double totPosPrev = 0.0;
|
||||
double totNegPrev = 0.0;
|
||||
|
||||
int idx = num_thresholds;
|
||||
|
||||
while (idx >= 0) {
|
||||
totPosPrev = totPos;
|
||||
totNegPrev = totNeg;
|
||||
totPos += stat_pos[idx];
|
||||
totNeg += stat_neg[idx];
|
||||
*auc += (totNeg - totNegPrev) * (totPos + totPosPrev) / 2.0;
|
||||
--idx;
|
||||
}
|
||||
|
||||
if (totPos > 0.0 && totNeg > 0.0) {
|
||||
*auc = *auc / totPos / totNeg;
|
||||
}
|
||||
if (need_add_batch_num) {
|
||||
stat_pos[num_thresholds + 1] += 1;
|
||||
stat_neg[num_thresholds + 1] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
inline static double trapezoidArea(double X1, double X2, double Y1, double Y2) {
|
||||
return (X1 > X2 ? (X1 - X2) : (X2 - X1)) * (Y1 + Y2) / 2.0;
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void statAuc(const Context &dev_ctx,
|
||||
const DenseTensor &label,
|
||||
const DenseTensor &predict,
|
||||
const int num_thresholds,
|
||||
const int slide_steps,
|
||||
int64_t *origin_stat_pos,
|
||||
int64_t *origin_stat_neg,
|
||||
const bool is_fake_data) {
|
||||
size_t batch_size = predict.dims()[0];
|
||||
size_t inference_width = predict.dims()[1];
|
||||
const T *inference_data = predict.data<T>();
|
||||
const auto *label_data = label.data<int64_t>();
|
||||
const int bucket_length = num_thresholds + 1;
|
||||
|
||||
if (slide_steps == 0) {
|
||||
AddDataKernel<<<(batch_size + PADDLE_CUDA_NUM_THREADS - 1) /
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
0,
|
||||
dev_ctx.stream()>>>(label_data,
|
||||
inference_data,
|
||||
inference_width,
|
||||
num_thresholds,
|
||||
origin_stat_pos,
|
||||
origin_stat_neg,
|
||||
batch_size,
|
||||
slide_steps);
|
||||
return;
|
||||
}
|
||||
// the last number of origin_stat_pos store the index should be used in
|
||||
// current step
|
||||
int cur_step_index =
|
||||
static_cast<int>(origin_stat_pos[(slide_steps + 1) * bucket_length]) %
|
||||
slide_steps;
|
||||
int64_t cur_step_begin = static_cast<int64_t>(cur_step_index) * bucket_length;
|
||||
int64_t sum_step_begin = static_cast<int64_t>(slide_steps) * bucket_length;
|
||||
|
||||
ClearObsoleteDataKernel<<<(bucket_length + PADDLE_CUDA_NUM_THREADS - 1) /
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
origin_stat_pos, origin_stat_neg, bucket_length, slide_steps);
|
||||
|
||||
AddDataKernel<<<(batch_size + PADDLE_CUDA_NUM_THREADS - 1) /
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
0,
|
||||
dev_ctx.stream()>>>(label_data,
|
||||
inference_data,
|
||||
inference_width,
|
||||
num_thresholds,
|
||||
origin_stat_pos,
|
||||
origin_stat_neg,
|
||||
batch_size,
|
||||
slide_steps);
|
||||
if (!is_fake_data) {
|
||||
UpdateSumDataKernel<<<(bucket_length + PADDLE_CUDA_NUM_THREADS - 1) /
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
PADDLE_CUDA_NUM_THREADS,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
origin_stat_pos, origin_stat_neg, bucket_length, slide_steps);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void AucKernel(const Context &dev_ctx,
|
||||
const DenseTensor &input,
|
||||
const DenseTensor &label,
|
||||
const DenseTensor &stat_pos,
|
||||
const DenseTensor &stat_neg,
|
||||
const optional<DenseTensor> &ins_tag_weight,
|
||||
const std::string &curve,
|
||||
int num_thresholds,
|
||||
int slide_steps,
|
||||
DenseTensor *auc,
|
||||
DenseTensor *stat_pos_out,
|
||||
DenseTensor *stat_neg_out) {
|
||||
// Only use output var for now, make sure it's persistable and
|
||||
// not cleaned up for each batch.
|
||||
auto *origin_stat_pos = dev_ctx.template Alloc<int64_t>(stat_pos_out);
|
||||
auto *origin_stat_neg = dev_ctx.template Alloc<int64_t>(stat_neg_out);
|
||||
auto *auc_value = dev_ctx.template Alloc<double>(auc);
|
||||
|
||||
auto *stat_pos_in_tensor = &stat_pos;
|
||||
auto *stat_neg_in_tensor = &stat_neg;
|
||||
auto *pos_in_data = stat_pos.data<int64_t>();
|
||||
auto *neg_in_data = stat_neg.data<int64_t>();
|
||||
bool is_fake_data = false;
|
||||
if (ins_tag_weight.get_ptr() != nullptr) {
|
||||
const auto *ins_tag_weight_data = ins_tag_weight->data<float>();
|
||||
if (ins_tag_weight_data[0] == 0) {
|
||||
is_fake_data = true;
|
||||
}
|
||||
}
|
||||
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
if (stat_pos_in_tensor != stat_pos_out) {
|
||||
cudaMemcpyAsync(
|
||||
origin_stat_pos,
|
||||
pos_in_data,
|
||||
((1 + slide_steps) * (num_thresholds + 1) + (slide_steps > 0 ? 1 : 0)) *
|
||||
sizeof(int64_t),
|
||||
cudaMemcpyDeviceToDevice,
|
||||
dev_ctx.stream());
|
||||
}
|
||||
if (stat_neg_in_tensor != stat_neg_out) {
|
||||
cudaMemcpyAsync(
|
||||
origin_stat_neg,
|
||||
neg_in_data,
|
||||
((1 + slide_steps) * (num_thresholds + 1) + (slide_steps > 0 ? 1 : 0)) *
|
||||
sizeof(int64_t),
|
||||
cudaMemcpyDeviceToDevice,
|
||||
dev_ctx.stream());
|
||||
}
|
||||
#else
|
||||
if (stat_pos_in_tensor != stat_pos_out) {
|
||||
hipMemcpy(
|
||||
origin_stat_pos,
|
||||
pos_in_data,
|
||||
((1 + slide_steps) * (num_thresholds + 1) + (slide_steps > 0 ? 1 : 0)) *
|
||||
sizeof(int64_t),
|
||||
hipMemcpyDeviceToDevice);
|
||||
}
|
||||
if (stat_neg_in_tensor != stat_neg_out) {
|
||||
hipMemcpy(
|
||||
origin_stat_neg,
|
||||
neg_in_data,
|
||||
((1 + slide_steps) * (num_thresholds + 1) + (slide_steps > 0 ? 1 : 0)) *
|
||||
sizeof(int64_t),
|
||||
hipMemcpyDeviceToDevice);
|
||||
}
|
||||
#endif
|
||||
|
||||
// when calculate global_auc && is fake data, just do nothing
|
||||
if (slide_steps == 0 && is_fake_data) {
|
||||
return;
|
||||
}
|
||||
|
||||
statAuc<T, Context>(dev_ctx,
|
||||
label,
|
||||
input,
|
||||
num_thresholds,
|
||||
slide_steps,
|
||||
origin_stat_pos,
|
||||
origin_stat_neg,
|
||||
is_fake_data);
|
||||
int64_t sum_offset = static_cast<int64_t>(slide_steps) * (num_thresholds + 1);
|
||||
CalcAucKernel<<<1, 1, 0, dev_ctx.stream()>>>(origin_stat_pos + sum_offset,
|
||||
origin_stat_neg + sum_offset,
|
||||
num_thresholds,
|
||||
auc_value,
|
||||
slide_steps > 0);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(auc, GPU, ALL_LAYOUT, phi::AucKernel, float) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::FLOAT64);
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
|
||||
kernel->OutputAt(2).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
/* Copyright (c) 2022 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/average_accumulates_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/average_accumulates_kernel_impl.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <>
|
||||
void GetAccumulators<GPUContext>(const GPUContext& dev_ctx,
|
||||
const DenseTensor& in_num_accumulates,
|
||||
const DenseTensor& in_old_num_accumulates,
|
||||
const DenseTensor& in_num_updates,
|
||||
int64_t* num_updates,
|
||||
int64_t* num_accumulates,
|
||||
int64_t* old_num_accumulates) {
|
||||
auto stream = dev_ctx.stream();
|
||||
auto cuda_place = in_old_num_accumulates.place();
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
old_num_accumulates,
|
||||
cuda_place,
|
||||
in_old_num_accumulates.data<int64_t>(),
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
num_accumulates,
|
||||
cuda_place,
|
||||
in_num_accumulates.data<int64_t>(),
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
num_updates,
|
||||
cuda_place,
|
||||
in_num_updates.data<int64_t>(),
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
}
|
||||
|
||||
template <>
|
||||
void SetAccumulators<GPUContext>(const GPUContext& dev_ctx,
|
||||
int64_t num_updates,
|
||||
int64_t num_accumulates,
|
||||
int64_t old_num_accumulates,
|
||||
DenseTensor* out_num_accumulates,
|
||||
DenseTensor* out_old_num_accumulates,
|
||||
DenseTensor* out_num_updates) {
|
||||
int64_t* out_num_accumulates_ptr =
|
||||
dev_ctx.template Alloc<int64_t>(out_num_accumulates);
|
||||
int64_t* out_old_num_accumulates_ptr =
|
||||
dev_ctx.template Alloc<int64_t>(out_old_num_accumulates);
|
||||
int64_t* out_num_updates_ptr =
|
||||
dev_ctx.template Alloc<int64_t>(out_num_updates);
|
||||
|
||||
auto stream = dev_ctx.stream();
|
||||
|
||||
auto cuda_place = out_old_num_accumulates->place();
|
||||
const int64_t* stable_na =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(&num_accumulates, 1);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
out_num_accumulates_ptr,
|
||||
CPUPlace(),
|
||||
stable_na,
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
|
||||
const int64_t* stable_ona = backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
&old_num_accumulates, 1);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
out_old_num_accumulates_ptr,
|
||||
CPUPlace(),
|
||||
stable_ona,
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
|
||||
const int64_t* stable_nu =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(&num_updates, 1);
|
||||
memory_utils::Copy(cuda_place,
|
||||
out_num_updates_ptr,
|
||||
CPUPlace(),
|
||||
stable_nu,
|
||||
sizeof(int64_t),
|
||||
stream);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(average_accumulates,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::AverageAccumulatesKernel,
|
||||
float,
|
||||
double) {
|
||||
kernel->OutputAt(3).SetDataType(phi::DataType::INT64);
|
||||
kernel->OutputAt(4).SetDataType(phi::DataType::INT64);
|
||||
kernel->OutputAt(5).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
/* 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/baddbmm_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/baddbmm_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(baddbmm_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BaddbmmGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
/* 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/baddbmm_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/baddbmm_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(baddbmm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BaddbmmKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,48 @@
|
||||
// Copyright (c) 2024 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/barrier_kernel.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BarrierKernel(const Context &dev_ctx,
|
||||
const DenseTensor &x,
|
||||
DenseTensor *out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
auto in = &x;
|
||||
auto comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext *>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
auto stream = comm_ctx->GetStream();
|
||||
ncclRedOp_t nccl_red_type = ncclSum;
|
||||
comm_ctx->AllReduce(out, *in, nccl_red_type, stream);
|
||||
backends::gpu::GpuStreamSync(stream);
|
||||
#else
|
||||
PADDLE_THROW(
|
||||
common::errors::Unavailable("PaddlePaddle should compile with NCCL."));
|
||||
#endif
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(barrier, GPU, ALL_LAYOUT, phi::BarrierKernel, int) {}
|
||||
@@ -0,0 +1,150 @@
|
||||
// Copyright (c) 2024 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 <string>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
const int CUDA_NUM_THREADS = 1024;
|
||||
static inline int GET_BLOCKS(const int N) {
|
||||
return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void add_bias_grad_kernel(const T* dout_data,
|
||||
int slot_pairs_num,
|
||||
int ins_num,
|
||||
int out_dim,
|
||||
T* db_data) {
|
||||
CUDA_KERNEL_LOOP(idx, slot_pairs_num * out_dim) {
|
||||
int row = idx / out_dim;
|
||||
int col = idx % out_dim;
|
||||
T temp = static_cast<T>(0);
|
||||
for (int i = 0; i < ins_num; ++i) {
|
||||
int select_index = ((row + 1) * i + 1) * col;
|
||||
temp += dout_data[select_index];
|
||||
}
|
||||
db_data[idx] += temp;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void add_bias_grad(gpuStream_t stream,
|
||||
const T* dout_data,
|
||||
int slot_pairs_num,
|
||||
int ins_num,
|
||||
int out_dim,
|
||||
T* db_data) {
|
||||
add_bias_grad_kernel<<<GET_BLOCKS(slot_pairs_num * out_dim),
|
||||
CUDA_NUM_THREADS,
|
||||
0,
|
||||
stream>>>(
|
||||
dout_data, slot_pairs_num, ins_num, out_dim, db_data);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BatchFCGradOpCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input_in,
|
||||
const DenseTensor& w_in,
|
||||
const DenseTensor& bias_in UNUSED,
|
||||
const DenseTensor& out_grad,
|
||||
DenseTensor* input_grad,
|
||||
DenseTensor* w_grad,
|
||||
DenseTensor* bias_grad) {
|
||||
auto* input = &input_in;
|
||||
auto* w = &w_in;
|
||||
auto* dout = &out_grad;
|
||||
|
||||
auto* dx = input_grad;
|
||||
auto* dw = w_grad;
|
||||
auto* db = bias_grad;
|
||||
|
||||
auto input_dims = input->dims();
|
||||
auto w_dims = w->dims();
|
||||
auto slot_pairs_num = input_dims[0];
|
||||
auto ins_num = input_dims[1];
|
||||
auto in_dim = input_dims[2];
|
||||
auto out_dim = w_dims[2];
|
||||
|
||||
auto& place = *dev_ctx.eigen_device();
|
||||
// initialize
|
||||
dev_ctx.template Alloc<T>(dx);
|
||||
auto dx_eigen = EigenVector<T>::Flatten(*dx);
|
||||
dx_eigen.device(place) = dx_eigen.constant(static_cast<T>(0));
|
||||
|
||||
dev_ctx.template Alloc<T>(dw);
|
||||
auto dw_eigen = EigenVector<T>::Flatten(*dw);
|
||||
dw_eigen.device(place) = dw_eigen.constant(static_cast<T>(0));
|
||||
|
||||
// get data ptr
|
||||
const T* x_data = input->data<T>();
|
||||
const T* w_data = w->data<T>();
|
||||
const T* dout_data = dout->data<T>();
|
||||
T* dx_data = dx->data<T>();
|
||||
T* dw_data = dw->data<T>();
|
||||
|
||||
dev_ctx.template Alloc<T>(db);
|
||||
auto db_eigen = EigenVector<T>::Flatten(*db);
|
||||
db_eigen.device(place) = db_eigen.constant(static_cast<T>(0));
|
||||
T* db_data = db->data<T>();
|
||||
add_bias_grad<T>(
|
||||
dev_ctx.stream(), dout_data, slot_pairs_num, ins_num, out_dim, db_data);
|
||||
|
||||
auto blas = funcs::GetBlas<GPUContext, T>(dev_ctx);
|
||||
T alpha = 1;
|
||||
T beta = 0;
|
||||
|
||||
// dx = dout_data * y^T
|
||||
blas.BatchedGEMM(CblasNoTrans,
|
||||
CblasTrans,
|
||||
ins_num,
|
||||
in_dim,
|
||||
out_dim,
|
||||
alpha,
|
||||
dout_data,
|
||||
w_data,
|
||||
beta,
|
||||
dx_data,
|
||||
slot_pairs_num,
|
||||
ins_num * out_dim,
|
||||
out_dim * in_dim);
|
||||
// dy = x^T * dout_data
|
||||
blas.BatchedGEMM(CblasTrans,
|
||||
CblasNoTrans,
|
||||
in_dim,
|
||||
out_dim,
|
||||
ins_num,
|
||||
alpha,
|
||||
x_data,
|
||||
dout_data,
|
||||
beta,
|
||||
dw_data,
|
||||
slot_pairs_num,
|
||||
in_dim * ins_num,
|
||||
ins_num * out_dim);
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(batch_fc_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BatchFCGradOpCUDAKernel,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright (c) 2024 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 <string>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
const int CUDA_NUM_THREADS = 1024;
|
||||
static inline int GET_BLOCKS(const int N) {
|
||||
return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void add_bias_kernel(
|
||||
T* data, int slot_pairs_num, int ins_num, int out_dim, const T* bias) {
|
||||
CUDA_KERNEL_LOOP(idx, slot_pairs_num * ins_num * out_dim) {
|
||||
int block_len = ins_num * out_dim;
|
||||
int slot_index = idx / block_len;
|
||||
int out_dim_index = (idx % block_len) % out_dim;
|
||||
T temp = data[idx] + bias[slot_index * out_dim + out_dim_index];
|
||||
data[idx] = temp;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void add_bias(gpuStream_t stream,
|
||||
T* data,
|
||||
int slot_pairs_num,
|
||||
int ins_num,
|
||||
int out_dim,
|
||||
const T* bias) {
|
||||
add_bias_kernel<<<GET_BLOCKS(slot_pairs_num * ins_num * out_dim),
|
||||
CUDA_NUM_THREADS,
|
||||
0,
|
||||
stream>>>(data, slot_pairs_num, ins_num, out_dim, bias);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BatchFCCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input_in,
|
||||
const DenseTensor& w_in,
|
||||
const DenseTensor& bias_in,
|
||||
DenseTensor* out) {
|
||||
// X.dim = slot_pairs_num * ins_num * in_dim
|
||||
// W.dim = slot_pairs_num * in_dim * out_dim
|
||||
// b.dim = slot_pairs_num * out_dim
|
||||
// output.dim = slot_pairs_num * ins_num * out_dim
|
||||
auto* input = &input_in;
|
||||
auto* w = &w_in;
|
||||
auto* bias = &bias_in;
|
||||
auto* output = out;
|
||||
auto input_dims = input->dims();
|
||||
auto w_dims = w->dims();
|
||||
auto slot_pairs_num = input_dims[0];
|
||||
auto ins_num = input_dims[1];
|
||||
auto in_dim = input_dims[2];
|
||||
auto out_dim = w_dims[2];
|
||||
|
||||
// get data ptr
|
||||
const T* in_data = input->data<T>();
|
||||
const T* w_data = w->data<T>();
|
||||
const T* bias_data = bias->data<T>();
|
||||
|
||||
output->Resize({slot_pairs_num, ins_num, out_dim});
|
||||
T* out_data = dev_ctx.template Alloc<T>(output);
|
||||
// initialize
|
||||
auto out_eigen = EigenVector<T>::Flatten(*output);
|
||||
auto& place = *dev_ctx.eigen_device();
|
||||
out_eigen.device(place) = out_eigen.constant(static_cast<T>(0));
|
||||
|
||||
CBLAS_TRANSPOSE transA = CblasNoTrans;
|
||||
CBLAS_TRANSPOSE transB = CblasNoTrans;
|
||||
|
||||
T alpha = 1;
|
||||
T beta = 0;
|
||||
int64_t strideA = ins_num * in_dim;
|
||||
int64_t strideB = in_dim * out_dim;
|
||||
|
||||
auto blas = funcs::GetBlas<GPUContext, T>(dev_ctx);
|
||||
blas.BatchedGEMM(transA,
|
||||
transB,
|
||||
ins_num,
|
||||
out_dim,
|
||||
in_dim,
|
||||
alpha,
|
||||
in_data,
|
||||
w_data,
|
||||
beta,
|
||||
out_data,
|
||||
slot_pairs_num,
|
||||
strideA,
|
||||
strideB);
|
||||
add_bias<T>(
|
||||
dev_ctx.stream(), out_data, slot_pairs_num, ins_num, out_dim, bias_data);
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
batch_fc, GPU, ALL_LAYOUT, phi::BatchFCCUDAKernel, float, double) {}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,63 @@
|
||||
// Copyright (c) 2022 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/bce_loss_grad_kernel.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/common/hostdevice.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T>
|
||||
struct BCELossGradFunctor {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
MT one = static_cast<MT>(1.0f);
|
||||
MT eps = static_cast<MT>(1e-12);
|
||||
|
||||
HOSTDEVICE inline T operator()(const T x, const T label, const T dout) const {
|
||||
MT x_mt = static_cast<MT>(x);
|
||||
MT term1 = max((one - x_mt) * x_mt, eps);
|
||||
return static_cast<T>(static_cast<MT>(dout) *
|
||||
(x_mt - static_cast<MT>(label)) / term1);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BCELossGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const DenseTensor& label,
|
||||
const DenseTensor& out_grad,
|
||||
DenseTensor* input_grad) {
|
||||
dev_ctx.template Alloc<T>(input_grad);
|
||||
std::vector<const DenseTensor*> ins = {&input, &label, &out_grad};
|
||||
std::vector<DenseTensor*> outs = {input_grad};
|
||||
auto functor = BCELossGradFunctor<T>();
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(bce_loss_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BCELossGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,71 @@
|
||||
// Copyright (c) 2022 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/bce_loss_kernel.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/common/hostdevice.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T>
|
||||
struct BCELossFunctor {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
MT zero = static_cast<MT>(0);
|
||||
MT one = static_cast<MT>(1.0f);
|
||||
MT neg_100 = static_cast<MT>(-100.);
|
||||
|
||||
HOSTDEVICE inline T operator()(const T x, const T label) const {
|
||||
MT x_mt = static_cast<MT>(x);
|
||||
MT label_mt = static_cast<MT>(label);
|
||||
|
||||
PADDLE_ENFORCE(
|
||||
(x_mt >= zero) && (x_mt <= one),
|
||||
"Input is expected to be within the interval [0, 1], but received %f.",
|
||||
x_mt);
|
||||
|
||||
MT term1 = max(kps::details::Log(x_mt), neg_100);
|
||||
MT term2 = max(kps::details::Log(one - x_mt), neg_100);
|
||||
return static_cast<T>((label_mt - one) * term2 - label_mt * term1);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BCELossKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const DenseTensor& label,
|
||||
DenseTensor* out) {
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
std::vector<const DenseTensor*> ins = {&input, &label};
|
||||
std::vector<DenseTensor*> outs = {out};
|
||||
auto functor = BCELossFunctor<T>();
|
||||
funcs::ElementwiseKernel<T>(dev_ctx, ins, &outs, functor);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(bce_loss,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BCELossKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2024 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/common/data_type.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/beam_search_decode_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(beam_search_decode,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BeamSearchDecodeOpKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
int,
|
||||
int64_t) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
// Copyright (c) 2024 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/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/beam_search_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(beam_search,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BeamSearchOpKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t) {}
|
||||
@@ -0,0 +1,107 @@
|
||||
// Copyright (c) 2022 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/bernoulli_kernel.h"
|
||||
|
||||
#ifdef __NVCC__
|
||||
#include <curand_kernel.h>
|
||||
#endif
|
||||
#ifdef __HIPCC__
|
||||
#include <hiprand_kernel.h>
|
||||
#endif
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/distribution_helper.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
// 'curand_uniform4/hiprand_uniform4' generate 4 random number each time
|
||||
template <typename T>
|
||||
__global__ void bernoulli_cuda_kernel(
|
||||
size_t size, uint64_t seed, uint64_t offset, const T* x_data, T* out_data) {
|
||||
size_t thread_idx =
|
||||
static_cast<size_t>(blockIdx.x * blockDim.x + threadIdx.x);
|
||||
|
||||
#if defined(__NVCC__)
|
||||
curandStatePhilox4_32_10_t state;
|
||||
curand_init(seed, thread_idx, offset, &state);
|
||||
#else
|
||||
hiprandStatePhilox4_32_10_t state;
|
||||
hiprand_init(seed, thread_idx, offset, &state);
|
||||
#endif
|
||||
|
||||
size_t total_thread =
|
||||
static_cast<size_t>(gridDim.x) * static_cast<size_t>(blockDim.x);
|
||||
for (size_t i = 4 * thread_idx; i < size; i += total_thread * 4) {
|
||||
funcs::uniform_distribution<float> dist;
|
||||
float4 rand = dist(&state);
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
#pragma unroll
|
||||
for (size_t j = 0; j < 4; j++) {
|
||||
size_t idx = i + j;
|
||||
if (idx < size) {
|
||||
MT p = static_cast<MT>(x_data[idx]);
|
||||
PADDLE_ENFORCE(p >= 0 && p <= 1,
|
||||
"The probability should be in [0, 1], but got %f",
|
||||
p);
|
||||
out_data[idx] = static_cast<T>((&rand.x)[j] <= static_cast<MT>(p));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BernoulliKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DenseTensor* out) {
|
||||
const T* x_data = x.data<T>();
|
||||
T* out_data = dev_ctx.template Alloc<T>(out);
|
||||
auto numel = x.numel();
|
||||
|
||||
auto gen_cuda = dev_ctx.GetGenerator();
|
||||
|
||||
auto seed_offset = gen_cuda->IncrementOffset(12);
|
||||
uint64_t seed = seed_offset.first;
|
||||
uint64_t offset = seed_offset.second;
|
||||
|
||||
auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 4);
|
||||
size_t grid_size_64 = gpu_config.GetGridSize();
|
||||
size_t block_size_64 = gpu_config.GetBlockSize();
|
||||
PADDLE_ENFORCE_LE_UINT32_MAX(grid_size_64, "bernoulli grid.x");
|
||||
PADDLE_ENFORCE_LE_UINT32_MAX(block_size_64, "bernoulli block.x");
|
||||
uint32_t grid_size = static_cast<uint32_t>(grid_size_64);
|
||||
uint32_t block_size = static_cast<uint32_t>(block_size_64);
|
||||
|
||||
bernoulli_cuda_kernel<<<grid_size, block_size, 0, dev_ctx.stream()>>>(
|
||||
numel, seed, offset, x_data, out_data);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(bernoulli,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BernoulliKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,21 @@
|
||||
// Copyright (c) 2022 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/bilinear_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/bilinear_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
bilinear_grad, GPU, ALL_LAYOUT, phi::BilinearGradKernel, float, double) {}
|
||||
@@ -0,0 +1,21 @@
|
||||
// Copyright (c) 2022 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/bilinear_kernel.h"
|
||||
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/bilinear_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
bilinear, GPU, ALL_LAYOUT, phi::BilinearKernel, float, double) {}
|
||||
@@ -0,0 +1,207 @@
|
||||
// Copyright (c) 2022 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/bincount_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
namespace phi {
|
||||
|
||||
inline int64_t GET_BLOCKS(const int64_t N) {
|
||||
return (N + PADDLE_CUDA_NUM_THREADS - 1) / PADDLE_CUDA_NUM_THREADS;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void KernelReduceMinMax(const T* input,
|
||||
int64_t numel,
|
||||
T* min_out,
|
||||
T* max_out) {
|
||||
__shared__ T smin[PADDLE_CUDA_NUM_THREADS];
|
||||
__shared__ T smax[PADDLE_CUDA_NUM_THREADS];
|
||||
int tid = threadIdx.x;
|
||||
int64_t global_thread_id =
|
||||
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
|
||||
|
||||
T local_min = std::numeric_limits<T>::max();
|
||||
T local_max = std::numeric_limits<T>::lowest();
|
||||
|
||||
for (int64_t i = global_thread_id; i < numel; i += stride) {
|
||||
T val = input[i];
|
||||
local_min = min(local_min, val);
|
||||
local_max = max(local_max, val);
|
||||
}
|
||||
|
||||
smin[tid] = local_min;
|
||||
smax[tid] = local_max;
|
||||
__syncthreads();
|
||||
|
||||
for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
|
||||
if (tid < offset) {
|
||||
smin[tid] = min(smin[tid], smin[tid + offset]);
|
||||
smax[tid] = max(smax[tid], smax[tid + offset]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (tid == 0) {
|
||||
CudaAtomicMin(min_out, smin[0]);
|
||||
CudaAtomicMax(max_out, smax[0]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename InputT, typename OutT>
|
||||
__global__ void KernelBincount(const InputT* input,
|
||||
const int64_t total_elements,
|
||||
const bool has_weights,
|
||||
const T* weights,
|
||||
OutT* output) {
|
||||
int64_t global_tid =
|
||||
static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
||||
int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
|
||||
for (int64_t i = global_tid; i < total_elements; i += stride) {
|
||||
InputT index = input[i];
|
||||
if (!has_weights) {
|
||||
CudaAtomicAdd(&output[index], 1L);
|
||||
} else {
|
||||
CudaAtomicAdd(&output[index], static_cast<OutT>(weights[i]));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Context, typename T, typename InputT>
|
||||
void BincountCUDAInner(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const optional<DenseTensor>& weights,
|
||||
int64_t minlength,
|
||||
DenseTensor* out) {
|
||||
const DenseTensor* input = &x;
|
||||
DenseTensor* output = out;
|
||||
const InputT* input_data = input->data<InputT>();
|
||||
|
||||
int64_t input_numel = static_cast<int64_t>(input->numel());
|
||||
|
||||
if (input_data == nullptr) {
|
||||
DDim out_dim{minlength};
|
||||
output->Resize(out_dim);
|
||||
Full<int64_t, Context>(dev_ctx, output->dims(), 0, output);
|
||||
return;
|
||||
}
|
||||
|
||||
DenseTensor input_min_max_cpu;
|
||||
input_min_max_cpu.Resize({2});
|
||||
auto* input_min_max_cpu_data =
|
||||
dev_ctx.template HostAlloc<InputT>(&input_min_max_cpu);
|
||||
input_min_max_cpu.data<InputT>()[0] = std::numeric_limits<InputT>::max();
|
||||
input_min_max_cpu.data<InputT>()[1] = std::numeric_limits<InputT>::lowest();
|
||||
|
||||
DenseTensor input_min_max_t;
|
||||
input_min_max_t.Resize({2});
|
||||
auto* input_min_max_data = dev_ctx.template Alloc<InputT>(&input_min_max_t);
|
||||
|
||||
Copy(dev_ctx, input_min_max_cpu, dev_ctx.GetPlace(), true, &input_min_max_t);
|
||||
|
||||
int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
int64_t num_blocks = std::min(GET_BLOCKS(input_numel), max_grid_x);
|
||||
KernelReduceMinMax<InputT>
|
||||
<<<num_blocks, PADDLE_CUDA_NUM_THREADS, 0, dev_ctx.stream()>>>(
|
||||
input_data, input_numel, input_min_max_data, input_min_max_data + 1);
|
||||
|
||||
Copy(dev_ctx, input_min_max_t, CPUPlace(), true, &input_min_max_cpu);
|
||||
|
||||
InputT input_min = input_min_max_cpu.data<InputT>()[0];
|
||||
|
||||
PADDLE_ENFORCE_GE(
|
||||
input_min,
|
||||
static_cast<InputT>(0),
|
||||
common::errors::InvalidArgument(
|
||||
"The elements in input tensor must be non-negative ints"));
|
||||
|
||||
int64_t output_size =
|
||||
static_cast<int64_t>(input_min_max_cpu.data<InputT>()[1]) + 1L;
|
||||
|
||||
output_size = std::max(output_size, minlength);
|
||||
DDim out_dim{output_size};
|
||||
output->Resize(out_dim);
|
||||
|
||||
bool has_weights = weights.is_initialized();
|
||||
|
||||
const T* weights_data = has_weights ? weights->data<T>() : nullptr;
|
||||
auto stream = dev_ctx.stream();
|
||||
|
||||
if (!has_weights) {
|
||||
int64_t* output_data = dev_ctx.template Alloc<int64_t>(output);
|
||||
funcs::SetConstant<Context, int64_t>()(
|
||||
dev_ctx, output, static_cast<int64_t>(0));
|
||||
|
||||
KernelBincount<T, InputT, int64_t>
|
||||
<<<num_blocks, PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
|
||||
input_data, input_numel, has_weights, weights_data, output_data);
|
||||
} else {
|
||||
if (weights->dtype() == DataType::FLOAT32) {
|
||||
float* output_data = dev_ctx.template Alloc<float>(output);
|
||||
funcs::SetConstant<Context, float>()(
|
||||
dev_ctx, output, static_cast<float>(0));
|
||||
|
||||
KernelBincount<T, InputT, float>
|
||||
<<<num_blocks, PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
|
||||
input_data, input_numel, has_weights, weights_data, output_data);
|
||||
} else {
|
||||
double* output_data = dev_ctx.template Alloc<double>(output);
|
||||
funcs::SetConstant<Context, double>()(
|
||||
dev_ctx, output, static_cast<double>(0));
|
||||
KernelBincount<T, InputT, double>
|
||||
<<<num_blocks, PADDLE_CUDA_NUM_THREADS, 0, stream>>>(
|
||||
input_data, input_numel, has_weights, weights_data, output_data);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BincountKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const optional<DenseTensor>& weights,
|
||||
const Scalar& minlength,
|
||||
DenseTensor* out) {
|
||||
int64_t int_minlength = minlength.to<int64_t>();
|
||||
PADDLE_ENFORCE_GE(int_minlength,
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"The minlength should be greater than or equal to 0."
|
||||
"But received minlength is %d",
|
||||
int_minlength));
|
||||
|
||||
if (x.dtype() == DataType::INT32) {
|
||||
BincountCUDAInner<Context, T, int>(dev_ctx, x, weights, int_minlength, out);
|
||||
} else if (x.dtype() == DataType::INT64) {
|
||||
BincountCUDAInner<Context, T, int64_t>(
|
||||
dev_ctx, x, weights, int_minlength, out);
|
||||
}
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(bincount,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BincountKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
/* Copyright (c) 2023 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. */
|
||||
|
||||
#ifdef __NVCC__
|
||||
#include <curand_kernel.h>
|
||||
#endif
|
||||
#ifdef __HIPCC__
|
||||
#include <hiprand_kernel.h>
|
||||
#endif
|
||||
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/binomial_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/for_range.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
__device__ __constant__ float kTailValues[] = {0.0810614667953272,
|
||||
0.0413406959554092,
|
||||
0.0276779256849983,
|
||||
0.02079067210376509,
|
||||
0.0166446911898211,
|
||||
0.0138761288230707,
|
||||
0.0118967099458917,
|
||||
0.0104112652619720,
|
||||
0.00925546218271273,
|
||||
0.00833056343336287};
|
||||
|
||||
template <typename T>
|
||||
__device__ T stirling_approx_tail(int64_t k) {
|
||||
if (k <= 9) {
|
||||
return static_cast<T>(kTailValues[static_cast<size_t>(k)]);
|
||||
}
|
||||
T kp1sq = (k + 1) * (k + 1);
|
||||
return (1.0 / 12 - (1.0 / 360 - 1.0 / 1260 / kp1sq) / kp1sq) / (k + 1);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ int64_t btrs(
|
||||
const T n, const T p, int64_t idx, unsigned int seed, unsigned int offset) {
|
||||
int64_t k;
|
||||
T U, V, us;
|
||||
|
||||
#ifdef __NVCC__
|
||||
curandStatePhilox4_32_10_t state;
|
||||
curand_init(seed, idx, offset, &state);
|
||||
#elif __HIPCC__
|
||||
hiprandStatePhilox4_32_10_t state;
|
||||
hiprand_init(seed, idx, offset, &state);
|
||||
#endif
|
||||
|
||||
const T stddev = std::sqrt(n * p * (1 - p));
|
||||
|
||||
const T b = 1.15 + 2.53 * stddev;
|
||||
const T a = -0.0873 + 0.0248 * b + 0.01 * p;
|
||||
const T c = n * p + 0.5;
|
||||
const T v_r = 0.92 - 4.2 / b;
|
||||
const T r = p / (1 - p);
|
||||
|
||||
const T alpha = (2.83 + 5.1 / b) * stddev;
|
||||
const T m = std::floor((n + 1) * p);
|
||||
|
||||
while (1) {
|
||||
#ifdef __NVCC__
|
||||
U = static_cast<T>(curand_uniform(&state)) - 0.5;
|
||||
V = static_cast<T>(curand_uniform(&state));
|
||||
#elif __HIPCC__
|
||||
U = static_cast<T>(hiprand_uniform(&state)) - 0.5;
|
||||
V = static_cast<T>(hiprand_uniform(&state));
|
||||
#endif
|
||||
|
||||
us = 0.5 - std::abs(U);
|
||||
k = static_cast<int64_t>(std::floor((2 * a / us + b) * U + c));
|
||||
|
||||
if (k < 0 || k > n) {
|
||||
continue;
|
||||
}
|
||||
if (us >= 0.07 && V <= v_r) {
|
||||
return k;
|
||||
}
|
||||
|
||||
V = std::log(V * alpha / (a / (us * us) + b));
|
||||
T upperbound =
|
||||
((m + 0.5) * std::log((m + 1) / (r * (n - m + 1))) +
|
||||
(n + 1) * std::log((n - m + 1) / (n - k + 1)) +
|
||||
(k + 0.5) * std::log(r * (n - k + 1) / (k + 1)) +
|
||||
stirling_approx_tail<T>(m) + stirling_approx_tail<T>(n - m) -
|
||||
stirling_approx_tail<T>(k) - stirling_approx_tail<T>(n - k));
|
||||
|
||||
if (V <= upperbound) {
|
||||
return k;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ int64_t binomial_inversion(
|
||||
const T n, const T p, int64_t idx, unsigned int seed, unsigned int offset) {
|
||||
T unif;
|
||||
T geom_sum = 0.0;
|
||||
int64_t num_geom = 0;
|
||||
T logprob = std::log1p(-p);
|
||||
|
||||
#ifdef __NVCC__
|
||||
curandStatePhilox4_32_10_t state;
|
||||
curand_init(seed, idx, offset, &state);
|
||||
#elif __HIPCC__
|
||||
hiprandStatePhilox4_32_10_t state;
|
||||
hiprand_init(seed, idx, offset, &state);
|
||||
#endif
|
||||
|
||||
while (1) {
|
||||
#ifdef __NVCC__
|
||||
unif = static_cast<T>(curand_uniform(&state));
|
||||
#elif __HIPCC__
|
||||
unif = static_cast<T>(hiprand_uniform(&state));
|
||||
#endif
|
||||
T geom = std::ceil(std::log(unif) / logprob);
|
||||
geom_sum += geom;
|
||||
if (geom_sum > n) {
|
||||
break;
|
||||
}
|
||||
num_geom = num_geom + 1;
|
||||
}
|
||||
return num_geom;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void BinomialSampling(const T* n,
|
||||
const T* p,
|
||||
int64_t* out,
|
||||
const int N,
|
||||
unsigned int seed,
|
||||
unsigned int offset) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
CUDA_KERNEL_LOOP_TYPE(idx, N, int64_t) {
|
||||
MT nt = static_cast<MT>(n[idx]);
|
||||
MT pt = static_cast<MT>(p[idx]);
|
||||
if (nt <= 0.0 || pt <= 0.0) {
|
||||
out[idx] = 0;
|
||||
} else if (pt >= 1.0) {
|
||||
out[idx] = static_cast<int64_t>(nt);
|
||||
} else if (pt <= 0.5) {
|
||||
if (nt * pt >= 10.0) {
|
||||
out[idx] = btrs<MT>(nt, pt, idx, seed, offset);
|
||||
} else {
|
||||
out[idx] = binomial_inversion<MT>(nt, pt, idx, seed, offset);
|
||||
}
|
||||
} else {
|
||||
MT qprob = 1.0 - pt;
|
||||
if (nt * qprob >= 10.0) {
|
||||
out[idx] =
|
||||
static_cast<int64_t>(nt) - btrs<MT>(nt, qprob, idx, seed, offset);
|
||||
} else {
|
||||
out[idx] = static_cast<int64_t>(nt) -
|
||||
binomial_inversion<MT>(nt, qprob, idx, seed, offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BinomialKernel(const Context& dev_ctx,
|
||||
const DenseTensor& count,
|
||||
const DenseTensor& prob,
|
||||
DenseTensor* out) {
|
||||
const T* count_data = count.data<T>();
|
||||
const T* prob_data = prob.data<T>();
|
||||
int64_t* out_data = dev_ctx.template Alloc<int64_t>(out);
|
||||
// TODO(large-tensor): downstream functors may still use int; guard until
|
||||
// upgraded.
|
||||
int64_t size = count.numel();
|
||||
|
||||
const int kMaxBlockDim = 256;
|
||||
|
||||
int block_size = std::min(kMaxBlockDim, dev_ctx.GetMaxThreadsPerBlock());
|
||||
dim3 dim_block(block_size);
|
||||
int64_t grid_64 = (size + block_size - 1) / block_size;
|
||||
PADDLE_ENFORCE_LE_UINT32_MAX(grid_64, "binomial grid.x");
|
||||
dim3 dim_grid(static_cast<uint32_t>(grid_64));
|
||||
backends::gpu::LimitGridDim(dev_ctx, &dim_grid);
|
||||
|
||||
auto gen_cuda = dev_ctx.GetGenerator();
|
||||
auto seed_offset = gen_cuda->IncrementOffset(20);
|
||||
uint64_t seed = seed_offset.first;
|
||||
uint64_t offset = seed_offset.second;
|
||||
BinomialSampling<T><<<dim_grid, dim_block>>>(
|
||||
count_data, prob_data, out_data, size, seed, offset);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(binomial,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BinomialKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2022 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/bmm_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/bmm_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(bmm_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BmmGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2022 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/bmm_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/bmm_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(bmm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BmmKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,79 @@
|
||||
// Copyright (c) 2024 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 <algorithm>
|
||||
|
||||
#include "paddle/common/hostdevice.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/lod_utils.h"
|
||||
#include "paddle/phi/core/mixed_vector.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/gpu/box_clip_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/box_clip_kernel_impl.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int ImInfoSize = 3;
|
||||
|
||||
template <typename T, int BlockSize>
|
||||
static __global__ void GPUBoxClip(const T *input,
|
||||
const size_t *lod,
|
||||
const size_t width,
|
||||
const T *im_info,
|
||||
T *output) {
|
||||
T im_w = round(im_info[blockIdx.x * ImInfoSize + 1] /
|
||||
im_info[blockIdx.x * ImInfoSize + 2]);
|
||||
T im_h = round(im_info[blockIdx.x * ImInfoSize] /
|
||||
im_info[blockIdx.x * ImInfoSize + 2]);
|
||||
for (size_t i = threadIdx.x;
|
||||
i < (lod[blockIdx.x + 1] - lod[blockIdx.x]) * width;
|
||||
i += BlockSize) {
|
||||
size_t idx = lod[blockIdx.x] * width + i;
|
||||
T im_size = (idx % 2 == 0) ? im_w : im_h;
|
||||
output[idx] = max(min(input[idx], im_size - 1), T(0.));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void GPUBoxClipKernel(const Context &dev_ctx,
|
||||
const DenseTensor &input,
|
||||
const DenseTensor &im_info,
|
||||
DenseTensor *output) {
|
||||
auto *input_p = &input;
|
||||
auto *im_info_p = &im_info;
|
||||
|
||||
const int64_t num = input_p->dims()[0];
|
||||
const int64_t bbox_width = input_p->numel() / num;
|
||||
auto lod = input_p->lod();
|
||||
LegacyLoD abs_offset_lod = ToAbsOffset(lod);
|
||||
|
||||
auto stream = dev_ctx.stream();
|
||||
const size_t batch_size = lod.back().size() - 1;
|
||||
T *output_data = dev_ctx.template Alloc<T>(output);
|
||||
MixVector<size_t> mix_vector(&abs_offset_lod[0]);
|
||||
GPUBoxClip<T, 512><<<batch_size, 512, 0, stream>>>(
|
||||
input_p->data<T>(),
|
||||
mix_vector.CUDAMutableData(dev_ctx.GetPlace()),
|
||||
bbox_width,
|
||||
im_info_p->data<T>(),
|
||||
output_data);
|
||||
mix_vector.CopyToCPU();
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
box_clip, GPU, ALL_LAYOUT, phi::GPUBoxClipKernel, float, double) {}
|
||||
@@ -0,0 +1,26 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void GPUBoxClipKernel(const Context &dev_ctx,
|
||||
const DenseTensor &input,
|
||||
const DenseTensor &im_info,
|
||||
DenseTensor *output);
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,264 @@
|
||||
// Copyright (c) 2022 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/box_coder_kernel.h"
|
||||
|
||||
#include <thrust/device_vector.h>
|
||||
#include <thrust/host_vector.h>
|
||||
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/box_coder.h"
|
||||
namespace phi {
|
||||
|
||||
template <typename T>
|
||||
__global__ void EncodeCenterSizeKernel(const T *prior_box_data,
|
||||
const T *prior_box_var_data,
|
||||
const T *target_box_data,
|
||||
const int64_t row,
|
||||
const int64_t col,
|
||||
const int64_t len,
|
||||
const bool normalized,
|
||||
const T prior_box_var_size,
|
||||
const float *variance,
|
||||
const int64_t var_size,
|
||||
T *output) {
|
||||
const int64_t idx =
|
||||
threadIdx.x + static_cast<int64_t>(blockIdx.x) * blockDim.x;
|
||||
if (idx < row * col) {
|
||||
const int64_t row_idx = idx / col;
|
||||
const int64_t col_idx = idx % col;
|
||||
T prior_box_width = prior_box_data[col_idx * len + 2] -
|
||||
prior_box_data[col_idx * len] + (normalized == false);
|
||||
T prior_box_height = prior_box_data[col_idx * len + 3] -
|
||||
prior_box_data[col_idx * len + 1] +
|
||||
(normalized == false);
|
||||
T prior_box_center_x = prior_box_data[col_idx * len] + prior_box_width / 2;
|
||||
T prior_box_center_y =
|
||||
prior_box_data[col_idx * len + 1] + prior_box_height / 2;
|
||||
|
||||
T target_box_center_x =
|
||||
(target_box_data[row_idx * len + 2] + target_box_data[row_idx * len]) /
|
||||
2;
|
||||
T target_box_center_y = (target_box_data[row_idx * len + 3] +
|
||||
target_box_data[row_idx * len + 1]) /
|
||||
2;
|
||||
T target_box_width = target_box_data[row_idx * len + 2] -
|
||||
target_box_data[row_idx * len] + (normalized == false);
|
||||
T target_box_height = target_box_data[row_idx * len + 3] -
|
||||
target_box_data[row_idx * len + 1] +
|
||||
(normalized == false);
|
||||
|
||||
output[idx * len] =
|
||||
(target_box_center_x - prior_box_center_x) / prior_box_width;
|
||||
output[idx * len + 1] =
|
||||
(target_box_center_y - prior_box_center_y) / prior_box_height;
|
||||
output[idx * len + 2] = log(fabs(target_box_width / prior_box_width));
|
||||
output[idx * len + 3] = log(fabs(target_box_height / prior_box_height));
|
||||
if (prior_box_var_data) {
|
||||
int64_t prior_var_offset = col_idx * len;
|
||||
output[idx * len] /= prior_box_var_data[prior_var_offset];
|
||||
output[idx * len + 1] /= prior_box_var_data[prior_var_offset + 1];
|
||||
output[idx * len + 2] /= prior_box_var_data[prior_var_offset + 2];
|
||||
output[idx * len + 3] /= prior_box_var_data[prior_var_offset + 3];
|
||||
} else if (var_size == 4) {
|
||||
for (int k = 0; k < 4; ++k) {
|
||||
output[idx * len + k] /= static_cast<T>(variance[k]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void DecodeCenterSizeKernel(const T *prior_box_data,
|
||||
const T *prior_box_var_data,
|
||||
const T *target_box_data,
|
||||
const int64_t row,
|
||||
const int64_t col,
|
||||
const int64_t len,
|
||||
const bool normalized,
|
||||
const T prior_box_var_size,
|
||||
const float *variance,
|
||||
const int64_t var_size,
|
||||
const int axis,
|
||||
T *output) {
|
||||
const int64_t idx =
|
||||
threadIdx.x + static_cast<int64_t>(blockIdx.x) * blockDim.x;
|
||||
int64_t prior_box_offset = 0;
|
||||
if (idx < row * col) {
|
||||
const int64_t col_idx = idx % col;
|
||||
const int64_t row_idx = idx / col;
|
||||
prior_box_offset = axis == 0 ? col_idx * len : row_idx * len;
|
||||
T prior_box_width = prior_box_data[prior_box_offset + 2] -
|
||||
prior_box_data[prior_box_offset] +
|
||||
(normalized == false);
|
||||
T prior_box_height = prior_box_data[prior_box_offset + 3] -
|
||||
prior_box_data[prior_box_offset + 1] +
|
||||
(normalized == false);
|
||||
T prior_box_center_x =
|
||||
prior_box_data[prior_box_offset] + prior_box_width / 2;
|
||||
T prior_box_center_y =
|
||||
prior_box_data[prior_box_offset + 1] + prior_box_height / 2;
|
||||
T target_box_width, target_box_height;
|
||||
T target_box_center_x, target_box_center_y;
|
||||
T box_var_x = T(1), box_var_y = T(1);
|
||||
T box_var_w = T(1), box_var_h = T(1);
|
||||
if (prior_box_var_data) {
|
||||
int64_t prior_var_offset = axis == 0 ? col_idx * len : row_idx * len;
|
||||
box_var_x = prior_box_var_data[prior_var_offset];
|
||||
box_var_y = prior_box_var_data[prior_var_offset + 1];
|
||||
box_var_w = prior_box_var_data[prior_var_offset + 2];
|
||||
box_var_h = prior_box_var_data[prior_var_offset + 3];
|
||||
} else if (var_size == 4) {
|
||||
box_var_x = static_cast<T>(variance[0]);
|
||||
box_var_y = static_cast<T>(variance[1]);
|
||||
box_var_w = static_cast<T>(variance[2]);
|
||||
box_var_h = static_cast<T>(variance[3]);
|
||||
}
|
||||
target_box_width =
|
||||
exp(box_var_w * target_box_data[idx * len + 2]) * prior_box_width;
|
||||
target_box_height =
|
||||
exp(box_var_h * target_box_data[idx * len + 3]) * prior_box_height;
|
||||
target_box_center_x =
|
||||
box_var_x * target_box_data[idx * len] * prior_box_width +
|
||||
prior_box_center_x;
|
||||
target_box_center_y =
|
||||
box_var_y * target_box_data[idx * len + 1] * prior_box_height +
|
||||
prior_box_center_y;
|
||||
|
||||
output[idx * len] = target_box_center_x - target_box_width / 2;
|
||||
output[idx * len + 1] = target_box_center_y - target_box_height / 2;
|
||||
output[idx * len + 2] =
|
||||
target_box_center_x + target_box_width / 2 - (normalized == false);
|
||||
output[idx * len + 3] =
|
||||
target_box_center_y + target_box_height / 2 - (normalized == false);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BoxCoderKernel(const Context &dev_ctx,
|
||||
const DenseTensor &prior_box,
|
||||
const optional<DenseTensor> &prior_box_var,
|
||||
const DenseTensor &target_box,
|
||||
const std::string &code_type_str,
|
||||
bool normalized,
|
||||
int axis,
|
||||
const std::vector<float> &variance,
|
||||
DenseTensor *output_box) {
|
||||
// prior_box and prior_box_var have the same shape, so do not judge
|
||||
// prior_box_var
|
||||
if (prior_box.numel() == 0 || target_box.numel() == 0) {
|
||||
Full<T, Context>(dev_ctx, output_box->dims(), 0, output_box);
|
||||
return;
|
||||
}
|
||||
|
||||
const T *prior_box_data = prior_box.template data<T>();
|
||||
const T *target_box_data = target_box.template data<T>();
|
||||
const T *prior_box_var_data = nullptr;
|
||||
auto prior_box_var_size = 0;
|
||||
if (prior_box_var) {
|
||||
PADDLE_ENFORCE_EQ(variance.empty(),
|
||||
true,
|
||||
common::errors::InvalidArgument(
|
||||
"Input 'PriorBoxVar' and attribute 'variance'"
|
||||
" of BoxCoder operator should not be used at the "
|
||||
"same time."));
|
||||
prior_box_var_data = prior_box_var->data<T>();
|
||||
prior_box_var_size = prior_box_var->dims().size();
|
||||
}
|
||||
if (!(variance.empty())) {
|
||||
PADDLE_ENFORCE_EQ(static_cast<int>(variance.size()),
|
||||
4,
|
||||
common::errors::InvalidArgument(
|
||||
"Size of attribute 'variance' in BoxCoder operator"
|
||||
" should be 4. But received size is %d",
|
||||
variance.size()));
|
||||
}
|
||||
|
||||
if (target_box.lod().size()) {
|
||||
PADDLE_ENFORCE_EQ(target_box.lod().size(),
|
||||
1,
|
||||
common::errors::InvalidArgument(
|
||||
"Input 'TargetBox' of BoxCoder operator only"
|
||||
" supports LoD with one level."));
|
||||
}
|
||||
const int var_size = static_cast<int>(variance.size());
|
||||
auto code_type = funcs::GetBoxCodeType(code_type_str);
|
||||
int64_t row = target_box.dims()[0];
|
||||
int64_t col = prior_box.dims()[0];
|
||||
if (code_type == funcs::BoxCodeType::kDecodeCenterSize) {
|
||||
col = target_box.dims()[1];
|
||||
}
|
||||
int64_t len = prior_box.dims()[1];
|
||||
int block = 512;
|
||||
int64_t grid64 = (row * col + block - 1) / block;
|
||||
PADDLE_ENFORCE_LE_INT_MAX(grid64, "grid");
|
||||
int grid = static_cast<int>(grid64);
|
||||
|
||||
int64_t bytes = var_size * sizeof(float);
|
||||
auto dev_var =
|
||||
memory_utils::Alloc(dev_ctx.GetPlace(),
|
||||
bytes,
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
float *dev_var_data = reinterpret_cast<float *>(dev_var->ptr());
|
||||
auto cplace = CPUPlace();
|
||||
const auto gplace = dev_ctx.GetPlace();
|
||||
const float *stable_variance =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
const_cast<float *>(variance.data()), variance.size());
|
||||
memory_utils::Copy(
|
||||
gplace, dev_var_data, cplace, stable_variance, bytes, dev_ctx.stream());
|
||||
|
||||
output_box->Resize({row, col, len});
|
||||
dev_ctx.template Alloc<T>(output_box);
|
||||
T *output = output_box->data<T>();
|
||||
|
||||
if (code_type == funcs::BoxCodeType::kEncodeCenterSize) {
|
||||
EncodeCenterSizeKernel<T>
|
||||
<<<grid, block, 0, dev_ctx.stream()>>>(prior_box_data,
|
||||
prior_box_var_data,
|
||||
target_box_data,
|
||||
row,
|
||||
col,
|
||||
len,
|
||||
normalized,
|
||||
prior_box_var_size,
|
||||
dev_var_data,
|
||||
var_size,
|
||||
output);
|
||||
} else if (code_type == funcs::BoxCodeType::kDecodeCenterSize) {
|
||||
DecodeCenterSizeKernel<T>
|
||||
<<<grid, block, 0, dev_ctx.stream()>>>(prior_box_data,
|
||||
prior_box_var_data,
|
||||
target_box_data,
|
||||
row,
|
||||
col,
|
||||
len,
|
||||
normalized,
|
||||
prior_box_var_size,
|
||||
dev_var_data,
|
||||
var_size,
|
||||
axis,
|
||||
output);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
box_coder, GPU, ALL_LAYOUT, phi::BoxCoderKernel, float, double) {}
|
||||
@@ -0,0 +1,89 @@
|
||||
// Copyright (c) 2023 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/broadcast_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BroadcastKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
int root,
|
||||
DenseTensor* out) {
|
||||
PADDLE_ENFORCE_GT(x.numel(),
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"Tensor need be broadcast must not empty."));
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
gpuStream_t stream = dev_ctx.stream();
|
||||
auto comm_context =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(
|
||||
comm_context,
|
||||
nullptr,
|
||||
errors::Unavailable("NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
comm_context->Broadcast(out, x, root, stream);
|
||||
out->set_lod(x.lod());
|
||||
#else
|
||||
PADDLE_THROW(
|
||||
errors::PreconditionNotMet("PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
#if NCCL_VERSION_CODE >= 21000
|
||||
PD_REGISTER_KERNEL(broadcast,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BroadcastKernel,
|
||||
float,
|
||||
double,
|
||||
phi::bfloat16,
|
||||
int,
|
||||
bool,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#else
|
||||
PD_REGISTER_KERNEL(broadcast,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BroadcastKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
bool,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
#endif
|
||||
@@ -0,0 +1,116 @@
|
||||
// Copyright (c) 2022 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/broadcast_tensors_grad_kernel.h"
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
||||
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void BroadcastTensorsGradKernel(const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& inputs,
|
||||
const std::vector<const DenseTensor*>& dout,
|
||||
std::vector<DenseTensor*> dx) {
|
||||
(void)inputs;
|
||||
// Find reduce dimensions
|
||||
const auto& in_tensors = dout;
|
||||
auto& out_tensors = dx;
|
||||
|
||||
size_t num_ins = in_tensors.size();
|
||||
|
||||
PADDLE_ENFORCE_GT(
|
||||
num_ins,
|
||||
1,
|
||||
errors::InvalidArgument(
|
||||
"Expected at least 2 input tensors, but only received %d.",
|
||||
in_tensors.size()));
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
num_ins,
|
||||
out_tensors.size(),
|
||||
errors::InvalidArgument(
|
||||
"BroadcastTensorsOp expects equal number of inputs and outputs,"
|
||||
"but received: %d inputs v.s %d outputs",
|
||||
num_ins,
|
||||
out_tensors.size()));
|
||||
|
||||
// For each In-Out tensor pair,
|
||||
// Prepare and apply broadcast dims array
|
||||
for (size_t i = 0; i < num_ins; i++) {
|
||||
auto* input_tensor = in_tensors[i];
|
||||
auto* output_tensor = out_tensors[i];
|
||||
|
||||
const DDim& input_dims = input_tensor->dims();
|
||||
const DDim& output_dims = output_tensor->dims();
|
||||
|
||||
int in_rank = input_dims.size();
|
||||
int out_rank = output_dims.size();
|
||||
|
||||
// Collect reduce_dims
|
||||
// Example:
|
||||
// dX = [1,1,1,1]
|
||||
// dOut = [1,1,1,4]
|
||||
//
|
||||
// reduce_dims = [3] // reduce along the broadcasted axis
|
||||
std::vector<int> reduce_dims_vec;
|
||||
for (int j = 0; j < in_rank; j++) {
|
||||
int out_axis = out_rank - j - 1;
|
||||
int in_axis = in_rank - j - 1;
|
||||
|
||||
if (out_axis < 0 || output_dims[out_axis] != input_dims[in_axis]) {
|
||||
reduce_dims_vec.push_back(in_axis);
|
||||
}
|
||||
}
|
||||
|
||||
bool just_copy = (reduce_dims_vec.size() == 0);
|
||||
dev_ctx.template Alloc<T>(output_tensor);
|
||||
if (just_copy) {
|
||||
// Turns out to be a No-Op, simply copy tensors
|
||||
Copy(dev_ctx, *input_tensor, dev_ctx.GetPlace(), false, output_tensor);
|
||||
} else {
|
||||
// reduce_sum implementation on CUDA
|
||||
SumKernel<T, Context>(dev_ctx,
|
||||
*input_tensor,
|
||||
reduce_dims_vec,
|
||||
output_tensor->dtype(),
|
||||
false,
|
||||
output_tensor);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(broadcast_tensors_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BroadcastTensorsGradKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,33 @@
|
||||
// Copyright (c) 2022 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/broadcast_tensors_kernel.h"
|
||||
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/broadcast_tensors_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(broadcast_tensors,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::BroadcastTensorsKernel,
|
||||
bool,
|
||||
int,
|
||||
int64_t,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,124 @@
|
||||
// Copyright (c) 2024 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 <vector>
|
||||
|
||||
#include "paddle/phi/api/include/tensor.h"
|
||||
#include "paddle/phi/backends/context_pool.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/c_concat_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
#if defined(PADDLE_WITH_FLAGCX)
|
||||
#include "paddle/phi/core/distributed/flagcx_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CConcatKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x_in,
|
||||
int rank,
|
||||
int nranks,
|
||||
int ring_id UNUSED,
|
||||
bool use_calc_stream UNUSED,
|
||||
bool use_model_parallel UNUSED,
|
||||
DenseTensor* out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
auto x = &x_in;
|
||||
PADDLE_ENFORCE_GE(rank,
|
||||
0,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of rank (%d) for c_concat must be "
|
||||
"greater than or equal to 0.",
|
||||
rank));
|
||||
PADDLE_ENFORCE_GE(nranks,
|
||||
2,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of nranks (%d) for c_concat must be "
|
||||
"greater than or equal to 2.",
|
||||
nranks));
|
||||
PADDLE_ENFORCE_LT(rank,
|
||||
nranks,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of rank (%d) for c_concat must be "
|
||||
"less than that of nranks (%d).",
|
||||
rank,
|
||||
nranks));
|
||||
|
||||
DenseTensor temp_out;
|
||||
DDim temp_out_dims = x->dims();
|
||||
temp_out_dims[0] *= nranks;
|
||||
temp_out.Resize(temp_out_dims);
|
||||
dev_ctx.template Alloc<T>(&temp_out);
|
||||
|
||||
gpuStream_t stream = nullptr;
|
||||
|
||||
#if defined(PADDLE_WITH_FLAGCX) && defined(PADDLE_KERNEL_WITH_FLAGCX)
|
||||
distributed::FlagcxCommContext* comm_ctx = nullptr;
|
||||
comm_ctx =
|
||||
static_cast<distributed::FlagcxCommContext*>(dev_ctx.GetCommContext());
|
||||
#else
|
||||
distributed::NCCLCommContext* comm_ctx = nullptr;
|
||||
comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
#endif
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
stream = dev_ctx.stream();
|
||||
#if defined(PADDLE_WITH_FLAGCX) && defined(PADDLE_KERNEL_WITH_FLAGCX)
|
||||
comm_ctx->AllGather(&temp_out, *x, reinterpret_cast<flagcxStream_t>(&stream));
|
||||
#else
|
||||
comm_ctx->AllGather(&temp_out, *x, stream);
|
||||
#endif
|
||||
|
||||
std::vector<DenseTensor> inputs;
|
||||
int axis = x->dims().size() - 1;
|
||||
auto out_dims = x->dims();
|
||||
out_dims[out_dims.size() - 1] *= nranks;
|
||||
int64_t rows_per_tensor = x->dims()[0];
|
||||
int64_t offset = 0;
|
||||
for (int i = 0; i < nranks; i++) {
|
||||
DenseTensor temp = temp_out.Slice(offset, offset + rows_per_tensor);
|
||||
inputs.emplace_back(temp);
|
||||
offset += rows_per_tensor;
|
||||
}
|
||||
|
||||
funcs::ConcatFunctor<GPUContext, T> functor;
|
||||
out->Resize(out_dims);
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
functor(dev_ctx, inputs, axis, out);
|
||||
#else
|
||||
PADDLE_THROW(common::errors::PreconditionNotMet(
|
||||
"PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_concat,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CConcatKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,152 @@
|
||||
// Copyright (c) 2023 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/c_embedding_kernel.h"
|
||||
|
||||
#include "glog/logging.h"
|
||||
|
||||
#include "paddle/common/flags.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/embedding_grad.h"
|
||||
|
||||
COMMON_DECLARE_int64(embedding_deterministic);
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int NumBlocks(const int64_t N) {
|
||||
return static_cast<int>(std::min<int64_t>(
|
||||
(N + kNumCUDAThreads - 1) / kNumCUDAThreads, kNumMaximumNumBlocks));
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CEmbeddingGrad(T* table,
|
||||
const T* output,
|
||||
const IndexT* ids,
|
||||
const int64_t rows,
|
||||
const int64_t columns,
|
||||
const int64_t N,
|
||||
const int64_t start_idx,
|
||||
const int64_t end_idx,
|
||||
const int64_t limit) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, limit, int64_t) {
|
||||
int64_t row = i / columns;
|
||||
int64_t col = i % columns;
|
||||
auto id = ids[row];
|
||||
if (id >= start_idx && id < end_idx) {
|
||||
auto real_idx = id - start_idx;
|
||||
CudaAtomicAdd(&table[real_idx * columns + col], output[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CEmbeddingGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& w,
|
||||
const DenseTensor& ids,
|
||||
const DenseTensor& out_grad,
|
||||
int64_t start_index,
|
||||
DenseTensor* w_grad) {
|
||||
int64_t N = w_grad->dims()[0];
|
||||
int64_t D = w_grad->dims()[1];
|
||||
int64_t K = ids.numel();
|
||||
|
||||
auto limit = K * D;
|
||||
auto blocks = NumBlocks(limit);
|
||||
int threads = kNumCUDAThreads;
|
||||
|
||||
const T* d_output = out_grad.data<T>();
|
||||
T* d_table = dev_ctx.template Alloc<T>(w_grad);
|
||||
|
||||
auto t = EigenVector<T>::Flatten(*w_grad);
|
||||
t.device(*dev_ctx.eigen_device()) = t.constant(static_cast<T>(0));
|
||||
|
||||
const auto& index_type = ids.dtype();
|
||||
if (FLAGS_embedding_deterministic == 1) {
|
||||
if (index_type == DataType::INT32) {
|
||||
funcs::LaunchEmbeddingGradDeterministicKernel<T, int32_t>(
|
||||
dev_ctx,
|
||||
ids.data<int32_t>(),
|
||||
d_output,
|
||||
d_table,
|
||||
N,
|
||||
D,
|
||||
K,
|
||||
start_index);
|
||||
return;
|
||||
} else if (index_type == DataType::INT64) {
|
||||
funcs::LaunchEmbeddingGradDeterministicKernel<T, int64_t>(
|
||||
dev_ctx,
|
||||
ids.data<int64_t>(),
|
||||
d_output,
|
||||
d_table,
|
||||
N,
|
||||
D,
|
||||
K,
|
||||
start_index);
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
if (FLAGS_embedding_deterministic > 1) {
|
||||
VLOG(2) << "Run grad kernel of embedding with single thread.";
|
||||
blocks = 1;
|
||||
}
|
||||
const int64_t end_idx = start_index + N;
|
||||
if (index_type == DataType::INT32) {
|
||||
CEmbeddingGrad<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(d_table,
|
||||
d_output,
|
||||
ids.data<int32_t>(),
|
||||
K,
|
||||
D,
|
||||
N,
|
||||
start_index,
|
||||
end_idx,
|
||||
limit);
|
||||
return;
|
||||
} else if (index_type == DataType::INT64) {
|
||||
CEmbeddingGrad<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(d_table,
|
||||
d_output,
|
||||
ids.data<int64_t>(),
|
||||
K,
|
||||
D,
|
||||
N,
|
||||
start_index,
|
||||
end_idx,
|
||||
limit);
|
||||
return;
|
||||
}
|
||||
}
|
||||
PADDLE_THROW(common::errors::InvalidArgument(
|
||||
"The data type of Input(Ids) must be int32 or int64."));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_embedding_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CEmbeddingGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,125 @@
|
||||
// Copyright (c) 2023 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/c_embedding_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int NumBlocks(const int64_t N) {
|
||||
return static_cast<int>(std::min<int64_t>(
|
||||
(N + kNumCUDAThreads - 1) / kNumCUDAThreads, kNumMaximumNumBlocks));
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CEmbedding(T* out,
|
||||
const T* table,
|
||||
const IndexT* ids,
|
||||
const int64_t rows,
|
||||
const int64_t columns,
|
||||
const int64_t N,
|
||||
const int64_t start_idx,
|
||||
const int64_t end_idx,
|
||||
const int64_t limit,
|
||||
const int64_t vocab_size) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, limit, int64_t) {
|
||||
int64_t row = i / columns;
|
||||
int64_t col = i % columns;
|
||||
auto id = ids[row];
|
||||
|
||||
PADDLE_ENFORCE(
|
||||
id >= 0 && (vocab_size < 0 || id < vocab_size),
|
||||
"The index is out of bounds, "
|
||||
"please check whether the dimensions of index and "
|
||||
"input meet the requirements. It should "
|
||||
"be less than [%d] and greater than or equal to 0, but received [%d]",
|
||||
vocab_size,
|
||||
id);
|
||||
if (id >= start_idx && id < end_idx) {
|
||||
auto real_idx = id - start_idx;
|
||||
out[i] = table[real_idx * columns + col];
|
||||
} else {
|
||||
out[i] = static_cast<T>(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CEmbeddingKernel(const Context& dev_ctx,
|
||||
const DenseTensor& w,
|
||||
const DenseTensor& ids,
|
||||
int64_t start_index,
|
||||
int64_t vocab_size,
|
||||
DenseTensor* out) {
|
||||
int64_t N = w.dims()[0];
|
||||
int64_t D = w.dims()[1];
|
||||
int64_t K = ids.numel();
|
||||
|
||||
const int64_t end_idx = start_index + N;
|
||||
|
||||
auto* table = w.data<T>();
|
||||
auto* output = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
auto limit = K * D;
|
||||
auto blocks = NumBlocks(limit);
|
||||
int threads = kNumCUDAThreads;
|
||||
|
||||
const auto& index_type = ids.dtype();
|
||||
if (index_type == DataType::INT32) {
|
||||
CEmbedding<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(output,
|
||||
table,
|
||||
ids.data<int32_t>(),
|
||||
K,
|
||||
D,
|
||||
N,
|
||||
start_index,
|
||||
end_idx,
|
||||
limit,
|
||||
vocab_size);
|
||||
|
||||
} else if (index_type == DataType::INT64) {
|
||||
CEmbedding<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(output,
|
||||
table,
|
||||
ids.data<int64_t>(),
|
||||
K,
|
||||
D,
|
||||
N,
|
||||
start_index,
|
||||
end_idx,
|
||||
limit,
|
||||
vocab_size);
|
||||
} else {
|
||||
PADDLE_THROW(common::errors::Unavailable(
|
||||
"GPU c_embedding ids only support int32 or int64."));
|
||||
}
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_embedding,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CEmbeddingKernel,
|
||||
float,
|
||||
double,
|
||||
phi::bfloat16,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
/* Copyright (c) 2023 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/c_identity_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/c_identity_kernel_impl.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
PD_REGISTER_KERNEL(c_identity,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CIdentityKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,125 @@
|
||||
// Copyright (c) 2024 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/gpu/c_scatter_kernel.h"
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/phi/core/distributed/comm_context_manager.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CScatterOpCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
int ring_id,
|
||||
int root,
|
||||
int nranks,
|
||||
bool use_calc_stream,
|
||||
DenseTensor* out) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
auto x = &input;
|
||||
int64_t numel = x->numel();
|
||||
ncclDataType_t dtype = ToNCCLDataType(x->dtype());
|
||||
|
||||
int root_id = root;
|
||||
auto place = dev_ctx.GetPlace();
|
||||
gpuStream_t stream = nullptr;
|
||||
distributed::NCCLCommContext* comm_ctx = nullptr;
|
||||
PADDLE_ENFORCE_GE(
|
||||
root_id,
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"The root_id (%d) for c_scatter_op must be non-negative.", root_id));
|
||||
PADDLE_ENFORCE_GE(
|
||||
ring_id,
|
||||
0,
|
||||
common::errors::InvalidArgument(
|
||||
"The ring_id (%d) for c_scatter_op must be non-negative.", ring_id));
|
||||
|
||||
comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
PADDLE_ENFORCE_EQ(nranks,
|
||||
comm_ctx->GetSize(),
|
||||
common::errors::InvalidArgument(
|
||||
"The number of ranks (%d) you set of must "
|
||||
"be equal to comm_ctx->GetSize() (%d).",
|
||||
nranks,
|
||||
comm_ctx->GetSize()));
|
||||
|
||||
stream = comm_ctx->GetStream();
|
||||
VLOG(3) << "new comm_context_manager has ring_id " << ring_id;
|
||||
|
||||
if (use_calc_stream) {
|
||||
// should ExecutionContext for calc stream.
|
||||
stream = dev_ctx.stream();
|
||||
}
|
||||
|
||||
DDim x_dims = x->dims();
|
||||
DDim out_dims(x_dims);
|
||||
DenseTensor temp;
|
||||
temp.Resize(out_dims);
|
||||
auto out_ptr = dev_ctx.template Alloc<T>(&temp);
|
||||
|
||||
if (root_id == comm_ctx->GetRank()) {
|
||||
comm_ctx->Broadcast(const_cast<DenseTensor*>(x), *x, root_id, stream);
|
||||
Copy(dev_ctx,
|
||||
*static_cast<const DenseTensor*>(x),
|
||||
place,
|
||||
false,
|
||||
static_cast<DenseTensor*>(&temp));
|
||||
} else {
|
||||
comm_ctx->Broadcast(&temp, temp, root_id, stream);
|
||||
}
|
||||
|
||||
out_dims[0] = out_dims[0] / nranks;
|
||||
auto start_index = out_dims[0] * comm_ctx->GetRank();
|
||||
auto end_index = start_index + out_dims[0];
|
||||
temp = temp.Slice(start_index, end_index);
|
||||
temp.Resize(out_dims);
|
||||
out->Resize(out_dims);
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
Copy(dev_ctx,
|
||||
*static_cast<const DenseTensor*>(&temp),
|
||||
place,
|
||||
true,
|
||||
static_cast<DenseTensor*>(out));
|
||||
out->Resize(out_dims);
|
||||
#else
|
||||
PADDLE_ENFORCE_EQ(
|
||||
true,
|
||||
false,
|
||||
common::errors::Unavailable("PaddlePaddle should compile with GPU."));
|
||||
#endif
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_scatter,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CScatterOpCUDAKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,29 @@
|
||||
// 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CScatterOpCUDAKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
int ring_id,
|
||||
int root,
|
||||
int nranks,
|
||||
bool use_calc_stream,
|
||||
DenseTensor* out);
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,229 @@
|
||||
// Copyright (c) 2024 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/c_softmax_with_cross_entropy_grad_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/axis_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/cross_entropy.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/math.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax_impl.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/utils/string/string_helper.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int64_t NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CalculateSoftLogitsGrad(T* logits_grad,
|
||||
IndexT* is_ignore,
|
||||
const IndexT* labels,
|
||||
const IndexT ignore_index,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t C) {
|
||||
const T prob = static_cast<T>(1.0 / C);
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
is_ignore[i] = labels[i * C];
|
||||
for (int j = 0; j < C; ++j) {
|
||||
auto real_label = labels[i * C + j];
|
||||
if (real_label == ignore_index) {
|
||||
is_ignore[i] = real_label;
|
||||
}
|
||||
if (real_label >= start_index && real_label < end_index) {
|
||||
int64_t idx = i * D + real_label - start_index;
|
||||
logits_grad[idx] = logits_grad[idx] - prob;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void SoftMaskLabelByIndexGrad(T* logits_grad,
|
||||
const T* loss_grad,
|
||||
const IndexT* is_ignore,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t ignore_index) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N * D, int64_t) {
|
||||
auto row = i / D;
|
||||
auto col = i % D;
|
||||
auto lbl = static_cast<int64_t>(is_ignore[row]);
|
||||
if (lbl == ignore_index) {
|
||||
logits_grad[i] = static_cast<T>(0.0);
|
||||
} else {
|
||||
logits_grad[i] *= loss_grad[row];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void MaskLabelByIndexGrad(T* logits_grad,
|
||||
const T* loss_grad,
|
||||
const IndexT* labels,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t ignore_index) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N * D, int64_t) {
|
||||
auto row = i / D;
|
||||
auto col = i % D;
|
||||
auto lbl = static_cast<int64_t>(labels[row]);
|
||||
if (lbl == ignore_index) {
|
||||
logits_grad[i] = static_cast<T>(0.0);
|
||||
} else if ((col + start_index) == labels[row]) {
|
||||
logits_grad[i] = (logits_grad[i] - static_cast<T>(1.0)) * loss_grad[row];
|
||||
} else {
|
||||
logits_grad[i] *= loss_grad[row];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CSoftmaxWithCrossEntropyGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& softmax_in,
|
||||
const DenseTensor& label_in,
|
||||
const DenseTensor& loss_grad_in,
|
||||
int64_t ignore_index,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* logits_grad) {
|
||||
const DenseTensor* labels = &label_in;
|
||||
const DenseTensor* loss_grad = &loss_grad_in;
|
||||
const DenseTensor* softmax = &softmax_in;
|
||||
DenseTensor* logit_grad = logits_grad;
|
||||
|
||||
if (logit_grad != softmax) {
|
||||
Copy(dev_ctx, *softmax, dev_ctx.GetPlace(), false, logit_grad);
|
||||
}
|
||||
const auto softmax_dims = softmax->dims();
|
||||
const int axis = softmax_dims.size() - 1;
|
||||
const int64_t N = funcs::SizeToAxis<int64_t>(axis, softmax_dims);
|
||||
const int64_t D = funcs::SizeFromAxis<int64_t>(axis, softmax_dims);
|
||||
|
||||
const auto label_dims = labels->dims();
|
||||
const int64_t C = label_dims[axis];
|
||||
|
||||
DenseTensor logit_grad_2d;
|
||||
logit_grad_2d.ShareDataWith(*logit_grad).Resize({N, D});
|
||||
|
||||
int64_t blocks = NumBlocks(N * D);
|
||||
int64_t blocks_cal = NumBlocks(N);
|
||||
int threads = kNumCUDAThreads;
|
||||
const auto& label_type = labels->dtype();
|
||||
const int64_t start_index = rank * D;
|
||||
const int64_t end_index = start_index + D;
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
if (C > 1) {
|
||||
DenseTensor is_ignore;
|
||||
is_ignore.Resize({N, 1});
|
||||
dev_ctx.template Alloc<int32_t>(&is_ignore);
|
||||
|
||||
CalculateSoftLogitsGrad<T, int32_t>
|
||||
<<<blocks_cal, threads, 0, dev_ctx.stream()>>>(
|
||||
logit_grad_2d.data<T>(),
|
||||
is_ignore.data<int32_t>(),
|
||||
labels->data<int32_t>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C);
|
||||
|
||||
SoftMaskLabelByIndexGrad<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
is_ignore.data<int32_t>(),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
ignore_index);
|
||||
} else {
|
||||
MaskLabelByIndexGrad<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
ignore_index);
|
||||
}
|
||||
} else if (label_type == DataType::INT64) {
|
||||
if (C > 1) {
|
||||
DenseTensor is_ignore;
|
||||
is_ignore.Resize({N, 1});
|
||||
dev_ctx.template Alloc<int32_t>(&is_ignore);
|
||||
|
||||
CalculateSoftLogitsGrad<T, int64_t>
|
||||
<<<blocks_cal, threads, 0, dev_ctx.stream()>>>(
|
||||
logit_grad_2d.data<T>(),
|
||||
is_ignore.data<int64_t>(),
|
||||
labels->data<int64_t>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C);
|
||||
|
||||
SoftMaskLabelByIndexGrad<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
is_ignore.data<int64_t>(),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
ignore_index);
|
||||
} else {
|
||||
MaskLabelByIndexGrad<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
ignore_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_softmax_with_cross_entropy_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CSoftmaxWithCrossEntropyGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,379 @@
|
||||
// Copyright (c) 2024 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/platform/collective_helper.h"
|
||||
#include "paddle/phi/kernels/activation_kernel.h"
|
||||
#include "paddle/phi/kernels/full_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/axis_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/broadcast_function.h"
|
||||
#include "paddle/phi/kernels/funcs/cross_entropy.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
|
||||
#include "paddle/phi/kernels/funcs/math.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax_impl.h"
|
||||
#include "paddle/phi/kernels/reduce_max_kernel.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/utils/string/string_helper.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename Context, typename T>
|
||||
struct CSoftmaxWithCrossEntropyFunctor {
|
||||
void operator()(const Context& dev_ctx,
|
||||
const DenseTensor& logits,
|
||||
const DenseTensor& label,
|
||||
int64_t ignore_index,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss);
|
||||
};
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int64_t NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void MaskLabelByIndex(T* predicted_logits,
|
||||
const T* logit,
|
||||
const IndexT* label,
|
||||
const IndexT ignore_index,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int nranks) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
auto real_label = label[i];
|
||||
PADDLE_ENFORCE(((real_label < D * nranks) && (real_label >= 0)) ||
|
||||
(real_label == ignore_index),
|
||||
"The index is out of bounds, "
|
||||
"please check whether the value of label and "
|
||||
"input meet the class number. It should "
|
||||
"be less than [%ld] or equal to [%ld], but received [%ld]",
|
||||
static_cast<int64_t>(D * nranks),
|
||||
static_cast<int64_t>(ignore_index),
|
||||
static_cast<int64_t>(real_label));
|
||||
|
||||
if (real_label >= start_index && real_label < end_index) {
|
||||
predicted_logits[i] = logit[i * D + real_label - start_index];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void SoftMaskLabelByIndex(T* predicted_logits,
|
||||
const T* logit,
|
||||
const IndexT* label,
|
||||
const IndexT ignore_index,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t C,
|
||||
const int nranks) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
for (int j = 0; j < C; ++j) {
|
||||
auto real_label = label[i * C + j];
|
||||
PADDLE_ENFORCE(((real_label < D * nranks) && (real_label >= 0)) ||
|
||||
(real_label == ignore_index),
|
||||
"The index is out of bounds, "
|
||||
"please check whether the value of label and "
|
||||
"input meet the class number. It should "
|
||||
"be less than [%ld] or equal to [%ld], but received [%ld]",
|
||||
static_cast<int64_t>(D * nranks),
|
||||
static_cast<int64_t>(ignore_index),
|
||||
static_cast<int64_t>(real_label));
|
||||
|
||||
if (real_label >= start_index && real_label < end_index) {
|
||||
predicted_logits[i * C + j] = logit[i * D + real_label - start_index];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CalculateLoss(T* loss,
|
||||
const T* predict_logits,
|
||||
const T* sum_exp_logits,
|
||||
const IndexT* label,
|
||||
const int64_t ignore_index,
|
||||
const int64_t N) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
auto real_label = static_cast<int64_t>(label[i]);
|
||||
loss[i] = ignore_index == real_label
|
||||
? static_cast<T>(0)
|
||||
: funcs::TolerableValue<T>()(
|
||||
funcs::TolerableValue<T>()(
|
||||
funcs::real_log(sum_exp_logits[i])) -
|
||||
predict_logits[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CalculateSoftLoss(T* loss,
|
||||
const T* predict_logits,
|
||||
const T* sum_exp_logits,
|
||||
const IndexT* label,
|
||||
const int64_t ignore_index,
|
||||
const int64_t N,
|
||||
const int64_t C) {
|
||||
const T prob = static_cast<T>(1.0 / C);
|
||||
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
T tmp_loss = static_cast<T>(0);
|
||||
int ignore_num = 0;
|
||||
for (int j = 0; j < C; ++j) {
|
||||
auto real_label = static_cast<int64_t>(label[i * C + j]);
|
||||
tmp_loss += ignore_index == real_label
|
||||
? static_cast<T>(0)
|
||||
: funcs::TolerableValue<T>()(
|
||||
(funcs::TolerableValue<T>()(
|
||||
funcs::real_log(sum_exp_logits[i])) -
|
||||
predict_logits[i * C + j]) *
|
||||
prob);
|
||||
ignore_num += ignore_index == real_label ? 1 : 0;
|
||||
}
|
||||
loss[i] = ignore_num > 0 ? static_cast<T>(0) : tmp_loss;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CSoftmaxWithCrossEntropyKernel(const Context& dev_ctx,
|
||||
const DenseTensor& logits,
|
||||
const DenseTensor& label,
|
||||
int64_t ignore_index,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss) {
|
||||
CSoftmaxWithCrossEntropyFunctor<GPUContext, T> functor_;
|
||||
functor_(dev_ctx, logits, label, ignore_index, rank, nranks, softmax, loss);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct CSoftmaxWithCrossEntropyFunctor<GPUContext, T> {
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const DenseTensor& logits_in,
|
||||
const DenseTensor& label_in,
|
||||
int64_t ignore_index,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
const DenseTensor* logits = &logits_in;
|
||||
const DenseTensor* labels = &label_in;
|
||||
|
||||
gpuStream_t stream = nullptr;
|
||||
distributed::NCCLCommContext* comm_ctx = nullptr;
|
||||
|
||||
comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
|
||||
stream = dev_ctx.stream();
|
||||
|
||||
// allocate memory on device.
|
||||
dev_ctx.template Alloc<T>(softmax);
|
||||
dev_ctx.template Alloc<T>(loss);
|
||||
|
||||
const auto& logits_dims = logits->dims();
|
||||
const auto& labels_dims = labels->dims();
|
||||
|
||||
const int axis = logits_dims.size() - 1;
|
||||
const int64_t N = funcs::SizeToAxis<int64_t>(axis, logits_dims);
|
||||
const int64_t D = funcs::SizeFromAxis<int64_t>(axis, logits_dims);
|
||||
const int64_t C = funcs::SizeFromAxis<int64_t>(axis, labels_dims);
|
||||
|
||||
DenseTensor logits_2d, softmax_2d, loss_2d;
|
||||
logits_2d.ShareDataWith(*logits).Resize({N, D});
|
||||
softmax_2d.ShareDataWith(*softmax).Resize({N, D});
|
||||
loss_2d.ShareDataWith(*loss).Resize({N, 1});
|
||||
|
||||
// step 1, obtain logit_max
|
||||
DenseTensor logits_max;
|
||||
logits_max.Resize({N, 1});
|
||||
dev_ctx.template Alloc<T>(&logits_max);
|
||||
|
||||
MaxKernel<T, GPUContext>(dev_ctx, logits_2d, {-1}, true, &logits_max);
|
||||
|
||||
comm_ctx->AllReduce(&logits_max, logits_max, ncclMax, stream);
|
||||
|
||||
// step 2, obtain logit - logit_max
|
||||
std::vector<const DenseTensor*> inputs = {&logits_2d, &logits_max};
|
||||
std::vector<DenseTensor*> outputs = {&softmax_2d};
|
||||
funcs::BroadcastKernel<T>(
|
||||
dev_ctx, inputs, &outputs, funcs::SubtractFunctor<T>());
|
||||
|
||||
// step 3, obtain predict target
|
||||
DenseTensor predicted_logits;
|
||||
predicted_logits.Resize({N, 1});
|
||||
dev_ctx.template Alloc<T>(&predicted_logits);
|
||||
|
||||
Full<T, GPUContext>(dev_ctx, predicted_logits.dims(), 0, &predicted_logits);
|
||||
|
||||
const int64_t start_index = rank * D;
|
||||
const int64_t end_index = start_index + D;
|
||||
|
||||
int64_t blocks = NumBlocks(N);
|
||||
int threads = kNumCUDAThreads;
|
||||
const auto& label_type = labels->dtype();
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
if (C > 1) {
|
||||
SoftMaskLabelByIndex<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
static_cast<int32_t>(ignore_index),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
nranks);
|
||||
} else {
|
||||
MaskLabelByIndex<T, int32_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
static_cast<int32_t>(ignore_index),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
nranks);
|
||||
}
|
||||
} else if (label_type == DataType::INT64) {
|
||||
if (C > 1) {
|
||||
SoftMaskLabelByIndex<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
nranks);
|
||||
} else {
|
||||
MaskLabelByIndex<T, int64_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
nranks);
|
||||
}
|
||||
}
|
||||
|
||||
dev_ctx.template Alloc<T>(&predicted_logits);
|
||||
comm_ctx->AllReduce(&predicted_logits, predicted_logits, ncclSum, stream);
|
||||
|
||||
// step 4, obtain exp(logit)
|
||||
ExpKernel<T, GPUContext>(dev_ctx, softmax_2d, &softmax_2d);
|
||||
|
||||
// step 5, obtain sum_exp_logits
|
||||
DenseTensor sum_exp_logits;
|
||||
sum_exp_logits.Resize({N, 1});
|
||||
dev_ctx.template Alloc<T>(&sum_exp_logits);
|
||||
|
||||
SumKernel<T, GPUContext>(
|
||||
dev_ctx, softmax_2d, {-1}, softmax_2d.dtype(), true, &sum_exp_logits);
|
||||
|
||||
comm_ctx->AllReduce(&sum_exp_logits, sum_exp_logits, ncclSum, stream);
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
if (C > 1) {
|
||||
CalculateSoftLoss<T, int32_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
ignore_index,
|
||||
N,
|
||||
C);
|
||||
} else {
|
||||
CalculateLoss<T, int32_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
ignore_index,
|
||||
N);
|
||||
}
|
||||
|
||||
} else {
|
||||
if (C > 1) {
|
||||
CalculateSoftLoss<T, int64_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
ignore_index,
|
||||
N,
|
||||
C);
|
||||
} else {
|
||||
CalculateLoss<T, int64_t><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
ignore_index,
|
||||
N);
|
||||
}
|
||||
}
|
||||
|
||||
ReciprocalKernel<T, GPUContext>(dev_ctx, sum_exp_logits, &sum_exp_logits);
|
||||
|
||||
inputs = std::vector<const DenseTensor*>{&softmax_2d, &sum_exp_logits};
|
||||
outputs = std::vector<DenseTensor*>{&softmax_2d};
|
||||
funcs::BroadcastKernel<T>(
|
||||
dev_ctx, inputs, &outputs, funcs::MultiplyFunctor<T>());
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_softmax_with_cross_entropy,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CSoftmaxWithCrossEntropyKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,161 @@
|
||||
// Copyright (c) 2024 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/axis_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/cross_entropy.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/math.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax_impl.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/utils/string/string_helper.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int64_t NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CalculateSoftLogitsGrad(T* logits_grad,
|
||||
const IndexT* labels,
|
||||
const T* smooth_weight,
|
||||
const IndexT ignore_index,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t C) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
for (int j = 0; j < C; ++j) {
|
||||
auto real_label = labels[i * C + j];
|
||||
auto prob = smooth_weight[i * C + j];
|
||||
if (real_label >= start_index && real_label < end_index) {
|
||||
int64_t idx = i * D + real_label - start_index;
|
||||
logits_grad[idx] = real_label == ignore_index
|
||||
? static_cast<T>(0)
|
||||
: (logits_grad[idx] - prob);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void SoftMaskLabelByIndexGrad(T* logits_grad,
|
||||
const T* loss_grad,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t C,
|
||||
const bool sum_multi_label_loss) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N * D, int64_t) {
|
||||
auto loss_grad_idx = sum_multi_label_loss ? (i / D) : (i / D * C);
|
||||
logits_grad[i] *= loss_grad[loss_grad_idx];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CSoftmaxWithMultiLabelCrossEntropyGradKernel(
|
||||
const Context& dev_ctx,
|
||||
const DenseTensor& softmax_in,
|
||||
const DenseTensor& label_in,
|
||||
const DenseTensor& smooth_weight_in,
|
||||
const DenseTensor& loss_grad_in,
|
||||
int64_t ignore_index,
|
||||
bool sum_multi_label_loss,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* logits_grad) {
|
||||
const DenseTensor* labels = &label_in;
|
||||
const DenseTensor* smooth_weight = &smooth_weight_in;
|
||||
const DenseTensor* loss_grad = &loss_grad_in;
|
||||
const DenseTensor* softmax = &softmax_in;
|
||||
DenseTensor* logit_grad = logits_grad;
|
||||
|
||||
if (logit_grad != softmax) {
|
||||
Copy(dev_ctx, *softmax, dev_ctx.GetPlace(), false, logit_grad);
|
||||
}
|
||||
const auto softmax_dims = softmax->dims();
|
||||
const int axis = softmax_dims.size() - 1;
|
||||
const int64_t N = funcs::SizeToAxis<int64_t>(axis, softmax_dims);
|
||||
const int64_t D = funcs::SizeFromAxis<int64_t>(axis, softmax_dims);
|
||||
|
||||
const auto label_dims = labels->dims();
|
||||
const int64_t C = label_dims[axis];
|
||||
|
||||
DenseTensor logit_grad_2d;
|
||||
logit_grad_2d.ShareDataWith(*logit_grad).Resize({N, D});
|
||||
|
||||
int64_t blocks = NumBlocks(N * D);
|
||||
int64_t blocks_cal = NumBlocks(N);
|
||||
int threads = kNumCUDAThreads;
|
||||
const auto& label_type = labels->dtype();
|
||||
const int64_t start_index = rank * D;
|
||||
const int64_t end_index = start_index + D;
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
CalculateSoftLogitsGrad<T, int32_t>
|
||||
<<<blocks_cal, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
smooth_weight->data<T>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C);
|
||||
|
||||
SoftMaskLabelByIndexGrad<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
sum_multi_label_loss);
|
||||
} else if (label_type == DataType::INT64) {
|
||||
CalculateSoftLogitsGrad<T, int64_t>
|
||||
<<<blocks_cal, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
smooth_weight->data<T>(),
|
||||
ignore_index,
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C);
|
||||
|
||||
SoftMaskLabelByIndexGrad<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(logit_grad_2d.data<T>(),
|
||||
loss_grad->data<T>(),
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
sum_multi_label_loss);
|
||||
}
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_softmax_with_multi_label_cross_entropy_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CSoftmaxWithMultiLabelCrossEntropyGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,315 @@
|
||||
// Copyright (c) 2024 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/platform/collective_helper.h"
|
||||
#include "paddle/phi/kernels/funcs/axis_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/cross_entropy.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/math.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax.h"
|
||||
#include "paddle/phi/kernels/funcs/softmax_impl.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/utils/string/string_helper.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename Context, typename T>
|
||||
struct CSoftmaxWithMultiLabelCrossEntropyFunctor {
|
||||
void operator()(const Context& dev_ctx,
|
||||
const DenseTensor& logits,
|
||||
const DenseTensor& label,
|
||||
const DenseTensor& smooth_weight,
|
||||
int64_t ignore_index,
|
||||
bool sum_multi_label_loss,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss);
|
||||
};
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int64_t NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void SoftMaskLabelByIndex(T* predicted_logits,
|
||||
const T* logit,
|
||||
const IndexT* label,
|
||||
const IndexT ignore_index,
|
||||
const int64_t start_index,
|
||||
const int64_t end_index,
|
||||
const int64_t N,
|
||||
const int64_t D,
|
||||
const int64_t C,
|
||||
const int nranks) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
for (int j = 0; j < C; ++j) {
|
||||
auto real_label = label[i * C + j];
|
||||
PADDLE_ENFORCE(((real_label < D * nranks) && (real_label >= 0)) ||
|
||||
(real_label == ignore_index),
|
||||
"The index is out of bounds, "
|
||||
"please check whether the value of label and "
|
||||
"input meet the class number. It should "
|
||||
"be less than [%ld] or equal to [%ld], but received [%ld]",
|
||||
static_cast<int64_t>(D * nranks),
|
||||
static_cast<int64_t>(ignore_index),
|
||||
static_cast<int64_t>(real_label));
|
||||
|
||||
if (real_label >= start_index && real_label < end_index) {
|
||||
predicted_logits[i * C + j] = logit[i * D + real_label - start_index];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IndexT>
|
||||
__global__ void CalculateSoftLoss(T* loss,
|
||||
const T* predict_logits,
|
||||
const T* sum_exp_logits,
|
||||
const IndexT* label,
|
||||
const T* smooth_weight,
|
||||
const int64_t ignore_index,
|
||||
const int64_t N,
|
||||
const int64_t C,
|
||||
const bool sum_multi_label_loss) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, N, int64_t) {
|
||||
T tmp_loss = static_cast<T>(0);
|
||||
loss[i] = static_cast<T>(0);
|
||||
T log_sum_exp_logits =
|
||||
funcs::TolerableValue<T>()(funcs::real_log(sum_exp_logits[i]));
|
||||
for (int j = 0; j < C; ++j) {
|
||||
int64_t label_idx = i * C + j;
|
||||
auto real_label = static_cast<int64_t>(label[label_idx]);
|
||||
auto prob = static_cast<T>(smooth_weight[label_idx]);
|
||||
tmp_loss =
|
||||
ignore_index == real_label
|
||||
? static_cast<T>(0)
|
||||
: funcs::TolerableValue<T>()(
|
||||
(log_sum_exp_logits - predict_logits[label_idx]) * prob);
|
||||
if (sum_multi_label_loss) {
|
||||
loss[i] += tmp_loss;
|
||||
} else {
|
||||
loss[label_idx] = tmp_loss;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CSoftmaxWithMultiLabelCrossEntropyKernel(const Context& dev_ctx,
|
||||
const DenseTensor& logits,
|
||||
const DenseTensor& label,
|
||||
const DenseTensor& smooth_weight,
|
||||
int64_t ignore_index,
|
||||
bool sum_multi_label_loss,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss) {
|
||||
CSoftmaxWithMultiLabelCrossEntropyFunctor<GPUContext, T> functor_;
|
||||
functor_(dev_ctx,
|
||||
logits,
|
||||
label,
|
||||
smooth_weight,
|
||||
ignore_index,
|
||||
sum_multi_label_loss,
|
||||
rank,
|
||||
nranks,
|
||||
softmax,
|
||||
loss);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct CSoftmaxWithMultiLabelCrossEntropyFunctor<GPUContext, T> {
|
||||
void operator()(const GPUContext& dev_ctx,
|
||||
const DenseTensor& logits_in,
|
||||
const DenseTensor& label_in,
|
||||
const DenseTensor& smooth_weight_in,
|
||||
int64_t ignore_index,
|
||||
bool sum_multi_label_loss,
|
||||
int rank,
|
||||
int nranks,
|
||||
DenseTensor* softmax,
|
||||
DenseTensor* loss) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
const DenseTensor* logits = &logits_in;
|
||||
const DenseTensor* labels = &label_in;
|
||||
const DenseTensor* smooth_weight = &smooth_weight_in;
|
||||
|
||||
gpuStream_t stream = nullptr;
|
||||
distributed::NCCLCommContext* comm_ctx = nullptr;
|
||||
|
||||
comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
|
||||
stream = dev_ctx.stream();
|
||||
|
||||
// allocate memory on device.
|
||||
dev_ctx.template Alloc<T>(softmax);
|
||||
dev_ctx.template Alloc<T>(loss);
|
||||
|
||||
const auto& logits_dims = logits->dims();
|
||||
const auto& labels_dims = labels->dims();
|
||||
|
||||
const int axis = logits_dims.size() - 1;
|
||||
const int64_t N = funcs::SizeToAxis<int64_t>(axis, logits_dims);
|
||||
const int64_t D = funcs::SizeFromAxis<int64_t>(axis, logits_dims);
|
||||
const int64_t C = funcs::SizeFromAxis<int64_t>(axis, labels_dims);
|
||||
|
||||
DenseTensor logits_2d, softmax_2d, loss_2d;
|
||||
logits_2d.ShareDataWith(*logits).Resize({N, D});
|
||||
softmax_2d.ShareDataWith(*softmax).Resize({N, D});
|
||||
int64_t loss_last_dim = sum_multi_label_loss ? 1 : C;
|
||||
loss_2d.ShareDataWith(*loss).Resize({N, loss_last_dim});
|
||||
|
||||
auto eigen_logits = funcs::EigenMatrix<T>::From(logits_2d);
|
||||
auto eigen_softmax = funcs::EigenMatrix<T>::From(softmax_2d);
|
||||
|
||||
// step 1, obtain logit_max
|
||||
DenseTensor logits_max;
|
||||
logits_max.Resize({N, 1});
|
||||
dev_ctx.template Alloc<T>(&logits_max);
|
||||
|
||||
auto eigen_logits_max = funcs::EigenMatrix<T>::From(logits_max);
|
||||
Eigen::DSizes<int, 1> along_axis(1);
|
||||
eigen_logits_max.device(*dev_ctx.eigen_device()) =
|
||||
eigen_logits.maximum(along_axis);
|
||||
|
||||
comm_ctx->AllReduce(&logits_max, logits_max, ncclMax, stream);
|
||||
|
||||
// step 2, obtain logit - logit_max
|
||||
Eigen::DSizes<int, 2> batch_by_one(N, 1);
|
||||
Eigen::DSizes<int, 2> one_by_class(1, D);
|
||||
|
||||
eigen_softmax.device(*dev_ctx.eigen_device()) =
|
||||
(eigen_logits -
|
||||
eigen_logits_max.reshape(batch_by_one).broadcast(one_by_class));
|
||||
|
||||
// step 3, obtain predict target
|
||||
DenseTensor predicted_logits;
|
||||
predicted_logits.Resize({N, C});
|
||||
dev_ctx.template Alloc<T>(&predicted_logits);
|
||||
|
||||
auto t = EigenVector<T>::Flatten(predicted_logits);
|
||||
t.device(*dev_ctx.eigen_device()) = t.constant(static_cast<T>(0));
|
||||
|
||||
const int64_t start_index = rank * D;
|
||||
const int64_t end_index = start_index + D;
|
||||
|
||||
int64_t blocks = NumBlocks(N);
|
||||
int threads = kNumCUDAThreads;
|
||||
const auto& label_type = labels->dtype();
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
SoftMaskLabelByIndex<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
static_cast<int32_t>(ignore_index),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
nranks);
|
||||
} else if (label_type == DataType::INT64) {
|
||||
SoftMaskLabelByIndex<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
predicted_logits.data<T>(),
|
||||
softmax_2d.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
static_cast<int32_t>(ignore_index),
|
||||
start_index,
|
||||
end_index,
|
||||
N,
|
||||
D,
|
||||
C,
|
||||
nranks);
|
||||
}
|
||||
|
||||
dev_ctx.template Alloc<T>(&predicted_logits);
|
||||
comm_ctx->AllReduce(&predicted_logits, predicted_logits, ncclSum, stream);
|
||||
|
||||
// step 4, obtain exp(logit)
|
||||
eigen_softmax.device(*dev_ctx.eigen_device()) = eigen_softmax.exp();
|
||||
|
||||
// step 5, obtain sum_exp_logits
|
||||
DenseTensor sum_exp_logits;
|
||||
sum_exp_logits.Resize({N, 1});
|
||||
dev_ctx.template Alloc<T>(&sum_exp_logits);
|
||||
|
||||
SumKernel<T, GPUContext>(
|
||||
dev_ctx, softmax_2d, {-1}, softmax_2d.dtype(), true, &sum_exp_logits);
|
||||
|
||||
comm_ctx->AllReduce(&sum_exp_logits, sum_exp_logits, ncclSum, stream);
|
||||
|
||||
if (label_type == DataType::INT32) {
|
||||
CalculateSoftLoss<T, int32_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int32_t>(),
|
||||
smooth_weight->data<T>(),
|
||||
ignore_index,
|
||||
N,
|
||||
C,
|
||||
sum_multi_label_loss);
|
||||
|
||||
} else {
|
||||
CalculateSoftLoss<T, int64_t>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(loss_2d.data<T>(),
|
||||
predicted_logits.data<T>(),
|
||||
sum_exp_logits.data<T>(),
|
||||
labels->data<int64_t>(),
|
||||
smooth_weight->data<T>(),
|
||||
ignore_index,
|
||||
N,
|
||||
C,
|
||||
sum_multi_label_loss);
|
||||
}
|
||||
|
||||
auto eigen_sum_exp_logits = funcs::EigenMatrix<T>::From(sum_exp_logits);
|
||||
eigen_softmax.device(*dev_ctx.eigen_device()) =
|
||||
(eigen_softmax *
|
||||
eigen_sum_exp_logits.inverse().broadcast(one_by_class));
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_softmax_with_multi_label_cross_entropy,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CSoftmaxWithMultiLabelCrossEntropyKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,113 @@
|
||||
// Copyright (c) 2023 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/c_split_kernel.h"
|
||||
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int64_t kNumCUDAThreads = 512;
|
||||
static constexpr int64_t kNumMaximumNumBlocks = 4096;
|
||||
|
||||
static inline int64_t NumBlocks(const int64_t N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void SplitFromRank(const T* input,
|
||||
T* output,
|
||||
const int64_t rows,
|
||||
const int64_t columns,
|
||||
const int rank,
|
||||
const int nranks,
|
||||
const int64_t limit) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, limit, int64_t) {
|
||||
int64_t row = i / columns;
|
||||
int64_t col = i % columns;
|
||||
|
||||
int64_t block = columns / nranks;
|
||||
int64_t start = block * rank;
|
||||
int64_t end = start + block;
|
||||
|
||||
if (col >= start && col < end) {
|
||||
int64_t idx = block * row + col % block;
|
||||
output[idx] = input[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CSplitKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
int rank,
|
||||
int nranks,
|
||||
bool use_model_parallel,
|
||||
DenseTensor* out) {
|
||||
auto place = dev_ctx.GetPlace();
|
||||
|
||||
PADDLE_ENFORCE_GE(rank,
|
||||
0,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of rank (%d) for c_split must be "
|
||||
"greater than or equal to 0.",
|
||||
rank));
|
||||
PADDLE_ENFORCE_GE(nranks,
|
||||
2,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of nranks (%d) for c_split must be "
|
||||
"greater than or equal to 2.",
|
||||
nranks));
|
||||
PADDLE_ENFORCE_LT(rank,
|
||||
nranks,
|
||||
common::errors::PreconditionNotMet(
|
||||
"The value of rank (%d) for c_split must be "
|
||||
"less than that of nranks (%d).",
|
||||
rank,
|
||||
nranks));
|
||||
|
||||
auto dims = x.dims();
|
||||
auto dims_size = dims.size();
|
||||
// final dim
|
||||
int64_t end_size = dims[dims_size - 1];
|
||||
|
||||
// remain dim
|
||||
auto remain_ddim = slice_ddim(dims, 0, dims_size - 1);
|
||||
int64_t remain_numel = common::product(remain_ddim);
|
||||
|
||||
int64_t limit = x.numel();
|
||||
int64_t blocks = NumBlocks(limit);
|
||||
int64_t threads = kNumCUDAThreads;
|
||||
|
||||
dims[dims_size - 1] /= nranks;
|
||||
out->Resize(dims);
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
|
||||
SplitFromRank<T><<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
x.data<T>(), out->data<T>(), remain_numel, end_size, rank, nranks, limit);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(c_split,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CSplitKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,140 @@
|
||||
// Copyright (c) 2024 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/calc_reduced_attn_kernel.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
|
||||
struct CalcReducedAttnScoresParams : public FlashAttnParamsBase {
|
||||
bool return_softmax;
|
||||
DenseTensor* softmax;
|
||||
|
||||
CalcReducedAttnScoresParams(const GPUContext& dev_ctx,
|
||||
const int _batch_size,
|
||||
const int64_t _max_seqlen_q,
|
||||
const int64_t _max_seqlen_k,
|
||||
const int _num_heads,
|
||||
const int _num_heads_k,
|
||||
const int _head_size,
|
||||
const float _scale,
|
||||
const DataType q_dtype)
|
||||
: FlashAttnParamsBase(/*version=*/2,
|
||||
/*is_fwd=*/true,
|
||||
_batch_size,
|
||||
_max_seqlen_q,
|
||||
_max_seqlen_k,
|
||||
_num_heads,
|
||||
_num_heads_k,
|
||||
_head_size,
|
||||
_scale,
|
||||
/*_causal=*/false,
|
||||
q_dtype,
|
||||
optional<DenseTensor>{},
|
||||
optional<DenseTensor>{},
|
||||
/*_unpadded_lse=*/false,
|
||||
/*_total_q*/ 0) {}
|
||||
};
|
||||
#endif
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CalcReducedAttnScoresKernel(const Context& dev_ctx,
|
||||
const DenseTensor& q,
|
||||
const DenseTensor& k,
|
||||
const DenseTensor& softmax_lse,
|
||||
DenseTensor* reduced_scores) {
|
||||
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
|
||||
PADDLE_ENFORCE_EQ(q.dims().size(),
|
||||
4,
|
||||
common::errors::InvalidArgument(
|
||||
"calc_reduced_attention receive input with dim "
|
||||
"[batch_size, seq_len, num_heads, head_dim]"));
|
||||
|
||||
PADDLE_ENFORCE_EQ(k.dims().size(),
|
||||
4,
|
||||
common::errors::InvalidArgument(
|
||||
"calc_reduced_attention receive input with dim "
|
||||
"[batch_size, seq_len, num_heads, head_dim]"));
|
||||
|
||||
if (!reduced_scores->IsInitialized())
|
||||
dev_ctx.template Alloc<float>(reduced_scores);
|
||||
funcs::SetConstant<Context, float> set_zero;
|
||||
set_zero(dev_ctx, reduced_scores, 0.0f);
|
||||
// q, k, v [batch_size, seq_len, num_heads, head_dim]
|
||||
const int64_t batch_size = q.dims()[0];
|
||||
const int64_t seqlen_q = q.dims()[1];
|
||||
const int64_t num_heads = q.dims()[2];
|
||||
const int64_t head_size = q.dims()[3];
|
||||
const int64_t seqlen_k = k.dims()[1];
|
||||
const int64_t num_heads_k = k.dims()[2];
|
||||
|
||||
const float softmax_scale = 1.0f / std::sqrt(head_size);
|
||||
const float softmax_unscale = std::sqrt(head_size);
|
||||
|
||||
using Params = CalcReducedAttnScoresParams;
|
||||
|
||||
Params params = Params(dev_ctx,
|
||||
batch_size,
|
||||
seqlen_q,
|
||||
seqlen_k,
|
||||
num_heads,
|
||||
num_heads_k,
|
||||
head_size,
|
||||
softmax_scale,
|
||||
q.dtype());
|
||||
|
||||
cudaStream_t stream = dev_ctx.stream();
|
||||
|
||||
bool succ = dynload::calc_reduced_attn_scores(q.data(),
|
||||
k.data(),
|
||||
softmax_lse.data(),
|
||||
reduced_scores->data(),
|
||||
/*softmax_ptr=*/nullptr,
|
||||
params.batch_size,
|
||||
params.max_seqlen_q,
|
||||
params.max_seqlen_k,
|
||||
params.num_heads,
|
||||
params.num_heads_k,
|
||||
params.head_size,
|
||||
params.softmax_scale,
|
||||
/*return_softmax=*/false,
|
||||
params.is_bf16,
|
||||
/*num_splits=*/0,
|
||||
stream,
|
||||
q.strides()[1],
|
||||
k.strides()[1],
|
||||
reduced_scores->strides()[1],
|
||||
q.strides()[2],
|
||||
k.strides()[2],
|
||||
reduced_scores->strides()[2],
|
||||
q.strides()[0],
|
||||
k.strides()[0],
|
||||
reduced_scores->strides()[0]);
|
||||
CheckFlashAttnStatus(succ);
|
||||
#else
|
||||
RaiseNotSupportedError();
|
||||
#endif
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(calc_reduced_attn_scores,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CalcReducedAttnScoresKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,56 @@
|
||||
// Copyright (c) 2022 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename InT, typename OutT>
|
||||
struct CastFunctor {
|
||||
__device__ __forceinline__ OutT operator()(const InT x) const {
|
||||
return static_cast<OutT>(x);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename InT, typename OutT>
|
||||
void CastCUDAKernelImpl(const GPUContext& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DataType out_dtype,
|
||||
DenseTensor* out) {
|
||||
std::vector<const DenseTensor*> inputs;
|
||||
std::vector<DenseTensor*> outputs;
|
||||
inputs.emplace_back(&x);
|
||||
outputs.emplace_back(out);
|
||||
dev_ctx.Alloc<OutT>(out);
|
||||
out->set_type(out_dtype);
|
||||
if (out->numel() == 0) return;
|
||||
funcs::ElementwiseKernel<OutT>(
|
||||
dev_ctx, inputs, &outputs, CastFunctor<InT, OutT>());
|
||||
}
|
||||
|
||||
template <typename InT>
|
||||
void CastCUDAKernel(const GPUContext& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DataType out_dtype,
|
||||
DenseTensor* out) {
|
||||
PD_VISIT_ALL_TYPES(out_dtype, "CastCUDAKernelImpl", ([&] {
|
||||
CastCUDAKernelImpl<InT, data_t>(
|
||||
dev_ctx, x, out_dtype, out);
|
||||
}));
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,83 @@
|
||||
// Copyright (c) 2021 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/cast_kernel.h"
|
||||
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/visit_type.h"
|
||||
#include "paddle/phi/kernels/gpu/cast_impl.h"
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CastKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
DataType out_dtype,
|
||||
DenseTensor* out) {
|
||||
if (x.dtype() == out_dtype) {
|
||||
if (x.dims() == make_ddim({-1})) {
|
||||
*out = x;
|
||||
return;
|
||||
}
|
||||
if (!out->IsSharedWith(x)) {
|
||||
Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (out->IsSharedWith(x)) {
|
||||
auto x_origin = x;
|
||||
CastCUDAKernel<T>(dev_ctx, x_origin, out_dtype, out);
|
||||
} else {
|
||||
CastCUDAKernel<T>(dev_ctx, x, out_dtype, out);
|
||||
}
|
||||
}
|
||||
#ifdef _WIN32
|
||||
INSTANTIATE_CAST_KERNEL(float, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(double, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(int, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(int64_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(uint8_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(uint16_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(uint32_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(uint64_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(bool, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(int16_t, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(float16, GPUContext)
|
||||
INSTANTIATE_CAST_KERNEL(bfloat16, GPUContext)
|
||||
#endif
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(cast,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CastKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
int16_t,
|
||||
bool,
|
||||
int8_t,
|
||||
uint8_t,
|
||||
uint16_t,
|
||||
uint32_t,
|
||||
uint64_t,
|
||||
phi::float16,
|
||||
phi::complex64,
|
||||
phi::complex128,
|
||||
phi::bfloat16,
|
||||
phi::float8_e4m3fn,
|
||||
phi::float8_e5m2) {
|
||||
kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED);
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2022 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/channel_shuffle_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(channel_shuffle_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ChannelShuffleGradKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,28 @@
|
||||
// Copyright (c) 2022 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/channel_shuffle_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(channel_shuffle,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ChannelShuffleKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,526 @@
|
||||
/* Copyright (c) 2023 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/check_numerics_kernel.h"
|
||||
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/check_numerics_utils.h"
|
||||
#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static std::once_flag init_multi_gpu_op_var_map_flag;
|
||||
|
||||
// lazy init
|
||||
static std::vector<std::unordered_map<std::string, Allocator::AllocationPtr>>&
|
||||
multi_op_var2gpu_str() {
|
||||
static std::vector<std::unordered_map<std::string, Allocator::AllocationPtr>>
|
||||
_multi_op_var2gpu_str;
|
||||
return _multi_op_var2gpu_str;
|
||||
}
|
||||
|
||||
static std::vector<std::mutex>& multi_op_var2gpu_str_mutex() {
|
||||
static std::vector<std::mutex> _multi_op_var2gpu_str_mutex;
|
||||
return _multi_op_var2gpu_str_mutex;
|
||||
}
|
||||
|
||||
static void InitMultiGPUOpVarMap() {
|
||||
int dev_count = backends::gpu::GetGPUDeviceCount();
|
||||
PADDLE_ENFORCE_GT(dev_count,
|
||||
0,
|
||||
common::errors::NotFound(
|
||||
"cuda device must > 0, now dev_count=%d", dev_count));
|
||||
|
||||
// https://stackoverflow.com/questions/16465633/how-can-i-use-something-like-stdvectorstdmutex
|
||||
std::vector<std::unordered_map<std::string, Allocator::AllocationPtr>>
|
||||
tmp_multi(dev_count);
|
||||
std::vector<std::mutex> tmp_multi_mutex(dev_count);
|
||||
|
||||
multi_op_var2gpu_str().swap(tmp_multi);
|
||||
multi_op_var2gpu_str_mutex().swap(tmp_multi_mutex);
|
||||
}
|
||||
|
||||
template <typename T, int ReduceType>
|
||||
__device__ T BlockReduce(T value) {
|
||||
__shared__ T shared_mem[1024];
|
||||
|
||||
shared_mem[threadIdx.x] = value;
|
||||
__syncthreads();
|
||||
|
||||
for (int stride = blockDim.x >> 1; stride > 0; stride = stride >> 1) {
|
||||
if (threadIdx.x < stride) {
|
||||
T value0 = shared_mem[threadIdx.x];
|
||||
T value1 = shared_mem[threadIdx.x + stride];
|
||||
T reduce_value;
|
||||
if (ReduceType == 0) {
|
||||
// max
|
||||
reduce_value = value0 > value1 ? value0 : value1;
|
||||
} else if (ReduceType == 1) {
|
||||
// min
|
||||
reduce_value = value0 < value1 ? value0 : value1;
|
||||
} else if (ReduceType == 2) {
|
||||
// sum
|
||||
reduce_value = value0 + value1;
|
||||
}
|
||||
shared_mem[threadIdx.x] = reduce_value;
|
||||
}
|
||||
|
||||
if (stride > 16) {
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
return shared_mem[0];
|
||||
}
|
||||
|
||||
__device__ void BlockReduceNumNanInfAndWrite(const int64_t num_nan,
|
||||
const int64_t num_inf,
|
||||
const int64_t num_zero,
|
||||
int64_t offset,
|
||||
int64_t* num_nan_ptr,
|
||||
int64_t* num_inf_ptr,
|
||||
int64_t* num_zero_ptr) {
|
||||
int64_t block_num_nan = BlockReduce<int64_t, 2>(num_nan);
|
||||
int64_t block_num_inf = BlockReduce<int64_t, 2>(num_inf);
|
||||
int64_t block_num_zero = BlockReduce<int64_t, 2>(num_zero);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
num_nan_ptr[offset] = block_num_nan;
|
||||
num_inf_ptr[offset] = block_num_inf;
|
||||
num_zero_ptr[offset] = block_num_zero;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
std::enable_if_t<std::is_same<T, complex64>::value ||
|
||||
std::is_same<T, complex128>::value,
|
||||
bool> = true>
|
||||
__device__ void BlockReduceMaxMinAndWrite(const T max_value,
|
||||
const T min_value,
|
||||
const T mean_value,
|
||||
int64_t offset,
|
||||
T* max_ptr,
|
||||
T* min_ptr,
|
||||
T* mean_ptr) {
|
||||
// TODO(Xreki): support complex
|
||||
}
|
||||
|
||||
template <typename T,
|
||||
std::enable_if_t<!std::is_same<T, complex64>::value &&
|
||||
!std::is_same<T, complex128>::value,
|
||||
bool> = true>
|
||||
__device__ void BlockReduceMaxMinAndWrite(const T max_value,
|
||||
const T min_value,
|
||||
const T mean_value,
|
||||
int64_t offset,
|
||||
T* max_ptr,
|
||||
T* min_ptr,
|
||||
T* mean_ptr) {
|
||||
if (max_ptr && min_ptr && mean_ptr) {
|
||||
__syncthreads();
|
||||
|
||||
T block_max_value = funcs::BlockReduceMax<T>(max_value, FINAL_MASK);
|
||||
T block_min_value = funcs::BlockReduceMin<T>(min_value, FINAL_MASK);
|
||||
T block_mean_value = funcs::BlockReduceSum<T>(mean_value, FINAL_MASK);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
max_ptr[offset] = block_max_value;
|
||||
min_ptr[offset] = block_min_value;
|
||||
mean_ptr[offset] = block_mean_value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void FindNanInfAndBlockMaxMin(const T* value_ptr,
|
||||
const int64_t numel,
|
||||
int64_t* block_num_nan_ptr,
|
||||
int64_t* block_num_inf_ptr,
|
||||
int64_t* block_num_zero_ptr,
|
||||
MT* tensor_block_max_ptr,
|
||||
MT* tensor_block_min_ptr,
|
||||
MT* tensor_block_mean_ptr) {
|
||||
int64_t i =
|
||||
static_cast<int64_t>(threadIdx.x) +
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
|
||||
|
||||
int64_t num_nan = 0;
|
||||
int64_t num_inf = 0;
|
||||
int64_t num_zero = 0;
|
||||
|
||||
MT max_value = static_cast<MT>(i < numel ? value_ptr[i] : value_ptr[0]);
|
||||
MT min_value = static_cast<MT>(i < numel ? value_ptr[i] : value_ptr[0]);
|
||||
MT mean_value = static_cast<MT>(0);
|
||||
for (; i < numel; i += blockDim.x * gridDim.x) {
|
||||
MT value = static_cast<MT>(value_ptr[i]);
|
||||
|
||||
max_value = value > max_value ? value : max_value;
|
||||
min_value = value < min_value ? value : min_value;
|
||||
mean_value += value / static_cast<MT>(numel);
|
||||
|
||||
if (isnan(value)) {
|
||||
num_nan += 1;
|
||||
} else if (isinf(value)) {
|
||||
num_inf += 1;
|
||||
}
|
||||
if (value == static_cast<MT>(0)) {
|
||||
num_zero += 1;
|
||||
}
|
||||
}
|
||||
|
||||
BlockReduceNumNanInfAndWrite(num_nan,
|
||||
num_inf,
|
||||
num_zero,
|
||||
blockIdx.x,
|
||||
block_num_nan_ptr,
|
||||
block_num_inf_ptr,
|
||||
block_num_zero_ptr);
|
||||
|
||||
BlockReduceMaxMinAndWrite<MT>(max_value,
|
||||
min_value,
|
||||
mean_value,
|
||||
blockIdx.x,
|
||||
tensor_block_max_ptr,
|
||||
tensor_block_min_ptr,
|
||||
tensor_block_mean_ptr);
|
||||
}
|
||||
|
||||
template <typename T, typename MT>
|
||||
__global__ void FindGlobalMaxMinAndPrint(const int64_t* block_num_nan_ptr,
|
||||
const int64_t* block_num_inf_ptr,
|
||||
const int64_t* block_num_zero_ptr,
|
||||
const MT* tensor_block_max_ptr,
|
||||
const MT* tensor_block_min_ptr,
|
||||
const MT* tensor_block_mean_ptr,
|
||||
const char* debug_info,
|
||||
int64_t numel,
|
||||
int64_t numel_max_min,
|
||||
int check_nan_inf_level,
|
||||
int64_t* stats_ptr,
|
||||
float* values_ptr) {
|
||||
if (blockIdx.x == 0 && threadIdx.x == 0) {
|
||||
int64_t num_nan = 0;
|
||||
int64_t num_inf = 0;
|
||||
int64_t num_zero = 0;
|
||||
|
||||
// numel_max_min <= 128
|
||||
for (int64_t i = 0; i < numel_max_min; ++i) {
|
||||
num_nan += block_num_nan_ptr[i];
|
||||
num_inf += block_num_inf_ptr[i];
|
||||
num_zero += block_num_zero_ptr[i];
|
||||
}
|
||||
|
||||
MT max_value = static_cast<MT>(0);
|
||||
MT min_value = static_cast<MT>(0);
|
||||
MT mean_value = static_cast<MT>(0);
|
||||
if (tensor_block_max_ptr && tensor_block_min_ptr && tensor_block_mean_ptr) {
|
||||
max_value = tensor_block_max_ptr[0];
|
||||
min_value = tensor_block_min_ptr[0];
|
||||
mean_value = tensor_block_mean_ptr[0];
|
||||
|
||||
// numel_max_min <= 128
|
||||
for (int64_t i = 1; i < numel_max_min; ++i) {
|
||||
MT tmp_max_value = tensor_block_max_ptr[i];
|
||||
MT tmp_min_value = tensor_block_min_ptr[i];
|
||||
MT tmp_mean_value = tensor_block_mean_ptr[i];
|
||||
|
||||
max_value = tmp_max_value > max_value ? tmp_max_value : max_value;
|
||||
min_value = tmp_min_value < min_value ? tmp_min_value : min_value;
|
||||
mean_value += tmp_mean_value;
|
||||
}
|
||||
funcs::SaveStatsAndValues<MT>(num_nan,
|
||||
num_inf,
|
||||
num_zero,
|
||||
max_value,
|
||||
min_value,
|
||||
mean_value,
|
||||
stats_ptr,
|
||||
values_ptr);
|
||||
}
|
||||
|
||||
funcs::PrintForDifferentLevel<T, MT>(debug_info,
|
||||
numel,
|
||||
num_nan,
|
||||
num_inf,
|
||||
num_zero,
|
||||
max_value,
|
||||
min_value,
|
||||
mean_value,
|
||||
check_nan_inf_level);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline std::string GetHintString(const std::string& op_type,
|
||||
const std::string& var_name,
|
||||
const Place& place,
|
||||
int dev_id = -1) {
|
||||
std::string op_var =
|
||||
funcs::GetCpuHintString<T>(op_type, var_name, place, dev_id);
|
||||
PADDLE_ENFORCE_EQ(
|
||||
(dev_id >= 0 && dev_id < multi_op_var2gpu_str_mutex().size()),
|
||||
true,
|
||||
common::errors::OutOfRange("GPU dev_id must >=0 and < dev_count=%d",
|
||||
multi_op_var2gpu_str_mutex().size()));
|
||||
return op_var;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static char* GetGpuHintStringPtr(const GPUContext& dev_ctx,
|
||||
const std::string& op_type,
|
||||
const std::string& var_name,
|
||||
int dev_id) {
|
||||
std::call_once(init_multi_gpu_op_var_map_flag, InitMultiGPUOpVarMap);
|
||||
|
||||
std::string op_var =
|
||||
GetHintString<T>(op_type, var_name, dev_ctx.GetPlace(), dev_id);
|
||||
char* gpu_str_ptr = nullptr;
|
||||
|
||||
{
|
||||
auto& op_var2gpu_str_mutex = multi_op_var2gpu_str_mutex().at(dev_id);
|
||||
auto& op_var2gpu_str = multi_op_var2gpu_str().at(dev_id);
|
||||
|
||||
std::lock_guard<std::mutex> guard(op_var2gpu_str_mutex);
|
||||
if (op_var2gpu_str.find(op_var) == op_var2gpu_str.end()) { // insert
|
||||
auto gpu_str_tensor = memory_utils::Alloc(
|
||||
dev_ctx.GetPlace(),
|
||||
op_var.length() + 1,
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
gpu_str_ptr = reinterpret_cast<char*>(gpu_str_tensor->ptr());
|
||||
|
||||
op_var2gpu_str.emplace(op_var, std::move(gpu_str_tensor));
|
||||
|
||||
auto iter = op_var2gpu_str.find(op_var);
|
||||
PADDLE_ENFORCE_EQ(iter != op_var2gpu_str.end(),
|
||||
true,
|
||||
common::errors::PreconditionNotMet(
|
||||
"op_var=%s should be successfully insert into "
|
||||
"op_var2gpu_str, but now failed",
|
||||
op_var));
|
||||
|
||||
#ifdef __HIPCC__
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(hipMemcpyAsync(gpu_str_ptr,
|
||||
iter->first.c_str(),
|
||||
op_var.length() + 1,
|
||||
hipMemcpyHostToDevice,
|
||||
dev_ctx.stream()));
|
||||
#else
|
||||
const char* stable_str =
|
||||
backends::gpu::RestoreHostMemIfCapturingCUDAGraph(
|
||||
const_cast<char*>(iter->first.c_str()), op_var.length() + 1);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpyAsync(gpu_str_ptr,
|
||||
stable_str,
|
||||
op_var.length() + 1,
|
||||
cudaMemcpyHostToDevice,
|
||||
dev_ctx.stream()));
|
||||
#endif
|
||||
} else { // get
|
||||
auto iter = op_var2gpu_str.find(op_var);
|
||||
PADDLE_ENFORCE_EQ(iter != op_var2gpu_str.end(),
|
||||
true,
|
||||
common::errors::PreconditionNotMet(
|
||||
"op_var=%s should be in the op_var2gpu_str, but "
|
||||
"now can't find it",
|
||||
op_var));
|
||||
gpu_str_ptr = reinterpret_cast<char*>(iter->second->ptr());
|
||||
}
|
||||
}
|
||||
return gpu_str_ptr;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void PrintStack(const GPUContext& dev_ctx,
|
||||
const DenseTensor& stats,
|
||||
const std::string& op_type,
|
||||
const std::string& var_name,
|
||||
int dev_id) {
|
||||
auto cpu_stats = memory_utils::Alloc(CPUPlace(), sizeof(int64_t) * 3);
|
||||
int64_t* cpu_stats_ptr = reinterpret_cast<int64_t*>(cpu_stats->ptr());
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
cpu_stats_ptr,
|
||||
stats.place(),
|
||||
stats.data(),
|
||||
3 * sizeof(int64_t),
|
||||
dev_ctx.stream());
|
||||
dev_ctx.Wait();
|
||||
if (cpu_stats_ptr[0] > 0 || cpu_stats_ptr[1] > 0) {
|
||||
const std::string debug_info =
|
||||
GetHintString<T>(op_type, var_name, stats.place(), dev_id);
|
||||
funcs::PrintAndThrowError(debug_info.c_str(),
|
||||
cpu_stats_ptr[0],
|
||||
cpu_stats_ptr[1],
|
||||
cpu_stats_ptr[2]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename MT>
|
||||
static void WriteToOutputDir(const GPUContext& dev_ctx,
|
||||
const DenseTensor& tensor,
|
||||
const DenseTensor& stats,
|
||||
const DenseTensor& values,
|
||||
const std::string& op_type,
|
||||
const std::string& var_name,
|
||||
const std::string& output_dir,
|
||||
const int check_nan_inf_level) {
|
||||
// Copy stats and values from GPU to CPU.
|
||||
DenseTensor cpu_stats;
|
||||
cpu_stats.Resize({static_cast<int64_t>(3)});
|
||||
Copy(dev_ctx, stats, CPUPlace(), false, &cpu_stats);
|
||||
|
||||
DenseTensor cpu_values;
|
||||
cpu_values.Resize({static_cast<int64_t>(3)});
|
||||
Copy(dev_ctx, values, CPUPlace(), false, &cpu_values);
|
||||
dev_ctx.Wait();
|
||||
|
||||
int dev_id = tensor.place().device;
|
||||
const std::string debug_info =
|
||||
GetHintString<T>(op_type, var_name, tensor.place(), dev_id);
|
||||
std::string log_name = "gpu." + std::to_string(dev_id);
|
||||
int64_t* cpu_stats_ptr = cpu_stats.data<int64_t>();
|
||||
float* cpu_values_ptr = cpu_values.data<float>();
|
||||
funcs::WriteToFileForDifferentLevel<T, MT>(debug_info.c_str(),
|
||||
tensor.numel(),
|
||||
cpu_stats_ptr[0],
|
||||
cpu_stats_ptr[1],
|
||||
cpu_stats_ptr[2],
|
||||
cpu_values_ptr[0],
|
||||
cpu_values_ptr[1],
|
||||
cpu_values_ptr[2],
|
||||
check_nan_inf_level,
|
||||
log_name,
|
||||
output_dir);
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CheckNumericsKernel(const Context& dev_ctx,
|
||||
const DenseTensor& tensor,
|
||||
const std::string& op_type,
|
||||
const std::string& var_name,
|
||||
const int check_nan_inf_level,
|
||||
const int stack_height_limit,
|
||||
const std::string& output_dir,
|
||||
DenseTensor* stats,
|
||||
DenseTensor* values) {
|
||||
int dev_id = tensor.place().device;
|
||||
VLOG(6) << "op_type=" << op_type << ", var_name=" << var_name
|
||||
<< ", dev_id=gpu:" << dev_id << ", numel=" << tensor.numel()
|
||||
<< ", stack_height_limit=" << stack_height_limit
|
||||
<< ", output_dir=" << output_dir;
|
||||
|
||||
if (tensor.numel() <= 0) return;
|
||||
|
||||
// Print to the standard output.
|
||||
char* gpu_str_ptr =
|
||||
GetGpuHintStringPtr<T>(dev_ctx, op_type, var_name, dev_id);
|
||||
|
||||
const size_t threads = 1024;
|
||||
size_t blocks =
|
||||
std::min(static_cast<size_t>(128),
|
||||
static_cast<size_t>((tensor.numel() + threads - 1) / threads));
|
||||
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
|
||||
int64_t numel_max_min = blocks;
|
||||
|
||||
DenseTensor block_num_nan_inf_zero;
|
||||
block_num_nan_inf_zero.Resize({static_cast<int64_t>(3 * numel_max_min)});
|
||||
int64_t* block_num_nan_ptr =
|
||||
dev_ctx.template Alloc<int64_t>(&block_num_nan_inf_zero);
|
||||
int64_t* block_num_inf_ptr = block_num_nan_ptr + numel_max_min;
|
||||
int64_t* block_num_zero_ptr = block_num_inf_ptr + numel_max_min;
|
||||
|
||||
DenseTensor tensor_block_max_min;
|
||||
tensor_block_max_min.Resize({static_cast<int64_t>(3 * numel_max_min)});
|
||||
MT* tensor_block_max_ptr = dev_ctx.template Alloc<MT>(&tensor_block_max_min);
|
||||
MT* tensor_block_min_ptr = tensor_block_max_ptr + numel_max_min;
|
||||
MT* tensor_block_mean_ptr = tensor_block_max_ptr + 2 * numel_max_min;
|
||||
|
||||
FindNanInfAndBlockMaxMin<T, MT>
|
||||
<<<blocks, threads, 0, dev_ctx.stream()>>>(tensor.data<T>(),
|
||||
tensor.numel(),
|
||||
block_num_nan_ptr,
|
||||
block_num_inf_ptr,
|
||||
block_num_zero_ptr,
|
||||
tensor_block_max_ptr,
|
||||
tensor_block_min_ptr,
|
||||
tensor_block_mean_ptr);
|
||||
|
||||
// stats stores the checking result of num_nan, num_inf and num_zero.
|
||||
stats->Resize({static_cast<int64_t>(3)});
|
||||
int64_t* stats_ptr = dev_ctx.template Alloc<int64_t>(stats);
|
||||
|
||||
// values stores the max_value, min_value and mean_value.
|
||||
values->Resize({static_cast<int64_t>(3)});
|
||||
float* values_ptr = dev_ctx.template Alloc<float>(values);
|
||||
|
||||
FindGlobalMaxMinAndPrint<T, MT>
|
||||
<<<1, 1, 0, dev_ctx.stream()>>>(block_num_nan_ptr,
|
||||
block_num_inf_ptr,
|
||||
block_num_zero_ptr,
|
||||
tensor_block_max_ptr,
|
||||
tensor_block_min_ptr,
|
||||
tensor_block_mean_ptr,
|
||||
gpu_str_ptr,
|
||||
tensor.numel(),
|
||||
numel_max_min,
|
||||
check_nan_inf_level,
|
||||
stats_ptr,
|
||||
values_ptr);
|
||||
|
||||
if (output_dir.size() > 0) {
|
||||
// Write log to output_dir.
|
||||
WriteToOutputDir<T, MT>(dev_ctx,
|
||||
tensor,
|
||||
*stats,
|
||||
*values,
|
||||
op_type,
|
||||
var_name,
|
||||
output_dir,
|
||||
check_nan_inf_level);
|
||||
}
|
||||
|
||||
if (check_nan_inf_level == 0 && stack_height_limit > 0) {
|
||||
PrintStack<T>(dev_ctx, *stats, op_type, var_name, dev_id);
|
||||
}
|
||||
}
|
||||
#ifdef _WIN32
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(float, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(double, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(float16, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(bfloat16, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(complex64, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(complex128, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(float8_e4m3fn, GPUContext)
|
||||
INSTANTIATE_CHECKNUMBERICS_KERNEL(float8_e5m2, GPUContext)
|
||||
#endif
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(check_numerics,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CheckNumericsKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128,
|
||||
phi::float8_e4m3fn,
|
||||
phi::float8_e5m2) {}
|
||||
@@ -0,0 +1,22 @@
|
||||
// Copyright (c) 2022 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/cholesky_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/cholesky_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
cholesky_grad, GPU, ALL_LAYOUT, phi::CholeskyGradKernel, float, double) {}
|
||||
@@ -0,0 +1,317 @@
|
||||
/* Copyright (c) 2022 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. */
|
||||
|
||||
#ifndef PADDLE_WITH_HIP
|
||||
// HIP not support cusolver
|
||||
|
||||
#include "paddle/phi/kernels/cholesky_kernel.h"
|
||||
|
||||
#include <thrust/device_vector.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/phi/backends/dynload/cusolver.h"
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/for_range.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T>
|
||||
struct MatrixBandPartFunctor {
|
||||
/*! Set output as input value outside a central band and 0 inside that band.
|
||||
* That is: output[i, j, ..., m, n] = in_band(m, n) * input[i, j, ..., m, n]
|
||||
* where: in_band(m, n) = (num_lower < 0 || (m-n) <= num_lower)) && (num_upper
|
||||
* < 0 || (n-m) <= num_upper)
|
||||
*/
|
||||
MatrixBandPartFunctor(const int m,
|
||||
const int n,
|
||||
const int num_lower_diags,
|
||||
const int num_upper_diags,
|
||||
const T* input,
|
||||
T* output)
|
||||
: m_(m),
|
||||
n_(n),
|
||||
num_lower_diags_(num_lower_diags),
|
||||
num_upper_diags_(num_upper_diags),
|
||||
input_(input),
|
||||
output_(output) {}
|
||||
|
||||
HOSTDEVICE void operator()(size_t index) const {
|
||||
const int col = index % n_;
|
||||
const int row = (index / n_) % m_;
|
||||
const int band_start = (num_lower_diags_ < 0 ? 0 : row - num_lower_diags_);
|
||||
const int band_end =
|
||||
(num_upper_diags_ < 0 ? n_ : row + num_upper_diags_ + 1);
|
||||
if (col < band_start || col >= band_end) {
|
||||
output_[index] = static_cast<T>(0);
|
||||
} else {
|
||||
output_[index] = input_[index];
|
||||
}
|
||||
}
|
||||
|
||||
const int m_, n_, num_lower_diags_, num_upper_diags_;
|
||||
const T* input_;
|
||||
T* output_;
|
||||
};
|
||||
|
||||
#define FUNC_WITH_TYPES(m) m(float, S) m(double, D)
|
||||
|
||||
#define POTRF_INSTANCE(T, C) \
|
||||
void Potrf(const GPUContext& dev_ctx, \
|
||||
cublasFillMode_t uplo, \
|
||||
int n, \
|
||||
T* A, \
|
||||
int lda, \
|
||||
int* info) { \
|
||||
auto handle = dev_ctx.cusolver_dn_handle(); \
|
||||
int workspace_size = 0; \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDn##C##potrf_bufferSize( \
|
||||
handle, uplo, n, A, lda, &workspace_size)); \
|
||||
auto workspace = memory_utils::Alloc( \
|
||||
dev_ctx.GetPlace(), \
|
||||
workspace_size * sizeof(T), \
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream()))); \
|
||||
T* workspace_ptr = reinterpret_cast<T*>(workspace->ptr()); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDn##C##potrf( \
|
||||
handle, uplo, n, A, lda, workspace_ptr, workspace_size, info)); \
|
||||
}
|
||||
|
||||
FUNC_WITH_TYPES(POTRF_INSTANCE);
|
||||
|
||||
#if CUDA_VERSION >= 11040
|
||||
#define POTRF64_INSTANCE(T, C) \
|
||||
void Potrf64(const GPUContext& dev_ctx, \
|
||||
cublasFillMode_t uplo, \
|
||||
int64_t n, \
|
||||
T* A, \
|
||||
int64_t lda, \
|
||||
int* info) { \
|
||||
auto handle = dev_ctx.cusolver_dn_handle(); \
|
||||
cusolverDnParams_t params; \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnCreateParams(¶ms)); \
|
||||
size_t workspace_device_size = 0; \
|
||||
size_t workspace_host_size = 0; \
|
||||
cudaDataType_t data_type = \
|
||||
std::is_same<T, float>::value ? CUDA_R_32F : CUDA_R_64F; \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
dynload::cusolverDnXpotrf_bufferSize(handle, \
|
||||
params, \
|
||||
uplo, \
|
||||
n, \
|
||||
data_type, \
|
||||
A, \
|
||||
lda, \
|
||||
data_type, \
|
||||
&workspace_device_size, \
|
||||
&workspace_host_size)); \
|
||||
auto workspace_device = memory_utils::Alloc( \
|
||||
dev_ctx.GetPlace(), \
|
||||
workspace_device_size, \
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream()))); \
|
||||
auto workspace_host = \
|
||||
memory_utils::Alloc(CPUPlace(), workspace_host_size); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
dynload::cusolverDnXpotrf(handle, \
|
||||
params, \
|
||||
uplo, \
|
||||
n, \
|
||||
data_type, \
|
||||
A, \
|
||||
lda, \
|
||||
data_type, \
|
||||
workspace_device->ptr(), \
|
||||
workspace_device_size, \
|
||||
workspace_host->ptr(), \
|
||||
workspace_host_size, \
|
||||
info)); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDestroyParams(params)); \
|
||||
}
|
||||
|
||||
FUNC_WITH_TYPES(POTRF64_INSTANCE);
|
||||
#endif
|
||||
|
||||
#if CUDA_VERSION >= 9020 && !defined(_WIN32)
|
||||
#define POTRF_BATCH_INSTANCE(T, C) \
|
||||
void PotrfBatched(const GPUContext& dev_ctx, \
|
||||
cublasFillMode_t uplo, \
|
||||
int n, \
|
||||
T* Aarray[], \
|
||||
int lda, \
|
||||
int* info_array, \
|
||||
int batch_size) { \
|
||||
auto handle = dev_ctx.cusolver_dn_handle(); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDn##C##potrfBatched( \
|
||||
handle, uplo, n, Aarray, lda, info_array, batch_size)); \
|
||||
}
|
||||
|
||||
FUNC_WITH_TYPES(POTRF_BATCH_INSTANCE);
|
||||
#endif
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CholeskyKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
bool upper,
|
||||
DenseTensor* out) {
|
||||
if (x.numel() == 0) {
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
return;
|
||||
}
|
||||
|
||||
auto& dims = x.dims();
|
||||
int64_t batch_count64 = 1;
|
||||
for (int i = 0; i < dims.size() - 2; i++) {
|
||||
batch_count64 *= dims[i];
|
||||
}
|
||||
// TODO(large-tensor): cusolver batch_count not support int64
|
||||
PADDLE_ENFORCE_LE_INT_MAX(batch_count64, "batch_count");
|
||||
int batch_count = static_cast<int>(batch_count64);
|
||||
|
||||
int64_t m = dims[dims.size() - 1];
|
||||
// TODO(large-tensor): cusolver n not support int64
|
||||
PADDLE_ENFORCE_LE_INT_MAX(m, "m");
|
||||
int m_int = static_cast<int>(m);
|
||||
|
||||
int64_t tensor_size = batch_count * static_cast<int64_t>(m) * m;
|
||||
|
||||
const auto* x_data = x.data<T>();
|
||||
auto* out_data = dev_ctx.template Alloc<T>(out);
|
||||
|
||||
// matrices are assumed to be stored in column-major order in cusolver
|
||||
cublasFillMode_t uplo =
|
||||
upper ? CUBLAS_FILL_MODE_LOWER : CUBLAS_FILL_MODE_UPPER;
|
||||
// portf is inplace, thus copy the triangular part of the input matrices to
|
||||
// the output and set the other triangular part to 0 firstly
|
||||
|
||||
funcs::ForRange<GPUContext> for_range(dev_ctx, tensor_size);
|
||||
// Pre-processing
|
||||
if (upper) {
|
||||
MatrixBandPartFunctor<T> matrix_band_part_functor(
|
||||
m, m, 0, -1, x_data, out_data);
|
||||
for_range(matrix_band_part_functor);
|
||||
} else {
|
||||
MatrixBandPartFunctor<T> matrix_band_part_functor(
|
||||
m, m, -1, 0, x_data, out_data);
|
||||
for_range(matrix_band_part_functor);
|
||||
}
|
||||
|
||||
auto info =
|
||||
memory_utils::Alloc(dev_ctx.GetPlace(),
|
||||
sizeof(int) * batch_count,
|
||||
Stream(reinterpret_cast<StreamId>(dev_ctx.stream())));
|
||||
auto* info_ptr = reinterpret_cast<int*>(info->ptr());
|
||||
|
||||
#if CUDA_VERSION >= 9020 && !defined(_WIN32)
|
||||
if (batch_count > 1) {
|
||||
std::vector<T*> output_ptrs;
|
||||
for (int i = 0; i < batch_count; i++) {
|
||||
output_ptrs.emplace_back(out_data + static_cast<int64_t>(i) * m * m);
|
||||
}
|
||||
thrust::device_vector<T*> dev_output_ptrs(output_ptrs.begin(),
|
||||
output_ptrs.end());
|
||||
PotrfBatched(dev_ctx,
|
||||
uplo,
|
||||
m_int,
|
||||
thrust::raw_pointer_cast(dev_output_ptrs.data()),
|
||||
m_int,
|
||||
info_ptr,
|
||||
batch_count);
|
||||
// TODO(guosheng): There seems to a bug in cusolver potrfBatched and need
|
||||
// to clear the upper triangle of the output. Remove this workaround once
|
||||
// the bug is fixed.
|
||||
|
||||
if (!upper) {
|
||||
MatrixBandPartFunctor<T> matrix_band_part_functor(
|
||||
m, m, -1, 0, out_data, out_data);
|
||||
for_range(matrix_band_part_functor);
|
||||
}
|
||||
} else {
|
||||
#endif
|
||||
for (int i = 0; i < batch_count; i++) {
|
||||
int64_t offset = static_cast<int64_t>(i) * m * m;
|
||||
#if CUDA_VERSION >= 11040
|
||||
Potrf64(dev_ctx, uplo, m_int, out_data + offset, m_int, info_ptr + i);
|
||||
#else
|
||||
Potrf(dev_ctx, uplo, m_int, out_data + offset, m_int, info_ptr + i);
|
||||
#endif
|
||||
}
|
||||
#if CUDA_VERSION >= 9020 && !defined(_WIN32)
|
||||
}
|
||||
#endif
|
||||
// check the info
|
||||
PADDLE_ENFORCE_EQ(
|
||||
backends::gpu::IsCUDAGraphCapturing(),
|
||||
false,
|
||||
common::errors::InvalidArgument(
|
||||
"CholeskyKernel does not support CUDA Graph capture: async D2H copy "
|
||||
"to local vector 'error_info' will bake the destination address into "
|
||||
"the graph; on replay the vector is re-created at a different "
|
||||
"address, causing a dangling-pointer write."));
|
||||
std::vector<int> error_info;
|
||||
error_info.resize(batch_count);
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
error_info.data(),
|
||||
dev_ctx.GetPlace(),
|
||||
info_ptr,
|
||||
sizeof(int) * batch_count,
|
||||
dev_ctx.stream());
|
||||
|
||||
for (int i = 0; i < batch_count; ++i) {
|
||||
const int info = error_info[i];
|
||||
if (info == 0) {
|
||||
continue;
|
||||
}
|
||||
if (info < 0) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
info,
|
||||
0,
|
||||
errors::InvalidArgument("Cholesky kernel failed for batch %d: "
|
||||
"The %d-th argument was invalid, please "
|
||||
"check the kernel implementation.",
|
||||
i,
|
||||
-info));
|
||||
}
|
||||
PADDLE_ENFORCE_EQ(
|
||||
info,
|
||||
0,
|
||||
errors::PreconditionNotMet(
|
||||
"Cholesky decomposition failed for batch %d: "
|
||||
"The leading minor of order %d is not positive definite.",
|
||||
i,
|
||||
info));
|
||||
}
|
||||
|
||||
// Post-processing to clear the other triangle
|
||||
if (upper) {
|
||||
MatrixBandPartFunctor<T> band_part_post(m, m, 0, -1, out_data, out_data);
|
||||
for_range(band_part_post);
|
||||
} else {
|
||||
MatrixBandPartFunctor<T> band_part_post(m, m, -1, 0, out_data, out_data);
|
||||
for_range(band_part_post);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(cholesky, // cuda_only
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CholeskyKernel,
|
||||
float,
|
||||
double) {}
|
||||
|
||||
#endif // not PADDLE_WITH_HIP
|
||||
@@ -0,0 +1,24 @@
|
||||
// Copyright (c) 2022 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/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/cholesky_solve_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(cholesky_solve_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::CholeskySolveGradKernel,
|
||||
float,
|
||||
double) {}
|
||||
@@ -0,0 +1,181 @@
|
||||
// Copyright (c) 2022 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.
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
#include "paddle/phi/backends/dynload/rocsolver.h"
|
||||
#else
|
||||
#include "paddle/phi/backends/dynload/cusolver.h"
|
||||
#endif
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/lapack/lapack_function.h"
|
||||
#include "paddle/phi/kernels/impl/cholesky_solve_kernel_impl.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
template <typename T>
|
||||
void rocsolver_potrs(const solverHandle_t &handle,
|
||||
rocblas_fill uplo,
|
||||
int M,
|
||||
int N,
|
||||
T *Adata,
|
||||
int lda,
|
||||
T *Bdata,
|
||||
int ldb);
|
||||
|
||||
using dtype::complex;
|
||||
#define FUNC_WITH_TYPES(m) \
|
||||
m(float, s, float) m(double, d, double) \
|
||||
m(complex<float>, c, rocblas_float_complex) \
|
||||
m(complex<double>, z, rocblas_double_complex)
|
||||
|
||||
#define POTRS_INSTANCE(T, C, CastType) \
|
||||
template <> \
|
||||
void rocsolver_potrs<T>(const solverHandle_t &handle, \
|
||||
rocblas_fill uplo, \
|
||||
int M, \
|
||||
int N, \
|
||||
T *Adata, \
|
||||
int lda, \
|
||||
T *Bdata, \
|
||||
int ldb) { \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
dynload::rocsolver_##C##potrs(handle, \
|
||||
uplo, \
|
||||
M, \
|
||||
N, \
|
||||
reinterpret_cast<CastType *>(Adata), \
|
||||
lda, \
|
||||
reinterpret_cast<CastType *>(Bdata), \
|
||||
ldb)); \
|
||||
}
|
||||
|
||||
FUNC_WITH_TYPES(POTRS_INSTANCE);
|
||||
|
||||
#else
|
||||
template <typename T>
|
||||
void cusolver_potrs(const solverHandle_t &handle,
|
||||
cublasFillMode_t uplo,
|
||||
int M,
|
||||
int N,
|
||||
T *Adata,
|
||||
int lda,
|
||||
T *Bdata,
|
||||
int ldb,
|
||||
int *devInfo);
|
||||
|
||||
template <>
|
||||
void cusolver_potrs<float>(const solverHandle_t &handle,
|
||||
cublasFillMode_t uplo,
|
||||
int M,
|
||||
int N,
|
||||
float *Adata,
|
||||
int lda,
|
||||
float *Bdata,
|
||||
int ldb,
|
||||
int *devInfo) {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnSpotrs(
|
||||
handle, uplo, M, N, Adata, lda, Bdata, ldb, devInfo));
|
||||
}
|
||||
|
||||
template <>
|
||||
void cusolver_potrs<double>(const solverHandle_t &handle,
|
||||
cublasFillMode_t uplo,
|
||||
int M,
|
||||
int N,
|
||||
double *Adata,
|
||||
int lda,
|
||||
double *Bdata,
|
||||
int ldb,
|
||||
int *devInfo) {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnDpotrs(
|
||||
handle, uplo, M, N, Adata, lda, Bdata, ldb, devInfo));
|
||||
}
|
||||
|
||||
template <>
|
||||
void cusolver_potrs<complex64>(const solverHandle_t &handle,
|
||||
cublasFillMode_t uplo,
|
||||
int M,
|
||||
int N,
|
||||
complex64 *Adata,
|
||||
int lda,
|
||||
complex64 *Bdata,
|
||||
int ldb,
|
||||
int *devInfo) {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cusolverDnCpotrs(handle,
|
||||
uplo,
|
||||
M,
|
||||
N,
|
||||
reinterpret_cast<const cuComplex *>(Adata),
|
||||
lda,
|
||||
reinterpret_cast<cuComplex *>(Bdata),
|
||||
ldb,
|
||||
devInfo));
|
||||
}
|
||||
|
||||
template <>
|
||||
void cusolver_potrs<complex128>(const cusolverDnHandle_t &handle,
|
||||
cublasFillMode_t uplo,
|
||||
int M,
|
||||
int N,
|
||||
complex128 *Adata,
|
||||
int lda,
|
||||
complex128 *Bdata,
|
||||
int ldb,
|
||||
int *devInfo) {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cusolverDnZpotrs(
|
||||
handle,
|
||||
uplo,
|
||||
M,
|
||||
N,
|
||||
reinterpret_cast<const cuDoubleComplex *>(Adata),
|
||||
lda,
|
||||
reinterpret_cast<cuDoubleComplex *>(Bdata),
|
||||
ldb,
|
||||
devInfo));
|
||||
}
|
||||
|
||||
#endif // PADDLE_WITH_HIP
|
||||
|
||||
template <typename T>
|
||||
class CholeskySolveFunctor<T, GPUContext> {
|
||||
public:
|
||||
void operator()(const GPUContext &dev_ctx,
|
||||
bool upper,
|
||||
int M,
|
||||
int N,
|
||||
T *Adata,
|
||||
int lda,
|
||||
T *Bdata,
|
||||
int *devInfo) {
|
||||
auto handle = dev_ctx.cusolver_dn_handle();
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
rocblas_fill uplo = upper ? rocblas_fill_upper : rocblas_fill_lower;
|
||||
rocsolver_potrs<T>(handle, uplo, M, N, Adata, lda, Bdata, lda);
|
||||
#else
|
||||
cublasFillMode_t uplo =
|
||||
upper ? CUBLAS_FILL_MODE_UPPER : CUBLAS_FILL_MODE_LOWER;
|
||||
cusolver_potrs<T>(handle, uplo, M, N, Adata, lda, Bdata, lda, devInfo);
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
cholesky_solve, GPU, ALL_LAYOUT, phi::CholeskySolveKernel, float, double) {}
|
||||
@@ -0,0 +1,586 @@
|
||||
// Copyright (c) 2024 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.
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
#include <hiprand.h>
|
||||
#include <hiprand_kernel.h>
|
||||
typedef hiprandState curandState;
|
||||
#else
|
||||
#include <curand.h>
|
||||
#include <curand_kernel.h>
|
||||
#endif
|
||||
|
||||
#include <iterator>
|
||||
#include <random>
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/kernels/class_center_sample_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/common/flags.h"
|
||||
#include "paddle/phi/core/distributed/nccl_comm_context.h"
|
||||
#endif
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
|
||||
namespace phi {
|
||||
#define CUDA_KERNEL_LOOP_TYPE(i, n, index_type) \
|
||||
for (index_type i = blockIdx.x * blockDim.x + threadIdx.x, \
|
||||
step = blockDim.x * gridDim.x; \
|
||||
i < (n); \
|
||||
i += step)
|
||||
|
||||
#define CUDA_KERNEL_LOOP(i, n) CUDA_KERNEL_LOOP_TYPE(i, n, int32_t)
|
||||
|
||||
static constexpr int kNumCUDAThreads = 512;
|
||||
static constexpr int kNumMaximumNumBlocks = 4096;
|
||||
|
||||
inline int32_t NumBlocks(const int32_t n) {
|
||||
return std::min((n + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void RandomSampleClassCenter(const int64_t n,
|
||||
int64_t seed,
|
||||
int64_t increment,
|
||||
const int64_t max_val,
|
||||
T* buffer) {
|
||||
const int id = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
curandState localState;
|
||||
size_t local_seed =
|
||||
(static_cast<size_t>(seed) + 0x9E3779B9U +
|
||||
(static_cast<size_t>(id) << 6U) + (static_cast<size_t>(id) >> 2U));
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
hiprand_init(local_seed, id, increment, &localState);
|
||||
CUDA_KERNEL_LOOP(i, n) {
|
||||
buffer[i] = static_cast<T>(hiprand(&localState) % max_val);
|
||||
}
|
||||
#else
|
||||
curand_init(local_seed, id, increment, &localState);
|
||||
CUDA_KERNEL_LOOP(i, n) {
|
||||
buffer[i] = static_cast<T>(curand(&localState) % max_val);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void Range(const int64_t n, T* out) {
|
||||
CUDA_KERNEL_LOOP(i, n) { out[i] = static_cast<T>(i); }
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void MarkPositiveClassCenter(const int64_t n,
|
||||
const int64_t rank,
|
||||
const T* class_interval_ptr,
|
||||
const int num_classes,
|
||||
const T* labels,
|
||||
T* out) {
|
||||
CUDA_KERNEL_LOOP(i, n) {
|
||||
T label = labels[i] - class_interval_ptr[rank];
|
||||
if (label >= 0 && label < num_classes) {
|
||||
out[label] = label - num_classes;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ void FindIntervalIndex(const T* class_interval_ptr,
|
||||
const int64_t nranks,
|
||||
const T value,
|
||||
int64_t* find_index) {
|
||||
int64_t start = 0;
|
||||
int64_t end = nranks;
|
||||
int64_t mid = ((end - start) >> 1) + start + 1;
|
||||
while (start < end) {
|
||||
if (class_interval_ptr[mid] == value) break;
|
||||
if (class_interval_ptr[mid] > value)
|
||||
end = mid - 1;
|
||||
else
|
||||
start = mid;
|
||||
mid = ((end - start) >> 1) + start + 1;
|
||||
}
|
||||
*find_index = min(mid, end);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void GetClassCenterBound(const int64_t n,
|
||||
const int64_t nranks,
|
||||
const T* class_interval_ptr,
|
||||
const T* key_ptr,
|
||||
const T* value_ptr,
|
||||
T* bound_index,
|
||||
T* bound_value) {
|
||||
CUDA_KERNEL_LOOP(i, n) {
|
||||
if (i != 0) {
|
||||
int64_t cur_index, pre_index;
|
||||
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[i], &cur_index);
|
||||
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[i - 1], &pre_index);
|
||||
if (cur_index > pre_index) {
|
||||
assert(cur_index < nranks);
|
||||
#pragma unroll
|
||||
for (int32_t j = pre_index + 1; j <= cur_index; ++j) {
|
||||
bound_index[j] = static_cast<T>(i);
|
||||
bound_value[j] = value_ptr[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
CUDA_KERNEL_LOOP(i, nranks + 1) {
|
||||
int64_t first_index, last_index;
|
||||
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[0], &first_index);
|
||||
FindIntervalIndex(class_interval_ptr, nranks, key_ptr[n - 1], &last_index);
|
||||
if (i <= first_index) {
|
||||
bound_index[i] = 0;
|
||||
bound_value[i] = value_ptr[0];
|
||||
} else if (i > last_index) {
|
||||
bound_index[i] = n;
|
||||
bound_value[i] = value_ptr[n - 1] + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void GetRemappedLabel(const int64_t n,
|
||||
const int64_t nranks,
|
||||
const T* sampled_class_interval_ptr,
|
||||
const T* bound_index,
|
||||
const T* bound_value,
|
||||
const T* label_map_key,
|
||||
T* label_map_value,
|
||||
T* mapped_label) {
|
||||
CUDA_KERNEL_LOOP(i, n) {
|
||||
#pragma unroll
|
||||
for (int64_t j = 0; j < nranks; j++) {
|
||||
if (i >= bound_index[j] && i < bound_index[j + 1]) {
|
||||
label_map_value[i] =
|
||||
label_map_value[i] - bound_value[j] + sampled_class_interval_ptr[j];
|
||||
}
|
||||
}
|
||||
mapped_label[label_map_key[i]] = label_map_value[i];
|
||||
}
|
||||
}
|
||||
|
||||
// aligned vector generates vectorized load/store on CUDA
|
||||
template <typename T, int Size>
|
||||
struct alignas(sizeof(T) * Size) AlignedVector {
|
||||
T val[Size];
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
inline int VectorizedSize(const T* pointer) {
|
||||
uint64_t address = reinterpret_cast<uint64_t>(pointer);
|
||||
constexpr int vec4 = std::alignment_of<AlignedVector<T, 4>>::value; // NOLINT
|
||||
if (address % vec4 == 0) {
|
||||
return 4;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
#undef CUDA_KERNEL_LOOP
|
||||
|
||||
template <typename T>
|
||||
class NotEqualToPreviousAdjacentIterator {
|
||||
public:
|
||||
using self_type = NotEqualToPreviousAdjacentIterator;
|
||||
using value_type = T;
|
||||
using difference_type = std::ptrdiff_t;
|
||||
using pointer = T*;
|
||||
using reference = T;
|
||||
using iterator_category = std::input_iterator_tag;
|
||||
|
||||
public:
|
||||
__host__ __device__ __forceinline__
|
||||
NotEqualToPreviousAdjacentIterator(const T* arr, int64_t offset)
|
||||
: arr_(arr), offset_(offset) {}
|
||||
|
||||
__host__ __device__ __forceinline__ reference operator*() const {
|
||||
return offset_ == 0 ? 0 : (arr_[offset_] == arr_[offset_ - 1] ? 0 : 1);
|
||||
}
|
||||
|
||||
template <typename Distance>
|
||||
__host__ __device__ __forceinline__ self_type operator+(Distance n) const {
|
||||
self_type ret(arr_, offset_ + n);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename Distance>
|
||||
__host__ __device__ __forceinline__ self_type operator-(Distance n) const {
|
||||
self_type ret(arr_, offset_ - n);
|
||||
return ret;
|
||||
}
|
||||
|
||||
template <typename Distance>
|
||||
__host__ __device__ __forceinline__ reference operator[](Distance n) const {
|
||||
return *(*this + n);
|
||||
}
|
||||
|
||||
private:
|
||||
const T* arr_;
|
||||
int64_t offset_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct ActualNumSampledFunctor {
|
||||
__host__ __device__ __forceinline__ T operator()(const T& a,
|
||||
const T& b) const {
|
||||
return max(num_samples, (b - a));
|
||||
}
|
||||
T num_samples;
|
||||
explicit ActualNumSampledFunctor(const T num) : num_samples(num) {}
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
class MemoryBuffer {
|
||||
public:
|
||||
MemoryBuffer(const int num_buffer_ele,
|
||||
const int num_temp_ele,
|
||||
const int nranks,
|
||||
const Context& dev_ctx) {
|
||||
offset1 = 0;
|
||||
offset2 = offset1 + num_buffer_ele;
|
||||
offset3 = offset2 + num_buffer_ele;
|
||||
offset4 = offset3 + num_buffer_ele;
|
||||
offset5 = offset4 + num_buffer_ele;
|
||||
offset6 = offset5 + (nranks + 1);
|
||||
offset7 = offset6 + (nranks + 1);
|
||||
offset8 = offset7 + (nranks + 1);
|
||||
offset9 = offset8 + num_temp_ele;
|
||||
|
||||
buffer.Resize({4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele});
|
||||
buffer_ptr = dev_ctx.template Alloc<T>(&buffer);
|
||||
}
|
||||
|
||||
T* cub_sort_keys_ptr() { return buffer_ptr + offset1; }
|
||||
T* cub_sort_keys_out_ptr() { return buffer_ptr + offset2; }
|
||||
T* cub_sort_values_ptr() { return buffer_ptr + offset3; }
|
||||
T* cub_sort_values_out_ptr() { return buffer_ptr + offset4; }
|
||||
T* bound_index_ptr() { return buffer_ptr + offset5; }
|
||||
T* bound_value_ptr() { return buffer_ptr + offset6; }
|
||||
T* class_interval_ptr() { return buffer_ptr + offset7; }
|
||||
void* cub_temp_storage_ptr() {
|
||||
return reinterpret_cast<void*>(buffer_ptr + offset8);
|
||||
}
|
||||
|
||||
private:
|
||||
DenseTensor buffer;
|
||||
T* buffer_ptr;
|
||||
int offset1;
|
||||
int offset2;
|
||||
int offset3;
|
||||
int offset4;
|
||||
int offset5;
|
||||
int offset6;
|
||||
int offset7;
|
||||
int offset8;
|
||||
int offset9;
|
||||
};
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ClassCenterSampleKernel(const Context& dev_ctx,
|
||||
const DenseTensor& label,
|
||||
int num_classes,
|
||||
int num_samples,
|
||||
int ring_id,
|
||||
int rank,
|
||||
int nranks,
|
||||
bool fix_seed,
|
||||
int seed,
|
||||
DenseTensor* remapped_label,
|
||||
DenseTensor* sampled_local_class_center) {
|
||||
PADDLE_ENFORCE_GT(num_classes,
|
||||
0,
|
||||
errors::InvalidArgument(
|
||||
"The value 'num_classes' for Op(class_center_sample) "
|
||||
"must be greater than 0, "
|
||||
"but the value given is %d.",
|
||||
num_classes));
|
||||
|
||||
PADDLE_ENFORCE_GT(num_samples,
|
||||
0,
|
||||
errors::InvalidArgument(
|
||||
"The value 'num_samples' for Op(class_center_sample) "
|
||||
"must be greater than 0, "
|
||||
"but the value given is %d.",
|
||||
num_samples));
|
||||
|
||||
PADDLE_ENFORCE_LE(num_samples,
|
||||
num_classes,
|
||||
errors::InvalidArgument(
|
||||
"The value 'num_samples' for Op(class_center_sample) "
|
||||
"must be less than or equal to %d, "
|
||||
"but the value given is %d.",
|
||||
num_classes,
|
||||
num_samples));
|
||||
|
||||
auto place = dev_ctx.GetPlace();
|
||||
|
||||
int64_t batch_size = label.numel();
|
||||
// TODO(large-tensor): downstream functors may still use int
|
||||
PADDLE_ENFORCE_LE_INT_MAX(label.numel(), "label.numel()");
|
||||
// Algorithm:
|
||||
// We first randomly generate a value in [0, num_classes) on each position
|
||||
// in a array(shape[num_classes]). Then, we mark the element as negative
|
||||
// value in the array according input label. Now, we can sort the array
|
||||
// by ascending to ensure that the positive class center always in the
|
||||
// front of the sorted array. So, we can get the sampled class center
|
||||
// index by sorted keys. Finally, we can get the rempped label by remap
|
||||
// the input label according sampled class center.
|
||||
|
||||
// step 1: Calculate num classes per device using nccl all reduce
|
||||
std::vector<T> shard_dim_vec(nranks + 1, 0);
|
||||
shard_dim_vec[rank + 1] = num_classes;
|
||||
DenseTensor num_classes_per_device;
|
||||
TensorFromVector(shard_dim_vec, dev_ctx, &num_classes_per_device);
|
||||
T* num_classes_per_device_ptr = num_classes_per_device.data<T>();
|
||||
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
if (nranks > 1) {
|
||||
auto stream = dev_ctx.stream();
|
||||
distributed::NCCLCommContext* comm_ctx = nullptr;
|
||||
comm_ctx =
|
||||
static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
|
||||
PADDLE_ENFORCE_NE(comm_ctx,
|
||||
nullptr,
|
||||
common::errors::Unavailable(
|
||||
"NCCLCommContext is nullptr, collective op should "
|
||||
"has ring_id attr."));
|
||||
|
||||
comm_ctx->AllReduce(
|
||||
&num_classes_per_device, num_classes_per_device, ncclSum, stream);
|
||||
backends::gpu::GpuStreamSync(stream);
|
||||
}
|
||||
#endif
|
||||
|
||||
// step 2: Determine temporary device storage requirements
|
||||
int num_buffer_ele = std::max(static_cast<int>(batch_size), num_classes);
|
||||
size_t cub_sort_temp_store_size = 0;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceRadixSort::SortPairs<T, T>(nullptr,
|
||||
cub_sort_temp_store_size,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
nullptr,
|
||||
num_buffer_ele,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
dev_ctx.stream())));
|
||||
|
||||
size_t cub_sum_temp_store_size = 0;
|
||||
NotEqualToPreviousAdjacentIterator<T> unique_counting_iter_temp(nullptr, 0);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveSum<NotEqualToPreviousAdjacentIterator<T>, T*>(
|
||||
nullptr,
|
||||
cub_sum_temp_store_size,
|
||||
unique_counting_iter_temp,
|
||||
nullptr,
|
||||
batch_size,
|
||||
dev_ctx.stream())));
|
||||
|
||||
size_t cub_scan_temp_store_size = 0;
|
||||
ActualNumSampledFunctor<T> actual_num_sampled_op_temp(num_samples);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveScan(nullptr,
|
||||
cub_scan_temp_store_size,
|
||||
num_classes_per_device_ptr,
|
||||
num_classes_per_device_ptr,
|
||||
actual_num_sampled_op_temp,
|
||||
nranks + 1,
|
||||
dev_ctx.stream())));
|
||||
|
||||
size_t cub_temp_storage_bytes =
|
||||
std::max(std::max(cub_sort_temp_store_size, cub_scan_temp_store_size),
|
||||
cub_sum_temp_store_size);
|
||||
int num_temp_ele = cub_temp_storage_bytes / sizeof(T) + 1;
|
||||
PADDLE_ENFORCE_GT(
|
||||
(4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele),
|
||||
0,
|
||||
errors::InvalidArgument(
|
||||
"Illegal memory allocation, total allocated space must be greater "
|
||||
"than 0, "
|
||||
"but received %d."
|
||||
"This is mainly caused by the size of 'label' being too large.",
|
||||
(4 * num_buffer_ele + 3 * (nranks + 1) + num_temp_ele)));
|
||||
|
||||
// step 3: Alloc buffer memory so that we can reuse allocated memory
|
||||
MemoryBuffer<T, Context> memory_buffer =
|
||||
MemoryBuffer<T, Context>(num_buffer_ele, num_temp_ele, nranks, dev_ctx);
|
||||
|
||||
T* cub_sort_keys_ptr = memory_buffer.cub_sort_keys_ptr();
|
||||
T* cub_sort_keys_out_ptr = memory_buffer.cub_sort_keys_out_ptr();
|
||||
T* cub_sort_values_ptr = memory_buffer.cub_sort_values_ptr();
|
||||
T* cub_sort_values_out_ptr = memory_buffer.cub_sort_values_out_ptr();
|
||||
T* bound_index_ptr = memory_buffer.bound_index_ptr();
|
||||
T* bound_value_ptr = memory_buffer.bound_value_ptr();
|
||||
T* class_interval_ptr = memory_buffer.class_interval_ptr();
|
||||
void* cub_temp_storage_ptr = memory_buffer.cub_temp_storage_ptr();
|
||||
|
||||
// step 4: Calculate class interval among nranks
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveSum(cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
num_classes_per_device_ptr,
|
||||
class_interval_ptr,
|
||||
nranks + 1,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 5: random sample negative class center
|
||||
uint64_t seed_data;
|
||||
uint64_t increment;
|
||||
int vec_size = VectorizedSize<T>(cub_sort_keys_ptr);
|
||||
auto offset = ((num_classes - 1) /
|
||||
(NumBlocks(num_classes) * kNumCUDAThreads * vec_size) +
|
||||
1) *
|
||||
vec_size;
|
||||
// auto gen_cuda = DefaultCUDAGenerator(device_id);
|
||||
auto gen_cuda = dev_ctx.GetGenerator();
|
||||
if (!fix_seed) {
|
||||
auto seed_offset = gen_cuda->IncrementOffset(offset);
|
||||
seed_data = seed_offset.first;
|
||||
increment = seed_offset.second;
|
||||
} else {
|
||||
seed_data = seed + rank;
|
||||
increment = offset;
|
||||
}
|
||||
RandomSampleClassCenter<T>
|
||||
<<<NumBlocks(num_classes), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
||||
num_classes, seed_data, increment, num_classes, cub_sort_keys_ptr);
|
||||
|
||||
// step 6: mark positive class center as negative value
|
||||
// fill the sort values to index 0, 1, ..., batch_size-1
|
||||
MarkPositiveClassCenter<T>
|
||||
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
||||
batch_size,
|
||||
rank,
|
||||
class_interval_ptr,
|
||||
num_classes,
|
||||
label.data<T>(),
|
||||
cub_sort_keys_ptr);
|
||||
Range<T><<<NumBlocks(num_buffer_ele), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
||||
num_buffer_ele, cub_sort_values_ptr);
|
||||
|
||||
// step 7: sort class center by ascending, so that positive class center
|
||||
// always be sampled.
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceRadixSort::SortPairs<T, T>(cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
cub_sort_keys_ptr,
|
||||
cub_sort_keys_out_ptr,
|
||||
cub_sort_values_ptr,
|
||||
cub_sort_values_out_ptr,
|
||||
num_classes,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 8: sort input label ascending
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceRadixSort::SortPairs<T, T>(cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
label.data<T>(),
|
||||
cub_sort_keys_out_ptr,
|
||||
cub_sort_values_ptr,
|
||||
cub_sort_keys_ptr,
|
||||
batch_size,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 9: Calculate new index using InclusiveSum on ascending sorted input
|
||||
// label
|
||||
NotEqualToPreviousAdjacentIterator<T> unique_counting_iter(
|
||||
cub_sort_keys_out_ptr, 0);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveSum<NotEqualToPreviousAdjacentIterator<T>, T*>(
|
||||
cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
unique_counting_iter,
|
||||
cub_sort_values_ptr,
|
||||
batch_size,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 10: Calculate new class center bound among ranks
|
||||
GetClassCenterBound<T>
|
||||
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
||||
batch_size,
|
||||
nranks,
|
||||
class_interval_ptr,
|
||||
cub_sort_keys_out_ptr,
|
||||
cub_sort_values_ptr,
|
||||
bound_index_ptr,
|
||||
bound_value_ptr);
|
||||
|
||||
// step 11: Calculate actual number of sampled class per device.
|
||||
// Since maybe num_positive_class_center > num_samples,
|
||||
// we need to ensure all positive class center per device are sampled.
|
||||
ActualNumSampledFunctor<T> actual_num_sampled_op(num_samples);
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveScan(cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
bound_value_ptr,
|
||||
num_classes_per_device_ptr,
|
||||
actual_num_sampled_op,
|
||||
nranks + 1,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 12: Calculate actual sampled class interval among nranks
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
(cub::DeviceScan::InclusiveSum(cub_temp_storage_ptr,
|
||||
cub_temp_storage_bytes,
|
||||
num_classes_per_device_ptr,
|
||||
class_interval_ptr,
|
||||
nranks + 1,
|
||||
dev_ctx.stream())));
|
||||
|
||||
// step 13: Get remapped label for output
|
||||
GetRemappedLabel<T>
|
||||
<<<NumBlocks(batch_size), kNumCUDAThreads, 0, dev_ctx.stream()>>>(
|
||||
batch_size,
|
||||
nranks,
|
||||
class_interval_ptr,
|
||||
bound_index_ptr,
|
||||
bound_value_ptr,
|
||||
cub_sort_keys_ptr,
|
||||
cub_sort_values_ptr,
|
||||
dev_ctx.template Alloc<T>(remapped_label));
|
||||
|
||||
// step 14: Get sampled class center for output
|
||||
Copy<Context>(dev_ctx,
|
||||
num_classes_per_device,
|
||||
CPUPlace(),
|
||||
true,
|
||||
&num_classes_per_device);
|
||||
T actual_num_samples = num_classes_per_device.data<T>()[rank + 1];
|
||||
sampled_local_class_center->Resize({actual_num_samples});
|
||||
|
||||
T* sampled_local_class_center_ptr =
|
||||
dev_ctx.template Alloc<T>(sampled_local_class_center);
|
||||
memory_utils::Copy(dev_ctx.GetPlace(),
|
||||
sampled_local_class_center_ptr,
|
||||
dev_ctx.GetPlace(),
|
||||
cub_sort_values_out_ptr,
|
||||
actual_num_samples * sizeof(T),
|
||||
nullptr);
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(class_center_sample,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ClassCenterSampleKernel,
|
||||
int64_t,
|
||||
int) {}
|
||||
@@ -0,0 +1,83 @@
|
||||
// Copyright (c) 2022 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/clip_by_norm_kernel.h"
|
||||
|
||||
#include <typeinfo>
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
||||
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
||||
#include "paddle/phi/kernels/impl/clip_by_norm_kernel_impl.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ClipByNormKernel(const Context& dev_ctx,
|
||||
const DenseTensor& in,
|
||||
float max_norm,
|
||||
DenseTensor* output) {
|
||||
if (typeid(T) == typeid(float)) {
|
||||
return ClipByNormFunctor<float, Context>(dev_ctx, in, max_norm, output);
|
||||
}
|
||||
auto input = ∈
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
|
||||
PADDLE_ENFORCE_NOT_NULL(input,
|
||||
common::errors::InvalidArgument(
|
||||
"Input(X) of ClipByNormOp should not be null. "
|
||||
"Please check if it is created correctly."));
|
||||
std::vector<int> reduce_dims;
|
||||
reduce_dims.resize(input->dims().size());
|
||||
for (int i = 0; i < reduce_dims.size(); ++i) {
|
||||
reduce_dims[i] = i;
|
||||
}
|
||||
DenseTensor tmp_tensor;
|
||||
auto* tmp = &tmp_tensor;
|
||||
tmp->Resize({1});
|
||||
dev_ctx.template Alloc<float>(tmp);
|
||||
funcs::ReduceKernel<T, float, kps::AddFunctor, kps::SquareFunctor<T, float>>(
|
||||
dev_ctx, *input, tmp, kps::SquareFunctor<T, float>(), reduce_dims);
|
||||
auto tmp_eigen = EigenVector<float>::Flatten(*tmp);
|
||||
auto x_norm = tmp_eigen.sqrt();
|
||||
|
||||
auto x = EigenVector<T>::Flatten(*input);
|
||||
auto out = EigenVector<T>::Flatten(*output);
|
||||
auto* place = dev_ctx.eigen_device();
|
||||
|
||||
auto temp = (x_norm <= max_norm).template cast<float>();
|
||||
auto epsilon =
|
||||
((x_norm <= static_cast<float>(1e-30)).all().template cast<float>()) *
|
||||
static_cast<float>(1e-6);
|
||||
|
||||
auto scaling =
|
||||
(temp + (static_cast<float>(1) - temp) * max_norm / (x_norm + epsilon))
|
||||
.template cast<T>();
|
||||
Eigen::array<int, 1> one_dim{{1}};
|
||||
Eigen::DSizes<int, 1> m_dsize(input->numel());
|
||||
|
||||
out.device(*place) = x * scaling.reshape(one_dim).broadcast(m_dsize);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(clip_by_norm,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ClipByNormKernel,
|
||||
float,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2022 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/clip_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/clip_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(clip_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ClipGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::bfloat16,
|
||||
phi::float16) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2022 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/clip_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/clip_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(clip,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ClipKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
@@ -0,0 +1,288 @@
|
||||
// Copyright (c) 2024 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/gpu/collect_fpn_proposals_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_graph_with_memory_pool.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/allocator.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/mixed_vector.h"
|
||||
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
#include "paddle/phi/kernels/funcs/detection/bbox_util.h"
|
||||
#include "paddle/phi/kernels/funcs/for_range.h"
|
||||
#include "paddle/phi/kernels/funcs/gather.cu.h"
|
||||
#include "paddle/phi/kernels/funcs/strided_memcpy.h"
|
||||
#include "paddle/phi/kernels/impl/collect_fpn_proposals_kernel_impl.h"
|
||||
#include "paddle/utils/optional.h"
|
||||
|
||||
namespace phi {
|
||||
|
||||
static constexpr int kNumCUDAThreads = 64;
|
||||
static constexpr int kNumMaximumNumBlocks = 4096;
|
||||
|
||||
const int kBBoxSize = 4;
|
||||
|
||||
static inline int NumBlocks(const int N) {
|
||||
return std::min((N + kNumCUDAThreads - 1) / kNumCUDAThreads,
|
||||
kNumMaximumNumBlocks);
|
||||
}
|
||||
|
||||
static __global__ void GetLengthLoD(const int nthreads,
|
||||
const int* batch_ids,
|
||||
int* length_lod) {
|
||||
CUDA_KERNEL_LOOP(i, nthreads) { CudaAtomicAdd(length_lod + batch_ids[i], 1); }
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void GPUCollectFpnProposalsOpKernel(
|
||||
const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& multi_level_rois,
|
||||
const std::vector<const DenseTensor*>& multi_level_scores,
|
||||
const optional<std::vector<const DenseTensor*>>& multi_level_rois_num,
|
||||
int post_nms_topn,
|
||||
DenseTensor* fpn_rois_out,
|
||||
DenseTensor* rois_num_out) {
|
||||
const auto roi_ins = multi_level_rois;
|
||||
const auto score_ins = multi_level_scores;
|
||||
auto fpn_rois = fpn_rois_out;
|
||||
|
||||
const int post_nms_topN = post_nms_topn;
|
||||
|
||||
// concat inputs along axis = 0
|
||||
int roi_offset = 0;
|
||||
int score_offset = 0;
|
||||
int total_roi_num = 0;
|
||||
for (size_t i = 0; i < roi_ins.size(); ++i) {
|
||||
total_roi_num += roi_ins[i]->dims()[0];
|
||||
}
|
||||
|
||||
int real_post_num = min(post_nms_topN, total_roi_num);
|
||||
fpn_rois->Resize({real_post_num, kBBoxSize});
|
||||
dev_ctx.template Alloc<T>(fpn_rois);
|
||||
DenseTensor concat_rois;
|
||||
DenseTensor concat_scores;
|
||||
concat_rois.Resize({total_roi_num, kBBoxSize});
|
||||
T* concat_rois_data = dev_ctx.template Alloc<T>(&concat_rois);
|
||||
concat_scores.Resize({total_roi_num, 1});
|
||||
T* concat_scores_data = dev_ctx.template Alloc<T>(&concat_scores);
|
||||
DenseTensor roi_batch_id_list;
|
||||
roi_batch_id_list.Resize({total_roi_num});
|
||||
int* roi_batch_id_data = dev_ctx.template HostAlloc<int>(&roi_batch_id_list);
|
||||
int index = 0;
|
||||
int lod_size;
|
||||
auto place = dev_ctx.GetPlace();
|
||||
|
||||
auto multi_rois_num = multi_level_rois_num
|
||||
? multi_level_rois_num.get()
|
||||
: std::vector<const DenseTensor*>();
|
||||
for (size_t i = 0; i < roi_ins.size(); ++i) {
|
||||
auto roi_in = roi_ins[i];
|
||||
auto score_in = score_ins[i];
|
||||
if (multi_rois_num.size() > 0) {
|
||||
DenseTensor temp;
|
||||
Copy(dev_ctx, *multi_rois_num[i], CPUPlace(), true, &temp);
|
||||
const int* length_in = temp.data<int>();
|
||||
lod_size = multi_rois_num[i]->numel();
|
||||
for (size_t n = 0; n < lod_size; ++n) {
|
||||
for (size_t j = 0; j < length_in[n]; ++j) {
|
||||
roi_batch_id_data[index++] = n;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
auto length_in = roi_in->lod().back();
|
||||
lod_size = length_in.size() - 1;
|
||||
for (size_t n = 0; n < lod_size; ++n) {
|
||||
for (size_t j = length_in[n]; j < length_in[n + 1]; ++j) {
|
||||
roi_batch_id_data[index++] = n;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
memory_utils::Copy(place,
|
||||
concat_rois_data + roi_offset,
|
||||
place,
|
||||
roi_in->data<T>(),
|
||||
roi_in->numel() * sizeof(T),
|
||||
dev_ctx.stream());
|
||||
memory_utils::Copy(place,
|
||||
concat_scores_data + score_offset,
|
||||
place,
|
||||
score_in->data<T>(),
|
||||
score_in->numel() * sizeof(T),
|
||||
dev_ctx.stream());
|
||||
roi_offset += roi_in->numel();
|
||||
score_offset += score_in->numel();
|
||||
}
|
||||
|
||||
// copy batch id list to GPU
|
||||
DenseTensor roi_batch_id_list_gpu;
|
||||
Copy(dev_ctx,
|
||||
roi_batch_id_list,
|
||||
dev_ctx.GetPlace(),
|
||||
false,
|
||||
&roi_batch_id_list_gpu);
|
||||
|
||||
DenseTensor index_in_t;
|
||||
index_in_t.Resize({total_roi_num});
|
||||
int* idx_in = dev_ctx.template Alloc<int>(&index_in_t);
|
||||
funcs::ForRange<GPUContext> for_range_total(dev_ctx, total_roi_num);
|
||||
for_range_total(funcs::RangeInitFunctor{0, 1, idx_in});
|
||||
|
||||
DenseTensor keys_out_t;
|
||||
keys_out_t.Resize({total_roi_num});
|
||||
T* keys_out = dev_ctx.template Alloc<T>(&keys_out_t);
|
||||
DenseTensor index_out_t;
|
||||
index_out_t.Resize({total_roi_num});
|
||||
int* idx_out = dev_ctx.template Alloc<int>(&index_out_t);
|
||||
|
||||
// Determine temporary device storage requirements
|
||||
size_t temp_storage_bytes = 0;
|
||||
cub::DeviceRadixSort::SortPairsDescending<T, int>(nullptr,
|
||||
temp_storage_bytes,
|
||||
concat_scores.data<T>(),
|
||||
keys_out,
|
||||
idx_in,
|
||||
idx_out,
|
||||
total_roi_num,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
dev_ctx.stream());
|
||||
// Allocate temporary storage
|
||||
auto d_temp_storage = memory_utils::Alloc(place, temp_storage_bytes);
|
||||
|
||||
// Run sorting operation
|
||||
// sort score to get corresponding index
|
||||
cub::DeviceRadixSort::SortPairsDescending<T, int>(d_temp_storage->ptr(),
|
||||
temp_storage_bytes,
|
||||
concat_scores.data<T>(),
|
||||
keys_out,
|
||||
idx_in,
|
||||
idx_out,
|
||||
total_roi_num,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
dev_ctx.stream());
|
||||
index_out_t.Resize({real_post_num});
|
||||
DenseTensor sorted_rois;
|
||||
sorted_rois.Resize({real_post_num, kBBoxSize});
|
||||
dev_ctx.template Alloc<T>(&sorted_rois);
|
||||
DenseTensor sorted_batch_id;
|
||||
sorted_batch_id.Resize({real_post_num});
|
||||
dev_ctx.template Alloc<int>(&sorted_batch_id);
|
||||
funcs::GPUGather<T>(dev_ctx, concat_rois, index_out_t, &sorted_rois);
|
||||
funcs::GPUGather<int>(
|
||||
dev_ctx, roi_batch_id_list_gpu, index_out_t, &sorted_batch_id);
|
||||
|
||||
DenseTensor batch_index_t;
|
||||
batch_index_t.Resize({real_post_num});
|
||||
int* batch_idx_in = dev_ctx.template Alloc<int>(&batch_index_t);
|
||||
funcs::ForRange<GPUContext> for_range_post(dev_ctx, real_post_num);
|
||||
for_range_post(funcs::RangeInitFunctor{0, 1, batch_idx_in});
|
||||
|
||||
DenseTensor out_id_t;
|
||||
out_id_t.Resize({real_post_num});
|
||||
int* out_id_data = dev_ctx.template Alloc<int>(&out_id_t);
|
||||
// Determine temporary device storage requirements
|
||||
temp_storage_bytes = 0;
|
||||
cub::DeviceRadixSort::SortPairs<int, int>(nullptr,
|
||||
temp_storage_bytes,
|
||||
sorted_batch_id.data<int>(),
|
||||
out_id_data,
|
||||
batch_idx_in,
|
||||
index_out_t.data<int>(),
|
||||
real_post_num,
|
||||
0,
|
||||
sizeof(int) * 8,
|
||||
dev_ctx.stream());
|
||||
// Allocate temporary storage
|
||||
d_temp_storage = memory_utils::Alloc(place, temp_storage_bytes);
|
||||
|
||||
// Run sorting operation
|
||||
// sort batch_id to get corresponding index
|
||||
cub::DeviceRadixSort::SortPairs<int, int>(d_temp_storage->ptr(),
|
||||
temp_storage_bytes,
|
||||
sorted_batch_id.data<int>(),
|
||||
out_id_data,
|
||||
batch_idx_in,
|
||||
index_out_t.data<int>(),
|
||||
real_post_num,
|
||||
0,
|
||||
sizeof(int) * 8,
|
||||
dev_ctx.stream());
|
||||
|
||||
funcs::GPUGather<T>(dev_ctx, sorted_rois, index_out_t, fpn_rois);
|
||||
|
||||
DenseTensor length_lod;
|
||||
length_lod.Resize({lod_size});
|
||||
int* length_lod_data = dev_ctx.template Alloc<int>(&length_lod);
|
||||
funcs::SetConstant<GPUContext, int> set_zero;
|
||||
set_zero(dev_ctx, &length_lod, static_cast<int>(0));
|
||||
|
||||
int blocks = NumBlocks(real_post_num);
|
||||
int threads = kNumCUDAThreads;
|
||||
|
||||
// get length-based lod by batch ids
|
||||
GetLengthLoD<<<blocks, threads, 0, dev_ctx.stream()>>>(
|
||||
real_post_num, out_id_data, length_lod_data);
|
||||
PADDLE_ENFORCE_EQ(
|
||||
backends::gpu::IsCUDAGraphCapturing(),
|
||||
false,
|
||||
common::errors::InvalidArgument(
|
||||
"CollectFpnProposals does not support CUDA Graph capture: async D2H "
|
||||
"copy to local vector 'length_lod_cpu' will bake the destination "
|
||||
"address into the graph; on replay the vector is re-created at a "
|
||||
"different address, causing a dangling-pointer write."));
|
||||
std::vector<int> length_lod_cpu(lod_size);
|
||||
memory_utils::Copy(CPUPlace(),
|
||||
length_lod_cpu.data(),
|
||||
place,
|
||||
length_lod_data,
|
||||
sizeof(int) * lod_size,
|
||||
dev_ctx.stream());
|
||||
dev_ctx.Wait();
|
||||
|
||||
std::vector<size_t> offset(1, 0);
|
||||
for (int i = 0; i < lod_size; ++i) {
|
||||
offset.emplace_back(offset.back() + length_lod_cpu[i]);
|
||||
}
|
||||
|
||||
if (rois_num_out != nullptr) {
|
||||
auto* rois_num = rois_num_out;
|
||||
rois_num->Resize({lod_size});
|
||||
int* rois_num_data = dev_ctx.template Alloc<int>(rois_num);
|
||||
memory_utils::Copy(place,
|
||||
rois_num_data,
|
||||
place,
|
||||
length_lod_data,
|
||||
lod_size * sizeof(int),
|
||||
dev_ctx.stream());
|
||||
}
|
||||
|
||||
LegacyLoD lod;
|
||||
lod.emplace_back(offset);
|
||||
fpn_rois->set_lod(lod);
|
||||
}
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(collect_fpn_proposals,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::GPUCollectFpnProposalsOpKernel,
|
||||
float,
|
||||
double) {
|
||||
kernel->InputAt(2).SetDataType(phi::DataType::INT32);
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2024 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.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <vector>
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
|
||||
namespace phi {
|
||||
template <typename T, typename Context>
|
||||
void GPUCollectFpnProposalsOpKernel(
|
||||
const Context& dev_ctx,
|
||||
const std::vector<const DenseTensor*>& multi_level_rois,
|
||||
const std::vector<const DenseTensor*>& multi_level_scores,
|
||||
const optional<std::vector<const DenseTensor*>>& multi_level_rois_num,
|
||||
int post_nms_topn,
|
||||
DenseTensor* fpn_rois_out,
|
||||
DenseTensor* rois_num_out);
|
||||
} // namespace phi
|
||||
@@ -0,0 +1,43 @@
|
||||
// Copyright (c) 2024 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/comm_init_all_kernel.h"
|
||||
#include <string>
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
#include "paddle/phi/core/platform/collective_helper.h"
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
template <typename T, typename Context>
|
||||
void CommInitAllKernel(const Context& dev_ctx,
|
||||
const std::vector<int>& devices_input,
|
||||
int ring_id) {
|
||||
#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
|
||||
std::vector<int> devices = devices_input;
|
||||
if (devices.empty()) {
|
||||
devices = backends::gpu::GetSelectedDevices();
|
||||
}
|
||||
|
||||
paddle::platform::NCCLCommContext::Instance().CreateAllNCCLComms(devices,
|
||||
ring_id);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
comm_init_all, GPU, ALL_LAYOUT, phi::CommInitAllKernel, float) {}
|
||||
@@ -0,0 +1,41 @@
|
||||
// Copyright (c) 2022 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/complex_grad_kernel.h"
|
||||
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/complex_grad_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(imag_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ImagGradKernel,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(real_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::RealGradKernel,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
complex_grad, GPU, ALL_LAYOUT, phi::ComplexGradKernel, float, double) {
|
||||
kernel->InputAt(2).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright (c) 2021 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/complex_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/impl/complex_kernel_impl.h"
|
||||
|
||||
PD_REGISTER_KERNEL(conj,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ConjKernel,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t) {}
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
real, GPU, ALL_LAYOUT, phi::RealKernel, phi::complex64, phi::complex128) {
|
||||
kernel->OutputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
imag, GPU, ALL_LAYOUT, phi::ImagKernel, phi::complex64, phi::complex128) {
|
||||
kernel->OutputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(
|
||||
complex, GPU, ALL_LAYOUT, phi::ComplexKernel, float, double) {
|
||||
kernel->OutputAt(0).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user