wehub-resource-sync
549 lines
19 KiB
Plaintext
549 lines
19 KiB
Plaintext
// 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.
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/visit_type.h"
|
|
#include "paddle/phi/kernels/contiguous_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_divide_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/broadcast_function.h"
|
|
#include "paddle/phi/kernels/funcs/dense_tensor_iterator.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
|
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
|
|
#include "paddle/phi/kernels/funcs/index_elementwise.cu.h"
|
|
#include "paddle/phi/kernels/impl/elementwise_kernel_impl.h"
|
|
#include "paddle/phi/kernels/scale_kernel.h"
|
|
#include "paddle/phi/kernels/stride/elementwise_stride_base.cu.h"
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
|
|
#include "paddle/phi/kernels/funcs/dims_simplifier.h"
|
|
|
|
#endif
|
|
|
|
COMMON_DECLARE_bool(use_stride_kernel);
|
|
COMMON_DECLARE_bool(use_stride_compute_kernel);
|
|
COMMON_DECLARE_bool(force_stride_compute_contig_out);
|
|
|
|
namespace phi {
|
|
|
|
inline bool FastContiguous(const int64_t &numel,
|
|
const DDim &shape,
|
|
const DDim &stride,
|
|
const uint64_t &offset) {
|
|
if (offset != 0) {
|
|
return false;
|
|
}
|
|
|
|
// For large tensors (>16M elements), transpose + contiguous elementwise
|
|
// is faster than direct strided elementwise kernel
|
|
if (numel < 16777216LL) {
|
|
return false;
|
|
}
|
|
|
|
if (shape.size() < 2 || stride.size() < 2) {
|
|
return false;
|
|
}
|
|
|
|
auto tmp_shape = shape;
|
|
auto tmp_stride = stride;
|
|
auto vec_size = tmp_shape.size();
|
|
|
|
std::swap(tmp_shape[vec_size - 1], tmp_shape[vec_size - 2]);
|
|
std::swap(tmp_stride[vec_size - 1], tmp_stride[vec_size - 2]);
|
|
|
|
if (!(tmp_stride[vec_size - 1] == 1) ||
|
|
!(tmp_stride[vec_size - 2] == tmp_shape[vec_size - 1])) {
|
|
return false;
|
|
}
|
|
|
|
if (DenseTensorMeta::calc_strides(tmp_shape) == tmp_stride) {
|
|
return true;
|
|
} else {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
#define DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(name, functor_name) \
|
|
template <typename T, typename Context> \
|
|
void name##StrideKernel(const Context &dev_ctx, \
|
|
const DenseTensor &x, \
|
|
const DenseTensor &y, \
|
|
DenseTensor *out) { \
|
|
if (!FLAGS_use_stride_kernel) { \
|
|
PADDLE_THROW(common::errors::Fatal( \
|
|
"FLAGS_use_stride_kernel is closed. Strided kernel " \
|
|
"be called, something wrong has happened!")); \
|
|
} \
|
|
DenseTensor x_; \
|
|
DenseTensor y_; \
|
|
\
|
|
bool fast_contiguous = false; \
|
|
if (FLAGS_force_stride_compute_contig_out) { \
|
|
bool x_fast = \
|
|
FastContiguous(x.numel(), x.dims(), x.strides(), x.offset()); \
|
|
bool y_fast = \
|
|
FastContiguous(y.numel(), y.dims(), y.strides(), y.offset()); \
|
|
fast_contiguous = x_fast || y_fast; \
|
|
} \
|
|
bool zero_size = false; \
|
|
if (x.numel() == 0 || y.numel() == 0) { \
|
|
zero_size = true; \
|
|
} \
|
|
if (!FLAGS_use_stride_compute_kernel || fast_contiguous || zero_size) { \
|
|
if (!x.meta().is_contiguous()) { \
|
|
x_ = Tensor2Contiguous<Context>(dev_ctx, x); \
|
|
} else { \
|
|
x_ = x; \
|
|
} \
|
|
if (!y.meta().is_contiguous()) { \
|
|
y_ = Tensor2Contiguous<Context>(dev_ctx, y); \
|
|
} else { \
|
|
y_ = y; \
|
|
} \
|
|
} else { \
|
|
x_ = x; \
|
|
y_ = y; \
|
|
} \
|
|
if (x_.meta().is_contiguous() && y_.meta().is_contiguous()) { \
|
|
auto meta = out->meta(); \
|
|
meta.strides = meta.calc_strides(out->dims()); \
|
|
out->set_meta(meta); \
|
|
phi::name##Kernel<T, Context>(dev_ctx, x_, y_, out); \
|
|
return; \
|
|
} \
|
|
if (!FLAGS_use_stride_compute_kernel) { \
|
|
PADDLE_THROW( \
|
|
common::errors::Fatal("FLAGS_use_stride_compute_kernel is closed. " \
|
|
"Kernel using DenseTensorIterator " \
|
|
"be called, something wrong has happened!")); \
|
|
} \
|
|
\
|
|
if (FLAGS_force_stride_compute_contig_out) { \
|
|
auto meta = out->meta(); \
|
|
meta.strides = meta.calc_strides(out->dims()); \
|
|
out->set_meta(meta); \
|
|
} \
|
|
LaunchBinaryElementwiseStrideKernel<T, Context>( \
|
|
dev_ctx, x_, y_, funcs::functor_name##Functor<T>(), -1, out); \
|
|
}
|
|
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Subtract, Subtract)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Multiply, Multiply)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Divide, Divide)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(CopySign, CopySign)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Remainder, Remainder)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Maximum, Maximum)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Minimum, Minimum)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(FloorDivide, FloorDivide)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(Heaviside, ElementwiseHeaviside)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(FMax, FMax)
|
|
DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP(FMin, FMin)
|
|
#undef DEFINE_CUDA_BINARY_ELEMENTWISE_STRIDE_OP
|
|
|
|
template <typename T, typename Context>
|
|
void AddStrideKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &y,
|
|
DenseTensor *out) {
|
|
if (!FLAGS_use_stride_kernel) {
|
|
PADDLE_THROW(common::errors::Fatal(
|
|
"FLAGS_use_stride_kernel is closed. Strided kernel "
|
|
"be called, something wrong has happened!"));
|
|
}
|
|
DenseTensor x_;
|
|
DenseTensor y_;
|
|
bool zero_size = false;
|
|
if (x.numel() == 0 || y.numel() == 0) {
|
|
zero_size = true;
|
|
}
|
|
if (!FLAGS_use_stride_compute_kernel || x.dtype() != y.dtype() || zero_size) {
|
|
if (!x.meta().is_contiguous()) {
|
|
x_ = Tensor2Contiguous<Context>(dev_ctx, x);
|
|
} else {
|
|
x_ = x;
|
|
}
|
|
if (!y.meta().is_contiguous()) {
|
|
y_ = Tensor2Contiguous<Context>(dev_ctx, y);
|
|
} else {
|
|
y_ = y;
|
|
}
|
|
} else {
|
|
x_ = x;
|
|
y_ = y;
|
|
}
|
|
if (x_.meta().is_contiguous() && y_.meta().is_contiguous()) {
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
phi::AddKernel<T, Context>(dev_ctx, x_, y_, out);
|
|
return;
|
|
}
|
|
if (!FLAGS_use_stride_compute_kernel) {
|
|
PADDLE_THROW(
|
|
common::errors::Fatal("FLAGS_use_stride_compute_kernel is closed. "
|
|
"Kernel using DenseTensorIterator "
|
|
"be called, something wrong has happened!"));
|
|
}
|
|
|
|
if (FLAGS_force_stride_compute_contig_out) {
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
}
|
|
|
|
LaunchBinaryElementwiseStrideKernel<T, Context>(
|
|
dev_ctx, x_, y_, funcs::AddFunctor<T>(), -1, out);
|
|
}
|
|
|
|
template <typename DataT, typename ParamT>
|
|
struct ScaleFunctor {
|
|
ParamT bias;
|
|
ParamT scale;
|
|
bool bias_after_scale;
|
|
|
|
ScaleFunctor(ParamT scale_data, ParamT bias_data, bool is_bias_after_scale)
|
|
: bias(bias_data),
|
|
scale(scale_data),
|
|
bias_after_scale(is_bias_after_scale) {}
|
|
|
|
__device__ __forceinline__ DataT operator()(const DataT x) const {
|
|
if (bias_after_scale) {
|
|
return static_cast<DataT>(scale * static_cast<ParamT>(x) + bias);
|
|
} else {
|
|
return static_cast<DataT>(scale * (static_cast<ParamT>(x) + bias));
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void ScaleStrideKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const Scalar &scale,
|
|
const Scalar &bias,
|
|
bool bias_after_scale,
|
|
DenseTensor *out) {
|
|
if (!FLAGS_use_stride_kernel) {
|
|
PADDLE_THROW(common::errors::Fatal(
|
|
"FLAGS_use_stride_kernel is closed. Strided kernel "
|
|
"be called, something wrong has happened!"));
|
|
}
|
|
DenseTensor x_;
|
|
bool zero_size = false;
|
|
if (x.numel() == 0) {
|
|
zero_size = true;
|
|
}
|
|
if (!FLAGS_use_stride_compute_kernel || zero_size) {
|
|
if (!x.meta().is_contiguous()) {
|
|
x_ = Tensor2Contiguous<Context>(dev_ctx, x);
|
|
} else {
|
|
x_ = x;
|
|
}
|
|
} else {
|
|
x_ = x;
|
|
}
|
|
if (x_.meta().is_contiguous()) {
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
phi::ScaleKernel<T, Context>(
|
|
dev_ctx, x_, scale, bias, bias_after_scale, out);
|
|
return;
|
|
}
|
|
if (!FLAGS_use_stride_compute_kernel) {
|
|
PADDLE_THROW(
|
|
common::errors::Fatal("FLAGS_use_stride_compute_kernel is closed. "
|
|
"Kernel using DenseTensorIterator "
|
|
"be called, something wrong has happened!"));
|
|
}
|
|
|
|
if (FLAGS_force_stride_compute_contig_out) {
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
}
|
|
|
|
if (x.numel() <= 0 || (!x.IsInitialized())) {
|
|
dev_ctx.template Alloc<T>(out);
|
|
return;
|
|
}
|
|
|
|
using MT = typename phi::dtype::MPTypeTrait<T>::Type;
|
|
LaunchUnaryElementwiseStrideKernel<T, Context>(
|
|
dev_ctx,
|
|
x_,
|
|
ScaleFunctor<T, MT>(scale.to<MT>(), bias.to<MT>(), bias_after_scale),
|
|
out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FullStrideKernel(const Context &dev_ctx,
|
|
const IntArray &shape,
|
|
const Scalar &val,
|
|
DataType dtype,
|
|
DenseTensor *out) {
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
FullKernel<T, Context>(dev_ctx, shape, val, dtype, out);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FullLikeStrideKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const Scalar &val,
|
|
DataType dtype,
|
|
DenseTensor *out) {
|
|
// Is this correct?
|
|
// In fact, both ones_like and full_like can only generate contiguous tensors,
|
|
// which differs from common sense, where both strides and shapes are
|
|
// considered.
|
|
auto meta = out->meta();
|
|
meta.strides = meta.calc_strides(out->dims());
|
|
out->set_meta(meta);
|
|
FullLikeKernel<T, Context>(dev_ctx, x, val, dtype, out);
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
using float16 = phi::float16;
|
|
using bfloat16 = phi::bfloat16;
|
|
using complex64 = phi::complex64;
|
|
using complex128 = phi::complex128;
|
|
|
|
PD_REGISTER_KERNEL(scale,
|
|
GPU,
|
|
STRIDED,
|
|
phi::ScaleStrideKernel,
|
|
bool,
|
|
float,
|
|
double,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
phi::float8_e4m3fn,
|
|
phi::float8_e5m2,
|
|
uint8_t,
|
|
int8_t,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
phi::complex64,
|
|
phi::complex128) {}
|
|
|
|
PD_REGISTER_KERNEL(full,
|
|
GPU,
|
|
STRIDED,
|
|
phi::FullStrideKernel,
|
|
float,
|
|
double,
|
|
int8_t,
|
|
uint8_t,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
bool,
|
|
phi::float8_e4m3fn,
|
|
phi::float8_e5m2,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
phi::complex64,
|
|
phi::complex128) {}
|
|
|
|
PD_REGISTER_KERNEL(full_like,
|
|
GPU,
|
|
STRIDED,
|
|
phi::FullLikeStrideKernel,
|
|
bool,
|
|
float,
|
|
double,
|
|
int,
|
|
int8_t,
|
|
int64_t,
|
|
int16_t,
|
|
uint8_t,
|
|
phi::float8_e4m3fn,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
phi::complex64,
|
|
phi::complex128) {
|
|
kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(add,
|
|
GPU,
|
|
STRIDED,
|
|
phi::AddStrideKernel,
|
|
float,
|
|
double,
|
|
int16_t,
|
|
int,
|
|
bool,
|
|
uint8_t,
|
|
int8_t,
|
|
int64_t,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
complex64,
|
|
complex128) {}
|
|
|
|
PD_REGISTER_KERNEL(subtract,
|
|
GPU,
|
|
STRIDED,
|
|
phi::SubtractStrideKernel,
|
|
float,
|
|
double,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
float16,
|
|
bfloat16,
|
|
complex64,
|
|
complex128) {}
|
|
|
|
PD_REGISTER_KERNEL(multiply,
|
|
GPU,
|
|
STRIDED,
|
|
phi::MultiplyStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t,
|
|
bool,
|
|
float16,
|
|
complex64,
|
|
complex128,
|
|
bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(divide,
|
|
GPU,
|
|
STRIDED,
|
|
phi::DivideStrideKernel,
|
|
float,
|
|
double,
|
|
int8_t,
|
|
uint8_t,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
bool,
|
|
float16,
|
|
bfloat16,
|
|
complex64,
|
|
complex128) {}
|
|
|
|
PD_REGISTER_KERNEL(copysign,
|
|
GPU,
|
|
STRIDED,
|
|
phi::CopySignStrideKernel,
|
|
bool,
|
|
uint8_t,
|
|
int8_t,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
float,
|
|
double,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(remainder,
|
|
GPU,
|
|
STRIDED,
|
|
phi::RemainderStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t,
|
|
phi::float16,
|
|
phi::complex64,
|
|
phi::complex128,
|
|
phi::bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(maximum,
|
|
GPU,
|
|
STRIDED,
|
|
phi::MaximumStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(minimum,
|
|
GPU,
|
|
STRIDED,
|
|
phi::MinimumStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(floor_divide,
|
|
GPU,
|
|
STRIDED,
|
|
phi::FloorDivideStrideKernel,
|
|
uint8_t,
|
|
int8_t,
|
|
int16_t,
|
|
int,
|
|
int64_t,
|
|
float,
|
|
double,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
|
|
PD_REGISTER_KERNEL(heaviside,
|
|
GPU,
|
|
STRIDED,
|
|
phi::HeavisideStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
float16,
|
|
bfloat16,
|
|
int64_t) {}
|
|
|
|
PD_REGISTER_KERNEL(fmax,
|
|
GPU,
|
|
STRIDED,
|
|
phi::FMaxStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
float16,
|
|
bfloat16,
|
|
int64_t) {}
|
|
|
|
PD_REGISTER_KERNEL(fmin,
|
|
GPU,
|
|
STRIDED,
|
|
phi::FMinStrideKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
float16,
|
|
bfloat16,
|
|
int64_t) {}
|
|
|
|
#endif
|