GPUCodeForces/S1/42/LogCoshLoss_cuda.py

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

N, C, H, W = 32, 64, 56, 56


class ModelNew(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction
        self.red_map = {'none': 0, 'mean': 1, 'sum': 2}
        self.reduction_id = self.red_map[reduction]
        self.block_size = 256
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>

        torch::Tensor log_cosh_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int reduction);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE 256
        #define WARP_SIZE 32

        __inline__ __device__ float warp_reduce_sum(float val) {
            #pragma unroll
            for (int offset = 16; offset > 0; offset /= 2) {
                val += __shfl_down_sync(0xffffffff, val, offset);
            }
            return val;
        }

        __inline__ __device__ float block_reduce_sum(float val) {
            __shared__ float shared[32];
            int lane = threadIdx.x % 32;
            int wid = threadIdx.x / 32;

            val = warp_reduce_sum(val);
            if (lane == 0) shared[wid] = val;
            __syncthreads();

            val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f;
            if (wid == 0) val = warp_reduce_sum(val);
            return val;
        }

        __global__ void log_cosh_kernel(
            const float* __restrict__ input,
            const float* __restrict__ target,
            float* __restrict__ output,
            int n,
            int reduction
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            int stride = blockDim.x * gridDim.x;

            float local_sum = 0.0f;
            float log_2 = 0.69314718056f;

            float4* in_ptr = (float4*)input;
            float4* tgt_ptr = (float4*)target;
            float4* out_ptr = (float4*)output;

            int vec_n = n / 4;

            for (int i = idx; i < vec_n; i += stride) {
                float4 in_val = in_ptr[i];
                float4 tgt_val = tgt_ptr[i];

                float diff[4];
                diff[0] = fabsf(in_val.x - tgt_val.x);
                diff[1] = fabsf(in_val.y - tgt_val.y);
                diff[2] = fabsf(in_val.z - tgt_val.z);
                diff[3] = fabsf(in_val.w - tgt_val.w);

                float losses[4];
                #pragma unroll
                for(int k=0; k<4; ++k) {
                    losses[k] = diff[k] + log1pf(expf(-2.0f * diff[k])) - log_2;
                }

                if (reduction == 0) {
                    float4 res;
                    res.x = losses[0]; res.y = losses[1];
                    res.z = losses[2]; res.w = losses[3];
                    out_ptr[i] = res;
                } else {
                    local_sum += losses[0] + losses[1] + losses[2] + losses[3];
                }
            }

            int rem_start = vec_n * 4;
            for (int i = rem_start + idx; i < n; i += stride) {
                float diff = fabsf(input[i] - target[i]);
                float loss = diff + log1pf(expf(-2.0f * diff)) - log_2;

                if (reduction == 0) {
                    output[i] = loss;
                } else {
                    local_sum += loss;
                }
            }

            if (reduction != 0) {
                local_sum = block_reduce_sum(local_sum);
                if (threadIdx.x == 0) {
                    atomicAdd(output, local_sum);
                }
            }
        }

        torch::Tensor log_cosh_forward_cuda(
            torch::Tensor input,
            torch::Tensor target,
            int reduction) 
        {
            int64_t n = input.numel();
            auto options = input.options();

            torch::Tensor output;
            if (reduction == 0) {
                output = torch::empty_like(input);
            } else {
                output = torch::zeros({1}, options);
            }

            const int block_size = 256;
            const int grid_size = std::min((int)((n + block_size * 4 - 1) / (block_size * 4)), 1024);

            log_cosh_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                target.data_ptr<float>(),
                output.data_ptr<float>(),
                n,
                reduction
            );

            if (reduction == 1) {
                output.div_(n);
            }

            return output;
        }
        """

        self.op = load_inline(
            name='log_cosh_cuda_opt',
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['log_cosh_forward_cuda'],
            extra_cuda_cflags=['-O3', '--use_fast_math'],
            verbose=False
        )

    def forward(self, input, target):
        if not input.is_cuda: input = input.cuda()
        if not target.is_cuda: target = target.cuda()

        input = input.contiguous()
        target = target.contiguous()

        return self.op.log_cosh_forward_cuda(input, target, self.reduction_id)