| 346 | } |
| 347 | |
| 348 | class csrCalcOutNNZKernel { |
| 349 | public: |
| 350 | csrCalcOutNNZKernel(write_accessor<unsigned> nnzc, |
| 351 | write_accessor<int> oRowIdx, unsigned M, |
| 352 | read_accessor<int> lRowIdx, read_accessor<int> lColIdx, |
| 353 | read_accessor<int> rRowIdx, read_accessor<int> rColIdx, |
| 354 | sycl::local_accessor<unsigned, 1> blkNNZ) |
| 355 | : nnzc_(nnzc) |
| 356 | , oRowIdx_(oRowIdx) |
| 357 | , M_(M) |
| 358 | , lRowIdx_(lRowIdx) |
| 359 | , lColIdx_(lColIdx) |
| 360 | , rRowIdx_(rRowIdx) |
| 361 | , rColIdx_(rColIdx) |
| 362 | , blkNNZ_(blkNNZ) {} |
| 363 | |
| 364 | void operator()(sycl::nd_item<1> it) const { |
| 365 | sycl::group g = it.get_group(); |
| 366 | |
| 367 | const uint row = it.get_global_id(0); |
| 368 | const uint tid = it.get_local_id(0); |
| 369 | |
| 370 | const bool valid = row < M_; |
| 371 | |
| 372 | const uint lEnd = (valid ? lRowIdx_[row + 1] : 0); |
| 373 | const uint rEnd = (valid ? rRowIdx_[row + 1] : 0); |
| 374 | |
| 375 | blkNNZ_[tid] = 0; |
| 376 | it.barrier(); |
| 377 | |
| 378 | uint l = (valid ? lRowIdx_[row] : 0); |
| 379 | uint r = (valid ? rRowIdx_[row] : 0); |
| 380 | uint nnz = 0; |
| 381 | while (l < lEnd && r < rEnd) { |
| 382 | uint lci = lColIdx_[l]; |
| 383 | uint rci = rColIdx_[r]; |
| 384 | l += (lci <= rci); |
| 385 | r += (lci >= rci); |
| 386 | nnz++; |
| 387 | } |
| 388 | nnz += (lEnd - l); |
| 389 | nnz += (rEnd - r); |
| 390 | |
| 391 | blkNNZ_[tid] = nnz; |
| 392 | it.barrier(); |
| 393 | |
| 394 | if (valid) oRowIdx_[row + 1] = nnz; |
| 395 | |
| 396 | for (uint s = g.get_local_range(0) / 2; s > 0; s >>= 1) { |
| 397 | if (tid < s) { blkNNZ_[tid] += blkNNZ_[tid + s]; } |
| 398 | it.barrier(); |
| 399 | } |
| 400 | |
| 401 | if (tid == 0) { |
| 402 | nnz = blkNNZ_[0]; |
| 403 | global_atomic_ref<uint>(nnzc_[0]) += nnz; |
| 404 | } |
| 405 | } |