import torch
import torch.nn as nn


class Model(nn.Module):
    """
    Fused sparse attention-value aggregation in CSR format.

    For each destination row, compute a softmax over edge logits and then use
    the normalized weights to aggregate source node values. This is a common
    fused graph-attention pattern.
    """

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

    def forward(
        self,
        row_ptr: torch.Tensor,
        col_idx: torch.Tensor,
        edge_scores: torch.Tensor,
        node_value: torch.Tensor,
    ) -> torch.Tensor:
        num_nodes = row_ptr.numel() - 1
        feat_dim = node_value.shape[1]
        out = torch.zeros(num_nodes, feat_dim, dtype=node_value.dtype, device=node_value.device)

        for dst in range(num_nodes):
            start = int(row_ptr[dst].item())
            end = int(row_ptr[dst + 1].item())
            if end > start:
                weights = torch.softmax(edge_scores[start:end], dim=0)
                src = col_idx[start:end].long()
                out[dst] = (node_value[src] * weights.unsqueeze(-1)).sum(dim=0)

        return out


num_nodes = 384
avg_degree = 18
feat_dim = 64


def get_inputs():
    return list(
        popcorn_pri.make_csr_fused_attention_value_graph(
            num_nodes, avg_degree, feat_dim
        )
    )



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_attn_forward(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_scores, torch::Tensor node_value);
'''

_cuda_source = r'''
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

template<int KIND>
__device__ __forceinline__ int read_idx_t(const void* p, int i) {
    if constexpr (KIND == 0) return ((const int*)p)[i];
    else if constexpr (KIND == 1) return (int)((const long long*)p)[i];
    else return (int)((const float*)p)[i];
}

__device__ __forceinline__ float warp_reduce_sum0(float v) {
    v += __shfl_down_sync(0xffffffff, v, 16);
    v += __shfl_down_sync(0xffffffff, v, 8);
    v += __shfl_down_sync(0xffffffff, v, 4);
    v += __shfl_down_sync(0xffffffff, v, 2);
    v += __shfl_down_sync(0xffffffff, v, 1);
    return v;
}
__device__ __forceinline__ float warp_reduce_max0(float v) {
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 16));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 8));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 4));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 2));
    v = fmaxf(v, __shfl_down_sync(0xffffffff, v, 1));
    return v;
}

template<int ROWK, int COLK>
__global__ void csr_attn_kernel_t(const void* __restrict__ row_ptr,
                                  const void* __restrict__ col_idx,
                                  const float* __restrict__ edge_scores,
                                  const float* __restrict__ node_value,
                                  float* __restrict__ out,
                                  int num_nodes,
                                  int feat_dim) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid & 31;
    int start = read_idx_t<ROWK>(row_ptr, row);
    int end = read_idx_t<ROWK>(row_ptr, row + 1);
    int deg = end - start;
    __shared__ float weights[1024];
    __shared__ float inv_sum_s;
    if (deg <= 0) {
        for (int f = tid; f < feat_dim; f += blockDim.x) out[row * feat_dim + f] = 0.0f;
        return;
    }
    if (deg <= 32) {
        if (tid < 32) {
            float s = (lane < deg) ? edge_scores[start + lane] : -3.402823466e+38F;
            float m = warp_reduce_max0(s);
            m = __shfl_sync(0xffffffff, m, 0);
            float w = (lane < deg) ? expf(s - m) : 0.0f;
            float denom = warp_reduce_sum0(w);
            denom = __shfl_sync(0xffffffff, denom, 0);
            if (lane < deg) weights[lane] = w;
            if (lane == 0) inv_sum_s = 1.0f / denom;
        }
    } else {
        if (tid == 0) {
            float m = -3.402823466e+38F;
            for (int e = 0; e < deg; ++e) m = fmaxf(m, edge_scores[start + e]);
            float denom = 0.0f;
            for (int e = 0; e < deg; ++e) { float w = expf(edge_scores[start + e] - m); weights[e] = w; denom += w; }
            inv_sum_s = 1.0f / denom;
        }
    }
    __syncthreads();
    float inv_sum = inv_sum_s;
    for (int f = tid; f < feat_dim; f += blockDim.x) {
        float acc = 0.0f;
        #pragma unroll
        for (int e = 0; e < 32; ++e) {
            if (e >= deg) break;
            int src = read_idx_t<COLK>(col_idx, start + e);
            acc = fmaf(node_value[src * feat_dim + f], weights[e] * inv_sum, acc);
        }
        if (deg > 32) {
            for (int e = 32; e < deg; ++e) {
                int src = read_idx_t<COLK>(col_idx, start + e);
                acc = fmaf(node_value[src * feat_dim + f], weights[e] * inv_sum, acc);
            }
        }
        out[row * feat_dim + f] = acc;
    }
}

template<int RK>
void launch_col(int ck, const void* rp, const void* ci, const float* es, const float* nv, float* out, int n, int fd, cudaStream_t stream) {
    if (ck == 0) csr_attn_kernel_t<RK,0><<<dim3(n), dim3(64), 0, stream>>>(rp, ci, es, nv, out, n, fd);
    else if (ck == 1) csr_attn_kernel_t<RK,1><<<dim3(n), dim3(64), 0, stream>>>(rp, ci, es, nv, out, n, fd);
    else csr_attn_kernel_t<RK,2><<<dim3(n), dim3(64), 0, stream>>>(rp, ci, es, nv, out, n, fd);
}

torch::Tensor csr_attn_forward(torch::Tensor row_ptr, torch::Tensor col_idx, torch::Tensor edge_scores, torch::Tensor node_value) {
    int num_nodes = (int)row_ptr.numel() - 1;
    int feat_dim = (int)node_value.size(1);
    auto out = torch::empty({num_nodes, feat_dim}, node_value.options());
    int rk = row_ptr.scalar_type() == at::kInt ? 0 : (row_ptr.scalar_type() == at::kLong ? 1 : 2);
    int ck = col_idx.scalar_type() == at::kInt ? 0 : (col_idx.scalar_type() == at::kLong ? 1 : 2);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    if (rk == 0) launch_col<0>(ck, row_ptr.data_ptr(), col_idx.data_ptr(), edge_scores.data_ptr<float>(), node_value.data_ptr<float>(), out.data_ptr<float>(), num_nodes, feat_dim, stream);
    else if (rk == 1) launch_col<1>(ck, row_ptr.data_ptr(), col_idx.data_ptr(), edge_scores.data_ptr<float>(), node_value.data_ptr<float>(), out.data_ptr<float>(), num_nodes, feat_dim, stream);
    else launch_col<2>(ck, row_ptr.data_ptr(), col_idx.data_ptr(), edge_scores.data_ptr<float>(), node_value.data_ptr<float>(), out.data_ptr<float>(), num_nodes, feat_dim, stream);
    return out;
}
'''

_ext = load_inline(name='csr_attn_ext_final', cpp_sources=_cpp_source, cuda_sources=_cuda_source,
                   functions=['csr_attn_forward'], 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, col_idx: torch.Tensor, edge_scores: torch.Tensor, node_value: torch.Tensor) -> torch.Tensor:
        return _ext.csr_attn_forward(row_ptr, col_idx, edge_scores, node_value)