project_6/cccl_upstream/cub/cub/detail/env_dispatch.cuh

// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3-Clause
#pragma once

#include <cub/config.cuh>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <cub/detail/device_memory_resource.cuh>
#include <cub/detail/temporary_storage.cuh>

#include <cuda/__execution/tune.h>
#include <cuda/__functional/call_or.h>
#include <cuda/__memory_resource/get_memory_resource.h>
#include <cuda/__stream/get_stream.h>
#include <cuda/std/__execution/env.h>

CUB_NAMESPACE_BEGIN

namespace detail
{
//! @cond
//! Generic environment-based algorithm dispatch wrapper
//!
//! Handles common boilerplate for all env-based algorithms:
//! - Query stream, memory resource, and tuning from environment
//! - Two-phase call (query temp storage size, then execute)
//! - Temporary storage allocation/deallocation
//! - Memory resource querying from environment
//!
//! @param env The execution environment
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
template <typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
  // Query stream from environment
  auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);

  // Query memory resource from environment
  auto mr = ::cuda::__call_or(::cuda::mr::__get_memory_resource, detail::device_memory_resource{}, env);

  // Query tuning from environment
  const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);

  void* d_temp_storage      = nullptr;
  size_t temp_storage_bytes = 0;

  // Phase 1: Query temporary storage size
  if (const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get()))
  {
    return error;
  }

  // Allocate temporary storage
  if (const auto error = CubDebug(detail::temporary_storage::allocate(stream, d_temp_storage, temp_storage_bytes, mr)))
  {
    return error;
  }

  // Phase 2: Execute algorithm
  const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());

  // Deallocate temporary storage (always attempt, even on error)
  const auto deallocate_error =
    CubDebug(detail::temporary_storage::deallocate(stream, d_temp_storage, temp_storage_bytes, mr));

  // Algorithm error takes precedence over deallocation error
  return (error != cudaSuccess) ? error : deallocate_error;
}
//! @endcond

template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t
dispatch_with_env_and_tuning(const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
  return detail::dispatch_with_env(
    env,
    [&algorithm_callable](
      [[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
      using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
      using policy_selector =
        ::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
      return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
    });
}

//! @cond
//! Generic environment-based algorithm dispatch wrapper
//!
//! Handles common boilerplate for env-based algorithms with user provided memory:
//! - Query stream, and tuning from environment
//! - Single-phase call passing user provided memory and size
//!
//! @param env The execution environment
//! @param[in] d_temp_storage @devicestorage
//! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation
//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified
template <typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(
  void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
  // Query stream from environment
  auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);

  // Query tuning from environment
  const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env);

  return algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get());
}
//! @endcond

template <typename DefaultPolicySelector, typename EnvT, typename AlgorithmCallable>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env_and_tuning(
  void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable)
{
  return detail::dispatch_with_env(
    d_temp_storage,
    temp_storage_bytes,
    env,
    [&algorithm_callable](
      [[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) {
      using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{}));
      using policy_selector =
        ::cuda::std::execution::__query_result_or_t<decltype(tuning_env), policy_t, DefaultPolicySelector>;
      return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream);
    });
}
//! @endcond
} // namespace detail

CUB_NAMESPACE_END