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_source = r"""
#include <torch/extension.h>
torch::Tensor csr_softmax_forward(torch::Tensor row_ptr, torch::Tensor edge_scores);
"""

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

template <typename RowT>
__global__ void csr_softmax_8h_warprow(const RowT* __restrict__ row_ptr,
                                       const float* __restrict__ x,
                                       float* __restrict__ out,
                                       int num_nodes) {
    int lane = threadIdx.x;
    int row = blockIdx.x * blockDim.y + threadIdx.y;
    if (row >= num_nodes) return;
    int start = (int)row_ptr[row];
    int end = (int)row_ptr[row + 1];
    int deg = end - start;
    int base = (start + lane) << 3;
    float v0=-FLT_MAX,v1=-FLT_MAX,v2=-FLT_MAX,v3=-FLT_MAX,v4=-FLT_MAX,v5=-FLT_MAX,v6=-FLT_MAX,v7=-FLT_MAX;
    if (lane < deg) {
        v0=x[base]; v1=x[base+1]; v2=x[base+2]; v3=x[base+3];
        v4=x[base+4]; v5=x[base+5]; v6=x[base+6]; v7=x[base+7];
    }
    float m0=v0,m1=v1,m2=v2,m3=v3,m4=v4,m5=v5,m6=v6,m7=v7;
    #pragma unroll
    for(int off=16; off>0; off>>=1){
        m0=fmaxf(m0,__shfl_down_sync(0xffffffff,m0,off)); m1=fmaxf(m1,__shfl_down_sync(0xffffffff,m1,off));
        m2=fmaxf(m2,__shfl_down_sync(0xffffffff,m2,off)); m3=fmaxf(m3,__shfl_down_sync(0xffffffff,m3,off));
        m4=fmaxf(m4,__shfl_down_sync(0xffffffff,m4,off)); m5=fmaxf(m5,__shfl_down_sync(0xffffffff,m5,off));
        m6=fmaxf(m6,__shfl_down_sync(0xffffffff,m6,off)); m7=fmaxf(m7,__shfl_down_sync(0xffffffff,m7,off));
    }
    m0=__shfl_sync(0xffffffff,m0,0); m1=__shfl_sync(0xffffffff,m1,0); m2=__shfl_sync(0xffffffff,m2,0); m3=__shfl_sync(0xffffffff,m3,0);
    m4=__shfl_sync(0xffffffff,m4,0); m5=__shfl_sync(0xffffffff,m5,0); m6=__shfl_sync(0xffffffff,m6,0); m7=__shfl_sync(0xffffffff,m7,0);
    float e0=0,e1=0,e2=0,e3=0,e4=0,e5=0,e6=0,e7=0;
    if (lane < deg) {
        e0=__expf(v0-m0); e1=__expf(v1-m1); e2=__expf(v2-m2); e3=__expf(v3-m3);
        e4=__expf(v4-m4); e5=__expf(v5-m5); e6=__expf(v6-m6); e7=__expf(v7-m7);
    }
    float s0=e0,s1=e1,s2=e2,s3=e3,s4=e4,s5=e5,s6=e6,s7=e7;
    #pragma unroll
    for(int off=16; off>0; off>>=1){
        s0+=__shfl_down_sync(0xffffffff,s0,off); s1+=__shfl_down_sync(0xffffffff,s1,off);
        s2+=__shfl_down_sync(0xffffffff,s2,off); s3+=__shfl_down_sync(0xffffffff,s3,off);
        s4+=__shfl_down_sync(0xffffffff,s4,off); s5+=__shfl_down_sync(0xffffffff,s5,off);
        s6+=__shfl_down_sync(0xffffffff,s6,off); s7+=__shfl_down_sync(0xffffffff,s7,off);
    }
    s0=__shfl_sync(0xffffffff,s0,0); s1=__shfl_sync(0xffffffff,s1,0); s2=__shfl_sync(0xffffffff,s2,0); s3=__shfl_sync(0xffffffff,s3,0);
    s4=__shfl_sync(0xffffffff,s4,0); s5=__shfl_sync(0xffffffff,s5,0); s6=__shfl_sync(0xffffffff,s6,0); s7=__shfl_sync(0xffffffff,s7,0);
    if (lane < deg) {
        out[base]=e0/s0; out[base+1]=e1/s1; out[base+2]=e2/s2; out[base+3]=e3/s3;
        out[base+4]=e4/s4; out[base+5]=e5/s5; out[base+6]=e6/s6; out[base+7]=e7/s7;
    }
}

template <typename RowT>
void launch8(torch::Tensor row_ptr, torch::Tensor edge_scores, torch::Tensor out, int num_nodes){
    dim3 block(32, 2, 1);
    dim3 grid((num_nodes + 1) / 2);
    csr_softmax_8h_warprow<RowT><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(row_ptr.data_ptr<RowT>(), edge_scores.data_ptr<float>(), out.data_ptr<float>(), num_nodes);
}

torch::Tensor csr_softmax_forward(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    int num_nodes = row_ptr.numel() - 1;
    auto out = torch::empty_like(edge_scores);
    if (num_nodes <= 0 || edge_scores.numel() == 0) return out;
    if (row_ptr.scalar_type() == at::kInt) launch8<int>(row_ptr, edge_scores, out, num_nodes);
    else if (row_ptr.scalar_type() == at::kLong) launch8<int64_t>(row_ptr, edge_scores, out, num_nodes);
    else if (row_ptr.scalar_type() == at::kFloat) launch8<float>(row_ptr, edge_scores, out, num_nodes);
    return out;
}
"""

_ext = load_inline(name="csr_softmax_ext_final", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["csr_softmax_forward"], with_cuda=True, extra_cuda_cflags=["-O3", "--use_fast_math"], extra_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:
        return _ext.csr_softmax_forward(row_ptr, edge_scores)