Count Array Element

原始题目：LeetGPU - Count Array Element

题目描述

编写一个 GPU 程序，统计 32 位整数数组中值为 $k$ 的元素数量。给定长度为 $N$ 的输入数组 input 和目标整数 $k$，输出等于 $k$ 的元素个数。

实现要求

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

示例

示例 1

Input:  [1, 2, 3, 4, 1], k = 1
Output: 2

示例 2

Input:  [5, 10, 5, 2], k = 11
Output: 0

约束条件

• $1 \le N \le 100{,}000{,}000$。
• $1 \le input[i], k \le 100{,}000$。
• 性能测试在 $k = 501{,}010,\ N = 100{,}000{,}000$ 的规模下进行。

解题思路

计数操作 = 逐元素比较（map）+ 规约求和（reduce）。每个线程比较一个元素是否等于 $k$（产生 0 或 1），然后用分块规约 + warp shuffle 求和。与纯规约相比多了一步比较，但本质上还是内存带宽受限。对于极大的 $N$，可以使用 atomicAdd 让每个 block 贡献局部计数到全局变量，减少规约层数。

代码实现

CUDA

#include <cuda_runtime.h>
__global__ void count_kernel(const int* input, int* output, int N, int k) {
    __shared__ int bc; if(threadIdx.x==0)bc=0; __syncthreads();
    for(int i=blockIdx.x*blockDim.x+threadIdx.x;i<N;i+=gridDim.x*blockDim.x)
        if(input[i]==k)atomicAdd(&bc,1);
    __syncthreads(); if(threadIdx.x==0)atomicAdd(output,bc);
}
extern "C" void solve(const int* input, int* output, int N, int k) {
    count_kernel<<<min((N+255)/256,1024),256>>>(input,output,N,k);
    cudaDeviceSynchronize();
}

Triton

import triton, triton.language as tl
@triton.jit
def count_kernel(input_ptr,output_ptr,N,k,BLOCK:tl.constexpr):
    idx=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK); mask=idx<N
    x=tl.load(input_ptr+idx,mask=mask)
    tl.atomic_add(output_ptr,tl.sum(tl.where(x==k,1,0),axis=0))