Set execution status of CUB device functions to error code (NVIDIA#4511)

Author: oleksandr-pavlyk

c/parallel/src/for.cu

@@ -147,8 +147,9 @@ CUresult cccl_device_for(
 
   try
   {
-    pushed = try_push_context();
-    Invoke(d_data, num_items, op, build.cc, (CUfunction) build.static_kernel, stream);
+    pushed           = try_push_context();
+    auto exec_status = Invoke(d_data, num_items, op, build.cc, (CUfunction) build.static_kernel, stream);
+    error            = static_cast<CUresult>(exec_status);
   }
   catch (...)
   {

c/parallel/src/merge_sort.cu

@@ -493,7 +493,7 @@ CUresult cccl_device_merge_sort(
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
 
-    cub::DispatchMergeSort<
+    auto exec_status = cub::DispatchMergeSort<
       indirect_arg_t,
       indirect_arg_t,
       indirect_arg_t,
@@ -517,6 +517,8 @@ CUresult cccl_device_merge_sort(
                                 {build},
                                 cub::detail::CudaDriverLauncherFactory{cu_device, build.cc},
                                 {d_out_keys.value_type.size});
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/radix_sort.cu

@@ -724,14 +724,15 @@ CUresult cccl_device_radix_sort_impl(
     cub::DoubleBuffer<indirect_arg_t> d_values_buffer(
       *static_cast<indirect_arg_t**>(&val_arg_in), *static_cast<indirect_arg_t**>(&val_arg_out));
 
-    cub::DispatchRadixSort<Order,
-                           indirect_arg_t,
-                           indirect_arg_t,
-                           OffsetT,
-                           indirect_arg_t,
-                           radix_sort::dynamic_radix_sort_policy_t<&radix_sort::get_policy>,
-                           radix_sort::radix_sort_kernel_source,
-                           cub::detail::CudaDriverLauncherFactory>::
+    auto exec_status = cub::DispatchRadixSort<
+      Order,
+      indirect_arg_t,
+      indirect_arg_t,
+      OffsetT,
+      indirect_arg_t,
+      radix_sort::dynamic_radix_sort_policy_t<&radix_sort::get_policy>,
+      radix_sort::radix_sort_kernel_source,
+      cub::detail::CudaDriverLauncherFactory>::
       Dispatch(
         d_temp_storage,
         *temp_storage_bytes,
@@ -748,6 +749,7 @@ CUresult cccl_device_radix_sort_impl(
         {d_keys_in.value_type.size});
 
     *selector = d_keys_buffer.selector;
+    error     = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/reduce.cu

@@ -375,16 +375,17 @@ CUresult cccl_device_reduce(
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
 
-    cub::DispatchReduce<indirect_arg_t, // InputIteratorT
-                        indirect_arg_t, // OutputIteratorT
-                        ::cuda::std::size_t, // OffsetT
-                        indirect_arg_t, // ReductionOpT
-                        indirect_arg_t, // InitT
-                        void, // AccumT
-                        ::cuda::std::__identity, // TransformOpT
-                        reduce::dynamic_reduce_policy_t<&reduce::get_policy>, // PolicyHub
-                        reduce::reduce_kernel_source, // KernelSource
-                        cub::detail::CudaDriverLauncherFactory>:: // KernelLauncherFactory
+    auto exec_status = cub::DispatchReduce<
+      indirect_arg_t, // InputIteratorT
+      indirect_arg_t, // OutputIteratorT
+      ::cuda::std::size_t, // OffsetT
+      indirect_arg_t, // ReductionOpT
+      indirect_arg_t, // InitT
+      void, // AccumT
+      ::cuda::std::__identity, // TransformOpT
+      reduce::dynamic_reduce_policy_t<&reduce::get_policy>, // PolicyHub
+      reduce::reduce_kernel_source, // KernelSource
+      cub::detail::CudaDriverLauncherFactory>:: // KernelLauncherFactory
       Dispatch(
         d_temp_storage,
         *temp_storage_bytes,
@@ -398,6 +399,8 @@ CUresult cccl_device_reduce(
         {build},
         cub::detail::CudaDriverLauncherFactory{cu_device, build.cc},
         {reduce::get_accumulator_type(op, d_in, init)});
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/scan.cu

@@ -368,7 +368,8 @@ CUresult cccl_device_scan(
 
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
-    auto cuda_error = cub::DispatchScan<
+
+    auto exec_status = cub::DispatchScan<
       indirect_arg_t,
       indirect_arg_t,
       indirect_arg_t,
@@ -391,11 +392,8 @@ CUresult cccl_device_scan(
         {build},
         cub::detail::CudaDriverLauncherFactory{cu_device, build.cc},
         {scan::get_accumulator_type(op, d_in, init)});
-    if (cuda_error != cudaSuccess)
-    {
-      const char* errorString = cudaGetErrorString(cuda_error); // Get the error string
-      std::cerr << "CUDA error: " << errorString << std::endl;
-    }
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/segmented_reduce.cu

@@ -404,7 +404,7 @@ CUresult cccl_device_segmented_reduce(
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
 
-    cub::DispatchSegmentedReduce<
+    auto exec_status = cub::DispatchSegmentedReduce<
       indirect_arg_t, // InputIteratorT
       indirect_arg_t, // OutputIteratorT
       indirect_arg_t, // BeginSegmentIteratorT
@@ -430,6 +430,8 @@ CUresult cccl_device_segmented_reduce(
         /* kernel_source */ {build},
         /* launcher_factory &*/ cub::detail::CudaDriverLauncherFactory{cu_device, build.cc},
         /* policy */ {segmented_reduce::get_accumulator_type(op, d_in, init)});
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/transform.cu

@@ -529,7 +529,8 @@ CUresult cccl_device_binary_transform(
 
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
-    auto cuda_error = cub::detail::transform::dispatch_t<
+
+    auto exec_status = cub::detail::transform::dispatch_t<
       cub::detail::transform::requires_stable_address::no, // TODO implement yes
       OffsetT,
       ::cuda::std::tuple<indirect_arg_t, indirect_arg_t>,
@@ -545,11 +546,8 @@ CUresult cccl_device_binary_transform(
                stream,
                {build},
                cub::detail::CudaDriverLauncherFactory{cu_device, build.cc});
-    if (cuda_error != cudaSuccess)
-    {
-      const char* errorString = cudaGetErrorString(cuda_error); // Get the error string
-      std::cerr << "CUDA error: " << errorString << std::endl;
-    }
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {

c/parallel/src/unique_by_key.cu

@@ -440,7 +440,7 @@ CUresult cccl_device_unique_by_key(
     CUdevice cu_device;
     check(cuCtxGetDevice(&cu_device));
 
-    cub::DispatchUniqueByKey<
+    auto exec_status = cub::DispatchUniqueByKey<
       indirect_arg_t,
       indirect_arg_t,
       indirect_arg_t,
@@ -466,6 +466,8 @@ CUresult cccl_device_unique_by_key(
                                 {build},
                                 cub::detail::CudaDriverLauncherFactory{cu_device, build.cc},
                                 {d_keys_in.value_type.size});
+
+    error = static_cast<CUresult>(exec_status);
   }
   catch (const std::exception& exc)
   {