Add grid_cluster dispatch + CPU kernel

Port grid voxelization from pytorch_cluster into pyg-lib.
Each point is assigned a 1D cluster index based on its quantized
voxel position: floor((pos - start) / size), then flattened
with cumulative voxel counts.

Author: akihironitta

pyg_lib/csrc/ops/cluster.cpp

@@ -0,0 +1,47 @@
+#include "cluster.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/library.h>
+
+namespace pyg {
+namespace ops {
+
+PYG_API at::Tensor grid_cluster(const at::Tensor& pos,
+                                const at::Tensor& size,
+                                const std::optional<at::Tensor>& start,
+                                const std::optional<at::Tensor>& end) {
+  at::TensorArg pos_arg{pos, "pos", 0};
+  at::TensorArg size_arg{size, "size", 1};
+  at::CheckedFrom c{"grid_cluster"};
+
+  at::checkAllDefined(c, {pos_arg, size_arg});
+
+  auto pos_2d = pos.view({pos.size(0), -1}).contiguous();
+  auto size_c = size.contiguous();
+
+  TORCH_CHECK(size_c.numel() == pos_2d.size(1),
+              "size.numel() must equal pos dimension count");
+
+  if (start.has_value()) {
+    TORCH_CHECK(start.value().numel() == pos_2d.size(1),
+                "start.numel() must equal pos dimension count");
+  }
+  if (end.has_value()) {
+    TORCH_CHECK(end.value().numel() == pos_2d.size(1),
+                "end.numel() must equal pos dimension count");
+  }
+
+  static auto op = c10::Dispatcher::singleton()
+                       .findSchemaOrThrow("pyg::grid_cluster", "")
+                       .typed<decltype(grid_cluster)>();
+  return op.call(pos_2d, size_c, start, end);
+}
+
+TORCH_LIBRARY_FRAGMENT(pyg, m) {
+  m.def(TORCH_SELECTIVE_SCHEMA(
+      "pyg::grid_cluster(Tensor pos, Tensor size, "
+      "Tensor? start=None, Tensor? end=None) -> Tensor"));
+}
+
+}  // namespace ops
+}  // namespace pyg

pyg_lib/csrc/ops/cluster.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include <ATen/ATen.h>
+#include "pyg_lib/csrc/macros.h"
+
+namespace pyg {
+namespace ops {
+
+PYG_API at::Tensor grid_cluster(const at::Tensor& pos,
+                                const at::Tensor& size,
+                                const std::optional<at::Tensor>& start,
+                                const std::optional<at::Tensor>& end);
+
+}  // namespace ops
+}  // namespace pyg

pyg_lib/csrc/ops/cpu/cluster_kernel.cpp

@@ -0,0 +1,54 @@
+#include "../cluster.h"
+
+#include <ATen/ATen.h>
+#include <torch/library.h>
+
+namespace pyg {
+namespace ops {
+
+namespace {
+
+at::Tensor grid_cluster_kernel(const at::Tensor& pos,
+                               const at::Tensor& size,
+                               const std::optional<at::Tensor>& optional_start,
+                               const std::optional<at::Tensor>& optional_end) {
+  auto N = pos.size(0);
+  auto D = pos.size(1);
+
+  at::Tensor start;
+  if (optional_start.has_value())
+    start = optional_start.value().contiguous();
+  else
+    start = std::get<0>(pos.min(0));
+
+  at::Tensor end;
+  if (optional_end.has_value())
+    end = optional_end.value().contiguous();
+  else
+    end = std::get<0>(pos.max(0));
+
+  auto pos_shifted = pos - start.unsqueeze(0);
+
+  auto num_voxels =
+      (end - start).div(size, /*rounding_mode=*/"trunc").to(at::kLong) + 1;
+  num_voxels = num_voxels.cumprod(0);
+  num_voxels = at::cat({at::ones({1}, num_voxels.options()), num_voxels}, 0);
+  num_voxels = num_voxels.narrow(0, 0, D);
+
+  auto out = pos_shifted.div(size.view({1, -1}), /*rounding_mode=*/"trunc")
+                 .to(at::kLong);
+  out *= num_voxels.view({1, -1});
+  out = out.sum(1);
+
+  return out;
+}
+
+}  // namespace
+
+TORCH_LIBRARY_IMPL(pyg, CPU, m) {
+  m.impl(TORCH_SELECTIVE_NAME("pyg::grid_cluster"),
+         TORCH_FN(grid_cluster_kernel));
+}
+
+}  // namespace ops
+}  // namespace pyg

pyg_lib/ops/__init__.py

@@ -392,6 +392,32 @@ def spline_weighting(
     return torch.ops.pyg.spline_weighting(x, weight, basis, weight_index)
 
 
+def grid_cluster(
+    pos: Tensor,
+    size: Tensor,
+    start: Optional[Tensor] = None,
+    end: Optional[Tensor] = None,
+) -> Tensor:
+    r"""Clusters all points in :obj:`pos` into voxels of size :obj:`size`.
+
+    Each point is assigned a cluster index based on which voxel it falls into.
+    The voxel grid is defined by the :obj:`size` parameter and optionally
+    bounded by :obj:`start` and :obj:`end`.
+
+    Args:
+        pos: Point positions of shape :obj:`[N, D]`.
+        size: Voxel size in each dimension of shape :obj:`[D]`.
+        start: Start of the voxel grid in each dimension of shape :obj:`[D]`.
+            If :obj:`None`, uses the minimum of :obj:`pos`.
+        end: End of the voxel grid in each dimension of shape :obj:`[D]`.
+            If :obj:`None`, uses the maximum of :obj:`pos`.
+
+    Returns:
+        Cluster index for each point of shape :obj:`[N]`.
+    """
+    return torch.ops.pyg.grid_cluster(pos, size, start, end)
+
+
 __all__ = [
     'grouped_matmul',
     'segment_matmul',
@@ -403,4 +429,5 @@ def spline_weighting(
     'softmax_csr',
     'spline_basis',
     'spline_weighting',
+    'grid_cluster',
 ]

test/ops/test_cluster.py

@@ -0,0 +1,57 @@
+import pytest
+import torch
+
+import pyg_lib
+
+
+@pytest.mark.parametrize('dtype', [torch.float, torch.double])
+def test_grid_cluster_2d(dtype: torch.dtype) -> None:
+    pos = torch.tensor(
+        [[0.0, 0.0], [0.1, 0.1], [0.5, 0.5], [1.0, 1.0], [1.1, 1.1]],
+        dtype=dtype)
+    size = torch.tensor([0.5, 0.5], dtype=dtype)
+
+    out = pyg_lib.ops.grid_cluster(pos, size)
+
+    # Points (0,0) and (0.1,0.1) should be in the same cluster
+    assert out[0] == out[1]
+    # Point (0.5,0.5) should be in a different cluster from (0,0)
+    assert out[0] != out[2]
+    # Points (1.0,1.0) and (1.1,1.1) should be in the same cluster
+    assert out[3] == out[4]
+
+
+@pytest.mark.parametrize('dtype', [torch.float, torch.double])
+def test_grid_cluster_3d(dtype: torch.dtype) -> None:
+    pos = torch.tensor([[0.0, 0.0, 0.0], [0.1, 0.1, 0.1], [1.0, 1.0, 1.0]],
+                       dtype=dtype)
+    size = torch.tensor([0.5, 0.5, 0.5], dtype=dtype)
+
+    out = pyg_lib.ops.grid_cluster(pos, size)
+
+    assert out[0] == out[1]
+    assert out[0] != out[2]
+
+
+@pytest.mark.parametrize('dtype', [torch.float, torch.double])
+def test_grid_cluster_with_start_end(dtype: torch.dtype) -> None:
+    pos = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]], dtype=dtype)
+    size = torch.tensor([0.5, 0.5], dtype=dtype)
+    start = torch.tensor([0.0, 0.0], dtype=dtype)
+    end = torch.tensor([1.0, 1.0], dtype=dtype)
+
+    out = pyg_lib.ops.grid_cluster(pos, size, start, end)
+
+    assert out.shape == (3, )
+    assert out.dtype == torch.long
+
+
+def test_grid_cluster_defaults_match_explicit() -> None:
+    pos = torch.tensor([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]])
+    size = torch.tensor([0.5, 0.5])
+
+    out_default = pyg_lib.ops.grid_cluster(pos, size)
+    out_explicit = pyg_lib.ops.grid_cluster(pos, size, start=pos.min(0).values,
+                                            end=pos.max(0).values)
+
+    assert torch.equal(out_default, out_explicit)