search

fishmingyu / GeoT

GeoT: Tensor Centric Library for Graph Neural Network via Efficient Segment Reduction on GPU
https://arxiv.org/abs/2404.03019
17 stars 1 forks source link

`torch.compile(..., fullgraph=True)` compatible API ? #34

Open mitkotak opened 2 weeks ago

mitkotak commented 2 weeks ago

Thanks for the great work ! Coming here after reading this. Was wondering if there's something similar to this in the repo that I can use that is compatible with torch.compile ? Ideally would like to have only Pytorch as the dependency.

fishmingyu commented 2 weeks ago

Thank you so much for your interest! We are currently working on torch.compile support, where you can check the test_gcn after installing GeoT. However if you want PyTorch as the only dependency, things could be more challenging since PyTorch team has it's own backend. You can check our triton code at Triton SpMM and manually apply it to PyTorch using custom triton kernel Tutorial.