wehub-resource-sync
134 lines
4.7 KiB
C++
134 lines
4.7 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
// CUDA and HIP use ReduceGpuKernel API, XPU use ReduceKernel API.
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
|
|
defined(PADDLE_WITH_XPU_KP)
|
|
|
|
#include "paddle/phi/core/visit_type.h"
|
|
#include "paddle/phi/kernels/funcs/reduce_function.h"
|
|
#include "paddle/phi/kernels/funcs/reduce_gpu_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T,
|
|
template <typename>
|
|
class ReduceOp,
|
|
template <typename, typename>
|
|
class TransformOp,
|
|
bool IsMean = false>
|
|
void Reduce(const KPDevice& dev_ctx,
|
|
const DenseTensor& x,
|
|
bool reduce_all,
|
|
const std::vector<int64_t>& dims,
|
|
bool keep_dim, // unused
|
|
DataType out_dtype,
|
|
DenseTensor* out) {
|
|
reduce_all = recompute_reduce_all(x, dims, reduce_all);
|
|
std::vector<int> reduce_dims =
|
|
funcs::details::GetReduceDim(dims, x.dims().size(), reduce_all);
|
|
|
|
int64_t reduce_num = 1;
|
|
for (auto i : reduce_dims) {
|
|
reduce_num *= (x.dims())[i];
|
|
}
|
|
#ifdef PADDLE_WITH_XPU_KP
|
|
using MPType = typename MPTypeTrait<T>::Type;
|
|
funcs::ReduceKernel<T, T, ReduceOp, TransformOp<T, MPType>, IsMean>(
|
|
dev_ctx, x, out, TransformOp<T, MPType>(reduce_num), reduce_dims);
|
|
#else
|
|
if (out_dtype != DataType::UNDEFINED && out_dtype != x.dtype()) {
|
|
auto tmp_tensor = Cast<T>(dev_ctx, x, out_dtype);
|
|
PD_VISIT_BOOL_AND_FLOATING_AND_COMPLEX_AND_4_TYPES(
|
|
DataType::INT32,
|
|
DataType::INT64,
|
|
DataType::FLOAT16,
|
|
DataType::BFLOAT16,
|
|
out_dtype,
|
|
"ReduceKernel",
|
|
([&] {
|
|
using MPType = typename MPTypeTrait<data_t>::Type;
|
|
funcs::ReduceKernel<data_t,
|
|
data_t,
|
|
ReduceOp,
|
|
TransformOp<data_t, MPType>,
|
|
IsMean>(dev_ctx,
|
|
tmp_tensor,
|
|
out,
|
|
TransformOp<data_t, MPType>(reduce_num),
|
|
reduce_dims);
|
|
}));
|
|
} else {
|
|
using MPType = typename MPTypeTrait<T>::Type;
|
|
funcs::ReduceKernel<T, T, ReduceOp, TransformOp<T, MPType>, IsMean>(
|
|
dev_ctx, x, out, TransformOp<T, MPType>(reduce_num), reduce_dims);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
|
template <typename T, template <typename, typename, typename> class ReduceOp>
|
|
void Reduce(const KPDevice& dev_ctx,
|
|
const DenseTensor& x,
|
|
bool reduce_all,
|
|
const std::vector<int64_t>& dims,
|
|
DataType out_dtype,
|
|
DenseTensor* out) {
|
|
reduce_all = recompute_reduce_all(x, dims, reduce_all);
|
|
std::vector<int> reduce_dims =
|
|
funcs::details::GetReduceDim(dims, x.dims().size(), reduce_all);
|
|
|
|
int64_t reduce_num = 1;
|
|
for (auto i : reduce_dims) {
|
|
reduce_num *= (x.dims())[i];
|
|
}
|
|
|
|
if (out_dtype != DataType::UNDEFINED && out_dtype != x.dtype()) {
|
|
if (x.dtype() == DataType::BFLOAT16 && out_dtype == DataType::FLOAT32) {
|
|
phi::funcs::ReduceGpuKernel<phi::bfloat16, float, ReduceOp>(
|
|
dev_ctx, x, out, reduce_dims);
|
|
} else if (x.dtype() == DataType::FLOAT16 &&
|
|
out_dtype == DataType::FLOAT32) {
|
|
phi::funcs::ReduceGpuKernel<phi::float16, float, ReduceOp>(
|
|
dev_ctx, x, out, reduce_dims);
|
|
} else {
|
|
auto tmp_tensor = Cast<T>(dev_ctx, x, out_dtype);
|
|
tmp_tensor.set_strides(x.strides());
|
|
|
|
PD_VISIT_BOOL_AND_FLOATING_AND_COMPLEX_AND_4_TYPES(
|
|
DataType::INT32,
|
|
DataType::INT64,
|
|
DataType::FLOAT16,
|
|
DataType::BFLOAT16,
|
|
out_dtype,
|
|
"ReduceGpuKernel",
|
|
([&] {
|
|
using MPType = typename MPTypeTrait<data_t>::Type;
|
|
phi::funcs::ReduceGpuKernel<data_t, data_t, ReduceOp>(
|
|
dev_ctx, tmp_tensor, out, reduce_dims);
|
|
}));
|
|
}
|
|
} else {
|
|
using MPType = typename MPTypeTrait<T>::Type;
|
|
phi::funcs::ReduceGpuKernel<T, T, ReduceOp>(dev_ctx, x, out, reduce_dims);
|
|
}
|
|
}
|
|
#endif
|
|
|
|
} // namespace phi
|
|
|
|
#endif
|