/* Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * 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.
 *  * Neither the name of NVIDIA CORPORATION 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 ``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 OWNER 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 <cooperative_groups.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/random.h>

// for cuda::std::tuple
#include <cuda/std/tuple>

namespace cg = cooperative_groups;
#include <helper_cuda.h>

////////////////////////////////////////////////////////////////////////////////
// A structure of 2D points (structure of arrays).
////////////////////////////////////////////////////////////////////////////////
class Points
{
    float *m_x;
    float *m_y;

public:
    // Constructor.
    __host__ __device__ Points()
        : m_x(NULL)
        , m_y(NULL)
    {
    }

    // Constructor.
    __host__ __device__ Points(float *x, float *y)
        : m_x(x)
        , m_y(y)
    {
    }

    // Get a point.
    __host__ __device__ __forceinline__ float2 get_point(int idx) const { return make_float2(m_x[idx], m_y[idx]); }

    // Set a point.
    __host__ __device__ __forceinline__ void set_point(int idx, const float2 &p)
    {
        m_x[idx] = p.x;
        m_y[idx] = p.y;
    }

    // Set the pointers.
    __host__ __device__ __forceinline__ void set(float *x, float *y)
    {
        m_x = x;
        m_y = y;
    }
};

////////////////////////////////////////////////////////////////////////////////
// A 2D bounding box
////////////////////////////////////////////////////////////////////////////////
class Bounding_box
{
    // Extreme points of the bounding box.
    float2 m_p_min;
    float2 m_p_max;

public:
    // Constructor. Create a unit box.
    __host__ __device__ Bounding_box()
    {
        m_p_min = make_float2(0.0f, 0.0f);
        m_p_max = make_float2(1.0f, 1.0f);
    }

    // Compute the center of the bounding-box.
    __host__ __device__ void compute_center(float2 &center) const
    {
        center.x = 0.5f * (m_p_min.x + m_p_max.x);
        center.y = 0.5f * (m_p_min.y + m_p_max.y);
    }

    // The points of the box.
    __host__ __device__ __forceinline__ const float2 &get_max() const { return m_p_max; }

    __host__ __device__ __forceinline__ const float2 &get_min() const { return m_p_min; }

    // Does a box contain a point.
    __host__ __device__ bool contains(const float2 &p) const
    {
        return p.x >= m_p_min.x && p.x < m_p_max.x && p.y >= m_p_min.y && p.y < m_p_max.y;
    }

    // Define the bounding box.
    __host__ __device__ void set(float min_x, float min_y, float max_x, float max_y)
    {
        m_p_min.x = min_x;
        m_p_min.y = min_y;
        m_p_max.x = max_x;
        m_p_max.y = max_y;
    }
};

////////////////////////////////////////////////////////////////////////////////
// A node of a quadree.
////////////////////////////////////////////////////////////////////////////////
class Quadtree_node
{
    // The identifier of the node.
    int m_id;
    // The bounding box of the tree.
    Bounding_box m_bounding_box;
    // The range of points.
    int m_begin, m_end;

public:
    // Constructor.
    __host__ __device__ Quadtree_node()
        : m_id(0)
        , m_begin(0)
        , m_end(0)
    {
    }

    // The ID of a node at its level.
    __host__ __device__ int id() const { return m_id; }

    // The ID of a node at its level.
    __host__ __device__ void set_id(int new_id) { m_id = new_id; }

    // The bounding box.
    __host__ __device__ __forceinline__ const Bounding_box &bounding_box() const { return m_bounding_box; }

    // Set the bounding box.
    __host__ __device__ __forceinline__ void set_bounding_box(float min_x, float min_y, float max_x, float max_y)
    {
        m_bounding_box.set(min_x, min_y, max_x, max_y);
    }

    // The number of points in the tree.
    __host__ __device__ __forceinline__ int num_points() const { return m_end - m_begin; }

    // The range of points in the tree.
    __host__ __device__ __forceinline__ int points_begin() const { return m_begin; }

    __host__ __device__ __forceinline__ int points_end() const { return m_end; }

    // Define the range for that node.
    __host__ __device__ __forceinline__ void set_range(int begin, int end)
    {
        m_begin = begin;
        m_end   = end;
    }
};

////////////////////////////////////////////////////////////////////////////////
// Algorithm parameters.
////////////////////////////////////////////////////////////////////////////////
struct Parameters
{
    // Choose the right set of points to use as in/out.
    int point_selector;
    // The number of nodes at a given level (2^k for level k).
    int num_nodes_at_this_level;
    // The recursion depth.
    int depth;
    // The max value for depth.
    const int max_depth;
    // The minimum number of points in a node to stop recursion.
    const int min_points_per_node;

    // Constructor set to default values.
    __host__ __device__ Parameters(int max_depth, int min_points_per_node)
        : point_selector(0)
        , num_nodes_at_this_level(1)
        , depth(0)
        , max_depth(max_depth)
        , min_points_per_node(min_points_per_node)
    {
    }

    // Copy constructor. Changes the values for next iteration.
    __host__ __device__ Parameters(const Parameters &params, bool)
        : point_selector((params.point_selector + 1) % 2)
        , num_nodes_at_this_level(4 * params.num_nodes_at_this_level)
        , depth(params.depth + 1)
        , max_depth(params.max_depth)
        , min_points_per_node(params.min_points_per_node)
    {
    }
};

////////////////////////////////////////////////////////////////////////////////
// Build a quadtree on the GPU. Use CUDA Dynamic Parallelism.
//
// The algorithm works as follows. The host (CPU) launches one block of
// NUM_THREADS_PER_BLOCK threads. That block will do the following steps:
//
// 1- Check the number of points and its depth.
//
// We impose a maximum depth to the tree and a minimum number of points per
// node. If the maximum depth is exceeded or the minimum number of points is
// reached. The threads in the block exit.
//
// Before exiting, they perform a buffer swap if it is needed. Indeed, the
// algorithm uses two buffers to permute the points and make sure they are
// properly distributed in the quadtree. By design we want all points to be
// in the first buffer of points at the end of the algorithm. It is the reason
// why we may have to swap the buffer before leavin (if the points are in the
// 2nd buffer).
//
// 2- Count the number of points in each child.
//
// If the depth is not too high and the number of points is sufficient, the
// block has to dispatch the points into four geometrical buckets: Its
// children. For that purpose, we compute the center of the bounding box and
// count the number of points in each quadrant.
//
// The set of points is divided into sections. Each section is given to a
// warp of threads (32 threads). Warps use __ballot and __popc intrinsics
// to count the points. See the Programming Guide for more information about
// those functions.
//
// 3- Scan the warps' results to know the "global" numbers.
//
// Warps work independently from each other. At the end, each warp knows the
// number of points in its section. To know the numbers for the block, the
// block has to run a scan/reduce at the block level. It's a traditional
// approach. The implementation in that sample is not as optimized as what
// could be found in fast radix sorts, for example, but it relies on the same
// idea.
//
// 4- Move points.
//
// Now that the block knows how many points go in each of its 4 children, it
// remains to dispatch the points. It is straightforward.
//
// 5- Launch new blocks.
//
// The block launches four new blocks: One per children. Each of the four blocks
// will apply the same algorithm.
////////////////////////////////////////////////////////////////////////////////
template <int NUM_THREADS_PER_BLOCK>
__global__ void build_quadtree_kernel(Quadtree_node *nodes, Points *points, Parameters params)
{
    // Handle to thread block group
    cg::thread_block cta = cg::this_thread_block();
    // The number of warps in a block.
    const int NUM_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / warpSize;

    // Shared memory to store the number of points.
    extern __shared__ int smem[];

    // s_num_pts[4][NUM_WARPS_PER_BLOCK];
    // Addresses of shared memory.
    volatile int *s_num_pts[4];

    for (int i = 0; i < 4; ++i)
        s_num_pts[i] = (volatile int *)&smem[i * NUM_WARPS_PER_BLOCK];

    // Compute the coordinates of the threads in the block.
    const int warp_id = threadIdx.x / warpSize;
    const int lane_id = threadIdx.x % warpSize;

    // Mask for compaction.
    // Same as: asm( "mov.u32 %0, %%lanemask_lt;" : "=r"(lane_mask_lt) );
    int lane_mask_lt = (1 << lane_id) - 1;

    // The current node.
    Quadtree_node &node = nodes[blockIdx.x];

    // The number of points in the node.
    int num_points = node.num_points();

    float2 center;
    int    range_begin, range_end;
    int    warp_cnts[4] = {0, 0, 0, 0};
    //
    // 1- Check the number of points and its depth.
    //

    // Stop the recursion here. Make sure points[0] contains all the points.
    if (params.depth >= params.max_depth || num_points <= params.min_points_per_node) {
        if (params.point_selector == 1) {
            int it = node.points_begin(), end = node.points_end();

            for (it += threadIdx.x; it < end; it += NUM_THREADS_PER_BLOCK)
                if (it < end)
                    points[0].set_point(it, points[1].get_point(it));
        }

        return;
    }

    // Compute the center of the bounding box of the points.
    const Bounding_box &bbox = node.bounding_box();

    bbox.compute_center(center);

    // Find how many points to give to each warp.
    int num_points_per_warp = max(warpSize, (num_points + NUM_WARPS_PER_BLOCK - 1) / NUM_WARPS_PER_BLOCK);

    // Each warp of threads will compute the number of points to move to each
    // quadrant.
    range_begin = node.points_begin() + warp_id * num_points_per_warp;
    range_end   = min(range_begin + num_points_per_warp, node.points_end());

    //
    // 2- Count the number of points in each child.
    //

    // Input points.
    const Points &in_points = points[params.point_selector];

    cg::thread_block_tile<32> tile32 = cg::tiled_partition<32>(cta);
    // Compute the number of points.
    for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end); range_it += warpSize) {
        // Is it still an active thread?
        bool is_active = range_it < range_end;

        // Load the coordinates of the point.
        float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);

        // Count top-left points.
        int num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y >= center.y));
        warp_cnts[0] += tile32.shfl(num_pts, 0);

        // Count top-right points.
        num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y >= center.y));
        warp_cnts[1] += tile32.shfl(num_pts, 0);

        // Count bottom-left points.
        num_pts = __popc(tile32.ballot(is_active && p.x < center.x && p.y < center.y));
        warp_cnts[2] += tile32.shfl(num_pts, 0);

        // Count bottom-right points.
        num_pts = __popc(tile32.ballot(is_active && p.x >= center.x && p.y < center.y));
        warp_cnts[3] += tile32.shfl(num_pts, 0);
    }

    if (tile32.thread_rank() == 0) {
        s_num_pts[0][warp_id] = warp_cnts[0];
        s_num_pts[1][warp_id] = warp_cnts[1];
        s_num_pts[2][warp_id] = warp_cnts[2];
        s_num_pts[3][warp_id] = warp_cnts[3];
    }

    // Make sure warps have finished counting.
    cg::sync(cta);

    //
    // 3- Scan the warps' results to know the "global" numbers.
    //

    // First 4 warps scan the numbers of points per child (inclusive scan).
    if (warp_id < 4) {
        int num_pts = tile32.thread_rank() < NUM_WARPS_PER_BLOCK ? s_num_pts[warp_id][tile32.thread_rank()] : 0;
#pragma unroll

        for (int offset = 1; offset < NUM_WARPS_PER_BLOCK; offset *= 2) {
            int n = tile32.shfl_up(num_pts, offset);

            if (tile32.thread_rank() >= offset)
                num_pts += n;
        }

        if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)
            s_num_pts[warp_id][tile32.thread_rank()] = num_pts;
    }

    cg::sync(cta);

    // Compute global offsets.
    if (warp_id == 0) {
        int sum = s_num_pts[0][NUM_WARPS_PER_BLOCK - 1];

        for (int row = 1; row < 4; ++row) {
            int tmp = s_num_pts[row][NUM_WARPS_PER_BLOCK - 1];
            cg::sync(tile32);

            if (tile32.thread_rank() < NUM_WARPS_PER_BLOCK)
                s_num_pts[row][tile32.thread_rank()] += sum;

            cg::sync(tile32);
            sum += tmp;
        }
    }

    cg::sync(cta);

    // Make the scan exclusive.
    int val = 0;
    if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {
        val = threadIdx.x == 0 ? 0 : smem[threadIdx.x - 1];
        val += node.points_begin();
    }

    cg::sync(cta);

    if (threadIdx.x < 4 * NUM_WARPS_PER_BLOCK) {
        smem[threadIdx.x] = val;
    }

    cg::sync(cta);

    //
    // 4- Move points.
    //
    if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {
        // Output points.
        Points &out_points = points[(params.point_selector + 1) % 2];

        warp_cnts[0] = s_num_pts[0][warp_id];
        warp_cnts[1] = s_num_pts[1][warp_id];
        warp_cnts[2] = s_num_pts[2][warp_id];
        warp_cnts[3] = s_num_pts[3][warp_id];

        const Points &in_points = points[params.point_selector];
        // Reorder points.
        for (int range_it = range_begin + tile32.thread_rank(); tile32.any(range_it < range_end);
             range_it += warpSize) {
            // Is it still an active thread?
            bool is_active = range_it < range_end;

            // Load the coordinates of the point.
            float2 p = is_active ? in_points.get_point(range_it) : make_float2(0.0f, 0.0f);

            // Count top-left points.
            bool pred = is_active && p.x < center.x && p.y >= center.y;
            int  vote = tile32.ballot(pred);
            int  dest = warp_cnts[0] + __popc(vote & lane_mask_lt);

            if (pred)
                out_points.set_point(dest, p);

            warp_cnts[0] += tile32.shfl(__popc(vote), 0);

            // Count top-right points.
            pred = is_active && p.x >= center.x && p.y >= center.y;
            vote = tile32.ballot(pred);
            dest = warp_cnts[1] + __popc(vote & lane_mask_lt);

            if (pred)
                out_points.set_point(dest, p);

            warp_cnts[1] += tile32.shfl(__popc(vote), 0);

            // Count bottom-left points.
            pred = is_active && p.x < center.x && p.y < center.y;
            vote = tile32.ballot(pred);
            dest = warp_cnts[2] + __popc(vote & lane_mask_lt);

            if (pred)
                out_points.set_point(dest, p);

            warp_cnts[2] += tile32.shfl(__popc(vote), 0);

            // Count bottom-right points.
            pred = is_active && p.x >= center.x && p.y < center.y;
            vote = tile32.ballot(pred);
            dest = warp_cnts[3] + __popc(vote & lane_mask_lt);

            if (pred)
                out_points.set_point(dest, p);

            warp_cnts[3] += tile32.shfl(__popc(vote), 0);
        }
    }

    cg::sync(cta);

    if (tile32.thread_rank() == 0) {
        s_num_pts[0][warp_id] = warp_cnts[0];
        s_num_pts[1][warp_id] = warp_cnts[1];
        s_num_pts[2][warp_id] = warp_cnts[2];
        s_num_pts[3][warp_id] = warp_cnts[3];
    }

    cg::sync(cta);

    //
    // 5- Launch new blocks.
    //
    if (!(params.depth >= params.max_depth || num_points <= params.min_points_per_node)) {
        // The last thread launches new blocks.
        if (threadIdx.x == NUM_THREADS_PER_BLOCK - 1) {
            // The children.
            Quadtree_node *children = &nodes[params.num_nodes_at_this_level - (node.id() & ~3)];

            // The offsets of the children at their level.
            int child_offset = 4 * node.id();

            // Set IDs.
            children[child_offset + 0].set_id(4 * node.id() + 0);
            children[child_offset + 1].set_id(4 * node.id() + 1);
            children[child_offset + 2].set_id(4 * node.id() + 2);
            children[child_offset + 3].set_id(4 * node.id() + 3);

            const Bounding_box &bbox = node.bounding_box();
            // Points of the bounding-box.
            const float2 &p_min = bbox.get_min();
            const float2 &p_max = bbox.get_max();

            // Set the bounding boxes of the children.
            children[child_offset + 0].set_bounding_box(p_min.x, center.y, center.x,
                                                        p_max.y); // Top-left.
            children[child_offset + 1].set_bounding_box(center.x, center.y, p_max.x,
                                                        p_max.y); // Top-right.
            children[child_offset + 2].set_bounding_box(p_min.x, p_min.y, center.x,
                                                        center.y); // Bottom-left.
            children[child_offset + 3].set_bounding_box(center.x, p_min.y, p_max.x,
                                                        center.y); // Bottom-right.

            // Set the ranges of the children.

            children[child_offset + 0].set_range(node.points_begin(), s_num_pts[0][warp_id]);
            children[child_offset + 1].set_range(s_num_pts[0][warp_id], s_num_pts[1][warp_id]);
            children[child_offset + 2].set_range(s_num_pts[1][warp_id], s_num_pts[2][warp_id]);
            children[child_offset + 3].set_range(s_num_pts[2][warp_id], s_num_pts[3][warp_id]);

            // Launch 4 children.
            build_quadtree_kernel<NUM_THREADS_PER_BLOCK>
                <<<4, NUM_THREADS_PER_BLOCK, 4 * NUM_WARPS_PER_BLOCK * sizeof(int)>>>(
                    &children[child_offset], points, Parameters(params, true));
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Make sure a Quadtree is properly defined.
////////////////////////////////////////////////////////////////////////////////
bool check_quadtree(const Quadtree_node *nodes, int idx, int num_pts, Points *pts, Parameters params)
{
    const Quadtree_node &node       = nodes[idx];
    int                  num_points = node.num_points();

    if (!(params.depth == params.max_depth || num_points <= params.min_points_per_node)) {
        int num_points_in_children = 0;

        num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 0].num_points();
        num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 1].num_points();
        num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 2].num_points();
        num_points_in_children += nodes[params.num_nodes_at_this_level + 4 * idx + 3].num_points();

        if (num_points_in_children != node.num_points())
            return false;

        return check_quadtree(
                   &nodes[params.num_nodes_at_this_level], 4 * idx + 0, num_pts, pts, Parameters(params, true))
            && check_quadtree(
                   &nodes[params.num_nodes_at_this_level], 4 * idx + 1, num_pts, pts, Parameters(params, true))
            && check_quadtree(
                   &nodes[params.num_nodes_at_this_level], 4 * idx + 2, num_pts, pts, Parameters(params, true))
            && check_quadtree(
                   &nodes[params.num_nodes_at_this_level], 4 * idx + 3, num_pts, pts, Parameters(params, true));
    }

    const Bounding_box &bbox = node.bounding_box();

    for (int it = node.points_begin(); it < node.points_end(); ++it) {
        if (it >= num_pts)
            return false;

        float2 p = pts->get_point(it);

        if (!bbox.contains(p))
            return false;
    }

    return true;
}

////////////////////////////////////////////////////////////////////////////////
// Parallel random number generator.
////////////////////////////////////////////////////////////////////////////////
struct Random_generator
{
    int count;

    __host__ __device__ Random_generator()
        : count(0)
    {
    }
    __host__ __device__ unsigned int hash(unsigned int a)
    {
        a = (a + 0x7ed55d16) + (a << 12);
        a = (a ^ 0xc761c23c) ^ (a >> 19);
        a = (a + 0x165667b1) + (a << 5);
        a = (a + 0xd3a2646c) ^ (a << 9);
        a = (a + 0xfd7046c5) + (a << 3);
        a = (a ^ 0xb55a4f09) ^ (a >> 16);
        return a;
    }

    __host__ __device__ __forceinline__ cuda::std::tuple<float, float> operator()()
    {
#ifdef __CUDA_ARCH__
        unsigned seed = hash(blockIdx.x * blockDim.x + threadIdx.x + count);
        // thrust::generate may call operator() more than once per thread.
        // Hence, increment count by grid size to ensure uniqueness of seed
        count += blockDim.x * gridDim.x;
#else
        unsigned seed = hash(0);
#endif
        thrust::default_random_engine                    rng(seed);
        thrust::random::uniform_real_distribution<float> distrib;
        return cuda::std::make_tuple(distrib(rng), distrib(rng));
    }
};

////////////////////////////////////////////////////////////////////////////////
// Allocate GPU structs, launch kernel and clean up
////////////////////////////////////////////////////////////////////////////////
bool cdpQuadtree(int warp_size)
{
    // Constants to control the algorithm.
    const int num_points          = 1024;
    const int max_depth           = 8;
    const int min_points_per_node = 16;

    // Allocate memory for points.
    thrust::device_vector<float> x_d0(num_points);
    thrust::device_vector<float> x_d1(num_points);
    thrust::device_vector<float> y_d0(num_points);
    thrust::device_vector<float> y_d1(num_points);

    // Generate random points.
    Random_generator rnd;
    thrust::generate(
        thrust::make_zip_iterator(x_d0.begin(), y_d0.begin()), thrust::make_zip_iterator(x_d0.end(), y_d0.end()), rnd);

    // Host structures to analyze the device ones.
    Points points_init[2];
    points_init[0].set(thrust::raw_pointer_cast(&x_d0[0]), thrust::raw_pointer_cast(&y_d0[0]));
    points_init[1].set(thrust::raw_pointer_cast(&x_d1[0]), thrust::raw_pointer_cast(&y_d1[0]));

    // Allocate memory to store points.
    Points *points;
    checkCudaErrors(cudaMalloc((void **)&points, 2 * sizeof(Points)));
    checkCudaErrors(cudaMemcpy(points, points_init, 2 * sizeof(Points), cudaMemcpyHostToDevice));

    // We could use a close form...
    int max_nodes = 0;

    for (int i = 0, num_nodes_at_level = 1; i < max_depth; ++i, num_nodes_at_level *= 4)
        max_nodes += num_nodes_at_level;

    // Allocate memory to store the tree.
    Quadtree_node root;
    root.set_range(0, num_points);
    Quadtree_node *nodes;
    checkCudaErrors(cudaMalloc((void **)&nodes, max_nodes * sizeof(Quadtree_node)));
    checkCudaErrors(cudaMemcpy(nodes, &root, sizeof(Quadtree_node), cudaMemcpyHostToDevice));

    // Build the quadtree.
    Parameters params(max_depth, min_points_per_node);
    std::cout << "Launching CDP kernel to build the quadtree" << std::endl;
    const int    NUM_THREADS_PER_BLOCK = 128; // Do not use less than 128 threads.
    const int    NUM_WARPS_PER_BLOCK   = NUM_THREADS_PER_BLOCK / warp_size;
    const size_t smem_size             = 4 * NUM_WARPS_PER_BLOCK * sizeof(int);
    build_quadtree_kernel<NUM_THREADS_PER_BLOCK><<<1, NUM_THREADS_PER_BLOCK, smem_size>>>(nodes, points, params);
    checkCudaErrors(cudaGetLastError());

    // Copy points to CPU.
    thrust::host_vector<float> x_h(x_d0);
    thrust::host_vector<float> y_h(y_d0);
    Points                     host_points;
    host_points.set(thrust::raw_pointer_cast(&x_h[0]), thrust::raw_pointer_cast(&y_h[0]));

    // Copy nodes to CPU.
    Quadtree_node *host_nodes = new Quadtree_node[max_nodes];
    checkCudaErrors(cudaMemcpy(host_nodes, nodes, max_nodes * sizeof(Quadtree_node), cudaMemcpyDeviceToHost));

    // Validate the results.
    bool ok = check_quadtree(host_nodes, 0, num_points, &host_points, params);
    std::cout << "Results: " << (ok ? "OK" : "FAILED") << std::endl;

    // Free CPU memory.
    delete[] host_nodes;

    // Free memory.
    checkCudaErrors(cudaFree(nodes));
    checkCudaErrors(cudaFree(points));

    return ok;
}

////////////////////////////////////////////////////////////////////////////////
// Main entry point.
////////////////////////////////////////////////////////////////////////////////
int main(int argc, char **argv)
{
    // Find/set the device.
    // The test requires an architecture SM35 or greater (CDP capable).
    int            cuda_device = findCudaDevice(argc, (const char **)argv);
    cudaDeviceProp deviceProps;
    checkCudaErrors(cudaGetDeviceProperties(&deviceProps, cuda_device));
    int cdpCapable = (deviceProps.major == 3 && deviceProps.minor >= 5) || deviceProps.major >= 4;

    printf(
        "GPU device %s has compute capabilities (SM %d.%d)\n", deviceProps.name, deviceProps.major, deviceProps.minor);

    if (!cdpCapable) {
        std::cerr << "cdpQuadTree requires SM 3.5 or higher to use CUDA Dynamic "
                     "Parallelism.  Exiting...\n"
                  << std::endl;
        exit(EXIT_WAIVED);
    }

    bool ok = cdpQuadtree(deviceProps.warpSize);

    return (ok ? EXIT_SUCCESS : EXIT_FAILURE);
}
