import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Multi-head row-wise softmax over CSR edge scores.

    This extends scalar graph edge softmax to multiple attention heads while
    preserving sparse CSR layout. Each row is normalized independently for each
    head.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> torch.Tensor:
        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 = 384
avg_degree = 18
num_heads = 8




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



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

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

__global__ void csr_softmax_warp_kernel(const int* __restrict__ row_ptr,
                                        const float* __restrict__ scores,
                                        float* __restrict__ out,
                                        int num_nodes) {
    int node = blockIdx.x;
    int tid = threadIdx.x;
    int head = tid >> 2;   // 8 heads, four lanes per head
    int lane = tid & 3;

    int start = row_ptr[node];
    int deg = row_ptr[node + 1] - start;
    if (deg <= 0) return;

    constexpr int H = 8;
    float m = -FLT_MAX;
    for (int e = lane; e < deg; e += 4) {
        float v = scores[(start + e) * H + head];
        m = fmaxf(m, v);
    }
    m = fmaxf(m, __shfl_down_sync(0xffffffff, m, 2, 4));
    m = fmaxf(m, __shfl_down_sync(0xffffffff, m, 1, 4));
    m = __shfl_sync(0xffffffff, m, 0, 4);

    float s = 0.0f;
    for (int e = lane; e < deg; e += 4) {
        s += expf(scores[(start + e) * H + head] - m);
    }
    s += __shfl_down_sync(0xffffffff, s, 2, 4);
    s += __shfl_down_sync(0xffffffff, s, 1, 4);
    s = __shfl_sync(0xffffffff, s, 0, 4);

    float inv_s = 1.0f / s;
    for (int e = lane; e < deg; e += 4) {
        float ex = expf(scores[(start + e) * H + head] - m);
        out[(start + e) * H + head] = ex * inv_s;
    }
}

torch::Tensor csr_softmax_forward(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) {
        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
        csr_softmax_warp_kernel<<<num_nodes, 32, 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_softmax_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src,
                   functions=["csr_softmax_forward"], extra_cflags=["-O3"],
                   extra_cuda_cflags=["-O3"], verbose=False)

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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        if row_ptr.dtype is not torch.int32:
            row_ptr = row_ptr.to(torch.int32)
        return _ext.csr_softmax_forward(row_ptr, edge_scores)