wehub-resource-sync
1470 lines
47 KiB
C++
1470 lines
47 KiB
C++
// 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 <bitset>
|
|
#include <limits>
|
|
#include <set>
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/dense_tensor_iterator.h"
|
|
#include "paddle/phi/kernels/funcs/index_elementwise.cu.h"
|
|
#include "paddle/phi/kernels/funcs/scatter.cu.h"
|
|
#include "paddle/phi/kernels/gpu/reduce.h"
|
|
#include "paddle/phi/kernels/legacy/reduce_max_kernel.h"
|
|
#include "paddle/phi/kernels/prod_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_all_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_amin_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_any_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_max_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_mean_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_min_kernel.h"
|
|
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
|
|
|
#include "paddle/phi/kernels/funcs/function_traits.h"
|
|
#include "paddle/phi/kernels/primitive/reduce_primitives.h"
|
|
|
|
#ifdef PADDLE_WITH_HIP
|
|
#define WARP_SIZE 64
|
|
#else
|
|
#define WARP_SIZE 32
|
|
#endif
|
|
|
|
// The GPUReduceScheduler splits tensors with indices exceeding 32-bit range to
|
|
// ensure that all incoming tensors can be addressed within 32-bit index space.
|
|
using IndexType = uint32_t;
|
|
|
|
template <typename T>
|
|
struct LoadImpl {
|
|
HOSTDEVICE static T Apply(const void* src) {
|
|
return *reinterpret_cast<const T*>(src);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct LoadImpl<bool> {
|
|
HOSTDEVICE static bool Apply(const void* src) {
|
|
static_assert(sizeof(bool) == sizeof(char));
|
|
return *reinterpret_cast<const unsigned char*>(src);
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
HOSTDEVICE constexpr T LoadData(const void* src) {
|
|
return LoadImpl<T>::Apply(src);
|
|
}
|
|
|
|
template <typename ScalarT>
|
|
HOSTDEVICE constexpr ScalarT LoadData(const ScalarT* src) {
|
|
return LoadImpl<ScalarT>::Apply(src);
|
|
}
|
|
|
|
namespace phi {
|
|
inline std::bitset<64> DimListToBitset(std::vector<int> opt_dims,
|
|
size_t ndims) {
|
|
std::bitset<64> dim_mask;
|
|
|
|
if (opt_dims.size() > 0) {
|
|
for (int dim : opt_dims) {
|
|
dim_mask.set(dim, true);
|
|
}
|
|
} else {
|
|
for (size_t dim = 0; dim < ndims; dim++) {
|
|
dim_mask.set(dim, true);
|
|
}
|
|
}
|
|
return dim_mask;
|
|
}
|
|
|
|
inline std::vector<int> ConvertToPositiveDims(
|
|
const std::vector<int>& origin_reduce_dims, int64_t ndim) {
|
|
std::vector<int> positive_reduce_dims = origin_reduce_dims;
|
|
for (size_t i = 0; i < origin_reduce_dims.size(); ++i) {
|
|
PADDLE_ENFORCE_GE(
|
|
origin_reduce_dims[i],
|
|
-ndim,
|
|
common::errors::InvalidArgument(
|
|
"ReduceOp: invalid axis, when x_dims is %d, "
|
|
"axis[i] should be in the range of [-%d, %d), but got %d.",
|
|
ndim,
|
|
ndim,
|
|
ndim,
|
|
origin_reduce_dims[i]));
|
|
PADDLE_ENFORCE_LT(
|
|
origin_reduce_dims[i],
|
|
ndim,
|
|
common::errors::InvalidArgument(
|
|
"ReduceOp: invalid axis, when x_dims is %d, "
|
|
"axis[i] should be in the range of [-%d, %d), but got %d.",
|
|
ndim,
|
|
ndim,
|
|
ndim,
|
|
origin_reduce_dims[i]));
|
|
|
|
if (origin_reduce_dims[i] < 0) {
|
|
positive_reduce_dims[i] = ndim + origin_reduce_dims[i];
|
|
}
|
|
}
|
|
return positive_reduce_dims;
|
|
}
|
|
|
|
inline std::bitset<64> MakeDimMask(std::vector<int> opt_dims,
|
|
int64_t ndim,
|
|
bool allow_empty_dims = false) {
|
|
// flip() sets all bits to 1 (masking all dimensions for reduction).
|
|
if (opt_dims.empty() && !allow_empty_dims) {
|
|
return std::bitset<64>().flip();
|
|
}
|
|
|
|
// Otherwise, use the dimensions specified in opt_dims.
|
|
return DimListToBitset(opt_dims, ndim);
|
|
}
|
|
|
|
inline DenseTensor ReviewReduceResult(const DenseTensor& src,
|
|
const DenseTensor& result,
|
|
int ndim,
|
|
std::bitset<64> mask) {
|
|
std::vector<int64_t> shape;
|
|
std::vector<int64_t> stride;
|
|
|
|
int64_t cal_stride = 1;
|
|
const auto& src_dims = src.dims();
|
|
|
|
for (int dim = ndim - 1; dim >= 0; dim--) {
|
|
if (!mask[dim]) {
|
|
shape.insert(shape.begin(), src_dims[dim]);
|
|
stride.insert(stride.begin(), cal_stride);
|
|
cal_stride *= src_dims[dim];
|
|
} else {
|
|
shape.insert(shape.begin(), 1);
|
|
stride.insert(stride.begin(), cal_stride);
|
|
}
|
|
}
|
|
|
|
return funcs::as_strided(result, shape, stride);
|
|
}
|
|
|
|
template <typename T, int Size>
|
|
DEVICE AlignedVector<T, Size> LoadVector(const T* base_ptr, uint32_t offset) {
|
|
using vec_t = AlignedVector<T, Size>;
|
|
auto* from = reinterpret_cast<const vec_t*>(base_ptr);
|
|
return from[offset];
|
|
}
|
|
|
|
template <int Size>
|
|
DEVICE AlignedVector<bool, Size> LoadVector(const bool* base_ptr,
|
|
uint32_t offset) {
|
|
auto tmp = LoadVector<uint8_t, Size>(
|
|
reinterpret_cast<const uint8_t*>(base_ptr), offset);
|
|
AlignedVector<bool, Size> ret;
|
|
for (int i = 0; i < Size; ++i) {
|
|
ret.val[i] = static_cast<bool>(tmp.val[i]);
|
|
}
|
|
return ret;
|
|
}
|
|
|
|
// Chose max num threads.
|
|
template <typename T>
|
|
struct MaxThreadsConfig {
|
|
static constexpr int MAX_NUM_THREADS = 512;
|
|
};
|
|
|
|
template <>
|
|
struct MaxThreadsConfig<phi::dtype::complex<double>> {
|
|
static constexpr int MAX_NUM_THREADS = 256;
|
|
};
|
|
|
|
template <int kNumThreads, int kOutputVecSize, typename Reducer>
|
|
__launch_bounds__(kNumThreads, 4) __global__
|
|
void VecReduceKernel(Reducer reduction) {
|
|
reduction.template Run<kOutputVecSize>();
|
|
}
|
|
|
|
template <typename IndexType>
|
|
static funcs::OffsetCalculator<2, IndexType> MakeOutputOffsetCalculator(
|
|
const DenseTensorIterator& iter) {
|
|
int num_reduce_dims = iter.num_reduce_dims();
|
|
int num_output_dims = iter.ndim() - num_reduce_dims;
|
|
int input_index = iter.ntensors() - 1;
|
|
int output_index = 0;
|
|
|
|
std::array<const int64_t*, 2> stride_ptrs = {
|
|
iter.strides(output_index).data() + num_reduce_dims,
|
|
iter.strides(input_index).data() + num_reduce_dims,
|
|
};
|
|
|
|
auto output_shape_ptr = iter.shape().data() + num_reduce_dims;
|
|
|
|
return funcs::OffsetCalculator<2, IndexType>(
|
|
num_output_dims, output_shape_ptr, stride_ptrs.data());
|
|
}
|
|
|
|
template <typename IndexType>
|
|
static funcs::OffsetCalculator<1, IndexType> MakeInputOffsetCalculator(
|
|
const DenseTensorIterator& iter) {
|
|
int num_reduce_dims = iter.num_reduce_dims();
|
|
int input_index = iter.ntensors() - 1;
|
|
|
|
std::array<const int64_t*, 1> strides = {
|
|
iter.strides(input_index).data(),
|
|
};
|
|
|
|
auto input_shape_ptr = iter.shape().data();
|
|
|
|
return funcs::OffsetCalculator<1, IndexType>(
|
|
num_reduce_dims, input_shape_ptr, strides.data());
|
|
}
|
|
|
|
template <typename T>
|
|
int GetOutputVecSize(const DenseTensorIterator& iter) {
|
|
int vec_size = 4;
|
|
|
|
auto UpdateVectorSize = [&vec_size](uint64_t n) {
|
|
while (n % vec_size != 0) {
|
|
vec_size /= 2;
|
|
}
|
|
};
|
|
|
|
// Check base address alignment.
|
|
uint64_t base_address =
|
|
reinterpret_cast<uint64_t>(iter.data_ptr(iter.noutputs())) / sizeof(T);
|
|
UpdateVectorSize(base_address);
|
|
|
|
// Check output dimension size.
|
|
const int output_index = iter.num_reduce_dims();
|
|
UpdateVectorSize(iter.shape()[output_index]);
|
|
|
|
// Check strides alignment for all dimensions except output dimension.
|
|
auto input_tensor_index = iter.noutputs();
|
|
auto input_strides = iter.strides(input_tensor_index);
|
|
|
|
for (int dim = 0; dim < input_strides.size(); ++dim) {
|
|
if (dim != output_index) {
|
|
UpdateVectorSize(input_strides[dim] / sizeof(T));
|
|
}
|
|
}
|
|
|
|
return vec_size;
|
|
}
|
|
|
|
// Simplify fraction by dividing both numerator and denominator by their GCD
|
|
// (Greatest Common Divisor).
|
|
HOSTDEVICE static void ReduceFraction(size_t* numerator, size_t* denominator) {
|
|
size_t a = *denominator;
|
|
size_t b = *numerator;
|
|
while (b != 0) {
|
|
a %= b;
|
|
size_t tmp = a;
|
|
a = b;
|
|
b = tmp;
|
|
}
|
|
|
|
*numerator /= a;
|
|
*denominator /= a;
|
|
}
|
|
|
|
struct ReduceConfig {
|
|
static constexpr int BLOCK_X = 0;
|
|
static constexpr int BLOCK_Y = 1;
|
|
static constexpr int CTA = 2;
|
|
|
|
ReduceConfig(int element_size_bytes, int num_outputs, int num_inputs)
|
|
: element_size_bytes(element_size_bytes),
|
|
num_inputs(num_inputs),
|
|
num_outputs(num_outputs) {}
|
|
|
|
// Basic configuration.
|
|
int element_size_bytes;
|
|
int num_inputs;
|
|
int num_outputs;
|
|
|
|
// Parallelism control.
|
|
int step_input = 1;
|
|
int step_output = 1;
|
|
int ctas_per_output = 1;
|
|
|
|
// Multiplier arrays for index calculation.
|
|
int input_multiplier[3] = {0, 0, 0};
|
|
int output_multiplier[2] = {0, 0};
|
|
|
|
// Dimensions.
|
|
int block_width;
|
|
int block_height;
|
|
int num_threads;
|
|
|
|
// Vectorization control.
|
|
bool vectorize_input = false;
|
|
int output_vec_size = 1;
|
|
|
|
template <typename T>
|
|
void SetBlockDimensions(int64_t dim0, int64_t dim1) {
|
|
const int max_num_threads =
|
|
MaxThreadsConfig<T>::MAX_NUM_THREADS / output_vec_size;
|
|
|
|
int dim0_pow2 =
|
|
(dim0 < max_num_threads)
|
|
? static_cast<int>(phi::backends::gpu::GetLastPow2(dim0))
|
|
: max_num_threads;
|
|
int dim1_pow2 =
|
|
(dim1 < max_num_threads)
|
|
? static_cast<int>(phi::backends::gpu::GetLastPow2(dim1))
|
|
: max_num_threads;
|
|
block_width = std::min(dim0_pow2, WARP_SIZE);
|
|
block_height =
|
|
std::min(dim1_pow2, static_cast<int>(max_num_threads / block_width));
|
|
block_width =
|
|
std::min(dim0_pow2, static_cast<int>(max_num_threads / block_height));
|
|
num_threads = block_width * block_height;
|
|
}
|
|
|
|
int SplitInput(int parallelism) {
|
|
const int current_step = step_input;
|
|
const int64_t next_step = static_cast<int64_t>(step_input) * parallelism;
|
|
PADDLE_ENFORCE_LE_INT_MAX(next_step, "reduce step_input");
|
|
step_input = static_cast<int>(next_step);
|
|
return current_step;
|
|
}
|
|
|
|
int SplitOutput(int parallelism) {
|
|
const int current_step = step_output;
|
|
const int64_t next_step = static_cast<int64_t>(step_output) * parallelism;
|
|
PADDLE_ENFORCE_LE_INT_MAX(next_step, "reduce step_output");
|
|
step_output = static_cast<int>(next_step);
|
|
return current_step;
|
|
}
|
|
|
|
dim3 GetBlockDim() const { return dim3(block_width, block_height); }
|
|
|
|
dim3 GetGridDim() const {
|
|
const int64_t grid_x = phi::backends::gpu::DivUp<int64_t>(
|
|
num_outputs / output_vec_size, step_output);
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(grid_x, "grid.x");
|
|
PADDLE_ENFORCE_LE_UINT32_MAX(ctas_per_output, "grid.y");
|
|
return dim3(static_cast<uint32_t>(grid_x),
|
|
static_cast<uint32_t>(ctas_per_output));
|
|
}
|
|
|
|
HOSTDEVICE bool ShouldReduceBlockX() const {
|
|
return input_multiplier[BLOCK_X] != 0;
|
|
}
|
|
|
|
HOSTDEVICE bool ShouldReduceBlockY() const {
|
|
return input_multiplier[BLOCK_Y] != 0;
|
|
}
|
|
|
|
HOSTDEVICE bool ShouldReduceGlobal() const {
|
|
return input_multiplier[CTA] != 0;
|
|
}
|
|
|
|
DEVICE bool ShouldStore(int output_idx) const {
|
|
// 1. Boundary Check: Ensure the output index is within the valid range.
|
|
// If out of bounds, no storage is necessary.
|
|
if (output_idx >= num_outputs) {
|
|
return false;
|
|
}
|
|
|
|
// 2. X-Reduction Check: If block-wide X-reduction is active, only the
|
|
// thread with index 0 in the X-dimension (the "leader") is allowed to
|
|
// store.
|
|
if (ShouldReduceBlockX() && threadIdx.x != 0) {
|
|
return false;
|
|
}
|
|
|
|
// 3. Y-Reduction Check: If block-wide Y-reduction is active, only the
|
|
// thread with index 0 in the Y-dimension (the "leader") is allowed to
|
|
// store.
|
|
if (ShouldReduceBlockY() && threadIdx.y != 0) {
|
|
return false;
|
|
}
|
|
|
|
// If the thread passes all checks, it is the designated thread to store the
|
|
// result.
|
|
return true;
|
|
}
|
|
|
|
DEVICE bool ShouldReduceTail() const {
|
|
return (!ShouldReduceBlockY() || threadIdx.y == 0) &&
|
|
(!ShouldReduceGlobal() || blockIdx.y == 0);
|
|
}
|
|
|
|
HOSTDEVICE int GetInIdx() const {
|
|
int thread_x = threadIdx.x;
|
|
int thread_y = threadIdx.y;
|
|
int block_y = blockIdx.y;
|
|
return (thread_x * input_multiplier[BLOCK_X] +
|
|
thread_y * input_multiplier[BLOCK_Y] +
|
|
block_y * input_multiplier[CTA]);
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
HOSTDEVICE int GetOutIdx() const {
|
|
int thread_x = threadIdx.x;
|
|
int thread_y = threadIdx.y;
|
|
int block_x = blockIdx.x;
|
|
return (thread_x * output_multiplier[BLOCK_X] +
|
|
thread_y * output_multiplier[BLOCK_Y] + block_x * step_output) *
|
|
kOutputVecSize;
|
|
}
|
|
|
|
DEVICE int SharedMemoryOffset(int offset) const {
|
|
return threadIdx.x + (threadIdx.y + offset) * blockDim.x;
|
|
}
|
|
|
|
DEVICE int StagingMemoryOffset(int block_y) const {
|
|
IndexType offset = block_y + static_cast<IndexType>(blockIdx.x) *
|
|
static_cast<IndexType>(gridDim.y);
|
|
if (!ShouldReduceBlockX()) {
|
|
offset = threadIdx.x + offset * blockDim.x;
|
|
}
|
|
|
|
return offset;
|
|
}
|
|
|
|
int SharedMemorySize() const {
|
|
if (!ShouldReduceBlockY() &&
|
|
(!ShouldReduceBlockX() || block_width <= WARP_SIZE)) {
|
|
return 0;
|
|
}
|
|
|
|
return element_size_bytes * num_threads * output_vec_size;
|
|
}
|
|
|
|
int64_t GlobalMemorySize() const {
|
|
if (!ShouldReduceGlobal()) {
|
|
return 0;
|
|
}
|
|
|
|
int64_t size = static_cast<int64_t>(element_size_bytes) * num_outputs *
|
|
ctas_per_output;
|
|
if (!ShouldReduceBlockX()) {
|
|
size *= GetBlockDim().x * output_vec_size;
|
|
}
|
|
|
|
return size;
|
|
}
|
|
|
|
size_t SemaphoreSize() const {
|
|
if (!ShouldReduceGlobal()) {
|
|
return 0;
|
|
}
|
|
return sizeof(int) * static_cast<size_t>(GetGridDim().x);
|
|
}
|
|
|
|
int ValuesPerThread() const {
|
|
int64_t val = phi::backends::gpu::DivUp<int64_t>(num_inputs, step_input);
|
|
PADDLE_ENFORCE_LE_INT_MAX(val, "values per thread");
|
|
return static_cast<int>(val);
|
|
}
|
|
};
|
|
|
|
template <typename MPType,
|
|
typename ScalarT,
|
|
int kVecSize,
|
|
int kInputVecSize = kVecSize>
|
|
ReduceConfig SetReduceConfig(const DenseTensorIterator& iter) {
|
|
int device_id = paddle::platform::GetCurrentDeviceId();
|
|
|
|
int64_t num_outputs = iter.num_output_elements();
|
|
int64_t inputs_per_output = iter.numel() / num_outputs;
|
|
int input_index = iter.ntensors() - 1;
|
|
|
|
auto config = ReduceConfig(sizeof(MPType), num_outputs, inputs_per_output);
|
|
|
|
int64_t dim0;
|
|
int64_t dim1;
|
|
int64_t fastest_moving_stride;
|
|
bool reduce_fastest_dim;
|
|
|
|
if (iter.ndim() > 0) {
|
|
// Check if we're reducing along the fastest-changing dimension
|
|
// This affects memory access patterns for better performance.
|
|
reduce_fastest_dim = (iter.num_reduce_dims() == iter.ndim()) ||
|
|
(iter.strides(input_index)[0] <
|
|
iter.strides(input_index)[iter.num_reduce_dims()]);
|
|
|
|
// Set block dimensions based on reduction pattern.
|
|
if (reduce_fastest_dim) {
|
|
// Reducing along fastest dimension: use block.x for reduction.
|
|
// block.x handles reduction elements.
|
|
// block.y handles output elements.
|
|
dim0 = inputs_per_output;
|
|
dim1 = num_outputs;
|
|
fastest_moving_stride = iter.strides(input_index)[0];
|
|
} else {
|
|
// Not reducing along fastest dimension: use block.x for outputs.
|
|
// block.x handles output elements.
|
|
// block.y handles reduction elements.
|
|
dim0 = num_outputs;
|
|
dim1 = inputs_per_output;
|
|
fastest_moving_stride = iter.strides(input_index)[iter.num_reduce_dims()];
|
|
}
|
|
} else {
|
|
// Handle 0-dimensional case.
|
|
reduce_fastest_dim = true;
|
|
fastest_moving_stride = sizeof(ScalarT);
|
|
dim0 = 1;
|
|
dim1 = 1;
|
|
}
|
|
|
|
// Use vectorization for better memory access. Two cases:
|
|
// Case 1: "Vectorize along input" - when reducing on fastest dimension,
|
|
// data in same vector corresponds to the same output.
|
|
// Case 2: "Vectorize along output" - when fastest dimension is not reduced,
|
|
// data in same vector corresponds to different outputs.
|
|
if (fastest_moving_stride == sizeof(ScalarT)) {
|
|
if (reduce_fastest_dim && dim0 >= 128 && iter.num_reduce_dims() == 1) {
|
|
// Case 1: Vectorize along input (load data for same output together).
|
|
config.vectorize_input = true;
|
|
dim0 /= kInputVecSize;
|
|
} else if (!reduce_fastest_dim) {
|
|
// Case 2: Vectorize along output (load data for multiple outputs
|
|
// together).
|
|
config.output_vec_size = GetOutputVecSize<ScalarT>(iter);
|
|
dim0 /= config.output_vec_size;
|
|
}
|
|
}
|
|
|
|
// Adjust block_width and block_height.
|
|
config.SetBlockDimensions<ScalarT>(dim0, dim1);
|
|
|
|
int block_width = config.block_width;
|
|
int block_height = config.block_height;
|
|
|
|
// Level 1 parallelization: split work at thread level.
|
|
if (iter.ndim() == 0 || reduce_fastest_dim) {
|
|
// Case 1: Split input across threads (requires thread synchronization).
|
|
config.input_multiplier[0] = config.SplitInput(block_width);
|
|
} else {
|
|
// Case 2: Split output across threads (each thread handles different
|
|
// output).
|
|
config.output_multiplier[0] = config.SplitOutput(block_width);
|
|
}
|
|
|
|
// Min elements per thread.
|
|
constexpr int min_values_per_thread = 16;
|
|
// Max elements per thread.
|
|
constexpr int max_values_per_thread = 256;
|
|
|
|
// Decide if we need to split work across warps.
|
|
const int warp_split_threshold =
|
|
std::min<int>(block_height * 16, max_values_per_thread);
|
|
bool split_across_warps = config.ValuesPerThread() >= warp_split_threshold;
|
|
|
|
const int num_mp = paddle::platform::GetGPUMultiProcessors(device_id);
|
|
|
|
// Level 2 parallelization: split work at warp level.
|
|
if (split_across_warps) {
|
|
// Case 1: Split input across warps (requires warp synchronization).
|
|
config.input_multiplier[1] = config.SplitInput(block_height);
|
|
} else {
|
|
// Case 2: Each warp handles independent outputs.
|
|
config.output_multiplier[1] = config.SplitOutput(block_height);
|
|
}
|
|
|
|
int max_threads_per_mp =
|
|
paddle::platform::GetGPUMaxThreadsPerMultiProcessor(device_id);
|
|
|
|
const int blocks_per_sm = max_threads_per_mp / config.num_threads;
|
|
const int target_grid_size = num_mp * blocks_per_sm;
|
|
int grid = config.GetGridDim().x;
|
|
|
|
// Level 3 parallelization: split work at block level (for large datasets).
|
|
if (config.input_multiplier[1] != 0 &&
|
|
config.ValuesPerThread() >= max_values_per_thread &&
|
|
grid <= target_grid_size) {
|
|
// Calculate optimal block splitting strategy.
|
|
// Based on SM utilization.
|
|
const int64_t ctas_per_output1 =
|
|
phi::backends::gpu::DivUp<int64_t>(target_grid_size, grid);
|
|
// Based on min workload.
|
|
const int64_t ctas_per_output2 = phi::backends::gpu::DivUp<int64_t>(
|
|
config.ValuesPerThread(), min_values_per_thread);
|
|
// Based on max workload.
|
|
const int64_t ctas_per_output3 = phi::backends::gpu::DivUp<int64_t>(
|
|
config.ValuesPerThread(), max_values_per_thread);
|
|
|
|
// Choose best splitting strategy to balance parallelism and per-thread
|
|
// workload.
|
|
const int64_t ctas_per_output = std::max(
|
|
std::min(ctas_per_output1, ctas_per_output2), ctas_per_output3);
|
|
PADDLE_ENFORCE_LE_INT_MAX(ctas_per_output, "reduce ctas_per_output");
|
|
config.ctas_per_output = static_cast<int>(ctas_per_output);
|
|
|
|
if (config.ctas_per_output > 1) {
|
|
// Case 3: Split input across blocks (requires global memory
|
|
// synchronization).
|
|
config.input_multiplier[2] = config.SplitInput(config.ctas_per_output);
|
|
}
|
|
}
|
|
return config;
|
|
}
|
|
|
|
template <typename ScalarT,
|
|
typename ReduceOp,
|
|
typename OutScalarT = ScalarT,
|
|
int kVecSize = 4,
|
|
int kInputVecSize = kVecSize>
|
|
struct ReduceExecutor {
|
|
using traits = phi::funcs::FunctionTraits<decltype(&ReduceOp::reduce)>;
|
|
using MPType =
|
|
typename std::decay<typename traits::template arg<0>::type>::type;
|
|
|
|
using InputCalculator = funcs::OffsetCalculator<1, IndexType>;
|
|
using OutputCalculator = funcs::OffsetCalculator<2, IndexType>;
|
|
|
|
static constexpr bool can_accumulate_in_output =
|
|
std::is_convertible_v<MPType, OutScalarT> &&
|
|
std::is_convertible_v<OutScalarT, MPType>;
|
|
|
|
// Core reduction algorithm configuration.
|
|
ReduceOp reducer;
|
|
ReduceConfig config;
|
|
MPType ident;
|
|
|
|
// Data access calculators for input and output indexing.
|
|
InputCalculator input_calc;
|
|
OutputCalculator output_calc;
|
|
|
|
// Data pointers for source, destination, and buffers.
|
|
const void* src;
|
|
char* dst[2];
|
|
void* acc_buf;
|
|
void* cta_buf;
|
|
|
|
// Parallel synchronization primitives.
|
|
int* semaphores;
|
|
|
|
// Runtime state and control flags.
|
|
int64_t base_idx;
|
|
bool accumulate;
|
|
bool final_output;
|
|
int noutputs;
|
|
|
|
ReduceExecutor(ReduceOp reducer,
|
|
ReduceConfig config,
|
|
MPType ident,
|
|
InputCalculator input_calc,
|
|
OutputCalculator output_calc,
|
|
const void* src,
|
|
char* dst0,
|
|
std::optional<char*> dst1,
|
|
void* acc_buf,
|
|
void* cta_buf,
|
|
int* semaphores,
|
|
int base_idx,
|
|
bool accumulate,
|
|
bool final_output,
|
|
int64_t noutputs)
|
|
: reducer(reducer),
|
|
config(config),
|
|
ident(ident),
|
|
input_calc(input_calc),
|
|
output_calc(output_calc),
|
|
src(src),
|
|
acc_buf(acc_buf),
|
|
cta_buf(cta_buf),
|
|
semaphores(semaphores),
|
|
base_idx(base_idx),
|
|
accumulate(accumulate),
|
|
final_output(final_output),
|
|
noutputs(noutputs) {
|
|
dst[0] = dst0;
|
|
if (dst1.has_value()) {
|
|
dst[1] = dst1.value();
|
|
}
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE void Run() const {
|
|
extern __shared__ char shared_memory[];
|
|
|
|
IndexType output_idx = config.GetOutIdx<kOutputVecSize>();
|
|
IndexType input_idx = config.GetInIdx();
|
|
auto base_offsets1 = output_calc.get(output_idx)[1];
|
|
|
|
using MPTypeVec = std::array<MPType, kOutputVecSize>;
|
|
MPTypeVec value;
|
|
|
|
if (output_idx < config.num_outputs && input_idx < config.num_inputs) {
|
|
const ScalarT* input_slice =
|
|
(const ScalarT*)((const char*)src + base_offsets1);
|
|
value = ThreadReduce<kOutputVecSize>(input_slice);
|
|
}
|
|
|
|
// Run BlockX reduction before BlockY to match PyTorch's accumulation
|
|
// order (>= 2.10, PR #165178), keeping broadcast-grad sum reductions
|
|
// bitwise-aligned with torch when vectorized loads are used.
|
|
if (config.ShouldReduceBlockX()) {
|
|
value = BlockXReduce<kOutputVecSize>(value, shared_memory);
|
|
}
|
|
|
|
if (config.ShouldReduceBlockY()) {
|
|
value = BlockYReduce<kOutputVecSize>(value, shared_memory);
|
|
}
|
|
|
|
using OutPtrVec = std::array<OutScalarT*, kOutputVecSize>;
|
|
using OffsetVec = std::array<IndexType, kOutputVecSize>;
|
|
|
|
OffsetVec base_offsets;
|
|
OutPtrVec out;
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
base_offsets[i] = output_calc.get(output_idx + i)[0];
|
|
out[i] = reinterpret_cast<OutScalarT*>(dst[0] + base_offsets[i]);
|
|
}
|
|
|
|
MPTypeVec* acc = nullptr;
|
|
if (acc_buf != nullptr) {
|
|
size_t numerator = sizeof(MPType);
|
|
size_t denominator = sizeof(OutScalarT);
|
|
ReduceFraction(&numerator, &denominator);
|
|
acc = reinterpret_cast<MPTypeVec*>(
|
|
reinterpret_cast<char*>(acc_buf) +
|
|
(base_offsets[0] * numerator / denominator));
|
|
}
|
|
|
|
if (config.ShouldReduceGlobal()) {
|
|
value = GlobalReduce<kOutputVecSize>(value, acc, shared_memory);
|
|
} else if (config.ShouldStore(output_idx)) {
|
|
if (acc == nullptr) {
|
|
if (accumulate) {
|
|
value = AccumulateInOutput<kOutputVecSize, can_accumulate_in_output>(
|
|
out, value);
|
|
}
|
|
if (final_output) {
|
|
SetResultsToOutput<kOutputVecSize>(value, base_offsets);
|
|
} else {
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
*(out[i]) = GetAccumulatedOutput<can_accumulate_in_output>(
|
|
out[i], value[i]);
|
|
}
|
|
}
|
|
} else {
|
|
if (accumulate) {
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce((*acc)[i], value[i]);
|
|
}
|
|
}
|
|
if (final_output) {
|
|
SetResultsToOutput<kOutputVecSize>(value, base_offsets);
|
|
} else {
|
|
*acc = value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE std::array<MPType, kOutputVecSize> ThreadReduce(
|
|
const ScalarT* data) const {
|
|
if (config.vectorize_input) {
|
|
return {InVectorizedThreadReduceImpl(data)};
|
|
} else {
|
|
IndexType element_stride = input_calc.strides_[0][0] / sizeof(ScalarT);
|
|
bool is_contiguous = (input_calc.dims == 1 && element_stride == 1);
|
|
if (is_contiguous) {
|
|
return ThreadReduceImpl<kOutputVecSize>(
|
|
data, [](IndexType idx) { return idx; });
|
|
} else if (input_calc.dims == 1) {
|
|
return ThreadReduceImpl<kOutputVecSize>(
|
|
data, [&](IndexType idx) { return idx * element_stride; });
|
|
} else {
|
|
return ThreadReduceImpl<kOutputVecSize>(data, [&](IndexType idx) {
|
|
return input_calc.get(idx)[0] / sizeof(ScalarT);
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
DEVICE MPType InVectorizedThreadReduceImpl(const ScalarT* data) const {
|
|
IndexType end = config.num_inputs;
|
|
MPType value = ident;
|
|
constexpr int align_bytes = alignof(AlignedVector<ScalarT, kInputVecSize>);
|
|
|
|
constexpr int align_elements = align_bytes / sizeof(ScalarT);
|
|
int shift = ((uint64_t)data) % align_bytes / sizeof(ScalarT);
|
|
|
|
if (shift > 0) {
|
|
data -= shift;
|
|
end += shift;
|
|
if (threadIdx.x >= shift && threadIdx.x < align_elements &&
|
|
config.ShouldReduceTail()) {
|
|
value = reducer.compute(value, LoadData(data + threadIdx.x));
|
|
}
|
|
end -= align_elements;
|
|
data += align_elements;
|
|
shift = align_elements - shift;
|
|
}
|
|
|
|
IndexType idx = config.GetInIdx();
|
|
const IndexType stride = config.step_input;
|
|
|
|
std::array<MPType, kInputVecSize> value_list;
|
|
value_list[0] = value;
|
|
|
|
#pragma unroll
|
|
for (int i = 1; i < kInputVecSize; i++) {
|
|
value_list[i] = ident;
|
|
}
|
|
|
|
using load_t = AlignedVector<ScalarT, kInputVecSize>;
|
|
|
|
while (idx * kInputVecSize + kInputVecSize - 1 < end) {
|
|
const auto values_vec = LoadVector<ScalarT, kInputVecSize>(data, idx);
|
|
|
|
#pragma unroll
|
|
for (IndexType i = 0; i < kInputVecSize; i++) {
|
|
value_list[i] = reducer.compute(value_list[i], values_vec.val[i]);
|
|
}
|
|
idx += stride;
|
|
}
|
|
|
|
// Tile processing.
|
|
IndexType tail_start = end - end % kInputVecSize;
|
|
|
|
if (config.ShouldReduceTail()) {
|
|
int idx = tail_start + threadIdx.x;
|
|
if (idx < end) {
|
|
const auto value = LoadData(data + idx);
|
|
value_list[0] = reducer.compute(value_list[0], value);
|
|
}
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int i = 1; i < kInputVecSize; i++) {
|
|
value_list[0] = reducer.reduce(value_list[0], value_list[i]);
|
|
}
|
|
|
|
return value_list[0];
|
|
}
|
|
|
|
template <int kOutputVecSize, typename offset_calc_t>
|
|
DEVICE std::array<MPType, kOutputVecSize> ThreadReduceImpl(
|
|
const ScalarT* data_, offset_calc_t calc) const {
|
|
IndexType idx = config.GetInIdx();
|
|
const IndexType end = config.num_inputs;
|
|
const IndexType stride = config.step_input;
|
|
|
|
using MPTypeVec = std::array<MPType, kOutputVecSize>;
|
|
using load_t = AlignedVector<ScalarT, kOutputVecSize>;
|
|
|
|
std::array<MPTypeVec, kVecSize> value_list;
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < kVecSize; i++) {
|
|
#pragma unroll
|
|
for (int j = 0; j < kOutputVecSize; j++) {
|
|
value_list[i][j] = ident;
|
|
}
|
|
}
|
|
|
|
std::array<load_t, kVecSize> values;
|
|
|
|
while (idx + (kVecSize - 1) * stride < end) {
|
|
#pragma unroll
|
|
for (IndexType i = 0; i < kVecSize; i++) {
|
|
const auto offset = calc(idx + i * stride) / kOutputVecSize;
|
|
values[i] = LoadVector<ScalarT, kOutputVecSize>(data_, offset);
|
|
}
|
|
#pragma unroll
|
|
for (IndexType i = 0; i < kVecSize; i++) {
|
|
#pragma unroll
|
|
for (IndexType j = 0; j < kOutputVecSize; j++) {
|
|
value_list[i][j] =
|
|
reducer.compute(value_list[i][j], values[i].val[j]);
|
|
}
|
|
}
|
|
idx += stride * kVecSize;
|
|
}
|
|
|
|
// tail
|
|
int idx_ = idx;
|
|
#pragma unroll
|
|
for (IndexType i = 0; i < kVecSize; i++) {
|
|
if (idx >= end) {
|
|
break;
|
|
}
|
|
const auto offset = calc(idx) / kOutputVecSize;
|
|
values[i] = LoadVector<ScalarT, kOutputVecSize>(data_, offset);
|
|
idx += stride;
|
|
}
|
|
idx = idx_;
|
|
#pragma unroll
|
|
for (IndexType i = 0; i < kVecSize; i++) {
|
|
if (idx >= end) {
|
|
break;
|
|
}
|
|
#pragma unroll
|
|
for (IndexType j = 0; j < kOutputVecSize; j++) {
|
|
value_list[i][j] = reducer.compute(value_list[i][j], values[i].val[j]);
|
|
}
|
|
idx += stride;
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int i = 1; i < kVecSize; i++) {
|
|
#pragma unroll
|
|
for (IndexType j = 0; j < kOutputVecSize; j++) {
|
|
value_list[0][j] = reducer.reduce(value_list[0][j], value_list[i][j]);
|
|
}
|
|
}
|
|
return value_list[0];
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE std::array<MPType, kOutputVecSize> BlockXReduce(
|
|
std::array<MPType, kOutputVecSize> value, char* shared_memory) const {
|
|
using MPTypeVec = std::array<MPType, kOutputVecSize>;
|
|
int dim_x = blockDim.x;
|
|
MPTypeVec* shared = reinterpret_cast<MPTypeVec*>(shared_memory);
|
|
|
|
unsigned mask = 0u;
|
|
CREATE_SHFL_MASK(mask, true);
|
|
|
|
if (dim_x > WARP_SIZE) {
|
|
IndexType address_base = static_cast<IndexType>(threadIdx.x) +
|
|
static_cast<IndexType>(threadIdx.y) *
|
|
static_cast<IndexType>(blockDim.x);
|
|
|
|
shared[address_base] = value;
|
|
for (int offset = dim_x / 2; offset >= WARP_SIZE; offset >>= 1) {
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x < offset && threadIdx.x + offset < blockDim.x) {
|
|
MPTypeVec other = shared[address_base + offset];
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce(value[i], other[i]);
|
|
}
|
|
shared[address_base] = value;
|
|
}
|
|
}
|
|
dim_x = WARP_SIZE;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Intra-warp reduction. Use decreasing offset on NVIDIA to match the
|
|
// accumulation order used by PyTorch (>= 2.10, PR #164790) and Triton,
|
|
// so that broadcast-grad sum reductions stay bitwise-aligned with torch.
|
|
// Keep the increasing-offset order on HIP/ROCm (matches torch's ROCm path).
|
|
#if defined(PADDLE_WITH_HIP)
|
|
for (int offset = 1; offset < dim_x; offset <<= 1) {
|
|
#else
|
|
for (int offset = dim_x >> 1; offset > 0; offset >>= 1) {
|
|
#endif
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
MPType other = reducer.shfl_sync(mask, value[i], offset);
|
|
value[i] = reducer.reduce(value[i], other);
|
|
}
|
|
}
|
|
return value;
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE std::array<MPType, kOutputVecSize> BlockYReduce(
|
|
std::array<MPType, kOutputVecSize> value, char* shared_memory) const {
|
|
using MPTypeVec = std::array<MPType, kOutputVecSize>;
|
|
MPTypeVec* shared = reinterpret_cast<MPTypeVec*>(shared_memory);
|
|
shared[config.SharedMemoryOffset(0)] = value;
|
|
|
|
for (int offset = blockDim.y / 2; offset > 0; offset >>= 1) {
|
|
__syncthreads();
|
|
if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {
|
|
MPTypeVec other = shared[config.SharedMemoryOffset(offset)];
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce(value[i], other[i]);
|
|
}
|
|
shared[config.SharedMemoryOffset(0)] = value;
|
|
}
|
|
}
|
|
return value;
|
|
}
|
|
|
|
DEVICE bool MarkBlockFinished() const {
|
|
__shared__ bool is_last_block_done_shared;
|
|
|
|
__syncthreads();
|
|
if (threadIdx.x == 0 && threadIdx.y == 0) {
|
|
int prev_blocks_finished = atomicAdd(&semaphores[blockIdx.x], 1);
|
|
is_last_block_done_shared = (prev_blocks_finished == gridDim.y - 1);
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
return is_last_block_done_shared;
|
|
}
|
|
|
|
template <int kOutputVecSize, bool can_acc>
|
|
DEVICE std::array<MPType, kOutputVecSize> AccumulateInOutput(
|
|
std::array<OutScalarT*, kOutputVecSize> out,
|
|
std::array<MPType, kOutputVecSize> value) const {
|
|
if constexpr (can_acc) {
|
|
std::array<MPType, kOutputVecSize> ret;
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
ret[i] = reducer.reduce(*(out[i]), value[i]);
|
|
}
|
|
return ret;
|
|
} else {
|
|
return {MPType{}};
|
|
}
|
|
}
|
|
|
|
template <bool can_acc>
|
|
DEVICE OutScalarT GetAccumulatedOutput(OutScalarT* out, MPType value) const {
|
|
if constexpr (can_acc) {
|
|
return (OutScalarT)value;
|
|
} else {
|
|
return *out;
|
|
}
|
|
}
|
|
|
|
template <class T>
|
|
DEVICE void SetResults(const T x, const IndexType base_offset) const {
|
|
auto res = reinterpret_cast<OutScalarT*>(dst[0] + base_offset);
|
|
*res = x;
|
|
}
|
|
|
|
template <class T1, class T2>
|
|
DEVICE void SetResults(const thrust::pair<T1, T2> x,
|
|
const IndexType base_offset) const {
|
|
if (noutputs >= 1) {
|
|
auto res0 = reinterpret_cast<T1*>(dst[0] + base_offset);
|
|
*res0 = x.first;
|
|
}
|
|
if (noutputs >= 2) {
|
|
auto res1 =
|
|
reinterpret_cast<T2*>(dst[1] + base_offset / sizeof(T1) * sizeof(T2));
|
|
*res1 = x.second;
|
|
}
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE void SetResultsToOutput(
|
|
std::array<MPType, kOutputVecSize> value,
|
|
std::array<IndexType, kOutputVecSize> base_offset) const {
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
SetResults(reducer.post_process(value[i]), base_offset[i]);
|
|
}
|
|
}
|
|
|
|
template <int kOutputVecSize>
|
|
DEVICE std::array<MPType, kOutputVecSize> GlobalReduce(
|
|
std::array<MPType, kOutputVecSize> value,
|
|
std::array<MPType, kOutputVecSize>* acc,
|
|
char* shared_memory) const {
|
|
using MPTypeVec = std::array<MPType, kOutputVecSize>;
|
|
using OutPtrVec = std::array<OutScalarT*, kOutputVecSize>;
|
|
using OffsetVec = std::array<IndexType, kOutputVecSize>;
|
|
|
|
MPTypeVec* reduce_buffer = reinterpret_cast<MPTypeVec*>(cta_buf);
|
|
IndexType output_idx = config.GetOutIdx<kOutputVecSize>();
|
|
OffsetVec base_offsets;
|
|
OutPtrVec out;
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
base_offsets[i] = output_calc.get(output_idx + i)[0];
|
|
out[i] = reinterpret_cast<OutScalarT*>(dst[0] + base_offsets[i]);
|
|
}
|
|
|
|
bool should_store = config.ShouldStore(output_idx);
|
|
if (should_store) {
|
|
IndexType offset = config.StagingMemoryOffset(blockIdx.y);
|
|
reduce_buffer[offset] = value;
|
|
}
|
|
|
|
__threadfence();
|
|
|
|
__syncthreads();
|
|
|
|
bool is_last_block_done = MarkBlockFinished();
|
|
|
|
if (is_last_block_done) {
|
|
__threadfence();
|
|
|
|
for (auto& v : value) {
|
|
v = ident;
|
|
}
|
|
|
|
if (config.ShouldReduceBlockX()) {
|
|
IndexType input_offset = static_cast<IndexType>(threadIdx.x) +
|
|
static_cast<IndexType>(threadIdx.y) *
|
|
static_cast<IndexType>(blockDim.x);
|
|
IndexType step = static_cast<IndexType>(blockDim.x) *
|
|
static_cast<IndexType>(blockDim.y);
|
|
|
|
for (; input_offset < config.ctas_per_output; input_offset += step) {
|
|
IndexType idx = config.StagingMemoryOffset(input_offset);
|
|
MPTypeVec next = reduce_buffer[idx];
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce(value[i], next[i]);
|
|
}
|
|
}
|
|
} else {
|
|
IndexType input_offset = threadIdx.y;
|
|
IndexType step = blockDim.y;
|
|
|
|
for (; input_offset < config.ctas_per_output; input_offset += step) {
|
|
IndexType idx = config.StagingMemoryOffset(input_offset);
|
|
MPTypeVec next = reduce_buffer[idx];
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce(value[i], next[i]);
|
|
}
|
|
}
|
|
}
|
|
value = BlockYReduce<kOutputVecSize>(value, shared_memory);
|
|
if (config.ShouldReduceBlockX()) {
|
|
value = BlockXReduce<kOutputVecSize>(value, shared_memory);
|
|
}
|
|
if (should_store) {
|
|
if (acc == nullptr) {
|
|
if (accumulate) {
|
|
value =
|
|
AccumulateInOutput<kOutputVecSize, can_accumulate_in_output>(
|
|
out, value);
|
|
}
|
|
if (final_output) {
|
|
SetResultsToOutput<kOutputVecSize>(value, base_offsets);
|
|
} else {
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
*(out[i]) = GetAccumulatedOutput<can_accumulate_in_output>(
|
|
out[i], value[i]);
|
|
}
|
|
}
|
|
} else {
|
|
if (accumulate) {
|
|
#pragma unroll
|
|
for (int i = 0; i < kOutputVecSize; i++) {
|
|
value[i] = reducer.reduce((*acc)[i], value[i]);
|
|
}
|
|
}
|
|
if (final_output) {
|
|
SetResultsToOutput<kOutputVecSize>(value, base_offsets);
|
|
} else {
|
|
*acc = value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return value;
|
|
}
|
|
}; // NOLINT
|
|
|
|
class AccumulationBuffer {
|
|
public:
|
|
AccumulationBuffer() {}
|
|
|
|
AccumulationBuffer(const KPDevice& dev_ctx,
|
|
size_t acc_t_size,
|
|
size_t out_t_size,
|
|
char* out_ptr,
|
|
int64_t size) {
|
|
out_ptr_ = reinterpret_cast<char*>(out_ptr);
|
|
if (out_t_size >= acc_t_size) {
|
|
acc_ptr_ = reinterpret_cast<char*>(out_ptr);
|
|
numerator_ = 1;
|
|
denominator_ = 1;
|
|
} else {
|
|
phi::Allocator* allocator =
|
|
const_cast<phi::Allocator*>(&(dev_ctx.GetAllocator())); // NOLINT
|
|
buffer_ = allocator->Allocate(size);
|
|
acc_ptr_ = reinterpret_cast<char*>(buffer_->ptr());
|
|
numerator_ = acc_t_size;
|
|
denominator_ = out_t_size;
|
|
ReduceFraction(&numerator_, &denominator_);
|
|
}
|
|
}
|
|
|
|
char* GetAccSlice(char* out_ptr) {
|
|
if (acc_ptr_ == nullptr) {
|
|
return nullptr;
|
|
}
|
|
return acc_ptr_ + ((out_ptr - out_ptr_) * numerator_ / denominator_);
|
|
}
|
|
|
|
private:
|
|
char* acc_ptr_ = nullptr;
|
|
char* out_ptr_ = nullptr;
|
|
size_t numerator_;
|
|
size_t denominator_;
|
|
Allocator::AllocationPtr buffer_;
|
|
};
|
|
|
|
template <int max_threads, typename R>
|
|
static void LaunchReduceKernel(const KPDevice& dev_ctx,
|
|
const ReduceConfig& config,
|
|
const R& reduction) {
|
|
dim3 block = config.GetBlockDim();
|
|
dim3 grid = config.GetGridDim();
|
|
int shared_memory = config.SharedMemorySize();
|
|
|
|
auto stream = dev_ctx.stream();
|
|
|
|
switch (config.output_vec_size) {
|
|
case 4:
|
|
VecReduceKernel<max_threads / 4, 4, R>
|
|
<<<grid, block, shared_memory, stream>>>(reduction);
|
|
break;
|
|
case 2:
|
|
VecReduceKernel<max_threads / 2, 2, R>
|
|
<<<grid, block, shared_memory, stream>>>(reduction);
|
|
break;
|
|
default:
|
|
VecReduceKernel<max_threads / 1, 1, R>
|
|
<<<grid, block, shared_memory, stream>>>(reduction);
|
|
break;
|
|
}
|
|
}
|
|
|
|
template <typename Tx,
|
|
typename Ty,
|
|
int kVecSize = 4,
|
|
int kInputVecSize = kVecSize,
|
|
typename ReduceOp,
|
|
typename ident_t = double>
|
|
inline void GPUReduceScheduler(const KPDevice& dev_ctx,
|
|
const DenseTensorIterator& iter,
|
|
const ReduceOp& reducer,
|
|
ident_t ident = 0,
|
|
AccumulationBuffer* acc_buf_ptr = nullptr,
|
|
int64_t base_idx = 0) {
|
|
auto stream = dev_ctx.stream();
|
|
|
|
using traits = phi::funcs::FunctionTraits<decltype(&ReduceOp::reduce)>;
|
|
using MPType = typename traits::template arg<0>::type;
|
|
|
|
static constexpr bool is_inp_out_type_half_or_chalf =
|
|
(std::is_same_v<phi::float16, Tx> && std::is_same_v<phi::float16, Ty>) ||
|
|
(std::is_same_v<phi::dtype::complex<float16>, Tx> &&
|
|
std::is_same_v<phi::dtype::complex<float16>, Ty>);
|
|
static constexpr bool is_inp_out_type_bfloat16 =
|
|
(std::is_same_v<phi::bfloat16, Tx> && std::is_same_v<phi::bfloat16, Ty>);
|
|
static constexpr bool can_accumulate_in_output =
|
|
std::is_convertible_v<MPType, Ty> &&
|
|
!(is_inp_out_type_half_or_chalf || is_inp_out_type_bfloat16);
|
|
|
|
bool can_use_32bit_indexing = iter.can_use_32bit_indexing();
|
|
std::unique_ptr<AccumulationBuffer> owned_buf_ptr;
|
|
if (acc_buf_ptr == NULL) {
|
|
if (!can_accumulate_in_output && !can_use_32bit_indexing) {
|
|
int64_t output_memory_size = phi::SizeOf(iter.dtype(0));
|
|
for (int dim = 0; dim < iter.ndim(); dim++) {
|
|
output_memory_size = std::max(output_memory_size,
|
|
iter.shape()[dim] * iter.strides(0)[dim]);
|
|
}
|
|
output_memory_size /= phi::SizeOf(iter.dtype(0));
|
|
owned_buf_ptr.reset(
|
|
new AccumulationBuffer(dev_ctx,
|
|
sizeof(MPType),
|
|
sizeof(Ty),
|
|
reinterpret_cast<char*>(iter.data_ptr(0)),
|
|
output_memory_size * sizeof(MPType)));
|
|
} else {
|
|
owned_buf_ptr.reset(new AccumulationBuffer());
|
|
}
|
|
acc_buf_ptr = owned_buf_ptr.get();
|
|
}
|
|
|
|
// Split iter if index exceeds 32-bit range.
|
|
if (!can_use_32bit_indexing) {
|
|
for (auto& sub_iter : iter.with_32bit_indexing()) {
|
|
int64_t sub_iter_base_idx = sub_iter.view_offsets()[0];
|
|
GPUReduceScheduler<Tx, Ty, kVecSize, kInputVecSize, ReduceOp>(
|
|
dev_ctx, sub_iter, reducer, ident, acc_buf_ptr, sub_iter_base_idx);
|
|
}
|
|
return;
|
|
}
|
|
|
|
const char* in_data =
|
|
reinterpret_cast<const char*>(iter.data_ptr(iter.ntensors() - 1));
|
|
char* out_data = reinterpret_cast<char*>(iter.data_ptr(0));
|
|
const auto noutputs = iter.noutputs();
|
|
|
|
std::optional<char*> out_data_extra;
|
|
if (noutputs > 1) {
|
|
out_data_extra = reinterpret_cast<char*>(iter.data_ptr(1));
|
|
} else {
|
|
out_data_extra = std::nullopt;
|
|
}
|
|
|
|
char* acc_data = acc_buf_ptr->GetAccSlice(out_data);
|
|
|
|
ReduceConfig config =
|
|
SetReduceConfig<MPType, Tx, kVecSize, kInputVecSize>(iter);
|
|
|
|
Allocator::AllocationPtr buffer;
|
|
Allocator::AllocationPtr semaphores;
|
|
void* buffer_ptr;
|
|
void* semaphores_ptr;
|
|
|
|
if (config.ShouldReduceGlobal()) {
|
|
phi::Allocator* allocator =
|
|
const_cast<phi::Allocator*>(&(dev_ctx.GetAllocator())); // NOLINT
|
|
buffer = allocator->Allocate(config.GlobalMemorySize());
|
|
semaphores = allocator->Allocate(config.SemaphoreSize());
|
|
buffer_ptr = buffer->ptr();
|
|
semaphores_ptr = semaphores->ptr();
|
|
|
|
phi::backends::gpu::GpuMemsetAsync(
|
|
semaphores_ptr, 0, config.SemaphoreSize(), stream);
|
|
}
|
|
|
|
auto output_calc = MakeOutputOffsetCalculator<uint32_t>(iter);
|
|
auto input_calc = MakeInputOffsetCalculator<uint32_t>(iter);
|
|
auto should_accumulate = iter.should_accumulate();
|
|
auto is_final_output = iter.is_final_output();
|
|
|
|
auto reduce = ReduceExecutor<Tx, ReduceOp, Ty, kVecSize, kInputVecSize>(
|
|
reducer,
|
|
config,
|
|
ident,
|
|
input_calc,
|
|
output_calc,
|
|
in_data,
|
|
out_data,
|
|
out_data_extra,
|
|
acc_data,
|
|
buffer_ptr,
|
|
reinterpret_cast<int*>(semaphores_ptr),
|
|
base_idx,
|
|
should_accumulate,
|
|
is_final_output,
|
|
noutputs);
|
|
|
|
LaunchReduceKernel<MaxThreadsConfig<Tx>::MAX_NUM_THREADS>(
|
|
dev_ctx, config, reduce);
|
|
|
|
return;
|
|
}
|
|
|
|
namespace funcs {
|
|
template <typename Tx,
|
|
typename Ty,
|
|
template <typename, typename, typename>
|
|
class ReduceOp>
|
|
void ReduceGpuKernel(const KPDevice& dev_ctx,
|
|
const phi::DenseTensor& x,
|
|
phi::DenseTensor* y,
|
|
const std::vector<int>& origin_reduce_dims,
|
|
const float norm_p = 1.0f) {
|
|
if (x.numel() == 0) {
|
|
dev_ctx.Alloc<Ty>(y);
|
|
return;
|
|
}
|
|
|
|
dev_ctx.Alloc<Ty>(y);
|
|
|
|
int64_t ndim = x.dims().size();
|
|
auto positive_reduce_dims = ConvertToPositiveDims(origin_reduce_dims, ndim);
|
|
auto mask = MakeDimMask(positive_reduce_dims, ndim);
|
|
auto viewed_result = ReviewReduceResult(x, *(y), ndim, mask);
|
|
|
|
auto x_dim = vectorize<int64_t>(x.dims());
|
|
|
|
DenseTensorIteratorConfig dense_iter_config;
|
|
dense_iter_config.is_reduction(true);
|
|
dense_iter_config.add_output(viewed_result);
|
|
dense_iter_config.add_const_input(x);
|
|
DenseTensorIterator iter = dense_iter_config.build();
|
|
|
|
using MPType = typename MPTypeTrait<Ty>::Type;
|
|
|
|
// TODO(baoqiwen): When ReduceOp is WelfordOps, kVecSize is 2.
|
|
constexpr int kVecSize = 4;
|
|
// For 16-bit (fp16/bf16) sum/mean, vectorize input loads by 8 elements
|
|
// (DWORDX4) to match PyTorch's accumulation grouping (>= 2.10, PR #165055),
|
|
// keeping broadcast-grad sum reductions bitwise-aligned with torch. Other
|
|
// dtypes / reduce ops keep the default input vector width.
|
|
constexpr bool kIsSumOrMeanOp =
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>, kps::SumOps<Tx, MPType, Ty>> ||
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>, kps::MeanOps<Tx, MPType, Ty>>;
|
|
constexpr int kInputVecSize =
|
|
(kIsSumOrMeanOp && sizeof(Tx) == 2) ? 8 : kVecSize;
|
|
|
|
// Initialize reducer.
|
|
ReduceOp reducer = [&iter, &norm_p]() {
|
|
constexpr bool kIsMeanOp =
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>, kps::MeanOps<Tx, MPType, Ty>>;
|
|
|
|
constexpr bool kIsPNormOp =
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::GenericPNormOps<Tx, MPType, Ty>>;
|
|
|
|
if constexpr (kIsMeanOp) {
|
|
MPType factor = static_cast<MPType>(iter.num_output_elements()) /
|
|
static_cast<MPType>(iter.numel());
|
|
return ReduceOp<Tx, MPType, Ty>{factor};
|
|
} else if constexpr (kIsPNormOp) {
|
|
return ReduceOp<Tx, MPType, Ty>{norm_p};
|
|
} else {
|
|
return ReduceOp<Tx, MPType, Ty>{};
|
|
}
|
|
}();
|
|
|
|
// Initialize ident value.
|
|
Tx ident = []() {
|
|
if constexpr (std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::MaxOps<Tx, MPType, Ty>> ||
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::AbsMaxOps<Tx, MPType, Ty>>) {
|
|
return std::numeric_limits<Tx>::lowest();
|
|
}
|
|
|
|
if constexpr (std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::MinOps<Tx, MPType, Ty>> ||
|
|
std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::AbsMinOps<Tx, MPType, Ty>>) {
|
|
return std::numeric_limits<Tx>::max();
|
|
}
|
|
|
|
if constexpr (std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::LogicalAndOps<Tx, MPType, Ty>>) {
|
|
return Tx{1};
|
|
}
|
|
|
|
if constexpr (std::is_same_v<ReduceOp<Tx, MPType, Ty>,
|
|
kps::ProdOps<Tx, MPType, Ty>>) {
|
|
return Tx{1};
|
|
}
|
|
|
|
// SumOps, MeanOps, LogicalOrOps and others
|
|
return Tx{0};
|
|
}();
|
|
|
|
GPUReduceScheduler<Tx, Ty, kVecSize, kInputVecSize, ReduceOp<Tx, MPType, Ty>>(
|
|
dev_ctx, iter, reducer, ident);
|
|
|
|
return;
|
|
}
|
|
} // namespace funcs
|
|
} // namespace phi
|