| 119 | |
| 120 | template<typename T> |
| 121 | void topkDim0(Param<T> ovals, Param<uint> oidxs, CParam<T> ivals, const int k, |
| 122 | const af::topkFunction order) { |
| 123 | dim3 threads(TOPK_THRDS_PER_BLK, 1); |
| 124 | const int thrdLoad = TOPK_IDX_THRD_LOAD; |
| 125 | |
| 126 | int numBlocksX = divup(ivals.dims[0], threads.x * thrdLoad); |
| 127 | dim3 blocks(numBlocksX, ivals.dims[1] * ivals.dims[3], ivals.dims[2]); |
| 128 | |
| 129 | // The algorithm is to iteratively find top k elements among each block |
| 130 | // of threads until there is only one block to launch. |
| 131 | // The additional memory used for values and indices is allocated only |
| 132 | // before the first iteration and reused for further iterations. |
| 133 | |
| 134 | // Temporary storage allocation for iterations |
| 135 | Array<T> tvals = createEmptyArray<T>(dim4()); |
| 136 | Array<uint> tidxs = createEmptyArray<uint>(dim4()); |
| 137 | |
| 138 | if (numBlocksX > 1) { |
| 139 | tvals = createEmptyArray<T>(dim4(k * numBlocksX, ivals.dims[1])); |
| 140 | // TODO(umar): this can be smaller because the first iteration is not |
| 141 | // reading this array. |
| 142 | tidxs = createEmptyArray<uint>(dim4(k * numBlocksX, ivals.dims[1])); |
| 143 | } |
| 144 | |
| 145 | int prevBlocksX = 1; |
| 146 | |
| 147 | CParam<T> iivals = ivals; |
| 148 | CParam<uint> iiidxs = tidxs; |
| 149 | |
| 150 | int dims0 = tvals.dims()[0]; |
| 151 | bool first_run = true; |
| 152 | do { |
| 153 | if (blocks.x == 1) { |
| 154 | tvals = createParamArray(ovals, false); |
| 155 | tidxs = createParamArray(oidxs, false); |
| 156 | } |
| 157 | |
| 158 | if (first_run) { |
| 159 | // Launch topk which doesn't read the indice values from global |
| 160 | // memory |
| 161 | CUDA_LAUNCH((kerTopkDim0<T, false>), blocks, threads, tvals, tidxs, |
| 162 | iivals, iiidxs, k, order, ivals.dims[1]); |
| 163 | first_run = false; |
| 164 | } else { |
| 165 | CUDA_LAUNCH((kerTopkDim0<T, true>), blocks, threads, tvals, tidxs, |
| 166 | iivals, iiidxs, k, order, ivals.dims[1]); |
| 167 | } |
| 168 | |
| 169 | POST_LAUNCH_CHECK(); |
| 170 | |
| 171 | prevBlocksX = blocks.x; |
| 172 | blocks.x = divup(dims0, threads.x * thrdLoad); |
| 173 | |
| 174 | // set output of current iteration as input for the next iteration |
| 175 | iivals = tvals; |
| 176 | iiidxs = tidxs; |
| 177 | |
| 178 | dims0 = blocks.x * k; |
nothing calls this directly
no test coverage detected