288 lines
12 KiB
Plaintext
288 lines
12 KiB
Plaintext
/* Copyright (c) 2026 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 <vector>
|
|
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/common/layout.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
#define CUDA_KERNEL_LOOP_TYPE(i, n, index_type) \
|
|
int64_t index_ = ((int64_t)blockIdx.x) * blockDim.x + threadIdx.x; \
|
|
for (index_type i = index_; index_ < (n); \
|
|
index_ += blockDim.x * gridDim.x, i = index_)
|
|
|
|
#define CUDA_KERNEL_LOOP(i, n) CUDA_KERNEL_LOOP_TYPE(i, n, int)
|
|
|
|
constexpr int CUDA_NUM_THREADS = 1024;
|
|
|
|
inline int GET_BLOCKS(const int64_t N,
|
|
const int64_t max_threads_per_block = CUDA_NUM_THREADS) {
|
|
auto block_num = (N - 1) / max_threads_per_block + 1;
|
|
return static_cast<int>(block_num);
|
|
}
|
|
|
|
#if defined(__CUDACC__) || defined(__HIPCC__)
|
|
|
|
template <typename dt>
|
|
__global__ void Im2colKernel(const int64_t n,
|
|
const dt* data_im,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t kernel_height,
|
|
const int64_t kernel_width,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
dt* data_col) {
|
|
CUDA_KERNEL_LOOP_TYPE(index, n, int64_t) {
|
|
int64_t w_out = index % width_col;
|
|
|
|
int64_t idx = index / width_col;
|
|
|
|
int64_t h_out = idx % height_col;
|
|
int64_t channel_in = idx / height_col;
|
|
int64_t channel_out = channel_in * kernel_height * kernel_width;
|
|
int64_t h_in = h_out * stride_height - pad_height;
|
|
int64_t w_in = w_out * stride_width - pad_width;
|
|
|
|
dt* col = data_col + (channel_out * height_col + h_out) * width_col + w_out;
|
|
const dt* im = data_im + (channel_in * height + h_in) * width + w_in;
|
|
|
|
for (int64_t i = 0; i < kernel_height; ++i) {
|
|
for (int64_t j = 0; j < kernel_width; ++j) {
|
|
int64_t h = h_in + i * dilation_height;
|
|
int64_t w = w_in + j * dilation_width;
|
|
*col = (h >= 0 && w >= 0 && h < height && w < width)
|
|
? im[i * dilation_height * width + j * dilation_width]
|
|
: static_cast<dt>(0);
|
|
col += height_col * width_col;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename accT, typename dt>
|
|
__forceinline__ __device__ void Col2imKernelImp(const int64_t index,
|
|
const dt* data_col,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t kernel_h,
|
|
const int64_t kernel_w,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
dt* data_im) {
|
|
accT val = static_cast<accT>(0);
|
|
const int64_t w_im = index % width + pad_width;
|
|
const int64_t h_im = (index / width) % height + pad_height;
|
|
const int64_t c_im = index / (width * height);
|
|
int64_t kernel_extent_w = (kernel_w - 1) * dilation_width + 1;
|
|
int64_t kernel_extent_h = (kernel_h - 1) * dilation_height + 1;
|
|
// compute the start and end of the output
|
|
const int64_t w_col_start = (w_im < kernel_extent_w)
|
|
? 0
|
|
: (w_im - kernel_extent_w) / stride_width + 1;
|
|
const int64_t w_col_end = ::min(w_im / stride_width + 1, width_col);
|
|
const int64_t h_col_start =
|
|
(h_im < kernel_extent_h) ? 0
|
|
: (h_im - kernel_extent_h) / stride_height + 1;
|
|
const int64_t h_col_end = ::min(h_im / stride_height + 1, height_col);
|
|
|
|
// TODO(dev): use the stride/dilation LCM to reduce loops
|
|
for (int64_t h_col = h_col_start; h_col < h_col_end; h_col += 1) {
|
|
for (int64_t w_col = w_col_start; w_col < w_col_end; w_col += 1) {
|
|
int64_t h_k = (h_im - h_col * stride_height);
|
|
int64_t w_k = (w_im - w_col * stride_width);
|
|
if (h_k % dilation_height == 0 && w_k % dilation_width == 0) {
|
|
h_k /= dilation_height;
|
|
w_k /= dilation_width;
|
|
int64_t data_col_index =
|
|
(((c_im * kernel_h + h_k) * kernel_w + w_k) * height_col + h_col) *
|
|
width_col +
|
|
w_col;
|
|
val += data_col[data_col_index];
|
|
}
|
|
}
|
|
}
|
|
data_im[index] = static_cast<dt>(val);
|
|
}
|
|
|
|
template <typename dt, typename accT>
|
|
__global__ void Col2imKernel(const int64_t n,
|
|
const dt* data_col,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t kernel_h,
|
|
const int64_t kernel_w,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
dt* data_im) {
|
|
CUDA_KERNEL_LOOP(index, n) {
|
|
Col2imKernelImp<accT>(index,
|
|
data_col,
|
|
height,
|
|
width,
|
|
kernel_h,
|
|
kernel_w,
|
|
pad_height,
|
|
pad_width,
|
|
stride_height,
|
|
stride_width,
|
|
dilation_height,
|
|
dilation_width,
|
|
height_col,
|
|
width_col,
|
|
data_im);
|
|
}
|
|
}
|
|
|
|
template <typename dt, typename Context>
|
|
void im2col_slow(const Context& dev_ctx,
|
|
const dt* data_im,
|
|
const int64_t channels,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
const int64_t kernel_height,
|
|
const int64_t kernel_width,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
dt* data_col) {
|
|
auto stream = dev_ctx.stream();
|
|
int64_t num_kernels = channels * height_col * width_col;
|
|
Im2colKernel<<<GET_BLOCKS(num_kernels), 1024, 0, stream>>>(num_kernels,
|
|
data_im,
|
|
height,
|
|
width,
|
|
kernel_height,
|
|
kernel_width,
|
|
pad_height,
|
|
pad_width,
|
|
stride_height,
|
|
stride_width,
|
|
dilation_height,
|
|
dilation_width,
|
|
height_col,
|
|
width_col,
|
|
data_col);
|
|
}
|
|
|
|
template <typename dt, typename accT, typename Context>
|
|
void col2im_slow(const Context& dev_ctx,
|
|
const dt* data_col,
|
|
const int64_t channels,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
const int64_t patch_height,
|
|
const int64_t patch_width,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
dt* data_im) {
|
|
auto stream = dev_ctx.stream();
|
|
int64_t num_kernels = channels * height * width;
|
|
Col2imKernel<dt, accT>
|
|
<<<GET_BLOCKS(num_kernels, 512), 512, 0, stream>>>(num_kernels,
|
|
data_col,
|
|
height,
|
|
width,
|
|
patch_height,
|
|
patch_width,
|
|
pad_height,
|
|
pad_width,
|
|
stride_height,
|
|
stride_width,
|
|
dilation_height,
|
|
dilation_width,
|
|
height_col,
|
|
width_col,
|
|
data_im);
|
|
}
|
|
#endif // __CUDACC__
|
|
|
|
template <typename dt>
|
|
void im2col_slow(cudaStream_t stream,
|
|
const dt* data_im,
|
|
const int64_t channels,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
const int64_t kernel_height,
|
|
const int64_t kernel_width,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
dt* data_col);
|
|
|
|
template <typename dt, typename accT>
|
|
void col2im_slow(cudaStream_t stream,
|
|
const dt* data_col,
|
|
const int64_t channels,
|
|
const int64_t height,
|
|
const int64_t width,
|
|
const int64_t height_col,
|
|
const int64_t width_col,
|
|
const int64_t patch_height,
|
|
const int64_t patch_width,
|
|
const int64_t pad_height,
|
|
const int64_t pad_width,
|
|
const int64_t stride_height,
|
|
const int64_t stride_width,
|
|
const int64_t dilation_height,
|
|
const int64_t dilation_width,
|
|
dt* data_im);
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|