#pragma once

#include <cooperative_groups.h>

#include <gsl-lite/gsl-lite.hpp>

#include <cuda/runtime_api.hpp>

#include <cub/cub.cuh>

#include <makeshift/variant.hpp>

#include <thrustshift/constant.h>
#include <thrustshift/copy.h>
#include <thrustshift/fill.h>
#include <thrustshift/histogram.h>
#include <thrustshift/math.h>
#include <thrustshift/not-a-vector.h>
#include <thrustshift/select-if.h>
#include <thrustshift/type-traits.h>

namespace thrustshift {

namespace device_function {

namespace implicit_unroll {

/*! \brief Sum an n x N array column-wise.
 *  \param N length of row
 *  \param n number of rows
 *  \param p input of length `N * n`
 *  \param result of length `N`. As a result you may pass `p`.
 */
template <typename T0, typename T1, typename I0, typename I1, typename I2>
CUDA_FD void sum_subsequent_into(const T0* p,
                                 T1* result,
                                 int tid,
                                 I0 num_threads,
                                 I1 N,
                                 I2 n) {

	auto num_columns_per_thread = N / num_threads;
	auto sum_column = [&](int col_id) {
		T0 x{};
#pragma unroll
		for (int j = 0; j < n; ++j) {
			x += p[col_id + N * j];
		}
		result[col_id] = x;
	};
#pragma unroll
	for (int i = 0; i < num_columns_per_thread; ++i) {
		const int col_id = i * num_threads + tid;
		sum_column(col_id);
	}
	auto num_rest = N % num_threads;
	if (tid < num_rest) {
		const int col_id = num_columns_per_thread * num_threads + tid;
		sum_column(col_id);
	}
}

} // namespace implicit_unroll

namespace explicit_unroll {

/*! \brief Sum an n x N array column-wise.
 *  \param N length of row
 *  \param n number of rows
 *  \param p input of length `N * n`
 *  \param result of length `N`. As a result you may pass `p`.
 */
template <typename T, int num_threads, int N, int n>
CUDA_FD void sum_subsequent_into(const T* p, T* result, int tid) {

	constexpr int num_columns_per_thread = N / num_threads;
	auto sum_column = [&](int col_id) {
		T x{};
#pragma unroll
		for (int j = 0; j < n; ++j) {
			x += p[col_id + N * j];
		}
		result[col_id] = x;
	};
#pragma unroll
	for (int i = 0; i < num_columns_per_thread; ++i) {
		const int col_id = i * num_threads + tid;
		sum_column(col_id);
	}
	constexpr int num_rest = N % num_threads;
	if (tid < num_rest) {
		const int col_id = num_columns_per_thread * num_threads + tid;
		sum_column(col_id);
	}
}

} // namespace explicit_unroll

/*! \brief Bin value `x` based on a subset of 8 bits into a 256 bin histogram with a full warp.
 *
 *  Only one atomic addition is executed for each bin, although multiple threads in a warp can
 *  have a value which belongs into the same bin. Via warp level primitives the bin incrementation is
 *  aggregated to one single thread.
 *
 *  ```
 *  x x x x x x x x x x x x x x x x x x x x x x x  <--- `x` represents one bit
 * |--bit_offset---|--bit_size---|--rest---------|
 * |-----prefix----|
 * ```
 *
 *  \param x value which is put into the bin. Can be different for each thread.
 *  \param bit_offset number of bits which are omitted from the left of the bit pattern of `x`.
 *  \param prefix bit pattern with `bit_offset` reasonable bits from the right. If the beginning of the bit pattern
 *      of `x` is not equal to `prefix`, the value is not binned.
 *  \param histogram length of 256. Does not have to be unique for each warp.
 *  \param valid_write set to false if a thread should not bin the value `x`.
 *  \param bin_index_transform lambda with `[] (int i) -> int {...}` signature. If set to the identity
 *      the values are binned from low to high. The lambda can be used to reverse the order in the
 *      histogram.
 */
template <typename T, typename IH, class F>
CUDA_FD void bin_value256(T x,
                          int bit_offset,
                          uint64_t prefix,
                          IH* histogram,
                          bool valid_write,
                          F bin_index_transform) {

	using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
	using K = uint8_t;
	constexpr int bit_size = sizeof(K) * 8;
	// extract the bits of length `bit_size` from the left with an offset of `bit_offset` to the right.
	//
	//  x x x x x x x x x x x x x x x x x x x x x x x  <--- `x` represents one bit
	// |--bit_offset---|--bit_size---|--rest---------|
	//
	const I i = *reinterpret_cast<I*>((void*) (&x));
	const K b = (i >> (sizeof(I) * 8 - bit_size - bit_offset)) &
	            I(std::numeric_limits<K>::max());

	int bi = bin_index_transform(b);

	int k = 1; // increment of the bin
	// Only count values which start with the given prefix pattern
	if ((i >> sizeof(I) * 8 - bit_offset) != static_cast<I>(prefix) ||
	    !valid_write) {
		bi = -1; // value has different prefix
		k = 0;
	}

	// Threads with higher ID take precedence over those with low ID and the
	// same bin.
	const int lane_id = threadIdx.x % warp_size;

	const int mask = __match_any_sync(0xffffffff, bi) &
	                 (~(1 << lane_id)); // set our own bit to zero

	const int lsbi = sizeof(int) * 8 - __clz(mask) -
	                 1; // get ID of highest thread which has this value
	const unsigned umask =
	    *reinterpret_cast<const unsigned*>((const void*) &mask);
	k = lsbi > lane_id ? 0 : (__popc(umask) + 1);

	if (k > 0 && bi >= 0) {
		atomicAdd(histogram + bi, k);
	}
}

namespace implicit_unroll {

/*! \brief Bin transformed values based on a subset of 8 bits into a 256 bin histogram.
 *
 *  After execution and synchronization, multiple histograms (`histograms`) are filled, which must subsequently be summed into
 *  one single histogram.
 *
 *  ```
 *  x x x x x x x x x x x x x x x x x x x x x x x  <--- `x` represents one bit
 * |--bit_offset---|--bit_size---|--rest---------|
 * |-----prefix----|
 * ```
 *
 *  \param values of length N.
 *  \param histograms of length `256 * block_dim / warp_size`. Contains the result after execution and must be
 *      initialized before function execution.
 *  \param num_histograms number of histograms.
 *  \param my_tile_start is equal to zero for intra block binning and equal to `blockIdx.x` for inter block binning.
 *  \param tile_increment is equal to one for intra block binning and equal to `gridDim.x` for inter block binning.
 *      For best performance, this should be a compile time constant.
 *  \param bit_offset number of bits which are omitted from the left of the bit pattern of `x`.
 *  \param num_threads the number of threads which are entering the function (num_threads % warp_size == 0).
 *  \param prefix bit pattern with `bit_offset` reasonable bits from the right. If the beginning of the bit pattern
 *      of `x` is not equal to `prefix`, the value is not binned.
 *  \param unary_functor lambda to transform the values before binning.
 *  \param bin_index_transform lambda with `[] (int i) -> int {...}` signature. If set to the identity
 *      the values are binned from low to high. The lambda can be used to reverse the order in the
 *      histogram.
 */
template <typename T,
          typename I0,
          typename I1,
          typename I2,
          typename I3,
          typename I4,
          typename I5,
          class F0,
          class F1>
CUDA_FHD void bin_values256(
    const T* values,
    I0 N,
    I1* histograms,
    I5 num_histograms,
    I2 my_tile_start, // =0 in case of intra block, =blockIdx.x in case of inter block binning
    I3 tile_increment, // =1 in case of intra block, =grid_dim in case of inter block binning
    int tid,
    I4 num_threads,
    int bit_offset,
    uint64_t prefix,
    F0 unary_functor,
    F1 bin_index_transform) {

	gsl_ExpectsAudit(num_threads % warp_size == 0);
	constexpr int histogram_length = 256;
	const int warp_id = tid / warp_size;
	auto num_warps = num_threads / warp_size;
	gsl_ExpectsAudit(num_warps % num_histograms == 0);
	gsl_ExpectsAudit(num_histograms <= num_warps);
	auto num_warps_per_histogram = num_warps / num_histograms;
	const int histogram_id = warp_id / num_warps_per_histogram;
	I1* my_histogram = histograms + histogram_id * histogram_length;

	const int num_tiles = thrustshift::ceil_divide(N, num_threads);
	auto tile_size = num_threads;
	int tile_id = my_tile_start;
	for (; tile_id < num_tiles - 1; tile_id += tile_increment) {
		const int tile_offset = tile_id * num_threads;
		const auto x = unary_functor(values[tile_offset + tid]);
		bin_value256(
		    x, bit_offset, prefix, my_histogram, true, bin_index_transform);
	}
	// last tile
	if (tile_id == num_tiles - 1) {
		const int tile_offset = tile_id * num_threads;
		const int curr_tile_size =
		    tile_offset + tile_size > N ? (N - tile_offset) : tile_size;
		bool valid_rw = tid < curr_tile_size;
		const auto x = [&] {
			if (valid_rw) {
				return unary_functor(values[tile_offset + tid]);
			}
			return T{};
		}();
		// every thread of the warp must enter this function as it contains
		// `*__sync` warp level primitives.
		bin_value256(
		    x, bit_offset, prefix, my_histogram, valid_rw, bin_index_transform);
	}
}

} // namespace implicit_unroll

template <int block_dim, int num_histograms, typename IH>
struct k_largest_values_abs_block {

	static constexpr int histogram_length = 256;
	static constexpr int num_warps = block_dim / warp_size;
	static_assert(block_dim % warp_size == 0);
	static_assert(histogram_length % block_dim == 0);
	static constexpr int num_scan_elements_per_thread =
	    histogram_length / block_dim;
	using BlockLoad = cub::BlockLoad<IH,
	                                 block_dim,
	                                 num_scan_elements_per_thread,
	                                 cub::BLOCK_LOAD_WARP_TRANSPOSE>;
	using BlockScan = cub::BlockScan<IH, block_dim>;
	using BlockStore = cub::BlockStore<IH,
	                                   block_dim,
	                                   num_scan_elements_per_thread,
	                                   cub::BLOCK_STORE_WARP_TRANSPOSE>;

	struct triplet_t {
		int k;
		uint64_t prefix;
		int selected_values_pos;
	};

	struct TempStorage {
		union {
			typename BlockLoad::TempStorage block_load;
			typename BlockScan::TempStorage block_scan;
			typename BlockStore::TempStorage block_store;
		};
		triplet_t triplet;
	};

	template <typename T, typename I0>
	static CUDA_FD thrust::tuple<uint64_t, int>
	k_largest_values_abs_radix_block(const T* values,
	                                 I0 N,
	                                 IH* uninitialized_histograms,
	                                 int k,
	                                 TempStorage& temp_storage) {

		const int tid = threadIdx.x;

		auto unary_functor = [](T x) {
			using std::abs;
			return abs(x);
		};

		// Create the histogram from large to small values
		auto bin_index_transform = [](auto i) {
			return histogram_length - i - 1;
		};

		uint64_t prefix = 0;

		for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8);
		     bit_offset += 8) {

			//
			// Initialize the histograms
			//
			device_function::implicit_unroll::fill(
			    uninitialized_histograms,
			    0,
			    tid,
			    block_dim,
			    num_warps * histogram_length);
			auto* histograms =
			    uninitialized_histograms; // rename for better readability
			__syncthreads();

			//
			// Bin values into one histogram per warp
			//
			thrustshift::device_function::implicit_unroll::bin_values256(
			    values,
			    N,
			    histograms,
			    num_histograms,
			    0, // tile start
			    1, // tile increment
			    tid,
			    block_dim,
			    bit_offset,
			    prefix,
			    unary_functor,
			    bin_index_transform);
			__syncthreads();

			//
			// Sum all histograms
			//
			if (num_histograms > 1) {
				thrustshift::device_function::implicit_unroll::
				    sum_subsequent_into(histograms,
				                        histograms,
				                        tid,
				                        block_dim,
				                        histogram_length,
				                        num_warps);
				__syncthreads();
			}
			// The first histogram is now the sum of all histograms in shared memory

			//
			// Scan the histogram
			//
			IH hvalues[num_scan_elements_per_thread];
			IH hcumulative_values[num_scan_elements_per_thread];
			BlockLoad(temp_storage.block_load).Load(histograms, hvalues);
			__syncthreads(); // necessary due to reuse of temporary memory
			BlockScan(temp_storage.block_scan)
			    .InclusiveSum(hvalues, hcumulative_values);
			__syncthreads();

			// Create helper array to have the value of our left neighbour
			IH hcumulative_values2[num_scan_elements_per_thread + 1];
#pragma unroll
			for (int j = 0; j < num_scan_elements_per_thread; ++j) {
				hcumulative_values2[j + 1] = hcumulative_values[j];
			}

			BlockStore(
			    temp_storage
			        .block_store) // block store changes the register values
			    .Store(histograms, hcumulative_values);
			__syncthreads();

			if (tid > 0) {
				hcumulative_values2[0] =
				    histograms[tid * num_scan_elements_per_thread - 1];
			}
			else {
				hcumulative_values2[0] = 0;
			}
#pragma unroll
			for (int j = 0; j < num_scan_elements_per_thread; ++j) {
				const int l = j + tid * num_scan_elements_per_thread;
				const int i = histogram_length - l - 1;
				if (hcumulative_values2[j + 1] >= k &&
				    hcumulative_values2[j] < k) {
					// Only one thread is expected to enter this branch
					prefix = (prefix << 8) | uint64_t(i);
					// all values with this prefix and larger are included
					temp_storage.triplet.k = hcumulative_values2[j + 1] == k
					                             ? 0
					                             : k - hcumulative_values2[j];
					temp_storage.triplet.prefix = prefix;
					// This is set for subsequent function execution to avoid one
					// additional __syncthreads();
					temp_storage.triplet.selected_values_pos = 0;
				}
			}

			__syncthreads();
			k = temp_storage.triplet.k;
			prefix = temp_storage.triplet.prefix;
			if (k == 0) {
				return {prefix, bit_offset + 8};
			}
		}
		return {prefix, sizeof(T) * 8};
	}

	template <typename It, typename ItSelected, typename I0>
	static CUDA_FD void select_k_largest_values_with_index_abs(
	    It values,
	    I0 N,
	    ItSelected selected_values,
	    IH* uninitialized_histograms,
	    int k,
	    TempStorage& temp_storage) {

		const int tid = threadIdx.x;

		auto tup = k_largest_values_abs_radix_block(
		    values, N, uninitialized_histograms, k, temp_storage);
		auto prefix = thrust::get<0>(tup);
		auto bit_offset = thrust::get<1>(tup);

		using T = typename std::remove_const<
		    typename std::iterator_traits<It>::value_type>::type;

		auto cit = thrust::make_counting_iterator(0);
		auto it = thrust::make_zip_iterator(thrust::make_tuple(values, cit));

		auto select_op = [prefix, bit_offset] __device__(
		                     const thrust::tuple<T, int>& tup) {
			using std::abs;
			auto x = abs(thrust::get<0>(tup));
			using I =
			    typename thrustshift::make_uintegral_of_equal_size<T>::type;
			const I i = *reinterpret_cast<I*>((void*) (&x));
			return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
		};

		device_function::implicit_unroll::select_if(
		    it,
		    N,
		    selected_values,
		    &temp_storage.triplet.selected_values_pos,
		    tid,
		    block_dim,
		    select_op);
	}
};

} // namespace device_function

namespace kernel {

template <typename T, int num_threads, int N, int n>
__global__ void sum_subsequent_into(const T* p, T* result) {
	const int tid = threadIdx.x;
	thrustshift::device_function::explicit_unroll::
	    sum_subsequent_into<T, num_threads, N, n>(p, result, tid);
}

template <typename T,
          int block_dim,
          int grid_dim,
          int num_sh_histograms,
          class F0,
          class F1>
__global__ void bin_values256(const T* data,
                              int N,
                              int* histograms,
                              int bit_offset,
                              uint64_t prefix,
                              F0 unary_functor,
                              F1 bin_index_transform) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	using histogram_value_type = unsigned;
	__shared__ histogram_value_type sh_histograms[all_sh_histograms_length];
	const int tid = threadIdx.x;

	device_function::explicit_unroll::
	    fill<histogram_value_type, block_dim, all_sh_histograms_length>(
	        sh_histograms, 0, tid);

	__syncthreads();

	device_function::implicit_unroll::bin_values256(data,
	                                                N,
	                                                sh_histograms,
	                                                num_sh_histograms,
	                                                blockIdx.x,
	                                                grid_dim,
	                                                tid,
	                                                block_dim,
	                                                bit_offset,
	                                                prefix,
	                                                unary_functor,
	                                                bin_index_transform);

	__syncthreads();

	//
	// Sum all histograms
	//
	if (num_sh_histograms > 1) {
		thrustshift::device_function::explicit_unroll::sum_subsequent_into<
		    histogram_value_type,
		    block_dim,
		    histogram_length,
		    num_sh_histograms>(sh_histograms, sh_histograms, tid);
		__syncthreads();
	}

	thrustshift::block_copy<block_dim, histogram_length>(
	    sh_histograms, histograms + blockIdx.x * histogram_length);
}

template <typename T,
          int block_dim,
          int grid_dim,
          int num_sh_histograms,
          class F0,
          class F1>
__global__ void bin_values256_threadfence(const T* data,
                                          int N,
                                          volatile int* histograms,
                                          int bit_offset,
                                          uint64_t prefix,
                                          F0 unary_functor,
                                          F1 bin_index_transform,
                                          unsigned* entry_ticket,
                                          unsigned* exit_ticket) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	using histogram_value_type = int;
	__shared__ histogram_value_type sh_histograms[all_sh_histograms_length];
	const int tid = threadIdx.x;

	device_function::explicit_unroll::
	    fill<histogram_value_type, block_dim, all_sh_histograms_length>(
	        sh_histograms, 0, tid);
	__shared__ int sh_entry_bid;
	__shared__ int sh_exit_bid;
	if (tid == 0) {
		sh_entry_bid = atomicInc(entry_ticket, grid_dim);
	}
	__syncthreads();
	int bid = sh_entry_bid;

	device_function::implicit_unroll::bin_values256(data,
	                                                N,
	                                                sh_histograms,
	                                                num_sh_histograms,
	                                                bid,
	                                                grid_dim,
	                                                tid,
	                                                block_dim,
	                                                bit_offset,
	                                                prefix,
	                                                unary_functor,
	                                                bin_index_transform);

	__syncthreads();

	//
	// Sum all histograms
	//
	if (num_sh_histograms > 1) {
		thrustshift::device_function::explicit_unroll::sum_subsequent_into<
		    histogram_value_type,
		    block_dim,
		    histogram_length,
		    num_sh_histograms>(sh_histograms, sh_histograms, tid);
		__syncthreads();
	}

	thrustshift::block_copy<block_dim, histogram_length>(
	    sh_histograms, histograms + bid * histogram_length);

	// To my understanding this sync threads is necessary because it might happen that
	// thread 0 does not observe the write of all threads of our histogram.
	__syncthreads();

	__threadfence();

	if (tid == 0) {
		sh_exit_bid = atomicInc(exit_ticket, grid_dim);
	}
	__syncthreads();
	if (sh_exit_bid == grid_dim - 1) {
		static_assert(block_dim >= histogram_length);
		if (tid < histogram_length) {
			histogram_value_type h = 0;
#pragma unroll
			for (int histo_id = 0; histo_id < grid_dim; ++histo_id) {
				h += histograms[histo_id * histogram_length + tid];
			}
			histograms[tid] = h;
		}
	}
}

template <typename T,
          int block_dim,
          int grid_dim,
          int num_sh_histograms,
          class F0,
          class F1>
__global__ void bin_values256_atomic(const T* data,
                                     int N,
                                     int* histogram,
                                     int bit_offset,
                                     uint64_t prefix,
                                     F0 unary_functor,
                                     F1 bin_index_transform) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	using histogram_value_type = int;

	__shared__ histogram_value_type sh_histograms[all_sh_histograms_length];

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	device_function::explicit_unroll::
	    fill<histogram_value_type, block_dim, all_sh_histograms_length>(
	        sh_histograms, 0, tid);

	__syncthreads();

	device_function::implicit_unroll::bin_values256(data,
	                                                N,
	                                                sh_histograms,
	                                                num_sh_histograms,
	                                                bid,
	                                                grid_dim,
	                                                tid,
	                                                block_dim,
	                                                bit_offset,
	                                                prefix,
	                                                unary_functor,
	                                                bin_index_transform);

	__syncthreads();

	//
	// Sum all histograms
	//
	if (num_sh_histograms > 1) {
		thrustshift::device_function::explicit_unroll::sum_subsequent_into<
		    histogram_value_type,
		    block_dim,
		    histogram_length,
		    num_sh_histograms>(sh_histograms, sh_histograms, tid);
		__syncthreads();
	}

	if (tid < histogram_length) {
		atomicAdd(histogram + tid, sh_histograms[tid]);
	}
}

template <typename T,
          int block_dim,
          int grid_dim,
          int num_sh_histograms,
          bool use_k0_and_zero_prefix,
          class F0,
          class F1>
__global__ void bin_values256_atomic_with_ptr(const T* data,
                                              int N,
                                              int* histogram,
                                              int bit_offset,
                                              uint64_t* prefix_,
                                              int* k_,
                                              int k0,
                                              F0 unary_functor,
                                              F1 bin_index_transform) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	using histogram_value_type = int;

	__shared__ histogram_value_type sh_histograms[all_sh_histograms_length];

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	const int k = [&] {
		if constexpr (use_k0_and_zero_prefix) {
			return k0;
		}
		else {
			return *k_;
		}
	}();
	if (k <= 0) {
		return;
	}

	const uint64_t prefix = [&]() -> uint64_t {
		if constexpr (use_k0_and_zero_prefix) {
			return 0;
		}
		else {
			return *prefix_;
		}
	}();

	device_function::explicit_unroll::
	    fill<histogram_value_type, block_dim, all_sh_histograms_length>(
	        sh_histograms, 0, tid);

	__syncthreads();

	device_function::implicit_unroll::bin_values256(data,
	                                                N,
	                                                sh_histograms,
	                                                num_sh_histograms,
	                                                bid,
	                                                grid_dim,
	                                                tid,
	                                                block_dim,
	                                                bit_offset,
	                                                prefix,
	                                                unary_functor,
	                                                bin_index_transform);

	__syncthreads();

	//
	// Sum all histograms
	//
	if (num_sh_histograms > 1) {
		thrustshift::device_function::explicit_unroll::sum_subsequent_into<
		    histogram_value_type,
		    block_dim,
		    histogram_length,
		    num_sh_histograms>(sh_histograms, sh_histograms, tid);
		__syncthreads();
	}

	if (tid < histogram_length) {
		atomicAdd(histogram + tid, sh_histograms[tid]);
	}
}

template <typename T,
          typename IH,
          int block_dim,
          int grid_dim,
          int num_sh_histograms>
__global__ void k_select_radix(const T* data,
                               int N,
                               IH* histograms,
                               int* bit_offset_,
                               uint64_t* prefix_,
                               int* k_) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	auto unary_functor = [](T x) {
		using std::abs;
		return abs(x);
	};

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histograms[all_sh_histograms_length];
	} temp_storage;

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	uint64_t prefix;
	int k;

	auto grid = cooperative_groups::this_grid();

	auto bin_index_transform = [](auto i) { return histogram_length - i - 1; };

	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {

		k = *k_;
		prefix = *prefix_;
		if (k == 0) {
			return;
		}

		device_function::explicit_unroll::
		    fill<IH, block_dim, all_sh_histograms_length>(
		        temp_storage.histograms, 0, tid);

		__syncthreads();

		device_function::implicit_unroll::bin_values256(data,
		                                                N,
		                                                temp_storage.histograms,
		                                                num_sh_histograms,
		                                                blockIdx.x,
		                                                grid_dim,
		                                                tid,
		                                                block_dim,
		                                                bit_offset,
		                                                prefix,
		                                                unary_functor,
		                                                bin_index_transform);

		__syncthreads();

		//
		// Sum all histograms
		//
		if (num_sh_histograms > 1) {
			thrustshift::device_function::explicit_unroll::sum_subsequent_into<
			    IH,
			    block_dim,
			    histogram_length,
			    num_sh_histograms>(
			    temp_storage.histograms, temp_storage.histograms, tid);
			__syncthreads();
		}

		IH h = temp_storage.histograms[tid];
		IH hcum;
		// it is faster if all blocks do already the scan before
		// the histograms are summed up.
		BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
		histograms[bid * histogram_length + tid] = hcum;

		//
		// Store block histogram to global memory
		//
		// thrustshift::block_copy<block_dim, histogram_length>(
		//     temp_storage.histograms, histograms + bid * histogram_length);

		grid.sync();
		if (bid == 0) {
			static_assert(histogram_length == block_dim);
#pragma unroll
			for (int histo_id = 1; histo_id < grid_dim; ++histo_id) {
				hcum += histograms[histo_id * histogram_length + tid];
			}
			//			BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
			//			__syncthreads();
			temp_storage.histograms[tid] = hcum;
			__syncthreads();
			IH hleft = 0;
			if (tid > 0) {
				hleft = temp_storage.histograms[tid - 1];
			}

			const int i = histogram_length - tid - 1;
			if (hcum >= k && hleft < k) {
				// Only one thread is expected to enter this branch
				prefix = (prefix << 8) | uint64_t(i);
				// all values with this prefix and larger are included
				*k_ = hcum == k ? 0 : k - hleft;
				*prefix_ = prefix;
				*bit_offset_ = bit_offset + 8;
			}

			__syncthreads();
		}
		grid.sync();
	}
}

// The summation is done redundantly
template <typename T,
          typename IH,
          int block_dim,
          int grid_dim,
          int num_sh_histograms>
__global__ void k_select_radix2(const T* data,
                                int N,
                                IH* histograms,
                                int* bit_offset_,
                                uint64_t* prefix_,
                                int k) {

	constexpr int histogram_length = 256;
	constexpr int all_sh_histograms_length =
	    histogram_length * num_sh_histograms;

	auto unary_functor = [](T x) {
		using std::abs;
		return abs(x);
	};

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histograms[all_sh_histograms_length];
	} temp_storage;

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	__shared__ uint64_t sh_prefix;
	__shared__ int sh_k;

	uint64_t prefix = 0;

	auto grid = cooperative_groups::this_grid();

	auto bin_index_transform = [](auto i) { return histogram_length - i - 1; };

	for (int bit_offset = 0, l = 0; bit_offset < int(sizeof(T) * 8);
	     bit_offset += 8, ++l) {

		if (k == 0) {
			return;
		}

		device_function::explicit_unroll::
		    fill<IH, block_dim, all_sh_histograms_length>(
		        temp_storage.histograms, 0, tid);

		__syncthreads();

		device_function::implicit_unroll::bin_values256(data,
		                                                N,
		                                                temp_storage.histograms,
		                                                num_sh_histograms,
		                                                blockIdx.x,
		                                                grid_dim,
		                                                tid,
		                                                block_dim,
		                                                bit_offset,
		                                                prefix,
		                                                unary_functor,
		                                                bin_index_transform);

		__syncthreads();

		//
		// Sum all histograms
		//
		if (num_sh_histograms > 1) {
			thrustshift::device_function::explicit_unroll::sum_subsequent_into<
			    IH,
			    block_dim,
			    histogram_length,
			    num_sh_histograms>(
			    temp_storage.histograms, temp_storage.histograms, tid);
			__syncthreads();
		}

		IH h = temp_storage.histograms[tid];
		IH hcum;
		// it is faster if all blocks do already the scan before
		// the histograms are summed up.
		BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
		histograms[l * grid_dim * histogram_length + bid * histogram_length +
		           tid] = hcum;

		//
		// Store block histogram to global memory
		//
		// thrustshift::block_copy<block_dim, histogram_length>(
		//     temp_storage.histograms, histograms + bid * histogram_length);

		grid.sync();
		static_assert(histogram_length == block_dim);
#pragma unroll
		for (int histo_id = 1; histo_id < grid_dim; ++histo_id) {
			hcum += histograms[l * grid_dim * histogram_length +
			                   histo_id * histogram_length + tid];
		}
		//			BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
		//			__syncthreads();
		temp_storage.histograms[tid] = hcum;
		__syncthreads();
		IH hleft = 0;
		if (tid > 0) {
			hleft = temp_storage.histograms[tid - 1];
		}

		const int i = histogram_length - tid - 1;
		if (hcum >= k && hleft < k) {
			// Only one thread is expected to enter this branch
			prefix = (prefix << 8) | uint64_t(i);
			// all values with this prefix and larger are included
			sh_k = hcum == k ? 0 : k - hleft;
			sh_prefix = prefix;
			if (bid == 0) {
				*prefix_ = prefix;
				*bit_offset_ = bit_offset + 8;
			}
		}

		__syncthreads();
		k = sh_k;
		prefix = sh_prefix;
	}
}

template <typename T,
          typename IH,
          int block_dim,
          int num_sh_histograms,
          int num_gl_histograms>
__global__ void k_select_radix_dynamic_parallelism(const T* data,
                                                   int N,
                                                   IH* histograms,
                                                   int* bit_offset_,
                                                   uint64_t* prefix_,
                                                   int k) {

	constexpr int histogram_length = 256;

	auto unary_functor = [](T x) {
		using std::abs;
		return abs(x);
	};

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histogram[histogram_length];
	} temp_storage;

	__shared__ int sh_k;
	__shared__ uint64_t sh_prefix;

	static_assert(block_dim == 256);

	constexpr int child_grid_dim = num_gl_histograms;
	constexpr int child_block_dim = 256;

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	uint64_t prefix = 0;

	auto bin_index_transform = [](auto i) { return histogram_length - i - 1; };

#pragma unroll
	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {

		if (k == 0) {
			return;
		}

		if (tid == 0) {

			kernel::bin_values256<T,
			                      child_block_dim,
			                      child_grid_dim,
			                      num_sh_histograms>
			    <<<child_grid_dim, child_block_dim>>>(data,
			                                          N,
			                                          histograms,
			                                          bit_offset,
			                                          prefix,
			                                          unary_functor,
			                                          bin_index_transform);
			cudaDeviceSynchronize();
		}
		__syncthreads();

		IH h = 0;
		IH hcum;

#pragma unroll
		for (int histo_id = 0; histo_id < num_gl_histograms; ++histo_id) {
			h += histograms[histo_id * histogram_length + tid];
		}

		BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
		__syncthreads();

		temp_storage.histogram[tid] = hcum;
		__syncthreads();

		IH hleft = 0;
		if (tid > 0) {
			hleft = temp_storage.histogram[tid - 1];
		}

		if (hcum >= k && hleft < k) {
			const int i = histogram_length - tid - 1;
			// Only one thread is expected to enter this branch
			prefix = (prefix << 8) | uint64_t(i);
			// all values with this prefix and larger are included
			k = hcum == k ? 0 : k - hleft;

			// Write to shared memory
			sh_k = k;
			sh_prefix = prefix;

			// Write to global memory
			//*k_ = k;
			*prefix_ = prefix;
			*bit_offset_ = bit_offset + 8;
		}
		__syncthreads();
		k = sh_k;
		prefix = sh_prefix;
	}
}

template <typename T,
          typename IH,
          int block_dim,
          int num_sh_histograms,
          int child_grid_dim>
__global__ void k_select_radix_dynamic_parallelism_atomic_binning(
    const T* data,
    int N,
    IH* histogram,
    int* bit_offset_,
    uint64_t* prefix_,
    int k) {

	constexpr int histogram_length = 256;

	auto unary_functor = [](T x) {
		using std::abs;
		return abs(x);
	};

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histogram[histogram_length];
	} temp_storage;

	__shared__ int sh_k;
	__shared__ uint64_t sh_prefix;

	static_assert(block_dim == 256);

	constexpr int child_block_dim = 256;

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;
	gsl_ExpectsAudit(bid == 0);

	uint64_t prefix = 0;

	auto bin_index_transform = [](auto i) { return histogram_length - i - 1; };

#pragma unroll
	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {

		if (k == 0) {
			return;
		}

		histogram[tid] = 0;
		__syncthreads();

		if (tid == 0) {

			kernel::bin_values256_atomic<T,
			                             child_block_dim,
			                             child_grid_dim,
			                             num_sh_histograms>
			    <<<child_grid_dim, child_block_dim>>>(data,
			                                          N,
			                                          histogram,
			                                          bit_offset,
			                                          prefix,
			                                          unary_functor,
			                                          bin_index_transform);
			cudaDeviceSynchronize();
		}
		__syncthreads();

		IH h = histogram[tid];
		IH hcum;

		BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
		__syncthreads();

		temp_storage.histogram[tid] = hcum;
		__syncthreads();

		IH hleft = 0;
		if (tid > 0) {
			hleft = temp_storage.histogram[tid - 1];
		}

		if (hcum >= k && hleft < k) {
			const int i = histogram_length - tid - 1;
			// Only one thread is expected to enter this branch
			prefix = (prefix << 8) | uint64_t(i);
			// all values with this prefix and larger are included
			k = hcum == k ? 0 : k - hleft;

			// Write to shared memory
			sh_k = k;
			sh_prefix = prefix;

			// Write to global memory
			//*k_ = k;
			*prefix_ = prefix;
			*bit_offset_ = bit_offset + 8;
		}
		__syncthreads();
		k = sh_k;
		prefix = sh_prefix;
	}
}

template <typename IH, int block_dim>
__global__ void k_select_radix_from_histogram(IH* histogram,
                                              int bit_offset,
                                              uint64_t* prefix_, // only output
                                              int* k_) {

	constexpr int histogram_length = 256;

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histogram[histogram_length];
	} temp_storage;

	static_assert(block_dim == 256);

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	uint64_t prefix = 0;
	int k = *k_;

	IH h = histogram[tid];
	IH hcum;

	BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
	__syncthreads();

	temp_storage.histogram[tid] = hcum;
	__syncthreads();

	IH hleft = 0;
	if (tid > 0) {
		hleft = temp_storage.histogram[tid - 1];
	}

	if (hcum >= k && hleft < k) {
		const int i = histogram_length - tid - 1;
		// Only one thread is expected to enter this branch
		prefix = uint64_t(i);
		// all values with this prefix and larger are included
		k = hcum == k ? 0 : k - hleft;

		*k_ = k;
		*prefix_ = prefix;
	}
}

template <typename IH, int block_dim, bool initialize_offset_prefix_k>
__global__ void k_select_radix_from_histogram_with_ptr(IH* histogram,
                                                       int* bit_offset_,
                                                       uint64_t* prefix_,
                                                       int* k_,
                                                       int k0) {

	constexpr int histogram_length = 256;

	using BlockScan = cub::BlockScan<IH, block_dim>;

	__shared__ union {
		typename BlockScan::TempStorage block_scan;
		IH histogram[histogram_length];
	} temp_storage;

	static_assert(block_dim == 256);

	const int tid = threadIdx.x;
	const int bid = blockIdx.x;

	uint64_t prefix = [&]() -> uint64_t {
		if constexpr (initialize_offset_prefix_k) {
			return 0;
		}
		else {
			return *prefix_;
		}
	}();
	int k = [&] {
		if constexpr (initialize_offset_prefix_k) {
			return k0;
		}
		else {
			return *k_;
		}
	}();
	int bit_offset = [&] {
		if constexpr (initialize_offset_prefix_k) {
			return 0;
		}
		else {
			return *bit_offset_;
		}
	}();

	if (k <= 0) {
		return;
	}

	IH h = histogram[tid];
	IH hcum;

	BlockScan(temp_storage.block_scan).InclusiveSum(h, hcum);
	__syncthreads();

	temp_storage.histogram[tid] = hcum;
	__syncthreads();

	IH hleft = 0;
	if (tid > 0) {
		hleft = temp_storage.histogram[tid - 1];
	}

	if (hcum >= k && hleft < k) {
		const int i = histogram_length - tid - 1;
		// Only one thread is expected to enter this branch
		prefix = (prefix << 8) | uint64_t(i);
		// all values with this prefix and larger are included
		k = hcum == k ? 0 : k - hleft;

		*k_ = k;
		*prefix_ = prefix;
		*bit_offset_ = bit_offset + 8;
	}
}

} // namespace kernel

namespace async {

template <typename T, class MemoryResource, class F0, class F1>
void bin_values256(cuda::stream_t& stream,
                   gsl_lite::span<const T> values,
                   gsl_lite::span<int> histogram,
                   int bit_offset,
                   uint64_t prefix,
                   F0 unary_functor,
                   F1 bin_index_transform,
                   MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	gsl_Expects(histogram.size() == histogram_length);

	constexpr int block_dim = 256;
	// RTX 2080 Ti has 68 SMs, `4*` to get best occupancy in practice
	constexpr int num_histograms = 4 * 68;
	constexpr int num_sh_histograms = 1;

	auto c = cuda::make_launch_config(num_histograms, block_dim);

	auto tmp_mem = make_not_a_vector<int>(num_histograms * histogram_length,
	                                      delayed_memory_resource);

	std::cout << "bit_offset = " << bit_offset << ", prefix = " << prefix
	          << std::endl;

	auto histograms = tmp_mem.to_span();

	cuda::enqueue_launch(kernel::bin_values256<T,
	                                           block_dim,
	                                           num_histograms,
	                                           num_sh_histograms,
	                                           F0,
	                                           F1>,
	                     stream,
	                     c,
	                     values.data(),
	                     values.size(),
	                     histograms.data(),
	                     bit_offset,
	                     prefix,
	                     unary_functor,
	                     bin_index_transform);

	cuda::enqueue_launch(kernel::sum_subsequent_into<int,
	                                                 block_dim,
	                                                 histogram_length,
	                                                 num_histograms>,
	                     stream,
	                     cuda::make_launch_config(1, 256),
	                     histograms.data(),
	                     histogram.data());
}

template <typename T, class MemoryResource, class F0, class F1>
void bin_values256(cuda::stream_t& stream,
                   gsl_lite::span<const T> values,
                   gsl_lite::span<int> histogram,
                   int bit_offset,
                   uint64_t prefix,
                   int block_dim,
                   int num_histograms,
                   int num_sh_histograms,
                   F0 unary_functor,
                   F1 bin_index_transform,
                   MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	gsl_Expects(histogram.size() == histogram_length);

	auto block_dim_v = makeshift::expand(
	    block_dim, MAKESHIFT_CONSTVAL(std::array{64, 128, 256, 512}));
	auto num_histograms_v = makeshift::expand(
	    num_histograms,
	    MAKESHIFT_CONSTVAL(std::array{68, 2 * 68, 3 * 68, 4 * 68, 8 * 68}));
	auto num_sh_histograms_v = makeshift::expand(
	    num_sh_histograms, MAKESHIFT_CONSTVAL(std::array{1, 2, 3, 4, 8}));

	std::visit(
	    [&](auto block_dim, auto num_histograms, auto num_sh_histograms) {
		    auto num_warps = block_dim / warp_size;
		    gsl_Expects(num_warps % num_sh_histograms == 0);
		    gsl_Expects(num_warps >= num_sh_histograms);
		    auto c =
		        cuda::make_launch_config(int(num_histograms), int(block_dim));

		    auto tmp_mem = make_not_a_vector<int>(
		        num_histograms * histogram_length, delayed_memory_resource);

		    auto histograms = tmp_mem.to_span();

		    cuda::enqueue_launch(kernel::bin_values256<T,
		                                               block_dim,
		                                               num_histograms,
		                                               num_sh_histograms,
		                                               F0,
		                                               F1>,
		                         stream,
		                         c,
		                         values.data(),
		                         values.size(),
		                         histograms.data(),
		                         bit_offset,
		                         prefix,
		                         unary_functor,
		                         bin_index_transform);

		    cuda::enqueue_launch(kernel::sum_subsequent_into<int,
		                                                     block_dim,
		                                                     histogram_length,
		                                                     num_histograms>,
		                         stream,
		                         cuda::make_launch_config(1, 256),
		                         histograms.data(),
		                         histogram.data());
	    },
	    block_dim_v,
	    num_histograms_v,
	    num_sh_histograms_v);
}

template <typename T, class MemoryResource, class F0, class F1>
void bin_values256_threadfence(cuda::stream_t& stream,
                               gsl_lite::span<const T> values,
                               gsl_lite::span<int> histogram,
                               int bit_offset,
                               uint64_t prefix,
                               int block_dim,
                               int num_histograms,
                               int num_sh_histograms,
                               F0 unary_functor,
                               F1 bin_index_transform,
                               MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	gsl_Expects(histogram.size() == histogram_length);

	auto block_dim_v =
	    makeshift::expand(block_dim, MAKESHIFT_CONSTVAL(std::array{256, 512}));
	auto num_histograms_v = makeshift::expand(
	    num_histograms,
	    MAKESHIFT_CONSTVAL(std::array{68, 2 * 68, 3 * 68, 4 * 68, 8 * 68}));
	auto num_sh_histograms_v = makeshift::expand(
	    num_sh_histograms, MAKESHIFT_CONSTVAL(std::array{1, 2, 3, 4, 8}));

	auto tmp = make_not_a_vector<unsigned>(2, delayed_memory_resource);
	auto tickets = tmp.to_span();

	std::visit(
	    [&](auto block_dim, auto num_histograms, auto num_sh_histograms) {
		    async::fill(stream, tickets, 0);
		    auto num_warps = block_dim / warp_size;
		    gsl_Expects(num_warps % num_sh_histograms == 0);
		    gsl_Expects(num_warps >= num_sh_histograms);
		    auto c =
		        cuda::make_launch_config(int(num_histograms), int(block_dim));

		    auto tmp_mem = make_not_a_vector<int>(
		        num_histograms * histogram_length, delayed_memory_resource);

		    auto histograms = tmp_mem.to_span();

		    cuda::enqueue_launch(
		        kernel::bin_values256_threadfence<T,
		                                          block_dim,
		                                          num_histograms,
		                                          num_sh_histograms,
		                                          F0,
		                                          F1>,
		        stream,
		        c,
		        values.data(),
		        values.size(),
		        histograms.data(),
		        bit_offset,
		        prefix,
		        unary_functor,
		        bin_index_transform,
		        tickets.data(),
		        tickets.data() + 1);
	    },
	    block_dim_v,
	    num_histograms_v,
	    num_sh_histograms_v);
}

template <typename T, class MemoryResource, class F0, class F1>
void bin_values256_atomic(cuda::stream_t& stream,
                          gsl_lite::span<const T> values,
                          gsl_lite::span<int> histogram,
                          int bit_offset,
                          uint64_t prefix,
                          int block_dim,
                          int grid_dim,
                          int num_sh_histograms,
                          F0 unary_functor,
                          F1 bin_index_transform,
                          MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	gsl_Expects(histogram.size() == histogram_length);

	auto block_dim_v =
	    makeshift::expand(block_dim, MAKESHIFT_CONSTVAL(std::array{256, 512}));
	auto grid_dim_v = makeshift::expand(
	    grid_dim,
	    MAKESHIFT_CONSTVAL(std::array{68, 2 * 68, 3 * 68, 4 * 68, 8 * 68}));
	auto num_sh_histograms_v = makeshift::expand(
	    num_sh_histograms, MAKESHIFT_CONSTVAL(std::array{1, 2, 3, 4, 8}));

	std::visit(
	    [&](auto block_dim, auto grid_dim, auto num_sh_histograms) {
		    async::fill(stream, histogram, 0);
		    auto num_warps = block_dim / warp_size;
		    gsl_Expects(num_warps % num_sh_histograms == 0);
		    gsl_Expects(num_warps >= num_sh_histograms);
		    auto c = cuda::make_launch_config(int(grid_dim), int(block_dim));

		    cuda::enqueue_launch(kernel::bin_values256_atomic<T,
		                                                      block_dim,
		                                                      grid_dim,
		                                                      num_sh_histograms,
		                                                      F0,
		                                                      F1>,
		                         stream,
		                         c,
		                         values.data(),
		                         values.size(),
		                         histogram.data(),
		                         bit_offset,
		                         prefix,
		                         unary_functor,
		                         bin_index_transform);
	    },
	    block_dim_v,
	    grid_dim_v,
	    num_sh_histograms_v);
}

template <typename T, class MemoryResource, class F0, class F1>
void bin_values256_atomic_with_ptr(cuda::stream_t& stream,
                                   gsl_lite::span<const T> values,
                                   gsl_lite::span<int> histogram,
                                   int bit_offset,
                                   uint64_t* prefix,
                                   int* k,
                                   int k0,
                                   bool use_k0_and_zero_prefix,
                                   int block_dim,
                                   int grid_dim,
                                   int num_sh_histograms,
                                   F0 unary_functor,
                                   F1 bin_index_transform,
                                   MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	gsl_Expects(histogram.size() == histogram_length);

	auto block_dim_v =
	    makeshift::expand(block_dim, MAKESHIFT_CONSTVAL(std::array{256, 512}));
	auto grid_dim_v = makeshift::expand(
	    grid_dim,
	    MAKESHIFT_CONSTVAL(std::array{68, 2 * 68, 3 * 68, 4 * 68, 8 * 68}));
	auto num_sh_histograms_v = makeshift::expand(
	    num_sh_histograms, MAKESHIFT_CONSTVAL(std::array{1, 2, 3, 4, 8}));
	auto use_k0_and_zero_prefix_v = makeshift::expand(
	    use_k0_and_zero_prefix, MAKESHIFT_CONSTVAL(std::array{0, 1}));

	std::visit(
	    [&](auto block_dim,
	        auto grid_dim,
	        auto num_sh_histograms,
	        auto use_k0_and_zero_prefix) {
		    async::fill(stream, histogram, 0);
		    auto num_warps = block_dim / warp_size;
		    gsl_Expects(num_warps % num_sh_histograms == 0);
		    gsl_Expects(num_warps >= num_sh_histograms);
		    auto c = cuda::make_launch_config(int(grid_dim), int(block_dim));

		    cuda::enqueue_launch(
		        kernel::bin_values256_atomic_with_ptr<T,
		                                              block_dim,
		                                              grid_dim,
		                                              num_sh_histograms,
		                                              use_k0_and_zero_prefix,
		                                              F0,
		                                              F1>,
		        stream,
		        c,
		        values.data(),
		        values.size(),
		        histogram.data(),
		        bit_offset,
		        prefix,
		        k,
		        k0,
		        unary_functor,
		        bin_index_transform);
	    },
	    block_dim_v,
	    grid_dim_v,
	    num_sh_histograms_v,
	    use_k0_and_zero_prefix_v);
}

} // namespace async

template <typename T, class MemoryResource>
std::tuple<uint64_t, int> k_largest_values_abs_radix(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    int k,
    MemoryResource& delayed_memory_resource) {

	auto unary_functor = [] __device__(T x) {
		using std::abs;
		return abs(x);
	};
	using IH = int;
	auto bin_index_transform = [] __device__(IH i) { return i; };

	constexpr int histogram_length = 256;
	auto tmp = make_not_a_vector<IH>(histogram_length, delayed_memory_resource);

	auto histogram = tmp.to_span();

	// {
	// 	const size_t s = histogram.size() * sizeof(IH);
	// 	void* ptr = reinterpret_cast<void*>(histogram.data());
	// 	cuda::memory::managed::region_t mr{ptr, s};
	// 	static auto device = stream.device();
	// 	static bool b = false;
	// 	if (!b) {
	// 		mr.set_preferred_location(device);
	// 		b = true;
	// 	}
	// }

	uint64_t prefix = 0;
	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {
		async::bin_values256<T>(
		    stream,
		    values,
		    histogram,
		    bit_offset,
		    prefix,
		    256, // block dim // params determined empirically by benchmarks
		    3 * 68, // num histos
		    1, // num sh histo
		    unary_functor,
		    bin_index_transform,
		    delayed_memory_resource);
		stream.synchronize();
		int acc = histogram[histogram_length - 1];
		int acc_prev = 0;
		for (int i = histogram_length - 2; i >= 0; --i) {
			acc += histogram[i];
			if (acc >= k) {
				prefix = (prefix << 8) | uint64_t(i);
				if (acc == k) {
					// all values with this prefix and larger are included
					return {prefix, bit_offset + 8};
				}
				k = k - acc_prev;
				break;
			}
			acc_prev = acc;
		}
	}
	return {prefix, sizeof(T) * 8};
}

template <typename T, class MemoryResource>
std::tuple<uint64_t, int> k_largest_values_abs_radix_atomic(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    int k,
    MemoryResource& delayed_memory_resource) {

	auto unary_functor = [] __device__(T x) {
		using std::abs;
		return abs(x);
	};
	using IH = int;
	auto bin_index_transform = [] __device__(IH i) { return i; };

	constexpr int histogram_length = 256;
	auto tmp = make_not_a_vector<IH>(histogram_length, delayed_memory_resource);

	auto histogram = tmp.to_span();

	//{
	//	const size_t s = histogram.size() * sizeof(IH);
	//	void* ptr = reinterpret_cast<void*>(histogram.data());
	//	cuda::memory::managed::region_t mr{ptr, s};
	//	static auto device = stream.device();
	//	static bool b = false;
	//	if (!b) {
	//		//mr.set_preferred_location(device);
	//		//cudaMemAdvise(ptr, s, cudaMemAdviseSetAccessedBy, cudaCpuDeviceId);
	//		cudaMemAdvise(ptr, s, cudaMemAdviseSetAccessedBy, device.id());
	//		b = true;
	//	}
	//}

	uint64_t prefix = 0;
	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {
		async::bin_values256_atomic<T>(
		    stream,
		    values,
		    histogram,
		    bit_offset,
		    prefix,
		    256, // block dim // params determined empirically by benchmarks
		    68 * 4, // grid dim
		    1, // num sh histo
		    unary_functor,
		    bin_index_transform,
		    delayed_memory_resource);
		stream.synchronize();
		int acc = histogram[histogram_length - 1];
		int acc_prev = 0;
		for (int i = histogram_length - 2; i >= 0; --i) {
			acc += histogram[i];
			if (acc >= k) {
				prefix = (prefix << 8) | uint64_t(i);
				if (acc == k) {
					// all values with this prefix and larger are included
					return {prefix, bit_offset + 8};
				}
				k = k - acc_prev;
				break;
			}
			acc_prev = acc;
		}
	}
	return {prefix, sizeof(T) * 8};

	// HACK BEGIN
	// constexpr std::array<uint64_t, 4> prefixes{0, 75, 19224, 4921493};
	// for (int bit_offset = 0, i = 0; bit_offset < int(sizeof(T) * 8);
	//      bit_offset += 8, ++i) {
	// 	uint64_t prefix = prefixes[i];
	// 	async::bin_values256_atomic<T>(
	// 	    stream,
	// 	    values,
	// 	    histogram,
	// 	    bit_offset,
	// 	    prefix,
	// 	    256, // block dim // params determined empirically by benchmarks
	// 	    68 * 4, // grid dim
	// 	    1, // num sh histo
	// 	    unary_functor,
	// 	    bin_index_transform,
	// 	    delayed_memory_resource);
	// 	stream.synchronize();
	// 	// int acc = histogram[histogram_length - 1];
	// 	// int acc_prev = 0;
	// 	// for (int i = histogram_length - 2; i >= 0; --i) {
	// 	// 	acc += histogram[i];
	// 	// 	if (acc >= k) {
	// 	// 		prefix = (prefix << 8) | uint64_t(i);
	// 	// 		if (acc == k) {
	// 	// 			// all values with this prefix and larger are included
	// 	// 			return {prefix, bit_offset + 8};
	// 	// 		}
	// 	// 		k = k - acc_prev;
	// 	// 		break;
	// 	// 	}
	// 	// 	acc_prev = acc;
	// 	// }
	// }
	// return {prefixes[3], sizeof(T) * 8};
	// HACK END
}

template <typename T, class MemoryResource>
std::tuple<uint64_t, int> k_largest_values_abs_radix_atomic_devicehisto(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    int k,
    MemoryResource& delayed_memory_resource) {

	auto unary_functor = [] __device__(T x) {
		using std::abs;
		return abs(x);
	};
	using IH = int;

	constexpr int histogram_length = 256;
	auto bin_index_transform = [] __device__(IH i) {
		return histogram_length - i - 1;
	};

	auto tmp = make_not_a_vector<IH>(histogram_length, delayed_memory_resource);
	auto tmp0 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<int>(1, delayed_memory_resource);
	auto prefix_s = tmp0.to_span();
	auto k_s = tmp1.to_span();
	async::fill(stream, k_s, k);

	auto histogram = tmp.to_span();
	uint64_t prefix = 0;

	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {
		async::bin_values256_atomic<T>(
		    stream,
		    values,
		    histogram,
		    bit_offset,
		    prefix,
		    256, // block dim // params determined empirically by benchmarks
		    68 * 4, // grid dim
		    1, // num sh histo
		    unary_functor,
		    bin_index_transform,
		    delayed_memory_resource);

		cuda::enqueue_launch(kernel::k_select_radix_from_histogram<IH, 256>,
		                     stream,
		                     cuda::make_launch_config(1, 256),
		                     histogram.data(),
		                     bit_offset,
		                     prefix_s.data(),
		                     k_s.data());
		stream.synchronize();
		prefix = (prefix << 8) | prefix_s[0];
		k = k_s[0];
		if (k <= 0) {
			return {prefix, bit_offset + 8};
		}
	}
	return {prefix, sizeof(T) * 8};
	// HACK BEGIN
	// constexpr std::array<uint64_t, 4> prefixes{0, 75, 19224, 4921493};
	// for (int bit_offset = 0, i = 0; bit_offset < int(sizeof(T) * 8);
	//      bit_offset += 8, ++i) {
	// 	uint64_t prefix = prefixes[i];
	// 	async::bin_values256_atomic<T>(
	// 	    stream,
	// 	    values,
	// 	    histogram,
	// 	    bit_offset,
	// 	    prefix,
	// 	    256, // block dim // params determined empirically by benchmarks
	// 	    68 * 4, // grid dim
	// 	    1, // num sh histo
	// 	    unary_functor,
	// 	    bin_index_transform,
	// 	    delayed_memory_resource);

	// 	cuda::enqueue_launch(kernel::k_select_radix_from_histogram<IH, 256>,
	// 	                     stream,
	// 	                     cuda::make_launch_config(1, 256),
	// 	                     histogram.data(),
	// 	                     bit_offset,
	// 	                     prefix_s.data(),
	// 	                     k_s.data());
	// 	// stream.synchronize();
	// 	// prefix = (prefix << 8) | prefix_s[0];
	// 	// k = k_s[0];
	// 	// if (k <= 0) {
	// 	// 	return {prefix, bit_offset + 8};
	// 	// }
	// }
	// return {prefixes[3], sizeof(T) * 8};
	// HACK END
}

namespace async {

template <typename T, class MemoryResource>
void k_largest_values_abs_radix_atomic_devicehisto_with_ptr(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    uint64_t* prefix, // only output
    int* bit_offset, //only output
    int k,
    MemoryResource& delayed_memory_resource) {

	auto unary_functor = [] __device__(T x) {
		using std::abs;
		return abs(x);
	};
	using IH = int;

	constexpr int histogram_length = 256;
	auto bin_index_transform = [] __device__(IH i) {
		return histogram_length - i - 1;
	};

	auto tmp = make_not_a_vector<IH>(histogram_length, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<int>(1, delayed_memory_resource);
	auto k_s = tmp1.to_span();
	gsl_lite::span<int> bit_offset_s({bit_offset, 1});
	gsl_lite::span<uint64_t> prefix_s({prefix, 1});

	auto histogram = tmp.to_span();

	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {

		async::bin_values256_atomic_with_ptr<T>(
		    stream,
		    values,
		    histogram,
		    bit_offset,
		    prefix_s.data(),
		    k_s.data(),
		    k,
		    bit_offset == 0,
		    256, // block dim // params determined empirically by benchmarks
		    68 * 4, // grid dim
		    1, // num sh histo
		    unary_functor,
		    bin_index_transform,
		    delayed_memory_resource);

		auto flag_v = makeshift::expand(bit_offset == 0,
		                                MAKESHIFT_CONSTVAL(std::array{0, 1}));
		std::visit(
		    [&](auto flag) {
			    cuda::enqueue_launch(
			        kernel::
			            k_select_radix_from_histogram_with_ptr<IH, 256, flag>,
			        stream,
			        cuda::make_launch_config(1, 256),
			        histogram.data(),
			        bit_offset_s.data(),
			        prefix_s.data(),
			        k_s.data(),
			        k);
		    },
		    flag_v);
	}
}

} // namespace async

template <typename T, class MemoryResource>
std::tuple<uint64_t, int> k_largest_values_abs_radix_with_cub(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    int k,
    MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	auto tmp =
	    make_not_a_vector<int>(histogram_length, delayed_memory_resource);
	auto histogram = tmp.to_span();

	using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;

	uint64_t prefix = 0;
	I lower_level = 0; // inclusive
	I upper_level = std::numeric_limits<I>::max(); // exclusive
	const int N = values.size();

	for (int bit_offset = 0; bit_offset < int(sizeof(T) * 8); bit_offset += 8) {

		auto sample_iterator = thrust::make_transform_iterator(
		    values.data(), [prefix, bit_offset] __device__(const T& x) {
			    using std::abs;
			    const T abs_x = abs(x);
			    const I i = *reinterpret_cast<I*>((void*) (&abs_x));
			    return i;
		    });
		async::bin_values_into_histogram(stream,
		                                 sample_iterator,
		                                 histogram.begin(),
		                                 histogram_length,
		                                 lower_level,
		                                 upper_level,
		                                 N,
		                                 delayed_memory_resource);
		stream.synchronize();
		int acc = histogram[histogram_length - 1];
		int acc_prev = 0;
		for (int i = histogram_length - 2; i >= 0; --i) {
			acc += histogram[i];
			if (acc >= k) {
				prefix = (prefix << 8) | uint64_t(i);
				if (acc == k) {
					// all values with this prefix and larger are included
					return {prefix, bit_offset + 8};
				}
				k = k - acc_prev;
				break;
			}
			acc_prev = acc;
		}
		lower_level = prefix << (sizeof(I) * 8 - (bit_offset + 8));
		upper_level = (prefix + 1) << (sizeof(I) * 8 - (bit_offset + 8));
	}
	return {prefix, sizeof(T) * 8};
}

// selected_values.size() == values.size() because CUB might select more values, if e.g. values are all equal
template <typename T, class MemoryResource>
void select_k_largest_values_abs(cuda::stream_t& stream,
                                 gsl_lite::span<const T> values,
                                 gsl_lite::span<T> selected_values,
                                 gsl_lite::span<int> selected_indices,
                                 int k,
                                 MemoryResource& delayed_memory_resource) {

	const auto tup = thrustshift::k_largest_values_abs_radix<T>(
	    stream, values, k, delayed_memory_resource);
	const auto prefix = std::get<0>(tup);
	const auto bit_offset = std::get<1>(tup);
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};
	auto tmp = make_not_a_vector<int>(1, delayed_memory_resource);
	auto num_selected = tmp.to_span();

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

template <typename T, class MemoryResource>
void select_k_largest_values_abs_atomic(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    gsl_lite::span<T> selected_values,
    gsl_lite::span<int> selected_indices,
    int k,
    MemoryResource& delayed_memory_resource) {

	const auto tup = thrustshift::k_largest_values_abs_radix_atomic<T>(
	    stream, values, k, delayed_memory_resource);
	const auto prefix = std::get<0>(tup);
	const auto bit_offset = std::get<1>(tup);
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};
	auto tmp = make_not_a_vector<int>(1, delayed_memory_resource);
	auto num_selected = tmp.to_span();

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

// selected_values.size() == values.size() because CUB might select more values, if e.g. values are all equal
template <typename T, class MemoryResource>
void select_k_largest_values_abs_with_cub(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    gsl_lite::span<T> selected_values,
    gsl_lite::span<int> selected_indices,
    int k,
    MemoryResource& delayed_memory_resource) {

	const auto tup = thrustshift::k_largest_values_abs_radix_with_cub<T>(
	    stream, values, k, delayed_memory_resource);
	const auto prefix = std::get<0>(tup);
	const auto bit_offset = std::get<1>(tup);
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};
	auto tmp = make_not_a_vector<int>(1, delayed_memory_resource);
	auto num_selected = tmp.to_span();

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

namespace cooperative {

template <typename T, class MemoryResource>
void select_k_largest_values_abs(cuda::stream_t& stream,
                                 gsl_lite::span<const T> values,
                                 gsl_lite::span<T> selected_values,
                                 gsl_lite::span<int> selected_indices,
                                 int k,
                                 MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	constexpr int block_dim = 256;
	constexpr int grid_dim = 2 * 68;
	constexpr int num_sh_histograms = 1;

	auto tmp = make_not_a_vector<int>(3, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto tmp2 = make_not_a_vector<unsigned>(histogram_length * grid_dim,
	                                        delayed_memory_resource);

	auto num_selected = tmp.to_span().first(1);
	auto bit_offset_s = tmp.to_span().subspan(1, 1);
	auto prefix_s = tmp1.to_span().first(1);
	auto histograms = tmp2.to_span();
	auto k_s = tmp.to_span().subspan(2, 1);
	k_s[0] = k;

	auto c = cuda::make_launch_config(grid_dim, block_dim);
	c.block_cooperation = true;

	const int N = values.size();

	cuda::enqueue_launch(
	    kernel::
	        k_select_radix<T, unsigned, block_dim, grid_dim, num_sh_histograms>,
	    stream,
	    c,
	    values.data(),
	    N,
	    histograms.data(),
	    bit_offset_s.data(),
	    prefix_s.data(),
	    k_s.data());
	stream.synchronize();

	const auto prefix = prefix_s[0];
	const auto bit_offset = bit_offset_s[0];
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

template <typename T, class MemoryResource>
void select_k_largest_values_abs2(cuda::stream_t& stream,
                                  gsl_lite::span<const T> values,
                                  gsl_lite::span<T> selected_values,
                                  gsl_lite::span<int> selected_indices,
                                  int k,
                                  MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	constexpr int block_dim = 256;
	constexpr int grid_dim = 2 * 68;
	constexpr int num_sh_histograms = 1;

	auto tmp = make_not_a_vector<int>(3, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto tmp2 = make_not_a_vector<unsigned>(histogram_length * grid_dim * 4,
	                                        delayed_memory_resource);

	auto num_selected = tmp.to_span().first(1);
	auto bit_offset_s = tmp.to_span().subspan(1, 1);
	auto prefix_s = tmp1.to_span().first(1);
	auto histograms = tmp2.to_span();
	auto k_s = tmp.to_span().subspan(2, 1);
	k_s[0] = k;

	auto c = cuda::make_launch_config(grid_dim, block_dim);
	c.block_cooperation = true;
	const int N = values.size();

	cuda::enqueue_launch(kernel::k_select_radix2<T,
	                                             unsigned,
	                                             block_dim,
	                                             grid_dim,
	                                             num_sh_histograms>,
	                     stream,
	                     c,
	                     values.data(),
	                     N,
	                     histograms.data(),
	                     bit_offset_s.data(),
	                     prefix_s.data(),
	                     k);
	stream.synchronize();

	const auto prefix = prefix_s[0];
	const auto bit_offset = bit_offset_s[0];
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

} // namespace cooperative

namespace dynamic_parallelism {

template <typename T, class MemoryResource>
std::tuple<uint64_t, int> k_largest_values_abs_radix_atomic_binning(
    cuda::stream_t& stream,
    gsl_lite::span<const T> values,
    int k,
    bool nosync,
    MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	constexpr int block_dim = 256;
	constexpr int child_grid_dim = 68 * 4;
	constexpr int num_sh_histograms = 1;

	using IH = int;
	auto tmp = make_not_a_vector<int>(1, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto tmp2 =
	    make_not_a_vector<IH>(histogram_length, delayed_memory_resource);

	auto bit_offset_s = tmp.to_span();
	auto prefix_s = tmp1.to_span();
	auto histogram = tmp2.to_span();

	auto c = cuda::make_launch_config(1, block_dim);
	const int N = values.size();

	cuda::enqueue_launch(
	    kernel::k_select_radix_dynamic_parallelism_atomic_binning<
	        T,
	        IH,
	        block_dim,
	        num_sh_histograms,
	        child_grid_dim>,
	    stream,
	    c,
	    values.data(),
	    N,
	    histogram.data(),
	    bit_offset_s.data(),
	    prefix_s.data(),
	    k);
	// Alternatively give CUB select if a lambda, which accesses the values, which are already on the GPU
	if (!nosync) {
		stream.synchronize();
		return {prefix_s[0], bit_offset_s[0]};
	}
	else {
		return {0, 0};
	}
}

template <typename T, class MemoryResource>
void select_k_largest_values_abs(cuda::stream_t& stream,
                                 gsl_lite::span<const T> values,
                                 gsl_lite::span<T> selected_values,
                                 gsl_lite::span<int> selected_indices,
                                 int k,
                                 MemoryResource& delayed_memory_resource) {

	constexpr int histogram_length = 256;
	constexpr int block_dim = 256;
	constexpr int num_sh_histograms = 1;
	constexpr int num_gl_histograms = 1 * 68;

	using IH = int;
	auto tmp = make_not_a_vector<int>(3, delayed_memory_resource);
	auto tmp1 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto tmp2 = make_not_a_vector<IH>(histogram_length * num_gl_histograms,
	                                  delayed_memory_resource);

	auto num_selected = tmp.to_span().first(1);
	auto bit_offset_s = tmp.to_span().subspan(1, 1);
	auto prefix_s = tmp1.to_span().first(1);
	auto histograms = tmp2.to_span();
	// auto k_s = tmp.to_span().subspan(2, 1);
	// k_s[0] = k;

	auto c = cuda::make_launch_config(1, block_dim);
	const int N = values.size();

	cuda::enqueue_launch(
	    kernel::k_select_radix_dynamic_parallelism<T,
	                                               IH,
	                                               block_dim,
	                                               num_sh_histograms,
	                                               num_gl_histograms>,
	    stream,
	    c,
	    values.data(),
	    N,
	    histograms.data(),
	    bit_offset_s.data(),
	    prefix_s.data(),
	    k);
	stream.synchronize();

	const auto prefix = prefix_s[0];
	const auto bit_offset = bit_offset_s[0];
	auto select_op = [prefix,
	                  bit_offset] __device__(const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
	stream.synchronize();
}

} // namespace dynamic_parallelism

namespace async {

template <typename T, class MemoryResource>
void select_k_largest_values_abs(cuda::stream_t& stream,
                                 gsl_lite::span<const T> values,
                                 gsl_lite::span<T> selected_values,
                                 gsl_lite::span<int> selected_indices,
                                 int k,
                                 MemoryResource& delayed_memory_resource) {

	auto tmp0 = make_not_a_vector<uint64_t>(1, delayed_memory_resource);
	auto prefix_s = tmp0.to_span();
	auto tmp1 = make_not_a_vector<int>(1, delayed_memory_resource);
	auto bit_offset_s = tmp1.to_span();

	uint64_t* prefix_ptr = prefix_s.data();
	int* bit_offset_ptr = bit_offset_s.data();

	async::k_largest_values_abs_radix_atomic_devicehisto_with_ptr<T>(
	    stream, values, prefix_ptr, bit_offset_ptr, k, delayed_memory_resource);

	auto select_op = [prefix_ptr, bit_offset_ptr] __device__(
	                     const thrust::tuple<T, int>& tup) {
		using std::abs;
		auto x = abs(thrust::get<0>(tup));
		using I = typename thrustshift::make_uintegral_of_equal_size<T>::type;
		const I i = *reinterpret_cast<I*>((void*) (&x));
		// No performance effects measured when we load these values here from
		// global memory. Probably they end up in the caches and can be loaded fast.
		uint64_t prefix = *prefix_ptr;
		int bit_offset = *bit_offset_ptr;
		// uint64_t prefix = 1259902258;
		// int bit_offset = 32;
		return (i >> sizeof(I) * 8 - bit_offset) >= static_cast<I>(prefix);
	};
	auto tmp = make_not_a_vector<int>(1, delayed_memory_resource);
	auto num_selected = tmp.to_span();

	async::select_if_with_index(stream,
	                            values,
	                            selected_values,
	                            selected_indices,
	                            num_selected.data(),
	                            select_op,
	                            delayed_memory_resource);
}

} // namespace async

} // namespace thrustshift


#include <gsl/gsl_test.h>
#include <gsl/gsl_ieee_utils.h>
#include <gsl/gsl_math.h>
#include "gsl_cblas.h"

#include "tests.h"

void
test_syr (void) {
const double flteps = 1e-4, dbleps = 1e-6;
  {
   int order = 101;
   int uplo = 121;
   int N = 1;
   int lda = 1;
   float alpha = 0.1f;
   float A[] = { -0.291f };
   float X[] = { 0.845f };
   int incX = -1;
   float A_expected[] = { -0.219597f };
   cblas_ssyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], flteps, "ssyr(case 1402)");
     }
   };
  };

  {
   int order = 101;
   int uplo = 122;
   int N = 1;
   int lda = 1;
   float alpha = 0.1f;
   float A[] = { -0.291f };
   float X[] = { 0.845f };
   int incX = -1;
   float A_expected[] = { -0.219597f };
   cblas_ssyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], flteps, "ssyr(case 1403)");
     }
   };
  };

  {
   int order = 102;
   int uplo = 121;
   int N = 1;
   int lda = 1;
   float alpha = 0.1f;
   float A[] = { -0.291f };
   float X[] = { 0.845f };
   int incX = -1;
   float A_expected[] = { -0.219597f };
   cblas_ssyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], flteps, "ssyr(case 1404)");
     }
   };
  };

  {
   int order = 102;
   int uplo = 122;
   int N = 1;
   int lda = 1;
   float alpha = 0.1f;
   float A[] = { -0.291f };
   float X[] = { 0.845f };
   int incX = -1;
   float A_expected[] = { -0.219597f };
   cblas_ssyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], flteps, "ssyr(case 1405)");
     }
   };
  };

  {
   int order = 101;
   int uplo = 121;
   int N = 1;
   int lda = 1;
   double alpha = -0.3;
   double A[] = { -0.65 };
   double X[] = { -0.891 };
   int incX = -1;
   double A_expected[] = { -0.8881643 };
   cblas_dsyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], dbleps, "dsyr(case 1406)");
     }
   };
  };

  {
   int order = 101;
   int uplo = 122;
   int N = 1;
   int lda = 1;
   double alpha = -0.3;
   double A[] = { -0.65 };
   double X[] = { -0.891 };
   int incX = -1;
   double A_expected[] = { -0.8881643 };
   cblas_dsyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], dbleps, "dsyr(case 1407)");
     }
   };
  };

  {
   int order = 102;
   int uplo = 121;
   int N = 1;
   int lda = 1;
   double alpha = -0.3;
   double A[] = { -0.65 };
   double X[] = { -0.891 };
   int incX = -1;
   double A_expected[] = { -0.8881643 };
   cblas_dsyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], dbleps, "dsyr(case 1408)");
     }
   };
  };

  {
   int order = 102;
   int uplo = 122;
   int N = 1;
   int lda = 1;
   double alpha = -0.3;
   double A[] = { -0.65 };
   double X[] = { -0.891 };
   int incX = -1;
   double A_expected[] = { -0.8881643 };
   cblas_dsyr(order, uplo, N, alpha, X, incX, A, lda);
   {
     int i;
     for (i = 0; i < 1; i++) {
       gsl_test_rel(A[i], A_expected[i], dbleps, "dsyr(case 1409)");
     }
   };
  };

}


#include <stdio.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_multimin.h>


/**
 *
 * @file core_dgeqp3_pivot.c
 *
 *  PLASMA core_blas kernel
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @version 2.6.0
 * @author Mark Gates
 * @date 2010-11-15
 * @generated d Tue Jan  7 11:44:49 2014
 *
 **/
#include <math.h>
#include <cblas.h>
#include "common.h"

#define A(m,n) BLKADDR( A, double, m, n )

/***************************************************************************//**
 *
 * @ingroup CORE_double
 *
 * CORE_dgeqp3_pivot finds next pivot, pvt, based on maximum column norm.
 * It applies the swap to the matrices A, F, and vectors jpvt, norms1, norms2.
 * If info != 0, it returns immediately, doing no work.
 *
 *******************************************************************************
 *
 * @param[in,out] A
 *         On entry, descriptor for m by n matrix A.
 *         On exit, column k of jj-th block column is swapped with column pvt.
 *
 * @param[in,out] F
 *         On entry, n by nb matrix F.
 *         On exit, row k is swapped with row pvt - jj*nb.
 *         Currently, F is stored column-wise, not tile-wise.
 *
 * @param[in] ldf
 *         Leading dimension of F. ldf >= max(1,A.n).
 *
 * @param[in] jj
 *         Index of current block column, 0 <= jj < A.nt.
 *
 * @param[in] k
 *         Index of current column within block column, 0 <= k < A.nb.
 *
 * @param[in,out] jpvt
 *         Permutation vector, dimension n.
 *         On exit, swaps entries jpvt[k+jj*nb] and jpvt[pvt].
 *
 * @param[in,out] norms1
 *         On entry, vector of partial column norms, dimension n.
 *         On exit, sets norms1[pvt] = norms1[k+jj*nb].
 *
 * @param[in,out] norms2
 *         On entry, vector of original column norms, dimension n.
 *         On exit, sets norms2[pvt] = norms2[k+jj*nb].
 *
 * @param[in] info
 *         Error code from dgeqp3_update; zero if no error.
 **/
#if defined(PLASMA_HAVE_WEAK)
#pragma weak CORE_dgeqp3_pivot = PCORE_dgeqp3_pivot
#define CORE_dgeqp3_pivot PCORE_dgeqp3_pivot
#endif
void CORE_dgeqp3_pivot( PLASMA_desc A, double *F, int ldf,
                        int jj, int k, int *jpvt,
                        double *norms1, double *norms2, int *info )
{
    double *Aij, *Aip;
    int pvt, ii, pp, p, mb, lda, tmp, jk;

    /* since pivot depends on all the norm updates, check info here
     * to detect errors from dgeqp3_updates. */
    if ( *info != 0 )
        return;

    /* jk and pvt are indices in global matrix */
    jk = jj*A.nb + k;
    pvt = jk + cblas_idamax( A.n - jk, &norms1[jk], 1 );
    if ( pvt != jk ) {
        tmp       = jpvt[jk];
        jpvt[jk]  = jpvt[pvt];
        jpvt[pvt] = tmp;

        norms1[pvt] = norms1[jk];  /* don't need to save norms[pvt] */
        norms2[pvt] = norms2[jk];

        cblas_dswap( A.nb, &F[k], ldf, &F[pvt-jj*A.nb], ldf );  /* rows k and pvt */

        pp = pvt / A.nb;  /* tile containing pvt  */
        p  = pvt % A.nb;  /* index within pp tile */
        for( ii = 0; ii < A.mt; ++ii ) {
            mb  = min( A.mb, A.m - ii*A.mb );
            lda = BLKLDD( A, ii );
            Aij = A(ii,jj);
            Aip = A(ii,pp);
            cblas_dswap( mb, &Aij[k*lda], 1, &Aip[p*lda], 1 );  /* cols k and pvt */
        }
    }
}


/*
Copyright (c) 2015, Patrick Weltevrede
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <gsl/gsl_sf.h>
#include "psrsalsa.h"
#define AmoebaAlgorithm 0
int readVonMisesModel(char *filename, vonMises_collection_definition *components, verbose_definition verbose)
{
  int i;
  FILE *fin;
  if(verbose.verbose) printf("Opening %s for reading\n", filename);
  fin = fopen(filename, "r");
  if(fin == NULL) {
    fflush(stdout);
    printerror(verbose.debug, "ERROR readVonMisesModel: Cannot open %s", filename);
    return 0;
  }
  for(components->nrcomponents = 0; components->nrcomponents < maxNrVonMisesComponents; (components->nrcomponents)++) {
    i = fscanf(fin, "%lf %lf %lf", &(components->centre[components->nrcomponents]), &(components->concentration[components->nrcomponents]), &(components->height[components->nrcomponents]));
    if(i != 3) {
      break;
    }else if(verbose.verbose) {
      printf("  component %d: phase=%e concentration=%e amplitude=%e\n", components->nrcomponents, components->centre[components->nrcomponents], components->concentration[components->nrcomponents], components->height[components->nrcomponents]);
    }
  }
  if(verbose.verbose) printf("Closing %s\n", filename);
  fclose(fin);
  if(components->nrcomponents == 0) {
    fflush(stdout);
    printerror(verbose.debug, "ERROR readVonMisesModel: No components found in %s", filename);
    return 0;
  }
  return 1;
}
double calcVonMisesFunction2(double centre, double concentration, double height, double phase, double shift)
{
  double y;
  y = exp((cos(2.0*M_PI*(phase-centre-shift))-1.0)*concentration) * height;
  return y;
}
double integrateVonMisesFunction2(double concentration, double height)
{
  double area;
  area = 2*M_PI*height*gsl_sf_bessel_I0_scaled(concentration);
  return area;
}
double widthVonMisesFunction2(double concentration, double ampfrac)
{
  double value = log(ampfrac)/concentration + 1.0;
  if(value < -1.0 || value >= 1.0) {
    return sqrt(-1);
  }
  return 2.0*acos(value);
}
double calcVonMisesFunction(vonMises_collection_definition *components, double phase, double shift)
{
  int n;
  double y;
  y = 0;
  if(components->nrcomponents > 0) {
    for(n = 0; n < components->nrcomponents; n++) {
      y += calcVonMisesFunction2(components->centre[n], components->concentration[n], components->height[n], phase, shift);
    }
  }
  return y;
}
void calcVonMisesProfile(vonMises_collection_definition *components, int nrbins, float *profile, double shift, int normalize)
{
  int i;
  double x, Imax = -1;
  for(i = 0; i < nrbins; i++) {
    profile[i] = 0;
  }
  for(i = 0; i < nrbins; i++) {
    x = i/(double)nrbins;
    profile[i] = calcVonMisesFunction(components, x, shift);
  }
  if(normalize) {
    for(i = 0; i < nrbins; i++) {
      if(profile[i] > Imax)
 Imax = profile[i];
    }
    for(i = 0; i < nrbins; i++) {
      profile[i] /= Imax;
    }
  }
}
float correlateVonMisesFunction(vonMises_collection_definition *components, int nrbins, float *profile, verbose_definition verbose)
{
  float correl_max, *profile2;
  int i;
  int ishift;
  if(verbose.verbose) {
    for(i = 0; i < verbose.indent; i++)
      printf(" ");
    printf("Correlating template with profile\n");
    verbose.indent += 2;
  }
  profile2 = (float *)malloc(nrbins*sizeof(float));
  if(profile2 == NULL) {
    fflush(stdout);
    printerror(verbose.debug, "ERROR correlateVonMisesFunction: Memory allocation error.");
    return 0;
  }
  calcVonMisesProfile(components, nrbins, profile2, 0, 0);
  find_peak_correlation(profile, profile2, nrbins, 0, 1, 1, 1, &ishift, &correl_max, verbose);
  if(verbose.verbose) {
    for(i = 0; i < verbose.indent; i++)
      printf(" ");
    printf("Found a shift of %d bins (%f phase) and max/min correlation is %f.\n", ishift, ishift/(float)nrbins, correl_max);
    for(i = 0; i < verbose.indent; i++)
      printf(" ");
    printf("so you could do a shift by %d bins (%f phase) to rotate profile to align with model.\n", -ishift, -ishift/(float)nrbins);
  }
  free(profile2);
  return ishift/(float)nrbins;
}


#include <config.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_complex.h>

#include <gsl/gsl_fft_complex_float.h>

#define BASE_FLOAT
#include "templates_on.h"
#include "c_pass.h"
#include "c_init.c"
#include "c_main.c"
#include "c_pass_2.c"
#include "c_pass_3.c"
#include "c_pass_4.c"
#include "c_pass_5.c"
#include "c_pass_6.c"
#include "c_pass_7.c"
#include "c_pass_n.c"
#include "bitreverse.c"
#include "c_radix2.c"
#include "templates_off.h"
#undef  BASE_FLOAT


/* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
/*
 *     Copyright 2016-Present Couchbase, Inc.
 *
 *   Use of this software is governed by the Business Source License included
 *   in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
 *   in that file, in accordance with the Business Source License, use of this
 *   software will be governed by the Apache License, Version 2.0, included in
 *   the file licenses/APL2.txt.
 */
#pragma once

#include <gsl/gsl-lite.hpp>
#include <nlohmann/json_fwd.hpp>
#include <cstdint>
#include <string>

namespace cb::crypto {
enum class Algorithm { MD5, SHA1, SHA256, SHA512 };

bool isSupported(Algorithm algorithm);

const int MD5_DIGEST_SIZE = 16;
const int SHA1_DIGEST_SIZE = 20;
const int SHA256_DIGEST_SIZE = 32;
const int SHA512_DIGEST_SIZE = 64;

/**
 * Generate a HMAC digest of the key and data by using the given
 * algorithm
 *
 * @throws std::invalid_argument - unsupported algorithm
 *         std::runtime_error - Failures generating the HMAC
 */
std::string HMAC(Algorithm algorithm,
                 std::string_view key,
                 std::string_view data);

/**
 * Generate a PBKDF2_HMAC digest of the key and data by using the given
 * algorithm
 *
 * @throws std::invalid_argument - unsupported algorithm
 *         std::runtime_error - Failures generating the HMAC
 */
std::string PBKDF2_HMAC(Algorithm algorithm,
                        std::string_view pass,
                        std::string_view salt,
                        unsigned int iterationCount);

/**
 * Generate a digest by using the requested algorithm
 */
std::string digest(Algorithm algorithm, std::string_view data);

} // namespace cb::crypto


#ifndef QDM_LOGL_H
#define QDM_LOGL_H 1

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>

#include <qdm/tau.h>

int
qdm_find_tau(
    double *result,

    double v,
    size_t spline_df,
    const gsl_vector *knots,
    const gsl_vector *mmm
);

int
qdm_logl(
    double *log_likelihood,
    double *tau,

    double v,
    size_t spline_df,
    const gsl_vector *knots,
    const gsl_vector *mmm,

    double tau_low,
    double tau_high,
    double xi_low,
    double xi_high
);

int
qdm_logl_2(
    double *log_likelihood,
    double *tau,

    double x,
    double y,

    double tau_low,
    double tau_high,

    double xi_low,
    double xi_high,

    size_t spline_df,

    const gsl_matrix *theta,
    const gsl_vector *knots
);

void
qdm_logl_3(
    double *log_likelihood,
    double *tau,

    double x,
    double y,

    const qdm_tau *t,
    const gsl_vector *xi,

    const gsl_matrix *theta
);

#endif /* QDM_LOGL_H */