Skip to content

Add GNN message-passing examples for the einsum JIT - #119

Closed
MesTTo wants to merge 2 commits into
trueagi-io:mainfrom
MesTTo:pr/linalg-gnn-examples
Closed

Add GNN message-passing examples for the einsum JIT#119
MesTTo wants to merge 2 commits into
trueagi-io:mainfrom
MesTTo:pr/linalg-gnn-examples

Conversation

@MesTTo

@MesTTo MesTTo commented Jun 24, 2026

Copy link
Copy Markdown
Contributor

This PR is stacked on #115.

Adds two examples that exercise the einsum-f32 JIT on graph neural-network message passing: gnn_stress (a single large layer, comparing the VM and JIT paths and checking they match) and gnn_million (a million-edge multi-layer sweep). Behind the jit feature.

They double as a sanity check on the JIT backend: gnn_stress asserts the JIT result matches the VM (max_abs_diff 0) and reports the speedup.

Validation:

  • CARGO_HOME=/home/user/Dev/.cargo-isolated RUSTFLAGS='-C target-cpu=native -Awarnings' cargo +nightly build -p linalg --features jit --examples
  • cargo +nightly run -p linalg --features jit --example gnn_stress (JIT matches VM, ~146x faster on this machine)
  • git diff --check

@Adam-Vandervorst

Copy link
Copy Markdown
Collaborator

Hmm I don't see why this PR has to be stacked; can't all /linalg changes be applied independently? I'd love more benchmarks here.

@MesTTo

MesTTo commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

the /linalg PRs are independent of each other... #111 (blas backend), #112 (einsum output-shape inference) and #115 (jit rework) each apply straight on main. this one only sits on #115 because the examples call the reworked api (EinsumF32Plan / einsum_auto), which doesn't exist before it, so once #115 lands this drops to just the two example files.

if you'd rather, i can retarget the examples at the current jit api so they stand alone, or fold them into #115... whichever reads better for you.

on more benchmarks: i have a crossover bench (openblas vs the jit as density increases) that i left out only because it needs both #111 and #115. i can add it the moment those are in, or bundle it now if you want it sooner.

@MesTTo

MesTTo commented Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

folded these into #115 so the jit rework ships self-contained with its examples and benchmarks, rather than as a separate stacked pr. closing this one... the gnn examples + graph/perf benches now live in #115 and apply independently on main. thanks for the nudge.

@MesTTo MesTTo closed this Jun 24, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants