| 71 | |
| 72 | template<typename Ti, typename To, af_op_t op> |
| 73 | static void scan_first(Param<To> out, CParam<Ti> in, bool inclusive_scan) { |
| 74 | uint threads_x = nextpow2(std::max(32u, (uint)out.dims[0])); |
| 75 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 76 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 77 | |
| 78 | uint blocks_x = divup(out.dims[0], threads_x * REPEAT); |
| 79 | uint blocks_y = divup(out.dims[1], threads_y); |
| 80 | |
| 81 | if (blocks_x == 1) { |
| 82 | scan_first_launcher<Ti, To, op>(out, out, in, blocks_x, blocks_y, |
| 83 | threads_x, true, inclusive_scan); |
| 84 | |
| 85 | } else { |
| 86 | Param<To> tmp = out; |
| 87 | |
| 88 | tmp.dims[0] = blocks_x; |
| 89 | tmp.strides[0] = 1; |
| 90 | for (int k = 1; k < 4; k++) |
| 91 | tmp.strides[k] = tmp.strides[k - 1] * tmp.dims[k - 1]; |
| 92 | |
| 93 | int tmp_elements = tmp.strides[3] * tmp.dims[3]; |
| 94 | auto tmp_alloc = memAlloc<To>(tmp_elements); |
| 95 | tmp.ptr = tmp_alloc.get(); |
| 96 | |
| 97 | scan_first_launcher<Ti, To, op>(out, tmp, in, blocks_x, blocks_y, |
| 98 | threads_x, false, inclusive_scan); |
| 99 | |
| 100 | // FIXME: Is there an alternative to the if condition ? |
| 101 | if (op == af_notzero_t) { |
| 102 | scan_first_launcher<To, To, af_add_t>(tmp, tmp, tmp, 1, blocks_y, |
| 103 | threads_x, true, true); |
| 104 | } else { |
| 105 | scan_first_launcher<To, To, op>(tmp, tmp, tmp, 1, blocks_y, |
| 106 | threads_x, true, true); |
| 107 | } |
| 108 | |
| 109 | bcast_first_launcher<To, op>(out, tmp, blocks_x, blocks_y, threads_x, |
| 110 | inclusive_scan); |
| 111 | } |
| 112 | } |
| 113 | |
| 114 | } // namespace kernel |
| 115 | } // namespace cuda |
no test coverage detected