未验证 提交 b97e6d13 编写于 作者: L Linjie Chen 提交者: GitHub

[phi] move viterbi_decode to phi (#40186)

* move viterbi to phi

* move infershape to phi

* update infershape

* fix

* resolve conflicts
上级 452c75b8
......@@ -9,8 +9,10 @@ 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. */
#include "paddle/fluid/operators/viterbi_decode_op.h"
#include "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/ternary.h"
namespace paddle {
namespace operators {
......@@ -19,47 +21,6 @@ class ViterbiDecodeOp : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
void InferShape(framework::InferShapeContext* ctx) const override {
OP_INOUT_CHECK(ctx->HasInput("Input"), "Input", "Input", "ViterbiDecode");
OP_INOUT_CHECK(ctx->HasInput("Transition"), "Input", "Transition",
"ViterbiDecode");
OP_INOUT_CHECK(ctx->HasInput("Length"), "Input", "Length", "ViterbiDecode");
OP_INOUT_CHECK(ctx->HasOutput("Scores"), "Output", "Scores",
"ViterbiDecode");
OP_INOUT_CHECK(ctx->HasOutput("Path"), "Output", "Path", "ViterbiDecode");
auto in_dims = ctx->GetInputDim("Input");
PADDLE_ENFORCE_EQ(in_dims.size(), 3,
platform::errors::InvalidArgument(
"The rank of Input in ViterbiDecode must be 3. But "
"received Input's rank is %d.",
in_dims.size()));
auto length_dims = ctx->GetInputDim("Length");
PADDLE_ENFORCE_EQ(length_dims.size(), 1,
platform::errors::InvalidArgument(
"The rank of Length in ViterbiDecode must be 1. But "
"received Length's rank is %d.",
length_dims.size()));
auto transition_dims = ctx->GetInputDim("Transition");
PADDLE_ENFORCE_EQ(
transition_dims.size(), 2,
platform::errors::InvalidArgument(
"The rank of Transition in ViterbiDecode must be 2. But "
"received Transition's rank is %d.",
transition_dims.size()));
if (ctx->IsRuntime()) {
PADDLE_ENFORCE_EQ(
in_dims[0], length_dims[0],
platform::errors::InvalidArgument(
"The batch size of Input and Length should be equal."));
PADDLE_ENFORCE_EQ(in_dims[2], transition_dims[0],
platform::errors::InvalidArgument(
"The number of tags of Input (%d) and Transition "
"(%d) should be equal.",
transition_dims[0], in_dims[2]));
}
ctx->SetOutputDim("Scores", length_dims);
}
protected:
framework::OpKernelType GetExpectedKernelType(
const framework::ExecutionContext& ctx) const override {
......@@ -102,8 +63,8 @@ class ViterbiDecodeOpMaker : public framework::OpProtoAndCheckerMaker {
namespace ops = paddle::operators;
namespace platform = paddle::platform;
DECLARE_INFER_SHAPE_FUNCTOR(viterbi_decode, ViterbiDecodeInferShapeFunctor,
PD_INFER_META(phi::ViterbiDecodeInferMeta));
REGISTER_OP_WITHOUT_GRADIENT(viterbi_decode, ops::ViterbiDecodeOp,
ops::ViterbiDecodeOpMaker);
REGISTER_OP_CPU_KERNEL(
viterbi_decode, ops::ViterbiDecodeKernel<platform::CPUDeviceContext, float>,
ops::ViterbiDecodeKernel<platform::CPUDeviceContext, double>);
ops::ViterbiDecodeOpMaker,
ViterbiDecodeInferShapeFunctor);
/* 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. */
#include "paddle/fluid/operators/elementwise/elementwise_functor.h"
#include "paddle/fluid/operators/elementwise/elementwise_op_broadcast.cu.h"
#include "paddle/fluid/operators/viterbi_decode_op.h"
#include "paddle/phi/kernels/funcs/gather.cu.h"
#ifdef __NVCC__
#include "cub/cub.cuh"
#endif
#ifdef __HIPCC__
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#endif
namespace paddle {
namespace operators {
#define FIXED_BLOCK_DIM_CASE_BASE(log2_block_dim, ...) \
case (1 << (log2_block_dim)): { \
constexpr auto kBlockDim = (1 << (log2_block_dim)); \
__VA_ARGS__; \
} break
#define FIXED_BLOCK_DIM_CASE(...) \
FIXED_BLOCK_DIM_CASE_BASE(10, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(9, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(8, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(7, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(6, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(5, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(4, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(3, ##__VA_ARGS__);
int64_t ComputeBlockSize(int64_t col) {
if (col > 512)
return 1024;
else if (col > 256)
return 512;
else if (col > 128)
return 256;
else if (col > 64)
return 128;
else if (col > 32)
return 64;
else if (col > 16)
return 32;
else if (col > 8)
return 16;
else
return 8;
}
template <template <typename T> typename BinaryFunctor, typename T>
struct BinaryOperation<platform::CUDADeviceContext, BinaryFunctor, T> {
void operator()(const platform::CUDADeviceContext& dev_ctx,
const framework::Tensor& lhs, const framework::Tensor& rhs,
framework::Tensor* output) {
std::vector<const framework::Tensor*> ins{&lhs, &rhs};
std::vector<framework::Tensor*> outs{output};
paddle::operators::LaunchElementwiseCudaKernel<ElementwiseType::kBinary, T,
T>(dev_ctx, ins, &outs, -1,
BinaryFunctor<T>());
}
};
template <template <typename InT, typename OutT> typename CompareFunctor,
typename T>
struct GetMask<platform::CUDADeviceContext, CompareFunctor, T> {
void operator()(const framework::ExecutionContext& ctx,
const framework::Tensor& lhs, const framework::Tensor& rhs,
framework::Tensor* mask) {
std::vector<const framework::Tensor*> ins = {&lhs, &rhs};
std::vector<framework::Tensor*> outs = {mask};
auto& dev_ctx = ctx.template device_context<platform::CUDADeviceContext>();
paddle::operators::LaunchSameDimsElementwiseCudaKernel<T>(
dev_ctx, ins, &outs, CompareFunctor<int64_t, T>());
}
};
template <typename T, typename IndType, size_t BlockDim>
__global__ void ArgmaxCUDAKernel(const int64_t height, // n * h
const int64_t width, // c
const int64_t post_size, // h
const T* in, IndType* out_idx, T* out) {
typedef cub::BlockReduce<cub::KeyValuePair<int, T>, BlockDim> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
cub::ArgMax reducer;
T init = (std::numeric_limits<T>::lowest)(); // for windows compile
for (int idx = blockIdx.x; idx < height; idx += gridDim.x) {
cub::KeyValuePair<int, T> kv_pair = {-1, init};
int h = idx / post_size;
int w = idx % post_size;
for (int k = threadIdx.x; k < width; k += blockDim.x) {
kv_pair =
reducer({k, in[h * width * post_size + k * post_size + w]}, kv_pair);
}
kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, reducer);
if (threadIdx.x == 0) {
// return max, argmax
if (out_idx != nullptr) out_idx[idx] = static_cast<IndType>(kv_pair.key);
if (out != nullptr) out[idx] = kv_pair.value;
}
__syncthreads();
}
}
__global__ void ARangeKernel(int64_t* data, int num, int64_t scale) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
for (int start = idx; idx < num; idx += gridDim.x) {
data[idx] = idx * scale;
}
}
template <>
struct ARange<platform::CUDADeviceContext> {
void operator()(const platform::CUDADeviceContext& dev_ctx, int64_t* data,
int num, int64_t scale) {
int64_t kBlockDim = ComputeBlockSize(num);
// kBlockDim > num at most of time, so we can set grid = 1
ARangeKernel<<<1, kBlockDim, 0, dev_ctx.stream()>>>(data, num, scale);
}
};
template <typename T, typename IndType>
struct Argmax<platform::CUDADeviceContext, T, IndType> {
void operator()(const framework::ExecutionContext& ctx,
const framework::Tensor& input, framework::Tensor* out_idx,
framework::Tensor* out, int axis) {
framework::DDim input_dims = input.dims();
int64_t numel = input.numel();
int64_t groups = numel / input_dims[axis];
int64_t pre = 1;
int64_t post = 1;
int64_t n = input_dims[axis];
for (int i = 0; i < axis; i++) {
pre *= input_dims[i];
}
for (int i = axis + 1; i < input_dims.size(); i++) {
post *= input_dims[i];
}
const auto& dev_ctx = ctx.cuda_device_context();
auto cu_stream = dev_ctx.stream();
int64_t max_grid_dimx = dev_ctx.GetCUDAMaxGridDimSize()[0];
int64_t height = pre * post;
int64_t width = n;
int64_t grid_size = height < max_grid_dimx ? height : max_grid_dimx;
const T* in_data = input.data<T>();
IndType* out_idx_data = out_idx->data<IndType>();
T* out_data = out->data<T>();
switch (ComputeBlockSize(width)) {
FIXED_BLOCK_DIM_CASE(
ArgmaxCUDAKernel<T, IndType,
kBlockDim><<<grid_size, kBlockDim, 0, cu_stream>>>(
height, width, post, in_data, out_idx_data, out_data));
}
}
};
template <typename T>
struct GetMaxValue<platform::CUDADeviceContext, T> {
void operator()(const platform::CUDADeviceContext& dev_ctx,
const framework::Tensor& input, T* max_value) {
framework::Tensor out_data;
out_data.Resize(phi::make_ddim({1}));
out_data.mutable_data<T>(platform::CUDAPlace());
switch (ComputeBlockSize(input.numel())) {
FIXED_BLOCK_DIM_CASE(
ArgmaxCUDAKernel<T, T,
kBlockDim><<<1, kBlockDim, 0, dev_ctx.stream()>>>(
1, input.numel(), 1, input.data<int64_t>(), nullptr,
out_data.data<int64_t>()));
}
framework::Tensor max_value_tensor;
framework::TensorCopy(out_data, platform::CPUPlace(), &max_value_tensor);
*max_value = max_value_tensor.data<T>()[0];
}
};
template <typename T, typename IndexT>
struct Gather<platform::CUDADeviceContext, T, IndexT> {
void operator()(const platform::CUDADeviceContext& ctx,
const framework::Tensor& src, const framework::Tensor& index,
framework::Tensor* output) {
phi::funcs::GPUGather<T, IndexT>(ctx, src, index, output);
}
};
} // namespace operators
} // namespace paddle
namespace ops = paddle::operators;
namespace platform = paddle::platform;
REGISTER_OP_CUDA_KERNEL(
viterbi_decode,
ops::ViterbiDecodeKernel<platform::CUDADeviceContext, float>,
ops::ViterbiDecodeKernel<platform::CUDADeviceContext, double>);
/* 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
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/elementwise/elementwise_functor.h"
#include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
#include "paddle/fluid/operators/math/concat_and_split.h"
#include "paddle/fluid/operators/transpose_op.h"
#include "paddle/fluid/operators/unique_op.h"
#include "paddle/phi/kernels/funcs/compare_functors.h"
#include "paddle/phi/kernels/funcs/gather.h"
#ifdef PADDLE_WITH_MKLML
#include <omp.h>
#endif
namespace paddle {
namespace operators {
template <typename DeviceContext, typename T, typename IndType>
struct Argmax {
void operator()(const framework::ExecutionContext& ctx,
const framework::Tensor& input, framework::Tensor* out_idx,
framework::Tensor* out, int axis) {
framework::DDim input_dims = input.dims();
int64_t pre = 1;
int64_t post = 1;
int64_t n = input_dims[axis];
for (int i = 0; i < axis; i++) {
pre *= input_dims[i];
}
for (int i = axis + 1; i < input_dims.size(); i++) {
post *= input_dims[i];
}
int64_t height = pre * post;
int64_t width = n;
const T* in_data = input.data<T>();
IndType* out_idx_data = out_idx->data<IndType>();
T* out_data = out->data<T>();
// Reduce
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int64_t i = 0; i < height; ++i) {
int64_t h = i / post;
int64_t w = i % post;
IndType max_idx = -1;
T max_value = (std::numeric_limits<T>::lowest)(); // for windows compile
for (int64_t j = 0; j < width; ++j) {
if (in_data[h * width * post + j * post + w] > max_value) {
max_value = in_data[h * width * post + j * post + w];
max_idx = j;
}
}
out_data[i] = max_value;
out_idx_data[i] = max_idx;
}
}
};
template <typename DeviceContext>
struct ARange {
void operator()(const DeviceContext& dev_ctx, int64_t* data, int end,
int64_t scale) {
for (int i = 0; i < end; ++i) {
data[i] = i * scale;
}
}
};
template <typename DeviceContext, typename T>
struct GetMaxValue {
void operator()(const DeviceContext& dev_ctx, const framework::Tensor& input,
T* max_value) {
auto input_ptr = input.data<T>();
auto num = input.numel();
*max_value = *std::max_element(input_ptr, input_ptr + num);
}
};
template <typename DeviceContext, typename T, typename IndexT = int>
struct Gather {
void operator()(const DeviceContext& ctx, const framework::Tensor& src,
const framework::Tensor& index, framework::Tensor* output) {
phi::funcs::CPUGather<T, IndexT>(ctx, src, index, output);
}
};
template <typename T, typename Functor, typename OutT = T>
void SameDimsBinaryOP(const framework::Tensor& lhs,
const framework::Tensor& rhs, framework::Tensor* out) {
const T* lhs_ptr = lhs.data<T>();
const T* rhs_ptr = rhs.data<T>();
OutT* out_ptr = out->data<OutT>();
Functor functor;
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int i = 0; i < out->numel(); ++i) {
out_ptr[i] = functor(lhs_ptr[i], rhs_ptr[i]);
}
}
template <typename DeviceContext,
template <typename InT, typename OutT> typename CompareFunctor,
typename T>
struct GetMask {
void operator()(const framework::ExecutionContext& ctx,
const framework::Tensor& lhs, const framework::Tensor& rhs,
framework::Tensor* mask) {
SameDimsBinaryOP<int64_t, CompareFunctor<int64_t, T>, T>(lhs, rhs, mask);
}
};
template <bool is_multi_threads>
struct GetInputIndex {
void operator()(const std::vector<int>& lhs_dims,
const std::vector<int>& rhs_dims,
const std::vector<int>& output_dims,
const std::vector<int>& lhs_strides,
const std::vector<int>& rhs_strides,
const std::vector<int>& output_strides, int output_idx,
int* index_array, int* lhs_idx, int* rhs_idx) {
int out_dims_size = output_strides.size();
for (int j = 0; j < out_dims_size; ++j) {
int curr_idx = output_idx / output_strides[j];
output_idx %= output_strides[j];
*lhs_idx += (lhs_dims[j] > 1) ? curr_idx * lhs_strides[j] : 0;
*rhs_idx += (rhs_dims[j] > 1) ? curr_idx * rhs_strides[j] : 0;
}
}
};
template <>
struct GetInputIndex<false> {
void operator()(const std::vector<int>& lhs_dims,
const std::vector<int>& rhs_dims,
const std::vector<int>& output_dims,
const std::vector<int>& lhs_strides,
const std::vector<int>& rhs_strides,
const std::vector<int>& output_strides, int output_idx,
int* index_array, int* lhs_idx, int* rhs_idx) {
int out_dims_size = output_strides.size();
*lhs_idx = phi::funcs::GetElementwiseIndex(lhs_dims.data(), out_dims_size,
index_array);
*rhs_idx = phi::funcs::GetElementwiseIndex(rhs_dims.data(), out_dims_size,
index_array);
phi::funcs::UpdateElementwiseIndexArray(output_dims.data(), out_dims_size,
index_array);
}
};
template <typename T, typename Functor, bool is_multi_threads = false>
void SimpleBroadcastBinaryOP(const framework::Tensor& lhs,
const framework::Tensor& rhs,
framework::Tensor* out) {
const T* lhs_ptr = lhs.data<T>();
const T* rhs_ptr = rhs.data<T>();
T* out_ptr = out->data<T>();
int out_size = static_cast<int>(out->dims().size());
std::vector<int> out_dims(out_size);
std::vector<int> lhs_dims(out_size);
std::vector<int> rhs_dims(out_size);
std::copy(lhs.dims().Get(), lhs.dims().Get() + out_size, lhs_dims.data());
std::copy(rhs.dims().Get(), rhs.dims().Get() + out_size, rhs_dims.data());
std::copy(out->dims().Get(), out->dims().Get() + out_size, out_dims.data());
std::vector<int> output_strides(out_size, 1);
std::vector<int> lhs_strides(out_size, 1);
std::vector<int> rhs_strides(out_size, 1);
std::vector<int> index_array(out_size, 0);
// calculate strides
for (int i = out_size - 2; i >= 0; --i) {
output_strides[i] = output_strides[i + 1] * out_dims[i + 1];
lhs_strides[i] = lhs_strides[i + 1] * lhs_dims[i + 1];
rhs_strides[i] = rhs_strides[i + 1] * rhs_dims[i + 1];
}
Functor functor;
GetInputIndex<is_multi_threads> get_input_index;
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int i = 0; i < out->numel(); ++i) {
int lhs_idx = 0;
int rhs_idx = 0;
get_input_index(lhs_dims, rhs_dims, out_dims, lhs_strides, rhs_strides,
output_strides, i, index_array.data(), &lhs_idx, &rhs_idx);
out_ptr[i] = functor(lhs_ptr[lhs_idx], rhs_ptr[rhs_idx]);
}
}
template <typename DeviceContext, template <typename T> typename BinaryFunctor,
typename T>
struct BinaryOperation {
void operator()(const DeviceContext& dev_ctx, const framework::Tensor& lhs,
const framework::Tensor& rhs, framework::Tensor* output) {
if (lhs.dims() == rhs.dims()) {
SameDimsBinaryOP<T, BinaryFunctor<T>>(lhs, rhs, output);
} else {
bool is_multi_threads = false;
#ifdef PADDLE_WITH_MKLML
if (omp_get_max_threads() > 1) {
is_multi_threads = true;
}
#endif
if (is_multi_threads) {
SimpleBroadcastBinaryOP<T, BinaryFunctor<T>, true>(lhs, rhs, output);
} else {
SimpleBroadcastBinaryOP<T, BinaryFunctor<T>, false>(lhs, rhs, output);
}
}
}
};
class TensorBuffer {
public:
explicit TensorBuffer(const framework::LoDTensor& in)
: buffer_(in), offset_(0) {
buffer_.Resize({buffer_.numel()});
}
framework::Tensor GetBufferBlock(std::initializer_list<int64_t> shape) {
int64_t size = std::accumulate(shape.begin(), shape.end(), 1,
std::multiplies<int64_t>());
framework::Tensor block = buffer_.Slice(offset_, offset_ + size);
offset_ += size;
block.Resize(shape);
return block;
}
private:
framework::LoDTensor buffer_; // need to resize 1-D Tensor
int offset_;
};
template <typename DeviceContext, typename T>
class ViterbiDecodeKernel : public framework::OpKernel<T> {
public:
void Compute(const framework::ExecutionContext& ctx) const override {
bool include_bos_eos_tag = ctx.Attr<bool>("include_bos_eos_tag");
auto& dev_ctx = ctx.template device_context<DeviceContext>();
auto curr_place = ctx.GetPlace();
auto* input = ctx.Input<framework::Tensor>("Input");
auto batch_size = static_cast<int>(input->dims()[0]);
auto seq_len = static_cast<int>(input->dims()[1]);
auto n_labels = static_cast<int>(input->dims()[2]);
phi::funcs::SetConstant<DeviceContext, T> float_functor;
phi::funcs::SetConstant<DeviceContext, int64_t> int_functor;
std::vector<framework::Tensor> historys;
// We create tensor buffer in order to avoid allocating memory frequently
// 10 means allocate 10*batch_size bytes memory, such as int_mask, zero...
int buffer_size = batch_size * (n_labels + 1) * seq_len + 10 * batch_size;
framework::LoDTensor int_buffer;
int_buffer.Resize(phi::make_ddim({buffer_size}));
int_buffer.mutable_data<int64_t>(ctx.GetPlace());
TensorBuffer int_tensor_buffer(int_buffer);
// create float tensor buffer
// 10 means allocate 10*batch_size*n_labels bytes, such as alpha, alpha_max
buffer_size = batch_size * (seq_len + 10) * n_labels +
(batch_size + 2) * n_labels * n_labels;
framework::LoDTensor float_buffer;
float_buffer.Resize(phi::make_ddim({buffer_size}));
float_buffer.mutable_data<T>(ctx.GetPlace());
TensorBuffer float_tensor_buffer(float_buffer);
auto* length = ctx.Input<framework::Tensor>("Length");
framework::Tensor left_length =
int_tensor_buffer.GetBufferBlock({batch_size, 1});
framework::TensorCopy(*length, curr_place, dev_ctx, &left_length);
int64_t max_seq_len = 0;
GetMaxValue<DeviceContext, int64_t> get_max_value;
get_max_value(dev_ctx, left_length, &max_seq_len);
auto* scores = ctx.Output<framework::Tensor>("Scores");
scores->mutable_data<T>(curr_place);
auto* path = ctx.Output<framework::Tensor>("Path");
path->Resize({batch_size, max_seq_len});
path->mutable_data<int64_t>(curr_place);
framework::Tensor tpath =
int_tensor_buffer.GetBufferBlock({max_seq_len, batch_size});
auto batch_path = Unbind(tpath);
for (auto it = batch_path.begin(); it != batch_path.end(); ++it) {
it->Resize({batch_size});
}
// create and init required tensor
framework::Tensor input_exp =
float_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
TransCompute<DeviceContext, T>(3, dev_ctx, *input, &input_exp, {1, 0, 2});
auto* transition = ctx.Input<framework::Tensor>("Transition");
framework::Tensor trans_exp =
float_tensor_buffer.GetBufferBlock({n_labels, n_labels});
framework::TensorCopy(*transition, curr_place, dev_ctx, &trans_exp);
trans_exp.Resize({1, n_labels, n_labels});
framework::Tensor alpha =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
framework::Tensor zero = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &zero, 0);
framework::Tensor one = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &one, 1);
framework::Tensor float_one =
float_tensor_buffer.GetBufferBlock({batch_size, 1});
float_functor(dev_ctx, &float_one, static_cast<T>(1.0));
framework::Tensor alpha_trn_sum =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels, n_labels});
framework::Tensor alpha_max =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
framework::Tensor alpha_argmax =
int_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
auto alpha_argmax_unbind = Unbind(alpha_argmax);
framework::Tensor alpha_nxt =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
framework::Tensor int_mask = int_tensor_buffer.GetBufferBlock({batch_size});
framework::Tensor zero_len_mask =
int_tensor_buffer.GetBufferBlock({batch_size});
framework::Tensor float_mask =
float_tensor_buffer.GetBufferBlock({batch_size, 1});
framework::Tensor stop_trans =
float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
framework::Tensor start_trans =
float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
framework::Tensor rest_trans =
float_tensor_buffer.GetBufferBlock({1, n_labels - 2, n_labels});
framework::Tensor last_ids = int_tensor_buffer.GetBufferBlock({batch_size});
framework::Tensor last_ids_tmp =
int_tensor_buffer.GetBufferBlock({batch_size});
framework::Tensor batch_offset =
int_tensor_buffer.GetBufferBlock({batch_size});
framework::Tensor gather_idx =
int_tensor_buffer.GetBufferBlock({batch_size});
std::vector<const framework::Tensor*> shape{&rest_trans, &stop_trans,
&start_trans};
std::vector<framework::Tensor*> outputs{&rest_trans, &stop_trans,
&start_trans};
math::SplitFunctor<DeviceContext, T> split_functor;
split_functor(dev_ctx, trans_exp, shape, 1, &outputs);
stop_trans.Resize({1, n_labels});
start_trans.Resize({1, n_labels});
auto logit0 = input_exp.Slice(0, 1);
logit0.Resize({batch_size, n_labels});
BinaryOperation<DeviceContext, AddFunctor, T> AddFloat;
BinaryOperation<DeviceContext, AddFunctor, int64_t> AddInt;
BinaryOperation<DeviceContext, MulFunctor, T> MulFloat;
BinaryOperation<DeviceContext, MulFunctor, int64_t> MulInt;
BinaryOperation<DeviceContext, SubFunctor, T> SubFloat;
BinaryOperation<DeviceContext, SubFunctor, int64_t> SubInt;
if (include_bos_eos_tag) {
AddFloat(dev_ctx, logit0, start_trans, &alpha);
GetMask<DeviceContext, phi::funcs::EqualFunctor, T>()(ctx, left_length,
one, &float_mask);
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
} else {
alpha = logit0;
}
SubInt(dev_ctx, left_length, one, &left_length);
Argmax<DeviceContext, T, int64_t> argmax;
for (int64_t i = 1; i < max_seq_len; ++i) {
framework::Tensor logit = input_exp.Slice(i, i + 1);
logit.Resize({batch_size, n_labels});
framework::Tensor& alpha_exp = alpha.Resize({batch_size, n_labels, 1});
AddFloat(dev_ctx, alpha_exp, trans_exp, &alpha_trn_sum);
auto alpha_argmax_temp = alpha_argmax_unbind[i - 1];
alpha_argmax_temp.Resize({batch_size, n_labels});
argmax(ctx, alpha_trn_sum, &alpha_argmax_temp, &alpha_max, 1);
historys.emplace_back(alpha_argmax_temp);
AddFloat(dev_ctx, alpha_max, logit, &alpha_nxt);
alpha.Resize({batch_size, n_labels});
// mask = paddle.cast((left_length > 0), dtype='float32')
// alpha = mask * alpha_nxt + (1 - mask) * alpha
GetMask<DeviceContext, phi::funcs::GreaterThanFunctor, T>()(
ctx, left_length, zero, &float_mask);
// alpha_nxt = mask * alpha_nxt
MulFloat(dev_ctx, alpha_nxt, float_mask, &alpha_nxt);
// inv_mask = 1 - mask
SubFloat(dev_ctx, float_one, float_mask, &float_mask);
// alpha = (1 - mask) * alpha
MulFloat(dev_ctx, alpha, float_mask, &alpha);
// alpha += alpha_nxt
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
if (include_bos_eos_tag) {
GetMask<DeviceContext, phi::funcs::EqualFunctor, T>()(ctx, left_length,
one, &float_mask);
// alpha += mask * trans_exp[:, self.stop_idx]
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
}
SubInt(dev_ctx, left_length, one, &left_length);
}
argmax(ctx, alpha, &last_ids, scores, 1);
left_length.Resize({batch_size});
GetMask<DeviceContext, phi::funcs::GreaterEqualFunctor, int64_t>()(
ctx, left_length, zero, &int_mask);
// last_ids_update = last_ids * tag_mask
int last_ids_index = 1;
int actual_len = (std::min)(seq_len, static_cast<int>(max_seq_len));
MulInt(dev_ctx, last_ids, int_mask,
&batch_path[actual_len - last_ids_index]);
// The algorithm below can refer to
// https://github.com/PaddlePaddle/PaddleNLP/blob/develop/paddlenlp/layers/crf.py#L438
ARange<DeviceContext> arange;
arange(dev_ctx, batch_offset.data<int64_t>(), batch_size, n_labels);
Gather<DeviceContext, int64_t, int64_t> gather;
for (auto hist = historys.rbegin(); hist != historys.rend(); ++hist) {
++last_ids_index;
AddInt(dev_ctx, left_length, one, &left_length);
AddInt(dev_ctx, batch_offset, last_ids, &gather_idx);
framework::Tensor& last_ids_update =
batch_path[actual_len - last_ids_index];
hist->Resize({batch_size * n_labels});
gather(dev_ctx, *hist, gather_idx, &last_ids_update);
GetMask<DeviceContext, phi::funcs::GreaterThanFunctor, int64_t>()(
ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids_update, int_mask, &last_ids_update);
GetMask<DeviceContext, phi::funcs::EqualFunctor, int64_t>()(
ctx, left_length, zero, &zero_len_mask);
MulInt(dev_ctx, last_ids, zero_len_mask, &last_ids_tmp);
SubInt(dev_ctx, one, zero_len_mask, &zero_len_mask);
MulInt(dev_ctx, last_ids_update, zero_len_mask, &last_ids_update);
AddInt(dev_ctx, last_ids_update, last_ids_tmp, &last_ids_update);
GetMask<DeviceContext, phi::funcs::LessThanFunctor, int64_t>()(
ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids, int_mask, &last_ids);
AddInt(dev_ctx, last_ids_update, last_ids, &last_ids);
}
TransCompute<DeviceContext, int64_t>(2, dev_ctx, tpath, path, {1, 0});
}
};
} // namespace operators
} // namespace paddle
......@@ -192,6 +192,53 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
out->set_dtype(x.dtype());
}
void ViterbiDecodeInferMeta(const MetaTensor& input,
const MetaTensor& transition,
const MetaTensor& length,
bool include_bos_eos_tag,
MetaTensor* scores,
MetaTensor* path,
MetaConfig config) {
auto in_dims = input.dims();
PADDLE_ENFORCE_EQ(in_dims.size(),
3,
phi::errors::InvalidArgument(
"The rank of Input in ViterbiDecode must be 3. But "
"received Input's rank is %d.",
in_dims.size()));
auto length_dims = length.dims();
PADDLE_ENFORCE_EQ(length_dims.size(),
1,
phi::errors::InvalidArgument(
"The rank of Length in ViterbiDecode must be 1. But "
"received Length's rank is %d.",
length_dims.size()));
auto transition_dims = transition.dims();
PADDLE_ENFORCE_EQ(
transition_dims.size(),
2,
phi::errors::InvalidArgument(
"The rank of Transition in ViterbiDecode must be 2. But "
"received Transition's rank is %d.",
transition_dims.size()));
if (config.is_runtime) {
PADDLE_ENFORCE_EQ(
in_dims[0],
length_dims[0],
phi::errors::InvalidArgument(
"The batch size of Input and Length should be equal."));
PADDLE_ENFORCE_EQ(in_dims[2],
transition_dims[0],
phi::errors::InvalidArgument(
"The number of tags of Input (%d) and Transition "
"(%d) should be equal.",
transition_dims[0],
in_dims[2]));
}
scores->set_dims(length_dims);
scores->set_dtype(length.dtype());
}
void LerpInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& weight,
......
......@@ -53,6 +53,14 @@ void ScatterNdAddInferMeta(const MetaTensor& x,
const MetaTensor& updates,
MetaTensor* out);
void ViterbiDecodeInferMeta(const MetaTensor& input,
const MetaTensor& transition,
const MetaTensor& length,
bool include_bos_eos_tag,
MetaTensor* scores,
MetaTensor* path,
MetaConfig config = MetaConfig());
void LerpInferMeta(const MetaTensor& x,
const MetaTensor& y,
const MetaTensor& weight,
......
// 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.
#include "paddle/phi/kernels/viterbi_decode_kernel.h"
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/copy_kernel.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/compare_functors.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/gather.h"
#include "paddle/phi/kernels/funcs/viterbi_decode_functor.h"
#include "paddle/phi/kernels/transpose_kernel.h"
namespace phi {
template <typename Context, typename T, typename IndType>
struct Argmax {
void operator()(const Context& dev_ctx,
const DenseTensor& input,
DenseTensor* out_idx,
DenseTensor* out,
int axis) {
phi::DDim input_dims = input.dims();
int64_t pre = 1;
int64_t post = 1;
int64_t n = input_dims[axis];
for (int i = 0; i < axis; i++) {
pre *= input_dims[i];
}
for (int i = axis + 1; i < input_dims.size(); i++) {
post *= input_dims[i];
}
int64_t height = pre * post;
int64_t width = n;
const T* in_data = input.data<T>();
IndType* out_idx_data = out_idx->data<IndType>();
T* out_data = out->data<T>();
// Reduce
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int64_t i = 0; i < height; ++i) {
int64_t h = i / post;
int64_t w = i % post;
IndType max_idx = -1;
T max_value = (std::numeric_limits<T>::lowest)(); // for windows compile
for (int64_t j = 0; j < width; ++j) {
if (in_data[h * width * post + j * post + w] > max_value) {
max_value = in_data[h * width * post + j * post + w];
max_idx = j;
}
}
out_data[i] = max_value;
out_idx_data[i] = max_idx;
}
}
};
template <typename Context>
struct ARange {
void operator()(const Context& dev_ctx,
int64_t* data,
int end,
int64_t scale) {
for (int i = 0; i < end; ++i) {
data[i] = i * scale;
}
}
};
template <typename Context, typename T>
struct GetMaxValue {
void operator()(const Context& dev_ctx,
const DenseTensor& input,
T* max_value) {
auto input_ptr = input.data<T>();
auto num = input.numel();
*max_value = *std::max_element(input_ptr, input_ptr + num);
}
};
template <typename Context, typename T, typename IndexT = int>
struct Gather {
void operator()(const Context& dev_ctx,
const DenseTensor& src,
const DenseTensor& index,
DenseTensor* output) {
phi::funcs::CPUGather<T, IndexT>(dev_ctx, src, index, output);
}
};
template <typename Context,
template <typename InT, typename OutT> typename CompareFunctor,
typename T>
struct GetMask {
void operator()(const Context& dev_ctx,
const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* mask) {
funcs::SameDimsBinaryOP<int64_t, CompareFunctor<int64_t, T>, T>(
lhs, rhs, mask);
}
};
template <typename Context,
template <typename T> typename BinaryFunctor,
typename T>
struct BinaryOperation {
void operator()(const Context& dev_ctx,
const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* output) {
if (lhs.dims() == rhs.dims()) {
funcs::SameDimsBinaryOP<T, BinaryFunctor<T>>(lhs, rhs, output);
} else {
bool is_multi_threads = false;
#ifdef PADDLE_WITH_MKLML
if (omp_get_max_threads() > 1) {
is_multi_threads = true;
}
#endif
if (is_multi_threads) {
funcs::SimpleBroadcastBinaryOP<T, BinaryFunctor<T>, true>(
lhs, rhs, output);
} else {
funcs::SimpleBroadcastBinaryOP<T, BinaryFunctor<T>, false>(
lhs, rhs, output);
}
}
}
};
template <typename T, typename Context>
void ViterbiDecodeKernel(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& transition,
const DenseTensor& length,
bool include_bos_eos_tag,
DenseTensor* scores,
DenseTensor* path) {
auto curr_place = dev_ctx.GetPlace();
auto batch_size = static_cast<int>(input.dims()[0]);
auto seq_len = static_cast<int>(input.dims()[1]);
auto n_labels = static_cast<int>(input.dims()[2]);
phi::funcs::SetConstant<Context, T> float_functor;
phi::funcs::SetConstant<Context, int64_t> int_functor;
std::vector<DenseTensor> historys;
// We create tensor buffer in order to avoid allocating memory frequently
// 10 means allocate 10*batch_size bytes memory, such as int_mask, zero...
int buffer_size = batch_size * (n_labels + 1) * seq_len + 10 * batch_size;
DenseTensor int_buffer = Empty<int64_t>(dev_ctx, {buffer_size});
funcs::TensorBuffer int_tensor_buffer(int_buffer);
// create float tensor buffer
// 10 means allocate 10*batch_size*n_labels bytes, such as alpha, alpha_max
buffer_size = batch_size * (seq_len + 10) * n_labels +
(batch_size + 2) * n_labels * n_labels;
DenseTensor float_buffer = Empty<T>(dev_ctx, {buffer_size});
funcs::TensorBuffer float_tensor_buffer(float_buffer);
DenseTensor left_length = int_tensor_buffer.GetBufferBlock({batch_size, 1});
phi::Copy(dev_ctx, length, curr_place, false, &left_length);
int64_t max_seq_len = 0;
GetMaxValue<Context, int64_t> get_max_value;
get_max_value(dev_ctx, left_length, &max_seq_len);
dev_ctx.template Alloc<T>(scores);
path->Resize({batch_size, max_seq_len});
dev_ctx.template Alloc<int64_t>(path);
DenseTensor tpath =
int_tensor_buffer.GetBufferBlock({max_seq_len, batch_size});
auto batch_path = funcs::Unbind(tpath);
for (auto it = batch_path.begin(); it != batch_path.end(); ++it) {
it->Resize({batch_size});
}
// create and init required tensor
DenseTensor input_exp =
float_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
TransposeKernel<T, Context>(dev_ctx, input, {1, 0, 2}, &input_exp);
DenseTensor trans_exp =
float_tensor_buffer.GetBufferBlock({n_labels, n_labels});
phi::Copy(dev_ctx, transition, curr_place, false, &trans_exp);
trans_exp.Resize({1, n_labels, n_labels});
DenseTensor alpha =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor zero = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &zero, 0);
DenseTensor one = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &one, 1);
DenseTensor float_one = float_tensor_buffer.GetBufferBlock({batch_size, 1});
float_functor(dev_ctx, &float_one, static_cast<T>(1.0));
DenseTensor alpha_trn_sum =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels, n_labels});
DenseTensor alpha_max =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor alpha_argmax =
int_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
auto alpha_argmax_unbind = funcs::Unbind(alpha_argmax);
DenseTensor alpha_nxt =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor int_mask = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor zero_len_mask = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor float_mask = float_tensor_buffer.GetBufferBlock({batch_size, 1});
DenseTensor stop_trans = float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
DenseTensor start_trans =
float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
DenseTensor rest_trans =
float_tensor_buffer.GetBufferBlock({1, n_labels - 2, n_labels});
DenseTensor last_ids = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor last_ids_tmp = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor batch_offset = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor gather_idx = int_tensor_buffer.GetBufferBlock({batch_size});
std::vector<const DenseTensor*> shape{&rest_trans, &stop_trans, &start_trans};
std::vector<DenseTensor*> outputs{&rest_trans, &stop_trans, &start_trans};
phi::funcs::SplitFunctor<Context, T> split_functor;
split_functor(dev_ctx, trans_exp, shape, 1, &outputs);
stop_trans.Resize({1, n_labels});
start_trans.Resize({1, n_labels});
auto logit0 = input_exp.Slice(0, 1);
logit0.Resize({batch_size, n_labels});
BinaryOperation<Context, phi::funcs::AddFunctor, T> AddFloat;
BinaryOperation<Context, phi::funcs::AddFunctor, int64_t> AddInt;
BinaryOperation<Context, phi::funcs::MultiplyFunctor, T> MulFloat;
BinaryOperation<Context, phi::funcs::MultiplyFunctor, int64_t> MulInt;
BinaryOperation<Context, phi::funcs::SubtractFunctor, T> SubFloat;
BinaryOperation<Context, phi::funcs::SubtractFunctor, int64_t> SubInt;
if (include_bos_eos_tag) {
AddFloat(dev_ctx, logit0, start_trans, &alpha);
GetMask<Context, phi::funcs::EqualFunctor, T>()(
dev_ctx, left_length, one, &float_mask);
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
} else {
alpha = logit0;
}
SubInt(dev_ctx, left_length, one, &left_length);
Argmax<Context, T, int64_t> argmax;
for (int64_t i = 1; i < max_seq_len; ++i) {
DenseTensor logit = input_exp.Slice(i, i + 1);
logit.Resize({batch_size, n_labels});
DenseTensor& alpha_exp = alpha.Resize({batch_size, n_labels, 1});
AddFloat(dev_ctx, alpha_exp, trans_exp, &alpha_trn_sum);
auto alpha_argmax_temp = alpha_argmax_unbind[i - 1];
alpha_argmax_temp.Resize({batch_size, n_labels});
argmax(dev_ctx, alpha_trn_sum, &alpha_argmax_temp, &alpha_max, 1);
historys.emplace_back(alpha_argmax_temp);
AddFloat(dev_ctx, alpha_max, logit, &alpha_nxt);
alpha.Resize({batch_size, n_labels});
GetMask<Context, phi::funcs::GreaterThanFunctor, T>()(
dev_ctx, left_length, zero, &float_mask);
MulFloat(dev_ctx, alpha_nxt, float_mask, &alpha_nxt);
SubFloat(dev_ctx, float_one, float_mask, &float_mask);
MulFloat(dev_ctx, alpha, float_mask, &alpha);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
if (include_bos_eos_tag) {
GetMask<Context, phi::funcs::EqualFunctor, T>()(
dev_ctx, left_length, one, &float_mask);
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
}
SubInt(dev_ctx, left_length, one, &left_length);
}
argmax(dev_ctx, alpha, &last_ids, scores, 1);
left_length.Resize({batch_size});
GetMask<Context, phi::funcs::GreaterEqualFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
// last_ids_update = last_ids * tag_mask
int last_ids_index = 1;
int actual_len = (std::min)(seq_len, static_cast<int>(max_seq_len));
MulInt(dev_ctx, last_ids, int_mask, &batch_path[actual_len - last_ids_index]);
// The algorithm below can refer to
// https://github.com/PaddlePaddle/PaddleNLP/blob/develop/paddlenlp/layers/crf.py#L438
ARange<Context> arange;
arange(dev_ctx, batch_offset.data<int64_t>(), batch_size, n_labels);
Gather<Context, int64_t, int64_t> gather;
for (auto hist = historys.rbegin(); hist != historys.rend(); ++hist) {
++last_ids_index;
AddInt(dev_ctx, left_length, one, &left_length);
AddInt(dev_ctx, batch_offset, last_ids, &gather_idx);
DenseTensor& last_ids_update = batch_path[actual_len - last_ids_index];
hist->Resize({batch_size * n_labels});
gather(dev_ctx, *hist, gather_idx, &last_ids_update);
GetMask<Context, phi::funcs::GreaterThanFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids_update, int_mask, &last_ids_update);
GetMask<Context, phi::funcs::EqualFunctor, int64_t>()(
dev_ctx, left_length, zero, &zero_len_mask);
MulInt(dev_ctx, last_ids, zero_len_mask, &last_ids_tmp);
SubInt(dev_ctx, one, zero_len_mask, &zero_len_mask);
MulInt(dev_ctx, last_ids_update, zero_len_mask, &last_ids_update);
AddInt(dev_ctx, last_ids_update, last_ids_tmp, &last_ids_update);
GetMask<Context, phi::funcs::LessThanFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids, int_mask, &last_ids);
AddInt(dev_ctx, last_ids_update, last_ids, &last_ids);
}
TransposeKernel<int64_t, Context>(dev_ctx, tpath, {1, 0}, path);
}
} // namespace phi
PD_REGISTER_KERNEL(
viterbi_decode, CPU, ALL_LAYOUT, phi::ViterbiDecodeKernel, float, double) {}
// 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
#ifdef PADDLE_WITH_MKLML
#include <omp.h>
#endif
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/kernels/funcs/math_function.h"
namespace phi {
namespace funcs {
static std::vector<DenseTensor> Unbind(const DenseTensor& in) {
int64_t size = in.dims()[0];
std::vector<DenseTensor> tensors(size);
for (int64_t i = 0; i < size; ++i) {
tensors[i] = in.Slice(i, i + 1);
}
return tensors;
}
template <typename T, typename Functor, typename OutT = T>
void SameDimsBinaryOP(const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* out) {
const T* lhs_ptr = lhs.data<T>();
const T* rhs_ptr = rhs.data<T>();
OutT* out_ptr = out->data<OutT>();
Functor functor;
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int i = 0; i < out->numel(); ++i) {
out_ptr[i] = functor(lhs_ptr[i], rhs_ptr[i]);
}
}
template <bool is_multi_threads>
struct GetInputIndex {
void operator()(const std::vector<int>& lhs_dims,
const std::vector<int>& rhs_dims,
const std::vector<int>& output_dims,
const std::vector<int>& lhs_strides,
const std::vector<int>& rhs_strides,
const std::vector<int>& output_strides,
int output_idx,
int* index_array,
int* lhs_idx,
int* rhs_idx) {
int out_dims_size = output_strides.size();
for (int j = 0; j < out_dims_size; ++j) {
int curr_idx = output_idx / output_strides[j];
output_idx %= output_strides[j];
*lhs_idx += (lhs_dims[j] > 1) ? curr_idx * lhs_strides[j] : 0;
*rhs_idx += (rhs_dims[j] > 1) ? curr_idx * rhs_strides[j] : 0;
}
}
};
template <typename T, typename Functor, bool is_multi_threads = false>
void SimpleBroadcastBinaryOP(const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* out) {
const T* lhs_ptr = lhs.data<T>();
const T* rhs_ptr = rhs.data<T>();
T* out_ptr = out->data<T>();
int out_size = static_cast<int>(out->dims().size());
std::vector<int> out_dims(out_size);
std::vector<int> lhs_dims(out_size);
std::vector<int> rhs_dims(out_size);
std::copy(lhs.dims().Get(), lhs.dims().Get() + out_size, lhs_dims.data());
std::copy(rhs.dims().Get(), rhs.dims().Get() + out_size, rhs_dims.data());
std::copy(out->dims().Get(), out->dims().Get() + out_size, out_dims.data());
std::vector<int> output_strides(out_size, 1);
std::vector<int> lhs_strides(out_size, 1);
std::vector<int> rhs_strides(out_size, 1);
std::vector<int> index_array(out_size, 0);
// calculate strides
for (int i = out_size - 2; i >= 0; --i) {
output_strides[i] = output_strides[i + 1] * out_dims[i + 1];
lhs_strides[i] = lhs_strides[i + 1] * lhs_dims[i + 1];
rhs_strides[i] = rhs_strides[i + 1] * rhs_dims[i + 1];
}
Functor functor;
GetInputIndex<is_multi_threads> get_input_index;
#ifdef PADDLE_WITH_MKLML
#pragma omp parallel for
#endif
for (int i = 0; i < out->numel(); ++i) {
int lhs_idx = 0;
int rhs_idx = 0;
get_input_index(lhs_dims,
rhs_dims,
out_dims,
lhs_strides,
rhs_strides,
output_strides,
i,
index_array.data(),
&lhs_idx,
&rhs_idx);
out_ptr[i] = functor(lhs_ptr[lhs_idx], rhs_ptr[rhs_idx]);
}
}
class TensorBuffer {
public:
explicit TensorBuffer(const DenseTensor& in) : buffer_(in), offset_(0) {
buffer_.Resize({buffer_.numel()});
}
DenseTensor GetBufferBlock(std::initializer_list<int64_t> shape) {
int64_t size = std::accumulate(
shape.begin(), shape.end(), 1, std::multiplies<int64_t>());
DenseTensor block = buffer_.Slice(offset_, offset_ + size);
offset_ += size;
block.Resize(shape);
return block;
}
private:
DenseTensor buffer_; // need to resize 1-D Tensor
int offset_;
};
} // namespace funcs
} // namespace phi
// 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.
#include "paddle/phi/kernels/viterbi_decode_kernel.h"
#ifdef __NVCC__
#include "cub/cub.cuh"
#endif
#ifdef __HIPCC__
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#endif
#ifdef PADDLE_WITH_MKLML
#include <omp.h>
#endif
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "paddle/fluid/operators/elementwise/elementwise_op_function.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/copy_kernel.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/compare_functors.h"
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/gather.cu.h"
#include "paddle/phi/kernels/funcs/viterbi_decode_functor.h"
#include "paddle/phi/kernels/transpose_kernel.h"
namespace phi {
#define FIXED_BLOCK_DIM_CASE_BASE(log2_block_dim, ...) \
case (1 << (log2_block_dim)): { \
constexpr auto kBlockDim = (1 << (log2_block_dim)); \
__VA_ARGS__; \
} break
#define FIXED_BLOCK_DIM_CASE(...) \
FIXED_BLOCK_DIM_CASE_BASE(10, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(9, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(8, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(7, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(6, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(5, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(4, ##__VA_ARGS__); \
FIXED_BLOCK_DIM_CASE_BASE(3, ##__VA_ARGS__);
int64_t ComputeBlockSize(int64_t col) {
if (col > 512)
return 1024;
else if (col > 256)
return 512;
else if (col > 128)
return 256;
else if (col > 64)
return 128;
else if (col > 32)
return 64;
else if (col > 16)
return 32;
else if (col > 8)
return 16;
else
return 8;
}
template <typename Context,
template <typename T> typename BinaryFunctor,
typename T>
struct BinaryOperation {
void operator()(const Context& dev_ctx,
const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* output) {
std::vector<const DenseTensor*> ins{&lhs, &rhs};
std::vector<DenseTensor*> outs{output};
paddle::operators::LaunchElementwiseCudaKernel<ElementwiseType::kBinary,
T,
T>(
dev_ctx, ins, &outs, -1, BinaryFunctor<T>());
}
};
template <typename Context,
template <typename InT, typename OutT> typename CompareFunctor,
typename T>
struct GetMask {
void operator()(const Context& dev_ctx,
const DenseTensor& lhs,
const DenseTensor& rhs,
DenseTensor* mask) {
std::vector<const DenseTensor*> ins = {&lhs, &rhs};
std::vector<DenseTensor*> outs = {mask};
paddle::operators::LaunchSameDimsElementwiseCudaKernel<T>(
dev_ctx, ins, &outs, CompareFunctor<int64_t, T>());
}
};
template <typename T, typename IndType, size_t BlockDim>
__global__ void ArgmaxCUDAKernel(const int64_t height, // n * h
const int64_t width, // c
const int64_t post_size, // h
const T* in,
IndType* out_idx,
T* out) {
typedef cub::BlockReduce<cub::KeyValuePair<int, T>, BlockDim> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
cub::ArgMax reducer;
T init = (std::numeric_limits<T>::lowest)(); // for windows compile
for (int idx = blockIdx.x; idx < height; idx += gridDim.x) {
cub::KeyValuePair<int, T> kv_pair = {-1, init};
int h = idx / post_size;
int w = idx % post_size;
for (int k = threadIdx.x; k < width; k += blockDim.x) {
kv_pair =
reducer({k, in[h * width * post_size + k * post_size + w]}, kv_pair);
}
kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, reducer);
if (threadIdx.x == 0) {
// return max, argmax
if (out_idx != nullptr) out_idx[idx] = static_cast<IndType>(kv_pair.key);
if (out != nullptr) out[idx] = kv_pair.value;
}
__syncthreads();
}
}
__global__ void ARangeKernel(int64_t* data, int num, int64_t scale) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
for (int start = idx; idx < num; idx += gridDim.x) {
data[idx] = idx * scale;
}
}
template <typename Context>
struct ARange {
void operator()(const Context& dev_ctx,
int64_t* data,
int num,
int64_t scale) {
int64_t kBlockDim = ComputeBlockSize(num);
// kBlockDim > num at most of time, so we can set grid = 1
ARangeKernel<<<1, kBlockDim, 0, dev_ctx.stream()>>>(data, num, scale);
}
};
template <typename Context, typename T, typename IndType>
struct Argmax {
void operator()(const Context& dev_ctx,
const DenseTensor& input,
DenseTensor* out_idx,
DenseTensor* out,
int axis) {
phi::DDim input_dims = input.dims();
int64_t numel = input.numel();
int64_t groups = numel / input_dims[axis];
int64_t pre = 1;
int64_t post = 1;
int64_t n = input_dims[axis];
for (int i = 0; i < axis; i++) {
pre *= input_dims[i];
}
for (int i = axis + 1; i < input_dims.size(); i++) {
post *= input_dims[i];
}
auto cu_stream = dev_ctx.stream();
int64_t max_grid_dimx = dev_ctx.GetCUDAMaxGridDimSize()[0];
int64_t height = pre * post;
int64_t width = n;
int64_t grid_size = height < max_grid_dimx ? height : max_grid_dimx;
const T* in_data = input.data<T>();
IndType* out_idx_data = out_idx->data<IndType>();
T* out_data = out->data<T>();
switch (ComputeBlockSize(width)) {
FIXED_BLOCK_DIM_CASE(
ArgmaxCUDAKernel<T,
IndType,
kBlockDim><<<grid_size, kBlockDim, 0, cu_stream>>>(
height, width, post, in_data, out_idx_data, out_data));
}
}
};
template <typename Context, typename T>
struct GetMaxValue {
void operator()(const Context& dev_ctx,
const DenseTensor& input,
T* max_value) {
DenseTensor out_data;
out_data.Resize(phi::make_ddim({1}));
dev_ctx.template Alloc<T>(&out_data);
switch (ComputeBlockSize(input.numel())) {
FIXED_BLOCK_DIM_CASE(
ArgmaxCUDAKernel<T,
T,
kBlockDim><<<1, kBlockDim, 0, dev_ctx.stream()>>>(
1,
input.numel(),
1,
input.data<int64_t>(),
nullptr,
out_data.data<int64_t>()));
}
DenseTensor max_value_tensor;
phi::Copy(dev_ctx, out_data, phi::CPUPlace(), false, &max_value_tensor);
*max_value = max_value_tensor.data<T>()[0];
}
};
template <typename Context, typename T, typename IndexT>
struct Gather {
void operator()(const Context& dev_ctx,
const DenseTensor& src,
const DenseTensor& index,
DenseTensor* output) {
phi::funcs::GPUGather<T, IndexT>(dev_ctx, src, index, output);
}
};
template <typename T, typename Context>
void ViterbiDecodeKernel(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& transition,
const DenseTensor& length,
bool include_bos_eos_tag,
DenseTensor* scores,
DenseTensor* path) {
auto curr_place = dev_ctx.GetPlace();
auto batch_size = static_cast<int>(input.dims()[0]);
auto seq_len = static_cast<int>(input.dims()[1]);
auto n_labels = static_cast<int>(input.dims()[2]);
phi::funcs::SetConstant<Context, T> float_functor;
phi::funcs::SetConstant<Context, int64_t> int_functor;
std::vector<DenseTensor> historys;
// We create tensor buffer in order to avoid allocating memory frequently
// 10 means allocate 10*batch_size bytes memory, such as int_mask, zero...
int buffer_size = batch_size * (n_labels + 1) * seq_len + 10 * batch_size;
DenseTensor int_buffer = Empty<int64_t>(dev_ctx, {buffer_size});
funcs::TensorBuffer int_tensor_buffer(int_buffer);
// create float tensor buffer
// 10 means allocate 10*batch_size*n_labels bytes, such as alpha, alpha_max
buffer_size = batch_size * (seq_len + 10) * n_labels +
(batch_size + 2) * n_labels * n_labels;
DenseTensor float_buffer = Empty<T>(dev_ctx, {buffer_size});
funcs::TensorBuffer float_tensor_buffer(float_buffer);
DenseTensor left_length = int_tensor_buffer.GetBufferBlock({batch_size, 1});
phi::Copy(dev_ctx, length, curr_place, false, &left_length);
int64_t max_seq_len = 0;
GetMaxValue<Context, int64_t> get_max_value;
get_max_value(dev_ctx, left_length, &max_seq_len);
dev_ctx.template Alloc<T>(scores);
path->Resize({batch_size, max_seq_len});
dev_ctx.template Alloc<int64_t>(path);
DenseTensor tpath =
int_tensor_buffer.GetBufferBlock({max_seq_len, batch_size});
auto batch_path = funcs::Unbind(tpath);
for (auto it = batch_path.begin(); it != batch_path.end(); ++it) {
it->Resize({batch_size});
}
// create and init required tensor
DenseTensor input_exp =
float_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
TransposeKernel<T, Context>(dev_ctx, input, {1, 0, 2}, &input_exp);
DenseTensor trans_exp =
float_tensor_buffer.GetBufferBlock({n_labels, n_labels});
phi::Copy(dev_ctx, transition, curr_place, false, &trans_exp);
trans_exp.Resize({1, n_labels, n_labels});
DenseTensor alpha =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor zero = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &zero, 0);
DenseTensor one = int_tensor_buffer.GetBufferBlock({batch_size, 1});
int_functor(dev_ctx, &one, 1);
DenseTensor float_one = float_tensor_buffer.GetBufferBlock({batch_size, 1});
float_functor(dev_ctx, &float_one, static_cast<T>(1.0));
DenseTensor alpha_trn_sum =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels, n_labels});
DenseTensor alpha_max =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor alpha_argmax =
int_tensor_buffer.GetBufferBlock({seq_len, batch_size, n_labels});
auto alpha_argmax_unbind = funcs::Unbind(alpha_argmax);
DenseTensor alpha_nxt =
float_tensor_buffer.GetBufferBlock({batch_size, n_labels});
DenseTensor int_mask = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor zero_len_mask = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor float_mask = float_tensor_buffer.GetBufferBlock({batch_size, 1});
DenseTensor stop_trans = float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
DenseTensor start_trans =
float_tensor_buffer.GetBufferBlock({1, 1, n_labels});
DenseTensor rest_trans =
float_tensor_buffer.GetBufferBlock({1, n_labels - 2, n_labels});
DenseTensor last_ids = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor last_ids_tmp = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor batch_offset = int_tensor_buffer.GetBufferBlock({batch_size});
DenseTensor gather_idx = int_tensor_buffer.GetBufferBlock({batch_size});
std::vector<const DenseTensor*> shape{&rest_trans, &stop_trans, &start_trans};
std::vector<DenseTensor*> outputs{&rest_trans, &stop_trans, &start_trans};
phi::funcs::SplitFunctor<Context, T> split_functor;
split_functor(dev_ctx, trans_exp, shape, 1, &outputs);
stop_trans.Resize({1, n_labels});
start_trans.Resize({1, n_labels});
auto logit0 = input_exp.Slice(0, 1);
logit0.Resize({batch_size, n_labels});
BinaryOperation<Context, phi::funcs::AddFunctor, T> AddFloat;
BinaryOperation<Context, phi::funcs::AddFunctor, int64_t> AddInt;
BinaryOperation<Context, phi::funcs::MultiplyFunctor, T> MulFloat;
BinaryOperation<Context, phi::funcs::MultiplyFunctor, int64_t> MulInt;
BinaryOperation<Context, phi::funcs::SubtractFunctor, T> SubFloat;
BinaryOperation<Context, phi::funcs::SubtractFunctor, int64_t> SubInt;
if (include_bos_eos_tag) {
AddFloat(dev_ctx, logit0, start_trans, &alpha);
GetMask<Context, phi::funcs::EqualFunctor, T>()(
dev_ctx, left_length, one, &float_mask);
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
} else {
alpha = logit0;
}
SubInt(dev_ctx, left_length, one, &left_length);
Argmax<Context, T, int64_t> argmax;
for (int64_t i = 1; i < max_seq_len; ++i) {
DenseTensor logit = input_exp.Slice(i, i + 1);
logit.Resize({batch_size, n_labels});
DenseTensor& alpha_exp = alpha.Resize({batch_size, n_labels, 1});
AddFloat(dev_ctx, alpha_exp, trans_exp, &alpha_trn_sum);
auto alpha_argmax_temp = alpha_argmax_unbind[i - 1];
alpha_argmax_temp.Resize({batch_size, n_labels});
argmax(dev_ctx, alpha_trn_sum, &alpha_argmax_temp, &alpha_max, 1);
historys.emplace_back(alpha_argmax_temp);
AddFloat(dev_ctx, alpha_max, logit, &alpha_nxt);
alpha.Resize({batch_size, n_labels});
GetMask<Context, phi::funcs::GreaterThanFunctor, T>()(
dev_ctx, left_length, zero, &float_mask);
MulFloat(dev_ctx, alpha_nxt, float_mask, &alpha_nxt);
SubFloat(dev_ctx, float_one, float_mask, &float_mask);
MulFloat(dev_ctx, alpha, float_mask, &alpha);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
if (include_bos_eos_tag) {
GetMask<Context, phi::funcs::EqualFunctor, T>()(
dev_ctx, left_length, one, &float_mask);
MulFloat(dev_ctx, stop_trans, float_mask, &alpha_nxt);
AddFloat(dev_ctx, alpha, alpha_nxt, &alpha);
}
SubInt(dev_ctx, left_length, one, &left_length);
}
argmax(dev_ctx, alpha, &last_ids, scores, 1);
left_length.Resize({batch_size});
GetMask<Context, phi::funcs::GreaterEqualFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
// last_ids_update = last_ids * tag_mask
int last_ids_index = 1;
int actual_len = (std::min)(seq_len, static_cast<int>(max_seq_len));
MulInt(dev_ctx, last_ids, int_mask, &batch_path[actual_len - last_ids_index]);
// The algorithm below can refer to
// https://github.com/PaddlePaddle/PaddleNLP/blob/develop/paddlenlp/layers/crf.py#L438
ARange<Context> arange;
arange(dev_ctx, batch_offset.data<int64_t>(), batch_size, n_labels);
Gather<Context, int64_t, int64_t> gather;
for (auto hist = historys.rbegin(); hist != historys.rend(); ++hist) {
++last_ids_index;
AddInt(dev_ctx, left_length, one, &left_length);
AddInt(dev_ctx, batch_offset, last_ids, &gather_idx);
DenseTensor& last_ids_update = batch_path[actual_len - last_ids_index];
hist->Resize({batch_size * n_labels});
gather(dev_ctx, *hist, gather_idx, &last_ids_update);
GetMask<Context, phi::funcs::GreaterThanFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids_update, int_mask, &last_ids_update);
GetMask<Context, phi::funcs::EqualFunctor, int64_t>()(
dev_ctx, left_length, zero, &zero_len_mask);
MulInt(dev_ctx, last_ids, zero_len_mask, &last_ids_tmp);
SubInt(dev_ctx, one, zero_len_mask, &zero_len_mask);
MulInt(dev_ctx, last_ids_update, zero_len_mask, &last_ids_update);
AddInt(dev_ctx, last_ids_update, last_ids_tmp, &last_ids_update);
GetMask<Context, phi::funcs::LessThanFunctor, int64_t>()(
dev_ctx, left_length, zero, &int_mask);
MulInt(dev_ctx, last_ids, int_mask, &last_ids);
AddInt(dev_ctx, last_ids_update, last_ids, &last_ids);
}
TransposeKernel<int64_t, Context>(dev_ctx, tpath, {1, 0}, path);
}
} // namespace phi
PD_REGISTER_KERNEL(
viterbi_decode, GPU, ALL_LAYOUT, phi::ViterbiDecodeKernel, float, double) {}
// 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 "paddle/phi/core/dense_tensor.h"
namespace phi {
template <typename T, typename Context>
void ViterbiDecodeKernel(const Context& dev_ctx,
const DenseTensor& input,
const DenseTensor& transition,
const DenseTensor& length,
bool include_bos_eos_tag,
DenseTensor* scores,
DenseTensor* path);
} // namespace phi
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册