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_CCU_CONTEXT_ALG_BASE_H_
      12              : #define HCCLV2_CCU_CONTEXT_ALG_BASE_H_
      13              : 
      14              : #include <vector>
      15              : #include <ios>
      16              : #include "log.h"
      17              : #include "ccu_ctx.h"
      18              : #include "ccu_datatype.h"
      19              : 
      20              : namespace Hccl {
      21              : 
      22              : class CcuContextAlgBase : public CcuContext {
      23              : public:
      24            5 :     CcuContextAlgBase(
      25              :         const CcuCtxArg& arg, const std::vector<CcuTransport*>& transports, const CcuTransportGroup& group)
      26            5 :         : CcuContext(arg, transports, group)
      27            5 :     {}
      28            9 :     ~CcuContextAlgBase() override {}
      29              : 
      30              :     // 子类实现
      31              :     void Algorithm() override = 0;
      32              :     std::vector<uint64_t> GeneArgs(const CcuTaskArg& arg) override = 0;
      33              : 
      34              : protected:
      35              :     uint64_t rankSize_{0};
      36              :     uint32_t rankId_{0};
      37              :     ReduceOp reduceOp_;
      38              : 
      39              :     struct GroupOpSizeV2 {
      40              :         CcuRep::Variable baseIterNum;
      41              :         CcuRep::Variable tailLoopId;
      42              :         CcuRep::Variable tail;
      43              :     };
      44              : 
      45              :     GroupOpSizeV2 CreateGroupOpSizeV2() { return GroupOpSizeV2{CreateVariable(), CreateVariable(), CreateVariable()}; }
      46              : 
      47              :     void GroupBroadcastV2(
      48              :         const std::vector<CcuTransport*>& transports, const std::vector<CcuRep::Memory>& dst, const CcuRep::Memory& src,
      49              :         const GroupOpSizeV2& goSize) const;
      50              :     void GroupReduceV2(
      51              :         const std::vector<CcuTransport*>& transports, const CcuRep::Memory& dst, const std::vector<CcuRep::Memory>& src,
      52              :         const GroupOpSizeV2& goSize, DataType dataType, DataType outputDataType, ReduceOp reduceType) const;
      53              : };
      54              : } // namespace Hccl
      55              : 
      56              : #endif // HCCLV2_CCU_CONTEXT_ALG_BASE_H_