log reduce function

src/common/device_helpers.hip.h

@@ -965,18 +965,32 @@ size_t SegmentedUniqueByKey(
 
 template <typename Policy, typename InputIt, typename Init, typename Func>
 auto Reduce(Policy policy, InputIt first, InputIt second, Init init, Func reduce_op) {
+  std::cerr << "Entering Reduce function" << std::endl;
   size_t constexpr kLimit = std::numeric_limits<int32_t>::max() / 2;
   size_t size = std::distance(first, second);
+  std::cerr << "Total size for reduction: " << size << std::endl;
+
   using Ty = std::remove_cv_t<Init>;
   Ty aggregate = init;
+
   for (size_t offset = 0; offset < size; offset += kLimit) {
     auto begin_it = first + offset;
     auto end_it = first + std::min(offset + kLimit, size);
     size_t batch_size = std::distance(begin_it, end_it);
     CHECK_LE(batch_size, size);
-    auto ret = thrust::reduce(policy, begin_it, end_it, init, reduce_op);
-    aggregate = reduce_op(aggregate, ret);
+
+    std::cerr << "Processing batch: offset=" << offset << ", batch_size=" << batch_size << std::endl;
+
+    try {
+      auto ret = thrust::reduce(policy, begin_it, end_it, init, reduce_op);
+      aggregate = reduce_op(aggregate, ret);
+    } catch (const std::exception& e) {
+      std::cerr << "Exception in thrust::reduce: " << e.what() << std::endl;
+      throw;
+    }
   }
+
+  std::cerr << "Exiting Reduce function" << std::endl;
   return aggregate;
 }