wehub-resource-sync
333 lines
10 KiB
Plaintext
333 lines
10 KiB
Plaintext
/**
|
|
* Copyright 2017-present, Facebook, Inc.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
#include "edit_dist.h"
|
|
#include <THC/THC.h>
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <device_launch_parameters.h>
|
|
#include <utility> // std::pair
|
|
|
|
template <typename scalar_t>
|
|
__global__ void generate_deletion_label_kernel(
|
|
const scalar_t* __restrict__ source,
|
|
const size_t source_size,
|
|
const size_t operation_size,
|
|
int* __restrict__ operations,
|
|
int* __restrict__ labels) {
|
|
|
|
const int index = blockIdx.x;
|
|
const int offset = index * operation_size;
|
|
const int offset_label = index * source_size;
|
|
|
|
for (int i = 0; i < source_size; i++) {
|
|
labels[offset_label + i] = 0;
|
|
}
|
|
|
|
int k = 0;
|
|
for (int i = 0; i < operation_size; i++){
|
|
if (operations[offset + i] == 0){
|
|
break;
|
|
} else if (operations[offset + i] == 1){
|
|
continue;
|
|
} else {
|
|
labels[offset_label + k] = 3 - operations[offset + i];
|
|
k++;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename scalar_t>
|
|
__global__ void generate_insertion_label_kernel(
|
|
const scalar_t* __restrict__ target,
|
|
const size_t target_size,
|
|
const size_t operation_size,
|
|
int* __restrict__ operations,
|
|
int* __restrict__ labels,
|
|
int* __restrict__ masks) {
|
|
|
|
const int index = blockIdx.x;
|
|
const int offset = index * operation_size;
|
|
const int offset_label = index * target_size;
|
|
|
|
int k = 0;
|
|
int u = 0;
|
|
int m = 0;
|
|
|
|
for (int i = 0; i < target_size; i++) {
|
|
labels[offset_label + i] = 0;
|
|
masks[offset_label + i] = 0;
|
|
}
|
|
|
|
for (int i = 0; i < operation_size-1; i++){
|
|
if (operations[offset + i] == 0){
|
|
break;
|
|
} else if (operations[offset + i] == 2){
|
|
continue;
|
|
} else if (operations[offset + i] == 1){
|
|
masks[offset_label + m] = 1;
|
|
u++; m++;
|
|
} else {
|
|
labels[offset_label + k] = u;
|
|
masks[offset_label + m] = 0;
|
|
k++; m++;
|
|
u = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename scalar_t>
|
|
__global__ void levenshtein_distance_kernel(
|
|
const scalar_t* __restrict__ source,
|
|
const scalar_t* __restrict__ target,
|
|
const int* __restrict__ source_length,
|
|
const int* __restrict__ target_length,
|
|
const size_t source_size,
|
|
const size_t target_size,
|
|
int* __restrict__ operations,
|
|
int* __restrict__ errors_curr) {
|
|
|
|
const int index = blockIdx.x;
|
|
const int offset = index * (source_size + target_size);
|
|
const int d = index * (source_size + 1) * (target_size + 1);
|
|
const int t = target_size + 1;
|
|
|
|
auto err_idx = [d, t](int i, int j) { return d + i * t + j; };
|
|
auto opt_idx = [offset](int k) { return offset + k; };
|
|
|
|
const int hyp_len = source_length[index];
|
|
const int ref_len = target_length[index];
|
|
const scalar_t* hyp_begin = source + index * source_size;
|
|
const scalar_t* ref_begin = target + index * target_size;
|
|
|
|
// dynamic programming
|
|
for (int i = 0; i <= hyp_len; i++){
|
|
errors_curr[err_idx(i, 0)] = i;
|
|
}
|
|
for (int j = 0; j <= ref_len; j++){
|
|
errors_curr[err_idx(0, j)] = j;
|
|
}
|
|
for (int i = 1; i <= hyp_len; i++){
|
|
for (int j = 1; j <= ref_len; j++){
|
|
errors_curr[err_idx(i, j)] = min(
|
|
min(
|
|
errors_curr[err_idx(i-1, j)],
|
|
errors_curr[err_idx(i, j-1)]
|
|
) + 1,
|
|
errors_curr[err_idx(i-1, j-1)] + 2 * (
|
|
*(hyp_begin+i-1) == *(ref_begin+j-1) ? 0 : 1
|
|
)
|
|
);
|
|
}
|
|
}
|
|
|
|
// back-tracing
|
|
int i = hyp_len;
|
|
int j = ref_len;
|
|
int o = hyp_len + ref_len;
|
|
|
|
for (int k = 0; k < source_size + target_size; k++) {
|
|
operations[opt_idx(k)] = 0;
|
|
}
|
|
|
|
while ((i >= 0) && (j >= 0)) {
|
|
if ((i == 0) && (j == 0)) {
|
|
break;
|
|
}
|
|
|
|
if ((j > 0) && (errors_curr[err_idx(i, j-1)] < errors_curr[err_idx(i, j)])) {
|
|
o--; operations[opt_idx(o)] = 1; j--; // insertion
|
|
} else if ((i > 0) && (errors_curr[err_idx(i-1, j)] < errors_curr[err_idx(i, j)])) {
|
|
o--; operations[opt_idx(o)] = 2; i--; // deletion
|
|
} else {
|
|
o--; operations[opt_idx(o)] = 3; i--; j--; // do nothing
|
|
}
|
|
}
|
|
|
|
// moving to the left
|
|
for (int k = 0; k < hyp_len + ref_len; k++) {
|
|
if (k + o < hyp_len + ref_len){
|
|
operations[opt_idx(k)] = operations[opt_idx(k+o)];
|
|
} else{
|
|
operations[opt_idx(k)] = 0; // padding
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
template <typename scalar_t>
|
|
__global__ void faster_levenshtein_distance_kernel(
|
|
const scalar_t* __restrict__ source,
|
|
const scalar_t* __restrict__ target,
|
|
const int* __restrict__ source_length,
|
|
const int* __restrict__ target_length,
|
|
const size_t source_size,
|
|
const size_t target_size,
|
|
int* __restrict__ operations) {
|
|
|
|
extern __shared__ short errors[];
|
|
auto errors_curr = errors;
|
|
|
|
const int index = blockIdx.x;
|
|
const int offset = index * (source_size + target_size);
|
|
const int t = target_size + 1;
|
|
|
|
auto err_idx = [t](int i, int j) { return i * t + j; };
|
|
auto opt_idx = [offset](int k) { return offset + k; };
|
|
|
|
const int hyp_len = source_length[index];
|
|
const int ref_len = target_length[index];
|
|
const scalar_t* hyp_begin = source + index * source_size;
|
|
const scalar_t* ref_begin = target + index * target_size;
|
|
|
|
// dynamic programming
|
|
for (int i = 0; i <= hyp_len; i++){
|
|
errors_curr[err_idx(i, 0)] = i;
|
|
}
|
|
for (int j = 0; j <= ref_len; j++){
|
|
errors_curr[err_idx(0, j)] = j;
|
|
}
|
|
for (int i = 1; i <= hyp_len; i++){
|
|
for (int j = 1; j <= ref_len; j++){
|
|
errors_curr[err_idx(i, j)] = min(
|
|
min(
|
|
errors_curr[err_idx(i-1, j)],
|
|
errors_curr[err_idx(i, j-1)]
|
|
) + 1,
|
|
errors_curr[err_idx(i-1, j-1)] + 2 * (
|
|
*(hyp_begin+i-1) == *(ref_begin+j-1) ? 0 : 1
|
|
)
|
|
);
|
|
}
|
|
}
|
|
|
|
// back-tracing
|
|
int i = hyp_len;
|
|
int j = ref_len;
|
|
int o = hyp_len + ref_len;
|
|
|
|
for (int k = 0; k < source_size + target_size; k++) {
|
|
operations[opt_idx(k)] = 0;
|
|
}
|
|
|
|
while ((i >= 0) && (j >= 0)) {
|
|
if ((i == 0) && (j == 0)) {
|
|
break;
|
|
}
|
|
|
|
if ((j > 0) && (errors_curr[err_idx(i, j-1)] < errors_curr[err_idx(i, j)])) {
|
|
o--; operations[opt_idx(o)] = 1; j--; // insertion
|
|
} else if ((i > 0) && (errors_curr[err_idx(i-1, j)] < errors_curr[err_idx(i, j)])) {
|
|
o--; operations[opt_idx(o)] = 2; i--; // deletion
|
|
} else {
|
|
o--; operations[opt_idx(o)] = 3; i--; j--; // do nothing
|
|
}
|
|
}
|
|
|
|
// moving to the left
|
|
for (int k = 0; k < hyp_len + ref_len; k++) {
|
|
if (k + o < hyp_len + ref_len){
|
|
operations[opt_idx(k)] = operations[opt_idx(k+o)];
|
|
} else{
|
|
operations[opt_idx(k)] = 0; // padding
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
|
|
torch::Tensor GenerateDeletionLabelCuda(
|
|
torch::Tensor source,
|
|
torch::Tensor operations) {
|
|
|
|
const auto batch_size = source.size(0);
|
|
at::TensorOptions options(source.device());
|
|
options = options.dtype(at::ScalarType::Int);
|
|
auto labels = torch::empty({batch_size, source.size(1)}, options);
|
|
auto stream = at::cuda::getCurrentCUDAStream(source.device().index());
|
|
|
|
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "generate_deletion_labels", ([&] {
|
|
generate_deletion_label_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
|
source.data_ptr<scalar_t>(),
|
|
source.size(1),
|
|
operations.size(1),
|
|
operations.data_ptr<int>(),
|
|
labels.data_ptr<int>());
|
|
}));
|
|
|
|
return labels;
|
|
}
|
|
|
|
std::pair<torch::Tensor, torch::Tensor> GenerateInsertionLabelCuda(
|
|
torch::Tensor target,
|
|
torch::Tensor operations) {
|
|
|
|
const auto batch_size = target.size(0);
|
|
at::TensorOptions options(target.device());
|
|
options = options.dtype(at::ScalarType::Int);
|
|
auto labels = torch::empty({batch_size, target.size(1)}, options);
|
|
auto masks = torch::empty({batch_size, target.size(1)}, options);
|
|
auto stream = at::cuda::getCurrentCUDAStream(target.device().index());
|
|
|
|
AT_DISPATCH_ALL_TYPES(target.scalar_type(), "generate_insertion_labels", ([&] {
|
|
generate_insertion_label_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
|
target.data_ptr<scalar_t>(),
|
|
target.size(1),
|
|
operations.size(1),
|
|
operations.data_ptr<int>(),
|
|
labels.data_ptr<int>(),
|
|
masks.data_ptr<int>());
|
|
}));
|
|
|
|
return std::make_pair(labels, masks);
|
|
}
|
|
|
|
|
|
torch::Tensor LevenshteinDistanceCuda(
|
|
torch::Tensor source,
|
|
torch::Tensor target,
|
|
torch::Tensor source_length,
|
|
torch::Tensor target_length) {
|
|
|
|
const auto batch_size = source.size(0);
|
|
const auto shared_size = (source.size(1) + 1) * (target.size(1) + 1) * sizeof(short);
|
|
|
|
at::TensorOptions options(source.device());
|
|
options = options.dtype(at::ScalarType::Int);
|
|
auto operations = torch::empty({batch_size, source.size(1) + target.size(1)}, options);
|
|
auto stream = at::cuda::getCurrentCUDAStream(source.device().index());
|
|
|
|
if (shared_size > 40000) {
|
|
auto distances = torch::empty({batch_size, (source.size(1) + 1) * (target.size(1) + 1)}, options);
|
|
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "levenshtein_distance", ([&] {
|
|
levenshtein_distance_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
|
source.data_ptr<scalar_t>(),
|
|
target.data_ptr<scalar_t>(),
|
|
source_length.data_ptr<int>(),
|
|
target_length.data_ptr<int>(),
|
|
source.size(1),
|
|
target.size(1),
|
|
operations.data_ptr<int>(),
|
|
distances.data_ptr<int>());
|
|
}));
|
|
} else {
|
|
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "faster_levenshtein_distance", ([&] {
|
|
faster_levenshtein_distance_kernel<scalar_t><<<batch_size, 1, shared_size, stream>>>(
|
|
source.data_ptr<scalar_t>(),
|
|
target.data_ptr<scalar_t>(),
|
|
source_length.data_ptr<int>(),
|
|
target_length.data_ptr<int>(),
|
|
source.size(1),
|
|
target.size(1),
|
|
operations.data_ptr<int>());
|
|
}));
|
|
}
|
|
|
|
return operations;
|
|
}
|