wehub-resource-sync
472 lines
17 KiB
C++
472 lines
17 KiB
C++
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
|
|
#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/kernels/contiguous_kernel.h"
|
|
#include "paddle/phi/kernels/elementwise_add_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"
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
|
|
#include "paddle/phi/kernels/funcs/dims_simplifier.h"
|
|
|
|
#endif
|
|
|
|
namespace phi {
|
|
|
|
// Not Support Vectorized Kernel For Now
|
|
#define STRIDE_VEC_SIZE 1
|
|
|
|
template <typename Functor,
|
|
typename OutT,
|
|
typename OffsetT,
|
|
int Arity,
|
|
int NumOuts,
|
|
int VecSize,
|
|
int vt>
|
|
__global__ void BinaryElementwiseKernel(
|
|
Array<const _ptr_ char *__restrict__, Arity> ins,
|
|
Array<_ptr_ OutT *, NumOuts> outs,
|
|
int64_t numel,
|
|
int read_lens,
|
|
Functor func,
|
|
funcs::OffsetCalculator<Arity + NumOuts, OffsetT, false> offset_calc) {
|
|
int64_t tid = THREAD_ID_X;
|
|
int64_t nv = BLOCK_NUM_X * vt;
|
|
int64_t idx = nv * BLOCK_ID_X + tid;
|
|
#pragma unroll
|
|
for (int i = 0; i < vt; i++) {
|
|
if (idx < numel) {
|
|
auto offsets = offset_calc.get(idx);
|
|
using Traits = funcs::FunctionTraits<Functor>;
|
|
using ArgsT = typename Traits::ArgsTuple;
|
|
__simd__ ArgsT args[VecSize];
|
|
__simd__ ConditionalT<OutT, NumOuts> result[VecSize];
|
|
std::get<0>(args[0]) =
|
|
*(reinterpret_cast<const _ptr_ std::tuple_element_t<0, ArgsT> *>(
|
|
reinterpret_cast<const _ptr_ char *>(ins[0]) + offsets[1]));
|
|
std::get<1>(args[0]) =
|
|
*(reinterpret_cast<const _ptr_ std::tuple_element_t<1, ArgsT> *>(
|
|
reinterpret_cast<const _ptr_ char *>(ins[1]) + offsets[2]));
|
|
funcs::SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
|
|
VecSize,
|
|
Functor,
|
|
ArgsT,
|
|
Arity>()(
|
|
func, args, result, read_lens);
|
|
char *out_ptr = reinterpret_cast<char *>(outs[0]) + offsets[0];
|
|
*reinterpret_cast<OutT *>(out_ptr) =
|
|
*reinterpret_cast<const OutT *>(&(result[0]));
|
|
idx += BLOCK_NUM_X;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename Functor,
|
|
typename OutT,
|
|
typename OffsetT,
|
|
int Arity,
|
|
int NumOuts,
|
|
int VecSize,
|
|
int vt>
|
|
__global__ void UnaryElementwiseKernel(
|
|
Array<const _ptr_ char *__restrict__, Arity> ins,
|
|
Array<_ptr_ OutT *, NumOuts> outs,
|
|
int64_t numel,
|
|
int read_lens,
|
|
Functor func,
|
|
funcs::OffsetCalculator<Arity + NumOuts, OffsetT, false> offset_calc) {
|
|
int64_t tid = THREAD_ID_X;
|
|
int64_t nv = BLOCK_NUM_X * vt;
|
|
int64_t idx = nv * BLOCK_ID_X + tid;
|
|
#pragma unroll
|
|
for (int i = 0; i < vt; i++) {
|
|
if (idx < numel) {
|
|
auto offsets = offset_calc.get(idx);
|
|
using Traits = funcs::FunctionTraits<Functor>;
|
|
using ArgsT = typename Traits::ArgsTuple;
|
|
__simd__ ArgsT args[VecSize];
|
|
__simd__ ConditionalT<OutT, NumOuts> result[VecSize];
|
|
std::get<0>(args[0]) =
|
|
*(reinterpret_cast<const _ptr_ std::tuple_element_t<0, ArgsT> *>(
|
|
reinterpret_cast<const _ptr_ char *>(ins[0]) + offsets[1]));
|
|
funcs::SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
|
|
VecSize,
|
|
Functor,
|
|
ArgsT,
|
|
Arity>()(
|
|
func, args, result, read_lens);
|
|
char *out_ptr = reinterpret_cast<char *>(outs[0]) + offsets[0];
|
|
*reinterpret_cast<OutT *>(out_ptr) =
|
|
*reinterpret_cast<const OutT *>(&(result[0]));
|
|
idx += BLOCK_NUM_X;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename OutT, typename Context, typename Functor, int NumOuts = 1>
|
|
void BinaryStrideBroadcastKernel(const Context &dev_ctx,
|
|
const std::vector<const DenseTensor *> &ins,
|
|
std::vector<DenseTensor *> *outs,
|
|
Functor func,
|
|
int axis = -1) {
|
|
using Traits = funcs::FunctionTraits<Functor>;
|
|
const int Arity = Traits::arity;
|
|
for (auto i = 0; i < outs->size(); ++i) {
|
|
if (i > 0) {
|
|
PADDLE_ENFORCE_EQ(
|
|
(*outs)[i]->dims(),
|
|
(*outs)[0]->dims(),
|
|
common::errors::InvalidArgument(
|
|
"The shape of each output tensor shall be identical yet, but "
|
|
"%d-th output tensor`s shape is not.",
|
|
i));
|
|
}
|
|
dev_ctx.template Alloc<OutT>((*outs)[i]);
|
|
}
|
|
if ((*outs)[0]->numel() == 0) {
|
|
return;
|
|
}
|
|
int max_rank = 0;
|
|
int min_rank = phi::DDim::kMaxRank;
|
|
for (auto *in : ins) {
|
|
max_rank = std::max(max_rank, in->dims().size());
|
|
min_rank = std::min(min_rank, in->dims().size());
|
|
}
|
|
if (ins.size() == 1) {
|
|
max_rank = std::max(max_rank, (*outs)[0]->dims().size());
|
|
}
|
|
axis = axis == -1 ? max_rank - min_rank : axis;
|
|
auto classifier =
|
|
funcs::BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(
|
|
ins, outs, axis);
|
|
DenseTensorIteratorConfig config;
|
|
config.add_output(*((*outs)[0]));
|
|
config.add_const_input(*(ins[0]));
|
|
config.add_const_input(*(ins[1]));
|
|
DenseTensorIterator iter = config.build();
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
const int64_t &numel = iter.numel();
|
|
|
|
constexpr int unroll_factor = sizeof(OutT) >= 4 ? 2 : 4;
|
|
auto stream = dev_ctx.stream();
|
|
auto threads = 128;
|
|
auto blocks = (numel + 128 * unroll_factor - 1) / (128 * unroll_factor);
|
|
int vec_size = STRIDE_VEC_SIZE;
|
|
|
|
bool is_big_tensor = false;
|
|
int64_t max_stride = 0;
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < iter.ndim(); j++) {
|
|
max_stride += iter.operands_[i].stride_bytes.data()[j] * iter.shape()[j];
|
|
}
|
|
}
|
|
|
|
if (!funcs::IsInUint32Range(max_stride * sizeof(OutT))) {
|
|
is_big_tensor = true;
|
|
}
|
|
|
|
if (is_big_tensor) {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<3, false, uint64_t>(iter);
|
|
BinaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint64_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
} else {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<3, false, uint32_t>(iter);
|
|
BinaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint32_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
}
|
|
}
|
|
|
|
template <typename OutT, typename Context, typename Functor, int NumOuts = 1>
|
|
void BinaryStrideElementwiseKernel(const Context &dev_ctx,
|
|
const std::vector<const DenseTensor *> &ins,
|
|
std::vector<DenseTensor *> *outs,
|
|
Functor func) {
|
|
using Traits = funcs::FunctionTraits<Functor>;
|
|
const int Arity = Traits::arity;
|
|
bool have_0_size = false;
|
|
for (int i = 0; i < outs->size(); ++i) {
|
|
if (outs->at(i)->numel() == 0) {
|
|
have_0_size = true;
|
|
}
|
|
if (i > 0) {
|
|
PADDLE_ENFORCE_EQ(
|
|
(*outs)[i]->dims(),
|
|
(*outs)[0]->dims(),
|
|
common::errors::InvalidArgument(
|
|
"The shape of each output tensor shall be identical yet, "
|
|
"but %dth output tensor`s shape is not.",
|
|
i));
|
|
}
|
|
dev_ctx.template Alloc<OutT>((*outs)[i]);
|
|
}
|
|
if (have_0_size) {
|
|
return;
|
|
}
|
|
int max_rank = 0;
|
|
int min_rank = phi::DDim::kMaxRank;
|
|
for (auto *in : ins) {
|
|
max_rank = std::max(max_rank, in->dims().size());
|
|
min_rank = std::min(min_rank, in->dims().size());
|
|
}
|
|
if (ins.size() == 1) {
|
|
max_rank = std::max(max_rank, (*outs)[0]->dims().size());
|
|
}
|
|
int axis = max_rank - min_rank;
|
|
auto classifier =
|
|
funcs::BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(
|
|
ins, outs, axis);
|
|
DenseTensorIteratorConfig config;
|
|
config.add_output(*((*outs)[0]));
|
|
config.add_const_input(*(ins[0]));
|
|
config.add_const_input(*(ins[1]));
|
|
DenseTensorIterator iter = config.build();
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
const int64_t &numel = iter.numel();
|
|
|
|
constexpr int unroll_factor = sizeof(OutT) >= 4 ? 2 : 4;
|
|
auto stream = dev_ctx.stream();
|
|
auto threads = 128;
|
|
auto blocks = (numel + 128 * unroll_factor - 1) / (128 * unroll_factor);
|
|
int vec_size = STRIDE_VEC_SIZE;
|
|
|
|
bool is_big_tensor = false;
|
|
int64_t max_stride = 0;
|
|
for (int i = 0; i < 3; i++) {
|
|
for (int j = 0; j < iter.ndim(); j++) {
|
|
max_stride += iter.operands_[i].stride_bytes.data()[j] * iter.shape()[j];
|
|
}
|
|
}
|
|
|
|
if (!funcs::IsInUint32Range(max_stride) * sizeof(OutT)) {
|
|
is_big_tensor = true;
|
|
}
|
|
|
|
if (is_big_tensor) {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<3, false, uint64_t>(iter);
|
|
BinaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint64_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
} else {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<3, false, uint32_t>(iter);
|
|
BinaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint32_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
}
|
|
}
|
|
|
|
template <typename OutT, typename Context, typename Functor, int NumOuts = 1>
|
|
void UnaryStrideElementwiseKernel(const Context &dev_ctx,
|
|
const std::vector<const DenseTensor *> &ins,
|
|
std::vector<DenseTensor *> *outs,
|
|
Functor func) {
|
|
using Traits = funcs::FunctionTraits<Functor>;
|
|
const int Arity = Traits::arity;
|
|
bool have_0_size = false;
|
|
for (int i = 0; i < outs->size(); ++i) {
|
|
if (outs->at(i)->numel() == 0) {
|
|
have_0_size = true;
|
|
}
|
|
if (i > 0) {
|
|
PADDLE_ENFORCE_EQ(
|
|
(*outs)[i]->dims(),
|
|
(*outs)[0]->dims(),
|
|
common::errors::InvalidArgument(
|
|
"The shape of each output tensor shall be identical yet, "
|
|
"but %dth output tensor`s shape is not.",
|
|
i));
|
|
}
|
|
dev_ctx.template Alloc<OutT>((*outs)[i]);
|
|
}
|
|
if (have_0_size) {
|
|
return;
|
|
}
|
|
int max_rank = 0;
|
|
int min_rank = phi::DDim::kMaxRank;
|
|
for (auto *in : ins) {
|
|
max_rank = std::max(max_rank, in->dims().size());
|
|
min_rank = std::min(min_rank, in->dims().size());
|
|
}
|
|
if (ins.size() == 1) {
|
|
max_rank = std::max(max_rank, (*outs)[0]->dims().size());
|
|
}
|
|
int axis = max_rank - min_rank;
|
|
auto classifier =
|
|
funcs::BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(
|
|
ins, outs, axis);
|
|
DenseTensorIteratorConfig config;
|
|
config.add_output(*((*outs)[0]));
|
|
config.add_const_input(*(ins[0]));
|
|
DenseTensorIterator iter = config.build();
|
|
// TODO(large-tensor): downstream functors may still use int; guard until
|
|
// upgraded.
|
|
const int64_t &numel = iter.numel();
|
|
|
|
funcs::OffsetCalculator offset_calc = funcs::make_offset_calculator<2>(iter);
|
|
constexpr int unroll_factor = sizeof(OutT) >= 4 ? 2 : 4;
|
|
auto stream = dev_ctx.stream();
|
|
auto threads = 128;
|
|
auto blocks = (numel + 128 * unroll_factor - 1) / (128 * unroll_factor);
|
|
int vec_size = STRIDE_VEC_SIZE;
|
|
|
|
bool is_big_tensor = false;
|
|
int64_t max_stride = 0;
|
|
for (int i = 0; i < 2; i++) {
|
|
for (int j = 0; j < iter.ndim(); j++) {
|
|
max_stride += iter.operands_[i].stride_bytes.data()[j] * iter.shape()[j];
|
|
}
|
|
}
|
|
|
|
if (!funcs::IsInUint32Range(max_stride * sizeof(OutT))) {
|
|
is_big_tensor = true;
|
|
}
|
|
|
|
if (is_big_tensor) {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<2, false, uint64_t>(iter);
|
|
UnaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint64_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
} else {
|
|
funcs::OffsetCalculator offset_calc =
|
|
funcs::make_offset_calculator<2, false, uint32_t>(iter);
|
|
UnaryElementwiseKernel<Functor,
|
|
OutT,
|
|
uint32_t,
|
|
Arity,
|
|
NumOuts,
|
|
STRIDE_VEC_SIZE,
|
|
unroll_factor>
|
|
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
|
classifier.outs_data,
|
|
numel,
|
|
vec_size,
|
|
func,
|
|
offset_calc);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context, typename Functor>
|
|
void LaunchUnaryElementwiseStrideKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
Functor func,
|
|
DenseTensor *out) {
|
|
std::vector<const DenseTensor *> inputs = {&x};
|
|
std::vector<DenseTensor *> outputs = {out};
|
|
dev_ctx.template Alloc<T>(out);
|
|
UnaryStrideElementwiseKernel<T, Context>(dev_ctx, inputs, &outputs, func);
|
|
}
|
|
|
|
template <typename T, typename Context, typename Functor>
|
|
void LaunchBinaryElementwiseStrideKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &y,
|
|
Functor func,
|
|
int axis,
|
|
DenseTensor *out) {
|
|
std::vector<const DenseTensor *> inputs = {&x, &y};
|
|
std::vector<DenseTensor *> outputs = {out};
|
|
dev_ctx.template Alloc<T>(out);
|
|
BinaryStrideBroadcastKernel<T, Context>(
|
|
dev_ctx, inputs, &outputs, func, axis);
|
|
}
|
|
|
|
template <typename Context>
|
|
DenseTensor Tensor2Contiguous(const Context &dev_ctx,
|
|
const DenseTensor &tensor) {
|
|
DenseTensor dense_out;
|
|
MetaTensor meta_input(tensor);
|
|
MetaTensor meta_out(&dense_out);
|
|
UnchangedInferMeta(meta_input, &meta_out);
|
|
PD_VISIT_ALL_TYPES(tensor.dtype(), "Tensor2Contiguous", ([&] {
|
|
ContiguousKernel<data_t, Context>(
|
|
dev_ctx, tensor, &dense_out);
|
|
}));
|
|
return dense_out;
|
|
}
|
|
|
|
#undef STRIDE_VEC_SIZE
|
|
|
|
} // namespace phi
|
|
|
|
#endif
|