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

# CUDA implementation of Conv2D (tiled + shared memory + output-channel blocking)
conv2d_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <math.h>

#ifndef CHECK_CUDA
#define CHECK_CUDA(x) TORCH_CHECK((x).is_cuda(), #x " must be a CUDA tensor")
#endif
#ifndef CHECK_CONTIGUOUS
#define CHECK_CONTIGUOUS(x) TORCH_CHECK((x).is_contiguous(), #x " must be contiguous")
#endif
#ifndef CHECK_FLOAT
#define CHECK_FLOAT(x) TORCH_CHECK((x).scalar_type() == at::kFloat, #x " must be float32")
#endif

// 每个block计算一个 (b, oc_group) 上的输出tile，复用输入tile计算 OC_TILE 个输出通道
template<int BLOCK_X, int BLOCK_Y, int OC_TILE>
__global__ void conv2d_tiled_kernel_oc(
    const float* __restrict__ input,   // [B, C_in, H, W]
    const float* __restrict__ weight,  // [C_out, C_in, K, K]
    const float* __restrict__ bias,    // [C_out] or nullptr
    float* __restrict__ output,        // [B, C_out, H_out, W_out]
    int B, int C_in, int C_out,
    int H, int W, int K, int H_out, int W_out,
    bool has_bias
) {
    // grid.z = B * ceil_div(C_out, OC_TILE)
    int groups = (C_out + OC_TILE - 1) / OC_TILE;
    int b  = blockIdx.z / groups;
    int og = blockIdx.z % groups;          // 输出通道组编号
    int co0 = og * OC_TILE;                // 本组起始输出通道

    int ow0 = blockIdx.x * BLOCK_X;
    int oh0 = blockIdx.y * BLOCK_Y;

    int ow = ow0 + threadIdx.x;
    int oh = oh0 + threadIdx.y;

    extern __shared__ float smem[];
    // 输入tile大小：(BLOCK_Y+K-1) x (BLOCK_X+K-1)
    int tile_w = BLOCK_X + K - 1;
    int tile_h = BLOCK_Y + K - 1;
    float* tile = smem;                                            // tile_h * tile_w
    float* w_sh = tile + tile_h * tile_w;                          // OC_TILE * K * K
    // w_sh 布局: [oc_local][K*K]

    // 累加器：每线程维护 OC_TILE 个通道
    float acc[OC_TILE];
    #pragma unroll
    for (int oc = 0; oc < OC_TILE; ++oc) {
        int co = co0 + oc;
        acc[oc] = (has_bias && co < C_out) ? bias[co] : 0.0f;
    }

    bool valid_xy = (oh < H_out) && (ow < W_out);

    // 遍历输入通道
    for (int ci = 0; ci < C_in; ++ci) {
        // 1) 加载本组 OC_TILE 的权重到共享内存
        int total_w = OC_TILE * K * K;
        for (int t = threadIdx.y * BLOCK_X + threadIdx.x; t < total_w; t += BLOCK_X * BLOCK_Y) {
            int oc = t / (K*K);
            int rem = t % (K*K);
            int kh = rem / K;
            int kw = rem % K;
            int co = co0 + oc;
            float wv = 0.0f;
            if (co < C_out) {
                int w_idx = ((co * C_in + ci) * K + kh) * K + kw;
                wv = weight[w_idx];
            }
            w_sh[t] = wv;
        }

        // 2) 加载输入tile到共享内存（该tile将被 OC_TILE 个输出通道复用）
        int ih0 = oh0;
        int iw0 = ow0;
        for (int th = threadIdx.y; th < tile_h; th += BLOCK_Y) {
            int ih = ih0 + th;
            bool in_h = (ih >= 0) && (ih < H);
            for (int tw = threadIdx.x; tw < tile_w; tw += BLOCK_X) {
                int iw = iw0 + tw;
                bool in_w = (iw >= 0) && (iw < W);
                float v = 0.0f;
                if (in_h && in_w) {
                    int in_idx = (((b * C_in + ci) * H + ih) * W + iw);
                    v = input[in_idx];
                }
                tile[th * tile_w + tw] = v;
            }
        }

        __syncthreads();

        // 3) 计算：同一输入tile，对 OC_TILE 个输出通道分别累加
        if (valid_xy) {
            int t_base = threadIdx.y * tile_w + threadIdx.x;
            #pragma unroll
            for (int kh = 0; kh < K; ++kh) {
                int t_row = t_base + kh * tile_w;
                #pragma unroll
                for (int kw = 0; kw < K; ++kw) {
                    float val = tile[t_row + kw];
                    #pragma unroll
                    for (int oc = 0; oc < OC_TILE; ++oc) {
                        float wv = w_sh[oc * (K*K) + kh * K + kw];
                        acc[oc] = fmaf(val, wv, acc[oc]);
                    }
                }
            }
        }

        __syncthreads(); // 保护下一个 ci 的加载
    }

    // 4) 写回输出
    if (valid_xy) {
        int base = (b * C_out) * (H_out * W_out);
        int out_offset = oh * W_out + ow;
        #pragma unroll
        for (int oc = 0; oc < OC_TILE; ++oc) {
            int co = co0 + oc;
            if (co < C_out) {
                int out_idx = base + co * (H_out * W_out) + out_offset;
                output[out_idx] = acc[oc];
            }
        }
    }
}

// C++ wrapper
torch::Tensor conv2d_cuda(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor bias
) {
    CHECK_CUDA(input);
    CHECK_CUDA(weight);
    CHECK_CUDA(bias);
    CHECK_CONTIGUOUS(input);
    CHECK_CONTIGUOUS(weight);
    CHECK_CONTIGUOUS(bias);
    CHECK_FLOAT(input);
    CHECK_FLOAT(weight);
    CHECK_FLOAT(bias);

    int B = input.size(0);
    int C_in = input.size(1);
    int H = input.size(2);
    int W = input.size(3);

    int C_out = weight.size(0);
    int K = weight.size(2);
    TORCH_CHECK(weight.size(3) == K, "Kernel must be square");

    int H_out = H - K + 1;
    int W_out = W - K + 1;

    auto output = torch::empty({B, C_out, H_out, W_out}, input.options());

    // 参数可按GPU微调：32x8, 16x16 等
    const int BLOCK_X = 16;
    const int BLOCK_Y = 16;
    const int OC_TILE = 4;

    dim3 block(BLOCK_X, BLOCK_Y, 1);
    int groups = (C_out + OC_TILE - 1) / OC_TILE;
    dim3 grid((W_out + BLOCK_X - 1) / BLOCK_X,
              (H_out + BLOCK_Y - 1) / BLOCK_Y,
              B * groups);

    size_t tile_w = BLOCK_X + K - 1;
    size_t tile_h = BLOCK_Y + K - 1;
    size_t shmem_elems = tile_w * tile_h + OC_TILE * K * K;
    size_t shmem_bytes = shmem_elems * sizeof(float);

    bool has_bias = bias.numel() > 0;

    conv2d_tiled_kernel_oc<BLOCK_X, BLOCK_Y, OC_TILE><<<grid, block, shmem_bytes>>>(
        input.data_ptr<float>(),
        weight.data_ptr<float>(),
        has_bias ? bias.data_ptr<float>() : nullptr,
        output.data_ptr<float>(),
        B, C_in, C_out, H, W, K, H_out, W_out, has_bias
    );

    auto err = cudaGetLastError();
    TORCH_CHECK(err == cudaSuccess, "conv2d kernel launch failed: ", cudaGetErrorString(err));
    return output;
}
"""

conv2d_cpp_source = r"""
torch::Tensor conv2d_cuda(torch::Tensor input, torch::Tensor weight, torch::Tensor bias);
"""

# Compile with O3 (no fast-math to keep FP32 parity)
conv2d = load_inline(
    name="conv2d_tiled_opt_oc",
    cpp_sources=conv2d_cpp_source,
    cuda_sources=conv2d_source,
    functions=["conv2d_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

class ModelNew(nn.Module):
    def __init__(self, weight, bias=None):
        super(ModelNew, self).__init__()
        self.weight = nn.Parameter(weight)
        self.bias = nn.Parameter(bias) if bias is not None else nn.Parameter(torch.empty(0, device=weight.device, dtype=weight.dtype))
        self.conv2d = conv2d

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.contiguous()
        w = self.weight.contiguous()
        b = self.bias.contiguous() if self.bias is not None else torch.empty(0, device=x.device, dtype=x.dtype)
        return self.conv2d.conv2d_cuda(x, w, b)