edit_distance_op.cu 4.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   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 <algorithm>
#include "paddle/framework/op_registry.h"
#include "paddle/platform/cuda_helper.h"
#include "paddle/platform/gpu_info.h"

namespace paddle {
namespace operators {

using platform::PADDLE_CUDA_NUM_THREADS;

template <typename T>
__global__ void FillFirstRow(T* dist, const int N) {
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  if (idx < N + 1) {
    dist[idx] = idx;
  }
}

template <typename T>
__global__ void FillFirstColumn(T* dist, const int M, const int N) {
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  if (idx < M + 1) {
    dist[idx * (N + 1)] = idx;
  }
}

template <typename T>
Y
Yibing Liu 已提交
42
__global__ void Levenshtein(T* dist, const int* x1, const int* x2, const int M,
43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
                            const int N, const int start) {
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  int offset = N;
  int index = start + idx * offset;
  int row = index / (N + 1);
  int col = index % (N + 1);
  if (row > 0 && col > 0 && row < M + 1 && col < N + 1) {
    int cost = x1[row - 1] == x2[col - 1] ? 0 : 1;
    int dels = dist[(row - 1) * (N + 1) + col] + 1;
    int ins = dist[row * (N + 1) + col - 1] + 1;
    int subs = dist[(row - 1) * (N + 1) + (col - 1)] + cost;
    dist[index] = min(dels, min(ins, subs));
  }
}

Y
Yibing Liu 已提交
58 59 60 61 62 63 64 65 66
template <typename T>
__global__ void SetOutput(T* out, const T* dist, const int M, const int N,
                          bool normalized) {
  int idx = blockDim.x * blockIdx.x + threadIdx.x;
  if (idx == 0) {
    out[0] = normalized ? dist[M * (N + 1) + N] / N : dist[M * (N + 1) + N];
  }
}

67
template <typename Place, typename T>
68
class EditDistanceGPUKernel : public framework::OpKernel<T> {
69 70 71 72 73 74 75
 public:
  void Compute(const framework::ExecutionContext& ctx) const {
    auto* out_t = ctx.Output<framework::Tensor>("Out");

    auto* x1_t = ctx.Input<framework::Tensor>("X1");
    auto* x2_t = ctx.Input<framework::Tensor>("X2");

Y
Yibing Liu 已提交
76 77
    out_t->mutable_data<T>(ctx.GetPlace());
    auto out = out_t->data<T>();
78 79 80 81 82 83 84 85

    auto normalized = ctx.Attr<bool>("normalized");
    auto stream = reinterpret_cast<const platform::CUDADeviceContext&>(
                      ctx.device_context())
                      .stream();

    auto m = x1_t->numel();
    auto n = x2_t->numel();
Y
Yibing Liu 已提交
86 87 88 89 90 91 92 93
    T distance = 0.0;
    if (m == 0 || n == 0) {
      distance = std::max(m, n);
      if (normalized) {
        distance = distance / n;
      }
      memory::Copy(boost::get<Place>(ctx.GetPlace()), out, platform::CPUPlace(),
                   &distance, sizeof(T), stream);
94 95 96 97 98
    } else {
      framework::Tensor dist_t;
      dist_t.Resize({m + 1, n + 1});
      dist_t.mutable_data<T>(ctx.GetPlace());
      auto dist = dist_t.data<T>();
Y
Yibing Liu 已提交
99 100
      auto x1 = x1_t->data<int>();
      auto x2 = x2_t->data<int>();
101 102 103 104 105 106

      FillFirstColumn<T><<<1 + m / PADDLE_CUDA_NUM_THREADS,
                           PADDLE_CUDA_NUM_THREADS, 0, stream>>>(dist, m, n);

      FillFirstRow<T><<<1 + n / PADDLE_CUDA_NUM_THREADS,
                        PADDLE_CUDA_NUM_THREADS, 0, stream>>>(dist, n);
Y
Yibing Liu 已提交
107 108
      // Compute the elements of distance matrix in the anti-diagonal diretion
      for (int64_t slice = 2; slice < m + n + 1; ++slice) {
109 110
        int z_m = slice < m + 1 ? 0 : slice - m;
        int z_n = slice < n + 1 ? 0 : slice - n;
Y
Yibing Liu 已提交
111 112
        int size = slice - (z_m + z_n) + 1;  // number of elments in the same
                                             // anti-diagonal line to update
113 114
        // the start index at which computes from
        int start = slice < n + 1 ? slice : (z_n + 1) * (n + 1) - 1;
115 116 117 118
        Levenshtein<T><<<1 + (size - 1) / PADDLE_CUDA_NUM_THREADS,
                         PADDLE_CUDA_NUM_THREADS, 0, stream>>>(dist, x1, x2, m,
                                                               n, start);
      }
Y
Yibing Liu 已提交
119
      SetOutput<T><<<1, 1, 0, stream>>>(out, dist, m, n, normalized);
120 121 122 123 124 125 126 127 128
    }
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

129 130 131
REGISTER_OP_CUDA_KERNEL(
    edit_distance,
    ops::EditDistanceGPUKernel<paddle::platform::CUDAPlace, float>);