`segment_matmul` integration from `pyg-lib`

[puririshi98]

🚀 The feature, motivation and pitch

• pyg-lib issue: https://github.com/pyg-team/pyg-lib/issues/49
• pyg-lib PR: https://github.com/pyg-team/pyg-lib/pull/51

Need to integrate this into the backend of PyG's heterogeneous GNNs.

[puririshi98]

Candidates for first round of integration:

• RGCNConv

• HeteroLinear

Candidates for second round(nvfuser+pyg-lib):

• HeteroConv

[rusty1s]

Important to note, this will require adding backward support of segment_matmul in pyg-lib.

[puririshi98]

Important to note, this will require adding backward support of segment_matmul in pyg-lib.

https://github.com/pyg-team/pyg-lib/pull/61 https://github.com/pyg-team/pyg-lib/pull/63

[puririshi98]

https://github.com/pyg-team/pytorch_geometric/pull/5052

[puririshi98]

speeding it up: https://github.com/pyg-team/pyg-lib/pull/73

[puririshi98]

heterolinear cutlass support: https://github.com/pyg-team/pytorch_geometric/pull/5330