Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 3aa101f

Browse files
a2d8a4va2d8a4v
authored
Add GPU implementation for median filter (OpenNMT#1917)
The current median filter operator has a CPU version alone. Adding a GPU implementation to satisfy the support for GPU computing.
1 parent 1251f7c commit 3aa101f

7 files changed

Lines changed: 246 additions & 53 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,8 @@ set(SOURCES
170170
src/ops/mean.cc
171171
src/ops/mean_cpu.cc
172172
src/ops/median_filter.cc
173+
src/ops/median_filter_cpu.cc
174+
src/ops/median_filter_gpu.cu
173175
src/ops/min_max.cc
174176
src/ops/mul.cc
175177
src/ops/multinomial.cc

‎include/ctranslate2/ops/median_filter.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
11
#pragma once
2-
32
#include "op.h"
43

54
namespace ctranslate2 {
65
namespace ops {
76

87
class MedianFilter : public Op {
98
public:
10-
MedianFilter(const dim_t width);
9+
explicit MedianFilter(dim_t width);
1110
void operator()(const StorageView& input, StorageView& output) const;
1211

1312
private:
1413
const dim_t _width;
14+
template <Device D, typename T>
15+
void compute(const StorageView& input, const dim_t axis_size, StorageView& output) const;
1516
};
1617

1718
}

‎src/ops/median_filter.cc‎

Lines changed: 8 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,57 +1,25 @@
11
#include "ctranslate2/ops/median_filter.h"
22

3-
#include <algorithm>
4-
5-
#include "cpu/parallel.h"
3+
#include "dispatch.h"
64

75
namespace ctranslate2 {
86
namespace ops {
97

10-
MedianFilter::MedianFilter(const dim_t width)
8+
MedianFilter::MedianFilter(dim_t width)
119
: _width(width)
12-
{
13-
}
10+
{
11+
}
1412

1513
void MedianFilter::operator()(const StorageView& input, StorageView& output) const {
1614
PROFILE("MedianFilter");
1715

18-
if (input.device() != Device::CPU)
19-
throw std::invalid_argument("MedianFilter currently only supports CPU execution");
16+
const dim_t axis = input.rank() - 1;
17+
const dim_t axis_size = input.dim(axis);
2018

2119
output.resize_as(input);
2220

23-
const dim_t depth = input.dim(-1);
24-
const dim_t batch_size = input.size() / depth;
25-
const dim_t rank = _width / 2;
26-
27-
if (depth <= rank)
28-
return;
29-
30-
const auto* src = input.data<float>();
31-
auto* dst = output.data<float>();
32-
33-
cpu::parallel_for(0, batch_size, 1, [&](dim_t begin, dim_t end) {
34-
StorageView window_storage({_width}, DataType::FLOAT32);
35-
auto* window = window_storage.data<float>();
36-
37-
for (dim_t i = begin; i < end; ++i) {
38-
const dim_t offset = i * depth;
39-
const auto* in = src + offset;
40-
auto* out = dst + offset;
41-
42-
for (dim_t j = 0; j < depth; ++j) {
43-
for (dim_t k = -rank; k <= rank; ++k) {
44-
dim_t read = std::abs(j + k);
45-
if (read >= depth)
46-
read = depth - (read - depth) - 2;
47-
window[k + rank] = in[read];
48-
}
49-
50-
std::nth_element(window, window + rank, window + _width);
51-
out[j] = window[rank];
52-
}
53-
}
54-
});
21+
DEVICE_AND_FLOAT_DISPATCH("MedianFilter", input.device(), input.dtype(),
22+
(compute<D, T>(input, axis_size, output)));
5523
}
5624

5725
}

‎src/ops/median_filter_cpu.cc‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
#include "ctranslate2/ops/median_filter.h"
2+
3+
#include <iostream>
4+
5+
#include <algorithm>
6+
#include "cpu/parallel.h"
7+
#include "type_dispatch.h"
8+
9+
namespace ctranslate2 {
10+
namespace ops {
11+
12+
template <Device D, typename T>
13+
void MedianFilter::compute(const StorageView& input,
14+
const dim_t axis_size,
15+
StorageView& output) const {
16+
const auto* src = input.data<T>();
17+
auto* dst = output.data<T>();
18+
19+
20+
const dim_t depth = axis_size;
21+
const dim_t batch_size = input.size() / depth;
22+
const dim_t rank = _width / 2;
23+
24+
if (depth <= rank)
25+
return;
26+
27+
cpu::parallel_for(0, batch_size, 1, [&](dim_t begin, dim_t end) {
28+
StorageView window_storage({_width}, DataType::FLOAT32);
29+
auto* window = window_storage.data<float>();
30+
31+
for (dim_t i = begin; i < end; ++i) {
32+
const dim_t offset = i * depth;
33+
const auto* in = src + offset;
34+
auto* out = dst + offset;
35+
36+
for (dim_t j = 0; j < depth; ++j) {
37+
for (dim_t k = -rank; k <= rank; ++k) {
38+
dim_t read = std::abs(j + k);
39+
if (read >= depth)
40+
read = depth - (read - depth) - 2;
41+
window[k + rank] = in[read];
42+
}
43+
44+
std::nth_element(window, window + rank, window + _width);
45+
out[j] = window[rank];
46+
}
47+
}
48+
});
49+
}
50+
51+
#define DECLARE_IMPL(T) \
52+
template void \
53+
MedianFilter::compute<Device::CPU, T>(const StorageView& input, \
54+
const dim_t axis_size, \
55+
StorageView& output) const;
56+
57+
DECLARE_IMPL(float)
58+
59+
}
60+
}

‎src/ops/median_filter_gpu.cu‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
#include "ctranslate2/ops/median_filter.h"
2+
3+
#include <cuda_fp16.h>
4+
#ifdef CUDA_BF16_AVAILABLE
5+
#include <cuda_bf16.h>
6+
#endif
7+
8+
#include "type_dispatch.h"
9+
#include "cuda/helpers.h"
10+
#include <type_traits>
11+
12+
namespace ctranslate2 {
13+
namespace ops {
14+
15+
constexpr dim_t num_threads = 256;
16+
17+
// Conversion helpers
18+
__device__ __forceinline__ float to_float(float v) { return v; }
19+
__device__ __forceinline__ float to_float(const half v) { return __half2float(v); }
20+
#ifdef CUDA_BF16_AVAILABLE
21+
__device__ __forceinline__ float to_float(const __nv_bfloat16 v) { return __bfloat162float(v); }
22+
#endif
23+
24+
__device__ __forceinline__ float from_float(float v) { return v; }
25+
__device__ __forceinline__ half from_float_half(float v) { return __float2half(v); }
26+
#ifdef CUDA_BF16_AVAILABLE
27+
__device__ __forceinline__ __nv_bfloat16 from_float_bf16(float v) { return __float2bfloat16(v); }
28+
#endif
29+
30+
namespace {
31+
constexpr int kMaxWindow = 129; // supports window widths up to 129 (rank 64)
32+
}
33+
34+
template <typename DeviceT, int kMax>
35+
__global__ void sliding_median_lastdim_kernel(const DeviceT* input,
36+
DeviceT* output,
37+
int rows,
38+
int depth,
39+
int width) {
40+
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
41+
const int total = rows * depth;
42+
if (tid >= total) return;
43+
44+
int row = tid / depth;
45+
int col = tid % depth;
46+
const int rank = width / 2;
47+
48+
if (depth <= rank) {
49+
output[tid] = input[tid];
50+
return;
51+
}
52+
if (width > kMax) {
53+
output[tid] = input[tid];
54+
return;
55+
}
56+
57+
float window[kMax];
58+
59+
const int row_offset = row * depth;
60+
// Reflection gather.
61+
for (int k = -rank; k <= rank; ++k) {
62+
int read = col + k;
63+
if (read < 0) read = -read;
64+
if (read >= depth) read = 2 * depth - read - 2;
65+
window[k + rank] = to_float(input[row_offset + read]);
66+
}
67+
68+
// Insertion sort (width is small: <= kMax, typically < 129).
69+
for (int i = 1; i < width; ++i) {
70+
float key = window[i];
71+
int j = i - 1;
72+
while (j >= 0 && window[j] > key) {
73+
window[j + 1] = window[j];
74+
--j;
75+
}
76+
window[j + 1] = key;
77+
}
78+
float median = window[rank];
79+
80+
if constexpr (std::is_same<DeviceT, float>::value) {
81+
output[tid] = median;
82+
} else if constexpr (std::is_same<DeviceT, half>::value) {
83+
output[tid] = from_float_half(median);
84+
#ifdef CUDA_BF16_AVAILABLE
85+
} else if constexpr (std::is_same<DeviceT, __nv_bfloat16>::value) {
86+
output[tid] = from_float_bf16(median);
87+
#endif
88+
}
89+
}
90+
91+
template <Device D, typename T>
92+
void MedianFilter::compute(const StorageView& input,
93+
const dim_t axis_size,
94+
StorageView& output) const {
95+
output.resize_as(input);
96+
const int depth = static_cast<int>(axis_size);
97+
const int rows = static_cast<int>(input.size() / depth);
98+
const int width = static_cast<int>(_width);
99+
const int rank = width / 2;
100+
101+
// Host-side guards and fallbacks.
102+
if (width <= 1) {
103+
if (&output != &input)
104+
output.copy_from(input);
105+
return;
106+
}
107+
if ((width & 1) == 0)
108+
throw std::invalid_argument("MedianFilter width must be odd");
109+
if (width > kMaxWindow)
110+
throw std::invalid_argument("MedianFilter width exceeds supported GPU max (" + std::to_string(kMaxWindow) + ")");
111+
if (depth <= rank) {
112+
if (&output != &input)
113+
output.copy_from(input);
114+
return;
115+
}
116+
117+
// Grid configuration
118+
const int total = rows * depth;
119+
int blocks = (total + num_threads - 1) / num_threads;
120+
if (blocks > cuda::max_blocks) {
121+
blocks = cuda::max_blocks;
122+
}
123+
124+
using device_t = cuda::device_type<T>;
125+
const device_t* in_ptr = cuda::device_cast(input.data<T>());
126+
device_t* out_ptr = cuda::device_cast(output.data<T>());
127+
sliding_median_lastdim_kernel<device_t, kMaxWindow><<<blocks, num_threads, 0, cuda::get_cuda_stream()>>>(
128+
in_ptr,
129+
out_ptr,
130+
rows,
131+
depth,
132+
width);
133+
CUDA_CHECK(cudaGetLastError());
134+
CUDA_CHECK(cudaDeviceSynchronize());
135+
}
136+
137+
#define DECLARE_IMPL(T) \
138+
template void \
139+
MedianFilter::compute<Device::CUDA, T>(const StorageView& input, \
140+
const dim_t axis_size, \
141+
StorageView& output) const;
142+
143+
DECLARE_IMPL(float)
144+
DECLARE_IMPL(float16_t)
145+
DECLARE_IMPL(bfloat16_t)
146+
147+
}
148+
}

‎tests/benchmark_ops.cc‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,15 @@ void benchmark_conv1d(Device device) {
116116
BENCHMARK(conv_op(x, weight, bias, y), 100);
117117
}
118118

119+
void benchmark_median_filter(Device device) {
120+
const dim_t width = 5;
121+
std::vector<float> x_ = rand_vector(100 * 512);
122+
StorageView x({100, 512}, x_, device);
123+
StorageView y(device);
124+
const ops::MedianFilter median_filter_op(width);
125+
BENCHMARK(median_filter_op(x, y), 10000);
126+
}
127+
119128
int main(int argc, char* argv[]) {
120129
if (argc < 3) {
121130
std::cerr << "usage: " << argv[0] << " op device [dtype]" << std::endl;
@@ -153,6 +162,8 @@ int main(int argc, char* argv[]) {
153162
benchmark_dequantize(device);
154163
else if (op == "conv1d")
155164
benchmark_conv1d(device);
165+
else if (op == "median_filter")
166+
benchmark_median_filter(device);
156167

157168
return 0;
158169
}

‎tests/ops_test.cc‎

Lines changed: 14 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -118,29 +118,32 @@ TEST(OpTest, QuantizeINT16) {
118118
expect_storage_eq(reverse, input);
119119
}
120120

121-
TEST(OpTest, MedianFilter) {
121+
class OpDeviceTest : public ::testing::TestWithParam<Device> {
122+
};
123+
124+
class OpDeviceFPTest : public ::testing::TestWithParam<FloatType> {
125+
};
126+
127+
128+
TEST_P(OpDeviceTest, MedianFilter) {
129+
Device device = GetParam();
122130
StorageView x({2, 8}, std::vector<float>{
123131
0.2556743323802948, 0.8028775453567505, 0.3514494299888611, 0.3542254865169525,
124132
0.5881291031837463, 0.1458204835653305, 0.6845740675926208, 0.543143630027771,
125133
0.9039326310157776, 0.38000917434692383, 0.9094009399414062, 0.4063926637172699,
126-
0.7943458557128906, 0.289182186126709, 0.9932224750518799, 0.01137143187224865});
134+
0.7943458557128906, 0.289182186126709, 0.9932224750518799, 0.01137143187224865},
135+
device);
127136
StorageView expected({2, 8}, std::vector<float>{
128137
0.3514494299888611, 0.3542254865169525, 0.3542254865169525, 0.3542254865169525,
129138
0.3542254865169525, 0.543143630027771, 0.5881291031837463, 0.543143630027771,
130139
0.9039326310157776, 0.4063926637172699, 0.7943458557128906, 0.4063926637172699,
131-
0.7943458557128906, 0.4063926637172699, 0.7943458557128906, 0.289182186126709});
132-
StorageView y;
140+
0.7943458557128906, 0.4063926637172699, 0.7943458557128906, 0.289182186126709},
141+
device);
142+
StorageView y(device);
133143
ops::MedianFilter(5)(x, y);
134144
expect_storage_eq(y, expected);
135145
}
136146

137-
class OpDeviceTest : public ::testing::TestWithParam<Device> {
138-
};
139-
140-
class OpDeviceFPTest : public ::testing::TestWithParam<FloatType> {
141-
};
142-
143-
144147
TEST_P(OpDeviceTest, Add) {
145148
Device device = GetParam();
146149
StorageView a({4}, std::vector<float>{1, 2, 3, 4}, device);

0 commit comments

Comments
 (0)