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
from torch.utils.cpp_extension import load_inline


_cpp_src = r'''
#include <torch/extension.h>

void relu_group_norm_launch(const float* x, const float* gamma, const float* beta,
                            float* y, float* stats,
                            int N, int C, int S, int G, cudaStream_t stream);

torch::Tensor relu_group_norm(torch::Tensor input, torch::Tensor gamma, torch::Tensor beta, int64_t groups) {
    TORCH_CHECK(input.is_cuda(), "input must be CUDA");
    TORCH_CHECK(gamma.is_cuda() && beta.is_cuda(), "affine params must be CUDA");
    TORCH_CHECK(input.scalar_type() == at::kFloat, "only float32 is supported");
    TORCH_CHECK(gamma.scalar_type() == at::kFloat && beta.scalar_type() == at::kFloat, "only float32 affine params are supported");
    TORCH_CHECK(input.dim() == 5, "expected NCDHW tensor");

    auto x = input.contiguous();
    auto g = gamma.contiguous();
    auto b = beta.contiguous();

    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int S = (int)(x.size(2) * x.size(3) * x.size(4));
    const int G = (int)groups;
    TORCH_CHECK(C % G == 0, "channels must be divisible by groups");

    auto y = torch::empty_like(x);
    auto stats = torch::empty({N, G, 2}, x.options());

    relu_group_norm_launch(x.data_ptr<float>(), g.data_ptr<float>(), b.data_ptr<float>(),
                           y.data_ptr<float>(), stats.data_ptr<float>(),
                           N, C, S, G, at::cuda::getCurrentCUDAStream());
    return y;
}
'''

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

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

static __inline__ __device__ float block_reduce_sum(float val) {
    __shared__ float shared[32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    val = warp_reduce_sum(val);
    if (lane == 0) shared[wid] = val;
    __syncthreads();
    val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : 0.0f;
    if (wid == 0) val = warp_reduce_sum(val);
    return val;
}

__global__ void relu_gn_stats_kernel(const float* __restrict__ x,
                                     float* __restrict__ stats,
                                     int N, int C, int S, int G) {
    int ng = blockIdx.x;
    int n = ng / G;
    int g = ng - n * G;
    int cpg = C / G;
    int elems = cpg * S;
    int c0 = g * cpg;

    float sum = 0.0f;
    float ss = 0.0f;

    for (int i = threadIdx.x; i < elems; i += blockDim.x) {
        int lc = i / S;
        int s = i - lc * S;
        int c = c0 + lc;
        float v = x[((n * C + c) * S) + s];
        v = v > 0.0f ? v : 0.0f;
        sum += v;
        ss = fmaf(v, v, ss);
    }

    float rsum = block_reduce_sum(sum);
    float rss = block_reduce_sum(ss);

    if (threadIdx.x == 0) {
        float inv_e = 1.0f / (float)elems;
        float mean = rsum * inv_e;
        float var = rss * inv_e - mean * mean;
        var = var > 0.0f ? var : 0.0f;
        stats[ng * 2] = mean;
        stats[ng * 2 + 1] = rsqrtf(var + 1.0e-5f);
    }
}

__global__ void relu_gn_norm_kernel(const float* __restrict__ x,
                                    const float* __restrict__ gamma,
                                    const float* __restrict__ beta,
                                    const float* __restrict__ stats,
                                    float* __restrict__ y,
                                    long long total, int C, int S, int G) {
    int cpg = C / G;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         idx < total;
         idx += (long long)blockDim.x * gridDim.x) {
        int s = (int)(idx % S);
        long long t = idx / S;
        int c = (int)(t % C);
        int n = (int)(t / C);
        int g = c / cpg;
        int ng = n * G + g;
        float v = x[idx];
        v = v > 0.0f ? v : 0.0f;
        float out = (v - stats[ng * 2]) * stats[ng * 2 + 1];
        y[idx] = out * gamma[c] + beta[c];
    }
}

void relu_group_norm_launch(const float* x, const float* gamma, const float* beta,
                            float* y, float* stats,
                            int N, int C, int S, int G, cudaStream_t stream) {
    const int threads = 256;
    relu_gn_stats_kernel<<<N * G, threads, 0, stream>>>(x, stats, N, C, S, G);
    long long total = (long long)N * C * S;
    int blocks = (int)((total + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    relu_gn_norm_kernel<<<blocks, threads, 0, stream>>>(x, gamma, beta, stats, y, total, C, S, G);
}
'''

_ext = load_inline(
    name="relu_group_norm_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["relu_group_norm"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    with_cuda=True,
    verbose=False,
)


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.relu = nn.ReLU()
        self.group_norm = nn.GroupNorm(num_groups=groups, num_channels=out_channels)
        self.groups = groups

    def forward(self, x):
        y = self.conv_transpose(x)
        return _ext.relu_group_norm(y, self.group_norm.weight, self.group_norm.bias, int(self.groups))