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          654 :     bool HasNode(const NodeId nodeId) const { return nodes.find(nodeId) != nodes.end(); }
      28              : 
      29          198 :     bool HasEdge(const NodeId srcNodeId, const NodeId dstNodeId) const
      30              :     {
      31          198 :         return edges.find(srcNodeId) != edges.end() && edges.at(srcNodeId).find(dstNodeId) != edges.at(srcNodeId).end();
      32              :     }
      33              : 
      34          196 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId, const NodeId dstNodeId) const
      35              :     {
      36          196 :         if (!HasEdge(srcNodeId, dstNodeId)) {
      37          120 :             return {};
      38              :         }
      39           76 :         return edges.at(srcNodeId).at(dstNodeId);
      40              :     }
      41              : 
      42            4 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId) const
      43              :     {
      44            4 :         if (edges.find(srcNodeId) == edges.end()) {
      45            0 :             return {};
      46              :         }
      47            4 :         std::vector<std::shared_ptr<EdgeType>> nodeEdges;
      48           10 :         for (const auto& it : edges.at(srcNodeId)) {
      49            6 :             nodeEdges.insert(nodeEdges.end(), it.second.begin(), it.second.end());
      50              :         }
      51            4 :         return nodeEdges;
      52            4 :     }
      53              : 
      54          115 :     void TraverseNode(std::function<void(NodeId nodeId, const std::shared_ptr<NodeType>&)> func) const
      55              :     {
      56          564 :         for (auto& node : nodes) {
      57          449 :             func(node.first, node.second);
      58              :         }
      59          115 :     }
      60              : 
      61           38 :     void TraverseNode(std::function<void(std::shared_ptr<NodeType>)> func) const
      62              :     {
      63          199 :         for (auto& node : nodes) {
      64          161 :             func(node.second);
      65              :         }
      66           38 :     }
      67              : 
      68          263 :     void TraverseEdge(const NodeId srcNodeId, std::function<void(std::shared_ptr<EdgeType>)> func) const
      69              :     {
      70          263 :         if (edges.find(srcNodeId) == edges.end()) {
      71           12 :             return;
      72              :         }
      73          701 :         for (auto& srcEdges : edges.at(srcNodeId)) {
      74          946 :             for (auto& edge : srcEdges.second) {
      75          496 :                 func(edge);
      76              :             }
      77              :         }
      78              :     }
      79              : 
      80         1315 :     void TraverseEdge(
      81              :         const NodeId srcNodeId, const NodeId dstNodeId, std::function<void(std::shared_ptr<EdgeType>)> func) const
      82              :     {
      83         1315 :         if (edges.find(srcNodeId) == edges.end() || edges.at(srcNodeId).find(dstNodeId) == edges.at(srcNodeId).end()) {
      84          788 :             return;
      85              :         }
      86         1078 :         for (auto& edge : edges.at(srcNodeId).at(dstNodeId)) {
      87          551 :             func(edge);
      88              :         }
      89              :     }
      90              : 
      91          596 :     void AddNode(const NodeId nodeId, std::shared_ptr<NodeType> node)
      92              :     {
      93          596 :         nodes[nodeId] = node;
      94          912 :         HCCL_DEBUG("[Graph]add node [%llu] success! node number is [%zu]", nodeId, nodes.size());
      95          596 :     }
      96              : 
      97          756 :     void AddEdge(const NodeId srcNodeId, const NodeId dstNodeId, std::shared_ptr<EdgeType> edge)
      98              :     {
      99          756 :         edges[srcNodeId][dstNodeId].push_back(edge);
     100         1544 :         HCCL_DEBUG(
     101              :             "[Graph]add edge from node [%llu] to node [%llu] success! edge number is [%zu]", srcNodeId, dstNodeId,
     102              :             edges[srcNodeId][dstNodeId].size());
     103          756 :     }
     104              : 
     105           16 :     void DeleteEdge(const NodeId srcNodeId, const NodeId dstNodeId)
     106              :     {
     107              :         // 在64+1中srcNode 和 dstNode代表fabric和db
     108           16 :         auto srcIt = edges.find(srcNodeId);
     109           16 :         if (srcIt == edges.end()) {
     110            8 :             return;
     111              :         }
     112              : 
     113           16 :         auto& dstMap = srcIt->second;
     114           16 :         auto dstIt = dstMap.find(dstNodeId);
     115           16 :         if (dstIt == dstMap.end()) {
     116            8 :             return;
     117              :         }
     118              : 
     119              :         // 删除对应的边
     120            8 :         dstMap.erase(dstIt);
     121              :     }
     122              : 
     123              : private:
     124              :     std::unordered_map<NodeId, std::unordered_map<NodeId, std::vector<std::shared_ptr<EdgeType>>>> edges;
     125              :     std::unordered_map<NodeId, std::shared_ptr<NodeType>> nodes;
     126              : };
     127              : } // namespace Hccl
     128              : 
     129              : #endif // GROUP_H