| |
| |
|
|
| #include "grid_sample.h" |
| #include "gpu_grid_sample_utils.cuh" |
|
|
| template<typename T> |
| void indirect_grid_sample_forward_bilinear(torch::TensorAccessor<T, 4> input, |
| torch::TensorAccessor<T, 4> grid, |
| torch::TensorAccessor<int64_t, 1> inputIndices, |
| torch::TensorAccessor<T, 4> output) |
| { |
| const int64_t N = inputIndices.size(0); |
| const int64_t C = output.size(1); |
|
|
| T fInputHeight = input.size(2); |
| T fInputWidth = input.size(3); |
| int64_t outputHeight = output.size(2); |
| int64_t outputWidth = output.size(3); |
|
|
| #pragma omp parallel for num_threads(8) |
| for (int64_t i = 0; i < N; ++i) { |
| int64_t inputIdx = inputIndices[i]; |
|
|
| for (int64_t c = 0; c < C; ++c) { |
| for (int64_t outY = 0; outY < outputHeight; ++outY) { |
| for (int64_t outX = 0; outX < outputWidth; ++outX) { |
| T u = grid[i][outY][outX][0]; |
| T v = grid[i][outY][outX][1]; |
|
|
| if (u < -1 || u > 1 || v < -1 || v > 1) { |
| output[i][c][outY][outX] = 0; |
| continue; |
| } |
|
|
| |
| u = (u + 1) * ((fInputWidth - 1) / 2); |
| v = (v + 1) * ((fInputHeight - 1) / 2); |
|
|
| |
| const T inX = u; |
| const T inXint = std::floor(inX); |
| const T inXfrac = inX - inXint; |
|
|
| const T inY = v; |
| const T inYint = std::floor(inY); |
| const T inYfrac = inY - inYint; |
|
|
| T ps[] = { 1 - inXfrac, inXfrac }; |
| T rs[] = { 1 - inYfrac, inYfrac }; |
| T opVal = 0; |
|
|
| #pragma unroll |
| for (int64_t row = 0; row < 2; ++row) { |
| #pragma unroll |
| for (int64_t col = 0; col < 2; ++col) { |
| T Tpx = utils::get_pixel_clamped(input, inputIdx, c, inXint + col, inYint + row); |
| opVal += rs[row] * ps[col] * Tpx; |
| } |
| } |
|
|
| output[i][c][outY][outX] = opVal; |
| } |
| } |
| } |
| } |
| } |
|
|
| torch::Tensor cpu_indirect_grid_sample_forward(torch::Tensor input, torch::Tensor grid, |
| torch::Tensor inputIndices, const std::string &method) |
| { |
| auto output = input.new_empty({ inputIndices.size(0), input.size(1), grid.size(1), grid.size(2) }); |
|
|
| AT_DISPATCH_FLOATING_TYPES( |
| input.scalar_type(), |
| "cpu_indirect_grid_sample_forward_impl", |
| ([&] { |
| typedef scalar_t T; |
| if (method == "bilinear") { |
| indirect_grid_sample_forward_bilinear( |
| input.accessor<T, 4>(), |
| grid.accessor<T, 4>(), |
| inputIndices.accessor<int64_t, 1>(), |
| output.accessor<T, 4>() |
| ); |
| } else { |
| throw std::runtime_error("Unsupported resample method: " + method); |
| } |
| }) |
| ); |
|
|
| return output; |
| } |
|
|
