Initial commit

csrc/hip/radius_hip.h

+#pragma once
+
+#include <torch/extension.h>
+
+torch::Tensor radius_cuda(torch::Tensor x, torch::Tensor y,
+                          torch::optional<torch::Tensor> ptr_x,
+                          torch::optional<torch::Tensor> ptr_y, double r,
+                          int64_t max_num_neighbors);

csrc/hip/rw_hip.h

+#pragma once
+
+#include <torch/extension.h>
+
+std::tuple<torch::Tensor, torch::Tensor>
+random_walk_cuda(torch::Tensor rowptr, torch::Tensor col, torch::Tensor start,
+                 int64_t walk_length, double p, double q);

csrc/hip/utils.cuh

+#pragma once
+
+#include <torch/extension.h>
+
+#define CHECK_CUDA(x)                                                          \
+  AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
+#define CHECK_INPUT(x) AT_ASSERTM(x, "Input mismatch")
+#define CHECK_CONTIGUOUS(x)                                                    \
+  AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+
+__device__ int64_t get_example_idx(int64_t idx, const int64_t *ptr,
+                                   const int64_t num_examples) {
+  for (int64_t i = 0; i < num_examples; i++) {
+    if (ptr[i + 1] > idx)
+      return i;
+  }
+  return num_examples - 1;
+}

csrc/knn.cpp

+#include <Python.h>
+#include <torch/script.h>
+
+#include "cpu/knn_cpu.h"
+
+#ifdef WITH_HIP
+#include "hip/knn_hip.h"
+#endif
+
+#ifdef _WIN32
+#ifdef WITH_HIP
+PyMODINIT_FUNC PyInit__knn_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__knn_cpu(void) { return NULL; }
+#endif
+#endif
+
+torch::Tensor knn(torch::Tensor x, torch::Tensor y,
+                  torch::optional<torch::Tensor> ptr_x,
+                  torch::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  int64_t num_workers) {
+  if (x.device().is_cuda()) {
+#ifdef WITH_HIP
+    return knn_cuda(x, y, ptr_x, ptr_y, k, cosine);
+#else
+    AT_ERROR("Not compiled with CUDA support");
+#endif
+  } else {
+    if (cosine)
+      AT_ERROR("`cosine` argument not supported on CPU");
+    return knn_cpu(x, y, ptr_x, ptr_y, k, num_workers);
+  }
+}
+
+static auto registry =
+    torch::RegisterOperators().op("torch_cluster::knn", &knn);

env.sh-22.10

+#!/bin/bash
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate torch1.10_py39_dtk22.10
+
+module purge
+module load compiler/devtoolset/7.3.1 mpi/hpcx/gcc-7.3.1 #compiler/dtk/22.10.1
+module list
+
+source ~/dtk-22.10.1/env.sh
+export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$C_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$CPLUS_INCLUDE_PATH
+export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$C_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$CPLUS_INCLUDE_PATH
+
+export C_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$C_INCLUDE_PATH
+export CPLUS_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$CPLUS_INCLUDE_PATH
+export LD_LIBRARY_PATH=$ROCM_PATH/rocrand/lib:$LD_LIBRARY_PATH
+export FORCE_ONLY_HIP=1
+export CC=hipcc
+export CXX=hipcc

torch_cluster.egg-info/dependency_links.txt

+