faisalhr1997/codeformer

1// from https://github.com/rosinality/stylegan2-pytorch/blob/master/op/fused_bias_act_kernel.cu2// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.3//4// This work is made available under the Nvidia Source Code License-NC.5// To view a copy of this license, visit6// https://nvlabs.github.io/stylegan2/license.html7 8#include <torch/types.h>9 10#include <ATen/ATen.h>11#include <ATen/AccumulateType.h>12#include <ATen/cuda/CUDAContext.h>13#include <ATen/cuda/CUDAApplyUtils.cuh>14 15#include <cuda.h>16#include <cuda_runtime.h>17 18 19template <typename scalar_t>20static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,21    int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {22    int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;23 24    scalar_t zero = 0.0;25 26    for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {27        scalar_t x = p_x[xi];28 29        if (use_bias) {30            x += p_b[(xi / step_b) % size_b];31        }32 33        scalar_t ref = use_ref ? p_ref[xi] : zero;34 35        scalar_t y;36 37        switch (act * 10 + grad) {38            default:39            case 10: y = x; break;40            case 11: y = x; break;41            case 12: y = 0.0; break;42 43            case 30: y = (x > 0.0) ? x : x * alpha; break;44            case 31: y = (ref > 0.0) ? x : x * alpha; break;45            case 32: y = 0.0; break;46        }47 48        out[xi] = y * scale;49    }50}51 52 53torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,54    int act, int grad, float alpha, float scale) {55    int curDevice = -1;56    cudaGetDevice(&curDevice);57    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);58 59    auto x = input.contiguous();60    auto b = bias.contiguous();61    auto ref = refer.contiguous();62 63    int use_bias = b.numel() ? 1 : 0;64    int use_ref = ref.numel() ? 1 : 0;65 66    int size_x = x.numel();67    int size_b = b.numel();68    int step_b = 1;69 70    for (int i = 1 + 1; i < x.dim(); i++) {71        step_b *= x.size(i);72    }73 74    int loop_x = 4;75    int block_size = 4 * 32;76    int grid_size = (size_x - 1) / (loop_x * block_size) + 1;77 78    auto y = torch::empty_like(x);79 80    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {81        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(82            y.data_ptr<scalar_t>(),83            x.data_ptr<scalar_t>(),84            b.data_ptr<scalar_t>(),85            ref.data_ptr<scalar_t>(),86            act,87            grad,88            alpha,89            scale,90            loop_x,91            size_x,92            step_b,93            size_b,94            use_bias,95            use_ref96        );97    });98 99    return y;100}101