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