44 lines
1.1 KiB
C++
44 lines
1.1 KiB
C++
#include <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <torch/extension.h>
|
|
#include <torch/python.h>
|
|
#include <vector>
|
|
|
|
template <typename input_t>
|
|
void pixel_norm_cuda(
|
|
const void* input,
|
|
const void* scale,
|
|
const void* shift,
|
|
void* output,
|
|
const float eps,
|
|
const int batch_size,
|
|
const int num_channels,
|
|
const int num_pixels,
|
|
cudaStream_t stream
|
|
);
|
|
|
|
void pixel_norm_inplace(at::Tensor &x, at::Tensor &scale, at::Tensor &shift, float eps) {
|
|
int num_pixels = x.size(2)*x.size(3)*x.size(4);
|
|
int batch_size = x.size(0);
|
|
int num_channels = x.size(1);
|
|
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
|
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
|
|
|
pixel_norm_cuda<at::BFloat16>(
|
|
x.data_ptr(),
|
|
scale.data_ptr(),
|
|
shift.data_ptr(),
|
|
x.data_ptr(),
|
|
eps,
|
|
batch_size,
|
|
num_channels,
|
|
num_pixels,
|
|
stream
|
|
);
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("pixel_norm_inplace", &pixel_norm_inplace,
|
|
"pixel norm inplace");
|
|
}
|