import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Performs a standard 3D convolution operation with square input and 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.
        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, kernel_size), 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, depth, width, height).

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

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

def get_inputs():
    x = torch.rand(batch_size, in_channels, depth, width, height)
    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_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  int64_t stride, int64_t padding, int64_t dilation, int64_t groups);

torch::Tensor conv3d_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                             int64_t stride, int64_t padding, int64_t dilation, int64_t groups) {
    return conv3d_forward_cuda(x, weight, bias, stride, padding, dilation, groups);
}
'''

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

__constant__ float c_weight[8192];
__constant__ float c_bias[1024];

__global__ __launch_bounds__(256, 4)
void conv3d_k3_c3_kernel(const float* __restrict__ x, float* __restrict__ y,
                         int N, int OC, int D, int H, int W,
                         int Do, int Ho, int Wo, int has_bias, int total) {
    const int S = Do * Ho * Wo;
    const int HW = H * W;
    const int stride_grid = blockDim.x * gridDim.x;
    for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += stride_grid) {
        int t = idx;
        const int s = t % S;  t /= S;
        const int oc = t % OC;
        const int n = t / OC;

        int r = s;
        const int ow = r % Wo; r /= Wo;
        const int oh = r % Ho;
        const int od = r / Ho;

        const float* base0 = x + ((n * 3) * D + od) * HW + oh * W + ow;
        const float* base1 = base0 + D * HW;
        const float* base2 = base1 + D * HW;
        const int wo = oc * 81;
        float acc = 0.0f;

        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int zoff = kd * HW;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int off = zoff + kh * W;
                const float x00 = base0[off + 0];
                const float x01 = base0[off + 1];
                const float x02 = base0[off + 2];
                const int w0 = wo + kd * 9 + kh * 3;
                acc = fmaf(x00, c_weight[w0 + 0], acc);
                acc = fmaf(x01, c_weight[w0 + 1], acc);
                acc = fmaf(x02, c_weight[w0 + 2], acc);
            }
        }
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int zoff = kd * HW;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int off = zoff + kh * W;
                const float x00 = base1[off + 0];
                const float x01 = base1[off + 1];
                const float x02 = base1[off + 2];
                const int w0 = wo + 27 + kd * 9 + kh * 3;
                acc = fmaf(x00, c_weight[w0 + 0], acc);
                acc = fmaf(x01, c_weight[w0 + 1], acc);
                acc = fmaf(x02, c_weight[w0 + 2], acc);
            }
        }
        #pragma unroll
        for (int kd = 0; kd < 3; ++kd) {
            const int zoff = kd * HW;
            #pragma unroll
            for (int kh = 0; kh < 3; ++kh) {
                const int off = zoff + kh * W;
                const float x00 = base2[off + 0];
                const float x01 = base2[off + 1];
                const float x02 = base2[off + 2];
                const int w0 = wo + 54 + kd * 9 + kh * 3;
                acc = fmaf(x00, c_weight[w0 + 0], acc);
                acc = fmaf(x01, c_weight[w0 + 1], acc);
                acc = fmaf(x02, c_weight[w0 + 2], acc);
            }
        }
        if (has_bias) acc += c_bias[oc];
        y[idx] = acc;
    }
}

__global__ __launch_bounds__(256, 3)
void conv3d_generic_kernel(const float* __restrict__ x, const float* __restrict__ wgt,
                           const float* __restrict__ bias, float* __restrict__ y,
                           int N, int C, int D, int H, int W, int OC, int K,
                           int stride, int pad, int dil, int groups,
                           int Do, int Ho, int Wo, int has_bias, long long total) {
    const int S = Do * Ho * Wo;
    const int Cg = C / groups;
    const int Og = OC / groups;
    const int HW = H * W;
    const long long step = (long long)blockDim.x * gridDim.x;
    for (long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; idx < total; idx += step) {
        long long t = idx;
        int s = (int)(t % S); t /= S;
        int oc = (int)(t % OC);
        int n = (int)(t / OC);
        int r = s;
        int ow = r % Wo; r /= Wo;
        int oh = r % Ho;
        int od = r / Ho;
        int g = oc / Og;
        float acc = 0.0f;
        for (int icg = 0; icg < Cg; ++icg) {
            int ic = g * Cg + icg;
            for (int kd = 0; kd < K; ++kd) {
                int id = od * stride - pad + kd * dil;
                if ((unsigned)id >= (unsigned)D) continue;
                for (int kh = 0; kh < K; ++kh) {
                    int ih = oh * stride - pad + kh * dil;
                    if ((unsigned)ih >= (unsigned)H) continue;
                    for (int kw = 0; kw < K; ++kw) {
                        int iw = ow * stride - pad + kw * dil;
                        if ((unsigned)iw >= (unsigned)W) continue;
                        float xv = x[((n * C + ic) * D + id) * HW + ih * W + iw];
                        float wv = wgt[(((oc * Cg + icg) * K + kd) * K + kh) * K + kw];
                        acc = fmaf(xv, wv, acc);
                    }
                }
            }
        }
        if (has_bias) acc += bias[oc];
        y[idx] = acc;
    }
}

torch::Tensor conv3d_forward_cuda(torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
                                  int64_t stride64, int64_t padding64, int64_t dilation64, int64_t groups64) {
    const int N = (int)x.size(0);
    const int C = (int)x.size(1);
    const int D = (int)x.size(2);
    const int H = (int)x.size(3);
    const int W = (int)x.size(4);
    const int OC = (int)weight.size(0);
    const int K = (int)weight.size(2);
    const int stride = (int)stride64;
    const int pad = (int)padding64;
    const int dil = (int)dilation64;
    const int groups = (int)groups64;
    const int Do = (D + 2 * pad - dil * (K - 1) - 1) / stride + 1;
    const int Ho = (H + 2 * pad - dil * (K - 1) - 1) / stride + 1;
    const int Wo = (W + 2 * pad - dil * (K - 1) - 1) / stride + 1;
    auto y = torch::empty({N, OC, Do, Ho, Wo}, x.options());
    const int has_bias = (bias.numel() != 0);
    const int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    if (C == 3 && K == 3 && stride == 1 && pad == 0 && dil == 1 && groups == 1 && weight.numel() <= 8192 && OC <= 1024) {
        size_t bytes = (size_t)weight.numel() * sizeof(float);
        C10_CUDA_CHECK(cudaMemcpyToSymbolAsync(c_weight, weight.data_ptr<float>(), bytes, 0, cudaMemcpyDeviceToDevice, stream));
        if (has_bias) {
            C10_CUDA_CHECK(cudaMemcpyToSymbolAsync(c_bias, bias.data_ptr<float>(), (size_t)OC * sizeof(float), 0, cudaMemcpyDeviceToDevice, stream));
        }
        int total = N * OC * Do * Ho * Wo;
        int blocks = (total + threads - 1) / threads;
        if (blocks > 65535) blocks = 65535;
        conv3d_k3_c3_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), y.data_ptr<float>(), N, OC, D, H, W, Do, Ho, Wo, has_bias, total);
    } else {
        long long total = (long long)N * OC * Do * Ho * Wo;
        int blocks = (int)((total + threads - 1) / threads);
        if (blocks > 65535) blocks = 65535;
        const float* bptr = has_bias ? bias.data_ptr<float>() : nullptr;
        conv3d_generic_kernel<<<blocks, threads, 0, stream>>>(x.data_ptr<float>(), weight.data_ptr<float>(), bptr, y.data_ptr<float>(),
            N, C, D, H, W, OC, K, stride, pad, dil, groups, Do, Ho, Wo, has_bias, total);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
'''

_ext = load_inline(
    name='conv3d_direct_k3c3_ext_v1',
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=['conv3d_forward'],
    extra_cflags=['-O3'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
    with_cuda=True,
    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, kernel_size), 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:
        # Inputs in the benchmark are contiguous CUDA fp32 tensors; keep the module state_dict-compatible with nn.Conv3d.
        b = self.conv3d.bias
        if b is None:
            b = torch.empty(0, device=x.device, dtype=x.dtype)
        return _ext.conv3d_forward(x.contiguous(), self.conv3d.weight.contiguous(), b, self._stride, self._padding, self._dilation, self._groups)