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              :         size_t hashMemType = std::hash<uint32_t>{}(static_cast<int>(mem.type));
      27              :         size_t hashMemaddr = std::hash<void*>{}(mem.addr);
      28              :         size_t hashMemSize = std::hash<uint64_t>{}(mem.size);
      29              :         return hashMemType ^ (hashMemaddr << MEM_ADDR_SHIFT) ^ (hashMemSize << MEM_SIZE_SHIFT);
      30              :     }
      31              : };
      32              : 
      33              : struct CommEngineHash {
      34          152 :     size_t operator()(const CommEngine& engine) const noexcept {
      35          152 :         return std::hash<int>{}(static_cast<int>(engine));
      36              :     }
      37              : };
      38              : 
      39              : class ContextManager {
      40              : public:
      41              :     ContextManager();
      42              :     ~ContextManager();
      43              :     HcclResult CreateCommEngineCtx(const std::string &tag, CommEngine engine, uint64_t size, void **ctx);
      44              :     HcclResult GetCommEngineCtx(const std::string &tag, CommEngine engine, void **ctx, uint64_t *size);
      45              :     HcclResult CopyCommEngineCtx(const std::string &tag, CommEngine engine, const void *srcCtx, uint64_t size,
      46              :         uint64_t dstCtxOffset);
      47              :     HcclResult DestroyCommEngineCtx(const std::string &tag, CommEngine engine);
      48              : private:
      49              :     std::unordered_map<std::string, std::unordered_map<CommEngine, HcclMem, CommEngineHash>> contextMap_;
      50              :     std::mutex mutex_;
      51              : };
      52              : }
      53              : #endif  // INDEPENDENT_OP_CONTEXT_MANAGER_H