From af67ddfb49a144ea80ed2de385f0d35482946794 Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Mon, 11 Sep 2023 17:35:02 -0700
Subject: [PATCH] Example device function with cuda::proclaim_return_type.

---
 cpp/include/cuspatial/detail/point_quadtree.cuh | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/cpp/include/cuspatial/detail/point_quadtree.cuh b/cpp/include/cuspatial/detail/point_quadtree.cuh
index 3e70d1fe2..ebe9b02e6 100644
--- a/cpp/include/cuspatial/detail/point_quadtree.cuh
+++ b/cpp/include/cuspatial/detail/point_quadtree.cuh
@@ -30,6 +30,8 @@
 #include <thrust/iterator/zip_iterator.h>
 #include <thrust/tuple.h>
 
+#include <cuda/functional>
+
 #include <tuple>
 
 namespace cuspatial {
@@ -108,9 +110,9 @@ inline point_quadtree make_quad_tree(rmm::device_uvector<uint32_t>& keys,
                       offsets_iter + num_valid_nodes,
                       offsets.begin(),
                       // return is_internal_node ? lhs : rhs
-                      [] __device__(auto const& t) {
+                      cuda::proclaim_return_type<uint32_t>([] __device__(auto const& t) {
                         return thrust::get<0>(t) ? thrust::get<1>(t) : thrust::get<2>(t);
-                      });
+                      }));
 
     return std::move(offsets);
   }();
@@ -126,9 +128,9 @@ inline point_quadtree make_quad_tree(rmm::device_uvector<uint32_t>& keys,
                     lengths_iter + num_valid_nodes,
                     lengths.begin(),
                     // return bool ? lhs : rhs
-                    [] __device__(auto const& t) {
+                    cuda::proclaim_return_type<uint32_t>([] __device__(auto const& t) {
                       return thrust::get<0>(t) ? thrust::get<1>(t) : thrust::get<2>(t);
-                    });
+                    }));
 
   // Shrink keys to the number of valid nodes
   keys.resize(num_valid_nodes, stream);