Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2025 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #ifndef GROUP_H
      12              : #define GROUP_H
      13              : 
      14              : #include <unordered_map>
      15              : #include <functional>
      16              : #include <memory>
      17              : #include <vector>
      18              : #include <algorithm>
      19              : #include "topo_common_types.h"
      20              : #include "log.h"
      21              : 
      22              : namespace Hccl {
      23              : 
      24              : template <typename NodeType, typename EdgeType> 
      25              : class Graph {
      26              : public:
      27          163 :     bool HasNode(const NodeId nodeId) const
      28              :     {
      29          163 :         return nodes.find(nodeId) != nodes.end();
      30              :     }
      31              : 
      32            8 :     bool HasEdge(const NodeId srcNodeId, const NodeId dstNodeId) const
      33              :     {
      34            8 :         return edges.find(srcNodeId) != edges.end() && edges.at(srcNodeId).find(dstNodeId) != edges.at(srcNodeId).end();
      35              :     }
      36              : 
      37            8 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId, const NodeId dstNodeId) const
      38              :     {
      39            8 :         if (!HasEdge(srcNodeId, dstNodeId)) {
      40            6 :             return {};
      41              :         }
      42            2 :         return edges.at(srcNodeId).at(dstNodeId);
      43              :     }
      44              : 
      45            3 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId) const
      46              :     {
      47            3 :         if (edges.find(srcNodeId) == edges.end()) {
      48            0 :             return {};
      49              :         }
      50            3 :         std::vector<std::shared_ptr<EdgeType>> nodeEdges;
      51            6 :         for (const auto &it : edges.at(srcNodeId)) {
      52            3 :             nodeEdges.insert(nodeEdges.end(), it.second.begin(), it.second.end());
      53              :         }
      54            3 :         return nodeEdges;
      55            3 :     }
      56              : 
      57           34 :     void TraverseNode(std::function<void(NodeId nodeId, const std::shared_ptr<NodeType> &)> func) const
      58              :     {
      59          137 :         for (auto &node : nodes) {
      60          103 :             func(node.first, node.second);
      61              :         }
      62           34 :     }
      63              : 
      64           14 :     void TraverseNode(std::function<void(std::shared_ptr<NodeType>)> func) const
      65              :     {
      66           51 :         for (auto &node : nodes) {
      67           37 :             func(node.second);
      68              :         }
      69           14 :     }
      70              : 
      71           72 :     void TraverseEdge(const NodeId srcNodeId, std::function<void(std::shared_ptr<EdgeType>)> func) const
      72              :     {
      73           72 :         if (edges.find(srcNodeId) == edges.end()) {
      74            6 :             return;
      75              :         }
      76          166 :         for (auto &srcEdges : edges.at(srcNodeId)) {
      77          204 :             for (auto &edge : srcEdges.second) {
      78          104 :                 func(edge);
      79              :             }
      80              :         }
      81              :     }
      82              : 
      83          720 :     void TraverseEdge(const NodeId srcNodeId, const NodeId dstNodeId,
      84              :                       std::function<void(std::shared_ptr<EdgeType>)> func) const
      85              :     {
      86          720 :         if (edges.find(srcNodeId) == edges.end() || edges.at(srcNodeId).find(dstNodeId) == edges.at(srcNodeId).end()) {
      87          432 :             return;
      88              :         }
      89          576 :         for (auto &edge : edges.at(srcNodeId).at(dstNodeId)) {
      90          288 :             func(edge);
      91              :         }
      92              :     }
      93              : 
      94          340 :     void AddNode(const NodeId nodeId, std::shared_ptr<NodeType> node)
      95              :     {
      96          340 :         nodes[nodeId] = node;
      97          340 :         HCCL_DEBUG("[Graph]add node [%llu] success! node number is [%zu]", nodeId, nodes.size());
      98          340 :     }
      99              : 
     100          313 :     void AddEdge(const NodeId srcNodeId, const NodeId dstNodeId, std::shared_ptr<EdgeType> edge)
     101              :     {
     102          313 :         edges[srcNodeId][dstNodeId].push_back(edge);
     103          313 :         HCCL_DEBUG("[Graph]add edge from node [%llu] to node [%llu] success! edge number is [%zu]", srcNodeId,
     104              :                    dstNodeId, edges[srcNodeId][dstNodeId].size());
     105          313 :     }
     106              : 
     107            0 :     void DeleteEdge(const NodeId srcNodeId, const NodeId dstNodeId)
     108              :     {
     109              :         // 在64+1中srcNode 和 dstNode代表fabric和db
     110            0 :         auto srcIt = edges.find(srcNodeId);
     111            0 :         if (srcIt == edges.end()) {
     112            0 :             return;
     113              :         }
     114              : 
     115            0 :         auto& dstMap = srcIt->second;
     116            0 :         auto dstIt = dstMap.find(dstNodeId);
     117            0 :         if (dstIt == dstMap.end()) {
     118            0 :             return;
     119              :         }
     120              : 
     121              :         // 删除对应的边
     122            0 :         dstMap.erase(dstIt);
     123              :     }
     124              : 
     125              : private:
     126              :     std::unordered_map<NodeId, std::unordered_map<NodeId, std::vector<std::shared_ptr<EdgeType>>>> edges;
     127              :     std::unordered_map<NodeId, std::shared_ptr<NodeType>>                                          nodes;
     128              : };
     129              : } // namespace Hccl
     130              : 
     131              : #endif // GROUP_H