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