Count 3D Array Element

原始题目：LeetGPU - Count 3D Array Element

题目描述

统计 $N \times M \times K$ 的 32 位整数三维数组中值为 $p$ 的元素数量。

实现要求

• 不允许使用外部库。
• solve 函数签名必须保持不变。

示例

Input:  input (2×2×2): [[[1,2],[3,1]],[[1,4],[5,1]]], p = 1
Output: 4

约束条件

• $1 \le N, M, K \le 100$，$1 \le input[i][j][k], p \le 100{,}000$。

解题思路

与一维/二维版本完全一致：逐元素比较 + 规约求和。三维索引线性化为 $i \cdot (M \cdot K) + j \cdot K + k$，每个线程处理一个元素。内存带宽受限内核。

代码实现

CUDA

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

Triton

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