/*******************************<GINKGO LICENSE>******************************
Copyright (c) 2017-2023, the Ginkgo authors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:

1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
******************************<GINKGO LICENSE>*******************************/

#ifndef GKO_PUBLIC_CORE_BASE_MACHINE_TOPOLOGY_HPP_
#define GKO_PUBLIC_CORE_BASE_MACHINE_TOPOLOGY_HPP_


#include <cassert>
#include <cstddef>
#include <functional>
#include <iostream>
#include <memory>
#include <mutex>
#include <string>
#include <type_traits>
#include <vector>


#include <ginkgo/config.hpp>
#include <ginkgo/core/base/exception.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>


#if GKO_HAVE_HWLOC

#include <hwloc.h>

#else

struct hwloc_obj_type_t {};
struct hwloc_obj_t {};

#endif


struct hwloc_topology;
struct hwloc_bitmap_s;


namespace gko {


/**
 * The machine topology class represents the hierarchical topology of a machine,
 * including NUMA nodes, cores and PCI Devices. Various information of the
 * machine are gathered with the help of the Hardware Locality library (hwloc).
 *
 * This class also provides functionalities to bind objects in the topology to
 * the execution objects. Binding can enhance performance by allowing data to be
 * closer to the executing object.
 *
 * See the hwloc documentation
 * (https://www.open-mpi.org/projects/hwloc/doc/) for more detailed
 * information on topology detection and binding interfaces.
 *
 * @note A global object of machine_topology type is created in a thread safe
 *       manner and only destroyed at the end of the program. This means that
 *       any subsequent queries will be from the same global object and hence
 *       use an extra atomic read.
 */
class machine_topology {
    template <typename T>
    using hwloc_manager = std::unique_ptr<T, std::function<void(T*)>>;

    /**
     * This struct holds the attributes for a normal non-IO object.
     */
    struct normal_obj_info {
        /**
         * The hwloc object.
         */
        hwloc_obj_t obj;

        /**
         * The logical_id assigned by HWLOC (assigned according to physical
         * proximity).
         *
         * @note Use this rather than os_id for all purposes other than binding.
         *       [Reference](https://www.open-mpi.org/projects/hwloc/doc/v2.4.0/a00364.php#faq_indexes)
         */
        size_type logical_id;

        /**
         * The os_id assigned by the OS (assigned arbitrarily by the OS)
         */
        size_type os_id;

        /**
         * The global persistent id assigned to the object by hwloc.
         */
        size_type gp_id;

        /**
         * The numa number of the object.
         */
        int numa;

        /**
         * The memory size of the object.
         */
        size_type memory_size;
    };


    /**
     * This struct holds the attributes for an IO/Misc object.
     *
     * Mainly used for PCI devices. The identifier important for PCI devices is
     * the PCI Bus ID, stored here as a string. PCI devices themselves usually
     * contain Hard disks, network components as well as other objects that are
     * not important for our case.
     *
     * In many cases, hwloc is able to identify the OS devices that belong to a
     * certain PCI Bus ID and here they are stored in the io children vector. A
     * list of their names are also additionally stored for easy access and
     * comparison.
     *
     * @note IO children can have names such as ibX for Infiniband cards, cudaX
     * for NVIDIA cards with CUDA and rsmiX for AMD cards.
     */
    struct io_obj_info {
        /**
         * The hwloc object.
         */
        hwloc_obj_t obj;

        /**
         * The logical_id assigned by HWLOC (assigned according to proximity).
         *
         * @note Use this rather than os_id for all purposes other than binding.
         * [Reference](https://www.open-mpi.org/projects/hwloc/doc/v2.4.0/a00364.php#faq_indexes)
         */
        size_type logical_id;

        /**
         * The os_id assigned by the OS (assigned arbitrarily by the OS)
         */
        size_type os_id;

        /**
         * The global persistent id assigned to the object by hwloc.
         */
        size_type gp_id;

        /**
         * The closest numa.
         */
        int closest_numa;

        /**
         * The non-io parent object.
         */
        hwloc_obj_t non_io_ancestor;

        /**
         * The ancestor local id.
         */
        int ancestor_local_id;

        /**
         * The ancestor type.
         */
        std::string ancestor_type;

        /**
         * The array of CPU ids closest to the object.
         */
        std::vector<int> closest_pu_ids;

        /**
         * The PCI Bus ID
         */
        std::string pci_bus_id;
    };

public:
    /**
     * Returns an instance of the machine_topology object.
     *
     * @return  the machine_topology instance
     */
    static machine_topology* get_instance()
    {
        static machine_topology instance;
        return &instance;
    }

    /**
     * Bind the calling process to the CPU cores associated with
     * the ids.
     *
     * @param ids  The ids of cores to be bound.
     * @param singlify  The ids of PUs are singlified to prevent possibly
     *                  expensive migrations by the OS. This means that the
     *                  binding is performed for only one of the ids in the
     *                  set of ids passed in.
     *                  See hwloc doc for
     *                  [singlify](https://www.open-mpi.org/projects/hwloc/doc/v2.4.0/a00175.php#gaa611a77c092e679246afdf9a60d5db8b)
     */
    void bind_to_cores(const std::vector<int>& ids,
                       const bool singlify = true) const
    {
        hwloc_binding_helper(this->cores_, ids, singlify);
    }

    /**
     * Bind to a single core
     *
     * @param ids  The ids of the core to be bound to the calling process.
     */
    void bind_to_core(const int& id) const
    {
        machine_topology::get_instance()->bind_to_cores(std::vector<int>{id});
    }

    /**
     * Bind the calling process to PUs associated with
     * the ids.
     *
     * @param ids  The ids of PUs to be bound.
     * @param singlify  The ids of PUs are singlified to prevent possibly
     *                  expensive migrations by the OS. This means that the
     *                  binding is performed for only one of the ids in the
     *                  set of ids passed in.
     *                  See hwloc doc for
     *                  [singlify](https://www.open-mpi.org/projects/hwloc/doc/v2.4.0/a00175.php#gaa611a77c092e679246afdf9a60d5db8b)
     */
    void bind_to_pus(const std::vector<int>& ids,
                     const bool singlify = true) const
    {
        hwloc_binding_helper(this->pus_, ids, singlify);
    }

    /**
     * Bind to a Processing unit (PU)
     *
     * @param ids  The ids of PUs to be bound to the calling process.
     */
    void bind_to_pu(const int& id) const
    {
        machine_topology::get_instance()->bind_to_pus(std::vector<int>{id});
    }

    /**
     * Get the object of type PU associated with the id.
     *
     * @param id  The id of the PU
     * @return  the PU object struct.
     */
    const normal_obj_info* get_pu(size_type id) const
    {
        GKO_ENSURE_IN_BOUNDS(id, this->pus_.size());
        return &this->pus_[id];
    }

    /**
     * Get the object of type core associated with the id.
     *
     * @param id  The id of the core
     * @return  the core object struct.
     */
    const normal_obj_info* get_core(size_type id) const
    {
        GKO_ENSURE_IN_BOUNDS(id, this->cores_.size());
        return &this->cores_[id];
    }

    /**
     * Get the object of type pci device associated with the id.
     *
     * @param id  The id of the pci device
     * @return  the PCI object struct.
     */
    const io_obj_info* get_pci_device(size_type id) const
    {
        GKO_ENSURE_IN_BOUNDS(id, this->pci_devices_.size());
        return &this->pci_devices_[id];
    }

    /**
     * Get the object of type pci device associated with the PCI bus id.
     *
     * @param pci_bus_id  The PCI bus id of the pci device
     * @return  the PCI object struct.
     */
    const io_obj_info* get_pci_device(const std::string& pci_bus_id) const;

    /**
     * Get the number of PU objects stored in this Topology tree.
     *
     * @return  the number of PUs.
     */
    size_type get_num_pus() const { return this->pus_.size(); }

    /**
     * Get the number of core objects stored in this Topology tree.
     *
     * @return  the number of cores.
     */
    size_type get_num_cores() const { return this->cores_.size(); }

    /**
     * Get the number of PCI device objects stored in this Topology tree.
     *
     * @return  the number of PCI devices.
     */
    size_type get_num_pci_devices() const { return this->pci_devices_.size(); }

    /**
     * Get the number of NUMA objects stored in this Topology tree.
     *
     * @return  the number of NUMA objects.
     */
    size_type get_num_numas() const { return this->num_numas_; }

    /**
     * @internal
     *
     * A helper function that binds the calling process with the ids of `obj`
     * object .
     */
    void hwloc_binding_helper(
        const std::vector<machine_topology::normal_obj_info>& obj,
        const std::vector<int>& ids, const bool singlify = true) const;

    /**
     * @internal
     *
     * Load the objects of a normal HWLOC type (Packages, cores, numa-nodes).
     *
     * @note The objects should be sorted by logical index since hwloc uses
     * logical index with these functions
     */
    void load_objects(hwloc_obj_type_t type,
                      std::vector<normal_obj_info>& objects) const;

    /**
     * @internal
     *
     * Load the objects of io type (PCI devices and OS devices).
     *
     * @note The objects should be sorted by logical index since hwloc uses
     * logical index with these functions
     */
    void load_objects(hwloc_obj_type_t type,
                      std::vector<io_obj_info>& vector) const;

    /**
     *
     * @internal
     *
     * Get object id from the os index
     */
    int get_obj_id_by_os_index(const std::vector<normal_obj_info>& objects,
                               size_type os_index) const;

    /**
     *
     * @internal
     *
     * Get object id from the hwloc index
     */
    int get_obj_id_by_gp_index(const std::vector<normal_obj_info>& objects,
                               size_type gp_index) const;

private:
    /**
     * Do not allow the machine_topology object to be copied/moved. There should
     * be only one global object per execution.
     */
    machine_topology();
    machine_topology(machine_topology&) = delete;
    machine_topology(machine_topology&&) = delete;
    machine_topology& operator=(machine_topology&) = delete;
    machine_topology& operator=(machine_topology&&) = delete;
    ~machine_topology() = default;

    std::vector<normal_obj_info> pus_;
    std::vector<normal_obj_info> cores_;
    std::vector<normal_obj_info> packages_;
    std::vector<normal_obj_info> numa_nodes_;
    std::vector<io_obj_info> pci_devices_;
    size_type num_numas_;

    hwloc_manager<hwloc_topology> topo_;
};


using MachineTopology GKO_DEPRECATED("please use machine_topology") =
    machine_topology;


}  // namespace gko


#endif  // GKO_PUBLIC_CORE_BASE_MACHINE_TOPOLOGY_HPP_
