Star627
210
代码
代码
Issues16
Pull Requests32
Actions
项目讨论
Wiki
分析
Star627
210
已合并
feat: flash_attn_lite_story v0 / kv_rms_norm_rope_cache_story / simt_histogram_story SetFlag/WaitFlag 改造为 Mutex 同步 #396
qq_51867290创建于 17 天前
feat: flash_attn_lite_story v0 / kv_rms_norm_rope_cache_story / simt_histogram_story SetFlag/WaitFlag 改造为 Mutex 同步 #396
已合并
qq_51867290创建于 17 天前
从qq_51867290/cann-samplesfeat/flash-attn-lite-v0-kv-simt-mutex合入到cann/cann-samplesmaster
共 6 个文件变更+106-60
MSamples/2_Performance/flash_attn_lite_story/src/v0/kernel/_shared_k.h+19-14
@@ -22,12 +22,14 @@ constexpr uint32_t VL_B32 = VECTOR_REG_WIDTH / sizeof(float);
22constexpr uint32_t VL_B16 = VECTOR_REG_WIDTH / sizeof(bfloat16_t);22constexpr uint32_t VL_B16 = VECTOR_REG_WIDTH / sizeof(bfloat16_t);
23constexpr uint32_t B16_PER_DATABLOCK = C0_BYTES / sizeof(bfloat16_t);23constexpr uint32_t B16_PER_DATABLOCK = C0_BYTES / sizeof(bfloat16_t);
24constexpr float FLOAT_LOWEST = -3.402823466e+38F;24constexpr float FLOAT_LOWEST = -3.402823466e+38F;
25-constexpr AscendC::TEventID STATIC_EVENT_ID0 = 0;
26-constexpr AscendC::TEventID STATIC_EVENT_ID1 = 1;
27-constexpr AscendC::TEventID STATIC_EVENT_ID2 = 2;
28constexpr AscendC::FixpipeConfig PFA_CFG_UB = {AscendC::CO2Layout::ROW_MAJOR, true};25constexpr AscendC::FixpipeConfig PFA_CFG_UB = {AscendC::CO2Layout::ROW_MAJOR, true};
29constexpr AscendC::FixpipeConfig PFA_CFG_GM = {AscendC::CO2Layout::ROW_MAJOR, false};26constexpr AscendC::FixpipeConfig PFA_CFG_GM = {AscendC::CO2Layout::ROW_MAJOR, false};
30constexpr uint8_t GROUP_CROSS_MODE = 2;27constexpr uint8_t GROUP_CROSS_MODE = 2;
28+ 
29+// Mutex IDs for intra-core pipeline synchronization (AIV side)
30+constexpr AscendC::MutexID MUTEX_S_UB = 8;
31+constexpr AscendC::MutexID MUTEX_OD_UB = 9;
32+constexpr AscendC::MutexID MUTEX_P_UB_GM = 10;
31#ifdef SIM_COMPATIBLE33#ifdef SIM_COMPATIBLE
32constexpr uint8_t PAIR_CROSS_MODE = 4;34constexpr uint8_t PAIR_CROSS_MODE = 4;
33constexpr uint16_t AIV1_FLAG_OFFSET = 16;35constexpr uint16_t AIV1_FLAG_OFFSET = 16;
@@ -41,13 +43,6 @@ static constexpr AscendC::Reg::CastTrait castTraitOne = {
41 AscendC::RoundMode::CAST_ROUND};43 AscendC::RoundMode::CAST_ROUND};
42 44 
43// ── Helper ──45// ── Helper ──
44-template <AscendC::HardEvent E>
45-__aicore__ inline void SetWaitFlag(const AscendC::TEventID eventId)
46-{
47- AscendC::SetFlag<E>(eventId);
48- AscendC::WaitFlag<E>(eventId);
49-}
50- 
51template <typename R, typename T1, typename T2>46template <typename R, typename T1, typename T2>
52__aicore__ inline R CeilDiv(T1 x, T2 y)47__aicore__ inline R CeilDiv(T1 x, T2 y)
53{48{
@@ -267,7 +262,10 @@ __aicore__ inline void SoftmaxAndCastP(
267 const uint16_t bc, const float scale, const bool isFirst)262 const uint16_t bc, const float scale, const bool isFirst)
268{263{
269 RunOnlineSoftmax(sUBLocal, mUBLocal, lUBLocal, alphaUBLocal, halfBr, bc, scale, isFirst);264 RunOnlineSoftmax(sUBLocal, mUBLocal, lUBLocal, alphaUBLocal, halfBr, bc, scale, isFirst);
265+ 
266+ AscendC::Mutex::Lock<PIPE_V>(MUTEX_P_UB_GM);
270 Cast<bfloat16_t, float>(pUBLocal, sUBLocal, AscendC::RoundMode::CAST_RINT, bc * halfBr);267 Cast<bfloat16_t, float>(pUBLocal, sUBLocal, AscendC::RoundMode::CAST_RINT, bc * halfBr);
268+ AscendC::Mutex::Unlock<PIPE_V>(MUTEX_P_UB_GM);
271}269}
272 270 
273// v0/v1 共享: Cast 后的 P(BF16) → GM. 用于 SoftmaxAndWriteP 的公共尾段.271// v0/v1 共享: Cast 后的 P(BF16) → GM. 用于 SoftmaxAndWriteP 的公共尾段.
@@ -276,13 +274,13 @@ __aicore__ inline void WritePToGM(
276 const uint16_t bc, const uint16_t halfBr, const uint16_t br)274 const uint16_t bc, const uint16_t halfBr, const uint16_t br)
277{275{
278 using namespace AscendC;276 using namespace AscendC;
279- SetWaitFlag<HardEvent::V_MTE3>(STATIC_EVENT_ID1);277+ Mutex::Lock<PIPE_MTE3>(MUTEX_P_UB_GM);
280 DataCopy(278 DataCopy(
281 pGlobal[pHead], pUBLocal,279 pGlobal[pHead], pUBLocal,
282 DataCopyParams(280 DataCopyParams(
283 static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr * sizeof(bfloat16_t) / C0_BYTES), 0,281 static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr * sizeof(bfloat16_t) / C0_BYTES), 0,
284 static_cast<uint16_t>((br - halfBr) * sizeof(bfloat16_t) / C0_BYTES)));282 static_cast<uint16_t>((br - halfBr) * sizeof(bfloat16_t) / C0_BYTES)));
285- SetWaitFlag<HardEvent::MTE3_V>(STATIC_EVENT_ID1);283+ Mutex::Unlock<PIPE_MTE3>(MUTEX_P_UB_GM);
286}284}
287 285 
288// v0/v1 共享: O_acc = alpha * O_acc + ΔO286// v0/v1 共享: O_acc = alpha * O_acc + ΔO
@@ -377,13 +375,20 @@ __aicore__ inline void FinalOutput(
377 const uint16_t halfBr, const uint16_t d)375 const uint16_t halfBr, const uint16_t d)
378{376{
379 using namespace AscendC;377 using namespace AscendC;
378+ 
379+ Mutex::Lock<PIPE_V>(MUTEX_P_UB_GM);
380 asc_vf_call<FusedDivCastVF>(380 asc_vf_call<FusedDivCastVF>(
381 reinterpret_cast<__ubuf__ bfloat16_t*>(pUBLocal.GetPhyAddr()),381 reinterpret_cast<__ubuf__ bfloat16_t*>(pUBLocal.GetPhyAddr()),
382 reinterpret_cast<__ubuf__ float*>(oAccUBLocal.GetPhyAddr()),382 reinterpret_cast<__ubuf__ float*>(oAccUBLocal.GetPhyAddr()),
383 reinterpret_cast<__ubuf__ float*>(lUBLocal.GetPhyAddr()), halfBr, d);383 reinterpret_cast<__ubuf__ float*>(lUBLocal.GetPhyAddr()), halfBr, d);
384- SetWaitFlag<HardEvent::V_MTE3>(STATIC_EVENT_ID1);384+ Mutex::Unlock<PIPE_V>(MUTEX_P_UB_GM);
385+ 
386+ Mutex::Lock<PIPE_MTE3>(MUTEX_P_UB_GM);
385 DataCopy(outGlobal[outOff], pUBLocal, halfBr * d);387 DataCopy(outGlobal[outOff], pUBLocal, halfBr * d);
386- SetWaitFlag<HardEvent::MTE3_V>(STATIC_EVENT_ID1);388+ Mutex::Unlock<PIPE_MTE3>(MUTEX_P_UB_GM);
389+ 
390+ Mutex::Lock<PIPE_V>(MUTEX_P_UB_GM);
391+ Mutex::Unlock<PIPE_V>(MUTEX_P_UB_GM);
387}392}
388 393 
389} // namespace FALite394} // namespace FALite
MSamples/2_Performance/flash_attn_lite_story/src/v0/kernel/falite_kernel_aic.h+49-16
@@ -14,14 +14,15 @@
14 14 
15namespace FALite {15namespace FALite {
16 16 
17-constexpr AscendC::TEventID EV_K_L1_READY = 1;17+// Mutex IDs for intra-core pipeline synchronization (AIC side)
18-constexpr AscendC::TEventID EV_QK_L0_READY = 0;18+constexpr AscendC::MutexID MUTEX_Q_L1 = 0;
19-constexpr AscendC::TEventID EV_S_L0C_READY = 0;19+constexpr AscendC::MutexID MUTEX_K_L1 = 1;
20-constexpr AscendC::TEventID EV_P_L1_DONE = 0;20+constexpr AscendC::MutexID MUTEX_QK_L0 = 2;
21-constexpr AscendC::TEventID EV_V_L1_READY = 2;21+constexpr AscendC::MutexID MUTEX_S_L0C = 3;
22-constexpr AscendC::TEventID EV_PV_L0_READY = 1;22+constexpr AscendC::MutexID MUTEX_P_L1 = 4;
23-constexpr AscendC::TEventID EV_DO_L0C_READY = 1;23+constexpr AscendC::MutexID MUTEX_V_L1 = 5;
24-constexpr AscendC::TEventID EV_Q_L1_READY = 0;24+constexpr AscendC::MutexID MUTEX_PV_L0 = 6;
25+constexpr AscendC::MutexID MUTEX_DO_L0C = 7;
25 26 
26// C1: K×Q^T → S^T→GM. Q 已在 L1.27// C1: K×Q^T → S^T→GM. Q 已在 L1.
27__aicore__ inline void CubeStage1(28__aicore__ inline void CubeStage1(
@@ -34,14 +35,29 @@ __aicore__ inline void CubeStage1(
34 if ASCEND_IS_AIC {35 if ASCEND_IS_AIC {
35 const uint32_t br = data.br, bc = data.bc, d = data.headDim;36 const uint32_t br = data.br, bc = data.bc, d = data.headDim;
36 const uint64_t kOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(j) * bc * d;37 const uint64_t kOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(j) * bc * d;
38+ Mutex::Lock<PIPE_MTE2>(MUTEX_K_L1);
37 CopyGmToL1<bfloat16_t>(kL1Local, kGlobal[kOff], bc, d, d);39 CopyGmToL1<bfloat16_t>(kL1Local, kGlobal[kOff], bc, d, d);
38- SetWaitFlag<HardEvent::MTE2_MTE1>(EV_K_L1_READY);40+ Mutex::Unlock<PIPE_MTE2>(MUTEX_K_L1);
41+ 
42+ Mutex::Lock<PIPE_MTE1>(MUTEX_K_L1);
43+ Mutex::Lock<PIPE_MTE1>(MUTEX_QK_L0);
39 CopyL1ToL0A<bfloat16_t>(aL0ALocal, kL1Local, bc, d, bc, d);44 CopyL1ToL0A<bfloat16_t>(aL0ALocal, kL1Local, bc, d, bc, d);
45+ Mutex::Unlock<PIPE_MTE1>(MUTEX_K_L1);
46+ 
47+ Mutex::Lock<PIPE_MTE1>(MUTEX_Q_L1);
40 CopyL1ToL0B<bfloat16_t>(bL0BLocal, qL1Local, br, d, d, br, false);48 CopyL1ToL0B<bfloat16_t>(bL0BLocal, qL1Local, br, d, d, br, false);
41- SetWaitFlag<HardEvent::MTE1_M>(EV_QK_L0_READY);49+ Mutex::Unlock<PIPE_MTE1>(MUTEX_Q_L1);
50+ Mutex::Unlock<PIPE_MTE1>(MUTEX_QK_L0);
51+ 
52+ Mutex::Lock<PIPE_M>(MUTEX_QK_L0);
53+ Mutex::Lock<PIPE_M>(MUTEX_S_L0C);
42 CubeMmad<float, bfloat16_t, bfloat16_t>(mmadL0CLocal, aL0ALocal, bL0BLocal, bc, br, d, true);54 CubeMmad<float, bfloat16_t, bfloat16_t>(mmadL0CLocal, aL0ALocal, bL0BLocal, bc, br, d, true);
43- SetWaitFlag<HardEvent::M_FIX>(EV_S_L0C_READY);55+ Mutex::Unlock<PIPE_M>(MUTEX_QK_L0);
56+ Mutex::Unlock<PIPE_M>(MUTEX_S_L0C);
57+ 
58+ Mutex::Lock<PIPE_FIX>(MUTEX_S_L0C);
44 FixpipeL0CToGM<float, float>(sGlobal[sOff], mmadL0CLocal, bc, br, data.br);59 FixpipeL0CToGM<float, float>(sGlobal[sOff], mmadL0CLocal, bc, br, data.br);
60+ Mutex::Unlock<PIPE_FIX>(MUTEX_S_L0C);
45 PipeBarrier<PIPE_FIX>();61 PipeBarrier<PIPE_FIX>();
46 }62 }
47}63}
@@ -58,16 +74,32 @@ __aicore__ inline void CubeStage2(
58 if ASCEND_IS_AIC {74 if ASCEND_IS_AIC {
59 const uint32_t br = data.br, bc = data.bc, d = data.headDim;75 const uint32_t br = data.br, bc = data.bc, d = data.headDim;
60 const uint64_t pTaskBase = static_cast<uint64_t>(taskId) * bc * br;76 const uint64_t pTaskBase = static_cast<uint64_t>(taskId) * bc * br;
77+ Mutex::Lock<PIPE_MTE1>(MUTEX_P_L1);
61 LoadPTransToL0A(pL1Local, aL0ALocal, pGlobal, pTaskBase, data);78 LoadPTransToL0A(pL1Local, aL0ALocal, pGlobal, pTaskBase, data);
62- SetWaitFlag<HardEvent::MTE1_MTE2>(EV_P_L1_DONE);79+ Mutex::Unlock<PIPE_MTE1>(MUTEX_P_L1);
80+ 
81+ Mutex::Lock<PIPE_MTE2>(MUTEX_P_L1);
82+ Mutex::Lock<PIPE_MTE2>(MUTEX_V_L1);
63 const uint64_t vOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(j) * bc * d;83 const uint64_t vOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(j) * bc * d;
64 CopyGmToL1<bfloat16_t>(vL1Local, vGlobal[vOff], bc, d, d);84 CopyGmToL1<bfloat16_t>(vL1Local, vGlobal[vOff], bc, d, d);
65- SetWaitFlag<HardEvent::MTE2_MTE1>(EV_V_L1_READY);85+ Mutex::Unlock<PIPE_MTE2>(MUTEX_V_L1);
86+ Mutex::Unlock<PIPE_MTE2>(MUTEX_P_L1);
87+ 
88+ Mutex::Lock<PIPE_MTE1>(MUTEX_V_L1);
89+ Mutex::Lock<PIPE_MTE1>(MUTEX_PV_L0);
66 CopyL1ToL0B<bfloat16_t>(bL0BLocal, vL1Local, bc, d, bc, d, true);90 CopyL1ToL0B<bfloat16_t>(bL0BLocal, vL1Local, bc, d, bc, d, true);
67- SetWaitFlag<HardEvent::MTE1_M>(EV_PV_L0_READY);91+ Mutex::Unlock<PIPE_MTE1>(MUTEX_V_L1);
92+ Mutex::Unlock<PIPE_MTE1>(MUTEX_PV_L0);
93+ 
94+ Mutex::Lock<PIPE_M>(MUTEX_PV_L0);
95+ Mutex::Lock<PIPE_M>(MUTEX_DO_L0C);
68 CubeMmad<float, bfloat16_t, bfloat16_t>(mmadL0CLocal, aL0ALocal, bL0BLocal, br, d, bc, true);96 CubeMmad<float, bfloat16_t, bfloat16_t>(mmadL0CLocal, aL0ALocal, bL0BLocal, br, d, bc, true);
69- SetWaitFlag<HardEvent::M_FIX>(EV_DO_L0C_READY);97+ Mutex::Unlock<PIPE_M>(MUTEX_PV_L0);
98+ Mutex::Unlock<PIPE_M>(MUTEX_DO_L0C);
99+ 
100+ Mutex::Lock<PIPE_FIX>(MUTEX_DO_L0C);
70 FixpipeL0CToGM<float, float>(dOGlobal[dOOff], mmadL0CLocal, br, d, d);101 FixpipeL0CToGM<float, float>(dOGlobal[dOOff], mmadL0CLocal, br, d, d);
102+ Mutex::Unlock<PIPE_FIX>(MUTEX_DO_L0C);
71 PipeBarrier<PIPE_FIX>();103 PipeBarrier<PIPE_FIX>();
72 }104 }
73}105}
@@ -83,8 +115,9 @@ __aicore__ inline void ProcessOneTaskAIC(
83 const uint32_t batchIdx = taskId / data.tr, tileIdx = taskId % data.tr;115 const uint32_t batchIdx = taskId / data.tr, tileIdx = taskId % data.tr;
84 const uint32_t br = data.br, d = data.headDim;116 const uint32_t br = data.br, d = data.headDim;
85 const uint64_t qOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(tileIdx) * br * d;117 const uint64_t qOff = static_cast<uint64_t>(batchIdx) * data.seqLen * d + static_cast<uint64_t>(tileIdx) * br * d;
118+ Mutex::Lock<PIPE_MTE2>(MUTEX_Q_L1);
86 CopyGmToL1<bfloat16_t>(ws.qL1, qGlobal[qOff], br, d, d);119 CopyGmToL1<bfloat16_t>(ws.qL1, qGlobal[qOff], br, d, d);
87- SetWaitFlag<HardEvent::MTE2_MTE1>(EV_Q_L1_READY);120+ Mutex::Unlock<PIPE_MTE2>(MUTEX_Q_L1);
88 121 
89 for (uint32_t j = 0; j < data.tc; ++j) {122 for (uint32_t j = 0; j < data.tc; ++j) {
90 if (j > 0) CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_DONE);123 if (j > 0) CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_DONE);
MSamples/2_Performance/flash_attn_lite_story/src/v0/kernel/falite_kernel_aiv.h+10-2
@@ -25,15 +25,19 @@ __aicore__ inline void SoftmaxAndWriteP(
25 using namespace AscendC;25 using namespace AscendC;
26 const uint32_t halfBr = data.br / 2, bc = data.bc;26 const uint32_t halfBr = data.br / 2, bc = data.bc;
27 CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_S_READY);27 CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_S_READY);
28+ Mutex::Lock<PIPE_MTE2>(MUTEX_S_UB);
28 DataCopy(29 DataCopy(
29 sUBLocal, sGlobal[sHead],30 sUBLocal, sGlobal[sHead],
30 DataCopyParams(31 DataCopyParams(
31 static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr * sizeof(float) / C0_BYTES),32 static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr * sizeof(float) / C0_BYTES),
32 static_cast<uint16_t>((data.br - halfBr) * sizeof(float) / C0_BYTES), 0));33 static_cast<uint16_t>((data.br - halfBr) * sizeof(float) / C0_BYTES), 0));
33- SetWaitFlag<HardEvent::MTE2_V>(STATIC_EVENT_ID0);34+ Mutex::Unlock<PIPE_MTE2>(MUTEX_S_UB);
35+ 
36+ Mutex::Lock<PIPE_V>(MUTEX_S_UB);
34 SoftmaxAndCastP(37 SoftmaxAndCastP(
35 sUBLocal, mUBLocal, lUBLocal, alphaUBLocal, pUBLocal, static_cast<uint16_t>(halfBr), static_cast<uint16_t>(bc),38 sUBLocal, mUBLocal, lUBLocal, alphaUBLocal, pUBLocal, static_cast<uint16_t>(halfBr), static_cast<uint16_t>(bc),
36 data.scale, j == 0);39 data.scale, j == 0);
40+ Mutex::Unlock<PIPE_V>(MUTEX_S_UB);
37 WritePToGM(41 WritePToGM(
38 pUBLocal, pGlobal, pHead, static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr),42 pUBLocal, pGlobal, pHead, static_cast<uint16_t>(bc), static_cast<uint16_t>(halfBr),
39 static_cast<uint16_t>(data.br));43 static_cast<uint16_t>(data.br));
@@ -49,10 +53,14 @@ __aicore__ inline void AccumulateDeltaO(
49 using namespace AscendC;53 using namespace AscendC;
50 const uint32_t halfBr = data.br / 2, d = data.headDim;54 const uint32_t halfBr = data.br / 2, d = data.headDim;
51 CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_O_READY);55 CrossCoreWaitFlag<GROUP_CROSS_MODE, PIPE_MTE2>(FLAG_O_READY);
56+ Mutex::Lock<PIPE_MTE2>(MUTEX_OD_UB);
52 DataCopy(oDeltaUBLocal, dOGlobal[oHead], halfBr * d);57 DataCopy(oDeltaUBLocal, dOGlobal[oHead], halfBr * d);
53- SetWaitFlag<HardEvent::MTE2_V>(STATIC_EVENT_ID0);58+ Mutex::Unlock<PIPE_MTE2>(MUTEX_OD_UB);
59+ 
60+ Mutex::Lock<PIPE_V>(MUTEX_OD_UB);
54 AccumulateDeltaOCore(61 AccumulateDeltaOCore(
55 oAccUBLocal, oDeltaUBLocal, alphaUBLocal, static_cast<uint16_t>(halfBr), static_cast<uint16_t>(d));62 oAccUBLocal, oDeltaUBLocal, alphaUBLocal, static_cast<uint16_t>(halfBr), static_cast<uint16_t>(d));
63+ Mutex::Unlock<PIPE_V>(MUTEX_OD_UB);
56 CrossCoreSetFlag<GROUP_CROSS_MODE, PIPE_V>(FLAG_DONE);64 CrossCoreSetFlag<GROUP_CROSS_MODE, PIPE_V>(FLAG_DONE);
57}65}
58 66 
MSamples/2_Performance/kv_rms_norm_rope_cache_story/membase/full_load.asc+8-9
@@ -70,6 +70,8 @@ public:
70 pipe_.InitBuffer(cosSinQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dk * NUM_TWO, sizeof(T)));70 pipe_.InitBuffer(cosSinQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dk * NUM_TWO, sizeof(T)));
71 pipe_.InitBuffer(outQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv, sizeof(T)));71 pipe_.InitBuffer(outQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv, sizeof(T)));
72 pipe_.InitBuffer(wsBuffer_, ubFactor_ * (tiling_->dv * NUM_THREE + tiling_->dk * NUM_EIGHT) * sizeof(float));72 pipe_.InitBuffer(wsBuffer_, ubFactor_ * (tiling_->dv * NUM_THREE + tiling_->dk * NUM_EIGHT) * sizeof(float));
73+ 
74+ sMte3Mutex = AscendC::AllocMutexID();
73 }75 }
74 76 
75 __aicore__ inline void Process()77 __aicore__ inline void Process()
@@ -91,6 +93,7 @@ public:
91 ProcessTile(loopIdx, rows, globalRow, gammaFp32);93 ProcessTile(loopIdx, rows, globalRow, gammaFp32);
92 }94 }
93 gammaQueue_.FreeTensor(gammaFp32);95 gammaQueue_.FreeTensor(gammaFp32);
96+ AscendC::ReleaseMutexID(sMte3Mutex);
94 }97 }
95 98 
96private:99private:
@@ -223,8 +226,9 @@ private:
223 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dk * sizeof(T)), 0, 0, 0};226 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dk * sizeof(T)), 0, 0, 0};
224 for (int64_t i = 0; i < rows; ++i) {227 for (int64_t i = 0; i < rows; ++i) {
225 int64_t token = globalRow + i;228 int64_t token = globalRow + i;
229+ Mutex::Lock<PIPE_MTE3>(sMte3Mutex);
226 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);230 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);
227- WaitScalarToMte3();231+ Mutex::Unlock<PIPE_MTE3>(sMte3Mutex);
228 if (cacheOffset >= 0) {232 if (cacheOffset >= 0) {
229 int64_t batch = token / tiling_->seqLength;233 int64_t batch = token / tiling_->seqLength;
230 DataCopyPad(kCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dk],234 DataCopyPad(kCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dk],
@@ -239,8 +243,9 @@ private:
239 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dv * sizeof(T)), 0, 0, 0};243 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dv * sizeof(T)), 0, 0, 0};
240 for (int64_t i = 0; i < rows; ++i) {244 for (int64_t i = 0; i < rows; ++i) {
241 int64_t token = globalRow + i;245 int64_t token = globalRow + i;
246+ Mutex::Lock<PIPE_MTE3>(sMte3Mutex);
242 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);247 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);
243- WaitScalarToMte3();248+ Mutex::Unlock<PIPE_MTE3>(sMte3Mutex);
244 if (cacheOffset >= 0) {249 if (cacheOffset >= 0) {
245 int64_t batch = token / tiling_->seqLength;250 int64_t batch = token / tiling_->seqLength;
246 DataCopyPad(vCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dv],251 DataCopyPad(vCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dv],
@@ -250,14 +255,8 @@ private:
250 }255 }
251 }256 }
252 257 
253- __aicore__ inline void WaitScalarToMte3()
254- {
255- event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
256- SetFlag<HardEvent::S_MTE3>(eventId);
257- WaitFlag<HardEvent::S_MTE3>(eventId);
258- }
259- 
260 TPipe pipe_;258 TPipe pipe_;
259+ uint8_t sMte3Mutex = 0;
261 TQue<QuePosition::VECIN, 1> gammaQueue_;260 TQue<QuePosition::VECIN, 1> gammaQueue_;
262 TQue<QuePosition::VECIN, 1> inQueue_;261 TQue<QuePosition::VECIN, 1> inQueue_;
263 TQue<QuePosition::VECIN, 1> cosSinQueue_;262 TQue<QuePosition::VECIN, 1> cosSinQueue_;
MSamples/2_Performance/kv_rms_norm_rope_cache_story/regbase/full_load.asc+8-9
@@ -199,6 +199,8 @@ public:
199 pipe_.InitBuffer(inQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv + tiling_->dk, sizeof(T)));199 pipe_.InitBuffer(inQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv + tiling_->dk, sizeof(T)));
200 pipe_.InitBuffer(cosSinQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dk * NUM_TWO, sizeof(T)));200 pipe_.InitBuffer(cosSinQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dk * NUM_TWO, sizeof(T)));
201 pipe_.InitBuffer(outQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv, sizeof(T)));201 pipe_.InitBuffer(outQueue_, BUFFER_NUM, ubFactor_ * AlignBytes(tiling_->dv, sizeof(T)));
202+ 
203+ sMte3Mutex = AscendC::AllocMutexID();
202 }204 }
203 205 
204 __aicore__ inline void Process()206 __aicore__ inline void Process()
@@ -217,6 +219,7 @@ public:
217 ProcessTile(loopIdx, rows, globalRow, gammaLocal);219 ProcessTile(loopIdx, rows, globalRow, gammaLocal);
218 }220 }
219 gammaQueue_.FreeTensor(gammaLocal);221 gammaQueue_.FreeTensor(gammaLocal);
222+ AscendC::ReleaseMutexID(sMte3Mutex);
220 }223 }
221 224 
222private:225private:
@@ -305,8 +308,9 @@ private:
305 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dk * sizeof(T)), 0, 0, 0};308 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dk * sizeof(T)), 0, 0, 0};
306 for (int64_t i = 0; i < rows; ++i) {309 for (int64_t i = 0; i < rows; ++i) {
307 int64_t token = globalRow + i;310 int64_t token = globalRow + i;
311+ Mutex::Lock<PIPE_MTE3>(sMte3Mutex);
308 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);312 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);
309- WaitScalarToMte3();313+ Mutex::Unlock<PIPE_MTE3>(sMte3Mutex);
310 if (cacheOffset >= 0) {314 if (cacheOffset >= 0) {
311 int64_t batch = token / tiling_->seqLength;315 int64_t batch = token / tiling_->seqLength;
312 DataCopyPad(kCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dk],316 DataCopyPad(kCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dk],
@@ -321,8 +325,9 @@ private:
321 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dv * sizeof(T)), 0, 0, 0};325 DataCopyExtParams copyParams{1, static_cast<uint32_t>(tiling_->dv * sizeof(T)), 0, 0, 0};
322 for (int64_t i = 0; i < rows; ++i) {326 for (int64_t i = 0; i < rows; ++i) {
323 int64_t token = globalRow + i;327 int64_t token = globalRow + i;
328+ Mutex::Lock<PIPE_MTE3>(sMte3Mutex);
324 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);329 int64_t cacheOffset = indexGm_(loopIdx * ubFactor_ + i);
325- WaitScalarToMte3();330+ Mutex::Unlock<PIPE_MTE3>(sMte3Mutex);
326 if (cacheOffset >= 0) {331 if (cacheOffset >= 0) {
327 int64_t batch = token / tiling_->seqLength;332 int64_t batch = token / tiling_->seqLength;
328 DataCopyPad(vCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dv],333 DataCopyPad(vCacheGm_[(batch * tiling_->cacheLength + cacheOffset) * tiling_->dv],
@@ -332,14 +337,8 @@ private:
332 }337 }
333 }338 }
334 339 
335- __aicore__ inline void WaitScalarToMte3()
336- {
337- event_t eventId = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
338- SetFlag<HardEvent::S_MTE3>(eventId);
339- WaitFlag<HardEvent::S_MTE3>(eventId);
340- }
341- 
342 TPipe pipe_;340 TPipe pipe_;
341+ uint8_t sMte3Mutex = 0;
343 TQue<QuePosition::VECIN, 1> gammaQueue_;342 TQue<QuePosition::VECIN, 1> gammaQueue_;
344 TQue<QuePosition::VECIN, 1> inQueue_;343 TQue<QuePosition::VECIN, 1> inQueue_;
345 TQue<QuePosition::VECIN, 1> cosSinQueue_;344 TQue<QuePosition::VECIN, 1> cosSinQueue_;
MSamples/2_Performance/simt_histogram_story/src/0_serial_scalar.asc+12-10
@@ -30,6 +30,8 @@ __global__ __aicore__ void simt_histogram_serial_scalar_kernel(
30 // 初始化30 // 初始化
31 // ================================================================31 // ================================================================
32 AscendC::InitSocState();32 AscendC::InitSocState();
33+ uint8_t mte2SMutex = AscendC::AllocMutexID();
34+ uint8_t vSMutex = AscendC::AllocMutexID();
33 int32_t coreId = static_cast<int32_t>(AscendC::GetBlockIdx());35 int32_t coreId = static_cast<int32_t>(AscendC::GetBlockIdx());
34 int32_t coreNum = static_cast<int32_t>(AscendC::GetBlockNum());36 int32_t coreNum = static_cast<int32_t>(AscendC::GetBlockNum());
35 37 
@@ -62,6 +64,7 @@ __global__ __aicore__ void simt_histogram_serial_scalar_kernel(
62 LocalTensor<XType> maxLocal = maxQue.AllocTensor<XType>();64 LocalTensor<XType> maxLocal = maxQue.AllocTensor<XType>();
63 65 
64 {66 {
67+ Mutex::Lock<PIPE_MTE2>(mte2SMutex);
65 DataCopyParams cpMin{1, static_cast<uint16_t>(sizeof(XType)), 0, 0};68 DataCopyParams cpMin{1, static_cast<uint16_t>(sizeof(XType)), 0, 0};
66 DataCopyPadParams pad{true, 0, 0, 0};69 DataCopyPadParams pad{true, 0, 0, 0};
67 DataCopyPad(minLocal, minGm, cpMin, pad);70 DataCopyPad(minLocal, minGm, cpMin, pad);
@@ -73,14 +76,12 @@ __global__ __aicore__ void simt_histogram_serial_scalar_kernel(
73 }76 }
74 77 
75 // 等待 MTE 搬运完成（MTE2 → Scalar 同步）78 // 等待 MTE 搬运完成（MTE2 → Scalar 同步）
76- {79+ Mutex::Unlock<PIPE_MTE2>(mte2SMutex);
77- event_t evMte = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));80+ Mutex::Lock<PIPE_S>(mte2SMutex);
78- SetFlag<HardEvent::MTE2_S>(evMte);
79- WaitFlag<HardEvent::MTE2_S>(evMte);
80- }
81 81 
82 XType minVal = minLocal.GetValue(0);82 XType minVal = minLocal.GetValue(0);
83 XType maxVal = maxLocal.GetValue(0);83 XType maxVal = maxLocal.GetValue(0);
84+ Mutex::Unlock<PIPE_S>(mte2SMutex);
84 85 
85 minQue.FreeTensor(minLocal);86 minQue.FreeTensor(minLocal);
86 maxQue.FreeTensor(maxLocal);87 maxQue.FreeTensor(maxLocal);
@@ -116,14 +117,13 @@ __global__ __aicore__ void simt_histogram_serial_scalar_kernel(
116 // 在 UB 上分配本地直方图缓冲区并清零117 // 在 UB 上分配本地直方图缓冲区并清零
117 // ================================================================118 // ================================================================
118 LocalTensor<OutType> yLocal = yQue.AllocTensor<OutType>();119 LocalTensor<OutType> yLocal = yQue.AllocTensor<OutType>();
120+ Mutex::Lock<PIPE_V>(vSMutex);
119 Duplicate<OutType>(yLocal, 0, bins + ALIGNED_NUM);121 Duplicate<OutType>(yLocal, 0, bins + ALIGNED_NUM);
120 122 
121 // 等待 Vector 的 Duplicate 指令执行完毕（Vector → Scalar 同步）123 // 等待 Vector 的 Duplicate 指令执行完毕（Vector → Scalar 同步）
122- {124+ Mutex::Unlock<PIPE_V>(vSMutex);
123- event_t evV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));125+ Mutex::Lock<PIPE_S>(vSMutex);
124- SetFlag<HardEvent::V_S>(evV);126+ Mutex::Unlock<PIPE_S>(vSMutex);
125- WaitFlag<HardEvent::V_S>(evV);
126- }
127 127 
128 // ================================================================128 // ================================================================
129 // 分 tile 处理本 core 的数据段（双缓冲流水）129 // 分 tile 处理本 core 的数据段（双缓冲流水）
@@ -208,6 +208,8 @@ __global__ __aicore__ void simt_histogram_serial_scalar_kernel(
208 208 
209 // --- 流水线收尾：等待所有硬件队列操作完毕 ---209 // --- 流水线收尾：等待所有硬件队列操作完毕 ---
210 AscendC::PipeBarrier<PIPE_ALL>();210 AscendC::PipeBarrier<PIPE_ALL>();
211+ AscendC::ReleaseMutexID(mte2SMutex);
212+ AscendC::ReleaseMutexID(vSMutex);
211}213}
212 214 
213void LaunchHistogramKernel(uint32_t blocks, aclrtStream stream, XType* x, XType* min, XType* max, OutType* y,215void LaunchHistogramKernel(uint32_t blocks, aclrtStream stream, XType* x, XType* min, XType* max, OutType* y,