// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "HotPixelsGPU.h"

#include <cub/device/device_radix_sort.cuh>

#include "../../common/JFJochException.h"

namespace {

void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

constexpr int THREADS = 256;

// Masked or error value, and saturated: neither has a count to rank.
__device__ bool is_valid(int32_t v) {
    return v != INT32_MIN && v != INT32_MAX;
}

// The rank-th smallest of the keys a block's elements j in [begin, end) give, where key_of(j, key)
// says whether element j has one. Radix selection, eight bits at a time from the top: a histogram of
// the next digit among the keys that share the digits already fixed says which digit the rank-th
// falls in. Exact for any keys, so no range has to be assumed. Called by all THREADS of the block.
template <class F>
__device__ uint32_t block_select(uint32_t begin, uint32_t end, uint32_t rank, F key_of) {
    __shared__ uint32_t hist[256];
    __shared__ uint32_t s_prefix, s_rank;
    uint32_t prefix = 0;
    for (int shift = 24; shift >= 0; shift -= 8) {
        hist[threadIdx.x] = 0;
        __syncthreads();
        const uint32_t fixed = shift == 24 ? 0u : ~0u << (shift + 8);
        for (uint32_t j = begin + threadIdx.x; j < end; j += THREADS) {
            uint32_t k;
            if (key_of(j, k) && (k & fixed) == prefix)
                atomicAdd(&hist[(k >> shift) & 255u], 1u);
        }
        __syncthreads();
        if (threadIdx.x == 0) {
            uint32_t d = 0;
            while (rank >= hist[d]) {
                rank -= hist[d];
                d++;
            }
            s_prefix = prefix | (d << shift);
            s_rank = rank;
        }
        __syncthreads();
        prefix = s_prefix;
        rank = s_rank;
    }
    return prefix;
}

// int32 to a uint32 that sorts the same way, and back.
__device__ uint32_t order_key(int32_t v) { return static_cast<uint32_t>(v) ^ 0x80000000u; }
__device__ int32_t from_order_key(uint32_t k) { return static_cast<int32_t>(k ^ 0x80000000u); }

// One block per ring-sector: its valid pixels and their lower median.
__global__ void sector_kernel(const int32_t *__restrict__ image, const uint32_t *__restrict__ pixels,
                              const uint32_t *__restrict__ key_begin, uint32_t *__restrict__ count,
                              int32_t *__restrict__ median) {
    const uint32_t begin = key_begin[blockIdx.x], end = key_begin[blockIdx.x + 1];
    __shared__ uint32_t s_n;
    if (threadIdx.x == 0) s_n = 0;
    __syncthreads();
    uint32_t n = 0;
    for (uint32_t j = begin + threadIdx.x; j < end; j += THREADS)
        n += is_valid(image[pixels[j]]);
    atomicAdd(&s_n, n);
    __syncthreads();
    n = s_n;
    if (threadIdx.x == 0) {
        count[blockIdx.x] = n;
        median[blockIdx.x] = 0;
    }
    if (n == 0)
        return;
    const uint32_t m = block_select(begin, end, n / 2, [&](uint32_t j, uint32_t &k) {
        const int32_t v = image[pixels[j]];
        k = order_key(v);
        return is_valid(v);
    });
    if (threadIdx.x == 0)
        median[blockIdx.x] = from_order_key(m);
}

// One block per ring: the lower median of its valid pixels, and the lower median of their absolute
// deviations from it - written as the host writes it, |v - median| in int.
__global__ void ring_kernel(const int32_t *__restrict__ image, const uint32_t *__restrict__ pixels,
                            const uint32_t *__restrict__ key_begin, const uint32_t *__restrict__ count,
                            int sectors, int32_t *__restrict__ ring_median, int32_t *__restrict__ ring_mad) {
    const int first = blockIdx.x * sectors;
    const uint32_t begin = key_begin[first], end = key_begin[first + sectors];
    uint32_t n = 0;
    for (int s = 0; s < sectors; s++)
        n += count[first + s];
    if (threadIdx.x == 0) {
        ring_median[blockIdx.x] = 0;
        ring_mad[blockIdx.x] = 0;
    }
    if (n == 0)
        return;
    const int32_t m = from_order_key(block_select(begin, end, n / 2, [&](uint32_t j, uint32_t &k) {
        const int32_t v = image[pixels[j]];
        k = order_key(v);
        return is_valid(v);
    }));
    const uint32_t mad = block_select(begin, end, n / 2, [&](uint32_t j, uint32_t &k) {
        const int32_t v = image[pixels[j]];
        k = static_cast<uint32_t>(abs(v - m));
        return is_valid(v);
    });
    if (threadIdx.x == 0) {
        ring_median[blockIdx.x] = m;
        ring_mad[blockIdx.x] = static_cast<int32_t>(mad);
    }
}

// HotPixelFinder::AddImage's per-pixel loop, line for line.
__global__ void accumulate_kernel(const int32_t *__restrict__ image, const int32_t *__restrict__ key,
                                  size_t npixels, int sectors, const int32_t *__restrict__ level,
                                  const float *__restrict__ threshold, const char *__restrict__ ring_ok,
                                  uint16_t *__restrict__ n_lit, uint16_t *__restrict__ n_error,
                                  int64_t *__restrict__ sum_value, uint16_t *__restrict__ n_error_ring_ok,
                                  int64_t *__restrict__ error_level_sum) {
    for (size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x; i < npixels;
         i += static_cast<size_t>(blockDim.x) * gridDim.x) {
        const int32_t k = key[i], v = image[i];
        if (k < 0) continue;
        if (v == INT32_MIN || v == INT32_MAX) {   // the error value or the saturation limit: no count
            n_error[i]++;
            if (ring_ok[k / sectors]) {
                n_error_ring_ok[i]++;
                error_level_sum[i] += level[k];
            }
            continue;
        }
        if (!ring_ok[k / sectors]) continue;
        sum_value[i] += v;
        if (static_cast<float>(v) > threshold[k])
            n_lit[i]++;
    }
}

// Each key's level and lit threshold from the frame's order statistics, and the frame's share of the
// per-key sums - HotPixelFinder::AddImage and AddLevels, step for step. The threshold is
// fma(nsigma, noise, level) + offset, the two roundings the host takes (see AddLevels).
__global__ void levels_kernel(size_t nkeys, int sectors, HotPixelLevelRules rules, const uint32_t *__restrict__ count,
                              const int32_t *__restrict__ sector_median, const int32_t *__restrict__ ring_median,
                              const int32_t *__restrict__ ring_mad, int32_t *__restrict__ level,
                              float *__restrict__ threshold, char *__restrict__ ring_ok,
                              uint32_t *__restrict__ key_frames, int64_t *__restrict__ key_level_sum,
                              int64_t *__restrict__ key_local_level_sum) {
    const size_t k = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (k >= nkeys) return;
    const size_t r = k / sectors;
    uint32_t n = 0;
    for (int s = 0; s < sectors; s++)
        n += count[r * sectors + s];
    const bool ok = n >= static_cast<uint32_t>(rules.min_ring_pixels);
    const int32_t ring_level = ok ? ring_median[r] : 0;
    const float ring_spread = ok ? 1.4826f * static_cast<float>(ring_mad[r]) : 0.0f;
    const bool sector_ok = count[k] >= static_cast<uint32_t>(rules.min_sector_pixels);
    const int32_t sector_level = sector_ok ? sector_median[k] : 0;
    const int32_t local_level = sector_ok ? sector_median[k] : ring_level;
    const int32_t lv = max(ring_level, sector_level);
    const float root = sqrtf(rules.gain * static_cast<float>(max(lv, 0)));
    const float noise = root < ring_spread ? ring_spread : root;
    level[k] = lv;
    threshold[k] = __fadd_rn(__fmaf_rn(rules.lit_nsigma, noise, static_cast<float>(lv)), rules.lit_offset);
    if (k % sectors == 0)
        ring_ok[r] = ok;
    if (ok) {
        atomicAdd(&key_frames[k], 1u);
        atomicAdd(reinterpret_cast<unsigned long long *>(&key_level_sum[k]),
                  static_cast<unsigned long long>(static_cast<int64_t>(lv)));
        atomicAdd(reinterpret_cast<unsigned long long *>(&key_local_level_sum[k]),
                  static_cast<unsigned long long>(static_cast<int64_t>(local_level)));
    }
}

__device__ int valid_frames(const int32_t *key, const uint32_t *key_frames, const uint16_t *n_error_ring_ok, size_t i) {
    return static_cast<int>(key_frames[key[i]]) - static_cast<int>(n_error_ring_ok[i]);
}

__global__ void chance_kernel(size_t npixels, int sectors, const int32_t *__restrict__ key,
                              const uint32_t *__restrict__ key_frames, const uint16_t *__restrict__ n_lit,
                              const uint16_t *__restrict__ n_error_ring_ok, unsigned long long *__restrict__ lit,
                              unsigned long long *__restrict__ seen) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i >= npixels || key[i] < 0) return;
    const int nv = valid_frames(key, key_frames, n_error_ring_ok, i);
    if (nv > 0 && 2 * static_cast<int>(n_lit[i]) <= nv) {
        atomicAdd(&lit[key[i] / sectors], static_cast<unsigned long long>(n_lit[i]));
        atomicAdd(&seen[key[i] / sectors], static_cast<unsigned long long>(nv));
    }
}

// Writes the candidates' indices from `out` on when `out` is given, and counts them either way.
__global__ void candidate_kernel(size_t npixels, int sectors, uint32_t frames, int min_valid, bool spacing_ok,
                                 const int32_t *__restrict__ key, const uint32_t *__restrict__ key_frames,
                                 const uint16_t *__restrict__ n_lit, const uint16_t *__restrict__ n_error,
                                 const uint16_t *__restrict__ n_error_ring_ok, const int64_t *__restrict__ sum_value,
                                 const int64_t *__restrict__ key_local_level_sum, const int *__restrict__ k_chance,
                                 uint32_t *__restrict__ count, uint32_t *__restrict__ out) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i >= npixels || key[i] < 0) return;
    const int nv = valid_frames(key, key_frames, n_error_ring_ok, i);
    const bool error = 2 * static_cast<uint32_t>(n_error[i]) > frames;
    const int lit = n_lit[i];
    const bool persistent = spacing_ok && nv >= min_valid && lit > 0
                            && lit >= min(max(2, k_chance[key[i] / sectors]), nv);
    // Can be dead (HotPixelFinder::GetMask): at most half of its sector's level, no error frames.
    const bool dark = n_error[i] == 0 && nv >= min_valid && 2 * sum_value[i] <= key_local_level_sum[key[i]];
    if (!error && !persistent && !dark) return;
    const uint32_t slot = atomicAdd(count, 1u);
    if (out) out[slot] = static_cast<uint32_t>(i);
}

template <class T>
__global__ void add_kernel(size_t n, const T *__restrict__ in, T *__restrict__ out) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i < n) out[i] += in[i];
}

__global__ void iota_kernel(size_t n, uint32_t *__restrict__ out) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i < n) out[i] = static_cast<uint32_t>(i);
}

template <class T>
__global__ void gather_kernel(size_t n, const uint32_t *__restrict__ index, const T *__restrict__ in, T *__restrict__ out) {
    const size_t j = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (j < n) out[j] = in[index[j]];
}

// The shared tables and sums are filled on a stream of their own, added to on the workers' streams
// and downloaded on the NULL stream, so they are allocated synchronously rather than from the pool:
// a pooled buffer is freed on the thread's allocation stream, which none of those is ordered before.
// Each of those steps is waited for on the host, so this costs nothing but the device-wide sync of
// eight allocations per run, and it leaves compute-sanitizer --track-stream-ordered-races nothing
// to report here, where it flagged every one of them.
constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous;

}   // namespace

HotPixelFinderGPU::HotPixelFinderGPU(const int32_t *host_key, size_t npixels,
                                     const std::vector<uint32_t> &host_key_begin, int nrings, int sectors,
                                     const HotPixelLevelRules &rules)
    : npixels(npixels), nkeys(static_cast<size_t>(nrings) * sectors), nrings(nrings),
      sectors(sectors), rules(rules),
      key(npixels, ALLOC), pixels_by_key(std::max<size_t>(host_key_begin.back(), 1), ALLOC),
      key_begin(host_key_begin.size(), ALLOC),
      n_lit(npixels, ALLOC), n_error(npixels, ALLOC), n_error_ring_ok(npixels, ALLOC),
      sum_value(npixels, ALLOC), error_level_sum(npixels, ALLOC),
      key_frames(std::max<size_t>(nkeys, 1), ALLOC), key_level_sum(std::max<size_t>(nkeys, 1), ALLOC),
      key_local_level_sum(std::max<size_t>(nkeys, 1), ALLOC) {
    cuda_err(cudaGetDevice(&device));
    cuda_err(cudaEventCreateWithFlags(&last_accumulate, cudaEventDisableTiming));
    CudaStream stream;
    cuda_err(cudaMemcpyAsync(key, host_key, npixels * sizeof(int32_t), cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemcpyAsync(key_begin, host_key_begin.data(), host_key_begin.size() * sizeof(uint32_t),
                             cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemsetAsync(n_lit, 0, npixels * sizeof(uint16_t), stream));
    cuda_err(cudaMemsetAsync(n_error, 0, npixels * sizeof(uint16_t), stream));
    cuda_err(cudaMemsetAsync(n_error_ring_ok, 0, npixels * sizeof(uint16_t), stream));
    cuda_err(cudaMemsetAsync(sum_value, 0, npixels * sizeof(int64_t), stream));
    cuda_err(cudaMemsetAsync(error_level_sum, 0, npixels * sizeof(int64_t), stream));
    cuda_err(cudaMemsetAsync(key_frames, 0, std::max<size_t>(nkeys, 1) * sizeof(uint32_t), stream));
    cuda_err(cudaMemsetAsync(key_level_sum, 0, std::max<size_t>(nkeys, 1) * sizeof(int64_t), stream));
    cuda_err(cudaMemsetAsync(key_local_level_sum, 0, std::max<size_t>(nkeys, 1) * sizeof(int64_t), stream));

    // The unmasked pixels grouped by key: a stable sort of the pixel indices by key, so each key's
    // pixels stay in pixel order. A masked pixel's key, -1, is the largest as unsigned and sorts last.
    {
        CudaDevicePtr<uint32_t> index(npixels, ALLOC), sorted_key(npixels, ALLOC), sorted_index(npixels, ALLOC);
        iota_kernel<<<static_cast<unsigned>((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>(npixels, index);
        cuda_err(cudaGetLastError());
        const auto *keys_in = reinterpret_cast<const uint32_t *>(key.get());
        size_t bytes = 0;
        cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys_in, sorted_key.get(), index.get(),
                                                 sorted_index.get(), npixels, 0, 32, stream));
        CudaDevicePtr<uint8_t> scratch(bytes, ALLOC);
        cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, keys_in, sorted_key.get(), index.get(),
                                                 sorted_index.get(), npixels, 0, 32, stream));
        if (host_key_begin.back() > 0)
            cuda_err(cudaMemcpyAsync(pixels_by_key, sorted_index, host_key_begin.back() * sizeof(uint32_t),
                                     cudaMemcpyDeviceToDevice, stream));
        cuda_err(cudaStreamSynchronize(stream));
    }
}

HotPixelFinderGPU::~HotPixelFinderGPU() {
    if (last_accumulate) cudaEventDestroy(last_accumulate);
}

void HotPixelFinderGPU::Add(const int32_t *device_image, Frame &frame) {
    if (nrings == 0)
        return;
    if (!frame.count.get()) {
        frame.count = CudaDevicePtr<uint32_t>(nkeys);
        frame.sector_median = CudaDevicePtr<int32_t>(nkeys);
        frame.ring_median = CudaDevicePtr<int32_t>(nrings);
        frame.ring_mad = CudaDevicePtr<int32_t>(nrings);
        frame.level = CudaDevicePtr<int32_t>(nkeys);
        frame.threshold = CudaDevicePtr<float>(nkeys);
        frame.ring_ok = CudaDevicePtr<char>(nrings);
    }
    const cudaStream_t stream = *frame.stream;
    sector_kernel<<<static_cast<unsigned>(nkeys), THREADS, 0, stream>>>(device_image, pixels_by_key, key_begin,
                                                                        frame.count, frame.sector_median);
    cuda_err(cudaGetLastError());
    ring_kernel<<<nrings, THREADS, 0, stream>>>(device_image, pixels_by_key, key_begin, frame.count, sectors,
                                                frame.ring_median, frame.ring_mad);
    cuda_err(cudaGetLastError());
    levels_kernel<<<static_cast<unsigned>((nkeys + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
            nkeys, sectors, rules, frame.count, frame.sector_median, frame.ring_median, frame.ring_mad,
            frame.level, frame.threshold, frame.ring_ok, key_frames, key_level_sum, key_local_level_sum);
    cuda_err(cudaGetLastError());

    std::lock_guard lock(accumulate_mutex);
    cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0));
    accumulate_kernel<<<static_cast<unsigned>((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
            device_image, key, npixels, sectors, frame.level, frame.threshold, frame.ring_ok,
            n_lit, n_error, sum_value, n_error_ring_ok, error_level_sum);
    cuda_err(cudaGetLastError());
    cuda_err(cudaEventRecord(last_accumulate, stream));
}

void HotPixelFinderGPU::Fold(HotPixelFinderGPU &other) {
    CudaStream stream;
    cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0));
    cuda_err(cudaStreamWaitEvent(stream, other.last_accumulate, 0));
    // One scratch buffer for every array in turn: the copies and additions are ordered on one stream.
    CudaDevicePtr<uint8_t> scratch(std::max(npixels, std::max<size_t>(nkeys, 1)) * sizeof(int64_t), ALLOC);
    const auto add = [&]<class T>(T *mine, const T *theirs, size_t n) {
        if (n == 0) return;
        cuda_err(cudaMemcpyPeerAsync(scratch, device, theirs, other.device, n * sizeof(T), stream));
        add_kernel<T><<<static_cast<unsigned>((n + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
                n, reinterpret_cast<const T *>(scratch.get()), mine);
        cuda_err(cudaGetLastError());
    };
    add(n_lit.get(), other.n_lit.get(), npixels);
    add(n_error.get(), other.n_error.get(), npixels);
    add(n_error_ring_ok.get(), other.n_error_ring_ok.get(), npixels);
    add(sum_value.get(), other.sum_value.get(), npixels);
    add(error_level_sum.get(), other.error_level_sum.get(), npixels);
    add(key_frames.get(), other.key_frames.get(), nkeys);
    add(key_level_sum.get(), other.key_level_sum.get(), nkeys);
    add(key_local_level_sum.get(), other.key_local_level_sum.get(), nkeys);
    cuda_err(cudaEventRecord(last_accumulate, stream));
    cuda_err(cudaStreamSynchronize(stream));
}

void HotPixelFinderGPU::ChanceCounts(std::vector<int64_t> &lit, std::vector<int64_t> &seen) {
    CudaStream stream;
    cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0));
    const size_t rings = std::max(nrings, 1);
    CudaDevicePtr<unsigned long long> d_lit(rings, ALLOC), d_seen(rings, ALLOC);
    cuda_err(cudaMemsetAsync(d_lit, 0, rings * sizeof(unsigned long long), stream));
    cuda_err(cudaMemsetAsync(d_seen, 0, rings * sizeof(unsigned long long), stream));
    chance_kernel<<<static_cast<unsigned>((npixels + THREADS - 1) / THREADS), THREADS, 0, stream>>>(
            npixels, sectors, key, key_frames, n_lit, n_error_ring_ok, d_lit, d_seen);
    cuda_err(cudaGetLastError());
    lit.resize(nrings);
    seen.resize(nrings);
    cuda_err(cudaMemcpyAsync(lit.data(), d_lit, nrings * sizeof(int64_t), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaMemcpyAsync(seen.data(), d_seen, nrings * sizeof(int64_t), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
}

HotPixelFinderGPU::Candidates HotPixelFinderGPU::GetCandidates(uint32_t frames, int min_valid, bool spacing_ok,
                                                               const std::vector<int> &k_chance) {
    CudaStream stream;
    cuda_err(cudaStreamWaitEvent(stream, last_accumulate, 0));
    const unsigned blocks = static_cast<unsigned>((npixels + THREADS - 1) / THREADS);
    CudaDevicePtr<int> d_k_chance(std::max<size_t>(k_chance.size(), 1), ALLOC);
    CudaDevicePtr<uint32_t> d_count(1, ALLOC);
    if (!k_chance.empty())
        cuda_err(cudaMemcpyAsync(d_k_chance, k_chance.data(), k_chance.size() * sizeof(int), cudaMemcpyHostToDevice,
                                 stream));
    cuda_err(cudaMemsetAsync(d_count, 0, sizeof(uint32_t), stream));
    candidate_kernel<<<blocks, THREADS, 0, stream>>>(npixels, sectors, frames, min_valid, spacing_ok, key, key_frames,
                                                     n_lit, n_error, n_error_ring_ok, sum_value, key_local_level_sum,
                                                     d_k_chance, d_count, nullptr);
    cuda_err(cudaGetLastError());
    uint32_t n = 0;
    cuda_err(cudaMemcpyAsync(&n, d_count, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));

    Candidates c;
    c.key_frames.resize(nkeys);
    c.key_level_sum.resize(nkeys);
    c.key_local_level_sum.resize(nkeys);
    if (nkeys > 0) {
        cuda_err(cudaMemcpyAsync(c.key_frames.data(), key_frames, nkeys * sizeof(uint32_t), cudaMemcpyDeviceToHost,
                                 stream));
        cuda_err(cudaMemcpyAsync(c.key_level_sum.data(), key_level_sum, nkeys * sizeof(int64_t),
                                 cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.key_local_level_sum.data(), key_local_level_sum, nkeys * sizeof(int64_t),
                                 cudaMemcpyDeviceToHost, stream));
    }
    if (n > 0) {
        CudaDevicePtr<uint32_t> index(n, ALLOC);
        CudaDevicePtr<uint16_t> g_lit(n, ALLOC), g_error(n, ALLOC), g_error_ring_ok(n, ALLOC);
        CudaDevicePtr<int64_t> g_sum(n, ALLOC), g_error_level(n, ALLOC);
        cuda_err(cudaMemsetAsync(d_count, 0, sizeof(uint32_t), stream));
        candidate_kernel<<<blocks, THREADS, 0, stream>>>(npixels, sectors, frames, min_valid, spacing_ok, key,
                                                         key_frames, n_lit, n_error, n_error_ring_ok, sum_value,
                                                         key_local_level_sum, d_k_chance, d_count, index);
        cuda_err(cudaGetLastError());
        const unsigned gb = (n + THREADS - 1) / THREADS;
        gather_kernel<<<gb, THREADS, 0, stream>>>(n, index, n_lit.get(), g_lit.get());
        gather_kernel<<<gb, THREADS, 0, stream>>>(n, index, n_error.get(), g_error.get());
        gather_kernel<<<gb, THREADS, 0, stream>>>(n, index, n_error_ring_ok.get(), g_error_ring_ok.get());
        gather_kernel<<<gb, THREADS, 0, stream>>>(n, index, sum_value.get(), g_sum.get());
        gather_kernel<<<gb, THREADS, 0, stream>>>(n, index, error_level_sum.get(), g_error_level.get());
        cuda_err(cudaGetLastError());
        c.index.resize(n);
        c.n_lit.resize(n);
        c.n_error.resize(n);
        c.n_error_ring_ok.resize(n);
        c.sum_value.resize(n);
        c.error_level_sum.resize(n);
        cuda_err(cudaMemcpyAsync(c.index.data(), index, n * sizeof(uint32_t), cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.n_lit.data(), g_lit, n * sizeof(uint16_t), cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.n_error.data(), g_error, n * sizeof(uint16_t), cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.n_error_ring_ok.data(), g_error_ring_ok, n * sizeof(uint16_t),
                                 cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.sum_value.data(), g_sum, n * sizeof(int64_t), cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaMemcpyAsync(c.error_level_sum.data(), g_error_level, n * sizeof(int64_t),
                                 cudaMemcpyDeviceToHost, stream));
    }
    cuda_err(cudaStreamSynchronize(stream));
    return c;
}