/*
 *  Copyright 2008-2013 NVIDIA Corporation
 *  Modifications Copyright© 2019-2025 Advanced Micro Devices, Inc. All rights reserved.
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include <thrust/detail/config.h>

#include <thrust/async/copy.h>
#include <thrust/async/reduce.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>

#include "test_real_assertions.hpp"
#include "test_param_fixtures.hpp"
#include "test_utils.hpp"

TESTS_DEFINE(AsyncReduceTests, NumericalTestsParams);

template <typename T>
struct custom_plus
{
  __host__ __device__ T operator()(T lhs, T rhs) const
  {
    return lhs + rhs;
  }
};

#define DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(NAME, MEMBERS, CTOR, DTOR, VALIDATE, ...) \
  template <typename T>                                                                \
  struct NAME                                                                          \
  {                                                                                    \
    MEMBERS                                                                            \
                                                                                       \
    NAME()                                                                             \
    {                                                                                  \
      CTOR                                                                             \
    }                                                                                  \
                                                                                       \
    ~NAME()                                                                            \
    {                                                                                  \
      DTOR                                                                             \
    }                                                                                  \
                                                                                       \
    template <typename Event>                                                          \
    void validate_event(Event& e)                                                      \
    {                                                                                  \
      THRUST_UNUSED_VAR(e);                                                            \
      VALIDATE                                                                         \
    }                                                                                  \
                                                                                       \
    template <typename ForwardIt, typename Sentinel>                                   \
    __host__ auto operator()(ForwardIt&& first, Sentinel&& last)                       \
      THRUST_DECLTYPE_RETURNS(::thrust::async::reduce(__VA_ARGS__))                    \
  };                                                                                   \
  /**/

#define DEFINE_ASYNC_REDUCE_INVOKER(NAME, ...)                                                     \
  DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(                                                            \
    NAME, THRUST_PP_EMPTY(), THRUST_PP_EMPTY(), THRUST_PP_EMPTY(), THRUST_PP_EMPTY(), __VA_ARGS__) \
  /**/

#define DEFINE_SYNC_REDUCE_INVOKER(NAME, ...)                    \
  template <typename T>                                          \
  struct NAME                                                    \
  {                                                              \
    template <typename ForwardIt, typename Sentinel>             \
    __host__ auto operator()(ForwardIt&& first, Sentinel&& last) \
      THRUST_DECLTYPE_RETURNS(::thrust::reduce(__VA_ARGS__))     \
  };                                                             \
  /**/

DEFINE_ASYNC_REDUCE_INVOKER(reduce_async_invoker, THRUST_FWD(first), THRUST_FWD(last));
DEFINE_ASYNC_REDUCE_INVOKER(reduce_async_invoker_device, thrust::device, THRUST_FWD(first), THRUST_FWD(last));
DEFINE_ASYNC_REDUCE_INVOKER(reduce_async_invoker_device_allocator,
                            thrust::device(thrust::device_allocator<void>{}),
                            THRUST_FWD(first),
                            THRUST_FWD(last));
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_on
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device.on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last));
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_on
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device(thrust::device_allocator<void>{}).on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last));

DEFINE_SYNC_REDUCE_INVOKER(reduce_sync_invoker, THRUST_FWD(first), THRUST_FWD(last));

DEFINE_ASYNC_REDUCE_INVOKER(reduce_async_invoker_init, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_init, thrust::device, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_init,
  thrust::device(thrust::device_allocator<void>{}),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_on_init
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device.on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_on_init
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device(thrust::device_allocator<void>{}).on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>());

DEFINE_SYNC_REDUCE_INVOKER(reduce_sync_invoker_init, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>());

DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_init_plus, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>(), thrust::plus<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_init_plus,
  thrust::device,
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  thrust::plus<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_init_plus,
  thrust::device(thrust::device_allocator<void>{}),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  thrust::plus<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_on_init_plus
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device.on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  thrust::plus<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_on_init_plus
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device(thrust::device_allocator<void>{}).on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  thrust::plus<T>());

DEFINE_SYNC_REDUCE_INVOKER(
  reduce_sync_invoker_init_plus, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>(), thrust::plus<T>());

DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_init_custom_plus, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>(), custom_plus<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_init_custom_plus,
  thrust::device,
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  custom_plus<T>());
DEFINE_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_init_custom_plus,
  thrust::device(thrust::device_allocator<void>{}),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  custom_plus<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_on_init_custom_plus
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device.on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  custom_plus<T>());
DEFINE_STATEFUL_ASYNC_REDUCE_INVOKER(
  reduce_async_invoker_device_allocator_on_init_custom_plus
  // Members.
  ,
  hipStream_t stream_;
  // Constructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream_, hipStreamNonBlocking));
  // Destructor.
  ,
  thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream_));
  // `validate_event` member.
  ,
  ASSERT_EQ_QUIET(stream_, e.stream().native_handle());
  // Arguments to `thrust::async::reduce`.
  ,
  thrust::device(thrust::device_allocator<void>{}).on(stream_),
  THRUST_FWD(first),
  THRUST_FWD(last),
  random_integer<T>(),
  custom_plus<T>());

DEFINE_SYNC_REDUCE_INVOKER(
  reduce_sync_invoker_init_custom_plus, THRUST_FWD(first), THRUST_FWD(last), random_integer<T>(), custom_plus<T>());

///////////////////////////////////////////////////////////////////////////////

template <typename T, template <typename> class AsyncReduceInvoker, template <typename> class SyncReduceInvoker>
void testAsyncReduce()
{
  for (auto size : get_sizes())
  {
    SCOPED_TRACE(testing::Message() << "with size = " << size);
    for (auto seed : get_seeds())
    {
      SCOPED_TRACE(testing::Message() << "with seed= " << seed);

      thrust::host_vector<T> h0 = get_random_data<T>(size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

      thrust::device_vector<T> d0a(h0);
      thrust::device_vector<T> d0b(h0);
      thrust::device_vector<T> d0c(h0);
      thrust::device_vector<T> d0d(h0);

      AsyncReduceInvoker<T> invoke_async;
      SyncReduceInvoker<T> invoke_sync;

      ASSERT_EQ(h0, d0a);
      ASSERT_EQ(h0, d0b);
      ASSERT_EQ(h0, d0c);
      ASSERT_EQ(h0, d0d);

      auto f0a = invoke_async(d0a.begin(), d0a.end());
      auto f0b = invoke_async(d0b.begin(), d0b.end());
      auto f0c = invoke_async(d0c.begin(), d0c.end());
      auto f0d = invoke_async(d0d.begin(), d0d.end());

      invoke_async.validate_event(f0a);
      invoke_async.validate_event(f0b);
      invoke_async.validate_event(f0c);
      invoke_async.validate_event(f0d);

      // This potentially runs concurrently with the copies.
      auto const r0 = invoke_sync(h0.begin(), h0.end());

      T r1a;
      T r1b;
      T r1c;
      T r1d;

      test_future_value_retrieval(f0a, r1a);
      test_future_value_retrieval(f0b, r1b);
      test_future_value_retrieval(f0c, r1c);
      test_future_value_retrieval(f0d, r1d);

      auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

      ASSERT_NEAR(r0, r1a, tolerance);
      ASSERT_NEAR(r0, r1b, tolerance);
      ASSERT_NEAR(r0, r1c, tolerance);
      ASSERT_NEAR(r0, r1d, tolerance);
    }
  }
};

TYPED_TEST(AsyncReduceTests, TestAsyncReduce)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicy)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocator)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_on, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyOn)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_on, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorOn)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_on, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyOnInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_on_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorOnInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_on_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyOnInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_on_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorOnInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_on_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyOnInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_on_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyAllocatorOnInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduce<T, reduce_async_invoker_device_allocator_on_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

///////////////////////////////////////////////////////////////////////////////

template <typename T, template <typename> class AsyncReduceInvoker, template <typename> class SyncReduceInvoker>
void testAsyncReduceCountingIterator()
{
  constexpr std::size_t n = 15 * sizeof(T);

  ASSERT_LE(T(n), truncate_to_max_representable<T>(n));

  thrust::counting_iterator<T> first(0);
  thrust::counting_iterator<T> last(n);

  AsyncReduceInvoker<T> invoke_async;
  SyncReduceInvoker<T> invoke_sync;

  auto f0a = invoke_async(first, last);
  auto f0b = invoke_async(first, last);
  auto f0c = invoke_async(first, last);
  auto f0d = invoke_async(first, last);

  invoke_async.validate_event(f0a);
  invoke_async.validate_event(f0b);
  invoke_async.validate_event(f0c);
  invoke_async.validate_event(f0d);

  // This potentially runs concurrently with the copies.
  auto const r0 = invoke_sync(first, last);

  T r1a;
  T r1b;
  T r1c;
  T r1d;

  test_future_value_retrieval(f0a, r1a);
  test_future_value_retrieval(f0b, r1b);
  test_future_value_retrieval(f0c, r1c);
  test_future_value_retrieval(f0d, r1d);

  auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

  ASSERT_NEAR(r0, r1a, tolerance);
  ASSERT_NEAR(r0, r1b, tolerance);
  ASSERT_NEAR(r0, r1c, tolerance);
  ASSERT_NEAR(r0, r1d, tolerance);
};

TYPED_TEST(AsyncReduceTests, TestAsyncReduceCountingIterator)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyCountingIterator)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_device, reduce_sync_invoker>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceCountingIteratorInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyCountingIteratorInit)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_device_init, reduce_sync_invoker_init>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceCountingIteratorInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyCountingIteratorInitPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_device_init_plus, reduce_sync_invoker_init_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReduceCountingIteratorInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

TYPED_TEST(AsyncReduceTests, TestAsyncReducePolicyCountingIteratorInitCustomPlus)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  testAsyncReduceCountingIterator<T, reduce_async_invoker_device_init_custom_plus, reduce_sync_invoker_init_custom_plus>();
}

///////////////////////////////////////////////////////////////////////////////

TYPED_TEST(AsyncReduceTests, TestAsyncReduceUsing)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  for (auto size : get_sizes())
  {
    SCOPED_TRACE(testing::Message() << "with size = " << size);
    for (auto seed : get_seeds())
    {
      SCOPED_TRACE(testing::Message() << "with seed= " << seed);

      thrust::host_vector<T> h0 = get_random_data<T>(size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

      thrust::device_vector<T> d0a(h0);
      thrust::device_vector<T> d0b(h0);

      ASSERT_EQ(h0, d0a);
      ASSERT_EQ(h0, d0b);

      thrust::device_future<T> f0a;
      thrust::device_future<T> f0b;

      // When you import the customization points into the global namespace,
      // they should be selected instead of the synchronous algorithms.
      {
        using namespace thrust::async;
        f0a = reduce(d0a.begin(), d0a.end());
      }
      {
        using thrust::async::reduce;
        f0b = reduce(d0b.begin(), d0b.end());
      }

      // ADL should find the synchronous algorithms.
      // This potentially runs concurrently with the copies.
      T const r0 = reduce(h0.begin(), h0.end());

      T r1a;
      T r1b;

      test_future_value_retrieval(f0a, r1a);
      test_future_value_retrieval(f0b, r1b);

      auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

      ASSERT_NEAR(r0, r1a, tolerance);
      ASSERT_NEAR(r0, r1b, tolerance);
    }
  }
};

///////////////////////////////////////////////////////////////////////////////

TYPED_TEST(AsyncReduceTests, TestAsyncReduceAfter)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  for (auto size : get_sizes())
  {
    SCOPED_TRACE(testing::Message() << "with size = " << size);
    for (auto seed : get_seeds())
    {
      SCOPED_TRACE(testing::Message() << "with seed= " << seed);

      thrust::host_vector<T> h0 = get_random_data<T>(size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

      thrust::device_vector<T> d0(h0);

      ASSERT_EQ(h0, d0);

      auto f0 = thrust::async::reduce(d0.begin(), d0.end());

      ASSERT_EQ(true, f0.valid_stream());

      auto const f0_stream = f0.stream().native_handle();

      auto f1 = thrust::async::reduce(thrust::device.after(f0), d0.begin(), d0.end());

      // Verify that double consumption of a future produces an exception.
      ASSERT_THROW(auto x = thrust::async::reduce(thrust::device.after(f0), d0.begin(), d0.end());
                   THRUST_UNUSED_VAR(x), thrust::event_error);

      ASSERT_EQ_QUIET(f0_stream, f1.stream().native_handle());

      auto after_policy2 = thrust::device.after(f1);

      auto f2 = thrust::async::reduce(after_policy2, d0.begin(), d0.end());

      // Verify that double consumption of a policy produces an exception.
      ASSERT_THROW(auto x = thrust::async::reduce(after_policy2, d0.begin(), d0.end());
                   THRUST_UNUSED_VAR(x), thrust::event_error);

      ASSERT_EQ_QUIET(f0_stream, f2.stream().native_handle());

      // This potentially runs concurrently with the copies.
      T const r0 = thrust::reduce(h0.begin(), h0.end());

      T r1;

      test_future_value_retrieval(f2, r1);

      auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

      ASSERT_NEAR(r0, r1, tolerance);
    }
  }
};

///////////////////////////////////////////////////////////////////////////////

TYPED_TEST(AsyncReduceTests, TestAsyncReduceOnThenAfter)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  for (auto size : get_sizes())
  {
    SCOPED_TRACE(testing::Message() << "with size = " << size);
    for (auto seed : get_seeds())
    {
      SCOPED_TRACE(testing::Message() << "with seed= " << seed);

      thrust::host_vector<T> h0 = get_random_data<T>(size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

      thrust::device_vector<T> d0(h0);

      ASSERT_EQ(h0, d0);

      hipStream_t stream;
      thrust::hip_rocprim::throw_on_error(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));

      auto f0 = thrust::async::reduce(thrust::device.on(stream), d0.begin(), d0.end());

      ASSERT_EQ_QUIET(stream, f0.stream().native_handle());

      auto f1 = thrust::async::reduce(thrust::device.after(f0), d0.begin(), d0.end());

      // Verify that double consumption of a future produces an exception.
      ASSERT_THROW(auto x = thrust::async::reduce(thrust::device.after(f0), d0.begin(), d0.end());
                   THRUST_UNUSED_VAR(x), thrust::event_error);

      ASSERT_EQ_QUIET(stream, f1.stream().native_handle());

      auto after_policy2 = thrust::device.after(f1);

      auto f2 = thrust::async::reduce(after_policy2, d0.begin(), d0.end());

      // Verify that double consumption of a policy produces an exception.
      ASSERT_THROW(auto x = thrust::async::reduce(after_policy2, d0.begin(), d0.end());
                   THRUST_UNUSED_VAR(x), thrust::event_error);

      ASSERT_EQ_QUIET(stream, f2.stream().native_handle());

      // This potentially runs concurrently with the copies.
      T const r0 = thrust::reduce(h0.begin(), h0.end());

      T r1;

      test_future_value_retrieval(f2, r1);

      auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

      ASSERT_NEAR(r0, r1, tolerance);

      thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream));
    }
  }
};

///////////////////////////////////////////////////////////////////////////////

/*
TYPED_TEST(AsyncReduceTests, TestAsyncReduceAllocatorOnThenAfter)
{
    SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
    using T = typename TestFixture::input_type;
    for(auto size : get_sizes())
    {
        SCOPED_TRACE(testing::Message() << "with size = " << size);
        for(size_t seed_index = 0; seed_index < random_seeds_count + seed_size; seed_index++)
        {
            unsigned int seed
                = seed_index < random_seeds_count ? rand() : seeds[seed_index - random_seeds_count];
            SCOPED_TRACE(testing::Message() << "with seed= " << seed);

            thrust::host_vector<T> h0 = get_random_data<T>(
                size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

            thrust::device_vector<T> d0(h0);

            ASSERT_EQ(h0, d0);

            hipStream_t stream0;
            thrust::hip_rocprim::throw_on_error(
                hipStreamCreateWithFlags(&stream0, hipStreamNonBlocking)
            );

            hipStream_t stream1;
            thrust::hip_rocprim::throw_on_error(
                hipStreamCreateWithFlags(&stream1, hipStreamNonBlocking)
            );

            auto f0 = thrust::async::reduce(
              thrust::device(thrust::device_allocator<void>{}).on(stream0)
            , d0.begin(), d0.end()
            );

            ASSERT_EQ_QUIET(stream0, f0.stream().native_handle());

            auto f1 = thrust::async::reduce(
              thrust::device(thrust::device_allocator<void>{}).after(f0)
            , d0.begin(), d0.end()
            );

            ASSERT_THROW(
              auto x = thrust::async::reduce(
                  thrust::device(thrust::device_allocator<void>{}).after(f0)
              , d0.begin(), d0.end()
              );
              THRUST_UNUSED_VAR(x)
            , thrust::event_error
            );

            ASSERT_EQ_QUIET(stream0, f1.stream().native_handle());

            auto f2 = thrust::async::reduce(
              thrust::device(thrust::device_allocator<void>{}).on(stream1).after(f1)
            , d0.begin(), d0.end()
            );

            ASSERT_THROW(
              auto x = thrust::async::reduce(
                  thrust::device(thrust::device_allocator<void>{}).on(stream1).after(f1)
              , d0.begin(), d0.end()
              );
              THRUST_UNUSED_VAR(x)
            , thrust::event_error
            );

            KNOWN_FAILURE;
            // FIXME: The below fails because you can't combine allocator attachment,
            // `.on`, and `.after`.
            ASSERT_EQ_QUIET(stream1, f2.stream().native_handle());

            // This potentially runs concurrently with the copies.
            T const r0 = thrust::reduce(h0.begin(), h0.end());

            T const r1 = TEST_FUTURE_VALUE_RETRIEVAL(f2);

            ASSERT_EQ(r0, r1);

            thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream0));
            thrust::hip_rocprim::throw_on_error(hipStreamDestroy(stream1));
        }
    }
}

///////////////////////////////////////////////////////////////////////////////

TYPED_TEST(AsyncReduceTests, TestAsyncReduceCaching)
{
    SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
    using T = typename TestFixture::input_type;
    for(auto size : get_sizes())
    {
        SCOPED_TRACE(testing::Message() << "with size = " << size);
        for(size_t seed_index = 0; seed_index < random_seeds_count + seed_size; seed_index++)
        {
            unsigned int seed
                = seed_index < random_seeds_count ? rand() : seeds[seed_index - random_seeds_count];
            SCOPED_TRACE(testing::Message() << "with seed= " << seed);

            thrust::host_vector<T> h0 = get_random_data<T>(
                size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

            constexpr std::int64_t m = 32;
            thrust::device_vector<T> d0(h0);

            ASSERT_EQ(h0, d0);

            T const* f0_raw_data;

            {
            // Perform one reduction to ensure there's an entry in the caching
            // allocator.
            auto f0 = thrust::async::reduce(d0.begin(), d0.end());

            TEST_EVENT_WAIT(f0);

            f0_raw_data = f0.raw_data();
            }

            for (std::int64_t i = 0; i < m; ++i)
            {
                auto f1 = thrust::async::reduce(d0.begin(), d0.end());

                ASSERT_EQ(true, f1.valid_stream());
                ASSERT_EQ(true, f1.valid_content());

                ASSERT_EQ_QUIET(f0_raw_data, f1.raw_data());

                // This potentially runs concurrently with the copies.
                T const r0 = thrust::reduce(h0.begin(), h0.end());

                T const r1 = TEST_FUTURE_VALUE_RETRIEVAL(f1);

                ASSERT_EQ(r0, r1);
            }
        }
    }
};
*/

///////////////////////////////////////////////////////////////////////////////

TYPED_TEST(AsyncReduceTests, TestAsyncCopyThenReduce)
{
  SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
  using T = typename TestFixture::input_type;
  for (auto size : get_sizes())
  {
    SCOPED_TRACE(testing::Message() << "with size = " << size);
    for (auto seed : get_seeds())
    {
      SCOPED_TRACE(testing::Message() << "with seed= " << seed);

      thrust::host_vector<T> h0 = get_random_data<T>(size, saturate_cast<T>(-1000), saturate_cast<T>(1000), seed);

      thrust::device_vector<T> d0a(h0);
      thrust::device_vector<T> d0b(h0);
      thrust::device_vector<T> d0c(h0);
      thrust::device_vector<T> d0d(h0);

      auto f0a = thrust::async::copy(h0.begin(), h0.end(), d0a.begin());
      auto f0b = thrust::async::copy(h0.begin(), h0.end(), d0b.begin());
      auto f0c = thrust::async::copy(h0.begin(), h0.end(), d0c.begin());
      auto f0d = thrust::async::copy(h0.begin(), h0.end(), d0d.begin());

      ASSERT_EQ(true, f0a.valid_stream());
      ASSERT_EQ(true, f0b.valid_stream());
      ASSERT_EQ(true, f0c.valid_stream());
      ASSERT_EQ(true, f0d.valid_stream());

      auto const f0a_stream = f0a.stream().native_handle();
      auto const f0b_stream = f0b.stream().native_handle();
      auto const f0c_stream = f0c.stream().native_handle();
      auto const f0d_stream = f0d.stream().native_handle();

      auto f1a = thrust::async::reduce(thrust::device.after(f0a), d0a.begin(), d0a.end());
      auto f1b = thrust::async::reduce(thrust::device.after(f0b), d0b.begin(), d0b.end());
      auto f1c = thrust::async::reduce(thrust::device.after(f0c), d0c.begin(), d0c.end());
      auto f1d = thrust::async::reduce(thrust::device.after(f0d), d0d.begin(), d0d.end());

      ASSERT_EQ(false, f0a.valid_stream());
      ASSERT_EQ(false, f0b.valid_stream());
      ASSERT_EQ(false, f0c.valid_stream());
      ASSERT_EQ(false, f0d.valid_stream());

      ASSERT_EQ(true, f1a.valid_stream());
      ASSERT_EQ(true, f1a.valid_content());
      ASSERT_EQ(true, f1b.valid_stream());
      ASSERT_EQ(true, f1b.valid_content());
      ASSERT_EQ(true, f1c.valid_stream());
      ASSERT_EQ(true, f1c.valid_content());
      ASSERT_EQ(true, f1d.valid_stream());
      ASSERT_EQ(true, f1d.valid_content());

      // Verify that streams were stolen.
      ASSERT_EQ_QUIET(f0a_stream, f1a.stream().native_handle());
      ASSERT_EQ_QUIET(f0b_stream, f1b.stream().native_handle());
      ASSERT_EQ_QUIET(f0c_stream, f1c.stream().native_handle());
      ASSERT_EQ_QUIET(f0d_stream, f1d.stream().native_handle());

      // This potentially runs concurrently with the copies.
      T const r0 = thrust::reduce(h0.begin(), h0.end());

      T r1a;
      T r1b;
      T r1c;
      T r1d;

      test_future_value_retrieval(f1a, r1a);
      test_future_value_retrieval(f1b, r1b);
      test_future_value_retrieval(f1c, r1c);
      test_future_value_retrieval(f1d, r1d);

      auto tolerance = std::max<T>(std::abs(0.1f * r0), T(precision<T>));

      ASSERT_NEAR(r0, r1a, tolerance);
      ASSERT_NEAR(r0, r1b, tolerance);
      ASSERT_NEAR(r0, r1c, tolerance);
      ASSERT_NEAR(r0, r1d, tolerance);
    }
  }
}

///////////////////////////////////////////////////////////////////////////////

// TODO: when_all from reductions.
