rustytorch/crates/specialized/rtx-geom/benches/gnn_bench.rs

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_geom::layers::{GATLayer, GCNLayer, GNNLayer, GraphSAGELayer};
use rtx_geom::{AggregationType, Edge, Graph, MessagePassing, Node};
use rtx_tensor::{Device, Tensor};

fn create_test_graph(num_nodes: usize, feature_dim: usize, edge_probability: f32) -> Graph {
    let mut graph = Graph::new();
    let mut nodes = Vec::new();
    let device = Device::cpu();

    // Add nodes
    for i in 0..num_nodes {
        let mut features = vec![0.0f32; feature_dim];
        for j in 0..feature_dim {
            features[j] = (i * 7 + j * 13) as f32 % 100.0 / 50.0 - 1.0; // [-1, 1] range
        }
        let node = graph.add_node(Node::new(
            Tensor::from_data(features, [feature_dim], &device).unwrap(),
        ));
        nodes.push(node);
    }

    // Add edges based on probability (deterministic for benchmarking)
    for i in 0..num_nodes {
        for j in 0..num_nodes {
            if i != j {
                let hash = (i * 1337 + j * 7919) % 10000;
                if (hash as f32 / 10000.0) < edge_probability {
                    let weight = Tensor::from_data(vec![1.0], [1], &device).unwrap();
                    let _ = graph.add_edge(nodes[i], nodes[j], Edge::new(weight));
                }
            }
        }
    }

    graph
}

fn bench_graph_creation(c: &mut Criterion) {
    let mut group = c.benchmark_group("graph_creation");

    for &size in &[10, 50, 100, 500] {
        group.throughput(Throughput::Elements(size as u64));
        group.bench_with_input(
            BenchmarkId::new("nodes_and_edges", size),
            &size,
            |b, &size| {
                b.iter(|| create_test_graph(size, 4, 0.1));
            },
        );
    }

    group.finish();
}

fn bench_message_passing(c: &mut Criterion) {
    let mut group = c.benchmark_group("message_passing");

    let graphs: Vec<_> = [10, 50, 100, 200]
        .iter()
        .map(|&size| (size, create_test_graph(size, 8, 0.2)))
        .collect();

    for (size, graph) in graphs {
        group.throughput(Throughput::Elements(size as u64));

        // Benchmark different aggregation types
        for &agg_type in &[
            AggregationType::Sum,
            AggregationType::Mean,
            AggregationType::Max,
            AggregationType::WeightedSum,
        ] {
            let mp = MessagePassing::new(agg_type);
            group.bench_with_input(
                BenchmarkId::new(format!("{:?}", agg_type), size),
                &(&graph, &mp),
                |b, (graph, mp)| {
                    b.iter(|| mp.compute_messages(graph));
                },
            );
        }
    }

    group.finish();
}

fn bench_gcn_layer(c: &mut Criterion) {
    let mut group = c.benchmark_group("gcn_layer");

    let test_cases = vec![
        (10, 4, 8),    // small graph
        (50, 8, 16),   // medium graph
        (100, 16, 32), // large graph
    ];

    for (num_nodes, input_dim, output_dim) in test_cases {
        let graph = create_test_graph(num_nodes, input_dim, 0.15);

        group.throughput(Throughput::Elements(num_nodes as u64));
        group.bench_function(
            BenchmarkId::new(
                "forward_pass",
                format!("{}x{}->{}", num_nodes, input_dim, output_dim),
            ),
            |b| {
                b.iter(|| {
                    let mut layer = GCNLayer::new(input_dim, output_dim).unwrap();
                    layer.forward(&graph).unwrap()
                });
            },
        );
    }

    group.finish();
}

fn bench_gat_layer(c: &mut Criterion) {
    let mut group = c.benchmark_group("gat_layer");

    let test_cases = vec![
        (10, 4, 8, 2),  // small graph, 2 heads
        (25, 6, 12, 4), // medium graph, 4 heads
        (50, 8, 16, 8), // large graph, 8 heads
    ];

    for (num_nodes, input_dim, output_dim, num_heads) in test_cases {
        let graph = create_test_graph(num_nodes, input_dim, 0.2);

        group.throughput(Throughput::Elements(num_nodes as u64));
        group.bench_function(
            BenchmarkId::new(
                "forward_pass",
                format!("{}x{}->{}_{}h", num_nodes, input_dim, output_dim, num_heads),
            ),
            |b| {
                b.iter(|| {
                    let mut layer = GATLayer::new(input_dim, output_dim, num_heads).unwrap();
                    layer.forward(&graph).unwrap()
                });
            },
        );
    }

    group.finish();
}

fn bench_graphsage_layer(c: &mut Criterion) {
    let mut group = c.benchmark_group("graphsage_layer");

    let test_cases = vec![
        (20, 4, 8, 2),     // small graph, sample_size=2
        (50, 8, 16, 5),    // medium graph, sample_size=5
        (100, 16, 32, 10), // large graph, sample_size=10
    ];

    for (num_nodes, input_dim, output_dim, sample_size) in test_cases {
        let graph = create_test_graph(num_nodes, input_dim, 0.25);

        group.throughput(Throughput::Elements(num_nodes as u64));
        group.bench_function(
            BenchmarkId::new(
                "forward_pass",
                format!(
                    "{}x{}->{}_s{}",
                    num_nodes, input_dim, output_dim, sample_size
                ),
            ),
            |b| {
                b.iter(|| {
                    let mut layer =
                        GraphSAGELayer::new(input_dim, output_dim, sample_size).unwrap();
                    layer.forward(&graph).unwrap()
                });
            },
        );
    }

    group.finish();
}

fn bench_adjacency_matrix(c: &mut Criterion) {
    let mut group = c.benchmark_group("adjacency_matrix");

    for &size in &[10, 25, 50, 100] {
        let graph = create_test_graph(size, 4, 0.15);

        group.throughput(Throughput::Elements((size * size) as u64));
        group.bench_with_input(BenchmarkId::new("computation", size), &graph, |b, graph| {
            b.iter(|| graph.adjacency_matrix());
        });
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_graph_creation,
    bench_message_passing,
    bench_gcn_layer,
    bench_gat_layer,
    bench_graphsage_layer,
    bench_adjacency_matrix
);

criterion_main!(benches);