Expand description
Sparse tensor operations for GNN (Graph Neural Network) workloads.
This module extends SpMV/SpMM primitives with operations tailored for graph neural network message passing and aggregation, including:
scatter_reduce– Scatter-reduce with configurable aggregationgather– Index-based gathering from source arrayssparse_message_passing– GNN message passing over adjacency matricessparse_attention_message– Attention-weighted message passing (GAT-style)compute_degree_matrix– Node degree computation from CSR offsetssymmetric_normalize– D^{-1/2} A D^{-1/2} normalization (GCN-style)add_self_loops– Add identity self-loops to adjacency matrixsparse_row_softmax– Row-wise softmax over sparse matrix valuesgenerate_message_passing_ptx– PTX kernel generation for GPU message passing
(C) 2026 COOLJAPAN OU (Team KitaSan)
Structs§
- Edge
Features - Edge features for attention-weighted message passing.
- GnnSparse
Config - Configuration for GNN sparse operations.
Enums§
- Message
Passing Op - Aggregation operation for GNN message passing.
Functions§
- add_
self_ loops - Add self-loops to an adjacency matrix: A’ = A + I.
- compute_
degree_ matrix - Compute the degree of each node from CSR row offsets.
- gather
- Gather operation.
- generate_
message_ passing_ ptx - Generate a PTX kernel for GPU-accelerated GNN message passing.
- scatter_
reduce - Scatter-reduce operation.
- sparse_
attention_ message - Attention-weighted message passing.
- sparse_
message_ passing - Sparse message passing over an adjacency matrix in CSR form.
- sparse_
row_ softmax - Row-wise softmax over sparse matrix values.
- symmetric_
normalize - Symmetric normalization of an adjacency matrix: D^{-1/2} A D^{-1/2}.