|
| 1 | +/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved. |
| 2 | +
|
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +you may not use this file except in compliance with the License. |
| 5 | +You may obtain a copy of the License at |
| 6 | +
|
| 7 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +
|
| 9 | +Unless required by applicable law or agreed to in writing, software |
| 10 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +See the License for the specific language governing permissions and |
| 13 | +limitations under the License. */ |
| 14 | + |
| 15 | +#include <algorithm> |
| 16 | +#include "paddle/fluid/framework/op_registry.h" |
| 17 | +#include "paddle/fluid/operators/gather_tree_op.h" |
| 18 | + |
| 19 | +namespace paddle { |
| 20 | +namespace operators { |
| 21 | + |
| 22 | +#define CUDA_1D_KERNEL_LOOP(i, n) \ |
| 23 | + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \ |
| 24 | + i += blockDim.x * gridDim.x) |
| 25 | + |
| 26 | +template <typename T> |
| 27 | +__global__ void GatherTree(const T *ids_data, const T *parents_data, |
| 28 | + T *out_data, const int64_t max_length, |
| 29 | + const int64_t batch_size, const int64_t beam_size) { |
| 30 | + CUDA_1D_KERNEL_LOOP(i, batch_size * beam_size) { |
| 31 | + int batch = i / beam_size; |
| 32 | + int beam = i % beam_size; |
| 33 | + auto idx = |
| 34 | + (max_length - 1) * batch_size * beam_size + batch * beam_size + beam; |
| 35 | + out_data[idx] = ids_data[idx]; |
| 36 | + auto parent = parents_data[idx]; |
| 37 | + for (int step = max_length - 2; step >= 0; step--) { |
| 38 | + idx = step * batch_size * beam_size + batch * beam_size; |
| 39 | + out_data[idx + beam] = ids_data[idx + parent]; |
| 40 | + parent = parents_data[idx + parent]; |
| 41 | + } |
| 42 | + } |
| 43 | +} |
| 44 | + |
| 45 | +template <typename T> |
| 46 | +class GatherTreeOpCUDAKernel : public framework::OpKernel<T> { |
| 47 | + public: |
| 48 | + void Compute(const framework::ExecutionContext &ctx) const override { |
| 49 | + auto *ids = ctx.Input<Tensor>("Ids"); |
| 50 | + auto *parents = ctx.Input<Tensor>("Parents"); |
| 51 | + auto *out = ctx.Output<Tensor>("Out"); |
| 52 | + |
| 53 | + const auto *ids_data = ids->data<T>(); |
| 54 | + const auto *parents_data = parents->data<T>(); |
| 55 | + auto *out_data = out->mutable_data<T>(ctx.GetPlace()); |
| 56 | + |
| 57 | + auto &ids_dims = ids->dims(); |
| 58 | + int64_t max_length = ids_dims[0]; |
| 59 | + int64_t batch_size = ids_dims[1]; |
| 60 | + int64_t beam_size = ids_dims[2]; |
| 61 | + |
| 62 | + auto &dev_ctx = ctx.cuda_device_context(); |
| 63 | + |
| 64 | + const int block = 512; |
| 65 | + int max_threads = |
| 66 | + std::min(static_cast<int64_t>(dev_ctx.GetMaxPhysicalThreadCount()), |
| 67 | + batch_size * beam_size); |
| 68 | + const int grid = std::max(max_threads / block, 1); |
| 69 | + GatherTree<<<grid, block>>>(ids_data, parents_data, out_data, max_length, |
| 70 | + batch_size, beam_size); |
| 71 | + } |
| 72 | +}; |
| 73 | + |
| 74 | +} // namespace operators |
| 75 | +} // namespace paddle |
| 76 | + |
| 77 | +namespace ops = paddle::operators; |
| 78 | + |
| 79 | +REGISTER_OP_CUDA_KERNEL(gather_tree, ops::GatherTreeOpCUDAKernel<int32_t>, |
| 80 | + ops::GatherTreeOpCUDAKernel<int64_t>); |
0 commit comments