Merge branch 'master' into triton

Author: denix56

torch_cluster/fps.py

@@ -5,6 +5,32 @@
 
 import torch_cluster.typing
 
+@torch.library.register_fake("torch_cluster::fps")
+def _(src, ptr, ratio, random_start = True):
+    torch._check(src.device == ptr.device)
+    torch._check(ptr.ndim == 1)
+
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return ptr.new_empty((nnz,))
+
+@torch.library.register_fake("torch_cluster::fps")
+def _(src, ptr, ratio, random_start=True):
+    torch._check(src.device == ptr.device)
+    torch._check(ptr.ndim == 1)
+
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return ptr.new_empty((nnz,))
+
+@torch.library.register_fake("torch_cluster::fps")
+def _(src, ptr, ratio, random_start=True):
+    torch._check(src.device == ptr.device)
+    torch._check(ptr.ndim == 1)
+
+    ctx = torch.library.get_ctx()
+    nnz = ctx.new_dynamic_size()
+    return ptr.new_empty((nnz,))
 
 @torch.library.register_fake("torch_cluster::fps")
 def _(src, ptr, ratio, random_start=True):

torch_cluster/triton/_kernels.py

@@ -99,16 +99,16 @@ def _nearest_kernel(
     offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)  # X indices.
     mask_x = offs_n < N  # Valid x rows.
     if EVEN_N:
-        tl.multiple_of(offs_n, 8)
-        tl.max_contiguous(offs_n, 8)
+        tl.multiple_of(offs_n, 8)  # Hint contiguous access.
+        tl.max_contiguous(offs_n, 8)  # Hint vectorization.
 
     best_dist = tl.full((BLOCK_N,), float('inf'), tl.float32)  # Best dist.
     best_idx = tl.zeros((BLOCK_N,), dtype=tl.int32)  # Best y index.
 
     if USE_BATCH:
-        batch_id = tl.load(batch_x_ptr + offs_n, mask=mask_x, other=0)
-        left = tl.load(ptr_y_ptr + batch_id, mask=mask_x, other=0)
-        right = tl.load(ptr_y_ptr + batch_id + 1, mask=mask_x, other=0)
+        batch_id = tl.load(batch_x_ptr + offs_n, mask=mask_x, other=0)  # Batch id per x.
+        left = tl.load(ptr_y_ptr + batch_id, mask=mask_x, other=0)  # y range start.
+        right = tl.load(ptr_y_ptr + batch_id + 1, mask=mask_x, other=0)  # y range end.
     else:
         left = 0  # Full y range.
         right = M
@@ -194,7 +194,7 @@ def _nearest_kernel(
             inv_y = tl.rsqrt(y_sq + EPS)  # 1/||y||.
             dist = 1.0 - acc * (inv_y[:, None] * inv_x[None, :])
         else:
-            dist = tl.fma(-2.0, acc, y_sq[:, None] + x_sq[None, :])
+            dist = tl.fma(-2.0, acc, y_sq[:, None] + x_sq[None, :])  # L2^2 distance.
 
         if full_y:
             valid = tl.broadcast_to(mask_x[None, :], (BLOCK_M, BLOCK_N))
@@ -218,7 +218,7 @@ def _nearest_kernel(
         best_dist = tl.where(better, block_min, best_dist)
         best_idx = tl.where(better, block_idx, best_idx)
 
-    tl.store(out_ptr + offs_n, best_idx.to(tl.int64), mask=mask_x)
+    tl.store(out_ptr + offs_n, best_idx.to(tl.int64), mask=mask_x)  # Write output.
 
 
 @triton.autotune(

torch_cluster/triton/knn.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
+import math
+import triton
 from typing import Optional
 
 import triton