| 185 | |
| 186 | template<typename Ti, typename To, af_op_t op, int dim> |
| 187 | void reduce_dim_default(Param<To> out, Param<Ti> in, bool change_nan, |
| 188 | double nanval) { |
| 189 | uint threads_y = std::min(creduce::THREADS_Y, nextpow2(in.info.dims[dim])); |
| 190 | uint threads_x = creduce::THREADS_X; |
| 191 | |
| 192 | dim_t blocks_dim[] = {divup(in.info.dims[0], threads_x), in.info.dims[1], |
| 193 | in.info.dims[2], in.info.dims[3]}; |
| 194 | blocks_dim[dim] = divup(in.info.dims[dim], threads_y * creduce::REPEAT); |
| 195 | |
| 196 | Param<To> tmp = out; |
| 197 | bufptr<To> tmp_alloc; |
| 198 | if (blocks_dim[dim] > 1) { |
| 199 | tmp.info.dims[dim] = blocks_dim[dim]; |
| 200 | int tmp_elements = tmp.info.dims[0] * tmp.info.dims[1] * |
| 201 | tmp.info.dims[2] * tmp.info.dims[3]; |
| 202 | |
| 203 | tmp_alloc = memAlloc<To>(tmp_elements); |
| 204 | tmp.data = tmp_alloc.get(); |
| 205 | |
| 206 | tmp.info.dims[dim] = blocks_dim[dim]; |
| 207 | for (int k = dim + 1; k < 4; k++) |
| 208 | tmp.info.strides[k] *= blocks_dim[dim]; |
| 209 | } |
| 210 | |
| 211 | reduce_dim_launcher_default<Ti, To, op, dim>(tmp, in, threads_y, blocks_dim, |
| 212 | change_nan, nanval); |
| 213 | |
| 214 | if (blocks_dim[dim] > 1) { |
| 215 | blocks_dim[dim] = 1; |
| 216 | |
| 217 | if (op == af_notzero_t) { |
| 218 | reduce_dim_launcher_default<To, To, af_add_t, dim>( |
| 219 | out, tmp, threads_y, blocks_dim, change_nan, nanval); |
| 220 | } else { |
| 221 | reduce_dim_launcher_default<To, To, op, dim>( |
| 222 | out, tmp, threads_y, blocks_dim, change_nan, nanval); |
| 223 | } |
| 224 | } |
| 225 | } |
| 226 | |
| 227 | } // namespace kernel |
| 228 | } // namespace oneapi |