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

# C++ 源代码
cpp_source = """
#include <torch/extension.h>

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

# CUDA 源代码
cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

template <typename T>
__global__ void cross_kernel_stride(
    T* output,
    const T* x,
    const T* y,
    const int64_t num_vectors,
    const int64_t dim_stride,
    const int64_t slice_stride)
{{
    // 每个线程负责一个完整的叉积计算
    const int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= num_vectors) return;

    // --- 使用 stride 计算基地址 ---
    const T* x_vec = x + idx * slice_stride;
    const T* y_vec = y + idx * slice_stride;
    T* out_vec = output + idx * slice_stride;

    // --- 使用 dim_stride 加载数据到寄存器 ---
    const T x0 = x_vec[0];
    const T x1 = x_vec[dim_stride];
    const T x2 = x_vec[2 * dim_stride];
    const T y0 = y_vec[0];
    const T y1 = y_vec[dim_stride];
    const T y2 = y_vec[2 * dim_stride];

    out_vec[0] = x1 * y2 - x2 * y1;
    out_vec[dim_stride] = x2 * y0 - x0 * y2;
    out_vec[2 * dim_stride] = x0 * y1 - x1 * y0;
}}

torch::Tensor cross_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 for this implementation");
    TORCH_CHECK(sizes[dim] == 3, "Dimension for cross product must be 3");

    // Stride 计算
    const int64_t dim_stride = x.stride(dim);
    const int64_t num_vectors = x.numel() / 3;

    // 创建一个临时视图来计算 slice_stride
    auto x_flat = x.transpose(dim, -1).flatten(0, -2);
    const int64_t slice_stride = x_flat.stride(0);
    
    auto output = torch::empty_like(x);
    
    const int block_size = 256;
    const int num_blocks = (num_vectors + block_size - 1) / block_size;

    AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "cross_kernel_stride", ([&] {{
        cross_kernel_stride<scalar_t><<<num_blocks, block_size>>>(
            output.data_ptr<scalar_t>(),
            x.data_ptr<scalar_t>(),
            y.data_ptr<scalar_t>(),
            num_vectors,
            dim_stride,
            slice_stride
        );
    }}));

    return output;
}
"""

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

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