Skip to content

Commit cd01935

Browse files
denix56denix56
authored
Merge pull request #4 from denix56/codex/implement-fused-k-nn-in-triton
Compute split-N tile counts via heuristics
2 parents 02cc6ae + 4bdd9b0 commit cd01935

2 files changed

Lines changed: 321 additions & 13 deletions

File tree

‎torch_cluster/triton/_kernels.py‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -372,3 +372,239 @@ def _pairwise_distances(
372372
y.stride(0), y.stride(1), out.stride(0), out.stride(1),
373373
ptr_x, batch_y, USE_BATCH=use_batch, COSINE=cosine, EPS=eps)
374374
return out
375+
376+
377+
@triton.autotune(
378+
configs=[
379+
triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'BLOCK_K': 32},
380+
num_warps=2,
381+
num_stages=2),
382+
triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'BLOCK_K': 32},
383+
num_warps=2,
384+
num_stages=2),
385+
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 16, 'BLOCK_K': 32},
386+
num_warps=2,
387+
num_stages=2),
388+
triton.Config({'BLOCK_M': 16, 'BLOCK_N': 64, 'BLOCK_K': 32},
389+
num_warps=2,
390+
num_stages=2),
391+
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32},
392+
num_warps=4,
393+
num_stages=3),
394+
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'BLOCK_K': 32},
395+
num_warps=4,
396+
num_stages=3),
397+
triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_K': 32},
398+
num_warps=4,
399+
num_stages=3),
400+
triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32},
401+
num_warps=8,
402+
num_stages=4),
403+
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 16},
404+
num_warps=8,
405+
num_stages=2),
406+
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32},
407+
num_warps=8,
408+
num_stages=4),
409+
triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32},
410+
num_warps=8,
411+
num_stages=4),
412+
],
413+
key=['M', 'N', 'D'],
414+
)
415+
@triton.heuristics({
416+
'EVEN_M': lambda args: args['M'] % args['BLOCK_M'] == 0,
417+
'EVEN_N': lambda args: args['N'] % args['BLOCK_N'] == 0,
418+
'NUM_TILES': lambda args: triton.cdiv(args['N'], args['BLOCK_N']),
419+
'TILES_PER_SPLIT': lambda args: triton.cdiv(triton.cdiv(args['N'], args['BLOCK_N']), args['SPLIT_N']),
420+
})
421+
@triton.jit
422+
def _knn_stage1_kernel(
423+
x_ptr,
424+
y_ptr,
425+
part_i_ptr,
426+
part_d_ptr,
427+
M,
428+
N,
429+
D,
430+
stride_xm,
431+
stride_xd,
432+
stride_ym,
433+
stride_yd,
434+
stride_part_s,
435+
stride_part_m,
436+
stride_part_k,
437+
ptr_x_ptr,
438+
batch_y_ptr,
439+
USE_BATCH: tl.constexpr,
440+
COSINE: tl.constexpr,
441+
EPS: tl.constexpr,
442+
BLOCK_M: tl.constexpr,
443+
BLOCK_N: tl.constexpr,
444+
BLOCK_K: tl.constexpr,
445+
EVEN_M: tl.constexpr,
446+
EVEN_N: tl.constexpr,
447+
K: tl.constexpr,
448+
SPLIT_N: tl.constexpr,
449+
NUM_TILES: tl.constexpr,
450+
TILES_PER_SPLIT: tl.constexpr,
451+
):
452+
r"""Split-N partial top-k kernel for fused KNN."""
453+
pid_m = tl.program_id(0)
454+
pid_s = tl.program_id(1)
455+
456+
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
457+
mask_m = offs_m < M
458+
if EVEN_M:
459+
tl.multiple_of(offs_m, 8)
460+
tl.max_contiguous(offs_m, 8)
461+
462+
if USE_BATCH:
463+
batch_ids = tl.load(batch_y_ptr + offs_m, mask=mask_m, other=0)
464+
left = tl.load(ptr_x_ptr + batch_ids, mask=mask_m, other=0)
465+
right = tl.load(ptr_x_ptr + batch_ids + 1, mask=mask_m, other=0)
466+
467+
best_d = tl.full((BLOCK_M, K), float('inf'), tl.float32)
468+
best_i = tl.full((BLOCK_M, K), -1, tl.int32)
469+
k_ids = tl.arange(0, K)
470+
471+
# Split-N: each program processes a contiguous range of N-tiles.
472+
tile_start = pid_s * TILES_PER_SPLIT
473+
for t in tl.static_range(0, TILES_PER_SPLIT):
474+
tile_idx = tile_start + t
475+
in_range = tile_idx < NUM_TILES
476+
offs_n = tile_idx * BLOCK_N + tl.arange(0, BLOCK_N)
477+
mask_n = (offs_n < N) & in_range
478+
if EVEN_N:
479+
tl.multiple_of(offs_n, 8)
480+
tl.max_contiguous(offs_n, 8)
481+
482+
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
483+
x_sq = tl.zeros((BLOCK_N,), dtype=tl.float32)
484+
y_sq = tl.zeros((BLOCK_M,), dtype=tl.float32)
485+
486+
y_block_ptr = tl.make_block_ptr(
487+
base=y_ptr,
488+
shape=(M, D),
489+
strides=(stride_ym, stride_yd),
490+
offsets=(pid_m * BLOCK_M, 0),
491+
block_shape=(BLOCK_M, BLOCK_K),
492+
order=(1, 0),
493+
)
494+
x_block_ptr = tl.make_block_ptr(
495+
base=x_ptr,
496+
shape=(D, N),
497+
strides=(stride_xd, stride_xm),
498+
offsets=(0, tile_idx * BLOCK_N),
499+
block_shape=(BLOCK_K, BLOCK_N),
500+
order=(0, 1),
501+
)
502+
503+
for _ in range(0, D, BLOCK_K):
504+
y = tl.load(y_block_ptr,
505+
boundary_check=(0, 1),
506+
padding_option="zero")
507+
x = tl.load(x_block_ptr,
508+
boundary_check=(0, 1),
509+
padding_option="zero")
510+
acc += tl.dot(y, x, input_precision="ieee")
511+
y_sq += tl.sum(y * y, axis=1)
512+
x_sq += tl.sum(x * x, axis=0)
513+
y_block_ptr = tl.advance(y_block_ptr, (0, BLOCK_K))
514+
x_block_ptr = tl.advance(x_block_ptr, (BLOCK_K, 0))
515+
516+
if COSINE:
517+
inv_y = tl.rsqrt(y_sq + EPS)
518+
inv_x = tl.rsqrt(x_sq + EPS)
519+
dist = 1.0 - acc * (inv_y[:, None] * inv_x[None, :])
520+
else:
521+
dist = tl.fma(-2.0, acc, y_sq[:, None] + x_sq[None, :])
522+
523+
dist += offs_n[None, :].to(tl.float32) * EPS
524+
525+
# Apply bounds and batch masking; invalid candidates become +inf.
526+
valid = mask_m[:, None] & mask_n[None, :]
527+
if USE_BATCH:
528+
valid &= (offs_n[None, :] >= left[:, None]) & (offs_n[None, :] < right[:, None])
529+
dist = tl.where(valid, dist, float('inf'))
530+
531+
# Maintain per-row partial top-k by replacing the current worst entry.
532+
worst_d, worst_idx = tl.max(best_d,
533+
axis=1,
534+
return_indices=True,
535+
return_indices_tie_break_left=True)
536+
for n in tl.static_range(0, BLOCK_N):
537+
cand_d = dist[:, n]
538+
cand_i = offs_n[n]
539+
replace = cand_d < worst_d
540+
repl_mask = replace[:, None] & (k_ids[None, :] == worst_idx[:, None])
541+
best_d = tl.where(repl_mask, cand_d[:, None], best_d)
542+
best_i = tl.where(repl_mask, cand_i, best_i)
543+
worst_d, worst_idx = tl.max(best_d,
544+
axis=1,
545+
return_indices=True,
546+
return_indices_tie_break_left=True)
547+
548+
offs_k = tl.arange(0, K)
549+
part_ptrs = (pid_s * stride_part_s +
550+
offs_m[:, None] * stride_part_m +
551+
offs_k[None, :] * stride_part_k)
552+
mask_out = mask_m[:, None]
553+
tl.store(part_d_ptr + part_ptrs, best_d, mask=mask_out)
554+
tl.store(part_i_ptr + part_ptrs, best_i, mask=mask_out)
555+
556+
557+
@triton.autotune(
558+
configs=[
559+
triton.Config({'BLOCK_M': 16}, num_warps=2, num_stages=2),
560+
triton.Config({'BLOCK_M': 32}, num_warps=2, num_stages=2),
561+
triton.Config({'BLOCK_M': 64}, num_warps=4, num_stages=3),
562+
triton.Config({'BLOCK_M': 128}, num_warps=8, num_stages=4),
563+
],
564+
key=['M'],
565+
)
566+
@triton.jit
567+
def _knn_stage2_merge_kernel(
568+
part_i_ptr,
569+
part_d_ptr,
570+
out_i_ptr,
571+
M,
572+
stride_part_s,
573+
stride_part_m,
574+
stride_part_k,
575+
stride_out_m,
576+
stride_out_k,
577+
BLOCK_M: tl.constexpr,
578+
K: tl.constexpr,
579+
SPLIT_N: tl.constexpr,
580+
):
581+
r"""Merge split-N partial top-k buffers into final neighbors."""
582+
pid_m = tl.program_id(0)
583+
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
584+
mask_m = offs_m < M
585+
586+
best_d = tl.full((BLOCK_M, K), float('inf'), tl.float32)
587+
best_i = tl.full((BLOCK_M, K), -1, tl.int32)
588+
k_ids = tl.arange(0, K)
589+
590+
# Merge all split-N partial candidates into the final top-k per row.
591+
for s in tl.static_range(0, SPLIT_N):
592+
part_base = (s * stride_part_s + offs_m[:, None] * stride_part_m +
593+
k_ids[None, :] * stride_part_k)
594+
cand_d = tl.load(part_d_ptr + part_base, mask=mask_m[:, None], other=float('inf'))
595+
cand_i = tl.load(part_i_ptr + part_base, mask=mask_m[:, None], other=-1)
596+
for kk in tl.static_range(0, K):
597+
cand_dk = cand_d[:, kk]
598+
cand_ik = cand_i[:, kk]
599+
worst_d, worst_idx = tl.max(best_d,
600+
axis=1,
601+
return_indices=True,
602+
return_indices_tie_break_left=True)
603+
replace = cand_dk < worst_d
604+
repl_mask = replace[:, None] & (k_ids[None, :] == worst_idx[:, None])
605+
best_d = tl.where(repl_mask, cand_dk[:, None], best_d)
606+
best_i = tl.where(repl_mask, cand_ik, best_i)
607+
608+
out_ptrs = (offs_m[:, None] * stride_out_m +
609+
k_ids[None, :] * stride_out_k)
610+
tl.store(out_i_ptr + out_ptrs, best_i, mask=mask_m[:, None])

‎torch_cluster/triton/knn.py‎

Lines changed: 85 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,9 @@
33

44
import torch
55
from torch import Tensor
6+
import triton
67

7-
from ._kernels import _pairwise_distances
8+
from ._kernels import _knn_stage1_kernel, _knn_stage2_merge_kernel, _pairwise_distances
89

910

1011
def knn(
@@ -32,18 +33,89 @@ def knn(
3233
Returns:
3334
Tensor: Edge index with shape [2, M * k].
3435
"""
35-
# Compute distances, optionally with batch masking.
36-
distances = _pairwise_distances(x, y, batch_size=batch_size, batch_x=batch_x,
37-
batch_y=batch_y, cosine=cosine)
38-
39-
# Break distance ties by preferring smaller indices for deterministic
40-
# results that match the reference operator.
41-
indices = torch.argsort(distances, stable=True)[..., :k]
42-
#eps = torch.finfo(distances.dtype).eps
43-
#tie_break = torch.arange(x.size(0), device=distances.device,
44-
# dtype=distances.dtype)
45-
#distances = distances + tie_break[None, :] * eps
46-
#_, indices = torch.topk(distances, k=k, largest=False)
36+
indices = knn_fused_splitN(x, y, k, batch_x, batch_y, cosine, batch_size)
4737
row = torch.arange(y.size(0), device=y.device).repeat_interleave(k)
4838
col = indices.reshape(-1)
4939
return torch.stack([row, col], dim=0)
40+
41+
42+
def knn_fused_splitN(
43+
x: Tensor,
44+
y: Tensor,
45+
k: int,
46+
batch_x: Optional[Tensor] = None,
47+
batch_y: Optional[Tensor] = None,
48+
cosine: bool = False,
49+
batch_size: Optional[int] = None,
50+
split_n: int = 4,
51+
) -> Tensor:
52+
r"""Compute k-NN indices using a split-N fused Triton implementation."""
53+
if k > 32:
54+
distances = _pairwise_distances(x, y, batch_size=batch_size, batch_x=batch_x,
55+
batch_y=batch_y, cosine=cosine)
56+
indices = torch.argsort(distances, stable=True)[..., :k]
57+
return indices
58+
59+
use_batch = (batch_size or 1) > 1
60+
if use_batch:
61+
assert batch_x is not None
62+
assert batch_y is not None
63+
arange = torch.arange((batch_size or 1) + 1, device=x.device)
64+
ptr_x = torch.bucketize(arange, batch_x)
65+
batch_y = batch_y.contiguous()
66+
else:
67+
ptr_x = x
68+
batch_y = y
69+
70+
M, N = y.size(0), x.size(0)
71+
part_i = torch.empty((split_n, M, k), device=x.device, dtype=torch.int32)
72+
part_d = torch.empty((split_n, M, k), device=x.device, dtype=torch.float32)
73+
out_i = torch.empty((M, k), device=x.device, dtype=torch.int32)
74+
eps = torch.finfo(part_d.dtype).eps
75+
76+
def stage1_grid(meta):
77+
return (triton.cdiv(M, meta['BLOCK_M']), split_n)
78+
79+
# Stage 1: split N into partitions and keep per-row partial top-k in registers.
80+
_knn_stage1_kernel[stage1_grid](
81+
x,
82+
y,
83+
part_i,
84+
part_d,
85+
M,
86+
N,
87+
x.size(1),
88+
x.stride(0),
89+
x.stride(1),
90+
y.stride(0),
91+
y.stride(1),
92+
part_i.stride(0),
93+
part_i.stride(1),
94+
part_i.stride(2),
95+
ptr_x,
96+
batch_y,
97+
USE_BATCH=use_batch,
98+
COSINE=cosine,
99+
EPS=eps,
100+
K=k,
101+
SPLIT_N=split_n,
102+
)
103+
104+
# Stage 2: merge partial candidates from all splits into the final top-k.
105+
def stage2_grid(meta):
106+
return (triton.cdiv(M, meta['BLOCK_M']),)
107+
108+
_knn_stage2_merge_kernel[stage2_grid](
109+
part_i,
110+
part_d,
111+
out_i,
112+
M,
113+
part_i.stride(0),
114+
part_i.stride(1),
115+
part_i.stride(2),
116+
out_i.stride(0),
117+
out_i.stride(1),
118+
K=k,
119+
SPLIT_N=split_n,
120+
)
121+
return out_i

0 commit comments

Comments
 (0)