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 <cuda_runtime.h>
#include <cuda_fp16.h>

__global__ void row_wise_softmax_kernel(
    const int* row_ptr,
    const float* edge_scores,
    float* 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 num_edges = end - start;

    if (num_edges == 0) return;

    edge_scores += start * num_heads;
    out += start * num_heads;

    for (int head = threadIdx.y; head < num_heads; head += blockDim.y) {
        float max_score = edge_scores[head];
        for (int i = threadIdx.x + 1; i < num_edges; i += blockDim.x) {
            max_score = fmaxf(max_score, edge_scores[i * num_heads + head]);
        }
        max_score = __shfl_sync(0xFFFFFFFF, max_score, 0);

        float sum_exp = __expf(edge_scores[head] - max_score);
        for (int i = threadIdx.x + 1; i < num_edges; i += blockDim.x) {
            sum_exp += __expf(edge_scores[i * num_heads + head] - max_score);
        }
#pragma unroll
        for (int offset = blockDim.x / 2; offset > 0; offset /= 2) {
            sum_exp += __shfl_down_sync(0xFFFFFFFF, sum_exp, offset);
        }
        sum_exp = __shfl_sync(0xFFFFFFFF, sum_exp, 0);

        out[head] = __expf(edge_scores[head] - max_score) / sum_exp;
        for (int i = threadIdx.x + 1; i < num_edges; i += blockDim.x) {
            out[i * num_heads + head] = __expf(edge_scores[i * num_heads + head] - max_score) / sum_exp;
        }
    }
}

torch::Tensor row_wise_softmax(
torch::Tensor row_ptr,
torch::Tensor edge_scores
) {
    int num_nodes = row_ptr.numel() - 1;
    int num_edges = row_ptr[num_nodes].item<int>();
    int num_heads = edge_scores.size(1);

    auto out = torch::empty_like(edge_scores);

    int blockSize = 32;
    dim3 block(blockSize, 8);
    dim3 grid(num_nodes);

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

    return out;
}
'''

module = load_inline(
    name='row_wise_softmax_cuda',
    cuda_sources=[cuda_source],
    functions=['row_wise_softmax'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        edge_scores: torch.Tensor,
    ) -> torch.Tensor:
        return module.row_wise_softmax(row_ptr, edge_scores)