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 INDEPENDENT_OP_CONTEXT_MANAGER_H
      12              : #define INDEPENDENT_OP_CONTEXT_MANAGER_H
      13              : 
      14              : #include "hccl/hccl_res.h"
      15              : #include "hccl_mem_defs.h"
      16              : #include <functional>
      17              : #include <unordered_map>
      18              : #include <mutex>
      19              : #include <string>
      20              : 
      21              : namespace hccl {
      22              : struct HcclMemHash {
      23              :     static constexpr size_t MEM_ADDR_SHIFT = 1;
      24              :     static constexpr size_t MEM_SIZE_SHIFT = 2;
      25              :     size_t operator()(const HcclMem& mem) const noexcept
      26              :     {
      27              :         size_t hashMemType = std::hash<uint32_t>{}(static_cast<int>(mem.type));
      28              :         size_t hashMemaddr = std::hash<void*>{}(mem.addr);
      29              :         size_t hashMemSize = std::hash<uint64_t>{}(mem.size);
      30              :         return hashMemType ^ (hashMemaddr << MEM_ADDR_SHIFT) ^ (hashMemSize << MEM_SIZE_SHIFT);
      31              :     }
      32              : };
      33              : 
      34              : struct CommEngineHash {
      35          152 :     size_t operator()(const CommEngine& engine) const noexcept { return std::hash<int>{}(static_cast<int>(engine)); }
      36              : };
      37              : 
      38              : class ContextManager {
      39              : public:
      40              :     ContextManager();
      41              :     ~ContextManager();
      42              :     HcclResult CreateCommEngineCtx(const std::string& tag, CommEngine engine, uint64_t size, void** ctx);
      43              :     HcclResult GetCommEngineCtx(const std::string& tag, CommEngine engine, void** ctx, uint64_t* size);
      44              :     HcclResult CopyCommEngineCtx(
      45              :         const std::string& tag, CommEngine engine, const void* srcCtx, uint64_t size, uint64_t dstCtxOffset);
      46              :     HcclResult DestroyCommEngineCtx(const std::string& tag, CommEngine engine);
      47              : 
      48              : private:
      49              :     std::unordered_map<std::string, std::unordered_map<CommEngine, HcclMem, CommEngineHash>> contextMap_;
      50              :     std::mutex mutex_;
      51              : };
      52              : } // namespace hccl
      53              : #endif // INDEPENDENT_OP_CONTEXT_MANAGER_H