Address review nits for DeviceSegmentedReduce env overloads

  - Add missing non-overlap precondition to env ArgMin and ArgMax docs
  - Reorder env tests: group all env tests before custom stream tests
  - Add not_guaranteed determinism test for Reduce env API

Author: gonidelis

cub/cub/device/device_segmented_reduce.cuh

@@ -1,5 +1,5 @@
 // SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
-// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved.
 // SPDX-License-Identifier: BSD-3
 
 //! @file
@@ -1284,6 +1284,11 @@ public:
   //!   ``segment_offsets`` (of length ``num_segments + 1``) can be aliased
   //!   for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where
   //!   the latter is specified as ``segment_offsets + 1``).
+  //! - Let ``s`` be in ``[0, num_segments)``. The range
+  //!   ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not
+  //!   overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``,
+  //!   ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor
+  //!   ``[d_end_offsets, d_end_offsets + num_segments)``.
   //! - Can use a specific stream or cuda memory resource through the ``env`` parameter
   //!
   //! Snippet
@@ -1989,6 +1994,11 @@ public:
   //!   ``segment_offsets`` (of length ``num_segments + 1``) can be aliased
   //!   for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where
   //!   the latter is specified as ``segment_offsets + 1``).
+  //! - Let ``s`` be in ``[0, num_segments)``. The range
+  //!   ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not
+  //!   overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``,
+  //!   ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor
+  //!   ``[d_end_offsets, d_end_offsets + num_segments)``.
   //! - Can use a specific stream or cuda memory resource through the ``env`` parameter
   //!
   //! Snippet

cub/test/catch2_test_device_segmented_reduce_env.cu

@@ -238,6 +238,60 @@ C2H_TEST("Device segmented max uses environment", "[segmented_reduce][device]")
   REQUIRE(d_out == expected);
 }
 
+C2H_TEST("Device segmented argmin uses environment", "[segmented_reduce][device]")
+{
+  int num_segments                     = 3;
+  thrust::device_vector<int> d_offsets = {0, 4, 7, 9};
+  auto d_offsets_it                    = thrust::raw_pointer_cast(d_offsets.data());
+  thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9, 1, 2};
+  thrust::device_vector<cub::KeyValuePair<int, int>> d_out(3);
+
+  size_t expected_bytes_allocated{};
+  REQUIRE(
+    cudaSuccess
+    == cub::DeviceSegmentedReduce::ArgMin(
+      nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1));
+
+  auto env = stdexec::env{expected_allocation_size(expected_bytes_allocated)};
+
+  device_segmented_reduce_argmin(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env);
+
+  thrust::host_vector<cub::KeyValuePair<int, int>> h_out(d_out);
+  REQUIRE(h_out[0].key == 3);
+  REQUIRE(h_out[0].value == 5);
+  REQUIRE(h_out[1].key == 1);
+  REQUIRE(h_out[1].value == 0);
+  REQUIRE(h_out[2].key == 0);
+  REQUIRE(h_out[2].value == 1);
+}
+
+C2H_TEST("Device segmented argmax uses environment", "[segmented_reduce][device]")
+{
+  int num_segments                     = 3;
+  thrust::device_vector<int> d_offsets = {0, 4, 7, 9};
+  auto d_offsets_it                    = thrust::raw_pointer_cast(d_offsets.data());
+  thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9, 1, 2};
+  thrust::device_vector<cub::KeyValuePair<int, int>> d_out(3);
+
+  size_t expected_bytes_allocated{};
+  REQUIRE(
+    cudaSuccess
+    == cub::DeviceSegmentedReduce::ArgMax(
+      nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1));
+
+  auto env = stdexec::env{expected_allocation_size(expected_bytes_allocated)};
+
+  device_segmented_reduce_argmax(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env);
+
+  thrust::host_vector<cub::KeyValuePair<int, int>> h_out(d_out);
+  REQUIRE(h_out[0].key == 0);
+  REQUIRE(h_out[0].value == 8);
+  REQUIRE(h_out[1].key == 2);
+  REQUIRE(h_out[1].value == 9);
+  REQUIRE(h_out[2].key == 1);
+  REQUIRE(h_out[2].value == 2);
+}
+
 TEST_CASE("Device segmented reduce uses custom stream", "[segmented_reduce][device]")
 {
   int num_segments                     = 3;
@@ -367,60 +421,6 @@ TEST_CASE("Device segmented max uses custom stream", "[segmented_reduce][device]
   REQUIRE(cudaSuccess == cudaStreamDestroy(custom_stream));
 }
 
-C2H_TEST("Device segmented argmin uses environment", "[segmented_reduce][device]")
-{
-  int num_segments                     = 3;
-  thrust::device_vector<int> d_offsets = {0, 4, 7, 9};
-  auto d_offsets_it                    = thrust::raw_pointer_cast(d_offsets.data());
-  thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9, 1, 2};
-  thrust::device_vector<cub::KeyValuePair<int, int>> d_out(3);
-
-  size_t expected_bytes_allocated{};
-  REQUIRE(
-    cudaSuccess
-    == cub::DeviceSegmentedReduce::ArgMin(
-      nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1));
-
-  auto env = stdexec::env{expected_allocation_size(expected_bytes_allocated)};
-
-  device_segmented_reduce_argmin(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env);
-
-  thrust::host_vector<cub::KeyValuePair<int, int>> h_out(d_out);
-  REQUIRE(h_out[0].key == 3);
-  REQUIRE(h_out[0].value == 5);
-  REQUIRE(h_out[1].key == 1);
-  REQUIRE(h_out[1].value == 0);
-  REQUIRE(h_out[2].key == 0);
-  REQUIRE(h_out[2].value == 1);
-}
-
-C2H_TEST("Device segmented argmax uses environment", "[segmented_reduce][device]")
-{
-  int num_segments                     = 3;
-  thrust::device_vector<int> d_offsets = {0, 4, 7, 9};
-  auto d_offsets_it                    = thrust::raw_pointer_cast(d_offsets.data());
-  thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9, 1, 2};
-  thrust::device_vector<cub::KeyValuePair<int, int>> d_out(3);
-
-  size_t expected_bytes_allocated{};
-  REQUIRE(
-    cudaSuccess
-    == cub::DeviceSegmentedReduce::ArgMax(
-      nullptr, expected_bytes_allocated, d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1));
-
-  auto env = stdexec::env{expected_allocation_size(expected_bytes_allocated)};
-
-  device_segmented_reduce_argmax(d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, env);
-
-  thrust::host_vector<cub::KeyValuePair<int, int>> h_out(d_out);
-  REQUIRE(h_out[0].key == 0);
-  REQUIRE(h_out[0].value == 8);
-  REQUIRE(h_out[1].key == 2);
-  REQUIRE(h_out[1].value == 9);
-  REQUIRE(h_out[2].key == 1);
-  REQUIRE(h_out[2].value == 2);
-}
-
 TEST_CASE("Device segmented argmin uses custom stream", "[segmented_reduce][device]")
 {
   int num_segments                     = 3;

cub/test/catch2_test_device_segmented_reduce_env_api.cu

@@ -291,3 +291,27 @@ C2H_TEST("cub::DeviceSegmentedReduce::Reduce accepts run_to_run determinism requ
   REQUIRE(d_out == expected);
   REQUIRE(error == cudaSuccess);
 }
+
+C2H_TEST("cub::DeviceSegmentedReduce::Reduce accepts not_guaranteed determinism requirements",
+         "[segmented_reduce][env]")
+{
+  int num_segments                     = 3;
+  thrust::device_vector<int> d_offsets = {0, 3, 3, 7};
+  auto d_offsets_it                    = thrust::raw_pointer_cast(d_offsets.data());
+  thrust::device_vector<int> d_in{8, 6, 7, 5, 3, 0, 9};
+  thrust::device_vector<int> d_out(3);
+
+  auto env = cuda::execution::require(cuda::execution::determinism::not_guaranteed);
+
+  auto error = cub::DeviceSegmentedReduce::Reduce(
+    d_in.begin(), d_out.begin(), num_segments, d_offsets_it, d_offsets_it + 1, ::cuda::std::plus<>{}, 0, env);
+  thrust::device_vector<int> expected{21, 0, 17};
+
+  if (error != cudaSuccess)
+  {
+    std::cerr << "cub::DeviceSegmentedReduce::Reduce failed with status: " << error << std::endl;
+  }
+
+  REQUIRE(d_out == expected);
+  REQUIRE(error == cudaSuccess);
+}