Speed-up torch.compile PyG models (#6980)

Currently, PyG models needs to be moved to `jittable` in order to see
any speed-ups.

CHANGELOG.md

@@ -36,7 +36,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `Pad` transform ([#5940](https://github.com/pyg-team/pytorch_geometric/pull/5940), [#6697](https://github.com/pyg-team/pytorch_geometric/pull/6697), [#6731](https://github.com/pyg-team/pytorch_geometric/pull/6731), [#6758](https://github.com/pyg-team/pytorch_geometric/pull/6758))
 - Added full batch mode to the inference benchmark ([#6631](https://github.com/pyg-team/pytorch_geometric/pull/6631))
 - Added `cat` aggregation type to the `HeteroConv` class so that features can be concatenated during grouping ([#6634](https://github.com/pyg-team/pytorch_geometric/pull/6634))
-- Added `torch.compile` support and benchmark study ([#6610](https://github.com/pyg-team/pytorch_geometric/pull/6610), [#6952](https://github.com/pyg-team/pytorch_geometric/pull/6952), [#6953](https://github.com/pyg-team/pytorch_geometric/pull/6953))
+- Added `torch.compile` support and benchmark study ([#6610](https://github.com/pyg-team/pytorch_geometric/pull/6610), [#6952](https://github.com/pyg-team/pytorch_geometric/pull/6952), [#6953](https://github.com/pyg-team/pytorch_geometric/pull/6953), [#6980](https://github.com/pyg-team/pytorch_geometric/pull/6980))
 - Added the `AntiSymmetricConv` layer ([#6577](https://github.com/pyg-team/pytorch_geometric/pull/6577))
 - Added a mixin for Huggingface model hub integration ([#5930](https://github.com/pyg-team/pytorch_geometric/pull/5930), [#6591](https://github.com/pyg-team/pytorch_geometric/pull/6591))
 - Added support for accelerated GNN layers in `nn.conv.cugraph` via `cugraph-ops` ([#6278](https://github.com/pyg-team/pytorch_geometric/pull/6278), [#6388](https://github.com/pyg-team/pytorch_geometric/pull/6388), [#6412](https://github.com/pyg-team/pytorch_geometric/pull/6412))

test/nn/test_compile_conv.py

@@ -50,14 +50,28 @@ def test_compile_conv(device, Conv):
     x = torch.randn(num_nodes, 64, device=args.device)
     edge_index = torch.randint(num_nodes, (2, num_edges), device=args.device)
 
-    for Conv in [GCNConv, SAGEConv, MySAGEConv, GATConv]:
+    conv = MySAGEConv(64, 64).to(args.device)
+    benchmark(
+        funcs=[conv, torch.compile(enable_compile()(conv))],
+        func_names=['Vanilla', 'Compiled'],
+        args=(x, edge_index),
+        num_steps=50 if args.device == 'cpu' else 500,
+        num_warmups=10 if args.device == 'cpu' else 100,
+        backward=args.backward,
+    )
+
+    for Conv in [GCNConv, SAGEConv]:
         print(f'Conv: {Conv.__name__}')
 
         conv = Conv(64, 64).to(args.device)
+        compiled_conv = torch.compile(enable_compile()(conv))
+
+        jit_conv = Conv(64, 64).jittable().to(args.device)
+        compiled_jit_conv = torch.compile(enable_compile()(jit_conv))
 
         benchmark(
-            funcs=[conv, torch.compile(enable_compile()(conv))],
-            func_names=['Vanilla', 'Compiled'],
+            funcs=[conv, compiled_conv, jit_conv, compiled_jit_conv],
+            func_names=['Vanilla', 'Compiled', 'JIT', 'JIT Compiled'],
             args=(x, edge_index),
             num_steps=50 if args.device == 'cpu' else 500,
             num_warmups=10 if args.device == 'cpu' else 100,