import torch
from torch.utils.cpp_extension import load_inline

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

// 融合L1Loss内核 - 单步完成所有计算
__global__ void l1_fused_kernel(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ loss_output,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 局部累加器
    float local_l1 = 0.0f;
    
    // 每个线程处理多个元素 - 融合diff+abs计算
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < size; i += stride) {
        float diff = input[i] - target[i];  // diff计算
        local_l1 += fabsf(diff);            // abs计算 + 累加
    }
    
    // 使用共享内存进行块内归约 - 融合sum计算
    extern __shared__ float shared_mem[];
    shared_mem[threadIdx.x] = local_l1;
    
    __syncthreads();
    
    // 块内归约
    for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
        if (threadIdx.x < stride) {
            shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
        }
        __syncthreads();
    }
    
    // 每个block写入部分和到全局内存
    if (threadIdx.x == 0) {
        atomicAdd(loss_output, shared_mem[0]);
    }
}

// 高级融合版本 - 多种优化策略
__global__ void l1_fused_advanced_kernel(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ loss_output,
    int size,
    bool use_warp_reduce
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int warp_id = threadIdx.x / 32;
    int lane_id = threadIdx.x % 32;
    
    // 局部累加器
    float local_l1 = 0.0f;
    
    // 每个线程处理多个元素
    int stride = blockDim.x * gridDim.x;
    for (int i = idx; i < size; i += stride) {
        float diff = input[i] - target[i];
        local_l1 += fabsf(diff);
    }
    
    if (use_warp_reduce) {
        // 使用warp级归约 - 更高效的融合
        for (int offset = 16; offset > 0; offset /= 2) {
            local_l1 += __shfl_down_sync(0xffffffff, local_l1, offset);
        }
        
        // 每个warp的第一个线程做原子操作
        if (lane_id == 0) {
            atomicAdd(loss_output, local_l1);
        }
    } else {
        // 使用共享内存归约
        extern __shared__ float shared_mem[];
        shared_mem[threadIdx.x] = local_l1;
        
        __syncthreads();
        
        // 块内归约
        for (int stride = blockDim.x / 2; stride > 0; stride /= 2) {
            if (threadIdx.x < stride) {
                shared_mem[threadIdx.x] += shared_mem[threadIdx.x + stride];
            }
            __syncthreads();
        }
        
        if (threadIdx.x == 0) {
            atomicAdd(loss_output, shared_mem[0]);
        }
    }
}

// 向量化融合版本 - 处理多个元素
__global__ void l1_fused_vectorized_kernel(
    const float* __restrict__ input,
    const float* __restrict__ target,
    float* __restrict__ loss_output,
    int size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // 向量化处理 - 每次处理4个元素
    int vectorized_size = (size / 4) * 4;
    int stride = blockDim.x * gridDim.x;
    
    float local_l1 = 0.0f;
    
    // 向量化主循环
    for (int i = idx * 4; i < vectorized_size; i += stride * 4) {
        // 加载4个元素
        float4 input_vec = *reinterpret_cast<const float4*>(&input[i]);
        float4 target_vec = *reinterpret_cast<const float4*>(&target[i]);
        
        // 计算L1距离
        local_l1 += fabsf(input_vec.x - target_vec.x);
        local_l1 += fabsf(input_vec.y - target_vec.y);
        local_l1 += fabsf(input_vec.z - target_vec.z);
        local_l1 += fabsf(input_vec.w - target_vec.w);
    }
    
    // 处理剩余元素
    for (int i = idx + vectorized_size; i < size; i += stride) {
        float diff = input[i] - target[i];
        local_l1 += fabsf(diff);
    }
    
    // Warp级归约
    for (int offset = 16; offset > 0; offset /= 2) {
        local_l1 += __shfl_down_sync(0xffffffff, local_l1, offset);
    }
    
    if (threadIdx.x % 32 == 0) {
        atomicAdd(loss_output, local_l1);
    }
}

// 主函数：融合L1Loss
torch::Tensor l1_fused_cuda(torch::Tensor input, torch::Tensor target, std::string mode = "fused") {
    TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32");
    TORCH_CHECK(target.scalar_type() == torch::kFloat32, "Target must be float32");
    TORCH_CHECK(input.sizes() == target.sizes(), "Input and target must have same shape");
    
    auto input_contig = input.contiguous();
    auto target_contig = target.contiguous();
    int size = input_contig.numel();
    
    // 创建输出张量
    auto l1_loss = torch::zeros(1, input.options());
    
    const int block_size = 256;
    int num_blocks = min(1024, (size + block_size - 1) / block_size);
    
    if (mode == "fused") {
        // 基础融合版本
        size_t shared_mem = block_size * sizeof(float);
        l1_fused_kernel<<<num_blocks, block_size, shared_mem>>>(
            input_contig.data_ptr<float>(),
            target_contig.data_ptr<float>(),
            l1_loss.data_ptr<float>(),
            size
        );
    } else if (mode == "advanced") {
        // 高级融合版本 - warp优化
        l1_fused_advanced_kernel<<<num_blocks, block_size>>>(
            input_contig.data_ptr<float>(),
            target_contig.data_ptr<float>(),
            l1_loss.data_ptr<float>(),
            size,
            true
        );
    } else if (mode == "vectorized") {
        // 向量化融合版本
        // 确保size是4的倍数
        int aligned_size = ((size + 3) / 4) * 4;
        int vectorized_num_blocks = min(1024, (aligned_size / 4 + block_size - 1) / block_size);
        
        l1_fused_vectorized_kernel<<<vectorized_num_blocks, block_size>>>(
            input_contig.data_ptr<float>(),
            target_contig.data_ptr<float>(),
            l1_loss.data_ptr<float>(),
            size
        );
    }
    
    return l1_loss;
}
"""

l1_fused_cpp_source = """
torch::Tensor l1_fused_cuda(torch::Tensor input, torch::Tensor target, std::string mode);
"""

# 编译融合CUDA代码
l1_fused = load_inline(
    name="l1_fused",
    cpp_sources=l1_fused_cpp_source,
    cuda_sources=l1_fused_source,
    functions=["l1_fused_cuda"],
    extra_cuda_cflags=[
        "-O3", 
        "--use_fast_math",
        "-gencode=arch=compute_80,code=sm_80"
    ],
    verbose=True
)

class ModelNew(torch.nn.Module):
    def __init__(self, mode="fused"):
        super(ModelNew, self).__init__()
        self.mode = mode
        self.l1_fused = l1_fused

    def forward(self, input, target):
        return self.l1_fused.l1_fused_cuda(input, target, self.mode)
