yzh119 / custom_op_benchmark

naive graph attention kernels

custom_op_benchmark

Currently I have only implemented the custom op for PyTorch.

TODO

Switch backend to dlpack.

Usage

cd graphop
python setup.py install
cd .. && python wrapper.py

Reference

• extension-cpp