chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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. */
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sstream>
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/kernels/funcs/elementwise_base.h"
|
||||
|
||||
#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
|
||||
#include "paddle/phi/kernels/funcs/dims_simplifier.h"
|
||||
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
namespace funcs {
|
||||
|
||||
#if defined(__NVCC__) || defined(__HIPCC__) || defined(__xpu__)
|
||||
|
||||
enum BroadcastType { kMixed = 1, kBroadcast = 2, kElementwise = 3 };
|
||||
|
||||
template <typename OutT, typename Functor, int Arity, int NumOuts>
|
||||
struct BroadcastTypeClassifier {
|
||||
int64_t numel{0};
|
||||
int broadcast_num{0}; // Not used for XPU
|
||||
bool all_elementwise{true}; // Not used for XPU
|
||||
Array<bool, Arity> use_broadcast; // Not used for XPU
|
||||
Array<kps::details::BroadcastConfig, Arity> configs;
|
||||
Array<const _ptr_ char *__restrict__, Arity> ins_data;
|
||||
Array<_ptr_ OutT *, NumOuts> outs_data;
|
||||
|
||||
BroadcastTypeClassifier() {}
|
||||
BroadcastTypeClassifier(const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis) {
|
||||
numel = (*outs)[0]->numel();
|
||||
|
||||
#ifndef PADDLE_WITH_XPU_KP
|
||||
for (size_t i = 0; i < ins.size(); ++i) {
|
||||
bool is_same_dim = ins[i]->numel() == numel;
|
||||
if (is_same_dim) {
|
||||
use_broadcast[i] = false;
|
||||
} else {
|
||||
use_broadcast[i] = true;
|
||||
broadcast_num++;
|
||||
}
|
||||
all_elementwise &= is_same_dim;
|
||||
}
|
||||
#endif
|
||||
|
||||
InitBroadcastConfigs(ins, outs, axis);
|
||||
|
||||
using Traits = funcs::FunctionTraits<Functor>;
|
||||
using ArgsT = typename Traits::ArgsTuple;
|
||||
ArgsT arg;
|
||||
UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, arg, &ins_data);
|
||||
for (int i = 0; i < NumOuts; ++i) {
|
||||
outs_data[i] = (*outs)[i]->data<OutT>();
|
||||
}
|
||||
}
|
||||
|
||||
void InitBroadcastConfigs(const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis) {
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
const auto dims_simplifier =
|
||||
BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
|
||||
if (VLOG_IS_ON(6)) {
|
||||
DimsSimplifiedLogger<int64_t>::Log(
|
||||
ins, outs, dims_simplifier, "BroadcastKernel");
|
||||
}
|
||||
configs[0] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
|
||||
dims_simplifier.in_dims[0],
|
||||
dims_simplifier.in_dims[1],
|
||||
dims_simplifier.rank);
|
||||
configs[1] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
|
||||
dims_simplifier.in_dims[1],
|
||||
dims_simplifier.in_dims[0],
|
||||
dims_simplifier.rank);
|
||||
#else
|
||||
if (!all_elementwise) {
|
||||
const auto dims_simplifier =
|
||||
BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
|
||||
if (VLOG_IS_ON(6)) {
|
||||
DimsSimplifiedLogger<int64_t>::Log(
|
||||
ins, outs, dims_simplifier, "BroadcastKernel");
|
||||
}
|
||||
for (int i = 0; i < Arity; ++i) {
|
||||
// if data shape is[m, n], then you should set data_dim = {n, m}
|
||||
// eg: out's shape [3, 45, 1]. then out_dims = {1, 45, 3}
|
||||
// if (ins[i]->numel() != (*outs)[0]->numel()) {
|
||||
if (ins[i]->numel()) {
|
||||
configs[i] = kps::details::BroadcastConfig(dims_simplifier.out_dims,
|
||||
dims_simplifier.in_dims[i],
|
||||
dims_simplifier.rank);
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
// Common broadcast/elementwise Loader.
|
||||
template <int Index, int VecSize, bool IsBoundary, int LoadType>
|
||||
struct BroadcastDataLoader {
|
||||
template <typename Array1, typename Array2, typename Array3, typename ArgsT>
|
||||
static __device__ __forceinline__ void Apply(const Array1 &ins,
|
||||
ArgsT *args,
|
||||
const Array2 &configs,
|
||||
const Array3 &use_broadcast,
|
||||
const int block_offset,
|
||||
const int num,
|
||||
const uint32_t numel,
|
||||
int read_lens) {
|
||||
using Type = std::tuple_element_t<Index, ArgsT>;
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
kps::Init<Type, ArgsT, Index, VecSize>(
|
||||
args, static_cast<Type>(1.0f), read_lens);
|
||||
if (use_broadcast[Index]) {
|
||||
kps::ReadDataBc<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
|
||||
args,
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index]),
|
||||
block_offset,
|
||||
configs[Index],
|
||||
numel,
|
||||
read_lens);
|
||||
} else {
|
||||
kps::ReadData<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
|
||||
args,
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index]) + block_offset,
|
||||
num,
|
||||
read_lens);
|
||||
}
|
||||
#else
|
||||
kps::Init<Type, ArgsT, Index, VecSize>(args, static_cast<Type>(1.0f));
|
||||
if (use_broadcast[Index]) {
|
||||
kps::ReadDataBc<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
|
||||
args,
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index]),
|
||||
block_offset,
|
||||
configs[Index],
|
||||
numel,
|
||||
VecSize);
|
||||
}
|
||||
// NOTE: If use if...else... with condition `use_broadcast[Index]` here,
|
||||
// there will be some errs with clang12 while compiling in ROCm.
|
||||
// When the compiler is upgraded, if...else... may be used.
|
||||
if (!use_broadcast[Index]) {
|
||||
kps::ReadData<Type, VecSize, 1, ArgsT, Index, IsBoundary>(
|
||||
args,
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index]) + block_offset,
|
||||
num,
|
||||
VecSize);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
/* BroadcastDataLoaders Partial specialization */
|
||||
#ifndef PADDLE_WITH_XPU_KP
|
||||
// Scalar elementwise Loader with consideration of IsBoundary.
|
||||
template <int Index, int VecSize>
|
||||
struct BroadcastDataLoader<Index, VecSize, true, kElementwise> {
|
||||
template <typename Array1, typename Array2, typename Array3, typename ArgsT>
|
||||
static __device__ __forceinline__ void Apply(const Array1 &ins,
|
||||
ArgsT *args,
|
||||
const Array2 &configs,
|
||||
const Array3 &use_broadcast,
|
||||
const int block_offset,
|
||||
const int num,
|
||||
const uint32_t numel,
|
||||
int read_lens) {
|
||||
using Type = std::tuple_element_t<Index, ArgsT>;
|
||||
int thread_offset = threadIdx.x * VecSize + block_offset;
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < VecSize; ++idx) {
|
||||
std::get<Index>(args[idx]) = static_cast<Type>(1);
|
||||
int index = thread_offset + idx;
|
||||
if (index < numel) {
|
||||
std::get<Index>(args[idx]) =
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index])[index];
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Vectorized elementwise Loader without consideration of IsBoundary.
|
||||
template <int Index, int VecSize>
|
||||
struct BroadcastDataLoader<Index, VecSize, false, kElementwise> {
|
||||
template <typename Array1, typename Array2, typename Array3, typename ArgsT>
|
||||
static __device__ __forceinline__ void Apply(const Array1 &ins,
|
||||
ArgsT *args,
|
||||
const Array2 &configs,
|
||||
const Array3 &use_broadcast,
|
||||
const int block_offset,
|
||||
const int num,
|
||||
const uint32_t numel,
|
||||
int read_lens) {
|
||||
using Type = std::tuple_element_t<Index, ArgsT>;
|
||||
using VecType = phi::kps::details::VectorType<Type, VecSize>;
|
||||
VecType vec_temp;
|
||||
|
||||
int64_t thread_offset =
|
||||
static_cast<int64_t>(threadIdx.x) +
|
||||
static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x);
|
||||
const VecType *__restrict__ vec_input =
|
||||
reinterpret_cast<const VecType *__restrict__>(ins[Index]);
|
||||
vec_temp = vec_input[thread_offset];
|
||||
#pragma unroll
|
||||
for (int idx = 0; idx < VecSize; ++idx) {
|
||||
std::get<Index>(args[idx]) = vec_temp.val[idx];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <int Index, int VecSize>
|
||||
struct BroadcastDataInit {
|
||||
template <typename ArgsT>
|
||||
static __device__ __forceinline__ void Apply(ArgsT *args) {
|
||||
using Type = std::tuple_element_t<Index, ArgsT>;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < VecSize; ++k) {
|
||||
std::get<Index>(args[k]) = static_cast<Type>(1);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <int Index, int VecSize>
|
||||
struct BroadcastDataSetter {
|
||||
template <typename Array, typename ArgsT>
|
||||
static __device__ __forceinline__ void Apply(const Array &ins,
|
||||
ArgsT *args,
|
||||
uint32_t index_bc[][VecSize]) {
|
||||
using Type = std::tuple_element_t<Index, ArgsT>;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < VecSize; ++k) {
|
||||
std::get<Index>(args[k]) =
|
||||
reinterpret_cast<const _ptr_ Type *>(ins[Index])[index_bc[Index][k]];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
|
||||
// static broadcast unroller
|
||||
template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
|
||||
typename Func,
|
||||
bool IsBoundary,
|
||||
int LoadType,
|
||||
int VecSize,
|
||||
int End,
|
||||
int Begin = 0>
|
||||
struct BcUnroller {
|
||||
template <typename... Args>
|
||||
static HOSTDEVICE inline void step(Args &&...args) {
|
||||
Func<Begin, VecSize, IsBoundary, LoadType>::Apply(
|
||||
std::forward<Args>(args)...);
|
||||
BcUnroller<Func, IsBoundary, LoadType, VecSize, End, Begin + 1>::step(
|
||||
args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <template <int Index, int VecSize, bool IsBoundary, int LoadType>
|
||||
typename Func,
|
||||
bool IsBoundary,
|
||||
int LoadType,
|
||||
int VecSize,
|
||||
int End>
|
||||
struct BcUnroller<Func, IsBoundary, LoadType, VecSize, End, End> {
|
||||
template <typename... Args>
|
||||
static HOSTDEVICE inline void step(Args &&...args) {}
|
||||
};
|
||||
|
||||
template <typename OutT,
|
||||
typename Functor,
|
||||
int Arity,
|
||||
int NumOuts,
|
||||
int VecSize,
|
||||
bool IsBoundary,
|
||||
int LoadType>
|
||||
__device__ void VectorizedBroadcastKernelImpl(
|
||||
const Array<const _ptr_ char *__restrict__, Arity> &ins,
|
||||
Array<_ptr_ OutT *, NumOuts> outs,
|
||||
const Array<bool, Arity> &use_broadcast,
|
||||
const uint32_t numel,
|
||||
const Array<kps::details::BroadcastConfig, Arity> &configs,
|
||||
uint32_t num,
|
||||
uint32_t block_offset,
|
||||
int read_lens,
|
||||
Functor func) {
|
||||
using Traits = funcs::FunctionTraits<Functor>;
|
||||
using ArgsT = typename Traits::ArgsTuple;
|
||||
__simd__ ArgsT args[VecSize];
|
||||
__simd__ ConditionalT<OutT, NumOuts> result[VecSize];
|
||||
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
|
||||
ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
|
||||
#else
|
||||
if (LoadType == kBroadcast) {
|
||||
uint32_t index_bc[Arity][VecSize] = {0};
|
||||
Unroller<BroadcastDataInit, VecSize, Arity>::step(args);
|
||||
uint32_t thread_offset = block_offset + threadIdx.x * VecSize;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < VecSize; ++k) {
|
||||
uint32_t idx = thread_offset + k;
|
||||
if (IsBoundary && idx == numel) break;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < DDim::kMaxRank; ++i) {
|
||||
if (i == configs[0].rank) break;
|
||||
auto fast_divmoder = configs[0].divmoders[i].Divmod(idx);
|
||||
idx = fast_divmoder.val[0];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < Arity; ++j) {
|
||||
index_bc[j][k] += fast_divmoder.val[1] * configs[j].strides[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
Unroller<BroadcastDataSetter, VecSize, Arity>::step(ins, args, index_bc);
|
||||
} else {
|
||||
BcUnroller<BroadcastDataLoader, IsBoundary, LoadType, VecSize, Arity>::step(
|
||||
ins, args, configs, use_broadcast, block_offset, num, numel, read_lens);
|
||||
}
|
||||
#endif
|
||||
SameDimsElementwisePrimitiveCaller<ConditionalT<OutT, NumOuts>,
|
||||
VecSize,
|
||||
Functor,
|
||||
ArgsT,
|
||||
Arity>()(func, args, result, read_lens);
|
||||
funcs::ElementwiseWriteDataCallerBc<OutT, VecSize, IsBoundary, NumOuts>()(
|
||||
outs, result, block_offset, num, read_lens);
|
||||
}
|
||||
|
||||
template <typename Functor,
|
||||
typename OutT,
|
||||
int Arity,
|
||||
int NumOuts,
|
||||
int VecSize,
|
||||
int LoadType>
|
||||
__global__ void VectorizedBroadcastKernel(
|
||||
Array<const _ptr_ char *__restrict__, Arity> ins,
|
||||
Array<_ptr_ OutT *, NumOuts> outs,
|
||||
Array<bool, Arity> use_broadcast,
|
||||
uint32_t numel,
|
||||
Array<kps::details::BroadcastConfig, Arity> configs,
|
||||
uint32_t main_offset,
|
||||
uint32_t tail_tid,
|
||||
int read_lens,
|
||||
Functor func) {
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
int64_t block_offset =
|
||||
static_cast<int64_t>(BLOCK_ID_X) * BLOCK_NUM_X * read_lens;
|
||||
int64_t stride = static_cast<int64_t>(BLOCK_NUM_X) * GRID_NUM_X * read_lens;
|
||||
for (; block_offset < main_offset; block_offset += stride) {
|
||||
VectorizedBroadcastKernelImpl<OutT,
|
||||
Functor,
|
||||
Arity,
|
||||
NumOuts,
|
||||
VecSize,
|
||||
false,
|
||||
LoadType>(ins,
|
||||
outs,
|
||||
use_broadcast,
|
||||
numel,
|
||||
configs,
|
||||
BLOCK_NUM_X * read_lens,
|
||||
block_offset,
|
||||
read_lens,
|
||||
func);
|
||||
}
|
||||
int64_t num = numel - block_offset;
|
||||
if (num > 0) {
|
||||
VectorizedBroadcastKernelImpl<OutT,
|
||||
Functor,
|
||||
Arity,
|
||||
NumOuts,
|
||||
VecSize,
|
||||
true,
|
||||
LoadType>(ins,
|
||||
outs,
|
||||
use_broadcast,
|
||||
numel,
|
||||
configs,
|
||||
num,
|
||||
block_offset,
|
||||
read_lens,
|
||||
func);
|
||||
}
|
||||
#else
|
||||
int64_t block_offset =
|
||||
static_cast<int64_t>(BLOCK_ID_X) * BLOCK_NUM_X * VecSize;
|
||||
if (block_offset < main_offset) {
|
||||
VectorizedBroadcastKernelImpl<OutT,
|
||||
Functor,
|
||||
Arity,
|
||||
NumOuts,
|
||||
VecSize,
|
||||
false,
|
||||
LoadType>(ins,
|
||||
outs,
|
||||
use_broadcast,
|
||||
numel,
|
||||
configs,
|
||||
BLOCK_NUM_X * VecSize,
|
||||
block_offset,
|
||||
read_lens,
|
||||
func);
|
||||
} else {
|
||||
VectorizedBroadcastKernelImpl<OutT,
|
||||
Functor,
|
||||
Arity,
|
||||
NumOuts,
|
||||
VecSize,
|
||||
true,
|
||||
LoadType>(ins,
|
||||
outs,
|
||||
use_broadcast,
|
||||
numel,
|
||||
configs,
|
||||
tail_tid,
|
||||
block_offset,
|
||||
read_lens,
|
||||
func);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
|
||||
void LaunchBroadcastKernel(
|
||||
const KPDevice &dev_ctx,
|
||||
const BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts> &classifier,
|
||||
Functor func) {
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
const int64_t numel_64 = classifier.numel;
|
||||
PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "XPU broadcast kernel numel");
|
||||
const uint32_t numel = static_cast<uint32_t>(numel_64);
|
||||
const int threads = 64;
|
||||
const int blocks = 8;
|
||||
int read_lens = configs[0].buf_len;
|
||||
auto stream = dev_ctx.x_context()->xpu_stream;
|
||||
const int64_t block_len = static_cast<int64_t>(read_lens) * threads;
|
||||
const int64_t main_offset_64 = (numel_64 / block_len) * block_len;
|
||||
const int64_t tail_tid_64 = numel_64 % block_len;
|
||||
const uint32_t main_offset = static_cast<uint32_t>(main_offset_64);
|
||||
const uint32_t tail_tid = static_cast<uint32_t>(tail_tid_64);
|
||||
|
||||
VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, false>
|
||||
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
||||
classifier.outs_data,
|
||||
classifier.use_broadcast,
|
||||
numel,
|
||||
classifier.configs,
|
||||
main_offset,
|
||||
tail_tid,
|
||||
read_lens,
|
||||
func);
|
||||
#else
|
||||
const int64_t numel_64 = classifier.numel;
|
||||
auto gpu_config =
|
||||
phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel_64, VecSize);
|
||||
auto stream = dev_ctx.stream();
|
||||
uint32_t threads = static_cast<uint32_t>(gpu_config.GetBlockSize());
|
||||
auto blocks = gpu_config.block_per_grid;
|
||||
PADDLE_ENFORCE_LE_UINT32_MAX(numel_64, "numel");
|
||||
const uint32_t numel = static_cast<uint32_t>(numel_64);
|
||||
const int64_t block_len = static_cast<int64_t>(VecSize) * threads;
|
||||
const int64_t main_offset_64 = (numel_64 / block_len) * block_len;
|
||||
const int64_t tail_tid_64 = numel_64 % block_len;
|
||||
uint32_t main_offset = static_cast<uint32_t>(main_offset_64);
|
||||
uint32_t tail_tid = static_cast<uint32_t>(tail_tid_64);
|
||||
|
||||
if (classifier.all_elementwise) {
|
||||
VectorizedBroadcastKernel<Functor,
|
||||
OutT,
|
||||
Arity,
|
||||
NumOuts,
|
||||
VecSize,
|
||||
kElementwise>
|
||||
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
||||
classifier.outs_data,
|
||||
classifier.use_broadcast,
|
||||
numel,
|
||||
classifier.configs,
|
||||
main_offset,
|
||||
tail_tid,
|
||||
VecSize,
|
||||
func);
|
||||
} else if (classifier.broadcast_num > (Arity >> 1)) {
|
||||
constexpr BroadcastType type_ = (Arity > 1) ? kBroadcast : kMixed;
|
||||
VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, type_>
|
||||
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
||||
classifier.outs_data,
|
||||
classifier.use_broadcast,
|
||||
numel,
|
||||
classifier.configs,
|
||||
main_offset,
|
||||
tail_tid,
|
||||
VecSize,
|
||||
func);
|
||||
} else {
|
||||
VectorizedBroadcastKernel<Functor, OutT, Arity, NumOuts, VecSize, kMixed>
|
||||
<<<blocks, threads, 0, stream>>>(classifier.ins_data,
|
||||
classifier.outs_data,
|
||||
classifier.use_broadcast,
|
||||
numel,
|
||||
classifier.configs,
|
||||
main_offset,
|
||||
tail_tid,
|
||||
VecSize,
|
||||
func);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
|
||||
typename std::enable_if<!NeedVectorized<OutT>::value, void>::type
|
||||
BroadcastKernelForDifferentVecSize(const KPDevice &dev_ctx,
|
||||
const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis,
|
||||
Functor func) {
|
||||
auto classifier =
|
||||
BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(ins, outs, axis);
|
||||
LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeS>(
|
||||
dev_ctx, classifier, func);
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
|
||||
typename std::enable_if<NeedVectorized<OutT>::value, void>::type
|
||||
BroadcastKernelForDifferentVecSize(const KPDevice &dev_ctx,
|
||||
const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis,
|
||||
Functor func) {
|
||||
auto classifier =
|
||||
BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(ins, outs, axis);
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
auto type = kps::details::OptType::CanNotOptimize;
|
||||
bool is_optimize = classifier.configs[0].cmp_type != type;
|
||||
int vec_size = is_optimize ? VecSizeL : VecSizeM;
|
||||
#else
|
||||
static int capability = dev_ctx.GetComputeCapability();
|
||||
// For Hopper and Blackwell, max vectorized size is VecSizeL(8).
|
||||
static int max_vec_size = capability >= 90 ? VecSizeVL : VecSizeL;
|
||||
// calculate the max vec_size for all ins and outs
|
||||
int vec_size = GetVectorizedSizeForTensors(ins, *outs, true);
|
||||
vec_size = std::min(vec_size, max_vec_size);
|
||||
int64_t numel = classifier.numel;
|
||||
// For small tensor, using VecSizeL can improve performance more than
|
||||
// VecSizeVL
|
||||
constexpr int64_t large_vect_threshold = 1024 * 1024 * 4;
|
||||
if (numel < large_vect_threshold) {
|
||||
vec_size = std::min(vec_size, VecSizeL);
|
||||
}
|
||||
#endif
|
||||
|
||||
switch (vec_size) {
|
||||
case VecSizeVL: {
|
||||
LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeVL>(
|
||||
dev_ctx, classifier, func);
|
||||
break;
|
||||
}
|
||||
case VecSizeL: {
|
||||
LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeL>(
|
||||
dev_ctx, classifier, func);
|
||||
break;
|
||||
}
|
||||
case VecSizeM: {
|
||||
LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeM>(
|
||||
dev_ctx, classifier, func);
|
||||
break;
|
||||
}
|
||||
case VecSizeS: {
|
||||
LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeS>(
|
||||
dev_ctx, classifier, func);
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
PADDLE_THROW(common::errors::Unimplemented(
|
||||
"Unsupported vectorized size: %d!", vec_size));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void updateStridesDims(std::vector<int64_t> *strides,
|
||||
std::vector<int64_t> *dims) {
|
||||
for (int i = 1; i < strides->size(); i++) {
|
||||
(*strides)[i] = (*strides)[i - 1] * (*dims)[i - 1];
|
||||
}
|
||||
// reverse origin_in_dim and origin_in_stride if in's dim_size > 0
|
||||
std::reverse(strides->begin(), strides->end());
|
||||
std::reverse(dims->begin(), dims->end());
|
||||
}
|
||||
|
||||
static void SliceTensor(DenseTensor *x,
|
||||
const DenseTensor *share,
|
||||
const std::vector<int64_t> &out_compute_dims,
|
||||
int64_t offset) {
|
||||
auto new_dim = make_ddim(out_compute_dims);
|
||||
DenseTensorMeta meta(share->dtype(),
|
||||
new_dim,
|
||||
share->layout(),
|
||||
offset * SizeOf(share->dtype()) + share->offset());
|
||||
x->set_meta(meta);
|
||||
x->ShareBufferWith(*(share), true);
|
||||
x->Resize(new_dim);
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int kArity, int NumOuts = 1>
|
||||
void BroadcastKernelSplit(const KPDevice &dev_ctx,
|
||||
const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis,
|
||||
Functor func,
|
||||
const int64_t compute_size) {
|
||||
const auto dims_simplifier =
|
||||
BroadcastDimsSimplifier(ins, (*outs)[0]->dims(), axis);
|
||||
if (VLOG_IS_ON(6)) {
|
||||
DimsSimplifiedLogger<int64_t>::Log(
|
||||
ins, outs, dims_simplifier, "GPU Broadcast");
|
||||
}
|
||||
|
||||
int all_rank = dims_simplifier.rank;
|
||||
std::vector<int64_t> origin_out_strides(all_rank, 1);
|
||||
auto origin_in_dims = dims_simplifier.in_dims;
|
||||
auto origin_out_dims = dims_simplifier.out_dims;
|
||||
auto origin_in_strides = dims_simplifier.in_dims;
|
||||
|
||||
// for split
|
||||
std::vector<int64_t> loop_num_out(all_rank, 1);
|
||||
std::vector<int64_t> loop_num_out_stride(all_rank, 1);
|
||||
|
||||
// for input's offset
|
||||
std::vector<int64_t> ins_offset(kArity, 0);
|
||||
std::vector<int64_t> ins_scale_for_dim(kArity, 0);
|
||||
|
||||
// init offset and check in's dim
|
||||
for (int k = 0; k < kArity; k++) {
|
||||
ins_scale_for_dim[k] = ins[k]->dims().size() == 0 ? 0 : 1;
|
||||
if (ins_scale_for_dim[k]) {
|
||||
origin_in_strides[k][0] = 1;
|
||||
}
|
||||
}
|
||||
|
||||
updateStridesDims(&origin_out_strides, &origin_out_dims);
|
||||
for (int k = 0; k < kArity; k++) {
|
||||
if (ins_scale_for_dim[k]) {
|
||||
updateStridesDims(&origin_in_strides[k], &origin_in_dims[k]);
|
||||
}
|
||||
}
|
||||
|
||||
// init out_split_dim and in_split_dims
|
||||
auto out_split_dim = origin_out_dims;
|
||||
auto in_split_dims = origin_in_dims;
|
||||
|
||||
// init
|
||||
int64_t loop_num = 1;
|
||||
int64_t split_idx = 0;
|
||||
|
||||
for (int r = 0; r < all_rank; r++) {
|
||||
// if the compute_size was too small the split_size must be 0, but the
|
||||
// dim_num must ge 1
|
||||
int64_t split_size = compute_size / origin_out_strides[r];
|
||||
out_split_dim[r] = std::max(split_size, static_cast<int64_t>(1));
|
||||
loop_num_out[r] =
|
||||
(origin_out_dims[r] + out_split_dim[r] - 1) / out_split_dim[r];
|
||||
loop_num *= loop_num_out[r];
|
||||
|
||||
for (int k = 0; k < kArity; k++) {
|
||||
if (ins_scale_for_dim[k]) {
|
||||
in_split_dims[k][r] = std::min(origin_in_dims[k][r], out_split_dim[r]);
|
||||
}
|
||||
}
|
||||
|
||||
// split_idx is the index for lash split dim
|
||||
if (split_size != 0) {
|
||||
split_idx = r;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
loop_num_out_stride[all_rank - 1] = 1;
|
||||
for (int r = all_rank - 2; r >= 0; r--) {
|
||||
loop_num_out_stride[r] = loop_num_out_stride[r + 1] * loop_num_out[r + 1];
|
||||
}
|
||||
|
||||
// compute
|
||||
|
||||
for (int iter = 0; iter < loop_num; iter++) {
|
||||
std::vector<const DenseTensor *> new_ins = {};
|
||||
std::vector<DenseTensor *> new_outs = {};
|
||||
DenseTensor tmp_in[kArity];
|
||||
DenseTensor tmp_out[NumOuts];
|
||||
|
||||
int64_t tmp_size = iter;
|
||||
int64_t out_offset = 0;
|
||||
// compute the offset before last split dim
|
||||
for (int i = 0; i < split_idx; i++) {
|
||||
auto repeat_times = tmp_size / loop_num_out_stride[i];
|
||||
out_offset += repeat_times * origin_out_strides[i];
|
||||
for (int k = 0; k < kArity; k++) {
|
||||
if (ins_scale_for_dim[k]) {
|
||||
ins_offset[k] +=
|
||||
(repeat_times % origin_in_dims[k][i]) * origin_in_strides[k][i];
|
||||
}
|
||||
}
|
||||
tmp_size = tmp_size % loop_num_out_stride[i];
|
||||
}
|
||||
// tmp_size is the last split_dims's repeat idx
|
||||
auto pre_deal_size = tmp_size * out_split_dim[split_idx];
|
||||
out_offset += pre_deal_size * origin_out_strides[split_idx];
|
||||
// compute_size
|
||||
auto remainder_size = origin_out_dims[split_idx] - pre_deal_size;
|
||||
|
||||
// get current compute size
|
||||
auto out_compute_dims = out_split_dim;
|
||||
out_compute_dims[split_idx] =
|
||||
std::min(out_split_dim[split_idx], remainder_size);
|
||||
|
||||
// in + compute_size
|
||||
auto in_compute_dims = in_split_dims;
|
||||
for (int k = 0; k < kArity; k++) {
|
||||
if (ins_scale_for_dim[k]) {
|
||||
auto split_repeat =
|
||||
origin_in_dims[k][split_idx] == origin_out_dims[split_idx]
|
||||
? tmp_size
|
||||
: 0;
|
||||
ins_offset[k] += split_repeat * in_split_dims[k][split_idx] *
|
||||
origin_in_strides[k][split_idx];
|
||||
in_compute_dims[k][split_idx] =
|
||||
std::min(in_split_dims[k][split_idx], out_compute_dims[split_idx]);
|
||||
}
|
||||
SliceTensor(&tmp_in[k],
|
||||
ins[k],
|
||||
in_compute_dims[k],
|
||||
ins_scale_for_dim[k] * ins_offset[k]);
|
||||
new_ins.emplace_back(&tmp_in[k]);
|
||||
ins_offset[k] = 0;
|
||||
}
|
||||
|
||||
for (int n = 0; n < NumOuts; n++) {
|
||||
SliceTensor(&tmp_out[n], (*outs)[n], out_compute_dims, out_offset);
|
||||
new_outs.emplace_back(&tmp_out[n]);
|
||||
}
|
||||
|
||||
BroadcastKernelForDifferentVecSize<OutT, Functor, kArity, NumOuts>(
|
||||
dev_ctx, new_ins, &new_outs, axis, func);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int kArity, int NumOuts = 1>
|
||||
void BroadcastKernelApply(const KPDevice &dev_ctx,
|
||||
const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
int axis,
|
||||
Functor func) {
|
||||
#ifndef PADDLE_WITH_XPU_KP
|
||||
auto compute_size =
|
||||
static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
|
||||
bool use_int64_index_kernel = false;
|
||||
for (auto *out : *outs) {
|
||||
if (out->numel() >= compute_size) {
|
||||
use_int64_index_kernel = true;
|
||||
}
|
||||
}
|
||||
if (use_int64_index_kernel) { // use_int64_index_kernel
|
||||
BroadcastKernelSplit<OutT, Functor, kArity, NumOuts>(
|
||||
dev_ctx, ins, outs, axis, func, compute_size);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
BroadcastKernelForDifferentVecSize<OutT, Functor, kArity, NumOuts>(
|
||||
dev_ctx, ins, outs, axis, func);
|
||||
}
|
||||
|
||||
template <typename OutT, typename Functor, int NumOuts = 1>
|
||||
void BroadcastKernel(const KPDevice &dev_ctx,
|
||||
const std::vector<const DenseTensor *> &ins,
|
||||
std::vector<DenseTensor *> *outs,
|
||||
Functor func,
|
||||
int axis = -1) {
|
||||
// When there are multiple inputs, the outputs's rank should be equal the
|
||||
// maximum rank of all inputs.
|
||||
using Traits = funcs::FunctionTraits<Functor>;
|
||||
const int kArity = Traits::arity;
|
||||
|
||||
#ifdef PADDLE_WITH_XPU_KP
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ins.size(),
|
||||
2,
|
||||
common::errors::InvalidArgument(
|
||||
"XPU only support inputs is 2, but received %d", ins.size()));
|
||||
#endif
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ins.size(),
|
||||
kArity,
|
||||
common::errors::InvalidArgument("The number of inputs is expected to be "
|
||||
"equal to the "
|
||||
"arity of functor. But received: the "
|
||||
"number of inputs "
|
||||
"is %d, the arity of functor is %d.",
|
||||
ins.size(),
|
||||
kArity));
|
||||
PADDLE_ENFORCE_EQ(
|
||||
outs->size(),
|
||||
NumOuts,
|
||||
common::errors::InvalidArgument("Number of outputs shall equal to number "
|
||||
"of functions, "
|
||||
"but number of outputs is %d, of "
|
||||
"functions is %d.",
|
||||
outs->size(),
|
||||
NumOuts));
|
||||
|
||||
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 = 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) {
|
||||
// When there is only 1 input, the input's rank may be less than outputs'
|
||||
// rank.
|
||||
max_rank = std::max(max_rank, (*outs)[0]->dims().size());
|
||||
}
|
||||
axis = axis == -1 ? max_rank - min_rank : axis;
|
||||
BroadcastKernelApply<OutT, Functor, kArity, NumOuts>(
|
||||
dev_ctx, ins, outs, axis, func);
|
||||
}
|
||||
|
||||
template <typename Functor, typename T, typename OutType = T>
|
||||
void ElementwiseCompute(const GPUContext &dev_ctx,
|
||||
const DenseTensor &x,
|
||||
const DenseTensor &y,
|
||||
Functor func,
|
||||
DenseTensor *z,
|
||||
int axis = -1) {
|
||||
std::vector<const DenseTensor *> ins = {&x, &y};
|
||||
std::vector<DenseTensor *> outs = {z};
|
||||
dev_ctx.template Alloc<OutType>(z);
|
||||
|
||||
BroadcastKernel<OutType, Functor, 1>(dev_ctx, ins, &outs, func, axis);
|
||||
}
|
||||
|
||||
template <typename DeviceContext,
|
||||
typename T,
|
||||
typename Functor,
|
||||
typename InverseFunctor>
|
||||
void DefaultElementwiseOperator(const DeviceContext &dev_ctx,
|
||||
const DenseTensor &x,
|
||||
const DenseTensor &y,
|
||||
DenseTensor *z,
|
||||
int axis = -1) {
|
||||
auto x_dims = x.dims();
|
||||
auto y_dims = y.dims();
|
||||
dev_ctx.template Alloc<T>(z);
|
||||
funcs::ElementwiseCompute<Functor, T>(dev_ctx, x, y, Functor(), z, axis);
|
||||
}
|
||||
|
||||
#else
|
||||
|
||||
template <typename DeviceContext,
|
||||
typename T,
|
||||
typename Functor,
|
||||
typename InverseFunctor>
|
||||
void DefaultElementwiseOperator(const DeviceContext &dev_ctx,
|
||||
const DenseTensor &x,
|
||||
const DenseTensor &y,
|
||||
DenseTensor *z,
|
||||
int axis = -1) {
|
||||
auto x_dims = x.dims();
|
||||
auto y_dims = y.dims();
|
||||
dev_ctx.template Alloc<T>(z);
|
||||
if (x_dims.size() >= y_dims.size()) {
|
||||
funcs::ElementwiseCompute<Functor, T>(dev_ctx, x, y, Functor(), z, axis);
|
||||
} else {
|
||||
funcs::ElementwiseCompute<InverseFunctor, T>(
|
||||
dev_ctx, x, y, InverseFunctor(), z, axis);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
} // namespace funcs
|
||||
} // namespace phi
|
||||
Reference in New Issue
Block a user