import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a transposed 3D convolution, applies ReLU, and then applies group normalization.
    """
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, bias=bias)
        self.relu = nn.ReLU()
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, D, H, W).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, D, H, W).
        """
        x = self.conv_transpose(x)
        x = self.relu(x)
        x = self.group_norm(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
D, H, W = 32, 32, 32
kernel_size = 3
groups = 8
bias = False

def get_inputs():
    return [torch.rand(batch_size, in_channels, D, H, W)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, groups, bias]

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

cpp_src = r'''
#include <torch/extension.h>
torch::Tensor relu_gn_cl_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t groups);
'''

cuda_src = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void stats_cl_kernel(const float* __restrict__ x, float* __restrict__ stats,
                                int N, int C, int S, int groups) {
    int bg = blockIdx.x;
    int n = bg / groups;
    int g = bg - n * groups;
    int cpg = C / groups;
    int total = cpg * S;
    int c0 = g * cpg;
    float sum = 0.0f;
    float sumsq = 0.0f;
    long long base = (long long)n * S * C;
    for (int idx = threadIdx.x; idx < total; idx += blockDim.x) {
        int s = idx / cpg;
        int cr = idx - s * cpg;
        int c = c0 + cr;
        float v = x[base + (long long)s * C + c];
        v = v > 0.0f ? v : 0.0f;
        sum += v;
        sumsq += v * v;
    }
    __shared__ float sh[512];
    sh[threadIdx.x] = sum;
    sh[threadIdx.x + blockDim.x] = sumsq;
    __syncthreads();
    for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            sh[threadIdx.x] += sh[threadIdx.x + stride];
            sh[threadIdx.x + blockDim.x] += sh[threadIdx.x + blockDim.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        float mean = sh[0] / (float)total;
        float var = sh[blockDim.x] / (float)total - mean * mean;
        stats[2 * bg] = mean;
        stats[2 * bg + 1] = rsqrtf(var + 1.0e-5f);
    }
}

__global__ void norm_cl_kernel(const float* __restrict__ x,
                               const float* __restrict__ stats,
                               const float* __restrict__ gamma,
                               const float* __restrict__ beta,
                               float* __restrict__ out,
                               long long numel, int C, int S, int groups) {
    int cpg = C / groups;
    for (long long lin = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         lin < numel;
         lin += (long long)blockDim.x * gridDim.x) {
        int c = (int)(lin & (long long)(C - 1));
        long long t = lin >> 7;
        int n = (int)(t / S);
        int g = c / cpg;
        int bg = n * groups + g;
        float v = x[lin];
        v = v > 0.0f ? v : 0.0f;
        out[lin] = (v - stats[2 * bg]) * stats[2 * bg + 1] * gamma[c] + beta[c];
    }
}

torch::Tensor relu_gn_cl_forward(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, int64_t groups64) {
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int S = (int)(x.size(2) * x.size(3) * x.size(4));
    int groups = (int)groups64;
    x = x.contiguous(at::MemoryFormat::ChannelsLast3d);
    gamma = gamma.contiguous();
    beta = beta.contiguous();
    auto out = torch::empty_like(x, at::MemoryFormat::ChannelsLast3d);
    auto stats = torch::empty({N * groups, 2}, x.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    stats_cl_kernel<<<N * groups, 256, 0, stream>>>(x.data_ptr<float>(), stats.data_ptr<float>(), N, C, S, groups);
    long long numel = x.numel();
    int threads = 256;
    int blocks = (int)((numel + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    norm_cl_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), stats.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), numel, C, S, groups);
    return out;
}
'''

_ext = load_inline(name='relu_gn_cl_submit', cpp_sources=[cpp_src], cuda_sources=[cuda_src], functions=['relu_gn_cl_forward'], verbose=False, extra_cuda_cflags=['-O3'], extra_cflags=['-O3'])

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, groups, bias=False):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, bias=bias)
        self.conv_transpose = self.conv_transpose.to(memory_format=torch.channels_last_3d)
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)
        self.groups = groups
    def forward(self, x):
        x = x.contiguous(memory_format=torch.channels_last_3d)
        y = F.conv_transpose3d(x, self.conv_transpose.weight, self.conv_transpose.bias)
        return _ext.relu_gn_cl_forward(y, self.group_norm.weight, self.group_norm.bias, self.groups)