wehub-resource-sync
157 lines
4.3 KiB
C++
157 lines
4.3 KiB
C++
#include "opencl_source_map.hpp"
|
|
namespace MNN {
|
|
#ifndef MNN_OPENCL_BUFFER_CLOSED
|
|
const char* topkv2_buf =
|
|
"#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"
|
|
"#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"
|
|
"#else\n"
|
|
"typedef FLOAT DTYPE;\n"
|
|
"#endif\n"
|
|
"inline bool afterAsc(DTYPE aValue,int aIndex,DTYPE bValue,int bIndex) {\n"
|
|
" if (aValue>bValue) {\n"
|
|
" return true;\n"
|
|
" }\n"
|
|
" if (aValue<bValue) {\n"
|
|
" return false;\n"
|
|
" }\n"
|
|
" return aIndex>bIndex;\n"
|
|
"}\n"
|
|
"inline bool better(DTYPE aValue,int aIndex,DTYPE bValue,int bIndex) {\n"
|
|
" if (bIndex<0) {\n"
|
|
" return true;\n"
|
|
" }\n"
|
|
" if (aIndex<0) {\n"
|
|
" return false;\n"
|
|
" }\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" if (aValue>bValue) {\n"
|
|
" return true;\n"
|
|
" }\n"
|
|
" if (aValue<bValue) {\n"
|
|
" return false;\n"
|
|
" }\n"
|
|
"#else\n"
|
|
" if (aValue<bValue) {\n"
|
|
" return true;\n"
|
|
" }\n"
|
|
" if (aValue>bValue) {\n"
|
|
" return false;\n"
|
|
" }\n"
|
|
"#endif\n"
|
|
" return aIndex<bIndex;\n"
|
|
"}\n"
|
|
"__kernel void topkv2_buf(GLOBAL_SIZE_3_DIMS\n"
|
|
" __global DTYPE *outValue,\n"
|
|
" __global int *outIndex,\n"
|
|
" __global const DTYPE *inValue,\n"
|
|
" __private const int rowSize,\n"
|
|
" __private const int k,\n"
|
|
" __private const int numRows) {\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"
|
|
" \n"
|
|
" DEAL_NON_UNIFORM_DIM3(gid0,gid1,gid2);\n"
|
|
" \n"
|
|
" const int tid=get_local_id(0);\n"
|
|
" const int row=get_group_id(1);\n"
|
|
" if (tid >= THREAD_NUMBER || row >= numRows) {\n"
|
|
" return;\n"
|
|
" }\n"
|
|
"#ifdef IS_INT\n"
|
|
" const DTYPE initWorst=(DTYPE)(2147483647);\n"
|
|
" const DTYPE initBestWorst=(DTYPE)(-2147483648);\n"
|
|
"#else\n"
|
|
" const DTYPE initWorst=(DTYPE)(FLT_MAX);\n"
|
|
" const DTYPE initBestWorst=(DTYPE)(-FLT_MAX);\n"
|
|
"#endif\n"
|
|
" DTYPE localValue[LOCAL_K];\n"
|
|
" int localIndex[LOCAL_K];\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" for (uint i=0; i<LOCAL_K; ++i) {\n"
|
|
" localValue[i]=initBestWorst;\n"
|
|
" localIndex[i]=-1;\n"
|
|
" }\n"
|
|
"#else\n"
|
|
" for (uint i=0; i<LOCAL_K; ++i) {\n"
|
|
" localValue[i]=initWorst;\n"
|
|
" localIndex[i]=-1;\n"
|
|
" }\n"
|
|
"#endif\n"
|
|
" const __global DTYPE *rowIn=inValue+row*rowSize;\n"
|
|
" for (int i=tid; i<rowSize; i += THREAD_NUMBER) {\n"
|
|
" const DTYPE value=rowIn[i];\n"
|
|
" if (!better(value,i,localValue[LOCAL_K-1],localIndex[LOCAL_K-1])) {\n"
|
|
" continue;\n"
|
|
" }\n"
|
|
" uint insertPos=LOCAL_K;\n"
|
|
" for (uint j=0; j<LOCAL_K; ++j) {\n"
|
|
" if (better(value,i,localValue[j],localIndex[j])) {\n"
|
|
" insertPos=j;\n"
|
|
" break;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" if (insertPos >= LOCAL_K) {\n"
|
|
" continue;\n"
|
|
" }\n"
|
|
" for (uint 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]=i;\n"
|
|
" }\n"
|
|
" __local DTYPE sharedValue[CANDIDATE_NUMBER];\n"
|
|
" __local int sharedIndex[CANDIDATE_NUMBER];\n"
|
|
" const uint base=tid*LOCAL_K;\n"
|
|
" for (uint 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"
|
|
" for (uint size=2; size <= CANDIDATE_NUMBER; size <<= 1) {\n"
|
|
" for (uint stride=size >> 1; stride>0; stride >>= 1) {\n"
|
|
" for (uint idx=tid; idx<CANDIDATE_NUMBER; idx += THREAD_NUMBER) {\n"
|
|
" const uint ixj=idx ^ stride;\n"
|
|
" if (ixj <= idx) {\n"
|
|
" continue;\n"
|
|
" }\n"
|
|
" bool up=((idx & size) == 0);\n"
|
|
"#ifdef SORT_DESC\n"
|
|
" up=!up;\n"
|
|
"#endif\n"
|
|
" const bool after=afterAsc(sharedValue[idx],sharedIndex[idx],sharedValue[ixj],sharedIndex[ixj]);\n"
|
|
" if (up == after) {\n"
|
|
" const DTYPE tValue=sharedValue[idx];\n"
|
|
" sharedValue[idx]=sharedValue[ixj];\n"
|
|
" sharedValue[ixj]=tValue;\n"
|
|
" const int tIndex=sharedIndex[idx];\n"
|
|
" sharedIndex[idx]=sharedIndex[ixj];\n"
|
|
" sharedIndex[ixj]=tIndex;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" barrier(CLK_LOCAL_MEM_FENCE);\n"
|
|
" }\n"
|
|
" }\n"
|
|
" if (tid == 0) {\n"
|
|
" __global DTYPE *rowOut=outValue+row*k;\n"
|
|
" __global int *rowIdx=outIndex+row*k;\n"
|
|
" const int realK=min(k,rowSize);\n"
|
|
" for (int i=0; i<realK; ++i) {\n"
|
|
" rowOut[i]=sharedValue[i];\n"
|
|
" rowIdx[i]=sharedIndex[i];\n"
|
|
" }\n"
|
|
" }\n"
|
|
"}\n"
|
|
;
|
|
#endif
|
|
}
|