import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, followed by two max pooling layers and a sum operation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True) 
        return x

batch_size = 16
in_channels = 32
out_channels = 64
depth, height, width = 32, 32, 32
kernel_size = 5
stride = 2
padding = 2

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding]

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

cpp_source = r'''
#include <torch/extension.h>
#include <ATen/ATen.h>
void pool8_launcher(const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,const torch::Tensor&,torch::Tensor&);

torch::Tensor fused_forward(torch::Tensor x, torch::Tensor weight, torch::Tensor bias) {
    auto opts = torch::TensorOptions().dtype(torch::kLong).device(weight.device());
    auto idx0 = torch::tensor({4,2,0}, opts);
    auto idx1 = torch::tensor({3,1}, opts);
    auto wp = weight.permute({1,0,2,3,4}).contiguous();
    c10::optional<torch::Tensor> nb = c10::nullopt;
    torch::Tensor ys[8];
    for(int pz=0; pz<2; ++pz) for(int py=0; py<2; ++py) for(int px=0; px<2; ++px) {
        auto wz = at::index_select(wp, 2, pz ? idx1 : idx0);
        auto wzy = at::index_select(wz, 3, py ? idx1 : idx0);
        auto wzyx = at::index_select(wzy, 4, px ? idx1 : idx0).contiguous();
        std::vector<int64_t> pad = {pz ? 0 : 1, py ? 0 : 1, px ? 0 : 1};
        ys[pz*4 + py*2 + px] = at::conv3d(x, wzyx, nb, {1,1,1}, pad, {1,1,1}, 1);
    }
    auto out = torch::empty({x.size(0), 1, 10, 10, 10}, x.options());
    pool8_launcher(ys[0],ys[1],ys[2],ys[3],ys[4],ys[5],ys[6],ys[7],bias,out);
    return out;
}
'''

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

__device__ __forceinline__ void scan3(const float* base, int H, int W, float &m) {
    #pragma unroll
    for(int dz=0; dz<3; ++dz) {
        #pragma unroll
        for(int dy=0; dy<3; ++dy) {
            const float* q = base + (dz*H + dy)*W;
            #pragma unroll
            for(int dx=0; dx<3; ++dx) m = fmaxf(m, q[dx]);
        }
    }
}

__global__ void pool8_kernel(const float* y0,const float* y1,const float* y2,const float* y3,const float* y4,const float* y5,const float* y6,const float* y7,const float* bias,float* out,int N) {
    int idx = blockIdx.x;
    int ox = idx % 10;
    int oy = (idx / 10) % 10;
    int oz = (idx / 100) % 10;
    int n = idx / 1000;
    int c = threadIdx.x;
    __shared__ float vals[64];
    float m = -FLT_MAX;
    const int C = 64;
    const int z = oz * 3, y = oy * 3, x = ox * 3;

    scan3(y0 + (((n*C+c)*32 + z)*32 + y)*32 + x, 32, 32, m);
    scan3(y1 + (((n*C+c)*32 + z)*32 + y)*31 + x, 32, 31, m);
    scan3(y2 + (((n*C+c)*32 + z)*31 + y)*32 + x, 31, 32, m);
    scan3(y3 + (((n*C+c)*32 + z)*31 + y)*31 + x, 31, 31, m);
    scan3(y4 + (((n*C+c)*31 + z)*32 + y)*32 + x, 32, 32, m);
    scan3(y5 + (((n*C+c)*31 + z)*32 + y)*31 + x, 32, 31, m);
    scan3(y6 + (((n*C+c)*31 + z)*31 + y)*32 + x, 31, 32, m);
    scan3(y7 + (((n*C+c)*31 + z)*31 + y)*31 + x, 31, 31, m);

    vals[c] = m + bias[c];
    __syncthreads();
    if(c < 32) vals[c] += vals[c+32];
    __syncthreads();
    if(c < 16) vals[c] += vals[c+16];
    __syncthreads();
    if(c < 8) vals[c] += vals[c+8];
    __syncthreads();
    if(c < 4) vals[c] += vals[c+4];
    __syncthreads();
    if(c < 2) vals[c] += vals[c+2];
    __syncthreads();
    if(c == 0) out[idx] = vals[0] + vals[1];
}

void pool8_launcher(const torch::Tensor& y0,const torch::Tensor& y1,const torch::Tensor& y2,const torch::Tensor& y3,const torch::Tensor& y4,const torch::Tensor& y5,const torch::Tensor& y6,const torch::Tensor& y7,const torch::Tensor& bias,torch::Tensor& out) {
    int N = y0.size(0);
    pool8_kernel<<<N*1000, 64, 0, at::cuda::getCurrentCUDAStream()>>>(y0.data_ptr<float>(),y1.data_ptr<float>(),y2.data_ptr<float>(),y3.data_ptr<float>(),y4.data_ptr<float>(),y5.data_ptr<float>(),y6.data_ptr<float>(),y7.data_ptr<float>(),bias.data_ptr<float>(),out.data_ptr<float>(),N);
}
'''

_ext = load_inline(name='ct3d_parity_conv_ext_submit', cpp_sources=cpp_source, cuda_sources=cuda_source, functions=['fused_forward'], extra_cflags=['-O3'], extra_cuda_cflags=['-O3'], verbose=False)

class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
    def forward(self, x):
        return _ext.fused_forward(x, self.conv_transpose.weight, self.conv_transpose.bias)