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>

torch::Tensor csr_softmax_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_cuda(row_ptr, edge_scores);
}
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAStream.h>
#include <cfloat>
#include <cstdint>

template <typename T>
__device__ __forceinline__ int to_int_offset(T v) { return (int)v; }

__inline__ __device__ float warp_reduce_max(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, v, offset);
        v = fmaxf(v, other);
    }
    return v;
}

__inline__ __device__ float warp_reduce_sum(float v) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        v += __shfl_down_sync(0xffffffff, v, offset);
    }
    return v;
}

template <typename PtrT>
__global__ __launch_bounds__(128, 4)
void csr_softmax_kernel(const PtrT* __restrict__ row_ptr,
                        const float* __restrict__ x,
                        float* __restrict__ out,
                        int num_nodes) {
    int tid = threadIdx.x;
    int lane = tid & 31;
    int warp = tid >> 5;
    int row = blockIdx.x * 4 + warp;
    if (row >= num_nodes) return;

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

    if (len <= 32) {
        float v = (lane < len) ? x[start + lane] : -FLT_MAX;
        float m = warp_reduce_max(v);
        m = __shfl_sync(0xffffffff, m, 0);
        float e = (lane < len) ? __expf(v - m) : 0.0f;
        float s = warp_reduce_sum(e);
        s = __shfl_sync(0xffffffff, s, 0);
        if (lane < len) out[start + lane] = e / s;
    } else {
        float m = -FLT_MAX;
        for (int i = lane; i < len; i += 32) {
            m = fmaxf(m, x[start + i]);
        }
        m = warp_reduce_max(m);
        m = __shfl_sync(0xffffffff, m, 0);

        float s = 0.0f;
        for (int i = lane; i < len; i += 32) {
            s += __expf(x[start + i] - m);
        }
        s = warp_reduce_sum(s);
        s = __shfl_sync(0xffffffff, s, 0);
        float inv_s = 1.0f / s;
        for (int i = lane; i < len; i += 32) {
            out[start + i] = __expf(x[start + i] - m) * inv_s;
        }
    }
}

template <typename PtrT>
void launch_kernel(torch::Tensor row_ptr, torch::Tensor edge_scores, torch::Tensor out, int num_nodes) {
    dim3 block(128);
    dim3 grid((num_nodes + 3) / 4);
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    csr_softmax_kernel<PtrT><<<grid, block, 0, stream>>>(
        row_ptr.data_ptr<PtrT>(), edge_scores.data_ptr<float>(), out.data_ptr<float>(), num_nodes);
}

torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    auto out = torch::empty_like(edge_scores);
    int64_t num_nodes64 = row_ptr.numel() - 1;
    if (num_nodes64 <= 0 || edge_scores.numel() == 0) {
        return out;
    }
    int num_nodes = static_cast<int>(num_nodes64);
    auto dtype = row_ptr.scalar_type();
    if (dtype == at::kInt) {
        launch_kernel<int>(row_ptr, edge_scores, out, num_nodes);
    } else if (dtype == at::kLong) {
        launch_kernel<int64_t>(row_ptr, edge_scores, out, num_nodes);
    } else if (dtype == at::kFloat) {
        launch_kernel<float>(row_ptr, edge_scores, out, num_nodes);
    } else if (dtype == at::kDouble) {
        launch_kernel<double>(row_ptr, edge_scores, out, num_nodes);
    } else {
        TORCH_CHECK(false, "unsupported row_ptr dtype");
    }
    return out;
}
"""

_ext = load_inline(
    name="csr_row_softmax_ext_v3",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["csr_softmax_forward"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    with_cuda=True,
    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)