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

from embeddingbag_torch import BATCH_SIZE, EMB_DIM, VOCAB_SIZE # 导入维度常量

N_ELEMENTS = BATCH_SIZE * EMB_DIM

class ModelNew(torch.nn.Module):
  
    def __init__(self, embedding_weight):
        super().__init__()
        self.register_buffer('weight', embedding_weight)
        self._compile_cuda_kernel()

    def _compile_cuda_kernel(self):

        cpp_source = """
        #include <torch/extension.h>
        torch::Tensor embbag_forward_cuda(torch::Tensor input, torch::Tensor offsets, torch::Tensor weight);
        """

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

        #define BLOCK_SIZE 256

        __global__ void embbag_safe_kernel(
            const long* input,        
            const long* offsets,       
            const float* weight,     
            float* output,
            int num_batches, int feature_dim, int max_total_indices
        ) {
            int batch_idx = blockIdx.x;
            if (batch_idx >= num_batches) return;
            
            int feature_idx = threadIdx.x;
            if (feature_idx >= feature_dim) return;

            long start_idx = offsets[batch_idx];
            long end_idx = (batch_idx == num_batches - 1) ? max_total_indices : offsets[batch_idx + 1];
            long seq_len = end_idx - start_idx;

            float final_sum = 0.0f;
            if (seq_len > 0) {
                for (long i = start_idx; i < end_idx; ++i) {
                    long vocab_index = input[i];
                    final_sum += weight[vocab_index * feature_dim + feature_idx];
                }
                output[batch_idx * feature_dim + feature_idx] = final_sum / (float)seq_len;
            } else {
                output[batch_idx * feature_dim + feature_idx] = 0.0f; // Handle empty sequence
            }
        }

        torch::Tensor embbag_forward_cuda(torch::Tensor input, torch::Tensor offsets, torch::Tensor weight) {
            TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor");
            TORCH_CHECK(input.dtype() == torch::kLong, "Input indices must be LongTensor");
            TORCH_CHECK(offsets.dtype() == torch::kLong, "Offsets must be LongTensor");

            input = input.contiguous();
            offsets = offsets.contiguous();
            weight = weight.contiguous();

            int num_batches = offsets.size(0);
            int feature_dim = weight.size(1);
            int max_total_indices = input.size(0);
            
            auto output = torch::empty({num_batches, feature_dim}, input.options().dtype(torch::kFloat32)); 
            
            embbag_safe_kernel<<<num_batches, feature_dim>>>(
                input.data_ptr<long>(),        
                offsets.data_ptr<long>(),     
                weight.data_ptr<float>(),
                output.data_ptr<float>(),
                num_batches, feature_dim, max_total_indices
            );
            
            return output;
        }

        """ 

        self.embbag_op = load_inline(
            name="embbag_fused_op_safe_final_v2",
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=["embbag_forward_cuda"],
            extra_cuda_cflags=["-O3", "--use_fast_math"],
            verbose=True
        )

    def forward(self, input: torch.Tensor, offsets: torch.Tensor) -> torch.Tensor:
        return self.embbag_op.embbag_forward_cuda(input, offsets, self.weight)