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 HCCLV2_OFFLOAD_STREAM_MANAGER_H
      12              : #define HCCLV2_OFFLOAD_STREAM_MANAGER_H
      13              : 
      14              : #include <vector>
      15              : #include <string>
      16              : #include <memory>
      17              : #include <unordered_map>
      18              : #include <unordered_set>
      19              : #include "hccl/base.h"
      20              : #include "stream.h"
      21              : 
      22              : namespace Hccl {
      23              : 
      24              : //可继续优化为公共数据结构
      25              : template <typename T>
      26              : class CountSet {
      27              : private:
      28              :     // 底层存储：键=唯一元素，值=计数
      29              :     using MapType = std::unordered_map<T, int>;
      30              :     MapType count_map;
      31              : public:
      32              :     using iterator = typename MapType::iterator;
      33              :     // 1. 添加元素（计数+1）
      34              :     // 返回std::pair<iterator, bool>,与std::set::insert返回值语义完全一致
      35            1 :     std::pair<iterator, bool> insert(const T& elem) {
      36            1 :         auto it = count_map.find(elem);
      37            1 :         bool inserted = false;
      38            1 :         if (it == count_map.end()) {
      39              :             // 首次插入，计数初始化为1
      40            1 :             it = count_map.emplace(elem, 1).first;
      41            1 :             inserted = true;
      42              :         } else {
      43              :             // 元素已存在，计数+1
      44            0 :             it->second++;
      45            0 :             inserted = false;
      46              :         }
      47            1 :         return {it, inserted};
      48              :     }
      49              : 
      50              :     // 2. 删除元素（计数-1，计数为0时移除该元素）
      51              :     // 返回值：删除后剩余的计数（-1表示元素不存在）
      52            0 :     int erase(const T& elem) {
      53            0 :         auto it = count_map.find(elem);
      54            0 :         if (it == count_map.end()) {
      55            0 :             return -1;  // 元素不存在
      56              :         }
      57            0 :         it->second--;   // 计数-1
      58            0 :         if (it->second == 0) {
      59            0 :             count_map.erase(it);  // 计数为0，移除键，避免枚举到空元素
      60            0 :             return 0;
      61              :         }
      62            0 :         return it->second;
      63              :     }
      64              : };
      65              : 
      66              : class OffloadStreamManager {
      67              : public:
      68              :     void RegisterMaster(const std::string &opTag, std::unique_ptr<Stream> stream);
      69              :     void RegisterSlaves(const std::string &opTag, const std::vector<void *> &slaveStreams);
      70              : 
      71              :     void Unregister(const std::string &opTag);
      72              : 
      73              :     Stream *GetMaster(const std::string &opTag);
      74              : 
      75              :     Stream *GetSlave(const std::string &opTag);
      76              : 
      77              :     void ResetIndex(const std::string &opTag, u32 index);
      78              : 
      79              :     u32 GetSlaveIndex(const std::string &opTag) const;
      80              : 
      81              :     Stream *GetSlave(const std::string &opTag, u32 index) const;
      82              :     HcclResult ClearOpStream(const std::string &opTag);
      83              : 
      84              : private:
      85              :     void ActivateSlaveStreams(const std::string &opTag, const Stream *masterStream);
      86              :     void CheckOpTag(const std::string &opTag) const;
      87              : 
      88              :     std::unordered_map<std::string, std::unique_ptr<Stream>>              masters;
      89              :     std::unordered_map<std::string, std::vector<std::unique_ptr<Stream>>> slaves;
      90              :     u32                                                                   slaveIndex{0};
      91              :     std::string                                                           currOpTag{""};
      92              :     std::unordered_map<u32, CountSet<u32>> streamActiveManager_{}; // set中存放当前进程中以已由hccl激活的stream
      93              : };
      94              : 
      95              : } // namespace Hccl
      96              : 
      97              : #endif // HCCLV2_OFFLOAD_STREAM_MANAGER_H