Add GNN message-passing examples for the einsum JIT - #119
Conversation
|
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. |
|
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 ( 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. |
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) andgnn_million(a million-edge multi-layer sweep). Behind thejitfeature.They double as a sanity check on the JIT backend:
gnn_stressasserts the JIT result matches the VM (max_abs_diff0) 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 --examplescargo +nightly run -p linalg --features jit --example gnn_stress(JIT matches VM, ~146x faster on this machine)git diff --check