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

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>
#include <cstdint>

__global__ void csr_softmax_kernel(
    const int32_t* __restrict__ row_ptr,
    const float* __restrict__ edge_scores,
    float* __restrict__ out,
    int num_nodes,
    int num_heads
) {
    int node = blockIdx.x;
    if (node >= num_nodes) return;

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

    for (int h = threadIdx.x; h < num_heads; h += blockDim.x) {
        float max_val = -FLT_MAX;
        for (int i = 0; i < degree; i++) {
            float val = edge_scores[(start + i) * num_heads + h];
            max_val = fmaxf(max_val, val);
        }
        float sum_exp = 0.0f;
        for (int i = 0; i < degree; i++) {
            float val = edge_scores[(start + i) * num_heads + h];
            sum_exp += expf(val - max_val);
        }
        float inv_sum = 1.0f / sum_exp;
        for (int i = 0; i < degree; i++) {
            float val = edge_scores[(start + i) * num_heads + h];
            out[(start + i) * num_heads + h] = expf(val - max_val) * inv_sum;
        }
    }
}

torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores) {
    int num_nodes = row_ptr.numel() - 1;
    auto num_heads = edge_scores.size(1);
    auto out = torch::empty_like(edge_scores);

    dim3 grid(num_nodes);
    dim3 block(8);

    csr_softmax_kernel<<<grid, block>>>(
        row_ptr.data_ptr<int32_t>(),
        edge_scores.data_ptr<float>(),
        out.data_ptr<float>(),
        num_nodes,
        num_heads
    );

    return out;
}
"""

cpp_source = """
torch::Tensor csr_softmax_cuda(torch::Tensor row_ptr, torch::Tensor edge_scores);
"""

module = load_inline(
    name="csr_softmax",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["csr_softmax_cuda"],
    verbose=False,
    with_cuda=True,
)

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

    def forward(self, row_ptr: torch.Tensor, edge_scores: torch.Tensor) -> torch.Tensor:
        row_ptr = row_ptr.cuda().to(torch.int32).contiguous()
        edge_scores = edge_scores.cuda().contiguous()
        return self.module.csr_softmax_cuda(row_ptr, edge_scores)