#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";
}