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

class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self.gamma = 2.2
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor gamma_correction_cuda(torch::Tensor input, float gamma);
        """
        
        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath>

        #define BLOCK_SIZE 256

        __global__ void gamma_f4_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int n_vecs,
            float gamma
        ) {
            int idx = blockIdx.x * blockDim.x + threadIdx.x;
            if (idx >= n_vecs) return;

         
            float4 v = reinterpret_cast<const float4*>(input)[idx];
            float4 out;

   
            out.x = powf(v.x, gamma);
            out.y = powf(v.y, gamma);
            out.z = powf(v.z, gamma);
            out.w = powf(v.w, gamma);

            reinterpret_cast<float4*>(output)[idx] = out;
        }

        torch::Tensor gamma_correction_cuda(torch::Tensor input, float gamma) {
            auto output = torch::empty_like(input);
            int numel = input.numel();
            

            if (numel % 4 != 0) return output;

            int n_vecs = numel / 4;
            const int grid_size = (n_vecs + BLOCK_SIZE - 1) / BLOCK_SIZE;

            gamma_f4_kernel<<<grid_size, BLOCK_SIZE>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                n_vecs,
                gamma
            );
            
            return output;
        }
        """

        self.op = load_inline(
            name="gamma_correction_f4_v1",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["gamma_correction_cuda"],
            extra_cuda_cflags=["-O3"], 
            verbose=False
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 确保连续内存
        if not x.is_contiguous(): x = x.contiguous()
        return self.op.gamma_correction_cuda(x, self.gamma)