wehub-resource-sync
221 lines
7.5 KiB
C++
221 lines
7.5 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 <cstdint>
|
|
|
|
#include "math.h" // NOLINT
|
|
|
|
#include "paddle/phi/core/cuda_stream.h"
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
template <int MaxTensorNumPerLaunch, int MaxChunkNumPerLaunch>
|
|
struct TensorMetaList {
|
|
static constexpr int kTensorNum = MaxTensorNumPerLaunch;
|
|
static constexpr int kChunkNum = MaxChunkNumPerLaunch;
|
|
|
|
static_assert(kTensorNum > 0 && kTensorNum < 256,
|
|
"kTensorNum must be inside (0, 256).");
|
|
static_assert(kChunkNum > 0 && kChunkNum < 65536,
|
|
"kChunkNum must be inside (0, 65536).");
|
|
|
|
/**
|
|
* The tensor numel offset of each tensor.
|
|
* The offsets[0] would be always 0 in the first launch,
|
|
* and then offsets[0] >= 0 in the following other launches.
|
|
* The numel of the i-th tensor would be offsets[i + 1] - offsets[i].
|
|
*/
|
|
int offsets[kTensorNum + 1];
|
|
|
|
/**
|
|
* The tensor id of each chunk. The tensor_ids[0] is always 0.
|
|
* Note that tensor_ids would be always in the ascending order.
|
|
* The actual tensor id is start_tensor_id + tensor_ids[i].
|
|
*
|
|
* The reason why we assume that the actual tensor id is
|
|
* start_tensor_id + tensor_ids[i] is to make tensor_ids to be
|
|
* a uint8_t array instead of an int array, making sizeof(TensorMetaList)
|
|
* smaller, so that kChunkNum can be larger.
|
|
*/
|
|
uint8_t tensor_ids[kChunkNum];
|
|
|
|
/**
|
|
* The chunk id of the chunk inside each tensor. It would be
|
|
* something like chunk_ids = [0, 1, 2, 0, 0, 1, 2, 3], meaning
|
|
* that there are 3 tensors and each tensor contains 3, 1 and 4
|
|
* chunks. Note that chunk_ids[0] is always 0 and the actual
|
|
* chunk id of the first tensor is always start_chunk_id + chunk_ids[i].
|
|
*
|
|
* The reason why we assume that the actual chunk id of the first
|
|
* tensor is always start_chunk_id + chunk_ids[i] is to make
|
|
* chunk_ids to be a uint16_t array instead of an int array, making
|
|
* sizeof(TensorMetaList) smaller, so that kChunkNum can be larger.
|
|
*/
|
|
uint16_t chunk_ids[kChunkNum];
|
|
|
|
/**
|
|
* The tensor_ids offset.
|
|
*/
|
|
int start_tensor_id;
|
|
|
|
/**
|
|
* The chunk_ids offset.
|
|
*/
|
|
int start_chunk_id;
|
|
};
|
|
|
|
template <typename Functor,
|
|
int MaxTensorNumPerLaunch,
|
|
int MaxChunkNumPerLaunch,
|
|
typename... Args>
|
|
static __global__ void MultiTensorApplyCUDAKernel(
|
|
Functor functor,
|
|
TensorMetaList<MaxTensorNumPerLaunch, MaxChunkNumPerLaunch> meta,
|
|
int chunk_size,
|
|
Args... args) {
|
|
const int block_id = blockIdx.x;
|
|
const int tensor_id = meta.tensor_ids[block_id];
|
|
const int chunk_id = static_cast<int>(meta.chunk_ids[block_id]) +
|
|
(tensor_id == 0) * meta.start_chunk_id;
|
|
const int prev_offset = meta.offsets[tensor_id];
|
|
const int next_offset = meta.offsets[tensor_id + 1];
|
|
const int ptr_offset = prev_offset + chunk_id * chunk_size;
|
|
const int size = min(next_offset - ptr_offset, chunk_size);
|
|
|
|
functor(
|
|
tensor_id + meta.start_tensor_id, chunk_id, ptr_offset, size, args...);
|
|
}
|
|
|
|
template <int MaxTensorNumPerLaunch, int MaxChunkNumPerLaunch>
|
|
class MultiTensorLauncher {
|
|
public:
|
|
MultiTensorLauncher(
|
|
const TensorMetaList<MaxTensorNumPerLaunch, MaxChunkNumPerLaunch> &meta,
|
|
const int &chunk_id,
|
|
const int &chunk_size,
|
|
const int &block_dim,
|
|
const gpuStream_t &stream)
|
|
: meta_(meta),
|
|
chunk_id_(chunk_id),
|
|
chunk_size_(chunk_size),
|
|
block_dim_(block_dim),
|
|
stream_(stream) {}
|
|
|
|
template <typename Functor, typename... Args>
|
|
void Launch(Functor &&functor, Args &&...args) const {
|
|
MultiTensorApplyCUDAKernel<Functor,
|
|
MaxTensorNumPerLaunch,
|
|
MaxChunkNumPerLaunch>
|
|
<<<chunk_id_, block_dim_, 0, stream_>>>(
|
|
functor, meta_, chunk_size_, args...);
|
|
}
|
|
|
|
private:
|
|
const TensorMetaList<MaxTensorNumPerLaunch, MaxChunkNumPerLaunch> &meta_;
|
|
const int &chunk_id_;
|
|
const int &chunk_size_;
|
|
const int &block_dim_;
|
|
const gpuStream_t &stream_;
|
|
};
|
|
|
|
template <int MaxTensorNumPerLaunch,
|
|
int MaxChunkNumPerLaunch,
|
|
typename Callback>
|
|
static void MultiTensorApplyWithCallback(gpuStream_t stream,
|
|
const int *offsets,
|
|
int n,
|
|
int chunk_size,
|
|
int block_dim,
|
|
Callback &&callback) {
|
|
if (n == 0) return;
|
|
|
|
constexpr auto NumTensor = MaxTensorNumPerLaunch;
|
|
constexpr auto NumChunk = MaxChunkNumPerLaunch;
|
|
TensorMetaList<NumTensor, NumChunk> metas;
|
|
|
|
int tensor_id = 0;
|
|
int chunk_id = 0;
|
|
int numel_offset = 0;
|
|
metas.start_tensor_id = 0;
|
|
metas.start_chunk_id = 0;
|
|
int launch_num = 0;
|
|
|
|
MultiTensorLauncher<MaxTensorNumPerLaunch, MaxChunkNumPerLaunch> launcher(
|
|
metas, chunk_id, chunk_size, block_dim, stream);
|
|
|
|
for (int i = 0; i < n; ++i) {
|
|
auto length = offsets[i + 1] - offsets[i];
|
|
if (tensor_id == 0) {
|
|
metas.start_tensor_id = i;
|
|
metas.offsets[0] = numel_offset;
|
|
}
|
|
metas.offsets[tensor_id + 1] = metas.offsets[tensor_id] + length;
|
|
++tensor_id;
|
|
numel_offset += length;
|
|
|
|
auto chunk_num = (length + chunk_size - 1) / chunk_size;
|
|
int last_launch_chunk_id = 0;
|
|
for (int j = 0; j < chunk_num; ++j) {
|
|
metas.chunk_ids[chunk_id] = j - last_launch_chunk_id;
|
|
metas.tensor_ids[chunk_id] = tensor_id - 1;
|
|
++chunk_id;
|
|
|
|
bool tensor_full = (tensor_id == NumTensor && j + 1 == chunk_num);
|
|
bool block_full = (chunk_id == NumChunk);
|
|
bool last_chunk = (i + 1 == n && j + 1 == chunk_num);
|
|
|
|
if (tensor_full || block_full || last_chunk) {
|
|
callback(launcher, launch_num);
|
|
++launch_num;
|
|
chunk_id = 0;
|
|
if (j + 1 == chunk_num) { // chunk for the current tensor is full
|
|
metas.start_chunk_id = 0;
|
|
tensor_id = 0;
|
|
} else {
|
|
metas.offsets[0] = metas.offsets[tensor_id - 1];
|
|
metas.offsets[1] = metas.offsets[tensor_id];
|
|
metas.start_tensor_id = i;
|
|
metas.start_chunk_id = j + 1;
|
|
last_launch_chunk_id = j + 1;
|
|
tensor_id = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename Functor,
|
|
int MaxTensorNumPerLaunch,
|
|
int MaxChunkNumPerLaunch,
|
|
typename... Args>
|
|
static void MultiTensorApply(Functor functor,
|
|
gpuStream_t stream,
|
|
const int *offsets,
|
|
int n,
|
|
int chunk_size,
|
|
int block_dim,
|
|
Args &&...args) {
|
|
auto callback = [&](const MultiTensorLauncher<MaxTensorNumPerLaunch,
|
|
MaxChunkNumPerLaunch> &launcher,
|
|
int i) { launcher.Launch(functor, args...); };
|
|
MultiTensorApplyWithCallback<MaxTensorNumPerLaunch, MaxChunkNumPerLaunch>(
|
|
stream, offsets, n, chunk_size, block_dim, callback);
|
|
}
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|