GenericPartitioner.h

// -*- tab-width: 2; indent-tabs-mode: nil; coding: utf-8-with-signature -*-
//-----------------------------------------------------------------------------
// Copyright 2000-2025 CEA (www.cea.fr) IFPEN (www.ifpenergiesnouvelles.com)
// See the top-level COPYRIGHT file for details.
// SPDX-License-Identifier: Apache-2.0
//-----------------------------------------------------------------------------
/*---------------------------------------------------------------------------*/
/* GenericPartitioner.h                                        (C) 2000-2025 */
/*                                                                           */
/* Algorithme de partitionnement de liste.                                   */
/*---------------------------------------------------------------------------*/
#ifndef ARCANE_ACCELERATOR_GENERICPARTITIONER_H
#define ARCANE_ACCELERATOR_GENERICPARTITIONER_H
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#include "arcane/utils/ArrayView.h"
#include "arcane/utils/FatalErrorException.h"
#include "arcane/utils/NotImplementedException.h"
#include "arcane/utils/NumArray.h"
 
#include "arcane/accelerator/AcceleratorGlobal.h"
#include "arcane/accelerator/core/RunQueue.h"
#include "arcane/accelerator/CommonUtils.h"
#include "arcane/accelerator/RunCommandLaunchInfo.h"
#if defined(ARCANE_COMPILING_SYCL)
#include "arcane/accelerator/RunCommandLoop.h"
#endif
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
namespace Arcane::Accelerator::impl
{
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class ARCANE_ACCELERATOR_EXPORT GenericPartitionerBase
{
  friend class GenericPartitionerIf;
 
 public:
 
  explicit GenericPartitionerBase(const RunQueue& queue);
 
 protected:
 
  Int32 _nbFirstPart() const;
  SmallSpan<const Int32> _nbParts() const;
  void _allocate();
  void _resetNbPart();
 
 protected:
 
  RunQueue m_queue;
  GenericDeviceStorage m_algo_storage;
  DeviceStorage<int, 2> m_device_nb_list1_storage;
  NumArray<Int32, MDDim1> m_host_nb_list1_storage;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class GenericPartitionerIf
{
 public:

  template <typename SelectLambda, typename InputIterator, typename OutputIterator>
  void apply(GenericPartitionerBase& s, Int32 nb_item, InputIterator input_iter, OutputIterator output_iter,
             const SelectLambda& select_lambda, const TraceInfo& trace_info = TraceInfo())
  {
    RunQueue queue = s.m_queue;
    eExecutionPolicy exec_policy = queue.executionPolicy();
    RunCommand command = makeCommand(queue);
    command << trace_info;
    Impl::RunCommandLaunchInfo launch_info(command, nb_item);
    launch_info.beginExecute();
    switch (exec_policy) {
#if defined(ARCANE_COMPILING_CUDA)
    case eExecutionPolicy::CUDA: {
      size_t temp_storage_size = 0;
      cudaStream_t stream = Impl::CudaUtils::toNativeStream(&queue);
      // Premier appel pour connaitre la taille pour l'allocation
      int* nb_list1_ptr = nullptr;
      ARCANE_CHECK_CUDA(::cub::DevicePartition::If(nullptr, temp_storage_size,
                                                   input_iter, output_iter, nb_list1_ptr, nb_item,
                                                   select_lambda, stream));
 
      s.m_algo_storage.allocate(temp_storage_size);
      nb_list1_ptr = s.m_device_nb_list1_storage.allocate();
      ARCANE_CHECK_CUDA(::cub::DevicePartition::If(s.m_algo_storage.address(), temp_storage_size,
                                                   input_iter, output_iter, nb_list1_ptr, nb_item,
                                                   select_lambda, stream));
      s.m_device_nb_list1_storage.copyToAsync(s.m_host_nb_list1_storage, queue);
    } break;
#endif
#if defined(ARCANE_COMPILING_HIP)
    case eExecutionPolicy::HIP: {
      size_t temp_storage_size = 0;
      // Premier appel pour connaitre la taille pour l'allocation
      hipStream_t stream = Impl::HipUtils::toNativeStream(&queue);
      int* nb_list1_ptr = nullptr;
      ARCANE_CHECK_HIP(rocprim::partition(nullptr, temp_storage_size, input_iter, output_iter,
                                          nb_list1_ptr, nb_item, select_lambda, stream));
 
      s.m_algo_storage.allocate(temp_storage_size);
      nb_list1_ptr = s.m_device_nb_list1_storage.allocate();
 
      ARCANE_CHECK_HIP(rocprim::partition(s.m_algo_storage.address(), temp_storage_size, input_iter, output_iter,
                                          nb_list1_ptr, nb_item, select_lambda, stream));
      s.m_device_nb_list1_storage.copyToAsync(s.m_host_nb_list1_storage, queue);
    } break;
#endif
#if defined(ARCANE_COMPILING_SYCL)
    case eExecutionPolicy::SYCL: {
#if defined(ARCANE_HAS_ONEDPL)
      // Seulement implémenté pour DPC++.
      // Actuellement (dpc++ 2025.0), il n'y a pas l'équivalent avec SYCL de
      // la méthode de partition de cub ou rocprim.
      // Utilise la fonction 'stable_partition'. Cependant, cette fonction
      // ne supporte pas si la mémoire uniquement sur accélérateur. De plus,
      // il faut que InputIterator et OutputIterator remplissent le concept
      // std::random_access_iterator ce qui n'est pas le cas (notamment car il n'y a
      // pas les opérateurs de copie ni de constructeur vide à cause de la lambda).
      // Pour éviter tous ces problèmes, on alloue donc des tableaux temporaires
      // pour l'appel à 'stable_sort' et on recopie les valeurs en sortie.
      using InputDataType = typename InputIterator::value_type;
      using DataType = typename OutputIterator::value_type;
      NumArray<DataType, MDDim1> tmp_output_numarray(nb_item);
      NumArray<bool, MDDim1> tmp_select_numarray(nb_item);
      auto tmp_output = tmp_output_numarray.to1DSmallSpan();
      auto tmp_select = tmp_select_numarray.to1DSmallSpan();
      {
        auto command = makeCommand(queue);
        command << RUNCOMMAND_LOOP1(iter, nb_item)
        {
          auto [i] = iter();
          tmp_output[i] = input_iter[i];
        };
      }
      auto tmp_select_lambda = [=](Int32 i) { return tmp_select[i]; };
      sycl::queue sycl_queue = Impl::SyclUtils::toNativeStream(queue);
      auto policy = oneapi::dpl::execution::make_device_policy(sycl_queue);
      auto output_after = oneapi::dpl::stable_partition(policy, tmp_output.begin(), tmp_output.end(), select_lambda);
      queue.barrier();
      Int32 nb_list1 = (output_after - tmp_output.begin());
      Int32 nb_list2 = nb_item - nb_list1;
      s.m_host_nb_list1_storage[0] = nb_list1;
      //std::cerr << "NbList1=" << nb_list1 << " NbList2=" << nb_list2 << "\n";
      {
        // Recopie dans output les valeurs filtrées.
        // Pour être cohérent avec 'cub' et 'rocprim', il faut inverser l'ordre des
        // des valeurs de la liste pour les éléments ne remplissant pas la condition.
        // Pour cela, on fait une boucle de taille (nb_list1 + nb_list/2) et chaque
        // itération pour i>=nb_list1 gère deux éléments.
        auto command = makeCommand(queue);
        Int32 nb_iter2 = (nb_list2 / 2) + (nb_list2 % 2);
        //std::cout << "NB_ITER2=" << nb_iter2 << "\n";
        command << RUNCOMMAND_LOOP1(iter, (nb_list1 + nb_iter2))
        {
          auto [i] = iter();
          if (i >= nb_list1) {
            // Partie de la liste pour les valeurs ne remplissant par le critère.
            Int32 j = i - nb_list1;
            Int32 reverse_i = (nb_item - (j + 1));
            auto x1 = tmp_output[i];
            auto x2 = tmp_output[reverse_i];
            output_iter[i] = tmp_output[reverse_i];
            output_iter[reverse_i] = tmp_output[i];
          }
          else
            output_iter[i] = tmp_output[i];
        };
      }
      queue.barrier();
#else // ARCANE_HAS_ONEDPL
      ARCANE_THROW(NotImplementedException,"Partition is only implemented for SYCL back-end using oneDPL");
#endif // ARCANE_HAS_ONEDPL
    } break;
#endif // ARCANE_COMPILING_SYCL
    case eExecutionPolicy::Thread:
      // Pas encore implémenté en multi-thread
      [[fallthrough]];
    case eExecutionPolicy::Sequential: {
      auto saved_output_iter = output_iter;
      auto output2_iter = output_iter + nb_item;
      for (Int32 i = 0; i < nb_item; ++i) {
        auto v = *input_iter;
        if (select_lambda(v)) {
          *output_iter = v;
          ++output_iter;
        }
        else {
          --output2_iter;
          *output2_iter = v;
        }
        ++input_iter;
      }
      Int32 nb_list1 = static_cast<Int32>(output_iter - saved_output_iter);
      s.m_host_nb_list1_storage[0] = nb_list1;
    } break;
    default:
      ARCANE_FATAL(getBadPolicyMessage(exec_policy));
    }
    launch_info.endExecute();
  }

  template <typename Select1Lambda, typename Select2Lambda,
            typename InputIterator, typename FirstOutputIterator,
            typename SecondOutputIterator, typename UnselectedIterator>
  void apply3(GenericPartitionerBase& s, Int32 nb_item,
              InputIterator input_iter,
              FirstOutputIterator first_output_iter,
              SecondOutputIterator second_output_iter,
              UnselectedIterator unselected_iter,
              const Select1Lambda& select1_lambda,
              const Select2Lambda& select2_lambda,
              const TraceInfo& trace_info = TraceInfo())
  {
    RunQueue queue = s.m_queue;
    eExecutionPolicy exec_policy = queue.executionPolicy();
    RunCommand command = makeCommand(queue);
    command << trace_info;
    Impl::RunCommandLaunchInfo launch_info(command, nb_item);
    launch_info.beginExecute();
    switch (exec_policy) {
#if defined(ARCANE_COMPILING_CUDA)
    case eExecutionPolicy::CUDA: {
      size_t temp_storage_size = 0;
      cudaStream_t stream = Impl::CudaUtils::toNativeStream(&queue);
      // Premier appel pour connaitre la taille pour l'allocation
      int* nb_list1_ptr = nullptr;
      ARCANE_CHECK_CUDA(::cub::DevicePartition::If(nullptr, temp_storage_size,
                                                   input_iter, first_output_iter, second_output_iter,
                                                   unselected_iter, nb_list1_ptr, nb_item,
                                                   select1_lambda, select2_lambda, stream));
 
      s.m_algo_storage.allocate(temp_storage_size);
      nb_list1_ptr = s.m_device_nb_list1_storage.allocate();
      ARCANE_CHECK_CUDA(::cub::DevicePartition::If(s.m_algo_storage.address(), temp_storage_size,
                                                   input_iter, first_output_iter, second_output_iter,
                                                   unselected_iter, nb_list1_ptr, nb_item,
                                                   select1_lambda, select2_lambda, stream));
      s.m_device_nb_list1_storage.copyToAsync(s.m_host_nb_list1_storage, queue);
    } break;
#endif
#if defined(ARCANE_COMPILING_HIP)
    case eExecutionPolicy::HIP: {
      size_t temp_storage_size = 0;
      // Premier appel pour connaitre la taille pour l'allocation
      hipStream_t stream = Impl::HipUtils::toNativeStream(&queue);
      int* nb_list1_ptr = nullptr;
      using namespace rocprim;
      ARCANE_CHECK_HIP(::rocprim::partition_three_way(nullptr, temp_storage_size, input_iter, first_output_iter,
                                                      second_output_iter, unselected_iter,
                                                      nb_list1_ptr, nb_item, select1_lambda, select2_lambda, stream));
 
      s.m_algo_storage.allocate(temp_storage_size);
      nb_list1_ptr = s.m_device_nb_list1_storage.allocate();
 
      ARCANE_CHECK_HIP(partition_three_way(s.m_algo_storage.address(), temp_storage_size, input_iter, first_output_iter,
                                           second_output_iter, unselected_iter, nb_list1_ptr, nb_item,
                                           select1_lambda, select2_lambda, stream));
      s.m_device_nb_list1_storage.copyToAsync(s.m_host_nb_list1_storage, queue);
    } break;
#endif
#if defined(ARCANE_COMPILING_SYCL)
    case eExecutionPolicy::SYCL: {
      ARCANE_THROW(NotImplementedException,"3-way partition is not implemented for SYCL back-end");
    }
    break;
#endif
    case eExecutionPolicy::Thread:
      // Pas encore implémenté en multi-thread
      [[fallthrough]];
    case eExecutionPolicy::Sequential: {
      Int32 nb_first = 0;
      Int32 nb_second = 0;
      for (Int32 i = 0; i < nb_item; ++i) {
        auto v = *input_iter;
        bool is_1 = select1_lambda(v);
        bool is_2 = select2_lambda(v);
        if (is_1) {
          *first_output_iter = v;
          ++first_output_iter;
          ++nb_first;
        }
        else {
          if (is_2) {
            *second_output_iter = v;
            ++second_output_iter;
            ++nb_second;
          }
          else {
            *unselected_iter = v;
            ++unselected_iter;
          }
        }
        // Incrémenter l'itérateur à la fin car il est utilisé pour le positionnement
        ++input_iter;
      }
      s.m_host_nb_list1_storage[0] = nb_first;
      s.m_host_nb_list1_storage[1] = nb_second;
    } break;
    default:
      ARCANE_FATAL(getBadPolicyMessage(exec_policy));
    }
    launch_info.endExecute();
  }
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
} // namespace Arcane::Accelerator::impl
 
namespace Arcane::Accelerator
{
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class GenericPartitioner
: private impl::GenericPartitionerBase
{
 public:
 
  explicit GenericPartitioner(const RunQueue& queue)
  : impl::GenericPartitionerBase(queue)
  {
    _allocate();
  }
 
 public:

  template <typename SelectLambda, typename SetterLambda>
  void applyWithIndex(Int32 nb_value, const SetterLambda& setter_lambda,
                      const SelectLambda& select_lambda, const TraceInfo& trace_info = TraceInfo())
  {
    if (_checkEmpty(nb_value))
      return;
    _setCalled();
    impl::GenericPartitionerBase* base_ptr = this;
    impl::GenericPartitionerIf gf;
 
    impl::IndexIterator input_iter;
    impl::SetterLambdaIterator<SetterLambda> out(setter_lambda);
    gf.apply(*base_ptr, nb_value, input_iter, out, select_lambda, trace_info);
  }

  template <typename InputIterator, typename OutputIterator, typename SelectLambda>
  void applyIf(Int32 nb_value, InputIterator input_iter, OutputIterator output_iter,
               const SelectLambda& select_lambda, const TraceInfo& trace_info = TraceInfo())
  {
    if (_checkEmpty(nb_value))
      return;
    _setCalled();
    impl::GenericPartitionerBase* base_ptr = this;
    impl::GenericPartitionerIf gf;
    gf.apply(*base_ptr, nb_value, input_iter, output_iter, select_lambda, trace_info);
  }

  template <typename Setter1Lambda, typename Setter2Lambda, typename UnselectedSetterLambda,
            typename Select1Lambda, typename Select2Lambda>
  void applyWithIndex(Int32 nb_value,
                      const Setter1Lambda setter1_lambda,
                      const Setter2Lambda setter2_lambda,
                      const UnselectedSetterLambda& unselected_setter_lambda,
                      const Select1Lambda& select1_lambda,
                      const Select2Lambda& select2_lambda,
                      const TraceInfo& trace_info = TraceInfo())
  {
    if (_checkEmpty(nb_value))
      return;
    _setCalled();
    impl::GenericPartitionerBase* base_ptr = this;
    impl::GenericPartitionerIf gf;
    impl::IndexIterator input_iter;
    impl::SetterLambdaIterator<Setter1Lambda> setter1_wrapper(setter1_lambda);
    impl::SetterLambdaIterator<Setter2Lambda> setter2_wrapper(setter2_lambda);
    impl::SetterLambdaIterator<UnselectedSetterLambda> unselected_setter_wrapper(unselected_setter_lambda);
    gf.apply3(*base_ptr, nb_value, input_iter, setter1_wrapper, setter2_wrapper,
              unselected_setter_wrapper, select1_lambda, select2_lambda, trace_info);
  }

  template <typename InputIterator, typename FirstOutputIterator,
            typename SecondOutputIterator, typename UnselectedIterator,
            typename Select1Lambda, typename Select2Lambda>
  void applyIf(Int32 nb_value, InputIterator input_iter,
               FirstOutputIterator first_output_iter,
               SecondOutputIterator second_output_iter,
               UnselectedIterator unselected_iter,
               const Select1Lambda& select1_lambda,
               const Select2Lambda& select2_lambda,
               const TraceInfo& trace_info = TraceInfo())
  {
    if (_checkEmpty(nb_value))
      return;
    _setCalled();
    impl::GenericPartitionerBase* base_ptr = this;
    impl::GenericPartitionerIf gf;
    gf.apply3(*base_ptr, nb_value, input_iter, first_output_iter, second_output_iter,
              unselected_iter, select1_lambda, select2_lambda, trace_info);
  }

  Int32 nbFirstPart()
  {
    m_is_already_called = false;
    return _nbFirstPart();
  }

  SmallSpan<const Int32> nbParts()
  {
    m_is_already_called = false;
    return _nbParts();
  }
 
 private:
 
  bool m_is_already_called = false;
 
 private:
 
  void _setCalled()
  {
    if (m_is_already_called)
      ARCANE_FATAL("apply() has already been called for this instance");
    m_is_already_called = true;
  }
  bool _checkEmpty(Int32 nb_value)
  {
    if (nb_value == 0) {
      _resetNbPart();
      return true;
    }
    return false;
  }
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
} // namespace Arcane::Accelerator
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#endif
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/