Star1298
1412
代码
代码
Issues108
Pull Requests216
Actions
项目讨论
Wiki
分析
Star1298
1412
已合并
cholesky算子支持大尾轴场景 #2255
xiu_ling_wang创建于 4月16日
cholesky算子支持大尾轴场景 #2255
已合并
xiu_ling_wang创建于 4月16日
从xiu_ling_wang/ops-mathcholesky-fix合入到cann/ops-mathmaster
共 5 个文件变更+454-296
Mmath/cholesky/op_host/cholesky_tiling.cpp+38-32
@@ -17,10 +17,10 @@
17namespace optiling {17namespace optiling {
18constexpr uint32_t TILING_KEY_FALSE = 1;18constexpr uint32_t TILING_KEY_FALSE = 1;
19constexpr uint32_t TILING_KEY_TRUE = 2;19constexpr uint32_t TILING_KEY_TRUE = 2;
20-constexpr uint32_t BYTE_LEN_4 = 4;
21constexpr uint32_t MINIMUM_DIMENSION = 2;20constexpr uint32_t MINIMUM_DIMENSION = 2;
22constexpr uint32_t UPPER_INDEX = 0;21constexpr uint32_t UPPER_INDEX = 0;
23constexpr uint32_t WS_SYS_SIZE = 16U * 1024U * 1024U;22constexpr uint32_t WS_SYS_SIZE = 16U * 1024U * 1024U;
23+constexpr uint32_t MAX_BLOCK_SIZE = 128;
24 24 
25class CholeskyTiling {25class CholeskyTiling {
26public:26public:
@@ -29,19 +29,18 @@ public:
29 ge::graphStatus RunBigKernelTiling();29 ge::graphStatus RunBigKernelTiling();
30 30 
31private:31private:
32- uint8_t GetDataTypeSize();32+ uint32_t GetTilingKeyVal() const;
33- uint64_t GetTilingKeyVal();33+ void PrintTilingData();
34- void FillTilingData();
35 34 
36private:35private:
37 gert::TilingContext* tilingContext = nullptr;36 gert::TilingContext* tilingContext = nullptr;
38- ge::DataType dataType = ge::DT_UNDEFINED;
39 CholeskyTilingData tilingData;37 CholeskyTilingData tilingData;
40- uint8_t dataTypeSize = 4;
41 uint32_t matSizeN = 0;38 uint32_t matSizeN = 0;
42- uint32_t matrixNumCount = 0;39+ uint64_t matrixNumCount = 1;
43 uint32_t needCoreNum = 0;40 uint32_t needCoreNum = 0;
44 bool upper = false;41 bool upper = false;
42+ uint32_t blockSize = 0;
43+ uint32_t blockNum = 0;
45};44};
46 45 
47ge::graphStatus CholeskyTiling::Init() {46ge::graphStatus CholeskyTiling::Init() {
@@ -50,14 +49,11 @@ ge::graphStatus CholeskyTiling::Init() {
50 return ge::GRAPH_FAILED;49 return ge::GRAPH_FAILED;
51 }50 }
52 51 
53- auto inputDtype = inputTensor->GetDataType();
54- if (dataType == ge::DT_UNDEFINED) {
55- dataType = inputDtype;
56- dataTypeSize = GetDataTypeSize();
57- }
58- 
59 auto attrs = tilingContext->GetAttrs();52 auto attrs = tilingContext->GetAttrs();
60 const bool* ptrUpper = attrs->GetAttrPointer<bool>(UPPER_INDEX);53 const bool* ptrUpper = attrs->GetAttrPointer<bool>(UPPER_INDEX);
54+ if (ptrUpper == nullptr) {
55+ return ge::GRAPH_FAILED;
56+ }
61 upper = *ptrUpper;57 upper = *ptrUpper;
62 58 
63 auto matAShape = tilingContext->GetInputShape(0)->GetOriginShape();59 auto matAShape = tilingContext->GetInputShape(0)->GetOriginShape();
@@ -66,13 +62,20 @@ ge::graphStatus CholeskyTiling::Init() {
66 return ge::GRAPH_FAILED;62 return ge::GRAPH_FAILED;
67 }63 }
68 matSizeN = static_cast<uint32_t>(matAShape[inputDim-1]);64 matSizeN = static_cast<uint32_t>(matAShape[inputDim-1]);
69- matrixNumCount = 1;
70 for (uint32_t i = 0; i < (inputDim - MINIMUM_DIMENSION); i++) {65 for (uint32_t i = 0; i < (inputDim - MINIMUM_DIMENSION); i++) {
71- matrixNumCount = matrixNumCount * static_cast<uint32_t>(matAShape[i]);66+ matrixNumCount = matrixNumCount * static_cast<uint64_t>(matAShape[i]);
67+ }
68+ 
69+ if (matSizeN <= MAX_BLOCK_SIZE) {
70+ blockSize = matSizeN;
71+ blockNum = 1;
72+ } else {
73+ blockSize = MAX_BLOCK_SIZE;
74+ blockNum = (matSizeN + blockSize - 1) / blockSize;
72 }75 }
73 76 
74 auto compileInfo = reinterpret_cast<const CholeskyCompileInfo*>(tilingContext->GetCompileInfo());77 auto compileInfo = reinterpret_cast<const CholeskyCompileInfo*>(tilingContext->GetCompileInfo());
75- int64_t coreNumPlatForm = compileInfo->coreNum;78+ uint32_t coreNumPlatForm = compileInfo->coreNum;
76 needCoreNum = coreNumPlatForm < matrixNumCount ? coreNumPlatForm : matrixNumCount;79 needCoreNum = coreNumPlatForm < matrixNumCount ? coreNumPlatForm : matrixNumCount;
77 80 
78 size_t* currentWorkSpace = tilingContext->GetWorkspaceSizes(1);81 size_t* currentWorkSpace = tilingContext->GetWorkspaceSizes(1);
@@ -86,20 +89,22 @@ ge::graphStatus CholeskyTiling::RunBigKernelTiling() {
86 tilingContext->SetBlockDim(needCoreNum);89 tilingContext->SetBlockDim(needCoreNum);
87 tilingContext->SetTilingKey(GetTilingKeyVal());90 tilingContext->SetTilingKey(GetTilingKeyVal());
88 tilingContext->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());91 tilingContext->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
89- FillTilingData();92+
93+ tilingData.set_matrixNumCount(matrixNumCount);
94+ tilingData.set_matSizeN(matSizeN);
95+ tilingData.set_blockSize(blockSize);
96+ tilingData.set_blockNum(blockNum);
97+
98+ if (tilingContext->GetRawTilingData() == nullptr) {
99+ return ge::GRAPH_FAILED;
100+ }
101+ tilingData.SaveToBuffer(tilingContext->GetRawTilingData()->GetData(), tilingContext->GetRawTilingData()->GetCapacity());
102+ 
103+ PrintTilingData();
90 return ge::GRAPH_SUCCESS;104 return ge::GRAPH_SUCCESS;
91}105}
92 106 
93-uint8_t CholeskyTiling::GetDataTypeSize() {107+uint32_t CholeskyTiling::GetTilingKeyVal() const {
94- switch (dataType) {
95- case ge::DT_FLOAT:
96- return BYTE_LEN_4;
97- default:
98- return BYTE_LEN_4;
99- }
100-}
101- 
102-uint64_t CholeskyTiling::GetTilingKeyVal() {
103 if (upper == true) {108 if (upper == true) {
104 return TILING_KEY_TRUE;109 return TILING_KEY_TRUE;
105 } else {110 } else {
@@ -107,10 +112,11 @@ uint64_t CholeskyTiling::GetTilingKeyVal() {
107 }112 }
108}113}
109 114 
110-void CholeskyTiling::FillTilingData() {115+void CholeskyTiling::PrintTilingData() {
111- tilingData.set_matrixNumCount(matrixNumCount);116+ OP_LOGD(tilingContext, "matSizeN: %ld", matSizeN);
112- tilingData.set_matSizeN(matSizeN);117+ OP_LOGD(tilingContext, "matrixNumCount: %ld", matrixNumCount);
113- tilingData.SaveToBuffer(tilingContext->GetRawTilingData()->GetData(), tilingContext->GetRawTilingData()->GetCapacity());118+ OP_LOGD(tilingContext, "blockSize: %ld", blockSize);
119+ OP_LOGD(tilingContext, "blockNum: %ld", blockNum);
114}120}
115 121 
116static ge::graphStatus CholeskyTilingFunc(gert::TilingContext* context)122static ge::graphStatus CholeskyTilingFunc(gert::TilingContext* context)
@@ -132,7 +138,7 @@ static ge::graphStatus tilingPrepareTiling(gert::TilingParseContext* context)
132 138
133 OP_CHECK_IF(139 OP_CHECK_IF(
134 (compileInfo->coreNum <= 0),140 (compileInfo->coreNum <= 0),
135- OP_LOGE(context->GetNodeName(), "Cholesky GetHardwareInfo Failed, vectorCoreNum: %d", compileInfo->coreNum), return ge::GRAPH_FAILED);141+ OP_LOGE(context->GetNodeName(), "Cholesky GetHardwareInfo Failed, vectorCoreNum: %u", compileInfo->coreNum), return ge::GRAPH_FAILED);
136 142 
137 return ge::GRAPH_SUCCESS;143 return ge::GRAPH_SUCCESS;
138}144}
Mmath/cholesky/op_host/cholesky_tiling.h+4-2
@@ -16,12 +16,14 @@
16namespace optiling {16namespace optiling {
17 17 
18struct CholeskyCompileInfo {18struct CholeskyCompileInfo {
19- int32_t coreNum = 0;19+ uint32_t coreNum = 0;
20};20};
21 21 
22BEGIN_TILING_DATA_DEF(CholeskyTilingData)22BEGIN_TILING_DATA_DEF(CholeskyTilingData)
23 TILING_DATA_FIELD_DEF(uint32_t, matSizeN);23 TILING_DATA_FIELD_DEF(uint32_t, matSizeN);
24- TILING_DATA_FIELD_DEF(uint32_t, matrixNumCount);24+ TILING_DATA_FIELD_DEF(uint64_t, matrixNumCount);
25+ TILING_DATA_FIELD_DEF(uint32_t, blockSize);
26+ TILING_DATA_FIELD_DEF(uint32_t, blockNum);
25END_TILING_DATA_DEF;27END_TILING_DATA_DEF;
26 28 
27REGISTER_TILING_DATA_CLASS(Cholesky, CholeskyTilingData)29REGISTER_TILING_DATA_CLASS(Cholesky, CholeskyTilingData)
Mmath/cholesky/op_host/op_api/aclnn_linalg_cholesky.cpp+39-28
@@ -28,6 +28,7 @@ extern "C" {
28 28 
29const int64_t SECOND_LAST_DIM_OFFSET = 2;29const int64_t SECOND_LAST_DIM_OFFSET = 2;
30const int64_t LAST_DIM_OFFSET = 1;30const int64_t LAST_DIM_OFFSET = 1;
31+const int64_t THRESHOLD_VALUE = 10;
31 32 
32static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST = {33static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_DTYPE_SUPPORT_LIST = {
33 op::DataType::DT_FLOAT, op::DataType::DT_BF16};34 op::DataType::DT_FLOAT, op::DataType::DT_BF16};
@@ -49,7 +50,7 @@ static bool CheckNotNull(const aclTensor* self, aclTensor* out) {
49 return true;50 return true;
50}51}
51 52 
52-static bool CheckDtypeValid(const aclTensor *self, aclTensor *out) {53+static bool CheckDtypeValid(const aclTensor *self, const aclTensor *out) {
53 auto supportList = GetDtypeSupportList();54 auto supportList = GetDtypeSupportList();
54 55 
55 //检查self与out的数据类型是否一致56 //检查self与out的数据类型是否一致
@@ -60,7 +61,7 @@ static bool CheckDtypeValid(const aclTensor *self, aclTensor *out) {
60 return true;61 return true;
61}62}
62 63 
63-static bool CheckFormat(const aclTensor *self, aclTensor *out) {64+static bool CheckFormat(const aclTensor *self, const aclTensor *out) {
64 // 输入输出的格式需要一致65 // 输入输出的格式需要一致
65 if (self->GetStorageFormat() != out->GetStorageFormat()) {66 if (self->GetStorageFormat() != out->GetStorageFormat()) {
66 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of input and output should be equal. self [%s], out [%s].",67 OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Format of input and output should be equal. self [%s], out [%s].",
@@ -77,7 +78,7 @@ static bool CheckFormat(const aclTensor *self, aclTensor *out) {
77 return true;78 return true;
78}79}
79 80 
80-static bool CheckShape(const aclTensor *self, aclTensor *out) {81+static bool CheckShape(const aclTensor *self, const aclTensor *out) {
81 // 维度不能超过882 // 维度不能超过8
82 OP_CHECK_MAX_DIM(self, ACLNN_MAX_SHAPE_RANK, return false);83 OP_CHECK_MAX_DIM(self, ACLNN_MAX_SHAPE_RANK, return false);
83 84 
@@ -85,16 +86,7 @@ static bool CheckShape(const aclTensor *self, aclTensor *out) {
85 OP_CHECK_MIN_DIM(self, 2, return false);86 OP_CHECK_MIN_DIM(self, 2, return false);
86 87 
87 // self和out的shape必须一致88 // self和out的shape必须一致
88- OP_CHECK_SHAPE_NOT_EQUAL(self, out, return false);89+ OP_CHECK_SHAPE_NOT_EQUAL(self, out, return false);
89- 
90- // self最后两维必须为相同
91- auto dims = static_cast<int64_t>(self->GetViewShape().GetDimNum());
92- int64_t last_dim_size = self->GetViewShape().GetDim(dims -1);
93- int64_t second_last_dim_size = self->GetViewShape().GetDim(dims -2);
94- if (last_dim_size != second_last_dim_size) {
95- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "self must be batches of square matrices, but they are [%ld] by [%ld] matrices", second_last_dim_size, last_dim_size);
96- return false;
97- }
98 90 
99 return true;91 return true;
100}92}
@@ -141,25 +133,50 @@ aclnnStatus aclnnLinalgCholeskyGetWorkspaceSize(const aclTensor *self, bool uppe
141 auto ret = CheckParams(self, out);133 auto ret = CheckParams(self, out);
142 CHECK_RET(ret == ACLNN_SUCCESS, ret);134 CHECK_RET(ret == ACLNN_SUCCESS, ret);
143 135 
136+ if (self->IsEmpty()) {
137+ // 根据实际支持情况补充
138+ *workspaceSize = 0;
139+ uniqueExecutor.ReleaseTo(executor);
140+ return ACLNN_SUCCESS;
141+ }
142+ 
143+ // self最后两维必须为相同
144+ auto dims = static_cast<int64_t>(self->GetViewShape().GetDimNum());
145+ int64_t last_dim_size = self->GetViewShape().GetDim(dims -1);
146+ int64_t second_last_dim_size = self->GetViewShape().GetDim(dims -2);
147+ if (last_dim_size != second_last_dim_size) {
148+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "self must be batches of square matrices, but they are [%ld] by [%ld] matrices", second_last_dim_size, last_dim_size);
149+ return ACLNN_ERR_PARAM_INVALID;
150+ }
151+ 
144 // self如果非连续，需要转换152 // self如果非连续，需要转换
145 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());153 auto selfContiguous = l0op::Contiguous(self, uniqueExecutor.get());
146 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);154 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
147 155 
148- // upper是false时，需要进行转置
149- if (!upper) {
150- selfContiguous = SwapDim(selfContiguous, uniqueExecutor.get());
151- CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
152- }
153- 
154 // 输入类型是bf16时，需要转换成fp32156 // 输入类型是bf16时，需要转换成fp32
155 if (self->GetDataType() == op::DataType::DT_BF16) {157 if (self->GetDataType() == op::DataType::DT_BF16) {
156 selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_FLOAT, uniqueExecutor.get());158 selfContiguous = l0op::Cast(selfContiguous, op::DataType::DT_FLOAT, uniqueExecutor.get());
157 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);159 CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
158 }160 }
159 161 
160- // 调用l0算子Cholesky进行计算162+ const aclTensor* choleskyResult = nullptr;
161- auto choleskyResult = l0op::Cholesky(selfContiguous, true, uniqueExecutor.get());163+ 
162- CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);164+ // upper是true并且最后一维大于10时，需要进行转置
165+ if (upper && last_dim_size > THRESHOLD_VALUE) {
166+ selfContiguous = SwapDim(selfContiguous, uniqueExecutor.get());
167+ CHECK_RET(selfContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
168+ 
169+ // 调用l0算子Cholesky进行计算
170+ choleskyResult = l0op::Cholesky(selfContiguous, false, uniqueExecutor.get());
171+ CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
172+ 
173+ choleskyResult = SwapDim(choleskyResult, uniqueExecutor.get());
174+ CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
175+ } else {
176+ // 调用l0算子Cholesky进行计算
177+ choleskyResult = l0op::Cholesky(selfContiguous, upper, uniqueExecutor.get());
178+ CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
179+ }
163 180 
164 // 输入类型是bf16时，需要将结果重新转换回去181 // 输入类型是bf16时，需要将结果重新转换回去
165 if (self->GetDataType() == op::DataType::DT_BF16) {182 if (self->GetDataType() == op::DataType::DT_BF16) {
@@ -167,12 +184,6 @@ aclnnStatus aclnnLinalgCholeskyGetWorkspaceSize(const aclTensor *self, bool uppe
167 CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);184 CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
168 }185 }
169 186 
170- // upper是false时，需要将结果重新转置回去
171- if (!upper) {
172- choleskyResult = SwapDim(choleskyResult, uniqueExecutor.get());
173- CHECK_RET(choleskyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
174- }
175- 
176 // 将结果拷贝到out187 // 将结果拷贝到out
177 auto viewCopyResult = l0op::ViewCopy(choleskyResult, out, uniqueExecutor.get());188 auto viewCopyResult = l0op::ViewCopy(choleskyResult, out, uniqueExecutor.get());
178 CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);189 CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
Mmath/cholesky/op_kernel/cholesky.h+372-233
@@ -16,7 +16,7 @@
16using namespace AscendC;16using namespace AscendC;
17 17 
18namespace Cholesky {18namespace Cholesky {
19-constexpr int32_t BUFFER_NUM = 1;19+constexpr uint32_t BUFFER_NUM = 1;
20constexpr uint32_t BASIC_BLOCK = 32;20constexpr uint32_t BASIC_BLOCK = 32;
21 21 
22template <typename T>22template <typename T>
@@ -29,23 +29,31 @@ public:
29 __aicore__ inline void ProcessTriu();29 __aicore__ inline void ProcessTriu();
30 30 
31private:31private:
32+ __aicore__ inline void PIPE_V_S();
33+ __aicore__ inline void PIPE_MTE2_S();
34+ __aicore__ inline void PIPE_MTE3_S();
35+ __aicore__ inline void PIPE_S_MTE3();
32 __aicore__ inline void GetTilingData(const CholeskyTilingData* tilingData);36 __aicore__ inline void GetTilingData(const CholeskyTilingData* tilingData);
33- __aicore__ inline void FirstColumn(uint32_t offsetPrefix, uint32_t matrixoffset);37+ __aicore__ inline void FirstColumn(uint64_t offsetPrefix, uint64_t offset);
34- __aicore__ inline void FirstRow(uint32_t offsetPrefix, uint32_t matrixoffset);38+ __aicore__ inline void SecondToNColumn(uint32_t index, uint64_t offsetPrefix, uint64_t offset);
35- __aicore__ inline void SecondToNColumn(uint32_t index, uint32_t offsetPrefix, uint32_t matrixoffset);39+ __aicore__ inline void FirstRow(uint64_t offsetPrefix, uint64_t offset);
36- __aicore__ inline void SecondToNRow(uint32_t index, uint32_t offsetPrefix, uint32_t matrixoffset);40+ __aicore__ inline void SecondToNRow(uint32_t index, uint64_t offsetPrefix, uint64_t offset);
37 41 
38 template <typename T1, typename T2>42 template <typename T1, typename T2>
39- __aicore__ inline T1 CeilA2B(T1 a, T2 b) {43+ __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
40 return b == 0 ? a : (a + b -1) / b;44 return b == 0 ? a : (a + b -1) / b;
41 }45 }
42 46 
43private:47private:
44- uint32_t matSizeN = 0;48+ uint32_t blockIdx_ = 0;
45- uint32_t matrixNumCount = 0;49+ uint32_t blockDim_ = 0;
46- uint32_t maxDataCount = 0;50+ uint32_t matSizeN_ = 0;
47- int32_t blockIdx = 0;51+ uint64_t matrixNumCount_ = 0;
48- int32_t numBlocks = 0;52+ uint64_t maxDataCount_ = 0;
53+ uint32_t blockSize_ = 0;
54+ uint32_t blockNum_ = 0;
55+ T inv_sqrt_A11_ = 0.0f; // 存储缩放因子，避免重复计算和直接访问GM内存
56+ T ZERO = 0.0f;
49 57 
50 TQue<QuePosition::VECIN, BUFFER_NUM> matAQueue;58 TQue<QuePosition::VECIN, BUFFER_NUM> matAQueue;
51 TQue<QuePosition::VECIN, BUFFER_NUM> matLeftQueue;59 TQue<QuePosition::VECIN, BUFFER_NUM> matLeftQueue;
@@ -55,19 +63,65 @@ private:
55 63 
56 GlobalTensor<T> matAGM;64 GlobalTensor<T> matAGM;
57 GlobalTensor<T> outGM;65 GlobalTensor<T> outGM;
66+ 
67+ // 辅助函数声明
68+ __aicore__ inline void ProcessColumnDotProduct(LocalTensor<T>& matLLocal, LocalTensor<T>& matLeftLocal, LocalTensor<T>& matRightLocal, LocalTensor<T>& matResultLocal,
69+ uint32_t index, uint64_t offset, uint32_t blockStart, uint32_t count);
70+
71+ __aicore__ inline void ProcessRowDotProduct(LocalTensor<T>& matLLocal, LocalTensor<T>& matLeftLocal, LocalTensor<T>& matRightLocal, LocalTensor<T>& matResultLocal,
72+ uint32_t index, uint64_t offset, uint32_t blockStart, uint32_t count);
73+
74+ __aicore__ inline T ComputeScaleFactor(LocalTensor<T>& matLLocal, uint64_t offsetPrefix, uint32_t index);
58};75};
59 76 
77+template <typename T>
78+__aicore__ inline void Cholesky<T>::PIPE_V_S() {
79+ event_t eventIDVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
80+ SetFlag<HardEvent::V_S>(eventIDVToS);
81+ WaitFlag<HardEvent::V_S>(eventIDVToS);
82+}
83+ 
84+template <typename T>
85+__aicore__ inline void Cholesky<T>::PIPE_MTE2_S() {
86+ event_t eventIDMTE2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
87+ SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
88+ WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS);
89+}
90+ 
91+template <typename T>
92+__aicore__ inline void Cholesky<T>::PIPE_MTE3_S() {
93+ event_t eventIDMTE3ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
94+ SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
95+ WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
96+}
97+ 
98+template <typename T>
99+__aicore__ inline void Cholesky<T>::PIPE_S_MTE3() {
100+ event_t eventIDSToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
101+ SetFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
102+ WaitFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
103+}
104+ 
105+template <typename T>
106+__aicore__ inline void Cholesky<T>::GetTilingData(const CholeskyTilingData* tilingData) {
107+ matSizeN_ = tilingData->matSizeN;
108+ matrixNumCount_ = tilingData->matrixNumCount;
109+ blockSize_ = tilingData->blockSize;
110+ blockNum_ = tilingData->blockNum;
111+}
112+ 
60template <typename T>113template <typename T>
61__aicore__ inline void Cholesky<T>::InitTril(GM_ADDR self, GM_ADDR out, GM_ADDR workspace, const CholeskyTilingData* tilingData, TPipe* pipe) {114__aicore__ inline void Cholesky<T>::InitTril(GM_ADDR self, GM_ADDR out, GM_ADDR workspace, const CholeskyTilingData* tilingData, TPipe* pipe) {
62- blockIdx = GetBlockIdx();115+ blockIdx_ = GetBlockIdx();
63- numBlocks = GetBlockNum();116+ blockDim_ = GetBlockNum();
64 GetTilingData(tilingData);117 GetTilingData(tilingData);
65 118
66- matAGM.SetGlobalBuffer((__gm__ T*)self, matSizeN * matSizeN);119+ matAGM.SetGlobalBuffer((__gm__ T*)self, matSizeN_ * matSizeN_);
67- outGM.SetGlobalBuffer((__gm__ T*)out, matSizeN * matSizeN);120+ outGM.SetGlobalBuffer((__gm__ T*)out, matSizeN_ * matSizeN_);
68 121 
69- uint32_t columnBufferSize = matSizeN * BASIC_BLOCK;122+ // 使用分块大小计算buffer，减少UB内存使用
70- uint32_t rowBufferSize = CeilA2B(matSizeN * sizeof(T), BASIC_BLOCK) * BASIC_BLOCK;123+ uint64_t columnBufferSize = blockSize_ * BASIC_BLOCK;
124+ uint64_t rowBufferSize = CeilDiv(blockSize_ * sizeof(T), BASIC_BLOCK) * BASIC_BLOCK;
71 125 
72 pipe->InitBuffer(matAQueue, BUFFER_NUM, columnBufferSize);126 pipe->InitBuffer(matAQueue, BUFFER_NUM, columnBufferSize);
73 pipe->InitBuffer(matLQueue, BUFFER_NUM, columnBufferSize);127 pipe->InitBuffer(matLQueue, BUFFER_NUM, columnBufferSize);
@@ -75,20 +129,225 @@ __aicore__ inline void Cholesky<T>::InitTril(GM_ADDR self, GM_ADDR out, GM_ADDR
75 pipe->InitBuffer(matRightQueue, BUFFER_NUM, rowBufferSize);129 pipe->InitBuffer(matRightQueue, BUFFER_NUM, rowBufferSize);
76 pipe->InitBuffer(matResultQueue, BUFFER_NUM, rowBufferSize);130 pipe->InitBuffer(matResultQueue, BUFFER_NUM, rowBufferSize);
77 131 
78- maxDataCount = CeilA2B(matSizeN, BASIC_BLOCK) * BASIC_BLOCK;132+ maxDataCount_ = CeilDiv(blockSize_, BASIC_BLOCK) * BASIC_BLOCK;
133+}
134+ 
135+template <typename T>
136+__aicore__ inline void Cholesky<T>::ProcessTril() {
137+ if (blockIdx_ < blockDim_) {
138+ auto loopTimes = matrixNumCount_ / blockDim_;
139+ for (uint64_t loopIndex = 0; loopIndex <= loopTimes; loopIndex++) {
140+ uint64_t offsetPrefix = blockIdx_ + blockDim_ * loopIndex;
141+ if (offsetPrefix < matrixNumCount_) {
142+ uint64_t offset = offsetPrefix * matSizeN_ * matSizeN_;
143+ FirstColumn(offsetPrefix, offset);
144+ for (uint32_t index = 1; index < matSizeN_; index++) {
145+ SecondToNColumn(index, offsetPrefix, offset);
146+ }
147+ }
148+ }
149+ }
150+}
151+ 
152+template <typename T>
153+__aicore__ inline void Cholesky<T>::FirstColumn(uint64_t offsetPrefix, uint64_t offset) {
154+ LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
155+
156+ // 核内分块处理，每次处理blockSize大小的数据
157+ for (uint32_t blockStart = 0; blockStart < matSizeN_; blockStart += blockSize_) {
158+ uint32_t count = (matSizeN_ - blockStart) > blockSize_ ? blockSize_ : (matSizeN_ - blockStart);
159+
160+ DataCopyExtParams copyParamsMatALocal {static_cast<uint16_t>(count), sizeof(T), static_cast<uint32_t>((matSizeN_ - 1) * sizeof(T)), 0, 0};
161+ DataCopyPadExtParams<T> padParamsMatALocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};
162+ DataCopyPad(matALocal, matAGM[offset + blockStart * matSizeN_], copyParamsMatALocal, padParamsMatALocal);
163+ PIPE_MTE2_S();
164+
165+ // 只在处理第一个元素时计算平方根并存储缩放因子
166+ if (blockStart == 0) {
167+ T A11 = matALocal.GetValue(0);
168+ PIPE_V_S();
169+ if (matrixNumCount_ > 1) {
170+ ascendc_assert(A11 > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n", offsetPrefix);
171+ } else {
172+ ascendc_assert(A11 > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n");
173+ }
174+ inv_sqrt_A11_ = T(1/sqrt(A11));
175+ Muls(matALocal, matALocal, inv_sqrt_A11_, count * BASIC_BLOCK / sizeof(T));
176+ } else {
177+ // 直接使用之前计算好的缩放因子，避免访问GM内存
178+ Muls(matALocal, matALocal, inv_sqrt_A11_, count * BASIC_BLOCK / sizeof(T));
179+ }
180+ 
181+ PIPE_S_MTE3();
182+ DataCopyExtParams dataCopyOutParams {static_cast<uint16_t>(count), sizeof(T), 0, static_cast<uint32_t>((matSizeN_ - 1) * sizeof(T)), 0};
183+ DataCopyPad(outGM[offset + blockStart * matSizeN_], matALocal, dataCopyOutParams);
184+ PIPE_MTE3_S();
185+ }
186+
187+ matAQueue.FreeTensor(matALocal);
188+}
189+ 
190+// 辅助函数：处理SecondToNColumn中的点积计算部分
191+template <typename T>
192+__aicore__ inline void Cholesky<T>::ProcessColumnDotProduct(LocalTensor<T>& matLLocal, LocalTensor<T>& matLeftLocal, LocalTensor<T>& matRightLocal, LocalTensor<T>& matResultLocal,
193+ uint32_t index, uint64_t offset, uint32_t blockStart, uint32_t count) {
194+ // 计算左侧分块的数量
195+ uint32_t leftBlockNum = (index + blockSize_ - 1) / blockSize_;
196+
197+ for (uint32_t leftBlockIdx = 0; leftBlockIdx < leftBlockNum; leftBlockIdx++) {
198+ // 计算leftlocal当前块的起始位置和大小，分块大小为blockSize_
199+ uint32_t leftBlockStart = leftBlockIdx * blockSize_;
200+ uint32_t leftBlockSize = (index - leftBlockStart) > blockSize_ ? blockSize_ : (index - leftBlockStart);
201+
202+ // 搬运当前块的matLeftLocal数据
203+ DataCopyExtParams copyParamsLeftLocal {1, static_cast<uint32_t>(sizeof(T) * leftBlockSize), 0, 0, 0};
204+ DataCopyPadExtParams<T> padParamsLeftLocal {false, 0, 0, 0};
205+ DataCopyPad(matLeftLocal, outGM[offset + index * matSizeN_ + leftBlockStart], copyParamsLeftLocal, padParamsLeftLocal);
206+ PIPE_MTE2_S();
207+ 
208+ // 对当前块中的每一行，分块搬运matRightLocal并计算
209+ for (uint32_t row_in_block = 0; row_in_block < count; row_in_block++) {
210+ uint32_t row_below_pivot = blockStart + row_in_block;
211+
212+ // 搬运当前块的matRightLocal数据
213+ DataCopyPad(matRightLocal, outGM[offset + (index + row_below_pivot) * matSizeN_ + leftBlockStart], copyParamsLeftLocal, padParamsLeftLocal);
214+ PIPE_MTE2_S();
215+
216+ // 计算当前块的点积并累加结果
217+ Mul(matResultLocal, matLeftLocal, matRightLocal, leftBlockSize);
218+ ReduceSum<T>(matResultLocal, matResultLocal, matResultLocal, leftBlockSize);
219+
220+ // 将当前块的结果累加到matLLocal中
221+ T currentSum = matResultLocal.GetValue(0);
222+ T existingSum = matLLocal.GetValue(row_in_block * BASIC_BLOCK / sizeof(T));
223+ matLLocal.SetValue(row_in_block * BASIC_BLOCK / sizeof(T), existingSum + currentSum);
224+ PIPE_V_S();
225+ }
226+ }
227+}
228+ 
229+// 辅助函数：计算缩放因子并进行正定性检查
230+template <typename T>
231+__aicore__ inline T Cholesky<T>::ComputeScaleFactor(LocalTensor<T>& matLLocal, uint64_t offsetPrefix, uint32_t index) {
232+ T b1 = matLLocal.GetValue(0);
233+ PIPE_V_S();
234+ if (matrixNumCount_ > 1) {
235+ ascendc_assert(b1 > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", offsetPrefix, index + 1);
236+ } else {
237+ ascendc_assert(b1 > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", index + 1);
238+ }
239+
240+ // 计算缩放因子
241+ return T(1/sqrt(b1));
242+}
243+ 
244+// 辅助函数：处理SecondToNRow中的点积计算部分
245+template <typename T>
246+__aicore__ inline void Cholesky<T>::ProcessRowDotProduct(LocalTensor<T>& matLLocal, LocalTensor<T>& matLeftLocal, LocalTensor<T>& matRightLocal, LocalTensor<T>& matResultLocal,
247+ uint32_t index, uint64_t offset, uint32_t blockStart, uint32_t count) {
248+ // 计算左侧分块的数量
249+ uint32_t leftBlockNum = (index + blockSize_ - 1) / blockSize_;
250+
251+ for (uint32_t leftBlockIdx = 0; leftBlockIdx < leftBlockNum; leftBlockIdx++) {
252+ // 计算leftlocal当前块的起始位置和大小，分块大小为blockSize_
253+ uint32_t leftBlockStart = leftBlockIdx * blockSize_;
254+ uint32_t leftBlockSize = (index - leftBlockStart) > blockSize_ ? blockSize_ : (index - leftBlockStart);
255+
256+ // 搬运当前块的matLeftLocal数据
257+ DataCopyExtParams copyParamsLeftLocal {static_cast<uint16_t>(leftBlockSize), sizeof(T), static_cast<uint32_t>((matSizeN_ - 1) * sizeof(T)), 0, 0};
258+ DataCopyPadExtParams<T> padParamsLeftLocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};
259+ DataCopyPad(matLeftLocal, outGM[offset + index + leftBlockStart * matSizeN_], copyParamsLeftLocal, padParamsLeftLocal);
260+ PIPE_MTE2_S();
261+ 
262+ // 对当前块中的每一列，分块搬运matRightLocal并计算
263+ for (uint32_t col_in_block = 0; col_in_block < count; col_in_block++) {
264+ uint32_t column_right_pivot = blockStart + col_in_block;
265+
266+ // 搬运当前块的matRightLocal数据
267+ DataCopyPad(matRightLocal, outGM[offset + (index + column_right_pivot) + leftBlockStart * matSizeN_], copyParamsLeftLocal, padParamsLeftLocal);
268+ PIPE_MTE2_S();
269+
270+ // 计算当前块的点积并累加结果
271+ Mul(matResultLocal, matLeftLocal, matRightLocal, leftBlockSize * BASIC_BLOCK / sizeof(T));
272+ ReduceSum<T>(matResultLocal, matResultLocal, matResultLocal, leftBlockSize * BASIC_BLOCK / sizeof(T));
273+
274+ // 将当前块的结果累加到matLLocal中
275+ T currentSum = matResultLocal.GetValue(0);
276+ T existingSum = matLLocal.GetValue(col_in_block);
277+ matLLocal.SetValue(col_in_block, existingSum + currentSum);
278+ PIPE_V_S();
279+ }
280+ }
281+}
282+ 
283+template <typename T>
284+__aicore__ inline void Cholesky<T>::SecondToNColumn(uint32_t index, uint64_t offsetPrefix, uint64_t offset) {
285+ LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
286+ LocalTensor<T> matLLocal = matLQueue.AllocTensor<T>();
287+ LocalTensor<T> matLeftLocal = matLeftQueue.AllocTensor<T>();
288+ LocalTensor<T> matRightLocal = matRightQueue.AllocTensor<T>();
289+ LocalTensor<T> matResultLocal = matResultQueue.AllocTensor<T>();
290+
291+ // 存储当前列的缩放因子，所有分块共享同一个缩放因子
292+ T column_scale_factor = 0.0f;
293+ bool scale_factor_computed = false;
294+
295+ // 对当前列的所有元素进行分块处理
296+ for (uint32_t blockStart = 0; blockStart < (matSizeN_ - index); blockStart += blockSize_) {
297+ // 计算当前块的大小
298+ uint32_t count = (matSizeN_ - index - blockStart) > blockSize_ ? blockSize_ : (matSizeN_ - index - blockStart);
299+
300+ // 1. 先搬运当前块的count个A元素进来
301+ DataCopyExtParams copyParamsMatALocal {static_cast<uint16_t>(count), sizeof(T), static_cast<uint32_t>((matSizeN_ - 1) * sizeof(T)), 0, 0};
302+ DataCopyPadExtParams<T> padParamsMatALocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};
303+ DataCopyPad(matALocal, matAGM[offset + index * matSizeN_ + index + blockStart * matSizeN_], copyParamsMatALocal, padParamsMatALocal);
304+ PIPE_MTE2_S();
305+
306+ // 2. 初始化当前块的L结果为0
307+ Duplicate(matLLocal, ZERO, count * BASIC_BLOCK / sizeof(T));
308+
309+ // 3. 调用辅助函数处理点积计算
310+ ProcessColumnDotProduct(matLLocal, matLeftLocal, matRightLocal, matResultLocal, index, offset, blockStart, count);
311+ 
312+ // 4. 执行计算操作
313+ Sub(matLLocal, matALocal, matLLocal, count * BASIC_BLOCK / sizeof(T));
314+
315+ // 只在第一次分块时计算缩放因子和进行正定性检查
316+ if (blockStart == 0) {
317+ column_scale_factor = ComputeScaleFactor(matLLocal, offsetPrefix, index);
318+ scale_factor_computed = true;
319+ }
320+
321+ // 对当前块的所有元素应用同一个缩放因子
322+ Muls(matLLocal, matLLocal, column_scale_factor, count * BASIC_BLOCK / sizeof(T));
323+ 
324+ // 5. 最后得到count个L元素并搬出
325+ PIPE_S_MTE3();
326+ DataCopyExtParams dataCopyOutParams {static_cast<uint16_t>(count), sizeof(T), 0, static_cast<uint32_t>((matSizeN_ - 1) * sizeof(T)), 0};
327+ DataCopyPad(outGM[offset + index * matSizeN_ + index + blockStart * matSizeN_], matLLocal, dataCopyOutParams);
328+ PIPE_MTE3_S();
329+ }
330+
331+ // 释放张量资源
332+ matResultQueue.FreeTensor(matResultLocal);
333+ matRightQueue.FreeTensor(matRightLocal);
334+ matLeftQueue.FreeTensor(matLeftLocal);
335+ matLQueue.FreeTensor(matLLocal);
336+ matAQueue.FreeTensor(matALocal);
79}337}
80 338 
81template <typename T>339template <typename T>
82__aicore__ inline void Cholesky<T>::InitTriu(GM_ADDR self, GM_ADDR out, GM_ADDR workspace, const CholeskyTilingData* tilingData, TPipe* pipe) {340__aicore__ inline void Cholesky<T>::InitTriu(GM_ADDR self, GM_ADDR out, GM_ADDR workspace, const CholeskyTilingData* tilingData, TPipe* pipe) {
83- blockIdx = GetBlockIdx();341+ blockIdx_ = GetBlockIdx();
84- numBlocks = GetBlockNum();342+ blockDim_ = GetBlockNum();
85 GetTilingData(tilingData);343 GetTilingData(tilingData);
86 344
87- matAGM.SetGlobalBuffer((__gm__ T*)self, matSizeN * matSizeN);345+ matAGM.SetGlobalBuffer((__gm__ T*)self, matSizeN_ * matSizeN_);
88- outGM.SetGlobalBuffer((__gm__ T*)out, matSizeN * matSizeN);346+ outGM.SetGlobalBuffer((__gm__ T*)out, matSizeN_ * matSizeN_);
89 347 
90- uint32_t columnBufferSize = matSizeN * BASIC_BLOCK;348+ // 使用分块大小计算buffer，减少UB内存使用
91- uint32_t rowBufferSize = CeilA2B(matSizeN * sizeof(T), BASIC_BLOCK) * BASIC_BLOCK;349+ uint32_t columnBufferSize = blockSize_ * BASIC_BLOCK;
350+ uint32_t rowBufferSize = CeilDiv(blockSize_ * sizeof(T), BASIC_BLOCK) * BASIC_BLOCK;
92 351 
93 pipe->InitBuffer(matAQueue, BUFFER_NUM, rowBufferSize);352 pipe->InitBuffer(matAQueue, BUFFER_NUM, rowBufferSize);
94 pipe->InitBuffer(matLQueue, BUFFER_NUM, rowBufferSize);353 pipe->InitBuffer(matLQueue, BUFFER_NUM, rowBufferSize);
@@ -96,79 +355,20 @@ __aicore__ inline void Cholesky<T>::InitTriu(GM_ADDR self, GM_ADDR out, GM_ADDR
96 pipe->InitBuffer(matRightQueue, BUFFER_NUM, columnBufferSize);355 pipe->InitBuffer(matRightQueue, BUFFER_NUM, columnBufferSize);
97 pipe->InitBuffer(matResultQueue, BUFFER_NUM, columnBufferSize);356 pipe->InitBuffer(matResultQueue, BUFFER_NUM, columnBufferSize);
98 357 
99- maxDataCount = CeilA2B(matSizeN, BASIC_BLOCK) * BASIC_BLOCK;358+ maxDataCount_ = CeilDiv(blockSize_, BASIC_BLOCK) * BASIC_BLOCK;
100-}
101- 
102-template <typename T>
103-__aicore__ inline void Cholesky<T>::ProcessTril() {
104- if (blockIdx < numBlocks) {
105- auto loopTimes = matrixNumCount / numBlocks;
106- for (uint32_t loopIndex = 0; loopIndex <= loopTimes; loopIndex++) {
107- uint32_t offsetPrefix = blockIdx + numBlocks * loopIndex;
108- if (offsetPrefix < matrixNumCount) {
109- uint32_t offset = offsetPrefix * matSizeN * matSizeN;
110- FirstColumn(offsetPrefix, offset);
111- for (uint32_t index = 1; index < matSizeN; index++) {
112- LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
113- matAQueue.EnQue(matALocal);
114- 
115- LocalTensor<T> matLLocal = matLQueue.AllocTensor<T>();
116- matLQueue.EnQue(matLLocal);
117- 
118- LocalTensor<T> matLeftLocal = matLeftQueue.AllocTensor<T>();
119- matLeftQueue.EnQue(matLeftLocal);
120- 
121- LocalTensor<T> matRightLocal = matRightQueue.AllocTensor<T>();
122- matRightQueue.EnQue(matRightLocal);
123- 
124- LocalTensor<T> matResultLocal = matResultQueue.AllocTensor<T>();
125- matResultQueue.EnQue(matResultLocal);
126- 
127- SecondToNColumn(index, offsetPrefix, offset);
128- 
129- matResultQueue.FreeTensor(matResultLocal);
130- matRightQueue.FreeTensor(matRightLocal);
131- matLeftQueue.FreeTensor(matLeftLocal);
132- matLQueue.FreeTensor(matLLocal);
133- matAQueue.FreeTensor(matALocal);
134- }
135- }
136- }
137- }
138}359}
139 360 
140template <typename T>361template <typename T>
141__aicore__ inline void Cholesky<T>::ProcessTriu() {362__aicore__ inline void Cholesky<T>::ProcessTriu() {
142- if (blockIdx < numBlocks) {363+ if (blockIdx_ < blockDim_) {
143- auto loopTimes = matrixNumCount / numBlocks;364+ auto loopTimes = matrixNumCount_ / blockDim_;
144- for (uint32_t loopIndex = 0; loopIndex <= loopTimes; loopIndex++) {365+ for (uint64_t loopIndex = 0; loopIndex <= loopTimes; loopIndex++) {
145- uint32_t offsetPrefix = blockIdx + numBlocks * loopIndex;366+ uint64_t offsetPrefix = blockIdx_ + blockDim_ * loopIndex;
146- if (offsetPrefix < matrixNumCount) {367+ if (offsetPrefix < matrixNumCount_) {
147- uint32_t offset = offsetPrefix * matSizeN * matSizeN;368+ uint64_t offset = offsetPrefix * matSizeN_ * matSizeN_;
148 FirstRow(offsetPrefix, offset);369 FirstRow(offsetPrefix, offset);
149- for (uint32_t index = 1; index < matSizeN; index++) {370+ for (uint32_t index = 1; index < matSizeN_; index++) {
150- LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
151- matAQueue.EnQue(matALocal);
152- 
153- LocalTensor<T> matLLocal = matLQueue.AllocTensor<T>();
154- matLQueue.EnQue(matLLocal);
155- 
156- LocalTensor<T> matLeftLocal = matLeftQueue.AllocTensor<T>();
157- matLeftQueue.EnQue(matLeftLocal);
158- 
159- LocalTensor<T> matRightLocal = matRightQueue.AllocTensor<T>();
160- matRightQueue.EnQue(matRightLocal);
161- 
162- LocalTensor<T> matResultLocal = matResultQueue.AllocTensor<T>();
163- matResultQueue.EnQue(matResultLocal);
164- 
165 SecondToNRow(index, offsetPrefix, offset);371 SecondToNRow(index, offsetPrefix, offset);
166- 
167- matResultQueue.FreeTensor(matResultLocal);
168- matRightQueue.FreeTensor(matRightLocal);
169- matLeftQueue.FreeTensor(matLeftLocal);
170- matLQueue.FreeTensor(matLLocal);
171- matAQueue.FreeTensor(matALocal);
172 }372 }
173 }373 }
174 }374 }
@@ -176,160 +376,99 @@ __aicore__ inline void Cholesky<T>::ProcessTriu() {
176}376}
177 377 
178template <typename T>378template <typename T>
179-__aicore__ inline void Cholesky<T>::GetTilingData(const CholeskyTilingData* tilingData) {379+__aicore__ inline void Cholesky<T>::FirstRow(uint64_t offsetPrefix, uint64_t offset) {
180- matSizeN = tilingData->matSizeN;
181- matrixNumCount = tilingData->matrixNumCount;
182-}
183- 
184-template <typename T>
185-__aicore__ inline void Cholesky<T>::FirstColumn(uint32_t offsetPrefix, uint32_t matrixoffset) {
186 LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();380 LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
187- DataCopyParams copyParamsMatALocal {static_cast<uint16_t>(matSizeN), sizeof(T), static_cast<uint16_t>((matSizeN - 1) * sizeof(T)), 0};381+
188- DataCopyPadParams padParamsMatALocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};382+ // 核内分块处理，每次处理blockSize大小的数据
189- DataCopyPad(matALocal, matAGM[matrixoffset], copyParamsMatALocal, padParamsMatALocal);383+ for (uint32_t blockStart = 0; blockStart < matSizeN_; blockStart += blockSize_) {
190- PipeBarrier<PIPE_ALL>();384+ uint32_t count = (matSizeN_ - blockStart) > blockSize_ ? blockSize_ : (matSizeN_ - blockStart);
191- matAQueue.EnQue(matALocal);385+
386+ DataCopyExtParams copyParamsMatALocal {1, static_cast<uint32_t>(sizeof(T) * count), 0, 0, 0};
387+ DataCopyPadExtParams<T> padParamsMatALocal {false, 0, 0, 0};
388+ DataCopyPad(matALocal, matAGM[offset + blockStart], copyParamsMatALocal, padParamsMatALocal);
389+ PIPE_MTE2_S();
390+
391+ // 只在处理第一个元素时计算平方根并存储缩放因子
392+ if (blockStart == 0) {
393+ T A11_sqrt = matALocal.GetValue(0);
394+ PIPE_V_S();
395+ if (matrixNumCount_ > 1) {
396+ ascendc_assert(A11_sqrt > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n", offsetPrefix);
397+ } else {
398+ ascendc_assert(A11_sqrt > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n");
399+ }
400+ inv_sqrt_A11_ = T(1/sqrt(A11_sqrt));
401+ Muls(matALocal, matALocal, inv_sqrt_A11_, count);
402+ } else {
403+ // 使用之前存储的缩放因子，避免重复计算和直接访问GM内存
404+ Muls(matALocal, matALocal, inv_sqrt_A11_, count);
405+ }
192 406 
193- matALocal = matAQueue.DeQue<T>();407+ // 搬出当前块的结果
194- T A11_sqrt = matALocal.GetValue(0);408+ PIPE_S_MTE3();
195- if (matrixNumCount > 1) {409+ DataCopyExtParams dataCopyOutParams {1, static_cast<uint32_t>(sizeof(T) * count), 0, 0, 0};
196- ascendc_assert(A11_sqrt > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n", offsetPrefix);410+ DataCopyPad(outGM[offset + blockStart], matALocal, dataCopyOutParams);
197- } else {411+ PIPE_MTE3_S();
198- ascendc_assert(A11_sqrt > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n");
199 }412 }
200- A11_sqrt = sqrt(A11_sqrt);413+
201- PipeBarrier<PIPE_ALL>();
202- Muls(matALocal, matALocal, T(1/A11_sqrt), matSizeN * BASIC_BLOCK / sizeof(T));
203- PipeBarrier<PIPE_ALL>();
204- DataCopyParams dataCopyOutParams {static_cast<uint16_t>(matSizeN), sizeof(T), 0, static_cast<uint16_t>((matSizeN - 1) * sizeof(T))};
205- DataCopyPad(outGM[matrixoffset], matALocal, dataCopyOutParams);
206- PipeBarrier<PIPE_ALL>();
207 matAQueue.FreeTensor(matALocal);414 matAQueue.FreeTensor(matALocal);
208}415}
209 416 
210template <typename T>417template <typename T>
211-__aicore__ inline void Cholesky<T>::FirstRow(uint32_t offsetPrefix, uint32_t matrixoffset) {418+__aicore__ inline void Cholesky<T>::SecondToNRow(uint32_t index, uint64_t offsetPrefix, uint64_t offset) {
212 LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();419 LocalTensor<T> matALocal = matAQueue.AllocTensor<T>();
213- DataCopy(matALocal, matAGM[matrixoffset], maxDataCount);420+ LocalTensor<T> matLLocal = matLQueue.AllocTensor<T>();
214- PipeBarrier<PIPE_ALL>();421+ LocalTensor<T> matLeftLocal = matLeftQueue.AllocTensor<T>();
215- matAQueue.EnQue(matALocal);422+ LocalTensor<T> matRightLocal = matRightQueue.AllocTensor<T>();
423+ LocalTensor<T> matResultLocal = matResultQueue.AllocTensor<T>();
424+
425+ // 存储当前行的缩放因子，所有分块共享同一个缩放因子
426+ T row_scale_factor = 0.0f;
427+ bool scale_factor_computed = false;
428+
429+ // 对当前行的所有元素进行分块处理
430+ for (uint32_t blockStart = 0; blockStart < (matSizeN_ - index); blockStart += blockSize_) {
431+ // 计算当前块的大小
432+ uint32_t count = (matSizeN_ - index - blockStart) > blockSize_ ? blockSize_ : (matSizeN_ - index - blockStart);
216 433 
217- matALocal = matAQueue.DeQue<T>();434+ // 1. 先搬运当前块的count个A元素进来
218- T A11_sqrt = matALocal.GetValue(0);435+ DataCopyExtParams copyParamsMatALocal {1, static_cast<uint32_t>(sizeof(T) * count), 0, 0, 0};
219- PipeBarrier<PIPE_ALL>();436+ DataCopyPadExtParams<T> padParamsMatALocal {false, 0, 0, 0};
220- if (matrixNumCount > 1) {437+ DataCopyPad(matALocal, matAGM[offset + index * matSizeN_ + index + blockStart], copyParamsMatALocal, padParamsMatALocal);
221- ascendc_assert(A11_sqrt > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n", offsetPrefix);438+ PIPE_MTE2_S();
222- } else {439+ 
223- ascendc_assert(A11_sqrt > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order 1 is not positive-definite).\n");440+ // 2. 初始化当前块的L结果为0
441+ Duplicate(matLLocal, ZERO, count);
442+
443+ // 3. 调用辅助函数处理点积计算
444+ ProcessRowDotProduct(matLLocal, matLeftLocal, matRightLocal, matResultLocal, index, offset, blockStart, count);
445+ 
446+ // 4. 执行计算操作
447+ Sub(matLLocal, matALocal, matLLocal, count);
448+
449+ // 只在第一次分块时计算缩放因子和进行正定性检查
450+ if (blockStart == 0) {
451+ row_scale_factor = ComputeScaleFactor(matLLocal, offsetPrefix, index);
452+ scale_factor_computed = true;
453+ }
454+
455+ // 对当前块的所有元素应用同一个缩放因子
456+ Muls(matLLocal, matLLocal, row_scale_factor, count);
457+ 
458+ // 5. 最后得到count个L元素并搬出
459+ PIPE_S_MTE3();
460+ DataCopyExtParams dataCopyOutParams {1, static_cast<uint32_t>(sizeof(T) * count), 0, 0, 0};
461+ DataCopyPad(outGM[offset + index * matSizeN_ + index + blockStart], matLLocal, dataCopyOutParams);
462+ PIPE_MTE3_S();
224 }463 }
225- PipeBarrier<PIPE_ALL>();464+
226- A11_sqrt = sqrt(A11_sqrt);465+ // 释放张量资源
227- PipeBarrier<PIPE_ALL>();466+ matResultQueue.FreeTensor(matResultLocal);
228- Muls(matALocal, matALocal, T(1/A11_sqrt), matSizeN);467+ matRightQueue.FreeTensor(matRightLocal);
229- PipeBarrier<PIPE_ALL>();468+ matLeftQueue.FreeTensor(matLeftLocal);
230- DataCopyParams dataCopyOutParams {1, static_cast<uint16_t>(sizeof(T) * matSizeN), 0, 0};469+ matLQueue.FreeTensor(matLLocal);
231- DataCopyPad(outGM[matrixoffset], matALocal, dataCopyOutParams);
232- PipeBarrier<PIPE_ALL>();
233 matAQueue.FreeTensor(matALocal);470 matAQueue.FreeTensor(matALocal);
234}471}
235 472 
236-template <typename T>
237-__aicore__ inline void Cholesky<T>::SecondToNColumn(uint32_t index, uint32_t offsetPrefix, uint32_t matrixoffset) {
238- LocalTensor<T> matALocal = matAQueue.DeQue<T>();
239- DataCopyParams copyParamsMatALocal {static_cast<uint16_t>(matSizeN - index), sizeof(T), static_cast<uint16_t>((matSizeN - 1) * sizeof(T)), 0};
240- DataCopyPadParams padParamsMatALocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};
241- DataCopyPad(matALocal, matAGM[matrixoffset + index * matSizeN + index], copyParamsMatALocal, padParamsMatALocal);
242- PipeBarrier<PIPE_ALL>();
243-
244- LocalTensor<T> matLeftLocal = matLeftQueue.DeQue<T>();
245- DataCopyParams copyParamsMatLocal {1, static_cast<uint16_t>(sizeof(T) * index), 0, 0};
246- DataCopyPadParams padParamsMatLocal {false, 0, 0, 0};
247- DataCopyPad(matLeftLocal, outGM[matrixoffset + index * matSizeN], copyParamsMatLocal, padParamsMatLocal);
248- PipeBarrier<PIPE_ALL>();
249-
250- LocalTensor<T> matRightLocal = matRightQueue.DeQue<T>();
251- LocalTensor<T> matResultLocal = matResultQueue.DeQue<T>();
252- 
253- LocalTensor<T> matLLocal = matLQueue.DeQue<T>();
254- for (uint32_t i = 0; i < matSizeN - index; i++) {
255- PipeBarrier<PIPE_ALL>();
256- DataCopyPad(matRightLocal, outGM[matrixoffset + (index + i) * matSizeN], copyParamsMatLocal, padParamsMatLocal);
257- PipeBarrier<PIPE_ALL>();
258- Mul(matResultLocal, matLeftLocal, matRightLocal, index);
259- PipeBarrier<PIPE_ALL>();
260- ReduceSum<T>(matResultLocal, matResultLocal, matResultLocal, index);
261- PipeBarrier<PIPE_ALL>();
262- matLLocal.SetValue(i * BASIC_BLOCK / sizeof(T), matResultLocal.GetValue(0));
263- }
264- 
265- PipeBarrier<PIPE_ALL>();
266- Sub(matLLocal, matALocal, matLLocal, (matSizeN - index) * BASIC_BLOCK / sizeof(T));
267- PipeBarrier<PIPE_ALL>();
268- T b1 = matLLocal.GetValue(0);
269- if (matrixNumCount > 1) {
270- ascendc_assert(b1 > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", offsetPrefix, index + 1);
271- } else {
272- ascendc_assert(b1 > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", index + 1);
273- }
274- b1 = sqrt(b1);
275- PipeBarrier<PIPE_ALL>();
276- Muls(matLLocal, matLLocal, T(1/b1), (matSizeN - index) * BASIC_BLOCK / sizeof(T));
277- PipeBarrier<PIPE_ALL>();
278- DataCopyParams dataCopyOutParams {static_cast<uint16_t>(matSizeN - index), sizeof(T), 0, static_cast<uint16_t>((matSizeN - 1) * sizeof(T))};
279- DataCopyPad(outGM[matrixoffset + index * matSizeN + index], matLLocal, dataCopyOutParams);
280- PipeBarrier<PIPE_ALL>();
281-}
282- 
283-template <typename T>
284-__aicore__ inline void Cholesky<T>::SecondToNRow(uint32_t index, uint32_t offsetPrefix, uint32_t matrixoffset) {
285- LocalTensor<T> matALocal = matAQueue.DeQue<T>();
286- PipeBarrier<PIPE_ALL>();
287- DataCopyParams copyParamsMatALocal {1, static_cast<uint16_t>(sizeof(T) * (matSizeN - index)), 0, 0};
288- DataCopyPadParams padParamsMatALocal {false, 0, 0, 0};
289- DataCopyPad(matALocal, matAGM[matrixoffset + index * matSizeN + index], copyParamsMatALocal, padParamsMatALocal);
290- PipeBarrier<PIPE_ALL>();
291- 
292- LocalTensor<T> matLeftLocal = matLeftQueue.DeQue<T>();
293- DataCopyParams copyParamsMatLocal {static_cast<uint16_t>(index), sizeof(T), static_cast<uint16_t>((matSizeN - 1) * sizeof(T)), 0};
294- DataCopyPadParams padParamsMatLocal {true, 0, BASIC_BLOCK / sizeof(T) - 1, 0};
295- DataCopyPad(matLeftLocal, outGM[matrixoffset + index], copyParamsMatLocal, padParamsMatLocal);
296- PipeBarrier<PIPE_ALL>();
297- 
298- LocalTensor<T> matRightLocal = matRightQueue.DeQue<T>();
299- LocalTensor<T> matResultLocal = matResultQueue.DeQue<T>();
300- LocalTensor<T> matLLocal = matLQueue.DeQue<T>();
301- for (uint32_t i = 0; i < matSizeN - index; i++) {
302- PipeBarrier<PIPE_ALL>();
303- DataCopyPad(matRightLocal, outGM[matrixoffset + index + i], copyParamsMatLocal, padParamsMatLocal);
304- PipeBarrier<PIPE_ALL>();
305- Mul(matResultLocal, matLeftLocal, matRightLocal, index * BASIC_BLOCK / sizeof(T));
306- PipeBarrier<PIPE_ALL>();
307- ReduceSum<T>(matResultLocal, matResultLocal, matResultLocal, index * BASIC_BLOCK / sizeof(T));
308- PipeBarrier<PIPE_ALL>();
309- matLLocal.SetValue(i, matResultLocal.GetValue(0));
310- PipeBarrier<PIPE_ALL>();
311- }
312- 
313- PipeBarrier<PIPE_ALL>();
314- Sub(matLLocal, matALocal, matLLocal, matSizeN - index);
315-
316- PipeBarrier<PIPE_ALL>();
317- T b1 = matLLocal.GetValue(0);
318- PipeBarrier<PIPE_ALL>();
319- if (matrixNumCount > 1) {
320- ascendc_assert(b1 > 0.0f, "(Batch element %d): The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", offsetPrefix, index + 1);
321- } else {
322- ascendc_assert(b1 > 0.0f, "The factorization could not be completed because the input is not positive-definite (the leading minor of order %d is not positive-definite).\n", index + 1);
323- }
324- PipeBarrier<PIPE_ALL>();
325- b1 = sqrt(b1);
326- PipeBarrier<PIPE_ALL>();
327- Muls(matLLocal, matLLocal, T(1/b1), matSizeN - index);
328- PipeBarrier<PIPE_ALL>();
329- DataCopyParams dataCopyOutParams {1, static_cast<uint16_t>(sizeof(T) * (matSizeN - index)), 0, 0};
330- DataCopyPad(outGM[matrixoffset + index * matSizeN + index], matLLocal, dataCopyOutParams);
331- PipeBarrier<PIPE_ALL>();
332-}
333- 
334}473}
335#endif474#endif
Mmath/cholesky/tests/ut/op_host/test_cholesky_tiling.cpp+1-1
@@ -55,7 +55,7 @@ TEST_F(CholeskyTiling, cholesky_test_tiling_case0)
55 },55 },
56 &compileInfo);56 &compileInfo);
57 uint64_t expectTilingKey = 2;57 uint64_t expectTilingKey = 2;
58- string expectTilingData = "12884901894 ";58+ string expectTilingData = "6 3 4294967302 ";
59 std::vector<size_t> expectWorkspaces = {16777216};59 std::vector<size_t> expectWorkspaces = {16777216};
60 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingData, expectWorkspaces);60 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingData, expectWorkspaces);
61}61}