From 24ac089454a4d766055b9327211e488c54a575a9 Mon Sep 17 00:00:00 2001
From: Nghia Truong <7416935+ttnghia@users.noreply.github.com>
Date: Mon, 11 Dec 2023 10:12:30 -0800
Subject: [PATCH] Remove `cuda::proclaim_return_type` from nested lambda
 (#14607)

This removes `cuda::proclaim_return_type` from a device lambda because that lambda is going to be nested inside another device lambda, which is in turn enclosed by `cuda::proclaim_return_type`.

This PR is to fix a compile issue that we encountered:
```
/usr/local/cuda/include/cuda/std/detail/libcxx/include/__functional/invoke.h(402): error:
calling a __device__ function("cudf::tdigest::detail::_NV_ANON_NAMESPACE::build_output_column(int,
  ::std::unique_ptr<   ::cudf::column,     ::std::default_delete<   ::cudf::column> >  &&,     ::std::unique_ptr<   ::cudf::column,     ::std::default_delete<   ::cudf::column> >  &&,     ::std::unique_ptr<   ::cudf::column,     ::std::default_delete<   ::cudf::column> >  &&,     ::std::unique_ptr<   ::cudf::column,     ::std::default_delete<   ::cudf::column> >  &&,     ::std::unique_ptr<   ::cudf::column,     ::std::default_delete<   ::cudf::column> >  &&, bool,  ::rmm::cuda_stream_view,  ::rmm::mr::device_memory_resource *)
::[lambda(int) (instance 2)]::operator ()(int) const") from a __host__ __device__ function("__invoke") is not allowed
```

Note: The issue is reproducible only in our build environment: ARM architecture, cuda 12 + rockylinux8.

Closes https://github.com/rapidsai/cudf/issues/14610.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Michael Schellenberger Costa (https://github.com/miscco)
  - Karthikeyan (https://github.com/karthikeyann)
  - https://github.com/nvdbaranec
  - Bradley Dice (https://github.com/bdice)

URL: https://github.com/rapidsai/cudf/pull/14607
---
 cpp/src/quantiles/tdigest/tdigest_aggregation.cu | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/cpp/src/quantiles/tdigest/tdigest_aggregation.cu b/cpp/src/quantiles/tdigest/tdigest_aggregation.cu
index 3ccef38715b..450996a43d2 100644
--- a/cpp/src/quantiles/tdigest/tdigest_aggregation.cu
+++ b/cpp/src/quantiles/tdigest/tdigest_aggregation.cu
@@ -587,15 +587,16 @@ std::unique_ptr<column> build_output_column(size_type num_rows,
   auto is_stub_weight = [weights = weights->view().begin<double>()] __device__(size_type i) {
     return weights[i] == 0;
   };
-  // whether or not this particular tdigest is a stub
-  auto is_stub_digest = cuda::proclaim_return_type<size_type>(
-    [offsets = offsets->view().begin<size_type>(), is_stub_weight] __device__(size_type i) {
-      return is_stub_weight(offsets[i]) ? 1 : 0;
-    });
+  // Whether or not this particular tdigest is a stub.
+  // This should not be wrapped in `proclaim_return_type` as it will be used inside another
+  // device lambda.
+  auto is_stub_digest = [offsets = offsets->view().begin<size_type>(), is_stub_weight] __device__(
+                          size_type i) { return is_stub_weight(offsets[i]) ? 1 : 0; };
 
   size_type const num_stubs = [&]() {
     if (!has_nulls) { return 0; }
-    auto iter = cudf::detail::make_counting_transform_iterator(0, is_stub_digest);
+    auto iter = cudf::detail::make_counting_transform_iterator(
+      0, cuda::proclaim_return_type<size_type>(is_stub_digest));
     return thrust::reduce(rmm::exec_policy(stream), iter, iter + num_rows);
   }();