# swiglu_cuda.py
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

from swiglu_torch import feature_dim


assert (feature_dim / 2) % 2 == 0, "feature_dim/2 must be a multiple of 2 for float2 vectorization"

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

    def _compile_cuda_kernel(self):
        cpp_source = """
        #include <torch/extension.h>
        #include <ATen/cuda/CUDAContext.h>
        torch::Tensor swiglu_forward_cuda(torch::Tensor input);
        """
        cuda_source = """
        #include <cuda_runtime.h>
        #include <cmath> // For expf


        __global__ void swiglu_fused_vectorized_kernel(
            const float* __restrict__ x,
            float* __restrict__ y,
            int feature_dim,      // 原始输入的特征维度
            int n_elements_out    // 输出张量的元素总数
        ) {
            const float2* x2 = reinterpret_cast<const float2*>(x);
            float2* y2 = reinterpret_cast<float2*>(y);

            int n_work_items = n_elements_out / 2;
            int grid_stride = gridDim.x * blockDim.x;

            for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
                 idx < n_work_items;
                 idx += grid_stride)
            {

                int feature_dim_out_f2 = (feature_dim / 2) / 2;
                int feature_dim_in_f2 = feature_dim / 2;

                int row = idx / feature_dim_out_f2;
                int col_f2 = idx % feature_dim_out_f2;

                int gate_idx_f2 = row * feature_dim_in_f2 + col_f2;
                int act_idx_f2 = gate_idx_f2 + feature_dim_out_f2;

                float2 gate_vec = x2[gate_idx_f2];
                float2 act_vec = x2[act_idx_f2];

                float sigmoid_gate_x = 1.0f / (1.0f + expf(-gate_vec.x));
                float silu_out_x = gate_vec.x * sigmoid_gate_x;

                float sigmoid_gate_y = 1.0f / (1.0f + expf(-gate_vec.y));
                float silu_out_y = gate_vec.y * sigmoid_gate_y;

                y2[idx] = make_float2(silu_out_x * act_vec.x, silu_out_y * act_vec.y);
            }
        }

        torch::Tensor swiglu_forward_cuda(torch::Tensor input) {
            input = input.contiguous();
            TORCH_CHECK(input.size(-1) % 2 == 0, "Last dimension must be even for SwiGLU");

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

            int n_elements_out = output.numel();
            TORCH_CHECK(n_elements_out % 2 == 0, "Output elements must be even for float2 vectorization");

            const int block_size = 256;
            const int n_work_items = n_elements_out / 2;

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

            swiglu_fused_vectorized_kernel<<<grid_size, block_size>>>(
                input.data_ptr<float>(),
                output.data_ptr<float>(),
                feature_dim,
                n_elements_out
            );
            return output;
        }
        """

        self.swiglu_op = load_inline(
            name="swiglu_fused_vectorized_op_fixed", 
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["swiglu_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.swiglu_op.swiglu_forward_cuda(x)