import torch
from torch.utils.cpp_extension import load_inline

has_element_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void has_element_vec4_kernel(
    const int* __restrict__ indices, 
    float* __restrict__ output, 
    int n) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int vec_idx = idx * 4; 

    // 1. Vectorized Path
    if (vec_idx + 4 <= n) {
        // Load int4
        const int4* in_ptr = reinterpret_cast<const int4*>(indices);
        int4 in_val = in_ptr[idx];
        
        float4 out_val;
    
        out_val.x = (in_val.x != -1) ? 1.0f : 0.0f;
        out_val.y = (in_val.y != -1) ? 1.0f : 0.0f;
        out_val.z = (in_val.z != -1) ? 1.0f : 0.0f;
        out_val.w = (in_val.w != -1) ? 1.0f : 0.0f;
  
        float4* out_ptr = reinterpret_cast<float4*>(output);
        out_ptr[idx] = out_val;
    }
}

__global__ void has_element_scalar_kernel(
    const int* __restrict__ indices, 
    float* __restrict__ output, 
    int n,
    int offset) 
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
    
    if (idx < n) {
        int val = indices[idx];
        output[idx] = (val != -1) ? 1.0f : 0.0f;
    }
}

torch::Tensor has_element_cuda(torch::Tensor indices) {
    int n = indices.numel();
 
    auto output = torch::empty_like(indices, torch::dtype(torch::kFloat32));
    

    int vec_elements = n / 4;
    if (vec_elements > 0) {
        const int block = 256;
        const int grid = (vec_elements + block - 1) / block;
        
        has_element_vec4_kernel<<<grid, block>>>(
            indices.data_ptr<int>(),
            output.data_ptr<float>(),
            n
        );
    }
    

    int remainder = n % 4;
    if (remainder > 0) {
        int offset = vec_elements * 4;
        has_element_scalar_kernel<<<1, 32>>>(
            indices.data_ptr<int>(),
            output.data_ptr<float>(),
            n,
            offset
        );
    }
    
    return output;
}
"""

cpp_source = "torch::Tensor has_element_cuda(torch::Tensor indices);"

has_element_module = load_inline(
    name="optional_has_element_extension",
    cpp_sources=cpp_source,
    cuda_sources=has_element_source,
    functions=["has_element_cuda"],
    verbose=True,
    with_cuda=True
)

class ModelNew(torch.nn.Module):
    def __init__(self):
        super(ModelNew, self).__init__()
        self.cuda_op = has_element_module

    def forward(self, indices):

        return self.cuda_op.has_element_cuda(indices.int().contiguous())