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          491 :     bool HasNode(const NodeId nodeId) const
      28              :     {
      29          491 :         return nodes.find(nodeId) != nodes.end();
      30              :     }
      31              : 
      32          190 :     bool HasEdge(const NodeId srcNodeId, const NodeId dstNodeId) const
      33              :     {
      34          190 :         return edges.find(srcNodeId) != edges.end() && edges.at(srcNodeId).find(dstNodeId) != edges.at(srcNodeId).end();
      35              :     }
      36              : 
      37          188 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId, const NodeId dstNodeId) const
      38              :     {
      39          188 :         if (!HasEdge(srcNodeId, dstNodeId)) {
      40          114 :             return {};
      41              :         }
      42           74 :         return edges.at(srcNodeId).at(dstNodeId);
      43              :     }
      44              : 
      45            1 :     std::vector<std::shared_ptr<EdgeType>> GetEdges(const NodeId srcNodeId) const
      46              :     {
      47            1 :         if (edges.find(srcNodeId) == edges.end()) {
      48            0 :             return {};
      49              :         }
      50            1 :         std::vector<std::shared_ptr<EdgeType>> nodeEdges;
      51            4 :         for (const auto &it : edges.at(srcNodeId)) {
      52            3 :             nodeEdges.insert(nodeEdges.end(), it.second.begin(), it.second.end());
      53              :         }
      54            1 :         return nodeEdges;
      55            1 :     }
      56              : 
      57           81 :     void TraverseNode(std::function<void(NodeId nodeId, const std::shared_ptr<NodeType> &)> func) const
      58              :     {
      59          427 :         for (auto &node : nodes) {
      60          346 :             func(node.first, node.second);
      61              :         }
      62           81 :     }
      63              : 
      64           24 :     void TraverseNode(std::function<void(std::shared_ptr<NodeType>)> func) const
      65              :     {
      66          148 :         for (auto &node : nodes) {
      67          124 :             func(node.second);
      68              :         }
      69           24 :     }
      70              : 
      71          191 :     void TraverseEdge(const NodeId srcNodeId, std::function<void(std::shared_ptr<EdgeType>)> func) const
      72              :     {
      73          191 :         if (edges.find(srcNodeId) == edges.end()) {
      74            6 :             return;
      75              :         }
      76          535 :         for (auto &srcEdges : edges.at(srcNodeId)) {
      77          742 :             for (auto &edge : srcEdges.second) {
      78          392 :                 func(edge);
      79              :             }
      80              :         }
      81              :     }
      82              : 
      83          546 :     void TraverseEdge(const NodeId srcNodeId, const NodeId dstNodeId,
      84              :                       std::function<void(std::shared_ptr<EdgeType>)> func) const
      85              :     {
      86          546 :         if (edges.find(srcNodeId) == edges.end() || edges.at(srcNodeId).find(dstNodeId) == edges.at(srcNodeId).end()) {
      87          307 :             return;
      88              :         }
      89          502 :         for (auto &edge : edges.at(srcNodeId).at(dstNodeId)) {
      90          263 :             func(edge);
      91              :         }
      92              :     }
      93              : 
      94          158 :     void AddNode(const NodeId nodeId, std::shared_ptr<NodeType> node)
      95              :     {
      96          158 :         nodes[nodeId] = node;
      97          474 :         HCCL_DEBUG("[Graph]add node [%llu] success! node number is [%zu]", nodeId, nodes.size());
      98          158 :     }
      99              : 
     100          394 :     void AddEdge(const NodeId srcNodeId, const NodeId dstNodeId, std::shared_ptr<EdgeType> edge)
     101              :     {
     102          394 :         edges[srcNodeId][dstNodeId].push_back(edge);
     103         1182 :         HCCL_DEBUG("[Graph]add edge from node [%llu] to node [%llu] success! edge number is [%zu]", srcNodeId,
     104              :                    dstNodeId, edges[srcNodeId][dstNodeId].size());
     105          394 :     }
     106              : 
     107           16 :     void DeleteEdge(const NodeId srcNodeId, const NodeId dstNodeId)
     108              :     {
     109              :         // 在64+1中srcNode 和 dstNode代表fabric和db
     110           16 :         auto srcIt = edges.find(srcNodeId);
     111           16 :         if (srcIt == edges.end()) {
     112            8 :             return;
     113              :         }
     114              : 
     115           16 :         auto& dstMap = srcIt->second;
     116           16 :         auto dstIt = dstMap.find(dstNodeId);
     117           16 :         if (dstIt == dstMap.end()) {
     118            8 :             return;
     119              :         }
     120              : 
     121              :         // 删除对应的边
     122            8 :         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