# upsample_cuda.py
import torch
from torch.utils.cpp_extension import load_inline

from upsample_torch import BATCH_SIZE, CHANNELS, H_IN, W_IN, SCALE_FACTOR # 导入维度常量


W_OUT = W_IN * SCALE_FACTOR
assert W_OUT % 4 == 0, "Output width (W_in * scale) must be a multiple of 4 for float4 vectorization"
VEC_SIZE = 4 

class ModelNew(torch.nn.Module):
  
    def __init__(self):
        super().__init__()
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor upsample_forward_cuda(torch::Tensor input);
        """

        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE 256
        #define SCALE_VAL {scale_factor}
        #define VEC_SIZE {vec_size}

        __global__ void upsample_fused_vectorized_kernel(
            const float* __restrict__ x,
            float* __restrict__ y,
            int N, int C, int H_in, int W_in
        ) {{

            int H_out = H_in * SCALE_VAL;
            int W_out = W_in * SCALE_VAL;
            int W_out_vec = W_out / VEC_SIZE; 
            int CHW_out_vec = C * H_out * W_out_vec;
            int C_HW_in = C * H_in * W_in; 
            
            int output_elements_vec = N * CHW_out_vec;
    
            int grid_stride = gridDim.x * blockDim.x;

            for (int idx_vec = blockIdx.x * blockDim.x + threadIdx.x;
                 idx_vec < output_elements_vec;
                 idx_vec += grid_stride)
            {{
              
                int n = idx_vec / CHW_out_vec;
                int rem_n = idx_vec % CHW_out_vec;
                
                int c = rem_n / (H_out * W_out_vec);
                int rem_c = rem_n % (H_out * W_out_vec);
                
                int h_out = rem_c / W_out_vec;
                int w_out_vec = rem_c % W_out_vec;
     
       
                int h_in = h_out / SCALE_VAL;

                long base_input_idx = (long)n * C_HW_in + 
                                     (long)c * (H_in * W_in) + 
                                     (long)h_in * W_in;
                                      
            
                float4* y4_ptr = reinterpret_cast<float4*>(y);

                int w_out_start = w_out_vec * VEC_SIZE; 

                int w_in_start = w_out_start / SCALE_VAL; 
       
                
                float val_in_0 = x[base_input_idx + w_in_start];
                float val_in_1 = x[base_input_idx + w_in_start + 1];

                float4 output_vec;

                output_vec.x = val_in_0;
                output_vec.y = val_in_0; 
                output_vec.z = val_in_1;
                output_vec.w = val_in_1;

                y4_ptr[idx_vec] = output_vec;
            }}
        }}

        torch::Tensor upsample_forward_cuda(torch::Tensor input) {{
            TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
            input = input.contiguous();

            int N = input.size(0);
            int C = input.size(1);
            int H_in = input.size(2);
            int W_in = input.size(3);
            int up_factor = SCALE_VAL;

            int H_out = H_in * up_factor;
            int W_out = W_in * up_factor;
 
            TORCH_CHECK(W_out % VEC_SIZE == 0, "Output width must be divisible by 4 for float4 vectorization");

            auto output = torch::empty({{N, C, H_out, W_out}}, input.options()); 
            
            int n_elements_vec = output.numel() / VEC_SIZE;

            const int block_size = 256;
            const int grid_size = (n_elements_vec + block_size - 1) / block_size; 

            upsample_fused_vectorized_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                N, C, H_in, W_in
            );
            
            return output;
        }}
        """.format(scale_factor=SCALE_FACTOR, vec_size=VEC_SIZE)
        
        self.ps_op = load_inline(
            name="upsample_fused_vectorized_op_final",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["upsample_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.ps_op.upsample_forward_cuda(input)