import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 3D convolution operation with an asymmetric input and a square kernel.

    Args:
        in_channels (int): Number of channels in the input tensor.
        out_channels (int): Number of channels produced by the convolution.
        kernel_size (int): Size of the square convolution kernel (kernel_size x kernel_size).
        stride (int, optional): Stride of the convolution. Defaults to 1.
        padding (int, optional): Padding applied to the input. Defaults to 0.
        dilation (int, optional): Spacing between kernel elements. Defaults to 1.
        groups (int, optional): Number of blocked connections from input channels to output channels. Defaults to 1.
        bias (bool, optional): If `True`, adds a learnable bias to the output. Defaults to `False`.
    """
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(Model, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, (kernel_size, kernel_size, 1), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs the 3D convolution.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_channels, height, width, depth).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, out_channels, height_out, width_out, depth_out).
        """
        return self.conv3d(x)

# Test code
batch_size = 16
in_channels = 3
out_channels = 64
kernel_size = 3
width = 256
height = 256
depth = 10

def get_inputs():
    x = torch.rand(batch_size, in_channels, height, width, depth)
    return [x]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size]  # Provide in_channels, out_channels, kernel_size for initialization

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

_tmp = os.path.join(os.getcwd(), 'nvcc_tmp')
os.makedirs(_tmp, exist_ok=True)
os.environ['TMPDIR'] = _tmp

_cpp = r'''
#include <torch/extension.h>
torch::Tensor conv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  int64_t stride, int64_t padding, int64_t dilation, int64_t groups);
'''

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

__constant__ float CW[8192];
__constant__ float CB[1024];

__global__ __launch_bounds__(256)
void conv_k331_oc16_fixed(const float* __restrict__ x, float* __restrict__ y, int has_bias, int N) {
    const int OW = 10;
    const int OH = 254;
    const int OD = 254;
    const int OC = 64;
    const int OCG = 16;
    const int GROUPS = 4;
    const int PLANE_IN = 2560;       // 256 * 10
    const int CHAN_IN  = 655360;     // 256 * 256 * 10
    const int SPATIAL_OUT = 645160;  // 254 * 254 * 10
    const int units = N * GROUPS * SPATIAL_OUT;
    int tid = (int)(blockIdx.x * blockDim.x + threadIdx.x);
    int step = (int)(blockDim.x * gridDim.x);
    for (int u = tid; u < units; u += step) {
        int ow = u % OW;
        int q = u / OW;
        int oh = q % OH; q /= OH;
        int od = q % OD; q /= OD;
        int ocg = q % GROUPS;
        int n = q / GROUPS;
        int oc0 = ocg * OCG;

        float acc[16];
        if (has_bias) {
            #pragma unroll
            for (int i = 0; i < 16; ++i) acc[i] = CB[oc0 + i];
        } else {
            #pragma unroll
            for (int i = 0; i < 16; ++i) acc[i] = 0.0f;
        }

        const int base = n * (3 * CHAN_IN) + od * PLANE_IN + oh * OW + ow;
        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    float xv = x[base + c * CHAN_IN + kd * PLANE_IN + kh * OW];
                    int koff = c * 9 + kd * 3 + kh;
                    #pragma unroll
                    for (int i = 0; i < 16; ++i) {
                        acc[i] = fmaf(xv, CW[(oc0 + i) * 27 + koff], acc[i]);
                    }
                }
            }
        }

        long out_sp = ((long)od * OH + oh) * OW + ow;
        long ybase = ((long)n * OC + oc0) * SPATIAL_OUT + out_sp;
        #pragma unroll
        for (int i = 0; i < 16; ++i) {
            y[ybase + (long)i * SPATIAL_OUT] = acc[i];
        }
    }
}

__global__ __launch_bounds__(256)
void conv_generic_kernel(const float* __restrict__ x, const float* __restrict__ w,
                         const float* __restrict__ b, float* __restrict__ y, long total,
                         int N, int C, int OC, int ID, int IH, int IW,
                         int KD, int KH, int KW, int OD, int OH, int OW,
                         int stride_v, int pad, int dil, int has_bias) {
    long tid = (long)blockIdx.x * blockDim.x + threadIdx.x;
    long step = (long)blockDim.x * gridDim.x;
    for (long idx = tid; idx < total; idx += step) {
        int ow = idx % OW;
        long q = idx / OW;
        int oh = q % OH; q /= OH;
        int od = q % OD; q /= OD;
        int oc = q % OC; q /= OC;
        int n = (int)q;
        float acc = has_bias ? b[oc] : 0.0f;
        for (int c = 0; c < C; ++c) {
            for (int kd = 0; kd < KD; ++kd) {
                int id = od * stride_v - pad + kd * dil;
                if ((unsigned)id >= (unsigned)ID) continue;
                for (int kh = 0; kh < KH; ++kh) {
                    int ih = oh * stride_v - pad + kh * dil;
                    if ((unsigned)ih >= (unsigned)IH) continue;
                    for (int kw = 0; kw < KW; ++kw) {
                        int iw = ow * stride_v - pad + kw * dil;
                        if ((unsigned)iw >= (unsigned)IW) continue;
                        long xi = (((long)n * C + c) * ID + id) * IH * IW + ih * IW + iw;
                        long wi = (((long)oc * C + c) * KD + kd) * KH * KW + kh * KW + kw;
                        acc = fmaf(x[xi], w[wi], acc);
                    }
                }
            }
        }
        y[idx] = acc;
    }
}

torch::Tensor conv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int ID = (int)x.size(2);
    int IH = (int)x.size(3);
    int IW = (int)x.size(4);
    int OC = (int)weight.size(0);
    int KD = (int)weight.size(2);
    int KH = (int)weight.size(3);
    int KW = (int)weight.size(4);
    int OD = (ID + 2 * (int)padding - (int)dilation * (KD - 1) - 1) / (int)stride + 1;
    int OH = (IH + 2 * (int)padding - (int)dilation * (KH - 1) - 1) / (int)stride + 1;
    int OW = (IW + 2 * (int)padding - (int)dilation * (KW - 1) - 1) / (int)stride + 1;
    auto y = torch::empty({N, OC, OD, OH, OW}, x.options());
    int has_bias = bias.numel() > 0;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, x.get_device());
    int threads = 256;

    if (C == 3 && OC == 64 && ID == 256 && IH == 256 && IW == 10 &&
        KD == 3 && KH == 3 && KW == 1 && stride == 1 && padding == 0 && dilation == 1 && groups == 1) {
        cudaMemcpyToSymbolAsync(CW, weight.data_ptr<float>(), 64 * 27 * sizeof(float), 0, cudaMemcpyDeviceToDevice, stream);
        if (has_bias) cudaMemcpyToSymbolAsync(CB, bias.data_ptr<float>(), 64 * sizeof(float), 0, cudaMemcpyDeviceToDevice, stream);
        int units = N * 4 * 254 * 254 * 10;
        int blocks = (units + threads - 1) / threads;
        int max_blocks = prop.multiProcessorCount * 32;
        if (blocks > max_blocks) blocks = max_blocks;
        conv_k331_oc16_fixed<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), has_bias, N);
    } else {
        long total = (long)N * OC * OD * OH * OW;
        int blocks = (int)((total + threads - 1) / threads);
        int max_blocks = prop.multiProcessorCount * 16;
        if (blocks > max_blocks) blocks = max_blocks;
        conv_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), has_bias ? bias.data_ptr<float>() : nullptr,
            y.data_ptr<float>(), total, N, C, OC, ID, IH, IW, KD, KH, KW, OD, OH, OW,
            (int)stride, (int)padding, (int)dilation, has_bias);
    }
    return y;
}
'''

_ext = load_inline(
    name='conv3d_k331_oc16_ext',
    cpp_sources=[_cpp],
    cuda_sources=[_cuda],
    functions=['conv3d_forward_cuda'],
    verbose=False,
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math']
)

class ModelNew(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = False):
        super(ModelNew, self).__init__()
        self.conv3d = nn.Conv3d(in_channels, out_channels, (kernel_size, kernel_size, 1), stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        b = self.conv3d.bias if self.conv3d.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv3d_forward_cuda(x, self.conv3d.weight, b, self.stride, self.padding, self.dilation, self.groups)