| 104 | |
| 105 | template<typename T> |
| 106 | static void where(Param<uint> &out, Param<T> in) { |
| 107 | uint threads_x = nextpow2(std::max(32u, (uint)in.info.dims[0])); |
| 108 | threads_x = std::min(threads_x, THREADS_PER_BLOCK); |
| 109 | uint threads_y = THREADS_PER_BLOCK / threads_x; |
| 110 | |
| 111 | uint groups_x = divup((uint)in.info.dims[0], (uint)(threads_x * REPEAT)); |
| 112 | uint groups_y = divup(in.info.dims[1], threads_y); |
| 113 | |
| 114 | Param<uint> rtmp; |
| 115 | Param<uint> otmp; |
| 116 | rtmp.info.dims[0] = groups_x; |
| 117 | otmp.info.dims[0] = in.info.dims[0]; |
| 118 | rtmp.info.strides[0] = 1; |
| 119 | otmp.info.strides[0] = 1; |
| 120 | |
| 121 | for (int k = 1; k < 4; k++) { |
| 122 | rtmp.info.dims[k] = in.info.dims[k]; |
| 123 | rtmp.info.strides[k] = rtmp.info.strides[k - 1] * rtmp.info.dims[k - 1]; |
| 124 | |
| 125 | otmp.info.dims[k] = in.info.dims[k]; |
| 126 | otmp.info.strides[k] = otmp.info.strides[k - 1] * otmp.info.dims[k - 1]; |
| 127 | } |
| 128 | |
| 129 | uintl rtmp_elements = rtmp.info.strides[3] * rtmp.info.dims[3]; |
| 130 | uintl otmp_elements = otmp.info.strides[3] * otmp.info.dims[3]; |
| 131 | auto rtmp_alloc = memAlloc<uint>(rtmp_elements); |
| 132 | auto otmp_alloc = memAlloc<uint>(otmp_elements); |
| 133 | rtmp.data = rtmp_alloc.get(); |
| 134 | otmp.data = otmp_alloc.get(); |
| 135 | |
| 136 | scan_first_launcher<T, uint, af_notzero_t>( |
| 137 | otmp, rtmp, in, groups_x, groups_y, threads_x, false, true); |
| 138 | |
| 139 | // Linearize the dimensions and perform scan |
| 140 | Param<uint> ltmp = rtmp; |
| 141 | ltmp.info.dims[0] = rtmp_elements; |
| 142 | for (int k = 1; k < 4; k++) { |
| 143 | ltmp.info.dims[k] = 1; |
| 144 | ltmp.info.strides[k] = rtmp_elements; |
| 145 | } |
| 146 | |
| 147 | scan_first<uint, uint, af_add_t>(ltmp, ltmp, true); |
| 148 | |
| 149 | // Get output size and allocate output |
| 150 | uint total; |
| 151 | |
| 152 | getQueue() |
| 153 | .submit([&](sycl::handler &h) { |
| 154 | auto acc_in = rtmp.data->get_access(h, sycl::range{1}, |
| 155 | sycl::id{rtmp_elements - 1}); |
| 156 | h.copy(acc_in, &total); |
| 157 | }) |
| 158 | .wait(); |
| 159 | |
| 160 | auto out_alloc = memAlloc<uint>(std::max(1U, total)); |
| 161 | out.data = out_alloc.get(); |
| 162 | |
| 163 | out.info.dims[0] = total; |