gputest

benchmark/ffi.py

+import torch
+from torch_cluster._ext import ffi
+
+cluster = torch.cuda.LongTensor(5)
+pos = torch.cuda.FloatTensor(5, 2)
+size = torch.cuda.FloatTensor(2)
+count = torch.cuda.LongTensor(2)
+
+func = ffi.THCCFloatGrid
+print(func)
+
+func(cluster, pos, size, count)

build_new.sh

+#!/bin/sh
+
+echo "Compiling kernel..."
+
+if [ -z "$1" ]; then TORCH=$(python -c "import os; import torch; print(os.path.dirname(torch.__file__))"); else TORCH="$1"; fi
+SRC_DIR=aten/THC
+BUILD_DIR=aten/build
+
+mkdir -p $BUILD_DIR
+for i in THCGreedy THCGrid; do
+  $(which nvcc) -c -o "$BUILD_DIR/$i.so" "$SRC_DIR/$i.cu" -arch=sm_52 -Xcompiler -fPIC -shared "-I$TORCH/lib/include/TH" "-I$TORCH/lib/include" "-I$SRC_DIR"
+done