breadlicker45/gpuGAN

1// Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.2//3// NVIDIA CORPORATION and its licensors retain all intellectual property4// and proprietary rights in and to this software, related documentation5// and any modifications thereto.  Any use, reproduction, disclosure or6// distribution of this software and related documentation without an express7// license agreement from NVIDIA CORPORATION is strictly prohibited.8 9#include <torch/extension.h>10#include <ATen/cuda/CUDAContext.h>11#include <c10/cuda/CUDAGuard.h>12#include "upfirdn2d.h"13 14//------------------------------------------------------------------------15 16static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx, int upy, int downx, int downy, int padx0, int padx1, int pady0, int pady1, bool flip, float gain)17{18    // Validate arguments.19    TORCH_CHECK(x.is_cuda(), "x must reside on CUDA device");20    TORCH_CHECK(f.device() == x.device(), "f must reside on the same device as x");21    TORCH_CHECK(f.dtype() == torch::kFloat, "f must be float32");22    TORCH_CHECK(x.numel() <= INT_MAX, "x is too large");23    TORCH_CHECK(f.numel() <= INT_MAX, "f is too large");24    TORCH_CHECK(x.dim() == 4, "x must be rank 4");25    TORCH_CHECK(f.dim() == 2, "f must be rank 2");26    TORCH_CHECK(f.size(0) >= 1 && f.size(1) >= 1, "f must be at least 1x1");27    TORCH_CHECK(upx >= 1 && upy >= 1, "upsampling factor must be at least 1");28    TORCH_CHECK(downx >= 1 && downy >= 1, "downsampling factor must be at least 1");29 30    // Create output tensor.31    const at::cuda::OptionalCUDAGuard device_guard(device_of(x));32    int outW = ((int)x.size(3) * upx + padx0 + padx1 - (int)f.size(1) + downx) / downx;33    int outH = ((int)x.size(2) * upy + pady0 + pady1 - (int)f.size(0) + downy) / downy;34    TORCH_CHECK(outW >= 1 && outH >= 1, "output must be at least 1x1");35    torch::Tensor y = torch::empty({x.size(0), x.size(1), outH, outW}, x.options(), x.suggest_memory_format());36    TORCH_CHECK(y.numel() <= INT_MAX, "output is too large");37 38    // Initialize CUDA kernel parameters.39    upfirdn2d_kernel_params p;40    p.x             = x.data_ptr();41    p.f             = f.data_ptr<float>();42    p.y             = y.data_ptr();43    p.up            = make_int2(upx, upy);44    p.down          = make_int2(downx, downy);45    p.pad0          = make_int2(padx0, pady0);46    p.flip          = (flip) ? 1 : 0;47    p.gain          = gain;48    p.inSize        = make_int4((int)x.size(3), (int)x.size(2), (int)x.size(1), (int)x.size(0));49    p.inStride      = make_int4((int)x.stride(3), (int)x.stride(2), (int)x.stride(1), (int)x.stride(0));50    p.filterSize    = make_int2((int)f.size(1), (int)f.size(0));51    p.filterStride  = make_int2((int)f.stride(1), (int)f.stride(0));52    p.outSize       = make_int4((int)y.size(3), (int)y.size(2), (int)y.size(1), (int)y.size(0));53    p.outStride     = make_int4((int)y.stride(3), (int)y.stride(2), (int)y.stride(1), (int)y.stride(0));54    p.sizeMajor     = (p.inStride.z == 1) ? p.inSize.w : p.inSize.w * p.inSize.z;55    p.sizeMinor     = (p.inStride.z == 1) ? p.inSize.z : 1;56 57    // Choose CUDA kernel.58    upfirdn2d_kernel_spec spec;59    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&]60    {61        spec = choose_upfirdn2d_kernel<scalar_t>(p);62    });63 64    // Set looping options.65    p.loopMajor     = (p.sizeMajor - 1) / 16384 + 1;66    p.loopMinor     = spec.loopMinor;67    p.loopX         = spec.loopX;68    p.launchMinor   = (p.sizeMinor - 1) / p.loopMinor + 1;69    p.launchMajor   = (p.sizeMajor - 1) / p.loopMajor + 1;70 71    // Compute grid size.72    dim3 blockSize, gridSize;73    if (spec.tileOutW < 0) // large74    {75        blockSize = dim3(4, 32, 1);76        gridSize = dim3(77            ((p.outSize.y - 1) / blockSize.x + 1) * p.launchMinor,78            (p.outSize.x - 1) / (blockSize.y * p.loopX) + 1,79            p.launchMajor);80    }81    else // small82    {83        blockSize = dim3(256, 1, 1);84        gridSize = dim3(85            ((p.outSize.y - 1) / spec.tileOutH + 1) * p.launchMinor,86            (p.outSize.x - 1) / (spec.tileOutW * p.loopX) + 1,87            p.launchMajor);88    }89 90    // Launch CUDA kernel.91    void* args[] = {&p};92    AT_CUDA_CHECK(cudaLaunchKernel(spec.kernel, gridSize, blockSize, args, 0, at::cuda::getCurrentCUDAStream()));93    return y;94}95 96//------------------------------------------------------------------------97 98PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)99{100    m.def("upfirdn2d", &upfirdn2d);101}102 103//------------------------------------------------------------------------104