import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Row-wise softmax over edge scores stored in CSR format.

    This is a core primitive in graph attention networks and sparse attention on
    irregular neighborhoods. Each node owns a contiguous edge segment in
    `edge_scores[row_ptr[i]:row_ptr[i + 1]]`, and the kernel normalizes scores
    independently within each segment.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            row_ptr:      (num_nodes + 1,) CSR row offsets
            edge_scores:  (num_edges,) unnormalized attention logits
        Returns:
            (num_edges,) softmax-normalized edge weights per source node
        """
        num_nodes = row_ptr.numel() - 1
        out = torch.empty_like(edge_scores)

        for node in range(num_nodes):
            start = int(row_ptr[node].item())
            end = int(row_ptr[node + 1].item())
            if end > start:
                out[start:end] = torch.softmax(edge_scores[start:end], dim=0)

        return out


num_nodes = 512
avg_degree = 24
batchless_num_edges = num_nodes * avg_degree


def _make_row_ptr(num_nodes: int, avg_degree: int) -> torch.Tensor:
    base = torch.full((num_nodes,), avg_degree, dtype=torch.int32)
    return row_ptr


def get_inputs():
    return list(popcorn_pri.make_csr_scalar_edge_scores(num_nodes, avg_degree))



def get_init_inputs():
    return []

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

_cpp_src = r"""
#include <torch/extension.h>

// CUDA launcher implemented in cuda source
torch::Tensor csr_softmax_forward_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores);

torch::Tensor csr_softmax_forward(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    return csr_softmax_forward_cuda(row_ptr, edge_scores);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

__device__ __forceinline__ float warp_reduce_max(float v) {
    unsigned mask = 0xffffffffu;
    v = fmaxf(v, __shfl_down_sync(mask, v, 16));
    v = fmaxf(v, __shfl_down_sync(mask, v, 8));
    v = fmaxf(v, __shfl_down_sync(mask, v, 4));
    v = fmaxf(v, __shfl_down_sync(mask, v, 2));
    v = fmaxf(v, __shfl_down_sync(mask, v, 1));
    return v;
}

__device__ __forceinline__ float warp_reduce_sum(float v) {
    unsigned mask = 0xffffffffu;
    v += __shfl_down_sync(mask, v, 16);
    v += __shfl_down_sync(mask, v, 8);
    v += __shfl_down_sync(mask, v, 4);
    v += __shfl_down_sync(mask, v, 2);
    v += __shfl_down_sync(mask, v, 1);
    return v;
}

template<int WARPS_PER_BLOCK>
__global__ void csr_softmax_warp_kernel(const int* __restrict__ row_ptr,
                                        const float* __restrict__ scores,
                                        float* __restrict__ out,
                                        int num_nodes) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    int row = blockIdx.x * WARPS_PER_BLOCK + warp;
    if (row >= num_nodes) return;

    int start = row_ptr[row];
    int end = row_ptr[row + 1];
    int n = end - start;
    if (n <= 0) return;

    float m = -FLT_MAX;
    for (int i = lane; i < n; i += 32) {
        m = fmaxf(m, scores[start + i]);
    }
    m = warp_reduce_max(m);

    float s = 0.0f;
    for (int i = lane; i < n; i += 32) {
        s += __expf(scores[start + i] - m);
    }
    s = warp_reduce_sum(s);
    float inv_s = 1.0f / s;

    for (int i = lane; i < n; i += 32) {
        out[start + i] = __expf(scores[start + i] - m) * inv_s;
    }
}

torch::Tensor csr_softmax_forward_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    auto out = torch::empty_like(edge_scores);
    int num_nodes = (int)row_ptr.numel() - 1;
    if (num_nodes <= 0 || edge_scores.numel() == 0) return out;

    const int warps = 4;
    dim3 block(32 * warps);
    dim3 grid((num_nodes + warps - 1) / warps);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    csr_softmax_warp_kernel<warps><<<grid, block, 0, stream>>>(
        row_ptr.data_ptr<int>(), edge_scores.data_ptr<float>(), out.data_ptr<float>(), num_nodes);
    return out;
}
"""

_ext = load_inline(
    name="csr_row_softmax_ext_v1",
    cpp_sources=[_cpp_src],
    cuda_sources=[_cuda_src],
    functions=["csr_softmax_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        return _ext.csr_softmax_forward(row_ptr, edge_scores)