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._compile_cuda_kernel()

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

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

        // 使用简单的向量化方法
        __global__ void reglu_vectorized_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int feature_dim, int total_elements) {

            const int tid = threadIdx.x + blockIdx.x * blockDim.x;
            const int stride = blockDim.x * gridDim.x;

            // 每个线程处理4个元素（向量化）
            const int elements_per_thread = 4;
            const int vectorized_elements = total_elements / elements_per_thread;

            // 处理向量化部分
            for (int i = tid; i < vectorized_elements; i += stride) {
                int base_idx = i * elements_per_thread;
                int row = base_idx / (feature_dim / 2);
                int base_col = base_idx % (feature_dim / 2);

                #pragma unroll
                for (int j = 0; j < elements_per_thread; j++) {
                    int col = base_col + j;
                    if (col < feature_dim / 2) {
                        int global_idx = base_idx + j;
                        int gate_offset = row * feature_dim + col;
                        int act_offset = gate_offset + (feature_dim / 2);

                        float gate_val = input[gate_offset];
                        float act_val = input[act_offset];
                        output[global_idx] = fmaxf(0.0f, gate_val) * act_val;
                    }
                }
            }

            // 处理剩余元素
            int remaining_start = vectorized_elements * elements_per_thread;
            for (int i = remaining_start + tid; i < total_elements; i += stride) {
                int row = i / (feature_dim / 2);
                int col = i % (feature_dim / 2);

                float gate_val = input[row * feature_dim + col];
                float act_val = input[row * feature_dim + col + (feature_dim / 2)];
                output[i] = fmaxf(0.0f, gate_val) * act_val;
            }
        }

        // 更稳定的版本 - 不使用向量化
        __global__ void reglu_simple_kernel(
            const float* __restrict__ input,
            float* __restrict__ output,
            int feature_dim, int total_elements) {

            const int tid = threadIdx.x + blockIdx.x * blockDim.x;
            const int stride = blockDim.x * gridDim.x;

            for (int i = tid; i < total_elements; i += stride) {
                int row = i / (feature_dim / 2);
                int col = i % (feature_dim / 2);

                float gate_val = input[row * feature_dim + col];
                float act_val = input[row * feature_dim + col + (feature_dim / 2)];

                // 使用fmaxf代替条件判断，性能更好
                output[i] = fmaxf(0.0f, gate_val) * act_val;
            }
        }

        torch::Tensor reglu_vectorized_parallel(torch::Tensor input) {
            input = input.contiguous();
            auto sizes = input.sizes().vec();
            int feature_dim = sizes.back();
            sizes.back() /= 2;
            auto output = torch::empty(sizes, input.options());

            int total_elements = output.numel();
            int threads = 256;
            int blocks = min((total_elements + threads - 1) / threads, 128); // 限制最大blocks

            // 使用简单稳定的内核
            reglu_simple_kernel<<<blocks, threads>>>(
                input.data_ptr<float>(), output.data_ptr<float>(),
                feature_dim, total_elements);

            cudaError_t err = cudaGetLastError();
            if (err != cudaSuccess) {
                AT_ERROR("CUDA error in reglu_vectorized_parallel: ", cudaGetErrorString(err));
            }

            return output;
        }
        """

        self.op = load_inline(
            name="reglu_vectorized_fixed",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["reglu_vectorized_parallel"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, x):
        return self.op.reglu_vectorized_parallel(x)