WorkGroupLoopRange.h

// -*- tab-width: 2; indent-tabs-mode: nil; coding: utf-8-with-signature -*-
//-----------------------------------------------------------------------------
// Copyright 2000-2024 CEA (www.cea.fr) IFPEN (www.ifpenergiesnouvelles.com)
// See the top-level COPYRIGHT file for details.
// SPDX-License-Identifier: Apache-2.0
//-----------------------------------------------------------------------------
/*---------------------------------------------------------------------------*/
/* WorkGroupLoopRange.h                                        (C) 2000-2025 */
/*                                                                           */
/* Boucle pour le parallélisme hiérarchique.                                 */
/*---------------------------------------------------------------------------*/
#ifndef ARCANE_ACCELERATOR_WORKGROUPLOOPRANGE_H
#define ARCANE_ACCELERATOR_WORKGROUPLOOPRANGE_H
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#include "arcane/accelerator/AcceleratorUtils.h"
 
#if defined(ARCANE_COMPILING_CUDA)
#include <cooperative_groups.h>
#endif
#if defined(ARCANE_COMPILING_HIP)
#include <hip/hip_cooperative_groups.h>
#endif
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
namespace Arcane::Accelerator
{
namespace Impl
{
  class WorkGroupSequentialForHelper;
} // namespace Impl
 
class WorkGroupLoopRange;
class WorkGroupLoopContext;
class HostWorkItemGroup;
class SyclDeviceWorkItemBlock;
class DeviceWorkItemBlock;
class SyclWorkGroupLoopContext;
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class WorkItem
{
  friend WorkGroupLoopContext;
  friend SyclDeviceWorkItemBlock;
  friend DeviceWorkItemBlock;
  friend HostWorkItemGroup;
 
 private:

  explicit constexpr ARCCORE_HOST_DEVICE WorkItem(Int32 loop_index)
  : m_loop_index(loop_index)
  {}
 
 public:

  constexpr Int32 linearIndex() const { return m_loop_index; }
 
 private:
 
  Int32 m_loop_index = 0;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class HostWorkItemGroup
{
  friend WorkGroupLoopContext;
  friend SyclDeviceWorkItemBlock;
  friend DeviceWorkItemBlock;
 
 private:

  explicit constexpr ARCCORE_HOST_DEVICE HostWorkItemGroup(Int32 loop_index, Int32 group_index, Int32 group_size, Int32 nb_active_item)
  : m_loop_index(loop_index)
  , m_group_size(group_size)
  , m_group_index(group_index)
  , m_nb_active_item(nb_active_item)
  {}
 
 public:

  constexpr Int32 groupRank() const { return m_group_index; }

  constexpr Int32 groupSize() const { return m_group_size; }

  constexpr Int32 activeWorkItemRankInGroup() const { return 0; }

  static constexpr bool isDevice() { return false; }

  void barrier() {}

  constexpr Int32 nbActiveItem() const { return m_nb_active_item; }

  WorkItem activeItem(Int32 index) const
  {
    ARCANE_CHECK_AT(index, m_nb_active_item);
    return WorkItem(m_loop_index);
  }
 
 private:
 
  Int32 m_loop_index = 0;
  Int32 m_group_size = 0;
  Int32 m_group_index = 0;
  Int32 m_nb_active_item = 0;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#if defined(ARCANE_COMPILING_CUDA) || defined(ARCANE_COMPILING_HIP)
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class DeviceWorkItemBlock
{
  friend WorkGroupLoopContext;
 
 private:

  explicit __device__ DeviceWorkItemBlock()
  : m_thread_block(cooperative_groups::this_thread_block())
  {}
 
 public:

  __device__ Int32 groupRank() const { return m_thread_block.group_index().x; }

  __device__ Int32 groupSize() { return m_thread_block.group_dim().x; }

  __device__ Int32 activeWorkItemRankInGroup() const { return m_thread_block.thread_index().x; }

  __device__ void barrier() { m_thread_block.sync(); }

  static constexpr __device__ bool isDevice() { return true; }

  constexpr __device__ Int32 nbActiveItem() const { return 1; }

  __device__ WorkItem activeItem(Int32 index)
  {
    // Seulement valide pour index==0
    ARCANE_CHECK_AT(index, 1);
    return WorkItem(blockDim.x * blockIdx.x + threadIdx.x);
  }
 
 private:
 
  cooperative_groups::thread_block m_thread_block;
};
#endif
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class WorkGroupLoopContext
{
  // Pour accéder aux constructeurs
  friend WorkGroupLoopRange;
  friend Impl::WorkGroupSequentialForHelper;
 
 private:

  explicit constexpr WorkGroupLoopContext(Int32 loop_index, Int32 group_index, Int32 group_size, Int32 nb_active_item)
  : m_loop_index(loop_index)
  , m_group_index(group_index)
  , m_group_size(group_size)
  , m_nb_active_item(nb_active_item)
  {
  }
 
  // Ce constructeur n'est utilisé que sur le device
  // Il ne fait rien car les valeurs utiles sont récupérées via cooperative_groups::this_thread_block()
  explicit constexpr ARCCORE_DEVICE WorkGroupLoopContext() {}
 
 public:
 
#if defined(ARCCORE_DEVICE_CODE) && !defined(ARCANE_COMPILING_SYCL)
  __device__ DeviceWorkItemBlock group() const { return DeviceWorkItemBlock(); }
#else
  HostWorkItemGroup group() const { return HostWorkItemGroup(m_loop_index, m_group_index, m_group_size, m_nb_active_item); }
#endif
 
 private:
 
  Int32 m_loop_index = 0;
  Int32 m_group_index = 0;
  Int32 m_group_size = 0;
  Int32 m_nb_active_item = 0;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
/*
 * Implémentation pour SYCL.
 *
 * L'équivalent de \a cooperative_groups::thread_group() avec SYCL
 * est le \a sycl::nd_item<1>. Il est plus compliqué à utiliser pour deux
 * raisons:
 *
 * - il n'y a pas dans SYCL un équivalent de
 * \a cooperative_groups::this_thread_block(). Il faut utiliser la valeur
 * de \a sycl::nb_item<1> passé en argument du noyau de calcul.
 * - il n'y a pas de constructeurs par défaut pour \a sycl::nb_item<1>.
 *
 * Pour contourner ces deux problèmes, on utilise un type spécifique pour
 * gérer les noyaux en SYCL. Heureusement, il est possible d'utiliser les
 * lambda template avec SYCL. On utilise donc deux types pour gérer
 * les noyaux selon qu'on s'exécute sur le device SYCL ou sur l'hôte.
 *
 * TODO: regarder si avec la macro SYCL_DEVICE_ONLY il n'est pas possible
 * d'avoir le même type comportant des champs différents
 */
#if defined(ARCANE_COMPILING_SYCL)

class SyclDeviceWorkItemBlock
{
  friend SyclWorkGroupLoopContext;
 
 private:
 
  explicit SyclDeviceWorkItemBlock(sycl::nd_item<1> n)
  : m_nd_item(n)
  {
  }
 
 public:

  Int32 groupRank() const { return static_cast<Int32>(m_nd_item.get_group(0)); }

  Int32 groupSize() { return static_cast<Int32>(m_nd_item.get_local_range(0)); }

  Int32 activeWorkItemRankInGroup() const { return static_cast<Int32>(m_nd_item.get_local_id(0)); }

  void barrier() { m_nd_item.barrier(); }

  static constexpr bool isDevice() { return true; }

  constexpr Int32 nbActiveItem() const { return 1; }

  WorkItem activeItem(Int32 index)
  {
    // Seulement valide pour index==0
    ARCANE_CHECK_AT(index, 1);
    return WorkItem(static_cast<Int32>(m_nd_item.get_group(0) * m_nd_item.get_local_range(0) + m_nd_item.get_local_id(0)));
  }
 
 private:
 
  sycl::nd_item<1> m_nd_item;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class SyclWorkGroupLoopContext
{
  friend WorkGroupLoopRange;
 
 private:
 
  // Ce constructeur n'est utilisé que sur le device
  explicit SyclWorkGroupLoopContext(sycl::nd_item<1> n)
  : m_nd_item(n)
  {
  }
 
 public:

  SyclDeviceWorkItemBlock group() const { return SyclDeviceWorkItemBlock(m_nd_item); }
 
 private:
 
  sycl::nd_item<1> m_nd_item;
};
 
#endif // ARCANE_COMPILING_SYCL
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
class ARCANE_ACCELERATOR_EXPORT WorkGroupLoopRange
{
 private:
 
  friend ARCANE_ACCELERATOR_EXPORT WorkGroupLoopRange
  makeWorkGroupLoopRange(RunCommand& command, Int32 nb_group, Int32 group_size);
  friend ARCANE_ACCELERATOR_EXPORT WorkGroupLoopRange
  makeWorkGroupLoopRange(RunCommand& command, Int32 nb_element, Int32 nb_group, Int32 group_size);
 
 public:
 
  using LoopIndexType = WorkGroupLoopContext;
 
 public:
 
  WorkGroupLoopRange() = default;
 
 private:

  WorkGroupLoopRange(Int32 total_nb_element, Int32 nb_group, Int32 group_size);
 
 public:

  constexpr Int32 nbElement() const { return m_total_size; }
  constexpr Int32 groupSize() const { return m_group_size; }
  constexpr Int32 nbGroup() const { return m_nb_group; }
  constexpr Int32 lastGroupSize() const { return m_last_group_size; }
  constexpr Int32 nbActiveItem(Int32 i) const
  {
    return ((i + 1) != m_nb_group) ? m_group_size : m_last_group_size;
  }
 
 public:
 
  //TODO rendre privé ou mettre en externe
#if defined(ARCANE_COMPILING_CUDA) || defined(ARCANE_COMPILING_HIP)
  constexpr ARCCORE_HOST_DEVICE WorkGroupLoopContext getIndices(Int32) const { return WorkGroupLoopContext(); }
#endif
 
#if defined(ARCANE_COMPILING_SYCL)
  //TODO rendre privé ou mettre en externe
  SyclWorkGroupLoopContext getIndices(sycl::nd_item<1> id) const
  {
    return SyclWorkGroupLoopContext(id);
  }
#endif
 
 private:
 
  Int32 m_total_size = 0;
  Int32 m_nb_group = 0;
  Int32 m_group_size = 0;
  Int32 m_last_group_size = 0;
};
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
} // namespace Arcane::Accelerator
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/
 
#endif
 
/*---------------------------------------------------------------------------*/
/*---------------------------------------------------------------------------*/