Add Triton support

Author: denix56

csrc/cluster.h

@@ -50,5 +50,3 @@ TORCH_LIBRARY(torch_cluster, m) {
     m.def("neighbor_sampler(Tensor start, Tensor rowptr, int count, float factor) -> Tensor");
     m.def("cuda_version() -> int");
 }
-
-}

csrc/fps.cpp

@@ -41,4 +41,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
   m.impl("fps", &fps_cuda);
 }
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("fps", &fps_cuda);
+}
 #endif

csrc/graclus.cpp

@@ -41,4 +41,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
   m.impl("graclus", &graclus_cuda);
 }
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("graclus", &graclus_cuda);
+}
 #endif

csrc/grid.cpp

@@ -42,4 +42,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
     m.impl("grid", &grid_cuda);
 }
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+    m.impl("grid", &grid_cuda);
+}
 #endif

csrc/knn.cpp

@@ -50,7 +50,16 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 }
 
 #ifdef WITH_CUDA
-    TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
-      m.impl("knn", &knn_cuda);
-    }
+torch::Tensor knn_cuda_wrap(torch::Tensor x, torch::Tensor y,
+                  std::optional<torch::Tensor> ptr_x,
+                  std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  int64_t num_workers) {
+    return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
+    m.impl("knn", &knn_cuda_wrap);
+}
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+    m.impl("knn", &knn_cuda_wrap);
+}
 #endif

csrc/nearest.cpp

@@ -37,4 +37,7 @@ CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tenso
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
     m.impl("nearest", &nearest_cuda);
 }
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("nearest", &nearest_cuda);
+}
 #endif

csrc/radius.cpp

@@ -44,4 +44,7 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
   m.impl("radius", &radius_cuda);
 }
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("radius", &radius_cuda);
+}
 #endif

csrc/rw.cpp

@@ -42,4 +42,8 @@ TORCH_LIBRARY_IMPL(torch_cluster, CPU, m) {
 TORCH_LIBRARY_IMPL(torch_cluster, CUDA, m) {
   m.impl("random_walk", &random_walk_cuda);
 }
+
+TORCH_LIBRARY_IMPL(torch_cluster, HIP, m) {
+  m.impl("random_walk", &random_walk_cuda);
+}
 #endif

csrc/version.cpp

@@ -1,10 +1,11 @@
 #ifdef WITH_PYTHON
 #include <Python.h>
 #endif
+
+#include <torch/library.h>
+
 #include "cluster.h"
 #include "macros.h"
-#include <torch/torch.h>
-#include <torch/library.h>
 
 #ifdef WITH_CUDA
 #ifdef USE_ROCM
@@ -38,6 +39,6 @@ CLUSTER_API int64_t cuda_version() noexcept {
 }
 } // namespace cluster
 
-TORCH_LIBRARY_IMPL(torch_cluster, m) {
+TORCH_LIBRARY_IMPL(torch_cluster, CompositeImplicitAutograd, m) {
   m.impl("cuda_version", [] { return cluster::cuda_version(); });
 }