Star1654
656
代码
代码
Issues133
Pull Requests225
Actions
项目讨论
Wiki
分析
Star1654
656
已合并
fix #704
WangShuying创建于 8月14日
fix #704
已合并
WangShuying创建于 8月14日
从WangShuying5/cann-learning-hub_8419test合入到cann/cann-learning-hubtest
共 11 个文件变更+6227-6039
Rtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/01.04_chapter_practice.ipynb→tutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/01.04_chapter_test.ipynb+1261-1093
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/include/block/block_scheduler_utils.h+48-48
@@ -1,49 +1,49 @@
1-#ifndef BLOCK_SCHEDULER_UTILS_H1+#ifndef BLOCK_SCHEDULER_UTILS_H
2-#define BLOCK_SCHEDULER_UTILS_H2+#define BLOCK_SCHEDULER_UTILS_H
3- 3+ 
4-#include "kernel_operator.h"4+#include "kernel_operator.h"
5- 5+ 
6-namespace MatmulAllGatherImpl {6+namespace MatmulAllGatherImpl {
7- 7+ 
8-constexpr int64_t MNK_M = 0;8+constexpr int64_t MNK_M = 0;
9-constexpr int64_t MNK_N = 1;9+constexpr int64_t MNK_N = 1;
10-constexpr int64_t MNK_K = 2;10+constexpr int64_t MNK_K = 2;
11-constexpr int64_t MNK_M_SPLIT = 3;11+constexpr int64_t MNK_M_SPLIT = 3;
12-constexpr int64_t MNK_N_SPLIT = 4;12+constexpr int64_t MNK_N_SPLIT = 4;
13- 13+ 
14-constexpr int64_t WINDOW_LEN = 4LL;14+constexpr int64_t WINDOW_LEN = 4LL;
15-constexpr int64_t INNER_AXIS_MIN_SPLIT_VAL = 128LL;15+constexpr int64_t INNER_AXIS_MIN_SPLIT_VAL = 128LL;
16- 16+ 
17-template <typename T>17+template <typename T>
18-__aicore__ inline T CeilDiv(T a, T b)18+__aicore__ inline T CeilDiv(T a, T b)
19-{19+{
20- return (a + b - 1) / b;20+ return (a + b - 1) / b;
21-}21+}
22- 22+ 
23-template <typename T>23+template <typename T>
24-__aicore__ inline T CeilAlign(T a, T b)24+__aicore__ inline T CeilAlign(T a, T b)
25-{25+{
26- return CeilDiv(a, b) * b;26+ return CeilDiv(a, b) * b;
27-}27+}
28- 28+ 
29-template <typename T>29+template <typename T>
30-__aicore__ inline T Align(T a, T b)30+__aicore__ inline T Align(T a, T b)
31-{31+{
32- return (a + b - 1) / b * b;32+ return (a + b - 1) / b * b;
33-}33+}
34- 34+ 
35-template <typename T>35+template <typename T>
36-__aicore__ inline T Min(T a, T b)36+__aicore__ inline T Min(T a, T b)
37-{37+{
38- return a < b ? a : b;38+ return a < b ? a : b;
39-}39+}
40- 40+ 
41-template <size_t I, typename Tp>41+template <size_t I, typename Tp>
42-__aicore__ constexpr inline decltype(auto) Get(Tp&& t)42+__aicore__ constexpr inline decltype(auto) Get(Tp&& t)
43-{43+{
44- return AscendC::Std::get<I>(AscendC::Std::forward<Tp>(t));44+ return AscendC::Std::get<I>(AscendC::Std::forward<Tp>(t));
45-}45+}
46- 46+ 
47-}47+}
48- 48+ 
49#endif49#endif
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/include/kernel/matmul_all_gather.h+125-126
@@ -1,127 +1,126 @@
1-#ifndef MATMUL_ALL_GATHER_KERNEL_H1+#ifndef MATMUL_ALL_GATHER_KERNEL_H
2-#define MATMUL_ALL_GATHER_KERNEL_H2+#define MATMUL_ALL_GATHER_KERNEL_H
3- 3+ 
4-#include "utils.h"4+#include "utils.h"
5-#include "mte_comm.h"5+#include "mte_comm.h"
6-#include "include/tensor.h"6+#include "include/tensor.h"
7-#include "block/simple_block_scheduler.h"7+#include "block/simple_block_scheduler.h"
8-#include "block/simple_block_mmad.h"8+#include "block/simple_block_mmad.h"
9- 9+ 
10-namespace MatmulAllGatherImpl {10+namespace MatmulAllGatherImpl {
11- 11+ 
12-using namespace AscendC;12+using namespace AscendC;
13-using namespace AscendC::Te;13+using namespace AscendC::Te;
14- 14+ 
15-template<typename DataType>15+template<typename DataType>
16-class MatmulAllGatherOp {16+class MatmulAllGatherOp {
17-public:17+public:
18- using AType = DataType;18+ using AType = DataType;
19- using BType = DataType;19+ using BType = DataType;
20- using CType = DataType;20+ using CType = DataType;
21- using BlockMmadType = SimpleBlockMmad<AType, BType, CType>;21+ using BlockMmadType = SimpleBlockMmad<AType, BType, CType>;
22- using BlockCoord = AscendC::Coord<int64_t, int64_t>;22+ using BlockCoord = AscendC::Coord<int64_t, int64_t>;
23- 23+ 
24- __aicore__ inline MatmulAllGatherOp() {};24+ __aicore__ inline MatmulAllGatherOp() {};
25- __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR mc2Context, GM_ADDR tilingGM, TPipe *tPipe);25+ __aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR mc2Context, GM_ADDR tilingGM, TPipe *tPipe);
26- __aicore__ inline void Process();26+ __aicore__ inline void Process();
27- 27+ 
28-private:28+private:
29- uint64_t mDim_{0};29+ uint64_t mDim_{0};
30- uint64_t kDim_{0};30+ uint64_t kDim_{0};
31- uint64_t nDim_{0};31+ uint64_t nDim_{0};
32- uint64_t totalWinSize_{0};32+ uint64_t totalWinSize_{0};
33- int64_t baseM_{128};33+ int64_t baseM_{128};
34- int64_t baseK_{128};34+ int64_t baseK_{128};
35- int64_t baseN_{128};35+ int64_t baseN_{128};
36- 36+ 
37- GM_ADDR aGm_;37+ GM_ADDR aGm_;
38- GM_ADDR bGm_;38+ GM_ADDR bGm_;
39- GM_ADDR cGm_;39+ GM_ADDR cGm_;
40- 40+ 
41- __gm__ MatmulAllGatherTilingData* tilingDataOut_;41+ __gm__ MatmulAllGatherTilingData* tilingDataOut_;
42- MTECommunication<DataType> mteComm_;42+ MTECommunication<DataType> mteComm_;
43- 43+ 
44- BlockMmadType mmadOp_;44+ BlockMmadType mmadOp_;
45-};45+};
46- 46+ 
47-template<typename DataType>47+template<typename DataType>
48-__aicore__ inline void MatmulAllGatherOp<DataType>::Init(GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR mc2Context, GM_ADDR tilingGM, TPipe *tPipe)48+__aicore__ inline void MatmulAllGatherOp<DataType>::Init(GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR mc2Context, GM_ADDR tilingGM, TPipe *tPipe)
49-{49+{
50- mteComm_.InitHcclContextByAddr(mc2Context);50+ mteComm_.InitHcclContextByAddr(mc2Context);
51- tilingDataOut_ = (__gm__ MatmulAllGatherTilingData*)tilingGM;51+ tilingDataOut_ = (__gm__ MatmulAllGatherTilingData*)tilingGM;
52- 52+ 
53- auto&& tInfo = tilingDataOut_->tilingInfo;53+ auto&& tInfo = tilingDataOut_->tilingInfo;
54- mDim_ = tInfo.mDim;54+ mDim_ = tInfo.mDim;
55- kDim_ = tInfo.kDim;55+ kDim_ = tInfo.kDim;
56- nDim_ = tInfo.nDim;56+ nDim_ = tInfo.nDim;
57- totalWinSize_ = tInfo.totalWinSize;57+ totalWinSize_ = tInfo.totalWinSize;
58- baseM_ = tInfo.baseM;58+ baseM_ = tInfo.baseM;
59- baseK_ = tInfo.baseK;59+ baseK_ = tInfo.baseK;
60- baseN_ = tInfo.baseN;60+ baseN_ = tInfo.baseN;
61- 61+ 
62- aGm_ = a;62+ aGm_ = a;
63- bGm_ = b;63+ bGm_ = b;
64- cGm_ = y;64+ cGm_ = y;
65- 65+
66- mteComm_.InitBuffer(tPipe);66+ mteComm_.InitBuffer(tPipe);
67- mteComm_.InitGMTensor(y, totalWinSize_);67+ mteComm_.InitGMTensor(y, totalWinSize_);
68- 68+
69- typename BlockMmadType::L1Params l1Params{static_cast<uint64_t>(baseK_), static_cast<uint64_t>(baseK_)};69+ typename BlockMmadType::L1Params l1Params{static_cast<uint64_t>(baseK_), static_cast<uint64_t>(baseK_)};
70- mmadOp_.Init(baseM_, baseN_, baseK_, l1Params);70+ mmadOp_.Init(baseM_, baseN_, baseK_, l1Params);
71- mmadOp_.UpdateK(kDim_);71+ mmadOp_.UpdateK(kDim_);
72-}72+}
73- 73+ 
74-template<typename DataType>74+template<typename DataType>
75-__aicore__ inline void MatmulAllGatherOp<DataType>::Process()75+__aicore__ inline void MatmulAllGatherOp<DataType>::Process()
76-{76+{
77- int64_t m = static_cast<int64_t>(mDim_);77+ int64_t m = static_cast<int64_t>(mDim_);
78- int64_t n = static_cast<int64_t>(nDim_);78+ int64_t n = static_cast<int64_t>(nDim_);
79- int64_t k = static_cast<int64_t>(kDim_);79+ int64_t k = static_cast<int64_t>(kDim_);
80- 80+
81- auto layoutA = AscendC::Te::MakeNDLayout<AType>(m, k);81+ if ASCEND_IS_AIC {
82- auto layoutB = AscendC::Te::MakeNDLayout<BType>(k, n);82+ auto layoutA = AscendC::Te::MakeNDLayout<AType>(m, k);
83- auto layoutC = AscendC::Te::MakeNDLayout<CType>(m, n);83+ auto layoutB = AscendC::Te::MakeNDLayout<BType>(k, n);
84- 84+ auto layoutC = AscendC::Te::MakeNDLayout<CType>(m, n);
85- auto gmA = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ AType*)aGm_), layoutA);85+
86- auto gmB = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ BType*)bGm_), layoutB);86+ auto gmA = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ AType*)aGm_), layoutA);
87- auto gmC = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ CType*)cGm_), layoutC);87+ auto gmB = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ BType*)bGm_), layoutB);
88- // auto gmC = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ CType*)mteComm_.localWinDataAddr_), layoutC);88+ auto gmC = AscendC::Te::MakeTensor(AscendC::Te::MakeGMmemPtr((__gm__ CType*)mteComm_.localWinDataAddr_), layoutC);
89- 89+
90- SimpleBlockScheduler bs(baseM_, baseN_);90+ SimpleBlockScheduler bs(baseM_, baseN_);
91- bs.UpdateNextProblem(m, n);91+ bs.UpdateNextProblem(m, n);
92- 92+
93- BlockCoord tileCoord;93+ BlockCoord tileCoord;
94- while (bs.GetTileIdx(tileCoord)) {94+ while (bs.GetTileIdx(tileCoord)) {
95- SimpleBlockScheduler::BlockShape singleShape = bs.GetBlockShape(tileCoord);95+ SimpleBlockScheduler::BlockShape singleShape = bs.GetBlockShape(tileCoord);
96- 96+
97- int64_t tileM = Get<MNK_M>(singleShape);97+ int64_t tileM = Get<MNK_M>(singleShape);
98- int64_t tileN = Get<MNK_N>(singleShape);98+ int64_t tileN = Get<MNK_N>(singleShape);
99- 99+
100- if (tileM <= 0 || tileN <= 0) {100+ if (tileM <= 0 || tileN <= 0) {
101- continue;101+ continue;
102- }102+ }
103- 103+
104- int64_t mOffset = Get<MNK_M>(tileCoord) * baseM_;104+ int64_t mOffset = Get<MNK_M>(tileCoord) * baseM_;
105- int64_t nOffset = Get<MNK_N>(tileCoord) * baseN_;105+ int64_t nOffset = Get<MNK_N>(tileCoord) * baseN_;
106- 106+
107- auto gmBlockA = gmA(AscendC::Te::MakeCoord(mOffset, 0), AscendC::Te::MakeShape(tileM, k));107+ auto gmBlockA = gmA(AscendC::Te::MakeCoord(mOffset, 0), AscendC::Te::MakeShape(tileM, k));
108- auto gmBlockB = gmB(AscendC::Te::MakeCoord(0, nOffset), AscendC::Te::MakeShape(k, tileN));108+ auto gmBlockB = gmB(AscendC::Te::MakeCoord(0, nOffset), AscendC::Te::MakeShape(k, tileN));
109- auto gmBlockC = gmC(AscendC::Te::MakeCoord(mOffset, nOffset), AscendC::Te::MakeShape(tileM, tileN));109+ auto gmBlockC = gmC(AscendC::Te::MakeCoord(mOffset, nOffset), AscendC::Te::MakeShape(tileM, tileN));
110- 110+
111- typename BlockMmadType::BlockShape mmadShape{tileM, tileN, k};111+ typename BlockMmadType::BlockShape mmadShape{tileM, tileN, k};
112- mmadOp_(gmBlockA, gmBlockB, gmBlockC, mmadShape);112+ mmadOp_(gmBlockA, gmBlockB, gmBlockC, mmadShape);
113- }113+ }
114- 114+ return;
115- if ASCEND_IS_AIC {115+ }
116- return;116+ 
117- }117+ PipeBarrier<PIPE_ALL>();
118- 118+ 
119- // // AllGather 通信119+ mteComm_.WriteStatusToWin();
120- // mteComm_.WriteStatusToWin();120+ mteComm_.ReadStatus();
121- // mteComm_.ReadStatus();121+ mteComm_.CopyDataFromWin(mDim_, nDim_);
122- // mteComm_.CopyDataFromWin(mDim_, nDim_);122+}
123-}123+ 
124- 124+}
125-}125+ 
126- 
127#endif126#endif
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/include/kernel/mte_comm.h+218-218
@@ -1,219 +1,219 @@
1-#ifndef MTE_COMM_H1+#ifndef MTE_COMM_H
2-#define MTE_COMM_H2+#define MTE_COMM_H
3- 3+ 
4-#include "utils.h"4+#include "utils.h"
5-#include "kernel_operator.h"5+#include "kernel_operator.h"
6- 6+ 
7-namespace MatmulAllGatherImpl {7+namespace MatmulAllGatherImpl {
8- 8+ 
9-using namespace AscendC;9+using namespace AscendC;
10- 10+ 
11-constexpr static uint32_t UB_ALIGN_BYTES = 32U;11+constexpr static uint32_t UB_ALIGN_BYTES = 32U;
12-constexpr static uint32_t FLOAT_ALIGN_NUM = 8U;12+constexpr static uint32_t FLOAT_ALIGN_NUM = 8U;
13-constexpr static uint32_t BUFFER_NUM = 2U;13+constexpr static uint32_t BUFFER_NUM = 2U;
14-constexpr static uint64_t STATE_WIN_SIZE = 1024UL * 1024UL;14+constexpr static uint64_t STATE_WIN_SIZE = 1024UL * 1024UL;
15-constexpr static uint64_t SINGLE_STATE_REGION_SIZE = 462UL * 1024UL;15+constexpr static uint64_t SINGLE_STATE_REGION_SIZE = 462UL * 1024UL;
16-constexpr static uint32_t ZERONE_STATE_POS = 0U;16+constexpr static uint32_t ZERONE_STATE_POS = 0U;
17-constexpr static uint64_t WIN_ADDR_ALIGN = 512UL;17+constexpr static uint64_t WIN_ADDR_ALIGN = 512UL;
18- 18+ 
19-template<typename DataType>19+template<typename DataType>
20-class MTECommunication {20+class MTECommunication {
21-public:21+public:
22- __aicore__ inline MTECommunication() {};22+ __aicore__ inline MTECommunication() {};
23- __aicore__ inline void InitHcclContextByAddr(GM_ADDR mc2Context);23+ __aicore__ inline void InitHcclContextByAddr(GM_ADDR mc2Context);
24- __aicore__ inline void InitGMTensor(GM_ADDR y, uint64_t winSpaceGmSize);24+ __aicore__ inline void InitGMTensor(GM_ADDR y, uint64_t winSpaceGmSize);
25- __aicore__ inline void InitBuffer(TPipe *tPipe);25+ __aicore__ inline void InitBuffer(TPipe *tPipe);
26- __aicore__ inline void WriteStatusToWin();26+ __aicore__ inline void WriteStatusToWin();
27- __aicore__ inline void ReadStatus();27+ __aicore__ inline void ReadStatus();
28- __aicore__ inline void CopyDataFromWin(uint64_t perRankM, uint64_t nDim);28+ __aicore__ inline void CopyDataFromWin(uint64_t perRankM, uint64_t nDim);
29- __aicore__ inline GM_ADDR GetWinAddrGm(uint32_t rankId, uint64_t offset = 0);29+ __aicore__ inline GM_ADDR GetWinAddrGm(uint32_t rankId, uint64_t offset = 0);
30- __aicore__ inline GM_ADDR GetWinDataAddrGm(uint32_t rankId, uint32_t winFlag);30+ __aicore__ inline GM_ADDR GetWinDataAddrGm(uint32_t rankId, uint32_t winFlag);
31- __aicore__ inline GM_ADDR GetWinStatusAddrGm(uint32_t rankId, uint32_t winFlag);31+ __aicore__ inline GM_ADDR GetWinStatusAddrGm(uint32_t rankId, uint32_t winFlag);
32- 32+ 
33- __gm__ HcclA5OpResParam *hcclContext_;33+ __gm__ HcclA5OpResParam *hcclContext_;
34- uint32_t aivId_{0};34+ uint32_t aivId_{0};
35- uint32_t winBufferFlags_{0};35+ uint32_t winBufferFlags_{0};
36- uint64_t totalLen_{0};36+ uint64_t totalLen_{0};
37- GlobalTensor<DataType> inputTensor_;37+ GlobalTensor<DataType> inputTensor_;
38- GlobalTensor<DataType> outputTensor_;38+ GlobalTensor<DataType> outputTensor_;
39- GlobalTensor<DataType> localWinOutTensor_;39+ GlobalTensor<DataType> localWinOutTensor_;
40- GM_ADDR localWinDataAddr_;40+ GM_ADDR localWinDataAddr_;
41- 41+ 
42-private:42+private:
43- __aicore__ inline void CopyDataBlockFromWin(uint32_t remoteRankId, uint64_t readOffset, uint64_t curOutOffset, uint32_t count);43+ __aicore__ inline void CopyDataBlockFromWin(uint32_t remoteRankId, uint64_t readOffset, uint64_t curOutOffset, uint32_t count);
44- 44+ 
45- TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> copyQueue_;45+ TQueBind<QuePosition::VECIN, QuePosition::VECOUT, 1> copyQueue_;
46- TBuf<> winFlagsBuf_;46+ TBuf<> winFlagsBuf_;
47- TBuf<> writeStateBuf_;47+ TBuf<> writeStateBuf_;
48- TBuf<> readStateBuf_;48+ TBuf<> readStateBuf_;
49- TBuf<> stateResetBuf_;49+ TBuf<> stateResetBuf_;
50- 50+ 
51- GlobalTensor<uint32_t> selfWinFlagGMTensor_;51+ GlobalTensor<uint32_t> selfWinFlagGMTensor_;
52- LocalTensor<DataType> copyTmpTensor_;52+ LocalTensor<DataType> copyTmpTensor_;
53- LocalTensor<float> stateResetTensor_;53+ LocalTensor<float> stateResetTensor_;
54-};54+};
55- 55+ 
56-template<typename DataType>56+template<typename DataType>
57-__aicore__ inline void MTECommunication<DataType>::InitHcclContextByAddr(GM_ADDR mc2Context)57+__aicore__ inline void MTECommunication<DataType>::InitHcclContextByAddr(GM_ADDR mc2Context)
58-{58+{
59- hcclContext_ = (__gm__ HcclA5OpResParam*)(mc2Context);59+ hcclContext_ = (__gm__ HcclA5OpResParam*)(mc2Context);
60- aivId_ = GetBlockIdx();60+ aivId_ = GetBlockIdx();
61-}61+}
62- 62+ 
63-template<typename DataType>63+template<typename DataType>
64-__aicore__ inline void MTECommunication<DataType>::InitGMTensor(GM_ADDR y, uint64_t winSpaceGmSize)64+__aicore__ inline void MTECommunication<DataType>::InitGMTensor(GM_ADDR y, uint64_t winSpaceGmSize)
65-{65+{
66- outputTensor_.SetGlobalBuffer((__gm__ DataType*)y);66+ outputTensor_.SetGlobalBuffer((__gm__ DataType*)y);
67- 67+ 
68- uint64_t currCoreFlagOffset = 2UL * SINGLE_STATE_REGION_SIZE + aivId_ * WIN_ADDR_ALIGN;68+ uint64_t currCoreFlagOffset = 2UL * SINGLE_STATE_REGION_SIZE + aivId_ * WIN_ADDR_ALIGN;
69- selfWinFlagGMTensor_.SetGlobalBuffer((__gm__ uint32_t*)GetWinAddrGm(hcclContext_->rankId, currCoreFlagOffset));69+ selfWinFlagGMTensor_.SetGlobalBuffer((__gm__ uint32_t*)GetWinAddrGm(hcclContext_->rankId, currCoreFlagOffset));
70- LocalTensor<uint32_t> winFlagLocalTensor = winFlagsBuf_.Get<uint32_t>();70+ LocalTensor<uint32_t> winFlagLocalTensor = winFlagsBuf_.Get<uint32_t>();
71- DataCopyExtParams flagCopyInParams{1, static_cast<uint32_t>(UB_ALIGN_BYTES), 0, 0, 0};71+ DataCopyExtParams flagCopyInParams{1, static_cast<uint32_t>(UB_ALIGN_BYTES), 0, 0, 0};
72- DataCopyPadExtParams<uint32_t> flagPadParams{false, 0, 0, 0};72+ DataCopyPadExtParams<uint32_t> flagPadParams{false, 0, 0, 0};
73- DataCopyPad(winFlagLocalTensor, selfWinFlagGMTensor_, flagCopyInParams, flagPadParams);73+ DataCopyPad(winFlagLocalTensor, selfWinFlagGMTensor_, flagCopyInParams, flagPadParams);
74- SyncFunc<AscendC::HardEvent::MTE2_S>();74+ SyncFunc<AscendC::HardEvent::MTE2_S>();
75- 75+ 
76- winBufferFlags_ = winFlagLocalTensor.GetValue(ZERONE_STATE_POS);76+ winBufferFlags_ = winFlagLocalTensor.GetValue(ZERONE_STATE_POS);
77- winFlagLocalTensor.SetValue(ZERONE_STATE_POS, 1 - winBufferFlags_);77+ winFlagLocalTensor.SetValue(ZERONE_STATE_POS, 1 - winBufferFlags_);
78- 78+ 
79- SyncFunc<AscendC::HardEvent::S_MTE3>();79+ SyncFunc<AscendC::HardEvent::S_MTE3>();
80- DataCopyExtParams flagCopyOutParams{1, static_cast<uint32_t>(UB_ALIGN_BYTES), 0, 0, 0};80+ DataCopyExtParams flagCopyOutParams{1, static_cast<uint32_t>(UB_ALIGN_BYTES), 0, 0, 0};
81- DataCopyPad(selfWinFlagGMTensor_, winFlagLocalTensor, flagCopyOutParams);81+ DataCopyPad(selfWinFlagGMTensor_, winFlagLocalTensor, flagCopyOutParams);
82- 82+ 
83- GM_ADDR localWinDataGm = GetWinDataAddrGm(hcclContext_->rankId, winBufferFlags_);83+ GM_ADDR localWinDataGm = GetWinDataAddrGm(hcclContext_->rankId, winBufferFlags_);
84- localWinDataAddr_ = localWinDataGm;84+ localWinDataAddr_ = localWinDataGm;
85- localWinOutTensor_.SetGlobalBuffer((__gm__ DataType*)localWinDataGm);85+ localWinOutTensor_.SetGlobalBuffer((__gm__ DataType*)localWinDataGm);
86-}86+}
87- 87+ 
88-template<typename DataType>88+template<typename DataType>
89-__aicore__ inline void MTECommunication<DataType>::InitBuffer(TPipe *tPipe)89+__aicore__ inline void MTECommunication<DataType>::InitBuffer(TPipe *tPipe)
90-{90+{
91- constexpr uint32_t COPY_BLOCK_SIZE = 1024U;91+ constexpr uint32_t COPY_BLOCK_SIZE = 1024U;
92- tPipe->InitBuffer(copyQueue_, BUFFER_NUM, COPY_BLOCK_SIZE * sizeof(DataType));92+ tPipe->InitBuffer(copyQueue_, BUFFER_NUM, COPY_BLOCK_SIZE * sizeof(DataType));
93- tPipe->InitBuffer(winFlagsBuf_, UB_ALIGN_BYTES);93+ tPipe->InitBuffer(winFlagsBuf_, UB_ALIGN_BYTES);
94- tPipe->InitBuffer(writeStateBuf_, UB_ALIGN_BYTES);94+ tPipe->InitBuffer(writeStateBuf_, UB_ALIGN_BYTES);
95- tPipe->InitBuffer(readStateBuf_, hcclContext_->rankDim * UB_ALIGN_BYTES);95+ tPipe->InitBuffer(readStateBuf_, hcclContext_->rankDim * UB_ALIGN_BYTES);
96- tPipe->InitBuffer(stateResetBuf_, hcclContext_->rankDim * UB_ALIGN_BYTES);96+ tPipe->InitBuffer(stateResetBuf_, hcclContext_->rankDim * UB_ALIGN_BYTES);
97- 97+ 
98- stateResetTensor_ = stateResetBuf_.Get<float>();98+ stateResetTensor_ = stateResetBuf_.Get<float>();
99- Duplicate<float>(stateResetTensor_, (float)0.0, static_cast<uint32_t>(hcclContext_->rankDim * FLOAT_ALIGN_NUM));99+ Duplicate<float>(stateResetTensor_, (float)0.0, static_cast<uint32_t>(hcclContext_->rankDim * FLOAT_ALIGN_NUM));
100-}100+}
101- 101+ 
102-template<typename DataType>102+template<typename DataType>
103-__aicore__ inline void MTECommunication<DataType>::WriteStatusToWin()103+__aicore__ inline void MTECommunication<DataType>::WriteStatusToWin()
104-{104+{
105- uint32_t coreOffset = aivId_ * hcclContext_->rankDim;105+ uint32_t coreOffset = aivId_ * hcclContext_->rankDim;
106- for (uint32_t curRank = 0; curRank < hcclContext_->rankDim; ++curRank) {106+ for (uint32_t curRank = 0; curRank < hcclContext_->rankDim; ++curRank) {
107- LocalTensor<float> statusTensor = writeStateBuf_.Get<float>();107+ LocalTensor<float> statusTensor = writeStateBuf_.Get<float>();
108- DataCopy<float>(statusTensor, stateResetTensor_, FLOAT_ALIGN_NUM);108+ DataCopy<float>(statusTensor, stateResetTensor_, FLOAT_ALIGN_NUM);
109- SyncFunc<AscendC::HardEvent::MTE2_S>();109+ SyncFunc<AscendC::HardEvent::MTE2_S>();
110- statusTensor.SetValue(0, (float)1.0);110+ statusTensor.SetValue(0, (float)1.0);
111- GM_ADDR remoteWinStateGM = GetWinStatusAddrGm(curRank, winBufferFlags_);111+ GM_ADDR remoteWinStateGM = GetWinStatusAddrGm(curRank, winBufferFlags_);
112- GlobalTensor<float> stateGMTensor;112+ GlobalTensor<float> stateGMTensor;
113- stateGMTensor.SetGlobalBuffer((__gm__ float*)remoteWinStateGM);113+ stateGMTensor.SetGlobalBuffer((__gm__ float*)remoteWinStateGM);
114- uint64_t curOffset = (coreOffset + hcclContext_->rankId) * FLOAT_ALIGN_NUM;114+ uint64_t curOffset = (coreOffset + hcclContext_->rankId) * FLOAT_ALIGN_NUM;
115- SyncFunc<AscendC::HardEvent::S_MTE3>();115+ SyncFunc<AscendC::HardEvent::S_MTE3>();
116- DataCopyExtParams stateCopyOutParams{1, static_cast<uint32_t>(FLOAT_ALIGN_NUM * sizeof(float)), 0, 0, 0};116+ DataCopyExtParams stateCopyOutParams{1, static_cast<uint32_t>(FLOAT_ALIGN_NUM * sizeof(float)), 0, 0, 0};
117- DataCopyPad(stateGMTensor[curOffset], statusTensor, stateCopyOutParams);117+ DataCopyPad(stateGMTensor[curOffset], statusTensor, stateCopyOutParams);
118- SyncFunc<AscendC::HardEvent::MTE3_S>();118+ SyncFunc<AscendC::HardEvent::MTE3_S>();
119- }119+ }
120-}120+}
121- 121+ 
122-template<typename DataType>122+template<typename DataType>
123-__aicore__ inline void MTECommunication<DataType>::ReadStatus()123+__aicore__ inline void MTECommunication<DataType>::ReadStatus()
124-{124+{
125- GM_ADDR stateGM = GetWinStatusAddrGm(hcclContext_->rankId, winBufferFlags_);125+ GM_ADDR stateGM = GetWinStatusAddrGm(hcclContext_->rankId, winBufferFlags_);
126- GlobalTensor<float> selfStatusWinTensor;126+ GlobalTensor<float> selfStatusWinTensor;
127- uint32_t offset = aivId_ * hcclContext_->rankDim * FLOAT_ALIGN_NUM;127+ uint32_t offset = aivId_ * hcclContext_->rankDim * FLOAT_ALIGN_NUM;
128- selfStatusWinTensor.SetGlobalBuffer((__gm__ float*)(stateGM));128+ selfStatusWinTensor.SetGlobalBuffer((__gm__ float*)(stateGM));
129- LocalTensor<float> statusTensor = readStateBuf_.Get<float>();129+ LocalTensor<float> statusTensor = readStateBuf_.Get<float>();
130- float flag = 0;130+ float flag = 0;
131- uint32_t statusCnt = hcclContext_->rankDim * FLOAT_ALIGN_NUM;131+ uint32_t statusCnt = hcclContext_->rankDim * FLOAT_ALIGN_NUM;
132- SumParams sumParams{1, statusCnt, statusCnt};132+ SumParams sumParams{1, statusCnt, statusCnt};
133- float minTarget = hcclContext_->rankDim - (float)0.5;133+ float minTarget = hcclContext_->rankDim - (float)0.5;
134- float maxTarget = hcclContext_->rankDim + (float)0.5;134+ float maxTarget = hcclContext_->rankDim + (float)0.5;
135- while ((flag < minTarget) || (flag > maxTarget)) {135+ while ((flag < minTarget) || (flag > maxTarget)) {
136- SyncFunc<AscendC::HardEvent::S_MTE2>();136+ SyncFunc<AscendC::HardEvent::S_MTE2>();
137- DataCopyExtParams statusCopyInParams{1, static_cast<uint32_t>(statusCnt * sizeof(float)), 0, 0, 0};137+ DataCopyExtParams statusCopyInParams{1, static_cast<uint32_t>(statusCnt * sizeof(float)), 0, 0, 0};
138- DataCopyPadExtParams<float> statusPadParams{false, 0, 0, 0};138+ DataCopyPadExtParams<float> statusPadParams{false, 0, 0, 0};
139- DataCopyPad(statusTensor, selfStatusWinTensor[offset], statusCopyInParams, statusPadParams);139+ DataCopyPad(statusTensor, selfStatusWinTensor[offset], statusCopyInParams, statusPadParams);
140- SyncFunc<AscendC::HardEvent::MTE2_V>();140+ SyncFunc<AscendC::HardEvent::MTE2_V>();
141- Sum(statusTensor, statusTensor, sumParams);141+ Sum(statusTensor, statusTensor, sumParams);
142- SyncFunc<AscendC::HardEvent::V_S>();142+ SyncFunc<AscendC::HardEvent::V_S>();
143- flag = statusTensor.GetValue(0);143+ flag = statusTensor.GetValue(0);
144- }144+ }
145- SyncFunc<AscendC::HardEvent::S_MTE3>();145+ SyncFunc<AscendC::HardEvent::S_MTE3>();
146- DataCopyExtParams statusCopyOutParams{1, static_cast<uint32_t>(statusCnt * sizeof(float)), 0, 0, 0};146+ DataCopyExtParams statusCopyOutParams{1, static_cast<uint32_t>(statusCnt * sizeof(float)), 0, 0, 0};
147- DataCopyPad(selfStatusWinTensor[offset], stateResetTensor_, statusCopyOutParams);147+ DataCopyPad(selfStatusWinTensor[offset], stateResetTensor_, statusCopyOutParams);
148-}148+}
149- 149+ 
150-template<typename DataType>150+template<typename DataType>
151-__aicore__ inline void MTECommunication<DataType>::CopyDataBlockFromWin(uint32_t remoteRankId, uint64_t readOffset, uint64_t curOutOffset, uint32_t count)151+__aicore__ inline void MTECommunication<DataType>::CopyDataBlockFromWin(uint32_t remoteRankId, uint64_t readOffset, uint64_t curOutOffset, uint32_t count)
152-{152+{
153- GM_ADDR remoteWinDataGm = GetWinDataAddrGm(remoteRankId, winBufferFlags_);153+ GM_ADDR remoteWinDataGm = GetWinDataAddrGm(remoteRankId, winBufferFlags_);
154- GlobalTensor<DataType> remoteWinXTensor;154+ GlobalTensor<DataType> remoteWinXTensor;
155- remoteWinXTensor.SetGlobalBuffer((__gm__ DataType*)remoteWinDataGm);155+ remoteWinXTensor.SetGlobalBuffer((__gm__ DataType*)remoteWinDataGm);
156- 156+ 
157- copyTmpTensor_ = copyQueue_.AllocTensor<DataType>();157+ copyTmpTensor_ = copyQueue_.AllocTensor<DataType>();
158- DataCopyExtParams copyInParams{1, static_cast<uint32_t>(count * sizeof(DataType)), 0, 0, 0};158+ DataCopyExtParams copyInParams{1, static_cast<uint32_t>(count * sizeof(DataType)), 0, 0, 0};
159- DataCopyPadExtParams<DataType> padParams{false, 0, 0, 0};159+ DataCopyPadExtParams<DataType> padParams{false, 0, 0, 0};
160- DataCopyPad(copyTmpTensor_, remoteWinXTensor[readOffset], copyInParams, padParams);160+ DataCopyPad(copyTmpTensor_, remoteWinXTensor[readOffset], copyInParams, padParams);
161- copyQueue_.EnQue(copyTmpTensor_);161+ copyQueue_.EnQue(copyTmpTensor_);
162- copyTmpTensor_ = copyQueue_.DeQue<DataType>();162+ copyTmpTensor_ = copyQueue_.DeQue<DataType>();
163- DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(count * sizeof(DataType)), 0, 0, 0};163+ DataCopyExtParams copyOutParams{1, static_cast<uint32_t>(count * sizeof(DataType)), 0, 0, 0};
164- DataCopyPad(outputTensor_[curOutOffset], copyTmpTensor_, copyOutParams);164+ DataCopyPad(outputTensor_[curOutOffset], copyTmpTensor_, copyOutParams);
165- copyQueue_.FreeTensor(copyTmpTensor_);165+ copyQueue_.FreeTensor(copyTmpTensor_);
166-}166+}
167- 167+ 
168-template<typename DataType>168+template<typename DataType>
169-__aicore__ inline void MTECommunication<DataType>::CopyDataFromWin(uint64_t perRankM, uint64_t nDim)169+__aicore__ inline void MTECommunication<DataType>::CopyDataFromWin(uint64_t perRankM, uint64_t nDim)
170-{170+{
171- constexpr uint32_t COPY_BLOCK_SIZE = 1024U;171+ constexpr uint32_t COPY_BLOCK_SIZE = 1024U;
172- uint64_t perRankOutSize = perRankM * nDim;172+ uint64_t perRankOutSize = perRankM * nDim;
173- uint64_t totalOutSize = perRankOutSize * hcclContext_->rankDim;173+ uint64_t totalOutSize = perRankOutSize * hcclContext_->rankDim;
174- 174+
175- for (uint32_t i = 0; i < hcclContext_->rankDim; ++i) {175+ for (uint32_t i = 0; i < hcclContext_->rankDim; ++i) {
176- uint32_t remoteRankId = (hcclContext_->rankId + i) % hcclContext_->rankDim;176+ uint32_t remoteRankId = (hcclContext_->rankId + i) % hcclContext_->rankDim;
177- uint64_t outBaseOffset = remoteRankId * perRankOutSize;177+ uint64_t outBaseOffset = remoteRankId * perRankOutSize;
178- 178+ 
179- uint64_t remainSize = perRankOutSize;179+ uint64_t remainSize = perRankOutSize;
180- uint64_t curOffset = 0;180+ uint64_t curOffset = 0;
181- while (remainSize > 0) {181+ while (remainSize > 0) {
182- uint32_t copySize = static_cast<uint32_t>(Min(remainSize, COPY_BLOCK_SIZE));182+ uint32_t copySize = static_cast<uint32_t>(Min(remainSize, COPY_BLOCK_SIZE));
183- CopyDataBlockFromWin(remoteRankId, curOffset, outBaseOffset + curOffset, copySize);183+ CopyDataBlockFromWin(remoteRankId, curOffset, outBaseOffset + curOffset, copySize);
184- curOffset += copySize;184+ curOffset += copySize;
185- remainSize -= copySize;185+ remainSize -= copySize;
186- }186+ }
187- }187+ }
188- PipeBarrier<PIPE_ALL>();188+ PipeBarrier<PIPE_ALL>();
189-}189+}
190- 190+ 
191-template<typename DataType>191+template<typename DataType>
192-__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinAddrGm(uint32_t rankId, uint64_t offset)192+__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinAddrGm(uint32_t rankId, uint64_t offset)
193-{193+{
194- return (GM_ADDR)(hcclContext_->windowsIn[rankId] + offset);194+ return (GM_ADDR)(hcclContext_->windowsIn[rankId] + offset);
195-}195+}
196- 196+ 
197-template<typename DataType>197+template<typename DataType>
198-__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinDataAddrGm(uint32_t rankId, uint32_t winFlag)198+__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinDataAddrGm(uint32_t rankId, uint32_t winFlag)
199-{199+{
200- if (winFlag == 0U) {200+ if (winFlag == 0U) {
201- return GetWinAddrGm(rankId, STATE_WIN_SIZE);201+ return GetWinAddrGm(rankId, STATE_WIN_SIZE);
202- } else {202+ } else {
203- return (GM_ADDR)(hcclContext_->windowsOut[rankId]);203+ return (GM_ADDR)(hcclContext_->windowsOut[rankId]);
204- }204+ }
205-}205+}
206- 206+ 
207-template<typename DataType>207+template<typename DataType>
208-__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinStatusAddrGm(uint32_t rankId, uint32_t winFlag)208+__aicore__ inline GM_ADDR MTECommunication<DataType>::GetWinStatusAddrGm(uint32_t rankId, uint32_t winFlag)
209-{209+{
210- if (winFlag == 0U) {210+ if (winFlag == 0U) {
211- return GetWinAddrGm(rankId);211+ return GetWinAddrGm(rankId);
212- } else {212+ } else {
213- return GetWinAddrGm(rankId, SINGLE_STATE_REGION_SIZE);213+ return GetWinAddrGm(rankId, SINGLE_STATE_REGION_SIZE);
214- }214+ }
215-}215+}
216- 216+ 
217-}217+}
218- 218+ 
219#endif219#endif
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/include/tiling_data.h+43-43
@@ -1,44 +1,44 @@
1-#ifndef TILING_DATA_H1+#ifndef TILING_DATA_H
2-#define TILING_DATA_H2+#define TILING_DATA_H
3- 3+ 
4-#include <stdint.h>4+#include <stdint.h>
5-#include <kernel_tiling/kernel_tiling.h>5+#include <kernel_tiling/kernel_tiling.h>
6- 6+ 
7-namespace MatmulAllGatherImpl {7+namespace MatmulAllGatherImpl {
8- 8+ 
9-constexpr static int64_t MAX_RANK_NUM = 64;9+constexpr static int64_t MAX_RANK_NUM = 64;
10- 10+ 
11-struct HcclA5OpResParam {11+struct HcclA5OpResParam {
12- uint64_t workSpace;12+ uint64_t workSpace;
13- uint64_t workSpaceSize;13+ uint64_t workSpaceSize;
14- uint32_t rankId;14+ uint32_t rankId;
15- uint32_t rankDim;15+ uint32_t rankDim;
16- uint64_t winSize;16+ uint64_t winSize;
17- uint64_t windowsIn[MAX_RANK_NUM];17+ uint64_t windowsIn[MAX_RANK_NUM];
18- uint64_t windowsOut[MAX_RANK_NUM];18+ uint64_t windowsOut[MAX_RANK_NUM];
19- uint64_t xnAddr;19+ uint64_t xnAddr;
20- uint64_t ckeAddr;20+ uint64_t ckeAddr;
21- uint64_t msAddr;21+ uint64_t msAddr;
22- uint64_t msSize;22+ uint64_t msSize;
23-};23+};
24- 24+ 
25-struct MatmulAllGatherTilingInfo {25+struct MatmulAllGatherTilingInfo {
26- uint64_t mDim; // M 维度（每个 rank 计算完整 M × K × N）26+ uint64_t mDim; // M 维度（每个 rank 计算完整 M × K × N）
27- uint64_t kDim;27+ uint64_t kDim;
28- uint64_t nDim;28+ uint64_t nDim;
29- uint64_t aivNum;29+ uint64_t aivNum;
30- uint64_t totalWinSize;30+ uint64_t totalWinSize;
31- uint64_t baseM;31+ uint64_t baseM;
32- uint64_t baseK;32+ uint64_t baseK;
33- uint64_t baseN;33+ uint64_t baseN;
34-};34+};
35- 35+ 
36-struct MatmulAllGatherTilingData {36+struct MatmulAllGatherTilingData {
37- Mc2InitTiling mc2InitTiling;37+ Mc2InitTiling mc2InitTiling;
38- Mc2CcTiling mc2CcTiling;38+ Mc2CcTiling mc2CcTiling;
39- MatmulAllGatherTilingInfo tilingInfo;39+ MatmulAllGatherTilingInfo tilingInfo;
40-};40+};
41- 41+ 
42-}42+}
43- 43+ 
44#endif44#endif
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/include/utils.h+50-50
@@ -1,51 +1,51 @@
1-#ifndef UTILS_H1+#ifndef UTILS_H
2-#define UTILS_H2+#define UTILS_H
3- 3+ 
4-#include "tiling_data.h"4+#include "tiling_data.h"
5-#include "kernel_operator.h"5+#include "kernel_operator.h"
6- 6+ 
7-namespace MatmulAllGatherImpl {7+namespace MatmulAllGatherImpl {
8- 8+ 
9-using namespace AscendC;9+using namespace AscendC;
10- 10+ 
11-constexpr static int64_t L0A_SIZE = 64 * 1024;11+constexpr static int64_t L0A_SIZE = 64 * 1024;
12-constexpr static int64_t L0B_SIZE = 64 * 1024;12+constexpr static int64_t L0B_SIZE = 64 * 1024;
13-constexpr static int64_t L0C_SIZE = 256 * 1024;13+constexpr static int64_t L0C_SIZE = 256 * 1024;
14-constexpr static int64_t L1_SIZE = 512 * 1024;14+constexpr static int64_t L1_SIZE = 512 * 1024;
15- 15+ 
16-constexpr static uint32_t CUBE_BASE_M = 16;16+constexpr static uint32_t CUBE_BASE_M = 16;
17-constexpr static uint32_t CUBE_BASE_K = 16;17+constexpr static uint32_t CUBE_BASE_K = 16;
18-constexpr static uint32_t CUBE_BASE_N = 16;18+constexpr static uint32_t CUBE_BASE_N = 16;
19- 19+ 
20-__aicore__ inline uint64_t CeilDiv(uint64_t a, uint64_t b)20+__aicore__ inline uint64_t CeilDiv(uint64_t a, uint64_t b)
21-{21+{
22- if (b == 0) {22+ if (b == 0) {
23- return 0;23+ return 0;
24- }24+ }
25- return (a + b - 1) / b;25+ return (a + b - 1) / b;
26-}26+}
27- 27+ 
28-__aicore__ inline uint64_t FloorAlign(uint64_t a, uint64_t b)28+__aicore__ inline uint64_t FloorAlign(uint64_t a, uint64_t b)
29-{29+{
30- if (b == 0) {30+ if (b == 0) {
31- return a;31+ return a;
32- }32+ }
33- return (a / b) * b;33+ return (a / b) * b;
34-}34+}
35- 35+ 
36-__aicore__ inline uint64_t Min(uint64_t a, uint64_t b)36+__aicore__ inline uint64_t Min(uint64_t a, uint64_t b)
37-{37+{
38- return (a < b) ? a : b;38+ return (a < b) ? a : b;
39-}39+}
40- 40+ 
41-template<AscendC::HardEvent event>41+template<AscendC::HardEvent event>
42-__aicore__ inline void SyncFunc()42+__aicore__ inline void SyncFunc()
43-{43+{
44- AscendC::TEventID eventID = GetTPipePtr()->FetchEventID<event>();44+ AscendC::TEventID eventID = GetTPipePtr()->FetchEventID<event>();
45- AscendC::SetFlag<event>(eventID);45+ AscendC::SetFlag<event>(eventID);
46- AscendC::WaitFlag<event>(eventID);46+ AscendC::WaitFlag<event>(eventID);
47-}47+}
48- 48+ 
49-}49+}
50- 50+ 
51#endif51#endif
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/src/matmul_all_gather.cpp+21-21
@@ -1,22 +1,22 @@
1-#include "kernel_operator.h"1+#include "kernel_operator.h"
2-#include "kernel/matmul_all_gather.h"2+#include "kernel/matmul_all_gather.h"
3- 3+ 
4-using namespace AscendC;4+using namespace AscendC;
5-using namespace MatmulAllGatherImpl;5+using namespace MatmulAllGatherImpl;
6- 6+ 
7-extern "C" __global__ __aicore__ void MatmulAllGather_Generic(7+extern "C" __global__ __aicore__ void MatmulAllGather_Generic(
8- GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR workspaceGM, GM_ADDR mc2Context, GM_ADDR tilingGM)8+ GM_ADDR a, GM_ADDR b, GM_ADDR y, GM_ADDR workspaceGM, GM_ADDR mc2Context, GM_ADDR tilingGM)
9-{9+{
10- KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);10+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
11- TPipe pipe;11+ TPipe pipe;
12- MatmulAllGatherOp<bfloat16_t> op;12+ MatmulAllGatherOp<bfloat16_t> op;
13- op.Init(a, b, y, mc2Context, tilingGM, &pipe);13+ op.Init(a, b, y, mc2Context, tilingGM, &pipe);
14- op.Process();14+ op.Process();
15-}15+}
16- 16+ 
17-extern "C" void matmul_all_gather_demo(uint32_t blockDim, void* stream,17+extern "C" void matmul_all_gather_demo(uint32_t blockDim, void* stream,
18- uint8_t* a, uint8_t* b, uint8_t* y, uint8_t* workspaceGM, uint8_t* mc2Context, uint8_t* tilingGM)18+ uint8_t* a, uint8_t* b, uint8_t* y, uint8_t* workspaceGM, uint8_t* mc2Context, uint8_t* tilingGM)
19-{19+{
20- MatmulAllGather_Generic<<<blockDim, nullptr, stream>>>(20+ MatmulAllGather_Generic<<<blockDim, nullptr, stream>>>(
21- a, b, y, workspaceGM, mc2Context, tilingGM);21+ a, b, y, workspaceGM, mc2Context, tilingGM);
22}22}
Mtutorials/MC2_fused_operator_development/01_MC2_basic_fused_operator_development/answer/01.04_answer/src/test_matmul_all_gather.cpp+326-326
@@ -1,327 +1,327 @@
1-#include <iostream>1+#include <iostream>
2-#include <vector>2+#include <vector>
3-#include <string>3+#include <string>
4-#include <cstring>4+#include <cstring>
5-#include <getopt.h>5+#include <getopt.h>
6-#include <cmath>6+#include <cmath>
7-#include <cstdlib>7+#include <cstdlib>
8-#include "hccl/hccl.h"8+#include "hccl/hccl.h"
9-#include "acl/acl.h"9+#include "acl/acl.h"
10-#include "hccl/hcom.h"10+#include "hccl/hcom.h"
11-#include "hccl/hccl_comm.h"11+#include "hccl/hccl_comm.h"
12-#include "securec.h"12+#include "securec.h"
13-#include "tiling/hccl/hccl_tiling.h"13+#include "tiling/hccl/hccl_tiling.h"
14-#include "tiling_data.h"14+#include "tiling_data.h"
15- 15+ 
16-using namespace std;16+using namespace std;
17-using namespace MatmulAllGatherImpl;17+using namespace MatmulAllGatherImpl;
18- 18+ 
19-static constexpr size_t BF16_SIZE = 2;19+static constexpr size_t BF16_SIZE = 2;
20- 20+ 
21-#define CHECK_RET(cond, return_expr) \21+#define CHECK_RET(cond, return_expr) \
22- do { \22+ do { \
23- if (!(cond)) { \23+ if (!(cond)) { \
24- return_expr; \24+ return_expr; \
25- } \25+ } \
26- } while (0)26+ } while (0)
27- 27+ 
28-#define LOG_PRINT(message, ...) \28+#define LOG_PRINT(message, ...) \
29- do { \29+ do { \
30- printf(message, ##__VA_ARGS__); \30+ printf(message, ##__VA_ARGS__); \
31- } while (0)31+ } while (0)
32- 32+ 
33-constexpr int64_t BLOCK_NUM = 64;33+constexpr int64_t BLOCK_NUM = 64;
34-int streamWithTimeout = 10000;34+int streamWithTimeout = 10000;
35- 35+ 
36-int gRankId = 0;36+int gRankId = 0;
37-int gRankSize = 2;37+int gRankSize = 2;
38-int64_t gM = 1024;38+int64_t gM = 1024;
39-int64_t gK = 1024;39+int64_t gK = 1024;
40-int64_t gN = 1024;40+int64_t gN = 1024;
41-int64_t gWorkSpaceSize = 256 * 1024 * 1024;41+int64_t gWorkSpaceSize = 256 * 1024 * 1024;
42-int64_t gHcclBufferSize = 0; // 动态计算42+int64_t gHcclBufferSize = 0; // 动态计算
43-std::string gOutputDir = ".";43+std::string gOutputDir = ".";
44- 44+ 
45-extern "C" void matmul_all_gather_demo(uint32_t blockDim, void* stream,45+extern "C" void matmul_all_gather_demo(uint32_t blockDim, void* stream,
46- uint8_t* a, uint8_t* b, uint8_t* y, uint8_t* workspaceGM, uint8_t* mc2Context, uint8_t* tilingGM);46+ uint8_t* a, uint8_t* b, uint8_t* y, uint8_t* workspaceGM, uint8_t* mc2Context, uint8_t* tilingGM);
47- 47+ 
48-extern "C" HcclResult HcclAllocComResourceByTiling(HcclComm comm, void *stream, void *mc2Tiling, void **commContext);48+extern "C" HcclResult HcclAllocComResourceByTiling(HcclComm comm, void *stream, void *mc2Tiling, void **commContext);
49- 49+ 
50-void GenerateRandomBF16Data(uint16_t* data, int64_t count, int seed, bool allOne = false)50+void GenerateRandomBF16Data(uint16_t* data, int64_t count, int seed, bool allOne = false)
51-{51+{
52- srand(seed);52+ srand(seed);
53- for (int64_t i = 0; i < count; ++i) {53+ for (int64_t i = 0; i < count; ++i) {
54- float val;54+ float val;
55- if (allOne) {55+ if (allOne) {
56- val = 1.0f;56+ val = 1.0f;
57- } else {57+ } else {
58- val = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;58+ val = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;
59- }59+ }
60- uint32_t x;60+ uint32_t x;
61- memcpy(&x, &val, sizeof(float));61+ memcpy(&x, &val, sizeof(float));
62- uint16_t bf16 = (x >> 16) & 0xFFFF;62+ uint16_t bf16 = (x >> 16) & 0xFFFF;
63- data[i] = bf16;63+ data[i] = bf16;
64- }64+ }
65-}65+}
66- 66+ 
67-int WriteFile(const std::string &filePath, const void *buffer, size_t size)67+int WriteFile(const std::string &filePath, const void *buffer, size_t size)
68-{68+{
69- if (buffer == nullptr) {69+ if (buffer == nullptr) {
70- printf("Write file failed. Buffer is nullptr.\n");70+ printf("Write file failed. Buffer is nullptr.\n");
71- return -1;71+ return -1;
72- }72+ }
73- FILE* fp = fopen(filePath.c_str(), "wb");73+ FILE* fp = fopen(filePath.c_str(), "wb");
74- if (fp == nullptr) {74+ if (fp == nullptr) {
75- printf("Open file failed. path = %s\n", filePath.c_str());75+ printf("Open file failed. path = %s\n", filePath.c_str());
76- return -1;76+ return -1;
77- }77+ }
78- size_t written = fwrite(buffer, 1, size, fp);78+ size_t written = fwrite(buffer, 1, size, fp);
79- if (written != size) {79+ if (written != size) {
80- fclose(fp);80+ fclose(fp);
81- return -1;81+ return -1;
82- }82+ }
83- fclose(fp);83+ fclose(fp);
84- return 0;84+ return 0;
85-}85+}
86- 86+ 
87-int CreateTilingDataAndContext(const char* hcomName, aclrtStream stream,87+int CreateTilingDataAndContext(const char* hcomName, aclrtStream stream,
88- void **deviceTilingAddr, void **deviceContextAddr)88+ void **deviceTilingAddr, void **deviceContextAddr)
89-{89+{
90- MatmulAllGatherTilingData *tilingData = new MatmulAllGatherTilingData();90+ MatmulAllGatherTilingData *tilingData = new MatmulAllGatherTilingData();
91- if (tilingData == nullptr) {91+ if (tilingData == nullptr) {
92- LOG_PRINT("[ERROR] tilingData is nullptr\n");92+ LOG_PRINT("[ERROR] tilingData is nullptr\n");
93- return -1;93+ return -1;
94- }94+ }
95- *tilingData = {};95+ *tilingData = {};
96- 96+ 
97- if (gHcclBufferSize <= 0) {97+ if (gHcclBufferSize <= 0) {
98- int64_t winDataSizeBytes = gM * gN * BF16_SIZE;98+ int64_t winDataSizeBytes = gM * gN * BF16_SIZE;
99- int64_t winDataSizeKB = (winDataSizeBytes + 1023) / 1024;99+ int64_t winDataSizeKB = (winDataSizeBytes + 1023) / 1024;
100- gHcclBufferSize = winDataSizeKB + 512; // Win 数据区 + 额外空间(状态区等)100+ gHcclBufferSize = winDataSizeKB + 1024; // Win 数据区 + 状态区(STATE_WIN_SIZE=1MB)