| 194 | |
| 195 | template<typename Ti, typename To, af_op_t op> |
| 196 | void reduce_first_default(Param<To> out, Param<Ti> in, bool change_nan, |
| 197 | double nanval) { |
| 198 | uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0])); |
| 199 | threads_x = std::min(threads_x, creduce::THREADS_PER_BLOCK); |
| 200 | uint threads_y = creduce::THREADS_PER_BLOCK / threads_x; |
| 201 | |
| 202 | uint blocks_x = divup(in.info.dims[0], threads_x * creduce::REPEAT); |
| 203 | uint blocks_y = divup(in.info.dims[1], threads_y); |
| 204 | |
| 205 | Param<To> tmp = out; |
| 206 | bufptr<To> tmp_alloc; |
| 207 | if (blocks_x > 1) { |
| 208 | tmp_alloc = memAlloc<To>(blocks_x * in.info.dims[1] * in.info.dims[2] * |
| 209 | in.info.dims[3]); |
| 210 | tmp.data = tmp_alloc.get(); |
| 211 | |
| 212 | tmp.info.dims[0] = blocks_x; |
| 213 | for (int k = 1; k < 4; k++) tmp.info.strides[k] *= blocks_x; |
| 214 | } |
| 215 | |
| 216 | reduce_first_launcher_default<Ti, To, op>(tmp, in, blocks_x, blocks_y, |
| 217 | threads_x, change_nan, nanval); |
| 218 | |
| 219 | if (blocks_x > 1) { |
| 220 | // FIXME: Is there an alternative to the if condition? |
| 221 | if (op == af_notzero_t) { |
| 222 | reduce_first_launcher_default<To, To, af_add_t>( |
| 223 | out, tmp, 1, blocks_y, threads_x, change_nan, nanval); |
| 224 | } else { |
| 225 | reduce_first_launcher_default<To, To, op>( |
| 226 | out, tmp, 1, blocks_y, threads_x, change_nan, nanval); |
| 227 | } |
| 228 | } |
| 229 | } |
| 230 | |
| 231 | } // namespace kernel |
| 232 | } // namespace oneapi |