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

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

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int stride, int padding, int dilation, int groups);
"""

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

__global__ void conv3d_k3c3_oc8_kernel(const float* __restrict__ x,
                                       const float* __restrict__ w,
                                       const float* __restrict__ b,
                                       float* __restrict__ y,
                                       int total, int N, int H, int W, int D,
                                       int O, int Hout, int Wout, int inner,
                                       bool has_bias) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int stride_grid = blockDim.x * gridDim.x;
    for (int idx = tid; idx < total; idx += stride_grid) {
        int spatial = idx % inner;
        int t0 = idx / inner;
        int og = t0 & 7;          // O/8 == 8 in the target case (O=64)
        int n = t0 >> 3;
        int oc0 = og << 3;

        int d = spatial % D;
        int t = spatial / D;
        int ow = t % Wout;
        int oh = t / Wout;

        float a0 = has_bias ? b[oc0 + 0] : 0.0f;
        float a1 = has_bias ? b[oc0 + 1] : 0.0f;
        float a2 = has_bias ? b[oc0 + 2] : 0.0f;
        float a3 = has_bias ? b[oc0 + 3] : 0.0f;
        float a4 = has_bias ? b[oc0 + 4] : 0.0f;
        float a5 = has_bias ? b[oc0 + 5] : 0.0f;
        float a6 = has_bias ? b[oc0 + 6] : 0.0f;
        float a7 = has_bias ? b[oc0 + 7] : 0.0f;

        #pragma unroll
        for (int c = 0; c < 3; ++c) {
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                int ih = oh + kh;
                #pragma unroll
                for (int kw = 0; kw < 3; ++kw) {
                    int iw = ow + kw;
                    float xv = __ldg(x + ((((n * 3 + c) * H + ih) * W + iw) * D + d));
                    int wo = ((oc0 * 3 + c) * 3 + kh) * 3 + kw;
                    a0 = fmaf(xv, __ldg(w + wo + 0 * 27), a0);
                    a1 = fmaf(xv, __ldg(w + wo + 1 * 27), a1);
                    a2 = fmaf(xv, __ldg(w + wo + 2 * 27), a2);
                    a3 = fmaf(xv, __ldg(w + wo + 3 * 27), a3);
                    a4 = fmaf(xv, __ldg(w + wo + 4 * 27), a4);
                    a5 = fmaf(xv, __ldg(w + wo + 5 * 27), a5);
                    a6 = fmaf(xv, __ldg(w + wo + 6 * 27), a6);
                    a7 = fmaf(xv, __ldg(w + wo + 7 * 27), a7);
                }
            }
        }

        int out = (n * O + oc0) * inner + spatial;
        y[out + 0 * inner] = a0;
        y[out + 1 * inner] = a1;
        y[out + 2 * inner] = a2;
        y[out + 3 * inner] = a3;
        y[out + 4 * inner] = a4;
        y[out + 5 * inner] = a5;
        y[out + 6 * inner] = a6;
        y[out + 7 * inner] = a7;
    }
}

__global__ void conv3d_generic_kernel(const float* __restrict__ x,
                                      const float* __restrict__ w,
                                      const float* __restrict__ b,
                                      float* __restrict__ y,
                                      int total, int N, int C, int H, int W, int D,
                                      int O, int Cg, int kH, int kW, int kD,
                                      int Hout, int Wout, int Dout,
                                      int stride, int pad, int dil, int groups,
                                      bool has_bias) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int grid_stride = blockDim.x * gridDim.x;
    for (int idx = tid; idx < total; idx += grid_stride) {
        int od = idx % Dout;
        int q = idx / Dout;
        int ow = q % Wout; q /= Wout;
        int oh = q % Hout; q /= Hout;
        int oc = q % O;
        int n = q / O;
        int g = oc / (O / groups);
        int cbase = g * Cg;
        float acc = has_bias ? b[oc] : 0.0f;
        for (int ci = 0; ci < Cg; ++ci) {
            int ic = cbase + ci;
            for (int kh = 0; kh < kH; ++kh) {
                int ih = oh * stride - pad + kh * dil;
                if ((unsigned)ih >= (unsigned)H) continue;
                for (int kw = 0; kw < kW; ++kw) {
                    int iw = ow * stride - pad + kw * dil;
                    if ((unsigned)iw >= (unsigned)W) continue;
                    for (int kd = 0; kd < kD; ++kd) {
                        int id = od * stride - pad + kd * dil;
                        if ((unsigned)id >= (unsigned)D) continue;
                        float xv = x[((((n * C + ic) * H + ih) * W + iw) * D + id)];
                        float wv = w[(((((oc * Cg + ci) * kH + kh) * kW + kw) * kD + kd))];
                        acc = fmaf(xv, wv, acc);
                    }
                }
            }
        }
        y[idx] = acc;
    }
}

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int stride, int padding, int dilation, int groups) {
    const c10::cuda::CUDAGuard device_guard(x.device());
    int N = (int)x.size(0);
    int C = (int)x.size(1);
    int H = (int)x.size(2);
    int W = (int)x.size(3);
    int D = (int)x.size(4);
    int O = (int)weight.size(0);
    int Cg = (int)weight.size(1);
    int kH = (int)weight.size(2);
    int kW = (int)weight.size(3);
    int kD = (int)weight.size(4);
    int Hout = (H + 2 * padding - dilation * (kH - 1) - 1) / stride + 1;
    int Wout = (W + 2 * padding - dilation * (kW - 1) - 1) / stride + 1;
    int Dout = (D + 2 * padding - dilation * (kD - 1) - 1) / stride + 1;
    auto y = torch::empty({N, O, Hout, Wout, Dout}, x.options());
    bool has_bias = bias.numel() != 0;
    const float* bp = has_bias ? bias.data_ptr<float>() : nullptr;
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (x.scalar_type() == at::kFloat && weight.scalar_type() == at::kFloat &&
        groups == 1 && stride == 1 && padding == 0 && dilation == 1 &&
        C == 3 && Cg == 3 && kH == 3 && kW == 3 && kD == 1 && (O % 8) == 0 && O == 64) {
        int inner = Hout * Wout * Dout;
        int total = N * (O / 8) * inner;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 65535) blocks = 65535;
        conv3d_k3c3_oc8_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            total, N, H, W, D, O, Hout, Wout, inner, has_bias);
    } else {
        int total = N * O * Hout * Wout * Dout;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 65535) blocks = 65535;
        conv3d_generic_kernel<<<blocks, threads, 0, stream>>>(
            x.data_ptr<float>(), weight.data_ptr<float>(), bp, y.data_ptr<float>(),
            total, N, C, H, W, D, O, Cg, kH, kW, kD, Hout, Wout, Dout,
            stride, padding, dilation, groups, has_bias);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
"""

_ext = load_inline(
    name="conv3d_asym_k3_fast_ext",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["conv3d_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)

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 if isinstance(stride, int) else stride[0]
        self._padding = padding if isinstance(padding, int) else padding[0]
        self._dilation = dilation if isinstance(dilation, int) else dilation[0]
        self._groups = groups

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not x.is_contiguous():
            x = x.contiguous()
        w = self.conv3d.weight
        if not w.is_contiguous():
            w = w.contiguous()
        b = self.conv3d.bias
        if b is None:
            b = x.new_empty((0,))
        return _ext.conv3d_forward(x, w, b, int(self._stride), int(self._padding), int(self._dilation), int(self._groups))