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

cpp_source = """
#include <torch/extension.h>

torch::Tensor vecdot_cuda_forward(const torch::Tensor& x, const torch::Tensor& y, int64_t dim);
"""

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

#define BLOCK_SIZE 512

template <typename T>
__global__ void vecdot_kernel(
    T* output,
    const T* x,
    const T* y,
    const int n_along_dim,
    const int stride_between_slices,
    const int n_slices)
{{
    const int slice_idx = blockIdx.x;
    if (slice_idx >= n_slices) return;

    __shared__ T sdata[BLOCK_SIZE];
    
    const int tid = threadIdx.x;
    T my_sum = 0.0; // Use double literal for double precision

    const T* x_slice = x + slice_idx * stride_between_slices;
    const T* y_slice = y + slice_idx * stride_between_slices;

    for (int i = tid; i < n_along_dim; i += blockDim.x) {{
        my_sum += x_slice[i] * y_slice[i];
    }}

    sdata[tid] = my_sum;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {{
        if (tid < s) {{
            sdata[tid] += sdata[tid + s];
        }}
        __syncthreads();
    }}
    
    if (tid == 0) {{
        output[slice_idx] = sdata[0];
    }}
}}

torch::Tensor vecdot_cuda_forward(const torch::Tensor& x, const torch::Tensor& y, int64_t dim) {
    auto sizes = x.sizes();
    int ndim = x.dim();
    if (dim < 0) dim += ndim;

    TORCH_CHECK(x.is_cuda() && y.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(x.sizes() == y.sizes(), "Input shapes must match");
    TORCH_CHECK(x.is_contiguous() && y.is_contiguous(), "Inputs must be contiguous");
    TORCH_CHECK(dim == ndim - 1, "This custom kernel only supports vecdot on the last dimension.");

    const int n_along_dim = sizes[dim];
    const int64_t n_slices = x.numel() / n_along_dim;
    const int stride_between_slices = n_along_dim;
    
    auto out_sizes = sizes.vec();
    out_sizes.erase(out_sizes.begin() + dim);
    auto output = torch::empty(out_sizes, x.options());
    
    const dim3 grid(n_slices);
    const dim3 block(BLOCK_SIZE);

    // AT_DISPATCH_FLOATING_TYPES will handle both float and double
    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "vecdot_kernel", ([&] {{
        vecdot_kernel<scalar_t><<<grid, block>>>(
            output.data_ptr<scalar_t>(),
            x.data_ptr<scalar_t>(),
            y.data_ptr<scalar_t>(),
            n_along_dim,
            stride_between_slices,
            n_slices
        );
    }}));

    return output;
}
"""

class ModelNew(nn.Module):
    def __init__(self, dim):
        super(ModelNew, self).__init__()
        self.dim = dim
        self.op = load_inline(
            name='vecdot_op_double', 
            cpp_sources=cpp_source,
            cuda_sources=cuda_source,
            functions=['vecdot_cuda_forward'],
            verbose=False
        )

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return self.op.vecdot_cuda_forward(x, y, self.dim)