K-Means Clustering

原始题目：LeetGPU - K-Means Clustering

题目描述

在 GPU 上实现 K-means 聚类算法（二维点）。给定数据点的 $x$ 和 $y$ 坐标数组、初始质心和迭代参数，将每个点分配到最近的质心，并迭代更新质心。最终质心和标签应存储在输出数组中。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。
• 最终结果必须存储在 labels、final_centroid_x 和 final_centroid_y 中。

示例

sample_size=4, k=2, max_iterations=10
data_x=[1,2,8,9], data_y=[1,1,8,8]
init_centroid_x=[1,8], init_centroid_y=[1,8]
→ 点0,1接近质心0, 点2,3接近质心1 → labels=[0,0,1,1]

约束条件

• $1 \le sample\_size \le 100{,}000$，$1 \le k \le 16$，$1 \le max\_iterations \le 100$。
• 坐标和质心为 32 位浮点数。

解题思路

K-means 的每次迭代包含两步：(1) 分配——每个点计算到所有 $k$ 个质心的距离并选最小的；(2) 更新——对每个 cluster 内的点坐标做平均。分配步骤是 $O(N \cdot k)$ 的，每个点完全独立，高度并行。更新步骤需要按 cluster 做坐标求和 + 计数，可以用 atomicAdd 或分块规约实现。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void kmeans_assign(const float* dx, const float* dy, const float* cx, const float* cy,
    int* labels, int N, int K) {
    int i=blockIdx.x*blockDim.x+threadIdx.x; if(i>=N)return;
    float xi=dx[i],yi=dy[i],min_d=INFINITY; int best=0;
    for(int k=0;k<K;k++){float dx=xi-cx[k],dy=yi-cy[k],d=dx*dx+dy*dy;if(d<min_d){min_d=d;best=k;}}
    labels[i]=best;
}
__global__ void kmeans_update(const float* dx, const float* dy, const int* labels,
    float* cx, float* cy, int* counts, int N, int K) {
    __shared__ float sx[16],sy[16];__shared__ int cnt[16];
    int tid=threadIdx.x; if(tid<K){sx[tid]=0;sy[tid]=0;cnt[tid]=0;}__syncthreads();
    for(int i=blockIdx.x*blockDim.x+tid;i<N;i+=gridDim.x*blockDim.x){
        int k=labels[i]; atomicAdd(&sx[k],dx[i]);atomicAdd(&sy[k],dy[i]);atomicAdd(&cnt[k],1);
    }
    __syncthreads();
    if(tid<K){atomicAdd(&cx[tid],sx[tid]);atomicAdd(&cy[tid],sy[tid]);atomicAdd(&counts[tid],cnt[tid]);}
}
extern "C" void solve(const float* dx,const float* dy,float* cx,float* cy,int* labels,int N,int K,int max_iter) {
    int *counts,*d_counts; cudaMalloc(&d_counts,K*sizeof(int));
    for(int iter=0;iter<max_iter;iter++){
        kmeans_assign<<<(N+255)/256,256>>>(dx,dy,cx,cy,labels,N,K);
        cudaMemset(d_counts,0,K*sizeof(int));
        float *ncx,*ncy; cudaMalloc(&ncx,K*sizeof(float));cudaMalloc(&ncy,K*sizeof(float));
        cudaMemset(ncx,0,K*sizeof(float));cudaMemset(ncy,0,K*sizeof(float));
        kmeans_update<<<min((N+255)/256,1024),256>>>(dx,dy,labels,ncx,ncy,d_counts,N,K);
        for(int k=0;k<K;k++){cx[k]=ncx[k]/max(d_counts[k],1);cy[k]=ncy[k]/max(d_counts[k],1);}
        cudaFree(ncx);cudaFree(ncy);
    }
    cudaFree(d_counts);cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def kmeans_assign(dx_ptr,dy_ptr,cx_ptr,cy_ptr,labels_ptr, N,K, BLOCK:tl.constexpr):
    i=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=i<N
    xi=tl.load(dx_ptr+i,mask=mask); yi=tl.load(dy_ptr+i,mask=mask)
    min_d=tl.full((BLOCK,),float('inf'),tl.float32); best=tl.zeros((BLOCK,),tl.int32)
    k_range=tl.arange(0,K)
    for k in range(K):
        cx=tl.load(cx_ptr+k);cy=tl.load(cy_ptr+k)
        d=(xi-cx)*(xi-cx)+(yi-cy)*(yi-cy)
        better=d<min_d; min_d=tl.where(better,d,min_d); best=tl.where(better,k,best)
    tl.store(labels_ptr+i,best,mask=mask)