new cuda layout

Author: rusty1s

aten/cuda/cluster.cpp

@@ -1,11 +1,15 @@
 #include <torch/torch.h>
 
-#define CHECK_CUDA(x) AT_ASSERT(x.type().is_cuda(), #x " must be a CUDA tensor")
+at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
+                at::Tensor end);
 
-#include "graclus.cpp"
-#include "grid.cpp"
+at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes);
+
+at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
+                            int num_nodes);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("graclus", &graclus, "Graclus (CUDA)");
   m.def("grid", &grid, "Grid (CUDA)");
+  m.def("graclus", &graclus, "Graclus (CUDA)");
+  m.def("weighted_graclus", &weighted_graclus, "Weightes Graclus (CUDA)");
 }

aten/cuda/color.cuh

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <ATen/ATen.h>
+
+#include "common.cuh"
+
+#define BLUE_PROB 0.53406
+
+__global__ void color_kernel(int64_t *cluster, size_t num_nodes) {
+  const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+  for (ptrdiff_t i = index; i < num_nodes; i += stride) {
+  }
+}
+
+inline bool color(at::Tensor cluster) {
+  color_kernel<scalar_t><<<BLOCKS(cluster.size(0)), THREADS>>>(
+      cluster.data<int64_t>(), cluster.size(0));
+
+  return true;
+}

aten/cuda/common.cuh

@@ -0,0 +1,12 @@
+#pragma once
+
+#include <ATen/ATen.h>
+
+#define THREADS 1024
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+inline at::Tensor degree(at::Tensor index, int num_nodes) {
+  auto zero = at::zeros(index.type(), {num_nodes});
+  auto one = at::ones(index.type(), {index.size(0)});
+  return zero.scatter_add_(0, index, one);
+}

aten/cuda/graclus.cpp

@@ -0,0 +1,61 @@
+#include <ATen/ATen.h>
+
+#include "color.cuh"
+#include "common.cuh"
+
+at::Tensor graclus(at::Tensor row, at::Tensor col, int num_nodes) {
+  // Remove self-loops.
+  auto mask = row != col;
+  row = row.masked_select(mask);
+  col.masked_select(mask);
+
+  // Sort by row index.
+  at::Tensor perm;
+  std::tie(row, perm) = row.sort();
+  col = col.index_select(0, perm);
+
+  // Generate helper vectors.
+  auto cluster = at::full(row.type(), {num_nodes}, -1);
+  auto prop = at::full(row.type(), {num_nodes}, -1);
+  auto deg = degree(row, num_nodes);
+  auto cum_deg = deg.cumsum(0);
+
+  color(cluster);
+
+  /* while (!color(cluster)) { */
+  /*   propose(cluster, prop, row, col, weight, deg, cum_deg); */
+  /*   response(cluster, prop, row, col, weight, deg, cum_deg); */
+  /* } */
+
+  return cluster;
+}
+
+at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
+                            int num_nodes) {
+  // Remove self-loops.
+  auto mask = row != col;
+  row = row.masked_select(mask);
+  col = col.masked_select(mask);
+  weight = weight.masked_select(mask);
+
+  // Sort by row index.
+  at::Tensor perm;
+  std::tie(row, perm) = row.sort();
+  col = col.index_select(0, perm);
+  weight = weight.index_select(0, perm);
+
+  // Generate helper vectors.
+  auto cluster = at::full(row.type(), {num_nodes}, -1);
+  auto prop = at::full(row.type(), {num_nodes}, -1);
+  auto deg = degree(row, num_nodes);
+  auto cum_deg = deg.cumsum(0);
+
+  color(cluster);
+
+  /* while (!color(cluster)) { */
+  /*   weighted_propose(cluster, prop, row, col, weight, deg, cum_deg); */
+  /*   weighted_response(cluster, prop, row, col, weight, deg, cum_deg); */
+  /* } */
+
+  return cluster;
+}

aten/cuda/graclus_kernel.cu

@@ -1,41 +1,38 @@
 #include <ATen/ATen.h>
 #include <ATen/cuda/detail/IndexUtils.cuh>
 
-#define THREADS 1024
-#define BLOCKS(N) (N + THREADS - 1) / THREADS
+#include "common.cuh"
 
 template <typename scalar_t>
 __global__ void
-grid_cuda_kernel(int64_t *cluster,
-                 at::cuda::detail::TensorInfo<scalar_t, int> pos,
-                 scalar_t *__restrict__ size, scalar_t *__restrict__ start,
-                 scalar_t *__restrict__ end, size_t num_nodes) {
+grid_kernel(int64_t *cluster, at::cuda::detail::TensorInfo<scalar_t, int> pos,
+            scalar_t *__restrict__ size, scalar_t *__restrict__ start,
+            scalar_t *__restrict__ end, size_t num_nodes) {
   const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
   const size_t stride = blockDim.x * gridDim.x;
   for (ptrdiff_t i = index; i < num_nodes; i += stride) {
     int64_t c = 0, k = 1;
     scalar_t tmp;
     for (ptrdiff_t d = 0; d < pos.sizes[1]; d++) {
-      tmp = (pos.data[i * pos.strides[0] + d * pos.strides[1]]) - start[d];
+      tmp = pos.data[i * pos.strides[0] + d * pos.strides[1]] - start[d];
       c += (int64_t)(tmp / size[d]) * k;
       k += (int64_t)((end[d] - start[d]) / size[d]);
     }
     cluster[i] = c;
   }
 }
 
-at::Tensor grid_cuda(at::Tensor pos, at::Tensor size, at::Tensor start,
-                     at::Tensor end) {
-  auto num_nodes = pos.size(0);
-  auto cluster = at::empty(pos.type().toScalarType(at::kLong), {num_nodes});
+at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
+                at::Tensor end) {
+  auto cluster = at::empty(pos.type().toScalarType(at::kLong), {pos.size(0)});
 
-  AT_DISPATCH_ALL_TYPES(pos.type(), "grid_cuda_kernel", [&] {
-    grid_cuda_kernel<scalar_t><<<BLOCKS(num_nodes), THREADS>>>(
+  AT_DISPATCH_ALL_TYPES(pos.type(), "grid_kernel", [&] {
+    grid_kernel<scalar_t><<<BLOCKS(pos.size(0)), THREADS>>>(
         cluster.data<int64_t>(),
         at::cuda::detail::getTensorInfo<scalar_t, int>(pos),
         size.toType(pos.type()).data<scalar_t>(),
         start.toType(pos.type()).data<scalar_t>(),
-        end.toType(pos.type()).data<scalar_t>(), num_nodes);
+        end.toType(pos.type()).data<scalar_t>(), pos.size(0));
   });
 
   return cluster;