wehub-resource-sync
181 lines
6.6 KiB
C++
181 lines
6.6 KiB
C++
#include "opencl_source_map.hpp"
|
|
namespace MNN {
|
|
const char* topkv2 =
|
|
"#ifdef MNN_SUPPORT_FP16\n"
|
|
"#pragma OPENCL EXTENSION cl_khr_fp16 : enable\n"
|
|
"#endif\n"
|
|
"#define GLOBAL_SIZE_3_DIMS "
|
|
" __private const int global_size_dim0,__private const int global_size_dim1,__private const int global_size_dim2,\n"
|
|
"#define DEAL_NON_UNIFORM_DIM3(input1, input2, input3) "
|
|
" if (input1 >= global_size_dim0 || input2 >= global_size_dim1 || input3 >= global_size_dim2) { "
|
|
" return; "
|
|
" }\n"
|
|
"__constant sampler_t SAMPLER=CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
|
|
"#define THREAD_NUMBER 128\n"
|
|
"#define LOCAL_K 8\n"
|
|
"#define CANDIDATE_NUMBER (THREAD_NUMBER*LOCAL_K)\n"
|
|
"#ifdef IS_INT\n"
|
|
"typedef int DTYPE;\n"
|
|
"typedef int4 DTYPE4;\n"
|
|
"#define READ_INPUT(img,coord) read_imagei(img,SAMPLER,coord)\n"
|
|
"#define WRITE_OUTPUT(img,coord,val) write_imagei(img,coord,val)\n"
|
|
"#else\n"
|
|
"typedef FLOAT DTYPE;\n"
|
|
"typedef FLOAT4 DTYPE4;\n"
|
|
"#define READ_INPUT(img,coord) RI_F(img,SAMPLER,coord)\n"
|
|
"#define WRITE_OUTPUT(img,coord,val) WI_F(img,coord,val)\n"
|
|
"#endif\n"
|
|
"inline bool afterAsc(DTYPE aValue,int aIndex,DTYPE bValue,int bIndex) {\n"
|
|
" if (aValue>bValue) return true;\n"
|
|
" if (aValue<bValue) return false;\n"
|
|
" return aIndex>bIndex;\n"
|
|
"}\n"
|
|
"inline bool better(DTYPE aValue,int aIndex,DTYPE bValue,int bIndex) {\n"
|
|
" if (bIndex<0) return true;\n"
|
|
" if (aIndex<0) return false;\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" if (aValue>bValue) return true;\n"
|
|
" if (aValue<bValue) return false;\n"
|
|
"#else\n"
|
|
" if (aValue<bValue) return true;\n"
|
|
" if (aValue>bValue) return false;\n"
|
|
"#endif\n"
|
|
" return aIndex<bIndex;\n"
|
|
"}\n"
|
|
"// Sort along channel dimension (last dim for 1D/2D tensors)\n"
|
|
"// Image layout: image_x=w+c4_idx*width,image_y=n*height+h\n"
|
|
"// For 2D [M,N]: width=1,height=1,numRows=M,rowSize=N\n"
|
|
"// For 1D [N]: width=1,height=1,numRows=1,rowSize=N\n"
|
|
"__kernel void topkv2_channel(GLOBAL_SIZE_3_DIMS\n"
|
|
" __read_only image2d_t input,\n"
|
|
" __write_only image2d_t outputValue,\n"
|
|
" __write_only image2d_t outputIndex,\n"
|
|
" __private const int rowSize,\n"
|
|
" __private const int k,\n"
|
|
" __private const int width,\n"
|
|
" __private const int channelBlocks\n"
|
|
") {\n"
|
|
" const int gid0=get_global_id(0);\n"
|
|
" const int gid1=get_global_id(1);\n"
|
|
" const int gid2=get_global_id(2);\n"
|
|
" DEAL_NON_UNIFORM_DIM3(gid0,gid1,gid2);\n"
|
|
" const int tid=get_local_id(0);\n"
|
|
" const int row=get_group_id(1);\n"
|
|
" // Decompose row into image coordinates\n"
|
|
" // row=image_y*width+w_pos (for sort-along-C)\n"
|
|
" const int w_pos=row % width;\n"
|
|
" const int image_y=row/width;\n"
|
|
" // Initialize per-thread top-k\n"
|
|
"#ifdef IS_INT\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" DTYPE initVal=(DTYPE)(-2147483647-1);\n"
|
|
"#else\n"
|
|
" DTYPE initVal=(DTYPE)(2147483647);\n"
|
|
"#endif\n"
|
|
"#else\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" DTYPE initVal=(DTYPE)(-FLT_MAX);\n"
|
|
"#else\n"
|
|
" DTYPE initVal=(DTYPE)(FLT_MAX);\n"
|
|
"#endif\n"
|
|
"#endif\n"
|
|
" DTYPE localValue[LOCAL_K];\n"
|
|
" int localIndex[LOCAL_K];\n"
|
|
" for (int i=0; i<LOCAL_K; ++i) {\n"
|
|
" localValue[i]=initVal;\n"
|
|
" localIndex[i]=-1;\n"
|
|
" }\n"
|
|
" // Read input image and build per-thread top-k\n"
|
|
" for (int c4=tid; c4<channelBlocks; c4 += THREAD_NUMBER) {\n"
|
|
" DTYPE4 pixel=READ_INPUT(input,(int2)(c4*width+w_pos,image_y));\n"
|
|
" int baseC=c4 << 2;\n"
|
|
" DTYPE vals[4];\n"
|
|
" vals[0]=pixel.x;\n"
|
|
" vals[1]=pixel.y;\n"
|
|
" vals[2]=pixel.z;\n"
|
|
" vals[3]=pixel.w;\n"
|
|
" for (int comp=0; comp<4 && baseC+comp<rowSize; ++comp) {\n"
|
|
" DTYPE value=vals[comp];\n"
|
|
" int idx=baseC+comp;\n"
|
|
" if (!better(value,idx,localValue[LOCAL_K-1],localIndex[LOCAL_K-1])) continue;\n"
|
|
" int insertPos=LOCAL_K;\n"
|
|
" for (int j=0; j<LOCAL_K; ++j) {\n"
|
|
" if (better(value,idx,localValue[j],localIndex[j])) {\n"
|
|
" insertPos=j;\n"
|
|
" break;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" if (insertPos >= LOCAL_K) continue;\n"
|
|
" for (int j=LOCAL_K-1; j>insertPos; --j) {\n"
|
|
" localValue[j]=localValue[j-1];\n"
|
|
" localIndex[j]=localIndex[j-1];\n"
|
|
" }\n"
|
|
" localValue[insertPos]=value;\n"
|
|
" localIndex[insertPos]=idx;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" // Scatter to shared memory\n"
|
|
" __local DTYPE sharedValue[CANDIDATE_NUMBER];\n"
|
|
" __local int sharedIndex[CANDIDATE_NUMBER];\n"
|
|
" const int base=tid*LOCAL_K;\n"
|
|
" for (int i=0; i<LOCAL_K; ++i) {\n"
|
|
" sharedValue[base+i]=localValue[i];\n"
|
|
" sharedIndex[base+i]=localIndex[i];\n"
|
|
" }\n"
|
|
" barrier(CLK_LOCAL_MEM_FENCE);\n"
|
|
" // Bitonic sort in shared memory\n"
|
|
" for (uint size=2; size <= CANDIDATE_NUMBER; size <<= 1) {\n"
|
|
" for (uint stride=size >> 1; stride>0; stride >>= 1) {\n"
|
|
" for (uint sIdx=tid; sIdx<CANDIDATE_NUMBER; sIdx += THREAD_NUMBER) {\n"
|
|
" const uint ixj=sIdx ^ stride;\n"
|
|
" if (ixj <= sIdx) continue;\n"
|
|
" bool up=((sIdx & size) == 0);\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" up=!up;\n"
|
|
"#endif\n"
|
|
" const bool after=afterAsc(sharedValue[sIdx],sharedIndex[sIdx],\n"
|
|
" sharedValue[ixj],sharedIndex[ixj]);\n"
|
|
" if (up == after) {\n"
|
|
" DTYPE tValue=sharedValue[sIdx];\n"
|
|
" sharedValue[sIdx]=sharedValue[ixj];\n"
|
|
" sharedValue[ixj]=tValue;\n"
|
|
" int tIndex=sharedIndex[sIdx];\n"
|
|
" sharedIndex[sIdx]=sharedIndex[ixj];\n"
|
|
" sharedIndex[ixj]=tIndex;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" barrier(CLK_LOCAL_MEM_FENCE);\n"
|
|
" }\n"
|
|
" }\n"
|
|
" // Thread 0 writes output\n"
|
|
" if (tid == 0) {\n"
|
|
" const int realK=min(k,rowSize);\n"
|
|
" const int outChannelBlocks=(k+3) >> 2;\n"
|
|
" for (int c4=0; c4<outChannelBlocks; ++c4) {\n"
|
|
" DTYPE4 valuePixel=(DTYPE4)(0);\n"
|
|
" int4 indexPixel=(int4)(0);\n"
|
|
" int baseIdx=c4 << 2;\n"
|
|
" if (baseIdx<realK) {\n"
|
|
" valuePixel.x=sharedValue[baseIdx];\n"
|
|
" indexPixel.x=sharedIndex[baseIdx];\n"
|
|
" }\n"
|
|
" if (baseIdx+1<realK) {\n"
|
|
" valuePixel.y=sharedValue[baseIdx+1];\n"
|
|
" indexPixel.y=sharedIndex[baseIdx+1];\n"
|
|
" }\n"
|
|
" if (baseIdx+2<realK) {\n"
|
|
" valuePixel.z=sharedValue[baseIdx+2];\n"
|
|
" indexPixel.z=sharedIndex[baseIdx+2];\n"
|
|
" }\n"
|
|
" if (baseIdx+3<realK) {\n"
|
|
" valuePixel.w=sharedValue[baseIdx+3];\n"
|
|
" indexPixel.w=sharedIndex[baseIdx+3];\n"
|
|
" }\n"
|
|
" int outX=c4*width+w_pos;\n"
|
|
" WRITE_OUTPUT(outputValue,(int2)(outX,image_y),valuePixel);\n"
|
|
" write_imagei(outputIndex,(int2)(outX,image_y),indexPixel);\n"
|
|
" }\n"
|
|
" }\n"
|
|
"}\n";
|
|
}
|