Star1662
3343
Code
Code
Issues453
Pull Requests701
Actions
Discussion
Wiki
Analysis
Star1662
3343
Merged
feat(mhc): 新增 TextGatedResidual 算子 #15380
kdy18482276080Add on 3 days ago
feat(mhc): 新增 TextGatedResidual 算子 #15380
Merged
kdy18482276080Add on 3 days ago
Merged fromkdy18482276080/ops-transformerfeat/text-gated-residual-deliveryintocann/ops-transformermaster
Total: 40 files changed+7704-0
Amhc/text_gated_residual/CMakeLists.txt+19-0
@@ -0,0 +1,19 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
11+if(NOT ENABLE_TEST AND NOT BENCHMARK)
12+ list(REMOVE_ITEM CURRENT_DIRS tests)
13+endif()
14+list(REMOVE_ITEM CURRENT_DIRS tmp)
15+foreach(SUB_DIR ${CURRENT_DIRS})
16+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
17+ add_subdirectory(${SUB_DIR})
18+ endif()
19+endforeach()
Amhc/text_gated_residual/examples/test_aclnn_text_gated_residual.cpp+90-0
@@ -0,0 +1,90 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <cstdint>
12+#include <cstdio>
13+#include <vector>
14+#include "acl/acl.h"
15+#include "aclnnop/aclnn_text_gated_residual.h"
16+ 
17+// FP16 zero inputs: mixed and inject must both be zero. No backward cache requested.
18+namespace {
19+constexpr int64_t EXAMPLE_SEQUENCE = 32;
20+constexpr int64_t EXAMPLE_CHANNELS = 4;
21+constexpr int64_t EXAMPLE_HIDDEN = 512;
22+constexpr size_t MIXED_ALLOCATION_INDEX = 4;
23+} // namespace
24+ 
25+int main()
26+{
27+ if (aclInit(nullptr) != ACL_SUCCESS || aclrtSetDevice(0) != ACL_SUCCESS)
28+ return 1;
29+ aclrtStream stream = nullptr;
30+ if (aclrtCreateStream(&stream) != ACL_SUCCESS)
31+ return 1;
32+ std::vector<aclTensor*> tensors;
33+ std::vector<void*> allocations;
34+ auto tensor = [&](std::vector<int64_t> shape) -> aclTensor* {
35+ std::vector<int64_t> strides(shape.size());
36+ int64_t elements = 1;
37+ for (int64_t i = static_cast<int64_t>(shape.size()) - 1; i >= 0; --i) {
38+ strides[i] = elements;
39+ elements *= shape[i];
40+ }
41+ void* data = nullptr;
42+ if (aclrtMalloc(&data, elements * 2, ACL_MEM_MALLOC_HUGE_FIRST) != ACL_SUCCESS)
43+ return nullptr;
44+ allocations.push_back(data);
45+ std::vector<uint16_t> zeros(elements, 0);
46+ if (aclrtMemcpy(data, elements * 2, zeros.data(), elements * 2, ACL_MEMCPY_HOST_TO_DEVICE) != ACL_SUCCESS)
47+ return nullptr;
48+ auto* t = aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT16, strides.data(), 0, ACL_FORMAT_ND,
49+ shape.data(), shape.size(), data);
50+ if (t)
51+ tensors.push_back(t);
52+ return t;
53+ };
54+ auto* x = tensor({1, 32, 2048});
55+ auto* down = tensor({64, 2048});
56+ auto* up = tensor({2048, 64});
57+ auto* gate = tensor({4, 2048});
58+ auto* mixed = tensor({1, 32, 512});
59+ auto* inject = tensor({1, 32, 4});
60+ uint64_t bytes = 0;
61+ aclOpExecutor* executor = nullptr;
62+ int result = 1;
63+ void* workspace = nullptr;
64+ if (x && down && up && gate && mixed && inject &&
65+ aclnnTextGatedResidualGetWorkspaceSize(x, down, up, gate, nullptr, EXAMPLE_CHANNELS, 1e-6, false, mixed, inject,
66+ nullptr, nullptr, nullptr, &bytes, &executor) == ACL_SUCCESS &&
67+ (bytes == 0 || aclrtMalloc(&workspace, bytes, ACL_MEM_MALLOC_HUGE_FIRST) == ACL_SUCCESS) &&
68+ aclnnTextGatedResidual(workspace, bytes, executor, stream) == ACL_SUCCESS &&
69+ aclrtSynchronizeStream(stream) == ACL_SUCCESS) {
70+ std::vector<uint16_t> output(EXAMPLE_SEQUENCE * EXAMPLE_HIDDEN);
71+ if (aclrtMemcpy(output.data(), output.size() * sizeof(uint16_t), allocations[MIXED_ALLOCATION_INDEX],
72+ output.size() * sizeof(uint16_t), ACL_MEMCPY_DEVICE_TO_HOST) == ACL_SUCCESS) {
73+ result = 0;
74+ for (auto value : output)
75+ if ((value & 0x7fff) != 0)
76+ result = 1;
77+ }
78+ }
79+ if (workspace)
80+ aclrtFree(workspace);
81+ for (auto* t : tensors)
82+ aclDestroyTensor(t);
83+ for (auto* p : allocations)
84+ aclrtFree(p);
85+ aclrtDestroyStream(stream);
86+ aclrtResetDevice(0);
87+ aclFinalize();
88+ std::printf("TextGatedResidual example: %s\n", result == 0 ? "PASS" : "FAIL");
89+ return result;
90+}
Amhc/text_gated_residual/op_host/CMakeLists.txt+26-0
@@ -0,0 +1,26 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+add_op_to_compiled_list()
11+if (BUILD_OPEN_PROJECT)
12+ target_sources(op_host_aclnnExc PRIVATE
13+ text_gated_residual_def.cpp
14+ )
15+endif()
16+ 
17+add_modules_sources(OPTYPE text_gated_residual ACLNNTYPE aclnn_exclude)
18+ 
19+set(HC_COMPILE_OPTIONS --cce-auto-sync=off -Wno-deprecated-declarations)
20+if (DEFINED CATLASS AND EXISTS "${CATLASS}/include/catlass/catlass.hpp")
21+ list(APPEND HC_COMPILE_OPTIONS "-I${CATLASS}/include")
22+endif()
23+add_ops_compile_options(
24+ OP_NAME TextGatedResidual
25+ OPTIONS ${HC_COMPILE_OPTIONS}
26+)
Amhc/text_gated_residual/op_host/op_api/aclnn_text_gated_residual.cpp+242-0
@@ -0,0 +1,242 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "aclnn_text_gated_residual.h"
12+#include "text_gated_residual.h"
13+ 
14+#include <cmath>
15+#include <limits>
16+#include "acl/acl.h"
17+#include "aclnn/aclnn_base.h"
18+#include "aclnn_kernels/common/op_error_check.h"
19+#include "aclnn_kernels/contiguous.h"
20+#include "opdev/make_op_executor.h"
21+#include "opdev/op_dfx.h"
22+#include "opdev/op_executor.h"
23+#include "opdev/op_log.h"
24+#include "opdev/platform.h"
25+#include "platform/soc_spec.h"
26+ 
27+namespace {
28+constexpr int64_t HC_ACTIVATION_RANK = 3;
29+constexpr int64_t HC_WEIGHT_RANK = 2;
30+constexpr int64_t HC_FEATURE_DIM = 2;
31+constexpr int64_t HC_INPUT_ELEMENT_BYTES = sizeof(uint16_t);
32+constexpr int64_t HC_ACCUMULATOR_BYTES = sizeof(float);
33+} // namespace
34+ 
35+using namespace op;
36+ 
37+namespace {
38+bool CheckNotNull(const aclTensor* x, const aclTensor* wDown, const aclTensor* wUp, const aclTensor* mixed)
39+{
40+ OP_CHECK_NULL(x, return false);
41+ OP_CHECK_NULL(wDown, return false);
42+ OP_CHECK_NULL(wUp, return false);
43+ OP_CHECK_NULL(mixed, return false);
44+ return true;
45+}
46+ 
47+aclnnStatus CheckParams(const aclTensor* x, const aclTensor* wDown, const aclTensor* wUp, const aclTensor* mixed,
48+ const aclTensor* wGate, bool useCombine, int64_t channels, bool needBackward,
49+ const aclTensor* invRmsOptional, const aclTensor* downPreActOptional,
50+ const aclTensor* upPreActOptional)
51+{
52+ CHECK_RET(CheckNotNull(x, wDown, wUp, mixed), ACLNN_ERR_PARAM_NULLPTR);
53+ CHECK_COND(channels > 0, ACLNN_ERR_PARAM_INVALID, "channels must be > 0.");
54+ CHECK_COND(x->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK, ACLNN_ERR_PARAM_INVALID,
55+ "x must be rank-3 [B,S,D].");
56+ CHECK_COND(wDown->GetViewShape().GetDimNum() == HC_WEIGHT_RANK, ACLNN_ERR_PARAM_INVALID,
57+ "w_down must be rank-2 [R,D].");
58+ CHECK_COND(wUp->GetViewShape().GetDimNum() == HC_WEIGHT_RANK, ACLNN_ERR_PARAM_INVALID,
59+ "w_up must be rank-2 [D,R].");
60+ const int64_t b = x->GetViewShape().GetDim(0);
61+ const int64_t s = x->GetViewShape().GetDim(1);
62+ const int64_t d = x->GetViewShape().GetDim(HC_FEATURE_DIM);
63+ const int64_t r = wDown->GetViewShape().GetDim(0);
64+ CHECK_COND(b >= 0 && s >= 0 && d > 0 && r > 0, ACLNN_ERR_PARAM_INVALID, "invalid nonpositive dimensions.");
65+ CHECK_COND(b <= UINT32_MAX && s <= UINT32_MAX && d <= UINT32_MAX && r <= UINT32_MAX && channels <= UINT32_MAX,
66+ ACLNN_ERR_PARAM_INVALID, "dimensions exceed tiling representation.");
67+ CHECK_COND(s == 0 || b <= UINT32_MAX / s, ACLNN_ERR_PARAM_INVALID, "B*S exceeds tiling representation.");
68+ const int64_t m = b * s; // B*S was checked against UINT32_MAX above.
69+ CHECK_COND(m <= INT64_MAX / (d * HC_ACCUMULATOR_BYTES) && m <= INT64_MAX / (r * HC_ACCUMULATOR_BYTES) &&
70+ r <= INT64_MAX / (d * HC_INPUT_ELEMENT_BYTES) &&
71+ channels <= INT64_MAX / (d * HC_INPUT_ELEMENT_BYTES),
72+ ACLNN_ERR_PARAM_INVALID, "tensor byte offsets exceed signed 64-bit representation.");
73+ CHECK_COND(d % channels == 0, ACLNN_ERR_PARAM_INVALID, "D must be divisible by channels.");
74+ CHECK_COND(mixed->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK && mixed->GetViewShape().GetDim(0) == b &&
75+ mixed->GetViewShape().GetDim(1) == s && mixed->GetViewShape().GetDim(HC_FEATURE_DIM) == d / channels,
76+ ACLNN_ERR_PARAM_INVALID, "mixed must have shape [B,S,H].");
77+ CHECK_COND(wDown->GetViewShape().GetDim(1) == d, ACLNN_ERR_PARAM_INVALID, "w_down last dim must equal D.");
78+ CHECK_COND(wUp->GetViewShape().GetDim(0) == d, ACLNN_ERR_PARAM_INVALID, "w_up first dim must equal D.");
79+ CHECK_COND(wUp->GetViewShape().GetDim(1) == r, ACLNN_ERR_PARAM_INVALID, "w_up last dim must equal R.");
80+ if (useCombine) {
81+ CHECK_COND(wGate != nullptr, ACLNN_ERR_PARAM_NULLPTR, "w_gate is required when w_gate is provided.");
82+ CHECK_COND(wGate->GetViewShape().GetDimNum() == HC_WEIGHT_RANK, ACLNN_ERR_PARAM_INVALID,
83+ "w_gate must be rank-2 [C,D].");
84+ CHECK_COND(wGate->GetViewShape().GetDim(0) == channels && wGate->GetViewShape().GetDim(1) == d,
85+ ACLNN_ERR_PARAM_INVALID, "w_gate must have shape [C,D].");
86+ }
87+ if (needBackward) {
88+ CHECK_COND(upPreActOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR,
89+ "upPreAct is required when needBackward is true.");
90+ CHECK_COND(upPreActOptional->GetDataType() == x->GetDataType(), ACLNN_ERR_PARAM_INVALID,
91+ "upPreAct dtype must match x.");
92+ CHECK_COND(upPreActOptional->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK, ACLNN_ERR_PARAM_INVALID,
93+ "upPreAct must be rank-3 [B,S,D].");
94+ CHECK_COND(upPreActOptional->GetViewShape().GetDim(0) == b && upPreActOptional->GetViewShape().GetDim(1) == s &&
95+ upPreActOptional->GetViewShape().GetDim(HC_FEATURE_DIM) == d,
96+ ACLNN_ERR_PARAM_INVALID, "upPreAct shape must be [B,S,D].");
97+ CHECK_COND(invRmsOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR, "invRms is required when needBackward is true.");
98+ CHECK_COND(downPreActOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR,
99+ "downPreAct is required when needBackward is true.");
100+ CHECK_COND(invRmsOptional->GetDataType() == DataType::DT_FLOAT, ACLNN_ERR_PARAM_INVALID,
101+ "invRms must be FLOAT32.");
102+ CHECK_COND(downPreActOptional->GetDataType() == x->GetDataType(), ACLNN_ERR_PARAM_INVALID,
103+ "downPreAct dtype must match x.");
104+ CHECK_COND(invRmsOptional->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK, ACLNN_ERR_PARAM_INVALID,
105+ "invRms must be rank-3 [B,S,C].");
106+ CHECK_COND(invRmsOptional->GetViewShape().GetDim(0) == b && invRmsOptional->GetViewShape().GetDim(1) == s &&
107+ invRmsOptional->GetViewShape().GetDim(HC_FEATURE_DIM) == channels,
108+ ACLNN_ERR_PARAM_INVALID, "invRms shape must be [B,S,C].");
109+ CHECK_COND(downPreActOptional->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK, ACLNN_ERR_PARAM_INVALID,
110+ "downPreAct must be rank-3 [B,S,R].");
111+ CHECK_COND(downPreActOptional->GetViewShape().GetDim(0) == b &&
112+ downPreActOptional->GetViewShape().GetDim(1) == s &&
113+ downPreActOptional->GetViewShape().GetDim(HC_FEATURE_DIM) == r,
114+ ACLNN_ERR_PARAM_INVALID, "downPreAct shape must be [B,S,R].");
115+ }
116+ return ACLNN_SUCCESS;
117+}
118+ 
119+bool IsDav2201()
120+{
121+ return GetCurrentPlatformInfo().GetCurNpuArch() == NpuArch::DAV_2201;
122+}
123+} // namespace
124+ 
125+#ifdef __cplusplus
126+extern "C" {
127+#endif
128+ 
129+aclnnStatus aclnnTextGatedResidualGetWorkspaceSize(const aclTensor* x, const aclTensor* wDown, const aclTensor* wUp,
130+ const aclTensor* wGateOptional, const aclTensor* gammaOptional,
131+ int64_t channels, double eps, bool needBackward,
132+ const aclTensor* mixed, const aclTensor* injectOptional,
133+ const aclTensor* invRmsOptional, const aclTensor* downPreActOptional,
134+ const aclTensor* upPreActOptional, uint64_t* workspaceSize,
135+ aclOpExecutor** executor)
136+{
137+ const bool useCombine = wGateOptional != nullptr;
138+ L2_DFX_PHASE_1(aclnnTextGatedResidual,
139+ DFX_IN(x, wDown, wUp, wGateOptional, gammaOptional, channels, eps, needBackward),
140+ DFX_OUT(mixed, injectOptional, invRmsOptional, downPreActOptional, upPreActOptional));
141+ const aclnnStatus paramStatus = CheckParams(x, wDown, wUp, mixed, wGateOptional, useCombine, channels, needBackward,
142+ invRmsOptional, downPreActOptional, upPreActOptional);
143+ if (paramStatus != ACLNN_SUCCESS) {
144+ return paramStatus;
145+ }
146+ if (gammaOptional != nullptr) {
147+ CHECK_COND(gammaOptional->GetDataType() == DataType::DT_FLOAT, ACLNN_ERR_PARAM_INVALID,
148+ "gamma must be FLOAT32.");
149+ CHECK_COND(gammaOptional->GetViewShape().GetDimNum() == 1 &&
150+ gammaOptional->GetViewShape().GetDim(0) == x->GetViewShape().GetDim(HC_FEATURE_DIM),
151+ ACLNN_ERR_PARAM_INVALID, "gamma must have shape [D=C*H].");
152+ }
153+ {
154+ CHECK_COND(std::isfinite(eps) && eps >= 0 && eps <= std::numeric_limits<float>::max(), ACLNN_ERR_PARAM_INVALID,
155+ "eps must be >= 0.");
156+ CHECK_COND(x->GetViewShape().GetDim(HC_FEATURE_DIM) != 0, ACLNN_ERR_PARAM_INVALID, "D must be > 0.");
157+ const auto xDtype = x->GetDataType();
158+ CHECK_COND(xDtype == DataType::DT_FLOAT16 || xDtype == DataType::DT_BF16, ACLNN_ERR_PARAM_INVALID,
159+ "x dtype must be FLOAT16 or BF16.");
160+ CHECK_COND(wDown->GetDataType() == xDtype && wUp->GetDataType() == xDtype && mixed->GetDataType() == xDtype,
161+ ACLNN_ERR_PARAM_INVALID, "weight and mixed dtype must match x.");
162+ if (wGateOptional != nullptr) {
163+ CHECK_COND(wGateOptional->GetDataType() == xDtype, ACLNN_ERR_PARAM_INVALID, "w_gate dtype must match x.");
164+ }
165+ if (gammaOptional != nullptr) {
166+ CHECK_COND(gammaOptional->GetDataType() == DataType::DT_FLOAT, ACLNN_ERR_PARAM_INVALID,
167+ "gamma must be FLOAT32.");
168+ }
169+ }
170+ if (x->GetViewShape().GetDim(0) == 0) {
171+ auto uniqueExecutor = CREATE_EXECUTOR();
172+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
173+ *workspaceSize = 0;
174+ uniqueExecutor.ReleaseTo(executor);
175+ return ACLNN_SUCCESS;
176+ }
177+ if (x->GetViewShape().GetDim(1) == 0) {
178+ auto uniqueExecutor = CREATE_EXECUTOR();
179+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
180+ *workspaceSize = 0;
181+ uniqueExecutor.ReleaseTo(executor);
182+ return ACLNN_SUCCESS;
183+ }
184+ CHECK_COND(mixed != nullptr, ACLNN_ERR_PARAM_NULLPTR, "mixed must not be nullptr.");
185+ if (useCombine) {
186+ CHECK_COND(injectOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR, "inject is required when w_gate is provided.");
187+ CHECK_COND(injectOptional->GetDataType() == x->GetDataType() &&
188+ injectOptional->GetViewShape().GetDimNum() == HC_ACTIVATION_RANK &&
189+ injectOptional->GetViewShape().GetDim(0) == x->GetViewShape().GetDim(0) &&
190+ injectOptional->GetViewShape().GetDim(1) == x->GetViewShape().GetDim(1) &&
191+ injectOptional->GetViewShape().GetDim(HC_FEATURE_DIM) == channels,
192+ ACLNN_ERR_PARAM_INVALID, "inject must match x dtype and shape [B,S,C].");
193+ }
194+ 
195+ auto uniqueExecutor = CREATE_EXECUTOR();
196+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
197+ auto executorPtr = uniqueExecutor.get();
198+ 
199+ auto xCont = l0op::Contiguous(x, executorPtr);
200+ auto wDownCont = l0op::Contiguous(wDown, executorPtr);
201+ auto wUpCont = l0op::Contiguous(wUp, executorPtr);
202+ CHECK_RET(xCont != nullptr && wDownCont != nullptr && wUpCont != nullptr, ACLNN_ERR_INNER_NULLPTR);
203+ const aclTensor* wGateCont = nullptr;
204+ if (wGateOptional != nullptr) {
205+ wGateCont = l0op::Contiguous(wGateOptional, executorPtr);
206+ CHECK_RET(wGateCont != nullptr, ACLNN_ERR_INNER_NULLPTR);
207+ }
208+ const aclTensor* gammaCont = nullptr;
209+ if (gammaOptional != nullptr) {
210+ gammaCont = l0op::Contiguous(gammaOptional, executorPtr);
211+ CHECK_RET(gammaCont != nullptr, ACLNN_ERR_INNER_NULLPTR);
212+ }
213+ 
214+ auto outputs = l0op::TextGatedResidual(xCont, wDownCont, wUpCont, wGateCont, gammaCont, channels,
215+ static_cast<float>(eps), needBackward, executorPtr);
216+ CHECK_RET(outputs[0] != nullptr, ACLNN_ERR_INNER_NULLPTR);
217+ CHECK_RET(l0op::ViewCopy(outputs[0], mixed, executorPtr) != nullptr, ACLNN_ERR_INNER_NULLPTR);
218+ if (useCombine) {
219+ CHECK_RET(l0op::ViewCopy(outputs[1], injectOptional, executorPtr) != nullptr, ACLNN_ERR_INNER_NULLPTR);
220+ }
221+ if (needBackward) {
222+ CHECK_RET(l0op::ViewCopy(outputs[2], invRmsOptional, executorPtr) != nullptr, ACLNN_ERR_INNER_NULLPTR);
223+ CHECK_RET(l0op::ViewCopy(outputs[3], downPreActOptional, executorPtr) != nullptr, ACLNN_ERR_INNER_NULLPTR);
224+ CHECK_RET(l0op::ViewCopy(outputs[4], upPreActOptional, executorPtr) != nullptr, ACLNN_ERR_INNER_NULLPTR);
225+ }
226+ 
227+ *workspaceSize = uniqueExecutor->GetWorkspaceSize();
228+ uniqueExecutor.ReleaseTo(executor);
229+ return ACLNN_SUCCESS;
230+}
231+ 
232+aclnnStatus aclnnTextGatedResidual(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)
233+{
234+ L2_DFX_PHASE_2(aclnnTextGatedResidual);
235+ CHECK_COND(CommonOpExecutorRun(workspace, workspaceSize, executor, stream) == ACLNN_SUCCESS, ACLNN_ERR_INNER,
236+ "TextGatedResidual launch failed.");
237+ return ACLNN_SUCCESS;
238+}
239+ 
240+#ifdef __cplusplus
241+}
242+#endif
Amhc/text_gated_residual/op_host/op_api/aclnn_text_gated_residual.h+64-0
@@ -0,0 +1,64 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#ifndef OP_API_INC_ACLNN_TEXT_GATED_RESIDUAL_H
11+#define OP_API_INC_ACLNN_TEXT_GATED_RESIDUAL_H
12+ 
13+#include "aclnn/aclnn_base.h"
14+#include "aclnn_util.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+/**
21+ * @brief 计算TextGatedResidual所需的工作空间大小。
22+ *
23+ * @param x [输入] 形状为[B, S, D]的ND张量，FLOAT16/BFLOAT16，D=C*H。
24+ * @param wDown [输入] 形状为[R, D]的ND张量，dtype须与x一致。
25+ * @param wUp [输入] 形状为[D, R]的ND张量，dtype须与x一致。
26+ * @param wGateOptional [输入] 可选的[C, D] ND张量，dtype须与x一致；为空时不计算inject。
27+ * @param gammaOptional [输入] 可选的[D] FLOAT32 ND张量；内部使用1+gamma缩放，为空时缩放因子为1。
28+ * @param channels [输入] 通道数C，须大于0且D可被C整除。
29+ * @param eps [输入] RMS归一化数值稳定性参数，须为非负有限DOUBLE且可表示为FLOAT，通常取1e-6。
30+ * @param needBackward [输入] 是否保存反向缓存；为true时invRmsOptional、downPreActOptional和upPreActOptional不能为空。
31+ * @param mixed [输出] 形状为[B, S, H]的ND张量，dtype须与x一致。
32+ * @param injectOptional [输出] 形状为[B, S, C]的ND张量，dtype须与x一致；wGateOptional非空时必须提供。
33+ * @param invRmsOptional [输出] 形状为[B, S, C]的FLOAT32 ND张量；needBackward为false时可传入nullptr。
34+ * @param downPreActOptional [输出] 形状为[B, S,
35+ * R]的ND张量，dtype须与x一致，保存未除以C的down投影；needBackward为false时可传入nullptr。
36+ * @param upPreActOptional [输出] 形状为[B, S,
37+ * D]的ND张量，dtype须与x一致，保存sigmoid之前的up投影；needBackward为false时可传入nullptr。
38+ * @param workspaceSize [输出] 设备侧所需的工作空间大小，单位为字节。
39+ * @param executor [输出] 第二段接口使用的执行器。
40+ * @return 成功时返回ACLNN_SUCCESS。
41+ */
42+ACLNN_API aclnnStatus aclnnTextGatedResidualGetWorkspaceSize(
43+ const aclTensor* x, const aclTensor* wDown, const aclTensor* wUp, const aclTensor* wGateOptional,
44+ const aclTensor* gammaOptional, int64_t channels, double eps, bool needBackward, const aclTensor* mixed,
45+ const aclTensor* injectOptional, const aclTensor* invRmsOptional, const aclTensor* downPreActOptional,
46+ const aclTensor* upPreActOptional, uint64_t* workspaceSize, aclOpExecutor** executor);
47+ 
48+/**
49+ * @brief 使用第一段接口生成的执行器运行TextGatedResidual。
50+ *
51+ * @param workspace [输入] Device侧workspace地址；workspaceSize为0时可传入nullptr。
52+ * @param workspaceSize [输入] 第一段接口返回的workspace大小，单位为字节。
53+ * @param executor [输入] 第一段接口返回的op执行器。
54+ * @param stream [输入] AscendCL Stream。
55+ * @return ACLNN_SUCCESS或对应错误码。
56+ */
57+ACLNN_API aclnnStatus aclnnTextGatedResidual(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
58+ aclrtStream stream);
59+ 
60+#ifdef __cplusplus
61+}
62+#endif
63+ 
64+#endif
Amhc/text_gated_residual/op_host/op_api/text_gated_residual.cpp+55-0
@@ -0,0 +1,55 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "text_gated_residual.h"
12+ 
13+#include "opdev/make_op_executor.h"
14+#include "opdev/op_dfx.h"
15+#include "opdev/op_log.h"
16+ 
17+using namespace op;
18+ 
19+namespace l0op {
20+OP_TYPE_REGISTER(TextGatedResidual);
21+ 
22+const std::array<const aclTensor*, 5> TextGatedResidual(const aclTensor* x, const aclTensor* wDown,
23+ const aclTensor* wUp, const aclTensor* wGateOptional,
24+ const aclTensor* gammaOptional, int64_t channels, float eps,
25+ bool needBackward, aclOpExecutor* executor)
26+{
27+ L0_DFX(TextGatedResidual, x, wDown, wUp, wGateOptional, gammaOptional, channels, eps, needBackward);
28+ const auto dtype = x->GetDataType();
29+ constexpr auto format = Format::FORMAT_ND;
30+ auto mixed = executor->AllocTensor(dtype, format, format);
31+ auto inject = executor->AllocTensor(dtype, format, format);
32+ auto invRms = executor->AllocTensor(DataType::DT_FLOAT, format, format);
33+ auto downPreAct = executor->AllocTensor(dtype, format, format);
34+ auto upPreAct = executor->AllocTensor(dtype, format, format);
35+ if (mixed == nullptr || inject == nullptr || invRms == nullptr || downPreAct == nullptr || upPreAct == nullptr) {
36+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "TextGatedResidual output allocation failed.");
37+ return {nullptr, nullptr, nullptr, nullptr, nullptr};
38+ }
39+ auto ret =
40+ INFER_SHAPE(TextGatedResidual, OP_INPUT(x, wDown, wUp, wGateOptional, gammaOptional),
41+ OP_OUTPUT(mixed, inject, invRms, downPreAct, upPreAct), OP_ATTR(channels, eps, needBackward));
42+ if (ret != ACLNN_SUCCESS) {
43+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "TextGatedResidual InferShape failed.");
44+ return {nullptr, nullptr, nullptr, nullptr, nullptr};
45+ }
46+ ret = ADD_TO_LAUNCHER_LIST_AICORE(TextGatedResidual, OP_INPUT(x, wDown, wUp, wGateOptional, gammaOptional),
47+ OP_OUTPUT(mixed, inject, invRms, downPreAct, upPreAct),
48+ OP_ATTR(channels, eps, needBackward));
49+ if (ret != ACLNN_SUCCESS) {
50+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "ADD_TO_LAUNCHER_LIST_AICORE TextGatedResidual failed.");
51+ return {nullptr, nullptr, nullptr, nullptr, nullptr};
52+ }
53+ return {mixed, inject, invRms, downPreAct, upPreAct};
54+}
55+} // namespace l0op
Amhc/text_gated_residual/op_host/op_api/text_gated_residual.h+23-0
@@ -0,0 +1,23 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#ifndef OP_API_INC_LEVEL0_OP_TEXT_GATED_RESIDUAL_OP_H
11+#define OP_API_INC_LEVEL0_OP_TEXT_GATED_RESIDUAL_OP_H
12+ 
13+#include <array>
14+#include "opdev/op_executor.h"
15+ 
16+namespace l0op {
17+const std::array<const aclTensor*, 5> TextGatedResidual(const aclTensor* x, const aclTensor* wDown,
18+ const aclTensor* wUp, const aclTensor* wGateOptional,
19+ const aclTensor* gammaOptional, int64_t channels, float eps,
20+ bool needBackward, aclOpExecutor* executor);
21+} // namespace l0op
22+ 
23+#endif
Amhc/text_gated_residual/op_host/text_gated_residual_def.cpp+96-0
@@ -0,0 +1,96 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "register/op_def_registry.h"
12+ 
13+namespace ops {
14+class TextGatedResidual : public OpDef {
15+public:
16+ explicit TextGatedResidual(const char* name)
17+ : OpDef(name)
18+ {
19+ this->Input("x")
20+ .ParamType(REQUIRED)
21+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
22+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
23+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
24+ .AutoContiguous();
25+ this->Input("w_down")
26+ .ParamType(REQUIRED)
27+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
28+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
29+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
30+ .AutoContiguous();
31+ this->Input("w_up")
32+ .ParamType(REQUIRED)
33+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
34+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
35+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
36+ .AutoContiguous();
37+ this->Input("w_gate")
38+ .ParamType(OPTIONAL)
39+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
40+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
41+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
42+ .AutoContiguous();
43+ this->Input("gamma")
44+ .ParamType(OPTIONAL)
45+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
46+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
47+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
48+ .AutoContiguous();
49+ this->Output("mixed")
50+ .ParamType(REQUIRED)
51+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
52+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
53+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
54+ this->Output("inject")
55+ .ParamType(OPTIONAL)
56+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
57+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
58+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
59+ this->Output("inv_rms")
60+ .ParamType(OPTIONAL)
61+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT})
62+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
63+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
64+ this->Output("down_pre_act")
65+ .ParamType(OPTIONAL)
66+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
67+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
68+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
69+ this->Output("up_pre_act")
70+ .ParamType(OPTIONAL)
71+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
72+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
73+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
74+ 
75+ this->Attr("channels").AttrType(REQUIRED).Int();
76+ this->Attr("eps").AttrType(OPTIONAL).Float(1e-6f);
77+ this->Attr("need_backward").AttrType(OPTIONAL).Bool(false);
78+ 
79+ OpAICoreConfig config;
80+ config.DynamicCompileStaticFlag(true)
81+ .DynamicFormatFlag(true)
82+ .DynamicRankSupportFlag(true)
83+ .DynamicShapeSupportFlag(true)
84+ .NeedCheckSupportFlag(false)
85+ .PrecisionReduceFlag(true)
86+ .ExtendCfgInfo("prebuildPattern.value", "Opaque")
87+ .ExtendCfgInfo("coreType.value", "AiCore")
88+ .ExtendCfgInfo("aclnnSupport.value", "support_aclnn");
89+ this->AICore().AddConfig("ascend950", config);
90+ this->AICore().AddConfig("ascend910b", config);
91+ this->AICore().AddConfig("ascend910_93", config);
92+ }
93+};
94+ 
95+OP_ADD(TextGatedResidual);
96+} // namespace ops
Amhc/text_gated_residual/op_host/text_gated_residual_infershape.cpp+110-0
@@ -0,0 +1,110 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "register/op_impl_registry.h"
12+#include "log/log.h"
13+ 
14+namespace {
15+constexpr int64_t HC_ACTIVATION_RANK = 3;
16+constexpr int64_t HC_WEIGHT_RANK = 2;
17+constexpr int64_t HC_FEATURE_DIM = 2;
18+constexpr int64_t HC_OUTPUT_UP = 4;
19+constexpr int64_t HC_OUTPUT_DOWN = 3;
20+constexpr int64_t HC_OUTPUT_INV = 2;
21+constexpr int64_t HC_OUTPUT_COUNT = 5;
22+} // namespace
23+ 
24+using namespace ge;
25+namespace ops {
26+static ge::graphStatus InferShape4TextGatedResidual(gert::InferShapeContext* context)
27+{
28+ const gert::Shape* xShape = context->GetInputShape(0);
29+ OP_CHECK_NULL_WITH_CONTEXT(context, xShape);
30+ auto attrs = context->GetAttrs();
31+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
32+ const int64_t* channelsPtr = attrs->GetAttrPointer<int64_t>(0);
33+ OP_CHECK_NULL_WITH_CONTEXT(context, channelsPtr);
34+ 
35+ gert::Shape* mixedShape = context->GetOutputShape(0);
36+ OP_CHECK_NULL_WITH_CONTEXT(context, mixedShape);
37+ if (xShape->GetDimNum() != HC_ACTIVATION_RANK || *channelsPtr <= 0) {
38+ return ge::GRAPH_FAILED;
39+ }
40+ const int64_t d = xShape->GetDim(HC_FEATURE_DIM);
41+ if (d % *channelsPtr != 0) {
42+ return ge::GRAPH_FAILED;
43+ }
44+ const auto* wd = context->GetInputShape(1);
45+ const auto* wu = context->GetInputShape(2);
46+ if (d <= 0 || wd == nullptr || wu == nullptr || wd->GetDimNum() != HC_WEIGHT_RANK ||
47+ wu->GetDimNum() != HC_WEIGHT_RANK || wd->GetDim(0) <= 0 || wd->GetDim(1) != d || wu->GetDim(0) != d ||
48+ wu->GetDim(1) != wd->GetDim(0))
49+ return ge::GRAPH_FAILED;
50+ const int64_t hidden = d / *channelsPtr;
51+ mixedShape->SetDimNum(HC_ACTIVATION_RANK);
52+ mixedShape->SetDim(0, xShape->GetDim(0));
53+ mixedShape->SetDim(1, xShape->GetDim(1));
54+ mixedShape->SetDim(HC_FEATURE_DIM, hidden);
55+ 
56+ const bool* needBackwardPtr = attrs->GetAttrPointer<bool>(2);
57+ const bool needBackward = needBackwardPtr != nullptr && *needBackwardPtr;
58+ const bool hasGate = context->GetOptionalInputShape(3) != nullptr;
59+ gert::Shape* injectShape = context->GetOutputShape(1);
60+ if (injectShape != nullptr) {
61+ injectShape->SetDimNum(HC_ACTIVATION_RANK);
62+ injectShape->SetDim(0, xShape->GetDim(0));
63+ injectShape->SetDim(1, xShape->GetDim(1));
64+ injectShape->SetDim(HC_FEATURE_DIM, *channelsPtr);
65+ }
66+ gert::Shape* invRmsShape = context->GetOutputShape(HC_OUTPUT_INV);
67+ if (invRmsShape != nullptr) {
68+ invRmsShape->SetDimNum(HC_ACTIVATION_RANK);
69+ invRmsShape->SetDim(0, xShape->GetDim(0));
70+ invRmsShape->SetDim(1, xShape->GetDim(1));
71+ invRmsShape->SetDim(HC_FEATURE_DIM, *channelsPtr);
72+ }
73+ const gert::Shape* wDownShape = context->GetInputShape(1);
74+ gert::Shape* downPreActShape = context->GetOutputShape(HC_OUTPUT_DOWN);
75+ if (downPreActShape != nullptr && wDownShape != nullptr && wDownShape->GetDimNum() >= 1) {
76+ downPreActShape->SetDimNum(HC_ACTIVATION_RANK);
77+ downPreActShape->SetDim(0, xShape->GetDim(0));
78+ downPreActShape->SetDim(1, xShape->GetDim(1));
79+ downPreActShape->SetDim(HC_FEATURE_DIM, wDownShape->GetDim(0));
80+ }
81+ auto* upPreActShape = context->GetOutputShape(HC_OUTPUT_UP);
82+ if (upPreActShape != nullptr) {
83+ *upPreActShape = *xShape;
84+ }
85+ for (size_t index = 1; index < HC_OUTPUT_COUNT; ++index) {
86+ if ((index == 1 && !hasGate) || (index >= HC_OUTPUT_INV && !needBackward)) {
87+ auto* disabledShape = context->GetOutputShape(index);
88+ if (disabledShape != nullptr) {
89+ disabledShape->SetDimNum(1);
90+ disabledShape->SetDim(0, 0);
91+ }
92+ }
93+ }
94+ return ge::GRAPH_SUCCESS;
95+}
96+ 
97+static ge::graphStatus InferDataType4TextGatedResidual(gert::InferDataTypeContext* context)
98+{
99+ context->SetOutputDataType(0, context->GetInputDataType(0));
100+ context->SetOutputDataType(1, context->GetInputDataType(0));
101+ context->SetOutputDataType(HC_OUTPUT_INV, ge::DT_FLOAT);
102+ context->SetOutputDataType(HC_OUTPUT_DOWN, context->GetInputDataType(0));
103+ context->SetOutputDataType(HC_OUTPUT_UP, context->GetInputDataType(0));
104+ return ge::GRAPH_SUCCESS;
105+}
106+ 
107+IMPL_OP_INFERSHAPE(TextGatedResidual)
108+ .InferShape(InferShape4TextGatedResidual)
109+ .InferDataType(InferDataType4TextGatedResidual);
110+} // namespace ops
Amhc/text_gated_residual/op_host/text_gated_residual_tiling.cpp+396-0
@@ -0,0 +1,396 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include "text_gated_residual_tiling.h"
12+#include "../op_kernel/catlass_hc_forward_mul_sum_tiling_data.h"
13+#include "platform/soc_spec.h"
14+ 
15+#include <cstdint>
16+#include <cmath>
17+#include <limits>
18+#include <initializer_list>
19+#include <register/op_impl_registry.h>
20+#include "tiling/platform/platform_ascendc.h"
21+ 
22+namespace {
23+constexpr int64_t HC_ACTIVATION_RANK = 3;
24+constexpr int64_t HC_WEIGHT_RANK = 2;
25+constexpr int64_t HC_FEATURE_DIM = 2;
26+constexpr int64_t HC_DTYPE_BF16 = 2;
27+constexpr int64_t HC_WORKSPACE_ALIGN_PADDING = 15;
28+} // namespace
29+ 
30+namespace optiling {
31+namespace {
32+constexpr size_t INPUT_X = 0;
33+constexpr size_t INPUT_W_DOWN = 1;
34+constexpr size_t INPUT_W_UP = 2;
35+constexpr size_t INPUT_W_GATE = 3;
36+constexpr size_t INPUT_GAMMA = 4;
37+constexpr size_t ATTR_CHANNELS = 0;
38+constexpr size_t ATTR_EPS = 1;
39+constexpr size_t ATTR_NEED_BACKWARD = 2;
40+constexpr uint32_t M_TASK_LARGE = HC_MTASK_LARGE;
41+constexpr uint32_t M_TASK_SMALL = HC_MTASK_SMALL;
42+constexpr uint32_t L1_M_SMALL = HC_L1M_SMALL;
43+constexpr uint32_t SMALL_M_THRESH = HC_SMALL_M_THRESH;
44+// A2：行任务 128 之后，小路径约 0.42µs/行，大路径在核没铺满时停在 430–450µs。
45+// M=896 小路径更快（407 vs 443），M=1024 大路径更快（433 vs 461）。
46+constexpr uint32_t SMALL_M_THRESH_A2 = 1024;
47+constexpr uint32_t L1_N0_LARGE = HC_L1N0_LARGE;
48+constexpr uint32_t L1_N0_LARGE_A2 = 256;
49+constexpr uint32_t L1_N1 = HC_L1N1;
50+ 
51+bool AddProduct(int64_t& total, std::initializer_list<int64_t> factors)
52+{
53+ int64_t product = 1;
54+ for (int64_t factor : factors) {
55+ if (factor < 0 || (factor != 0 && product > INT64_MAX / factor)) {
56+ return false;
57+ }
58+ product *= factor;
59+ }
60+ if (total > INT64_MAX - product) {
61+ return false;
62+ }
63+ total += product;
64+ return true;
65+}
66+ 
67+bool CheckedWorkspaceBytes(const HcForwardMulSumTilingData& t, bool isA2, bool fusedSmall, int64_t& bytes)
68+{
69+ bytes = 0;
70+ const int64_t m = t.M, d = t.D, r = t.R, cores = t.usedAic, task = t.mTask;
71+ const int64_t rawCols = r + t.nWsGate;
72+ if (fusedSmall) {
73+ return AddProduct(bytes, {m, d + r, t.szT}) && AddProduct(bytes, {2, m, rawCols, t.szT}) &&
74+ AddProduct(bytes, {2, m, rawCols, 4, t.splitkFactor}) &&
75+ (t.needBackward || AddProduct(bytes, {m, d, t.szT}));
76+ }
77+ const bool ring = t.smallN == 0;
78+ if (!AddProduct(bytes, {ring ? cores : 1, ring ? 3 : 1, ring ? task : m, d, t.szT}) ||
79+ !AddProduct(bytes, {ring ? cores : 1, ring ? 2 : 1, ring ? task : m, r, t.szT})) {
80+ return false;
81+ }
82+ if (isA2) {
83+ if (!AddProduct(bytes, {ring ? cores : 1, 2, ring ? t.L1M0 : task, r + HcC0GateCols(t), 4})) {
84+ return false;
85+ }
86+ if (t.smallN != 0 && t.splitkFactor > 1 && !AddProduct(bytes, {2, task, HcTilingNWs(t), 4, t.splitkFactor})) {
87+ return false;
88+ }
89+ } else if (!AddProduct(bytes, {cores, t.L1M0, HcA5TilingNWs(t), t.szT, 2})) {
90+ return false;
91+ }
92+ return t.needBackward || AddProduct(bytes, {m, d, t.szT});
93+}
94+ 
95+} // namespace
96+ 
97+bool XDtypeOk(ge::DataType dt)
98+{
99+ return dt == ge::DT_FLOAT16 || dt == ge::DT_BF16;
100+}
101+ 
102+ge::graphStatus Tiling4TextGatedResidual(gert::TilingContext* context)
103+{
104+ if (context == nullptr || context->GetPlatformInfo() == nullptr)
105+ return ge::GRAPH_FAILED;
106+ bool isDav2201 = false;
107+ if (context != nullptr && context->GetPlatformInfo() != nullptr) {
108+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
109+ isDav2201 = ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_2201;
110+ }
111+ const auto xShapePtr = context->GetInputShape(INPUT_X);
112+ const auto wDownShapePtr = context->GetInputShape(INPUT_W_DOWN);
113+ const auto wUpShapePtr = context->GetInputShape(INPUT_W_UP);
114+ const auto attrs = context->GetAttrs();
115+ if (xShapePtr == nullptr || wDownShapePtr == nullptr || wUpShapePtr == nullptr || attrs == nullptr) {
116+ return ge::GRAPH_FAILED;
117+ }
118+ const auto& xShape = xShapePtr->GetStorageShape();
119+ if (xShape.GetDimNum() != HC_ACTIVATION_RANK) {
120+ return ge::GRAPH_FAILED;
121+ }
122+ const int64_t* channelsPtr = attrs->GetAttrPointer<int64_t>(ATTR_CHANNELS);
123+ const float* epsPtr = attrs->GetAttrPointer<float>(ATTR_EPS);
124+ const bool hasWGate = context->GetOptionalInputShape(INPUT_W_GATE) != nullptr;
125+ const bool* needBackwardPtr = attrs->GetAttrPointer<bool>(ATTR_NEED_BACKWARD);
126+ const uint32_t needBackward = (needBackwardPtr != nullptr && *needBackwardPtr) ? 1u : 0u;
127+ if (channelsPtr == nullptr || *channelsPtr <= 0) {
128+ return ge::GRAPH_FAILED;
129+ }
130+ 
131+ const auto& wdShape = wDownShapePtr->GetStorageShape();
132+ const auto& wuShape = wUpShapePtr->GetStorageShape();
133+ if (wdShape.GetDimNum() != HC_WEIGHT_RANK || wuShape.GetDimNum() != HC_WEIGHT_RANK || epsPtr == nullptr ||
134+ !std::isfinite(*epsPtr) || *epsPtr < 0.f)
135+ return ge::GRAPH_FAILED;
136+ const int64_t b64 = xShape.GetDim(0), s64 = xShape.GetDim(1), d64 = xShape.GetDim(HC_FEATURE_DIM),
137+ r64 = wdShape.GetDim(0);
138+ if (b64 < 0 || s64 < 0 || d64 <= 0 || r64 <= 0 || b64 > UINT32_MAX || s64 > UINT32_MAX || d64 > UINT32_MAX ||
139+ r64 > UINT32_MAX || *channelsPtr > UINT32_MAX || (s64 != 0 && b64 > UINT32_MAX / s64) ||
140+ wdShape.GetDim(1) != d64 || wuShape.GetDim(0) != d64 || wuShape.GetDim(1) != r64)
141+ return ge::GRAPH_FAILED;
142+ int64_t tensorBytes = 0;
143+ if (!AddProduct(tensorBytes, {b64, s64, d64, 4}) || !AddProduct(tensorBytes, {b64, s64, r64, 4}) ||
144+ !AddProduct(tensorBytes, {r64, d64, 2}) || !AddProduct(tensorBytes, {*channelsPtr, d64, 2})) {
145+ return ge::GRAPH_FAILED;
146+ }
147+ const int64_t B = b64;
148+ const int64_t S = s64;
149+ const uint32_t D = static_cast<uint32_t>(xShape.GetDim(HC_FEATURE_DIM));
150+ const uint32_t C = static_cast<uint32_t>(*channelsPtr);
151+ if (C == 0 || D % C != 0) {
152+ return ge::GRAPH_FAILED;
153+ }
154+ const uint32_t H = D / C;
155+ const uint32_t R = static_cast<uint32_t>(wDownShapePtr->GetStorageShape().GetDim(0));
156+ const uint32_t M = static_cast<uint32_t>(B * S);
157+ if (isDav2201) {
158+ if (epsPtr == nullptr || *epsPtr < 0.f || D == 0 || R == 0) {
159+ return ge::GRAPH_FAILED;
160+ }
161+ const auto xDescEarly = context->GetInputDesc(INPUT_X);
162+ const auto downDesc = context->GetInputDesc(INPUT_W_DOWN);
163+ const auto upDesc = context->GetInputDesc(INPUT_W_UP);
164+ if (xDescEarly == nullptr || downDesc == nullptr || upDesc == nullptr || !XDtypeOk(xDescEarly->GetDataType()) ||
165+ downDesc->GetDataType() != xDescEarly->GetDataType() ||
166+ upDesc->GetDataType() != xDescEarly->GetDataType()) {
167+ return ge::GRAPH_FAILED;
168+ }
169+ const auto& wDownShape = wDownShapePtr->GetStorageShape();
170+ const auto& wUpShape = wUpShapePtr->GetStorageShape();
171+ if (wDownShape.GetDimNum() != HC_WEIGHT_RANK || wUpShape.GetDimNum() != HC_WEIGHT_RANK ||
172+ wDownShape.GetDim(1) != static_cast<int64_t>(D) || wUpShape.GetDim(0) != static_cast<int64_t>(D) ||
173+ wUpShape.GetDim(1) != static_cast<int64_t>(R)) {
174+ return ge::GRAPH_FAILED;
175+ }
176+ const bool useCombineNow = hasWGate;
177+ if (useCombineNow) {
178+ const auto wGateShapePtr = context->GetOptionalInputShape(INPUT_W_GATE);
179+ if (wGateShapePtr == nullptr || wGateShapePtr->GetStorageShape().GetDimNum() != HC_WEIGHT_RANK ||
180+ wGateShapePtr->GetStorageShape().GetDim(0) != static_cast<int64_t>(C) ||
181+ wGateShapePtr->GetStorageShape().GetDim(1) != static_cast<int64_t>(D)) {
182+ return ge::GRAPH_FAILED;
183+ }
184+ const auto gateDesc = context->GetInputDesc(INPUT_W_GATE);
185+ if (gateDesc == nullptr || gateDesc->GetDataType() != xDescEarly->GetDataType()) {
186+ return ge::GRAPH_FAILED;
187+ }
188+ }
189+ if (B == 0 || S == 0) {
190+ TextGatedResidualTilingData tiling;
191+ tiling.set_B(B);
192+ tiling.set_S(S);
193+ tiling.set_M(0);
194+ tiling.set_D(D);
195+ tiling.set_R(R);
196+ tiling.set_H(H);
197+ tiling.set_C(C);
198+ tiling.set_tilingKey(0);
199+ tiling.set_eps(epsPtr != nullptr ? *epsPtr : 0.f);
200+ tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
201+ context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
202+ context->SetBlockDim(1);
203+ context->SetTilingKey(0);
204+ context->GetWorkspaceSizes(1)[0] = 0;
205+ return ge::GRAPH_SUCCESS;
206+ }
207+ }
208+ const auto xDesc = context->GetInputDesc(INPUT_X);
209+ if (xDesc == nullptr || !XDtypeOk(xDesc->GetDataType())) {
210+ return ge::GRAPH_FAILED;
211+ }
212+ const auto* wdDesc = context->GetInputDesc(INPUT_W_DOWN);
213+ const auto* wuDesc = context->GetInputDesc(INPUT_W_UP);
214+ if (wdDesc == nullptr || wuDesc == nullptr || wdDesc->GetDataType() != xDesc->GetDataType() ||
215+ wuDesc->GetDataType() != xDesc->GetDataType())
216+ return ge::GRAPH_FAILED;
217+ if (hasWGate) {
218+ const auto* gateShape = context->GetOptionalInputShape(INPUT_W_GATE);
219+ const auto* gateDesc = context->GetInputDesc(INPUT_W_GATE);
220+ if (gateShape == nullptr || gateDesc == nullptr || gateDesc->GetDataType() != xDesc->GetDataType() ||
221+ gateShape->GetStorageShape().GetDimNum() != HC_WEIGHT_RANK || gateShape->GetStorageShape().GetDim(0) != C ||
222+ gateShape->GetStorageShape().GetDim(1) != D)
223+ return ge::GRAPH_FAILED;
224+ }
225+ uint32_t dtypeId = 1;
226+ uint32_t szT = 2;
227+ if (xDesc->GetDataType() == ge::DT_BF16) {
228+ dtypeId = HC_DTYPE_BF16;
229+ }
230+ const uint32_t smallThresh = isDav2201 ? SMALL_M_THRESH_A2 : SMALL_M_THRESH;
231+ const uint32_t smallN = (M < smallThresh) ? 1u : 0u;
232+ uint32_t mTask = smallN ? M_TASK_SMALL : M_TASK_LARGE;
233+ if (smallN != 0u) {
234+ if (isDav2201) {
235+ // M>=128 时拆成两段等长任务，行数上取整到 32。M=128/192/320/384 比固定
236+ // 128 或 256 快 7–13µs；M 更小时半段不足 64 行，单段更快。
237+ if (M >= 128u) {
238+ uint32_t half = M / 2u + M % 2u;
239+ mTask = (half + 31u) & ~31u;
240+ } else if (M >= 64u) {
241+ mTask = M_TASK_LARGE;
242+ }
243+ }
244+ }
245+ auto platformInfo = context->GetPlatformInfo();
246+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
247+ uint32_t aicNum = ascendcPlatform.GetCoreNumAic();
248+ // 大路径 128 行任务不足以覆盖 AIC 时，降到 64 行提高并行度。
249+ if (smallN == 0u && (M / mTask + (M % mTask != 0u)) < aicNum) {
250+ mTask = 64u;
251+ }
252+ // A5 large kernels keep compile-time L1M=128 for C0 slot addressing even
253+ // when the task has only 64 valid rows. Host workspace must use that stride.
254+ const uint32_t l1M = smallN ? L1_M_SMALL : (isDav2201 ? mTask : M_TASK_LARGE);
255+ const uint32_t l1N0Small = isDav2201 ? 64u : 32u;
256+ const uint32_t l1N0 = smallN ? l1N0Small : (isDav2201 ? L1_N0_LARGE_A2 : L1_N0_LARGE);
257+ const uint32_t taskNum = M / mTask + (M % mTask != 0u);
258+ const uint32_t c1Swizzle = (M < D) ? 1u : 0u;
259+ const uint32_t useCombine = (hasWGate) ? 1u : 0u;
260+ // Padding and concatenation remain uint32 tile counts; reject before narrowing or alignment.
261+ const int64_t paddedR = (static_cast<int64_t>(R) + 31) / 32 * 32;
262+ const int64_t paddedC = useCombine ? (static_cast<int64_t>(C) + 31) / 32 * 32 : 0;
263+ const int64_t concat = static_cast<int64_t>(R) + (useCombine ? C : 0u);
264+ if (paddedR > UINT32_MAX || paddedC > UINT32_MAX || paddedR + paddedC > UINT32_MAX ||
265+ concat > static_cast<int64_t>(UINT32_MAX) - HC_WORKSPACE_ALIGN_PADDING) {
266+ return ge::GRAPH_FAILED;
267+ }
268+ const uint32_t nWsGate = isDav2201 ? 0u : HcNWsGate(C, useCombine);
269+ const uint32_t nConcat = isDav2201 ? HcNConcat(R, C, useCombine) : HcNWsDown(R) + nWsGate;
270+ const uint32_t nWs = isDav2201 ? HcNWs(nConcat) : HcNWsDown(R);
271+ 
272+ uint32_t splitkFactor = 1;
273+ uint32_t usedAic = 1;
274+ if (smallN && isDav2201) {
275+ const uint32_t nSliceC0 = nConcat / l1N0 + (nConcat % l1N0 != 0u);
276+ const uint32_t nSliceC1 = (D + L1_N1 - 1) / L1_N1;
277+ splitkFactor = HcGetSplitkFactor(mTask, R, D, l1M, l1N0, HcC0L1K(szT), aicNum);
278+ // C0 每个 wave 的循环数是 mTile×nTile；A2 bf16 的 kernel 侧 L1M0=32。
279+ // 不计入 mTile 会把 usedAic 压到 nSliceC1=10，C0 峰值被抬到 3 块/核。
280+ const uint32_t c0TileM = isDav2201 ? 32u : l1M;
281+ uint32_t c0Loops = nSliceC0 * splitkFactor;
282+ if (isDav2201) {
283+ c0Loops *= (mTask + c0TileM - 1u) / c0TileM;
284+ }
285+ uint32_t nSliceMax = c0Loops > nSliceC1 ? c0Loops : nSliceC1;
286+ usedAic = aicNum < nSliceMax ? aicNum : nSliceMax;
287+ } else if (smallN == 0u) {
288+ usedAic = aicNum < taskNum ? aicNum : taskNum;
289+ }
290+ if (usedAic == 0) {
291+ usedAic = 1;
292+ }
293+ 
294+ const auto gammaTensor = context->GetOptionalInputTensor(INPUT_GAMMA);
295+ uint32_t hasGamma = 0u;
296+ if (gammaTensor != nullptr) {
297+ const auto gammaShape = context->GetOptionalInputShape(INPUT_GAMMA);
298+ if (gammaShape == nullptr || gammaShape->GetStorageShape().GetDimNum() != 1 ||
299+ gammaShape->GetStorageShape().GetShapeSize() != static_cast<int64_t>(D)) {
300+ return ge::GRAPH_FAILED;
301+ }
302+ hasGamma = 1u;
303+ }
304+ uint32_t tilingKey = c1Swizzle | (smallN << 1); // bit0 swizzle，bit1 smallN；dtype 由 DTYPE_X
305+ 
306+ const bool fusedSmall = ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510 && smallN != 0u && C >= 2u &&
307+ C <= 10u && H >= 256u && H <= 4096u && H % 256u == 0u && R >= 64u && R <= 896u &&
308+ R % 64u == 0u && aicNum > 0u;
309+ if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_3510 && M > 0u && smallN != 0u && !fusedSmall) {
310+ return ge::GRAPH_FAILED;
311+ }
312+ if (fusedSmall) {
313+ mTask = M;
314+ usedAic = aicNum;
315+ splitkFactor = 14u;
316+ }
317+ 
318+ TextGatedResidualTilingData tiling;
319+ tiling.set_B(B);
320+ tiling.set_S(S);
321+ tiling.set_M(M);
322+ tiling.set_D(D);
323+ tiling.set_R(R);
324+ tiling.set_H(H);
325+ tiling.set_C(C);
326+ tiling.set_usedAic(usedAic);
327+ tiling.set_mTask(mTask);
328+ tiling.set_taskNum(fusedSmall ? 1u : taskNum);
329+ tiling.set_L1M0(l1M);
330+ tiling.set_L1N0(l1N0);
331+ tiling.set_L1M1(l1M);
332+ tiling.set_L1N1(L1_N1);
333+ tiling.set_szT(szT);
334+ tiling.set_useCombine(useCombine);
335+ tiling.set_hasGamma(hasGamma);
336+ tiling.set_tilingKey(HcEncodeTilingKey(dtypeId, useCombine, hasGamma, c1Swizzle, smallN));
337+ tiling.set_eps(epsPtr != nullptr ? *epsPtr : 1e-6f);
338+ tiling.set_smallN(smallN);
339+ tiling.set_splitkFactor(smallN ? splitkFactor : 1u);
340+ tiling.set_nConcat(nConcat);
341+ tiling.set_nWs(nWs);
342+ tiling.set_nWsGate(nWsGate);
343+ tiling.set_needBackward(needBackward);
344+ tiling.set_pad0(0);
345+ tiling.set_pad1(0u);
346+ tiling.set_pad2(0);
347+ tiling.set_pad3(0u);
348+ 
349+ tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
350+ context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
351+ 
352+ HcForwardMulSumTilingData ws{};
353+ ws.M = M;
354+ ws.D = D;
355+ ws.R = R;
356+ ws.usedAic = usedAic;
357+ ws.mTask = mTask;
358+ ws.L1M0 = l1M;
359+ ws.L1N0 = l1N0;
360+ ws.L1M1 = l1M;
361+ ws.L1N1 = L1_N1;
362+ ws.szT = szT;
363+ ws.smallN = smallN;
364+ ws.splitkFactor = smallN ? splitkFactor : 1u;
365+ ws.C = C;
366+ ws.useCombine = useCombine;
367+ ws.nConcat = nConcat;
368+ ws.nWs = nWs;
369+ ws.nWsGate = nWsGate;
370+ ws.needBackward = needBackward;
371+ 
372+ context->SetBlockDim(usedAic);
373+ context->SetTilingKey(tilingKey);
374+ int64_t workspaceBytes = 0;
375+ if (!CheckedWorkspaceBytes(ws, isDav2201, fusedSmall, workspaceBytes)) {
376+ return ge::GRAPH_FAILED;
377+ }
378+ const auto libBytes = ascendcPlatform.GetLibApiWorkSpaceSize();
379+ if (libBytes > static_cast<uint64_t>(INT64_MAX) || !AddProduct(workspaceBytes, {static_cast<int64_t>(libBytes)}) ||
380+ static_cast<uint64_t>(workspaceBytes) > std::numeric_limits<size_t>::max()) {
381+ return ge::GRAPH_FAILED;
382+ }
383+ context->GetWorkspaceSizes(1)[0] = static_cast<size_t>(workspaceBytes);
384+ return ge::GRAPH_SUCCESS;
385+}
386+ 
387+ge::graphStatus TilingPrepare4TextGatedResidual(gert::TilingParseContext* context)
388+{
389+ (void)context;
390+ return ge::GRAPH_SUCCESS;
391+}
392+ 
393+IMPL_OP_OPTILING(TextGatedResidual)
394+ .Tiling(Tiling4TextGatedResidual)
395+ .TilingParse<TextGatedResidualCompileInfo>(TilingPrepare4TextGatedResidual);
396+} // namespace optiling
Amhc/text_gated_residual/op_host/text_gated_residual_tiling.h+52-0
@@ -0,0 +1,52 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+#pragma once
11+ 
12+#include <cstdint>
13+#include <register/tilingdata_base.h>
14+ 
15+namespace optiling {
16+ 
17+BEGIN_TILING_DATA_DEF(TextGatedResidualTilingData)
18+TILING_DATA_FIELD_DEF(int64_t, B);
19+TILING_DATA_FIELD_DEF(int64_t, S);
20+TILING_DATA_FIELD_DEF(uint32_t, M);
21+TILING_DATA_FIELD_DEF(uint32_t, D);
22+TILING_DATA_FIELD_DEF(uint32_t, R);
23+TILING_DATA_FIELD_DEF(uint32_t, H);
24+TILING_DATA_FIELD_DEF(uint32_t, C);
25+TILING_DATA_FIELD_DEF(uint32_t, usedAic);
26+TILING_DATA_FIELD_DEF(uint32_t, mTask);
27+TILING_DATA_FIELD_DEF(uint32_t, taskNum);
28+TILING_DATA_FIELD_DEF(uint32_t, L1M0);
29+TILING_DATA_FIELD_DEF(uint32_t, L1N0);
30+TILING_DATA_FIELD_DEF(uint32_t, L1M1);
31+TILING_DATA_FIELD_DEF(uint32_t, L1N1);
32+TILING_DATA_FIELD_DEF(uint32_t, szT);
33+TILING_DATA_FIELD_DEF(uint32_t, useCombine);
34+TILING_DATA_FIELD_DEF(uint32_t, hasGamma);
35+TILING_DATA_FIELD_DEF(uint32_t, tilingKey);
36+TILING_DATA_FIELD_DEF(float, eps);
37+TILING_DATA_FIELD_DEF(uint32_t, smallN);
38+TILING_DATA_FIELD_DEF(uint32_t, splitkFactor);
39+TILING_DATA_FIELD_DEF(uint32_t, nConcat);
40+TILING_DATA_FIELD_DEF(uint32_t, nWs);
41+TILING_DATA_FIELD_DEF(uint32_t, nWsGate);
42+TILING_DATA_FIELD_DEF(uint32_t, needBackward);
43+TILING_DATA_FIELD_DEF(uint32_t, pad0);
44+TILING_DATA_FIELD_DEF(uint32_t, pad1);
45+TILING_DATA_FIELD_DEF(uint32_t, pad2);
46+TILING_DATA_FIELD_DEF(uint32_t, pad3);
47+END_TILING_DATA_DEF;
48+ 
49+REGISTER_TILING_DATA_CLASS(TextGatedResidual, TextGatedResidualTilingData)
50+ 
51+struct TextGatedResidualCompileInfo {};
52+} // namespace optiling
Amhc/text_gated_residual/op_kernel/arch22/block/block_epilogue_hc_v0.hpp+419-0
@@ -0,0 +1,419 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+ 
13+#include "catlass/catlass.hpp"
14+#include "catlass/arch/resource.hpp"
15+#include "catlass/gemm/gemm_type.hpp"
16+#include "catlass/layout/layout.hpp"
17+#include "../hc_gm_copy.hpp"
18+ 
19+namespace Catlass::Epilogue {
20+ 
21+struct EpilogueHcV0Rms {
22+ using ArchTag = Arch::AtlasA2;
23+ static constexpr uint32_t UB_STAGES = 1;
24+};
25+ 
26+template <class Element, class ArchTag = Arch::AtlasA2>
27+class BlockEpilogueHcV0 {
28+public:
29+ using ElementCompute = float;
30+ static constexpr uint32_t TMP_ELEMS = 512;
31+ 
32+ CATLASS_DEVICE
33+ BlockEpilogueHcV0() {}
34+ 
35+ // GlobalTensor 只在 kernel 入口绑一次。后续按元素偏移取视图，L2 hint 随偏移保留。
36+ CATLASS_DEVICE
37+ void Init(GM_ADDR gmX, GM_ADDR gmXn, GM_ADDR gmGamma, GM_ADDR gmInvRms, bool l2HintX, bool l2HintXn)
38+ {
39+ using namespace AscendC;
40+ gX_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmX));
41+ if (l2HintX) {
42+ gX_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
43+ }
44+ gXn_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmXn));
45+ // xn 写出后只被本核 c0、v2 各读一次。大 shape 与 v2 一样不占 L2。
46+ if (l2HintXn) {
47+ gXn_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
48+ }
49+ hasGamma_ = gmGamma != nullptr;
50+ if (hasGamma_) {
51+ gGamma_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(gmGamma));
52+ }
53+ hasInv_ = gmInvRms != nullptr;
54+ if (hasInv_) {
55+ gInvRms_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(gmInvRms));
56+ }
57+ }
58+ 
59+ template <class Resource>
60+ CATLASS_DEVICE void operator()(Resource& resource, int64_t offX, int64_t offXn, uint32_t actualM, uint32_t d,
61+ uint32_t channels, uint32_t hidden, float eps, bool useCombine,
62+ bool disableSubSplit = false, int64_t offInv = 0)
63+ {
64+ using namespace AscendC;
65+ if (actualM == 0) {
66+ return;
67+ }
68+ const uint32_t subNum = GetSubBlockNum() == 0 ? 1 : GetSubBlockNum();
69+ const uint32_t subIdx = GetSubBlockIdx();
70+ const uint32_t row0 = (actualM + subNum - 1) / subNum;
71+ uint32_t rowStart = subIdx * row0;
72+ uint32_t rowEnd = rowStart + row0 < actualM ? rowStart + row0 : actualM;
73+ if (disableSubSplit) {
74+ rowStart = 0;
75+ rowEnd = actualM;
76+ }
77+ (void)useCombine;
78+ GlobalTensor<Element> gX = gX_[offX];
79+ GlobalTensor<Element> gXn = gXn_[offXn];
80+ GlobalTensor<float> gGamma;
81+ if (hasGamma_) {
82+ gGamma = gGamma_;
83+ }
84+ GlobalTensor<float> gInvRms;
85+ if (hasInv_) {
86+ gInvRms = gInvRms_[offInv];
87+ }
88+ 
89+ // A2 与 nn 仓 KernelRmsNorm NORMAL 一致：UB 按一个通道（一行 H）划分，不按整行 D。
90+ // inQueueX(H) + gamma(H) + outQueueY(H) + rstd(64) + x_fp32(H) + sqx(H) + reduce(64)。
91+ uint32_t off = 0;
92+ LocalTensor<uint8_t> ub = resource.ubBuf.template GetBufferByByte<uint8_t>(0);
93+ auto align32 = [](uint32_t n) { return (n + 31u) & ~31u; };
94+ uint32_t colBytesT = align32(hidden * static_cast<uint32_t>(sizeof(Element)));
95+ uint32_t colBytesF = align32(hidden * static_cast<uint32_t>(sizeof(float)));
96+ LocalTensor<Element> ubRowT = ub[off].template ReinterpretCast<Element>();
97+ LocalTensor<float> ubRowF;
98+ LocalTensor<float> ubXFp32;
99+ if constexpr (std::is_same_v<Element, float>) {
100+ ubRowF = ub[off + colBytesT].template ReinterpretCast<float>();
101+ ubXFp32 = ubRowT;
102+ off += colBytesT;
103+ } else {
104+ off += colBytesT;
105+ ubXFp32 = ub[off].template ReinterpretCast<float>();
106+ off += colBytesF;
107+ ubRowF = ub[off].template ReinterpretCast<float>();
108+ }
109+ off += colBytesF;
110+ LocalTensor<Element> ubY;
111+ ubY = ub[off].template ReinterpretCast<Element>();
112+ off += colBytesT;
113+ // 第二套输入/输出，让下一通道的 MTE2 叠在当前通道的 V 上，和 nn 仓队列腾出 x 之后的行为一致。
114+ LocalTensor<Element> ubIn1 = ub[off].template ReinterpretCast<Element>();
Zhengyi Zhang
Zhengyi ZhangZhengyi Zhang3 days ago

【review】 问题分类：正确性 位置：mhc/text_gated_residual/op_kernel/arch22/block/block_epilogue_hc_v0.hpp:94、105、114、116，代码：LocalTensor ubIn1 = ub[off].template ReinterpretCast();；mhc/text_gated_residual/op_host/text_gated_residual_tiling.cpp:129、145、148 问题：新增 A2 RMS 路径按完整 H 分配四份 H 元素的 FP16/BF16 缓冲区和两份 H 元素的 FP32 缓冲区，却未根据 UB 容量分块或拒绝超限维度。例如 x=[1,32,65536]、channels=4、w_down=[64,65536]、w_up=[65536,64]、gamma=None、w_gate=None、needBackward=false、eps=1e-6，满足当前 ACLNN/tiling 的 shape、dtype 和范围校验，进入 A2 小路径后 H=16384。仅上述主缓冲区就需要 262144 字节；项目固定依赖的 CATLASS v2.0.0 中 AtlasA2::UB_SIZE 为 196608 字节。此时 ubIn1 起始偏移已是 196608，后续预取仍向它搬入 32768 字节，整个输入槽均在 UB 范围之外。这是静态容量及数据流核对，未执行设备测试。 可能导致的结果：该合法尺寸下的下一通道预取和归一化访问超出 UB 范围，算子可能执行失败或产生错误结果；新增的五输出及反向缓存不能消除此越界。 建议的修改方法：按目标 UB 容量计算包含所有对齐、临时量和可选 gamma/cache 的完整占用，对超限 H 实现分段归约与写回；若暂不支持，则在 ACLNN/tiling 中明确拒绝并同步规格约束。增加上述尺寸及容量临界点两侧的用例。

likedislike
115+ off += colBytesT;
116+ LocalTensor<Element> ubOut1 = ub[off].template ReinterpretCast<Element>();
117+ off += colBytesT;
118+ LocalTensor<float> ubMean = ub[off].template ReinterpretCast<float>();
119+ off += 64 * sizeof(float);
120+ LocalTensor<float> ubTmp = ub[off].template ReinterpretCast<float>();
121+ off += 64 * sizeof(float);
122+ LocalTensor<float> ubGamma;
123+ bool gammaHoisted = false;
124+ if (hasGamma_) {
125+ // gamma 很小（C×H×4B），一次搬进 UB 常驻复用，省掉逐行逐通道的 GM 重搬与 +1。
126+ // 超大 C×H 放不下 UB 时退回逐通道搬运的旧路径。
127+ uint32_t gammaAllBytes = align32(channels * hidden * static_cast<uint32_t>(sizeof(float)));
128+ gammaHoisted = gammaAllBytes <= 65536u;
129+ ubGamma = ub[off].template ReinterpretCast<float>();
130+ off += gammaHoisted ? gammaAllBytes : align32(hidden * static_cast<uint32_t>(sizeof(float)));
131+ }
132+ // inv_rms 批量落盘：WholeReduceSum 直写非对齐紧凑槽攒「原始 sum」（替代 SetValue），
133+ // 攒满一批（最多 256 个 float）后先整批向量算 inv 再一次 DataCopy 落 GM，
134+ // 替代原来每行一次的 DataCopyPad + flag 往返。
135+ // gmInvRms 为空时不写不落盘。
136+ constexpr uint32_t ACC_MAX_ELEMS = 256;
137+ uint32_t accRows = hasInv_ && channels > 0 ? (ACC_MAX_ELEMS / channels > 0 ? ACC_MAX_ELEMS / channels : 1) : 1;
138+ uint32_t accElems = accRows * channels;
139+ LocalTensor<float> ubAcc = ub[off].template ReinterpretCast<float>();
140+ off += align32((hasInv_ && channels > 0 ? accElems : channels) * static_cast<uint32_t>(sizeof(float)));
141+ // 向量 inv 链的常数 1.0（Div 分子，整批大小时用）与 Brcb 广播块（1 个 32B block）。
142+ uint32_t onesElems = (hasInv_ && channels > 0) ? accElems : 8u;
143+ if (onesElems < 8u) {
144+ onesElems = 8u;
145+ }
146+ LocalTensor<float> ubOnes = ub[off].template ReinterpretCast<float>();
147+ off += align32(onesElems * static_cast<uint32_t>(sizeof(float)));
148+ LocalTensor<float> ubBrc = ub[off].template ReinterpretCast<float>();
149+ // Brcb repeat=1 读 1 个源 block（8 元素）、写 8 个目的 block（每元素独占一个 block）。
150+ off += 8u * 32u;
151+ Duplicate(ubOnes, 1.0f, onesElems);
152+ PipeBarrier<PIPE_V>();
153+ 
154+ if (hasGamma_ && gammaHoisted) {
155+ for (uint32_t c = 0; c < channels; ++c) {
156+ DataCopyExtParams gammaParams{1, hidden * static_cast<uint32_t>(sizeof(float)), 0, 0, 0};
157+ DataCopyPadExtParams<float> gammaPad{false, 0, 0, 0.f};
158+ DataCopyPad(ubGamma[c * hidden], gGamma[static_cast<int64_t>(c) * hidden], gammaParams, gammaPad);
Zhengyi Zhang
Zhengyi ZhangZhengyi Zhang3 days ago

【review】 问题分类：正确性 位置：mhc/text_gated_residual/op_kernel/arch22/block/block_epilogue_hc_v0.hpp:127、158、336，代码：DataCopyPad(ubGamma[c * hidden], gGamma[static_cast<int64_t>(c) * hidden], gammaParams, gammaPad);；mhc/text_gated_residual/op_host/text_gated_residual_tiling.cpp:142、145 问题：新增 gamma 常驻路径将各通道按未对齐的 H 紧密放在 UB，并直接用 c*H 作为搬运及 Mul 的起点，但 host 只校验 D%channels==0，未限制 H 为 8 的倍数。例如 A2、FP16、x=[1,1,1028]、channels=4、w_down=[64,1028]、w_up=[1028,64]、gamma 为 FP32[1028]、w_gate=None、eps=1e-6，可通过当前入口校验；H=257 且 gammaAllBytes=4128，因而启用 gammaHoisted。c=1 时 ubGamma[257] 相对 UB 基址偏移 1028 字节，不能满足 32 字节对齐。DataCopyPad 支持 GM 地址和长度非对齐，但其 UB 目的地址仍必须 32 字节对齐，后续 Mul 也直接使用该非对齐通道起点。依据为 Ascend C DataCopyPad（GMToUB）地址约束；未执行设备测试。 可能导致的结果：带 gamma 的非对齐 H 输入在归一化阶段违反指令地址约束，不能可靠地按每个通道的 gamma 计算结果，可能出现执行异常或结果错误。 建议的修改方法：对各通道使用独立的 32 字节对齐 UB 步长，并同步修正搬运、gamma 加一及乘法的索引；或在 H 非 8 的倍数时使用现有按通道搬入对齐临时槽的路径。补充 H=257 等非对齐宽度、各通道 gamma 不同的验证用例，并与无 gamma 分支分别覆盖。

likedislike
K
kdy18482276080
2 days ago Comments：
159+ }
160+ SetFlag<HardEvent::MTE2_V>(EVENT_ID3);
161+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID3);
162+ Adds(ubGamma, ubGamma, 1.0f, channels * hidden);
163+ PipeBarrier<PIPE_V>();
164+ }
165+ 
166+ // 对齐 nn 仓 KernelRmsNorm::Compute：Cast 一次、count 模式 Mul/Add 做半区间归约、
167+ // Muls 乘 rstd、再 Cast 回原 dtype。一行对齐时一次 DataCopy，不再按 32 字节拆开搬。
168+ // 行/通道循环里不变的量统一在循环外算一次：归约分解（body/tail 只由 hidden 决定）、
169+ // 搬运参数结构体、1/hidden；行内所有偏移都是 rowBytes 的整数倍，
170+ // 逐通道的对齐判定结果恒等于 rowBytesAligned，也一并外提。
171+ constexpr int32_t kRep = 64;
172+ auto powerTwo = [](int32_t n1) {
173+ n1 |= n1 >> 1;
174+ n1 |= n1 >> 2;
175+ n1 |= n1 >> 4;
176+ n1 |= n1 >> 8;
177+ n1 |= n1 >> 16;
178+ return (n1 + 1) >> 1;
179+ };
180+ const uint32_t rowBytes = hidden * static_cast<uint32_t>(sizeof(Element));
181+ const bool rowBytesAligned = (rowBytes & 31u) == 0u;
182+ const int32_t count = static_cast<int32_t>(hidden);
183+ const int32_t body0 = count > kRep ? powerTwo(count) : count;
184+ const int32_t tail0 = count > kRep ? count - body0 : 0;
185+ const uint64_t maskCount = count > kRep ? static_cast<uint64_t>(kRep) : static_cast<uint64_t>(count);
186+ const float invHidden = count > 0 ? 1.0f / static_cast<float>(count) : 0.0f;
187+ DataCopyExtParams xParams;
188+ xParams.blockCount = 1;
189+ xParams.blockLen = rowBytes;
190+ xParams.srcStride = 0;
191+ xParams.dstStride = 0;
192+ DataCopyPadExtParams<Element> xPad{false, 0, 0, static_cast<Element>(0)};
193+ DataCopyExtParams gammaParams{1, hidden * static_cast<uint32_t>(sizeof(float)), 0, 0, 0};
194+ DataCopyPadExtParams<float> gammaPad{false, 0, 0, 0.f};
195+ LocalTensor<Element> inSlot[2] = {ubRowT, ubIn1};
196+ LocalTensor<Element> outSlot[2] = {ubY, ubOut1};
197+ 
198+ uint32_t pendingRows = 0; // 自上次落盘以来攒下的行数
199+ uint32_t flushBase = rowStart; // 攒批首行的全局行号（GM 落盘基址）
200+ // 输入、输出各两块，按 m*C+c 单调 ping-pong。进循环前各 Set 两次，表示两块都空闲；
201+ // 下标不在通道边界清零，行末也不把队列排空。下一格可以落到下一行的 c=0。
202+ static constexpr uint32_t kEv[2] = {EVENT_ID0, EVENT_ID1};
203+ SetFlag<HardEvent::V_MTE2>(kEv[0]);
204+ SetFlag<HardEvent::V_MTE2>(kEv[1]);
205+ SetFlag<HardEvent::MTE3_V>(kEv[0]);
206+ SetFlag<HardEvent::MTE3_V>(kEv[1]);
207+ if (channels > 0 && rowStart < rowEnd) {
208+ uint32_t headSlot = (rowStart * channels) & 1u;
209+ event_t headEv = static_cast<event_t>(kEv[headSlot]);
210+ WaitFlag<HardEvent::V_MTE2>(headEv);
211+ if (rowBytesAligned) {
212+ DataCopy(inSlot[headSlot], gX[static_cast<int64_t>(rowStart) * d], hidden);
213+ } else {
214+ DataCopyPad(inSlot[headSlot], gX[static_cast<int64_t>(rowStart) * d], xParams, xPad);
215+ }
216+ SetFlag<HardEvent::MTE2_V>(headEv);
217+ }
218+ const BinaryRepeatParams mulBrcParams{1, 1, 0, 8, 8, 0};
219+ uint32_t c = 0;
220+ uint32_t m = rowStart;
221+ const uint32_t mcBegin = rowStart * channels;
222+ const uint32_t mcEnd = rowEnd * channels;
223+ for (uint32_t mc = mcBegin; mc < mcEnd; ++mc) {
224+ uint32_t slot = mc & 1u;
225+ event_t ev = static_cast<event_t>(kEv[slot]);
226+ WaitFlag<HardEvent::MTE2_V>(ev);
227+ bool hasNext = mc + 1u < mcEnd;
228+ if (hasNext) {
229+ uint32_t nextC = c + 1u;
230+ uint32_t nextM = m;
231+ if (nextC == channels) {
232+ nextC = 0;
233+ nextM = m + 1u;
234+ }
235+ uint32_t nxt = slot ^ 1u;
236+ event_t evN = static_cast<event_t>(kEv[nxt]);
237+ WaitFlag<HardEvent::V_MTE2>(evN);
238+ int64_t nxtOff = static_cast<int64_t>(nextM) * d + static_cast<int64_t>(nextC) * hidden;
239+ if (rowBytesAligned) {
240+ DataCopy(inSlot[nxt], gX[nxtOff], hidden);
241+ } else {
242+ DataCopyPad(inSlot[nxt], gX[nxtOff], xParams, xPad);
243+ }
244+ SetFlag<HardEvent::MTE2_V>(evN);
245+ }
246+ LocalTensor<Element> curIn = inSlot[slot];
247+ if constexpr (std::is_same_v<Element, float>) {
248+ Mul(ubRowF, curIn, curIn, hidden);
249+ } else {
250+ // bf16 下这是 inSlot 的最后一次向量读，之后只碰 ubXFp32 / ubRowF。
251+ Cast(ubXFp32, curIn, RoundMode::CAST_NONE, hidden);
252+ SetFlag<HardEvent::V_MTE2>(ev);
253+ PipeBarrier<PIPE_V>();
254+ Mul(ubRowF, ubXFp32, ubXFp32, hidden);
255+ }
256+ PipeBarrier<PIPE_V>();
257+ if (count > kRep) {
258+ if (tail0 > 0) {
259+ Add(ubRowF, ubRowF, ubRowF[body0], tail0);
260+ PipeBarrier<PIPE_V>();
261+ }
262+ int32_t body = body0;
263+ while (body > kRep) {
264+ body /= 2;
265+ Add(ubRowF, ubRowF, ubRowF[body], body);
266+ PipeBarrier<PIPE_V>();
267+ }
268+ }
269+ AscendCUtils::SetMask<float>(maskCount);
270+ if (g_coreType == AIV) {
271+ // 原始 sum 只归约一次：WholeReduceSum 直写非对齐紧凑槽 ubAcc（dst 可为任意
272+ // 元素偏移，同 block_attention_residuals 的 vecMeta_[metaIdx] 用法），替代 SetValue。
273+ // 无 backward（不落盘）时直写对齐 ubTmp，省一次回拷。
274+ if (hasInv_) {
275+ WholeReduceSum<float, false>(ubAcc[pendingRows * channels + c], ubRowF, static_cast<uint64_t>(kRep),
276+ 1, 0, 1, 0);
277+ } else {
278+ WholeReduceSum<float, false>(ubTmp, ubRowF, static_cast<uint64_t>(kRep), 1, 0, 1, 0);
279+ }
280+ }
281+ SetMaskNorm();
282+ ResetMask();
283+ PipeBarrier<PIPE_V>();
284+ // 广播用 inv（第一遍，广播后算）：Brcb 先行——把包含本元素的对齐 block 展开，
285+ // 展开后一个元素独占一个 block（8 lane 全是该值的重复），本元素落在对齐的
286+ // ubBrc[lane*8]，之后 Muls/Adds/Sqrt/Div 全部在 block 对齐地址上做（紧凑非对齐
287+ // 槽做不了这些链式操作）。其余 lane/block 是相邻通道或陈旧数据，链式结果无人读；
288+ // 块内 8 lane 均为本值副本，链式输出天然就是行缩放 Mul 的广播块。
289+ uint32_t brcLane = 0;
290+ if (hasInv_) {
291+ uint32_t accOff = pendingRows * channels + c;
292+ brcLane = accOff & 7u;
293+ Brcb(ubBrc, ubAcc[accOff - brcLane], 1, {1, 8});
294+ } else {
295+ Brcb(ubBrc, ubTmp, 1, {1, 8});
296+ }
297+ PipeBarrier<PIPE_V>();
298+ LocalTensor<float> brcBlk = ubBrc[brcLane * 8];
299+ // inv = 1/sqrt(sum*invHidden+eps)（同 reduce_common_arch22 InvRmsInPlace），
300+ // Sqrt+Div 精确路径，不判断 sum>0，全程向量无 GetValue/SetValue。
301+ Muls(brcBlk, brcBlk, invHidden, 8);
302+ PipeBarrier<PIPE_V>();
303+ Adds(brcBlk, brcBlk, eps, 8);
304+ PipeBarrier<PIPE_V>();
305+ Sqrt(brcBlk, brcBlk, 8);
306+ PipeBarrier<PIPE_V>();
307+ Div(brcBlk, ubOnes, brcBlk, 8);
308+ PipeBarrier<PIPE_V>();
309+ {
310+ // 行缩放：src1 以 blkStride/repStride=0 固定复用广播块 brcBlk
311+ // （同 MulRowByBrcBlock）。COUNTER mask 单次上限 4096，超长行分段。
312+ constexpr uint32_t kCntrMax = 4096;
313+ uint32_t rowOff = 0;
314+ do {
315+ uint32_t len = hidden - rowOff < kCntrMax ? hidden - rowOff : kCntrMax;
316+ SetMaskCount();
317+ SetVectorMask<float, MaskMode::COUNTER>(len);
318+ if constexpr (std::is_same_v<Element, float>) {
319+ Mul<float, false>(ubRowF[rowOff], curIn[rowOff], brcBlk, MASK_PLACEHOLDER, 1, mulBrcParams);
320+ } else {
321+ Mul<float, false>(ubRowF[rowOff], ubXFp32[rowOff], brcBlk, MASK_PLACEHOLDER, 1, mulBrcParams);
322+ }
323+ PipeBarrier<PIPE_V>();
324+ rowOff += len;
325+ } while (rowOff < hidden);
326+ SetMaskNorm();
327+ ResetMask();
328+ if constexpr (std::is_same_v<Element, float>) {
329+ // float 的行缩放还在读 curIn，这是 inSlot 的最后一次向量读。
330+ SetFlag<HardEvent::V_MTE2>(ev);
331+ }
332+ }
333+ if (hasGamma_) {
334+ if (gammaHoisted) {
335+ // gamma 已在启动时整体搬入 UB 并加过 1.0，直接按通道偏移复用。
336+ Mul(ubRowF, ubRowF, ubGamma[c * hidden], hidden);
337+ PipeBarrier<PIPE_V>();
338+ } else {
339+ // Reuse one H-element gamma tile; wait before overwriting it.
340+ PipeBarrier<PIPE_ALL>();
341+ DataCopyPad(ubGamma, gGamma[static_cast<int64_t>(c) * hidden], gammaParams, gammaPad);
342+ SetFlag<HardEvent::MTE2_V>(EVENT_ID3);
343+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID3);
344+ Adds(ubGamma, ubGamma, 1.0f, hidden);
345+ PipeBarrier<PIPE_V>();
346+ Mul(ubRowF, ubRowF, ubGamma, hidden);
347+ PipeBarrier<PIPE_V>();
348+ }
349+ }
350+ // 同一输出槽上一次 GM 搬运结束后才能改写。两块缓冲随平坦下标轮转，不按行重置。
351+ WaitFlag<HardEvent::MTE3_V>(ev);
352+ if constexpr (!std::is_same_v<Element, float>) {
353+ Cast(outSlot[slot], ubRowF, RoundMode::CAST_RINT, hidden);
354+ PipeBarrier<PIPE_V>();
355+ } else {
356+ DataCopy(outSlot[slot], ubRowF, hidden);
357+ PipeBarrier<PIPE_V>();
358+ }
359+ SetFlag<HardEvent::V_MTE3>(ev);
360+ WaitFlag<HardEvent::V_MTE3>(ev);
361+ int64_t outOff = static_cast<int64_t>(m) * d + static_cast<int64_t>(c) * hidden;
362+ if (rowBytesAligned) {
363+ DataCopy(gXn[outOff], outSlot[slot], hidden);
364+ } else {
365+ HcStoreGmElems(gXn, outOff, outSlot[slot], hidden, ubMean.template ReinterpretCast<Element>());
366+ }
367+ SetFlag<HardEvent::MTE3_V>(ev);
368+ if (c + 1u == channels) {
369+ if (hasInv_) {
370+ // 攒满 accRows 行（或到达本核末行）再批量落盘：一批最多 256 个 float。
371+ ++pendingRows;
372+ if (pendingRows == accRows || m + 1u == rowEnd) {
373+ uint32_t accNum = pendingRows * channels;
374+ // inv 搬出前算第二遍：ubAcc 攒的是原始 sum，这里整批向量算
375+ // inv = 1/sqrt(sum*invHidden+eps)，保证落盘缓冲连续、一次搬出。
376+ Muls(ubAcc, ubAcc, invHidden, accNum);
377+ PipeBarrier<PIPE_V>();
378+ Adds(ubAcc, ubAcc, eps, accNum);
379+ PipeBarrier<PIPE_V>();
380+ Sqrt(ubAcc, ubAcc, accNum);
381+ PipeBarrier<PIPE_V>();
382+ Div(ubAcc, ubOnes, ubAcc, accNum);
383+ PipeBarrier<PIPE_V>();
384+ SetFlag<HardEvent::V_MTE3>(EVENT_ID2);
385+ WaitFlag<HardEvent::V_MTE3>(EVENT_ID2);
386+ DataCopyExtParams flushParams;
387+ flushParams.blockCount = 1;
388+ flushParams.blockLen = accNum * static_cast<uint32_t>(sizeof(float));
389+ flushParams.srcStride = 0;
390+ flushParams.dstStride = 0;
391+ DataCopyPad(gInvRms[static_cast<int64_t>(flushBase) * channels], ubAcc, flushParams);
392+ SetFlag<HardEvent::MTE3_V>(EVENT_ID2);
393+ WaitFlag<HardEvent::MTE3_V>(EVENT_ID2);
394+ pendingRows = 0;
395+ flushBase = m + 1u;
396+ }
397+ }
398+ c = 0;
399+ ++m;
400+ } else {
401+ ++c;
402+ }
403+ }
404+ WaitFlag<HardEvent::V_MTE2>(kEv[0]);
405+ WaitFlag<HardEvent::V_MTE2>(kEv[1]);
406+ WaitFlag<HardEvent::MTE3_V>(kEv[0]);
407+ WaitFlag<HardEvent::MTE3_V>(kEv[1]);
408+ }
409+ 
410+private:
411+ AscendC::GlobalTensor<Element> gX_;
412+ AscendC::GlobalTensor<Element> gXn_;
413+ AscendC::GlobalTensor<float> gGamma_;
414+ AscendC::GlobalTensor<float> gInvRms_;
415+ bool hasGamma_ = false;
416+ bool hasInv_ = false;
417+};
418+ 
419+} // namespace Catlass::Epilogue
Amhc/text_gated_residual/op_kernel/arch22/block/block_epilogue_hc_v1.hpp+338-0
@@ -0,0 +1,338 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+ 
13+#include "catlass/catlass.hpp"
14+#include "catlass/arch/resource.hpp"
15+#include "catlass/gemm/gemm_type.hpp"
16+#include "catlass/layout/layout.hpp"
17+#include "../hc_gm_copy.hpp"
18+ 
19+namespace Catlass::Epilogue {
20+ 
21+// src、dst、work 等长。dst = 2 * sigmoid(src / C)，不读写 GM。
22+CATLASS_DEVICE void HcInjectCompute(AscendC::LocalTensor<float> const& dst, AscendC::LocalTensor<float> const& src,
23+ AscendC::LocalTensor<float> const& work, uint32_t n, float invC)
24+{
25+ using namespace AscendC;
26+ constexpr uint32_t kStep = 8192;
27+ for (uint32_t cursor = 0; cursor < n; cursor += kStep) {
28+ uint32_t len = n - cursor < kStep ? n - cursor : kStep;
29+ auto d = src[cursor];
30+ auto o = dst[cursor];
31+ auto e = work[cursor];
32+ Muls(d, d, invC, len);
33+ PipeBarrier<PIPE_V>();
34+ Muls(e, d, -1.0f, len);
35+ PipeBarrier<PIPE_V>();
36+ Exp(e, e, len);
37+ PipeBarrier<PIPE_V>();
38+ Adds(e, e, 1.0f, len);
39+ PipeBarrier<PIPE_V>();
40+ Duplicate(o, 1.0f, len);
41+ PipeBarrier<PIPE_V>();
42+ Div(o, o, e, len);
43+ PipeBarrier<PIPE_V>();
44+ Muls(o, o, 2.0f, len);
45+ PipeBarrier<PIPE_V>();
46+ }
47+}
48+ 
49+struct EpilogueHcV1DivSilu {
50+ using ArchTag = Arch::AtlasA2;
51+ static constexpr uint32_t UB_STAGES = 2;
52+};
53+ 
54+template <class Element, class ArchTag = Arch::AtlasA2>
55+class BlockEpilogueHcV1 {
56+public:
57+ using ElementCompute = float;
58+ static constexpr uint32_t COMPUTE_LENGTH = 8192;
59+ static constexpr uint32_t kUbBudget = 160u * 1024u;
60+ 
61+ CATLASS_DEVICE
62+ BlockEpilogueHcV1() {}
63+ 
64+ CATLASS_DEVICE
65+ void Init(GM_ADDR gmC0Base, GM_ADDR gmDownBase, GM_ADDR gmInjectBase, GM_ADDR gmDownPreActBase,
66+ bool l2HintDownPreAct, uint32_t channels, uint32_t rDim)
67+ {
68+ using namespace AscendC;
69+ gC0_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmC0Base));
70+ gGate_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmC0Base));
71+ gDown_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmDownBase));
72+ gInj_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmInjectBase));
73+ hasPreAct_ = gmDownPreActBase != nullptr;
74+ if (hasPreAct_) {
75+ gDownPreAct_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmDownPreActBase));
76+ gC0FromDpa_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmDownPreActBase));
77+ // 只给反向保存一次，不参与前向复用，不占用 L2。小 shape 该缓冲很小，关闭 L2
78+ // 属于大 shape 的专属优化，小路径不设置。
79+ if (l2HintDownPreAct) {
80+ gDownPreAct_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
81+ }
82+ }
83+ invC_ = channels == 0 ? 0.0f : 1.0f / static_cast<float>(static_cast<int32_t>(channels));
84+ uint32_t yRowBytes = std::is_same_v<Element, float> ? 0u : rDim * static_cast<uint32_t>(sizeof(Element));
85+ uint32_t perRow = rDim * static_cast<uint32_t>(sizeof(float)) * 5u + yRowBytes * 2u;
86+ rowsOnceMax_ = 1;
87+ if (perRow > 0 && perRow <= kUbBudget) {
88+ uint32_t fit = kUbBudget / perRow;
89+ if (fit > 0) {
90+ rowsOnceMax_ = fit;
91+ }
92+ }
93+ }
94+ 
95+ // c0FromPreAct：c0dp 路径下 C0 读入来自 down_pre_act（保持原实现不带 L2 hint），
96+ // 否则来自 c0 槽工作区。各 offset 均为相对 Init 基址的元素偏移。
97+ template <class Resource>
98+ CATLASS_DEVICE void operator()(Resource& resource, int64_t offC0, bool c0FromPreAct, int64_t offDown,
99+ int64_t offInject, uint32_t actualM, uint32_t nWs, uint32_t nConcat, uint32_t rDim,
100+ uint32_t channels, bool useCombine, bool disableSubSplit = false,
101+ bool savePreAct = false, int64_t offDownPreAct = 0, int64_t offGate = 0)
102+ {
103+ using namespace AscendC;
104+ if (actualM == 0 || rDim == 0) {
105+ return;
106+ }
107+ const uint32_t subNum = GetSubBlockNum() == 0 ? 1 : GetSubBlockNum();
108+ const uint32_t subIdx = GetSubBlockIdx();
109+ const uint32_t row0 = (actualM + subNum - 1) / subNum;
110+ uint32_t rowStart = subIdx * row0;
111+ uint32_t rowEnd = rowStart + row0 < actualM ? rowStart + row0 : actualM;
112+ if (disableSubSplit) {
113+ rowStart = 0;
114+ rowEnd = actualM;
115+ }
116+ const uint32_t rows = rowEnd > rowStart ? rowEnd - rowStart : 0;
117+ if (rows == 0) {
118+ return;
119+ }
120+ (void)nConcat;
121+ (void)nWs;
122+ const float invC = invC_;
123+ const bool doInject = useCombine && channels > 0;
124+ uint32_t rowsOnce = rowsOnceMax_;
125+ if (rowsOnce > rows) {
126+ rowsOnce = rows;
127+ }
128+ 
129+ // 每任务仅做基址 offset 平移（operator[] 返回带偏移的 GlobalTensor 拷贝，
130+ // L2 hint 编码在地址高位，随偏移保留），不再执行 SetGlobalBuffer。
131+ GlobalTensor<Element> gC0 = (c0FromPreAct && hasPreAct_) ? gC0FromDpa_[offC0] : gC0_[offC0];
132+ GlobalTensor<Element> gGate = gGate_[offGate];
133+ GlobalTensor<Element> gDown = gDown_[offDown];
134+ GlobalTensor<Element> gDownPreAct;
135+ if (savePreAct && hasPreAct_) {
136+ gDownPreAct = gDownPreAct_[offDownPreAct];
137+ }
138+ 
139+ uint32_t slotFloats = rowsOnce * rDim;
140+ uint32_t slotBytes = (slotFloats * static_cast<uint32_t>(sizeof(float)) + 31u) & ~31u;
141+ uint32_t slotStride = slotBytes / static_cast<uint32_t>(sizeof(float));
142+ uint32_t yBytes = 0;
143+ uint32_t yStride = 0;
144+ if constexpr (!std::is_same_v<Element, float>) {
145+ yBytes = (slotFloats * static_cast<uint32_t>(sizeof(Element)) + 31u) & ~31u;
146+ yStride = yBytes / static_cast<uint32_t>(sizeof(Element));
147+ }
148+ 
149+ LocalTensor<uint8_t> ub = resource.ubBuf.template GetBufferByByte<uint8_t>(0);
150+ uint32_t offb = 0;
151+ LocalTensor<float> ubIn = ub[offb].template ReinterpretCast<float>();
152+ offb += slotBytes * 2u;
153+ LocalTensor<float> ubOut = ub[offb].template ReinterpretCast<float>();
154+ offb += slotBytes * 2u;
155+ LocalTensor<float> work = ub[offb].template ReinterpretCast<float>();
156+ offb += slotBytes;
157+ LocalTensor<Element> y = ub[offb].template ReinterpretCast<Element>();
158+ if constexpr (!std::is_same_v<Element, float>) {
159+ offb += yBytes * 2u;
160+ }
161+ (void)offb;
162+ 
163+ DataCopyExtParams inParams;
164+ inParams.blockCount = 1;
165+ inParams.srcStride = 0;
166+ inParams.dstStride = 0;
167+ DataCopyPadExtParams<Element> inPad{false, 0, 0, static_cast<Element>(0)};
168+ DataCopyExtParams outParams;
169+ outParams.blockCount = 1;
170+ outParams.srcStride = 0;
171+ outParams.dstStride = 0;
172+ SetMaskNorm();
173+ ResetMask();
174+ static constexpr uint32_t kEv[2] = {EVENT_ID0, EVENT_ID1};
175+ SetFlag<HardEvent::V_MTE2>(kEv[0]);
176+ SetFlag<HardEvent::V_MTE2>(kEv[1]);
177+ SetFlag<HardEvent::MTE3_V>(kEv[0]);
178+ SetFlag<HardEvent::MTE3_V>(kEv[1]);
179+ uint32_t round = 0;
180+ for (uint32_t base = 0; base < rows; base += rowsOnce) {
181+ uint32_t stage = round & 1u;
182+ event_t ev = static_cast<event_t>(kEv[stage]);
183+ uint32_t nThis = rows - base < rowsOnce ? rows - base : rowsOnce;
184+ uint32_t nElem = nThis * rDim;
185+ uint32_t m0 = rowStart + base;
186+ LocalTensor<float> src = ubIn[stage * slotStride];
187+ LocalTensor<float> dst = ubOut[stage * slotStride];
188+ LocalTensor<Element> ySlot = y[stage * yStride];
189+ WaitFlag<HardEvent::V_MTE2>(ev);
190+ inParams.blockLen = nElem * static_cast<uint32_t>(sizeof(Element));
191+ DataCopyPad(src.template ReinterpretCast<Element>(), gC0[static_cast<int64_t>(m0) * rDim], inParams, inPad);
192+ SetFlag<HardEvent::MTE2_V>(ev);
193+ WaitFlag<HardEvent::MTE3_V>(ev);
194+ WaitFlag<HardEvent::MTE2_V>(ev);
195+ if constexpr (!std::is_same_v<Element, float>) {
196+ Cast(work, src.template ReinterpretCast<Element>(), RoundMode::CAST_NONE, nElem);
197+ PipeBarrier<PIPE_V>();
198+ if ((nElem & 7u) == 0u) {
199+ DataCopy(src, work, nElem);
200+ } else {
201+ Adds(src, work, 0.0f, nElem);
202+ }
203+ PipeBarrier<PIPE_V>();
204+ }
205+ // c0dp 语义：down_pre_act 统一保存 C0 原始输出（未除 C），与大分支 C0 直写一致，
206+ // 反向自行重算 /C。保存必须在下面 invC 的 Muls 之前做，否则存的是已除 C 的值。
207+ if (savePreAct) {
208+ if constexpr (!std::is_same_v<Element, float>) {
209+ Cast(ySlot, src, RoundMode::CAST_RINT, nElem);
210+ PipeBarrier<PIPE_V>();
211+ }
212+ SetFlag<HardEvent::V_MTE3>(EVENT_ID2);
213+ WaitFlag<HardEvent::V_MTE3>(EVENT_ID2);
214+ outParams.blockLen = nElem * static_cast<uint32_t>(sizeof(Element));
215+ if constexpr (std::is_same_v<Element, float>) {
216+ auto pre = src.template ReinterpretCast<Element>();
217+ DataCopyPad(gDownPreAct[static_cast<int64_t>(m0) * rDim], pre, outParams);
218+ } else {
219+ DataCopyPad(gDownPreAct[static_cast<int64_t>(m0) * rDim], ySlot, outParams);
220+ }
221+ SetFlag<HardEvent::MTE3_V>(EVENT_ID2);
222+ WaitFlag<HardEvent::MTE3_V>(EVENT_ID2);
223+ }
224+ for (uint32_t cursor = 0; cursor < nElem; cursor += COMPUTE_LENGTH) {
225+ uint32_t n = nElem - cursor < COMPUTE_LENGTH ? nElem - cursor : COMPUTE_LENGTH;
226+ Muls(src[cursor], src[cursor], invC, n);
227+ PipeBarrier<PIPE_V>();
228+ }
229+ for (uint32_t cursor = 0; cursor < nElem; cursor += COMPUTE_LENGTH) {
230+ uint32_t n = nElem - cursor < COMPUTE_LENGTH ? nElem - cursor : COMPUTE_LENGTH;
231+ auto d = src[cursor];
232+ auto o = dst[cursor];
233+ auto e = work[cursor];
234+ Muls(e, d, -1.0f, n);
235+ PipeBarrier<PIPE_V>();
236+ Exp(e, e, n);
237+ PipeBarrier<PIPE_V>();
238+ Adds(e, e, 1.0f, n);
239+ PipeBarrier<PIPE_V>();
240+ Div(o, d, e, n);
241+ PipeBarrier<PIPE_V>();
242+ }
243+ SetFlag<HardEvent::V_MTE2>(ev);
244+ if constexpr (!std::is_same_v<Element, float>) {
245+ Cast(ySlot, dst, RoundMode::CAST_RINT, nElem);
246+ PipeBarrier<PIPE_V>();
247+ }
248+ SetFlag<HardEvent::V_MTE3>(ev);
249+ WaitFlag<HardEvent::V_MTE3>(ev);
250+ outParams.blockLen = nElem * static_cast<uint32_t>(sizeof(Element));
251+ if constexpr (std::is_same_v<Element, float>) {
252+ auto out = dst.template ReinterpretCast<Element>();
253+ DataCopyPad(gDown[static_cast<int64_t>(m0) * rDim], out, outParams);
254+ } else {
255+ DataCopyPad(gDown[static_cast<int64_t>(m0) * rDim], ySlot, outParams);
256+ }
257+ SetFlag<HardEvent::MTE3_V>(ev);
258+ ++round;
259+ }
260+ if (doInject) {
261+ uint32_t cap = slotFloats / channels;
262+ if (cap == 0) {
263+ cap = 1;
264+ }
265+ if (cap > 65535u) {
266+ cap = 65535u;
267+ }
268+ GlobalTensor<Element> gInj = gInj_[offInject];
269+ DataCopyExtParams gateParams;
270+ gateParams.blockCount = 1;
271+ gateParams.srcStride = 0;
272+ gateParams.dstStride = 0;
273+ for (uint32_t base = 0; base < rows; base += cap) {
274+ uint32_t stage = round & 1u;
275+ event_t ev = static_cast<event_t>(kEv[stage]);
276+ uint32_t nInj = rows - base < cap ? rows - base : cap;
277+ uint32_t nElem = nInj * channels;
278+ uint32_t m0 = rowStart + base;
279+ LocalTensor<float> src = ubIn[stage * slotStride];
280+ LocalTensor<float> dst = ubOut[stage * slotStride];
281+ LocalTensor<Element> ySlot = y[stage * yStride];
282+ WaitFlag<HardEvent::V_MTE2>(ev);
283+ gateParams.blockLen = nElem * static_cast<uint32_t>(sizeof(Element));
284+ DataCopyPad(src.template ReinterpretCast<Element>(), gGate[static_cast<int64_t>(m0) * channels],
285+ gateParams, inPad);
286+ SetFlag<HardEvent::MTE2_V>(ev);
287+ WaitFlag<HardEvent::MTE3_V>(ev);
288+ WaitFlag<HardEvent::MTE2_V>(ev);
289+ if constexpr (!std::is_same_v<Element, float>) {
290+ Cast(work, src.template ReinterpretCast<Element>(), RoundMode::CAST_NONE, nElem);
291+ PipeBarrier<PIPE_V>();
292+ if ((nElem & 7u) == 0u) {
293+ DataCopy(src, work, nElem);
294+ } else {
295+ Adds(src, work, 0.0f, nElem);
296+ }
297+ PipeBarrier<PIPE_V>();
298+ }
299+ HcInjectCompute(dst, src, work, nElem, invC);
300+ SetFlag<HardEvent::V_MTE2>(ev);
301+ if constexpr (!std::is_same_v<Element, float>) {
302+ Cast(ySlot, dst, RoundMode::CAST_RINT, nElem);
303+ PipeBarrier<PIPE_V>();
304+ }
305+ SetFlag<HardEvent::V_MTE3>(ev);
306+ WaitFlag<HardEvent::V_MTE3>(ev);
307+ outParams.blockLen = nElem * static_cast<uint32_t>(sizeof(Element));
308+ if constexpr (std::is_same_v<Element, float>) {
309+ auto out = dst.template ReinterpretCast<Element>();
310+ DataCopyPad(gInj[static_cast<int64_t>(m0) * channels], out, outParams);
311+ } else {
312+ DataCopyPad(gInj[static_cast<int64_t>(m0) * channels], ySlot, outParams);
313+ }
314+ SetFlag<HardEvent::MTE3_V>(ev);
315+ ++round;
316+ }
317+ }
318+ WaitFlag<HardEvent::V_MTE2>(kEv[0]);
319+ WaitFlag<HardEvent::V_MTE2>(kEv[1]);
320+ WaitFlag<HardEvent::MTE3_V>(kEv[0]);
321+ WaitFlag<HardEvent::MTE3_V>(kEv[1]);
322+ }
323+ 
324+private:
325+ // kernel 入口 Init 一次的常驻基址；c0 槽工作区同时充当 gate 槽基址。
326+ AscendC::GlobalTensor<Element> gC0_;
327+ AscendC::GlobalTensor<Element> gGate_;
328+ AscendC::GlobalTensor<Element> gDown_;
329+ AscendC::GlobalTensor<Element> gInj_;
330+ // down_pre_act 双视图：保存路径带 CACHE_MODE_DISABLE，c0dp 读路径保持无 hint。
331+ AscendC::GlobalTensor<Element> gDownPreAct_;
332+ AscendC::GlobalTensor<Element> gC0FromDpa_;
333+ bool hasPreAct_ = false;
334+ float invC_ = 0.0f;
335+ uint32_t rowsOnceMax_ = 1;
336+};
337+ 
338+} // namespace Catlass::Epilogue
Amhc/text_gated_residual/op_kernel/arch22/block/block_epilogue_hc_v2.hpp+257-0
@@ -0,0 +1,257 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+#include "../../hc_vector_fallback.hpp"
13+ 
14+#include "catlass/catlass.hpp"
15+#include "catlass/arch/resource.hpp"
16+#include "../hc_gm_copy.hpp"
17+ 
18+namespace Catlass::Epilogue {
19+ 
20+struct EpilogueHcV2GateMean {
21+ using ArchTag = Arch::AtlasA2;
22+ static constexpr uint32_t UB_STAGES = 1;
23+};
24+ 
25+template <class Element, class ArchTag = Arch::AtlasA2>
26+class BlockEpilogueHcV2 {
27+public:
28+ CATLASS_DEVICE
29+ BlockEpilogueHcV2() {}
30+ 
31+ // up / xn / mixed 基址只绑一次。大 shape 的 L2 hint 在这里设置，小 shape 不关 L2。
32+ CATLASS_DEVICE
33+ void Init(GM_ADDR gmUp, GM_ADDR gmXn, GM_ADDR gmMixed, bool l2HintUp, bool l2HintXn)
34+ {
35+ using namespace AscendC;
36+ gUp_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmUp));
37+ if (l2HintUp) {
38+ gUp_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
39+ }
40+ gXn_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmXn));
41+ if (l2HintXn) {
42+ gXn_.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
43+ }
44+ gMixed_.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(gmMixed));
45+ }
46+ 
47+ template <class Resource>
48+ CATLASS_DEVICE void operator()(Resource& resource, int64_t offUp, int64_t offXn, int64_t offMixed, uint32_t actualM,
49+ uint32_t channels, uint32_t hidden, bool disableSubSplit = false)
50+ {
51+ using namespace AscendC;
52+ if (actualM == 0 || hidden == 0 || channels == 0) {
53+ return;
54+ }
55+ const uint32_t subNum = GetSubBlockNum() == 0 ? 1 : GetSubBlockNum();
56+ const uint32_t subIdx = GetSubBlockIdx();
57+ const uint32_t row0 = (actualM + subNum - 1) / subNum;
58+ uint32_t rowStart = subIdx * row0;
59+ uint32_t rowEnd = rowStart + row0 < actualM ? rowStart + row0 : actualM;
60+ if (disableSubSplit) {
61+ rowStart = 0;
62+ rowEnd = actualM;
63+ }
64+ const float invC = 1.0f / static_cast<float>(static_cast<int32_t>(channels));
65+ const uint32_t d = channels * hidden;
66+ 
67+ GlobalTensor<Element> gUp = gUp_[offUp];
68+ GlobalTensor<Element> gXn = gXn_[offXn];
69+ GlobalTensor<Element> gMixed = gMixed_[offMixed];
70+ 
71+ if (static_cast<uint64_t>(d) * 14u + static_cast<uint64_t>(hidden) * 8u + 256u > ArchTag::UB_SIZE) {
72+ HcV2Striped<Element>(resource, gUp, gXn, gMixed, rowStart, rowEnd, channels, hidden);
73+ return;
74+ }
75+ uint32_t off = 0;
76+ LocalTensor<uint8_t> ub = resource.ubBuf.template GetBufferByByte<uint8_t>(0);
77+ LocalTensor<float> ubUp = ub[off].template ReinterpretCast<float>();
78+ off += d * sizeof(float) + 32u;
79+ off = (off + 31u) & ~31u;
80+ LocalTensor<Element> ubXnT = ub[off].template ReinterpretCast<Element>();
81+ off += d * sizeof(Element) + 32u;
82+ off = (off + 31u) & ~31u;
83+ LocalTensor<float> ubXn;
84+ if constexpr (std::is_same_v<Element, float>) {
85+ ubXn = ubXnT.template ReinterpretCast<float>();
86+ } else {
87+ ubXn = ub[off].template ReinterpretCast<float>();
88+ off += d * sizeof(float);
89+ off = (off + 31u) & ~31u;
90+ }
91+ LocalTensor<float> ubTmp = ub[off].template ReinterpretCast<float>();
92+ off += d * sizeof(float);
93+ off = (off + 31u) & ~31u;
94+ LocalTensor<float> ubOnes = ub[off].template ReinterpretCast<float>();
95+ off += 8u * sizeof(float);
96+ off = (off + 31u) & ~31u;
97+ // 输出双缓冲：MTE3 读一份时，下一行写入另一份。计算用的 ubMixed 在 BF16 上仍是单缓冲，
98+ // Cast 读完就可以被下一行复用。
99+ constexpr uint32_t kOutStages = 2;
100+ uint32_t mixedBytes = (hidden * static_cast<uint32_t>(sizeof(float)) + 31u) & ~31u;
101+ uint32_t mixedStride = mixedBytes / static_cast<uint32_t>(sizeof(float));
102+ LocalTensor<float> ubMixed = ub[off].template ReinterpretCast<float>();
103+ off += mixedBytes * (std::is_same_v<Element, float> ? kOutStages : 1u);
104+ off = (off + 31u) & ~31u;
105+ uint32_t outBytes = (hidden * static_cast<uint32_t>(sizeof(Element)) + 31u) & ~31u;
106+ uint32_t outStride = outBytes / static_cast<uint32_t>(sizeof(Element));
107+ LocalTensor<Element> ubMixedT = ub[off].template ReinterpretCast<Element>();
108+ if constexpr (!std::is_same_v<Element, float>) {
109+ off += outBytes * kOutStages;
110+ }
111+ (void)off;
112+ (void)mixedStride;
113+ (void)outStride;
114+ 
115+ if (rowEnd <= rowStart) {
116+ return;
117+ }
118+ Duplicate(ubOnes, 1.0f, 8);
119+ PipeBarrier<PIPE_V>();
120+ // 输出槽用事件数组做 ping-pong。进入循环前先 Set 两次，循环里按 i&1 取槽，不再写 if。
121+ static constexpr uint32_t kOutEv[2] = {EVENT_ID0, EVENT_ID1};
122+ SetFlag<HardEvent::MTE3_V>(kOutEv[0]);
123+ SetFlag<HardEvent::MTE3_V>(kOutEv[1]);
124+ 
125+ DataCopyExtParams upParams;
126+ upParams.blockCount = 1;
127+ upParams.blockLen = d * static_cast<uint32_t>(sizeof(Element));
128+ upParams.srcStride = 0;
129+ upParams.dstStride = 0;
130+ DataCopyPadExtParams<Element> upPad{false, 0, 0, static_cast<Element>(0)};
131+ DataCopyExtParams xnParams;
132+ xnParams.blockCount = 1;
133+ xnParams.blockLen = d * static_cast<uint32_t>(sizeof(Element));
134+ xnParams.srcStride = 0;
135+ xnParams.dstStride = 0;
136+ DataCopyPadExtParams<Element> xnPad{false, 0, 0, static_cast<Element>(0)};
137+ DataCopyExtParams mixedParams;
138+ mixedParams.blockCount = 1;
139+ mixedParams.blockLen = hidden * static_cast<uint32_t>(sizeof(Element));
140+ mixedParams.srcStride = 0;
141+ mixedParams.dstStride = 0;
142+ 
143+ PipeBarrier<PIPE_ALL>();
144+ for (uint32_t m = rowStart; m < rowEnd; ++m) {
145+ uint32_t stage = (m - rowStart) & 1u;
146+ event_t slotEv = static_cast<event_t>(kOutEv[stage]);
147+ uint32_t mNext = m + 1u;
148+ bool hasNext = mNext < rowEnd;
149+ if (m == rowStart) {
150+ DataCopyPad(ubUp.template ReinterpretCast<Element>(), gUp[static_cast<int64_t>(m) * d], upParams,
151+ upPad);
152+ DataCopyPad(ubXnT, gXn[static_cast<int64_t>(m) * d], xnParams, xnPad);
153+ SetFlag<HardEvent::MTE2_V>(EVENT_ID0);
154+ SetFlag<HardEvent::MTE2_V>(EVENT_ID1);
155+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID0);
156+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID1);
157+ } else {
158+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID2);
159+ WaitFlag<HardEvent::MTE2_V>(EVENT_ID3);
160+ }
161+ if constexpr (std::is_same_v<Element, float>) {
162+ WaitFlag<HardEvent::MTE3_V>(slotEv);
163+ } else {
164+ Cast(ubXn, ubXnT, RoundMode::CAST_NONE, d);
165+ if (hasNext) {
166+ SetFlag<HardEvent::V_MTE2>(EVENT_ID1);
167+ WaitFlag<HardEvent::V_MTE2>(EVENT_ID1);
168+ DataCopyPad(ubXnT, gXn[static_cast<int64_t>(mNext) * d], xnParams, xnPad);
169+ SetFlag<HardEvent::MTE2_V>(EVENT_ID3);
170+ }
171+ }
172+ // 取负之后 up 不再被读。下一条 up 从这里开始搬，和后面的 Exp/Div 重叠。
173+ // Expand the Cube output in the existing float scratch.
174+ if constexpr (!std::is_same_v<Element, float>) {
175+ Cast(ubTmp, ubUp.template ReinterpretCast<Element>(), RoundMode::CAST_NONE, d);
176+ PipeBarrier<PIPE_V>();
177+ DataCopy(ubUp, ubTmp, d);
178+ PipeBarrier<PIPE_V>();
179+ }
180+ Muls(ubTmp, ubUp, -1.0f, d);
181+ if (hasNext) {
182+ SetFlag<HardEvent::V_MTE2>(EVENT_ID0);
183+ WaitFlag<HardEvent::V_MTE2>(EVENT_ID0);
184+ DataCopyPad(ubUp.template ReinterpretCast<Element>(), gUp[static_cast<int64_t>(mNext) * d], upParams,
185+ upPad);
186+ SetFlag<HardEvent::MTE2_V>(EVENT_ID2);
187+ }
188+ LocalTensor<float> mixedRow = ubMixed;
189+ if constexpr (std::is_same_v<Element, float>) {
190+ mixedRow = ubMixed[stage * mixedStride];
191+ }
192+ SigmoidMulAcc(mixedRow, ubOnes, ubXn, ubTmp, channels, hidden, invC);
193+ if constexpr (std::is_same_v<Element, float>) {
194+ if (hasNext) {
195+ SetFlag<HardEvent::V_MTE2>(EVENT_ID1);
196+ WaitFlag<HardEvent::V_MTE2>(EVENT_ID1);
197+ DataCopyPad(ubXnT, gXn[static_cast<int64_t>(mNext) * d], xnParams, xnPad);
198+ SetFlag<HardEvent::MTE2_V>(EVENT_ID3);
199+ }
200+ SetFlag<HardEvent::V_MTE3>(slotEv);
201+ WaitFlag<HardEvent::V_MTE3>(slotEv);
202+ DataCopyPad(gMixed[static_cast<int64_t>(m) * hidden], mixedRow, mixedParams);
203+ } else {
204+ WaitFlag<HardEvent::MTE3_V>(slotEv);
205+ LocalTensor<Element> mixedT = ubMixedT[stage * outStride];
206+ Cast(mixedT, mixedRow, RoundMode::CAST_RINT, hidden);
207+ SetFlag<HardEvent::V_MTE3>(slotEv);
208+ WaitFlag<HardEvent::V_MTE3>(slotEv);
209+ DataCopyPad(gMixed[static_cast<int64_t>(m) * hidden], mixedT, mixedParams);
210+ }
211+ SetFlag<HardEvent::MTE3_V>(slotEv);
212+ }
213+ WaitFlag<HardEvent::MTE3_V>(kOutEv[0]);
214+ WaitFlag<HardEvent::MTE3_V>(kOutEv[1]);
215+ PipeBarrier<PIPE_ALL>();
216+ }
217+ 
218+private:
219+ // ubTmp 进入时已经是 -up。ones 是 1 个 block 的 1，Counter 模式下 src0 步长为 0。
220+ CATLASS_DEVICE
221+ static void SigmoidMulAcc(AscendC::LocalTensor<float> const& mixedRow, AscendC::LocalTensor<float> const& ubOnes,
222+ AscendC::LocalTensor<float> const& ubXn, AscendC::LocalTensor<float> const& ubTmp,
223+ uint32_t channels, uint32_t hidden, float invC)
224+ {
225+ using namespace AscendC;
226+ const uint32_t d = channels * hidden;
227+ Exp(ubTmp, ubTmp, d);
228+ PipeBarrier<PIPE_V>();
229+ Adds(ubTmp, ubTmp, 1.0f, d);
230+ PipeBarrier<PIPE_V>();
231+ // Counter 一次 Div：mask=d，repeat=1。src0 是 1 个 block 的 1，blk/rep 步长为 0。
232+ const BinaryRepeatParams divRep{1, 0, 1, 8, 0, 8};
233+ SetMaskCount();
234+ SetVectorMask<float, MaskMode::COUNTER>(d);
235+ Div<float, false>(ubTmp, ubOnes, ubTmp, MASK_PLACEHOLDER, 1, divRep);
236+ PipeBarrier<PIPE_V>();
237+ SetMaskNorm();
238+ ResetMask();
239+ Mul(ubTmp, ubTmp, ubXn, d);
240+ PipeBarrier<PIPE_V>();
241+ Muls(ubTmp, ubTmp, invC, d);
242+ PipeBarrier<PIPE_V>();
243+ Duplicate(mixedRow, 0.0f, hidden);
244+ PipeBarrier<PIPE_V>();
245+ for (uint32_t c = 0; c < channels; ++c) {
246+ uint32_t off = c * hidden;
247+ Add(mixedRow, mixedRow, ubTmp[off], hidden);
248+ PipeBarrier<PIPE_V>();
249+ }
250+ }
251+ 
252+ AscendC::GlobalTensor<Element> gUp_;
253+ AscendC::GlobalTensor<Element> gXn_;
254+ AscendC::GlobalTensor<Element> gMixed_;
255+};
256+ 
257+} // namespace Catlass::Epilogue
Amhc/text_gated_residual/op_kernel/arch22/hc_gm_copy.hpp+117-0
@@ -0,0 +1,117 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+ 
13+#include "catlass/catlass.hpp"
14+ 
15+namespace Catlass::Epilogue {
16+ 
17+// A2 向量 Rsqrt 的结果只有约 8 bit 尾数，过不了 FP32 混合容差。标量牛顿迭代给出完整 FP32。
18+CATLASS_DEVICE
19+float HcRsqrtF32(float x)
20+{
21+ if (!(x > 0.0f)) {
22+ return 0.0f;
23+ }
24+ union {
25+ float f;
26+ uint32_t u;
27+ } bits;
28+ bits.f = x;
29+ bits.u = 0x5f375a86u - (bits.u >> 1);
30+ float y = bits.f;
31+ float halfX = 0.5f * x;
32+ y = y * (1.5f - halfX * y * y);
33+ y = y * (1.5f - halfX * y * y);
34+ y = y * (1.5f - halfX * y * y);
35+ y = y * (1.5f - halfX * y * y);
36+ return y;
37+}
38+ 
39+// Atlas A2：GM 地址和 UB 地址都要 32 字节对齐。长度不是 32 字节时用 DataCopyPad，
40+// blockLen 写实际有效字节数，不开启 isPad。后续计算按有效长度走，尾部 dummy 不参与。
41+ 
42+template <typename T>
43+CATLASS_DEVICE void HcCopyGmToUb(AscendC::LocalTensor<T> const& dst, AscendC::GlobalTensor<T> const& src, uint32_t n)
44+{
45+ using namespace AscendC;
46+ if (n == 0) {
47+ return;
48+ }
49+ uint32_t bytes = n * static_cast<uint32_t>(sizeof(T));
50+ if ((bytes % 32u) == 0u) {
51+ DataCopy(dst, src, n);
52+ return;
53+ }
54+ DataCopyExtParams p{1, bytes, 0, 0, 0};
55+ DataCopyPadExtParams<T> pad{false, 0, 0, static_cast<T>(0)};
56+ DataCopyPad(dst, src, p, pad);
57+}
58+ 
59+template <typename T>
60+CATLASS_DEVICE void HcCopyUbToGm(AscendC::GlobalTensor<T> const& dst, AscendC::LocalTensor<T> const& src, uint32_t n)
61+{
62+ using namespace AscendC;
63+ if (n == 0) {
64+ return;
65+ }
66+ uint32_t bytes = n * static_cast<uint32_t>(sizeof(T));
67+ if ((bytes % 32u) == 0u) {
68+ DataCopy(dst, src, n);
69+ return;
70+ }
71+ DataCopyExtParams p{1, bytes, 0, 0, 0};
72+ DataCopyPad(dst, src, p);
73+}
74+ 
75+// scratch 至少 32 字节且 32 字节对齐，不与 src 重叠。base 必须是 32 字节对齐的张量起点。
76+template <typename T>
77+CATLASS_DEVICE void HcStoreGmElems(AscendC::GlobalTensor<T> const& base, int64_t elemOff,
78+ AscendC::LocalTensor<T> const& src, uint32_t n,
79+ AscendC::LocalTensor<T> const& scratch)
80+{
81+ using namespace AscendC;
82+ if (n == 0) {
83+ return;
84+ }
85+ constexpr uint32_t kElem = static_cast<uint32_t>(sizeof(T));
86+ constexpr uint32_t kBlockElems = 32u / kElem;
87+ uint32_t done = 0;
88+ while (done < n) {
89+ int64_t absElem = elemOff + done;
90+ int64_t byteOff = absElem * kElem;
91+ int64_t alignedByte = byteOff & ~int64_t{31};
92+ uint32_t prefixElems = static_cast<uint32_t>((byteOff - alignedByte) / kElem);
93+ uint32_t room = kBlockElems - prefixElems;
94+ uint32_t chunk = (n - done) < room ? (n - done) : room;
95+ int64_t startElem = alignedByte / kElem;
96+ if (prefixElems == 0 && chunk == kBlockElems) {
97+ DataCopy(base[startElem], src[done], chunk);
98+ } else {
99+ DataCopy(scratch, base[startElem], kBlockElems);
100+ SetFlag<HardEvent::MTE2_S>(EVENT_ID3);
101+ WaitFlag<HardEvent::MTE2_S>(EVENT_ID3);
102+ SetFlag<HardEvent::V_S>(EVENT_ID3);
103+ WaitFlag<HardEvent::V_S>(EVENT_ID3);
104+ for (uint32_t i = 0; i < chunk; ++i) {
105+ scratch.SetValue(prefixElems + i, src.GetValue(done + i));
106+ }
107+ SetFlag<HardEvent::S_MTE3>(EVENT_ID3);
108+ WaitFlag<HardEvent::S_MTE3>(EVENT_ID3);
109+ DataCopy(base[startElem], scratch, kBlockElems);
110+ SetFlag<HardEvent::MTE3_MTE2>(EVENT_ID3);
111+ WaitFlag<HardEvent::MTE3_MTE2>(EVENT_ID3);
112+ }
113+ done += chunk;
114+ }
115+}
116+ 
117+} // namespace Catlass::Epilogue
Amhc/text_gated_residual/op_kernel/arch22/kernel/hc_forward_large_kernel.hpp+222-0
@@ -0,0 +1,222 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+#include "hc_forward_mul_sum_kernel.hpp"
13+namespace Catlass::Gemm::Kernel {
14+template <class Element, uint32_t C1SwizzleDir, class ArchTag_ = Arch::AtlasA2,
15+ class Tiling = HcForwardMulSumTilingData>
16+struct HcForwardLargeKernel : HcForwardCommon<Element, C1SwizzleDir, false, ArchTag_, Tiling> {
17+ using Base = HcForwardCommon<Element, C1SwizzleDir, false, ArchTag_, Tiling>;
18+ using Params = typename Base::Params;
19+ using ArchTag = typename Base::ArchTag;
20+ using LayoutTagA = typename Base::LayoutTagA;
21+ using LayoutTagC = typename Base::LayoutTagC;
22+ using BlockMmadC0 = typename Base::BlockMmadC0;
23+ using BlockMmadC1 = typename Base::BlockMmadC1;
24+ using BlockSchedulerC1 = typename Base::BlockSchedulerC1;
25+ using Base::L1M0;
26+ using Base::L1N0;
27+ using Base::L1M1;
28+ using Base::L1N1;
29+ template <class TensorA, class TensorBDown, class TensorBGate, class TensorC, class TensorGate>
30+ CATLASS_DEVICE void RunLargeC0NConcat(Arch::Resource<ArchTag>& resource, TensorA const& tensorA,
31+ TensorBDown const& tensorBDown, TensorBGate const& tensorBGate,
32+ TensorC const& tensorC0, TensorGate const& tensorGate, uint32_t row,
33+ uint32_t actualM, uint32_t rDim, uint32_t channels, uint32_t kDim,
34+ bool useCombine, uint32_t rowC0 = 0)
35+ {
36+ // BlockMmad 把整块 actualN 放进一块 L0C。L0C 只按 L1N0 分配，N 超过 L1N0 时 FIXP 读越界。
37+ // rowC0：down 输出张量的行偏移（c0dp 时 down 直写 down_pre_act 全区域张量，按任务行寻址）。
38+ // gate 恒写槽内（mTask, channels）张量，坐标保持槽内 mOff，不吃 rowC0。
39+ BlockMmadC0 mmad(resource);
40+ for (uint32_t mOff = 0; mOff < actualM; mOff += L1M0) {
41+ uint32_t tileM = actualM - mOff < L1M0 ? actualM - mOff : L1M0;
42+ auto tileA = GetTile(tensorA, tla::MakeCoord(row + mOff, 0), tla::MakeShape(tileM, kDim));
43+ for (uint32_t nOff = 0; nOff < rDim; nOff += L1N0) {
44+ uint32_t tileN = rDim - nOff < L1N0 ? rDim - nOff : L1N0;
45+ auto tileB = GetTile(tensorBDown, tla::MakeCoord(0, nOff), tla::MakeShape(kDim, tileN));
46+ auto tileC = GetTile(tensorC0, tla::MakeCoord(rowC0 + mOff, nOff), tla::MakeShape(tileM, tileN));
47+ mmad(tileA, tileB, tileC, GemmCoord{tileM, tileN, kDim});
48+ }
49+ if (useCombine && channels > 0) {
50+ auto tileBg = GetTile(tensorBGate, tla::MakeCoord(0, 0), tla::MakeShape(kDim, channels));
51+ auto tileCi = GetTile(tensorGate, tla::MakeCoord(mOff, 0), tla::MakeShape(tileM, channels));
52+ mmad(tileA, tileBg, tileCi, GemmCoord{tileM, channels, kDim});
53+ }
54+ }
55+ }
56+ 
57+ template <class Context>
58+ CATLASS_DEVICE void RunBody(Context& c)
59+ {
60+ using namespace AscendC;
61+ auto& params = c.f_params;
62+ auto& usedAic = c.f_usedAic;
63+ auto& mTask = c.f_mTask;
64+ auto& d = c.f_d;
65+ auto& r = c.f_r;
66+ auto& hidden = c.f_hidden;
67+ auto& channels = c.f_channels;
68+ auto& M = c.f_M;
69+ auto& eps = c.f_eps;
70+ auto& useCombine = c.f_useCombine;
71+ auto& nConcat = c.f_nConcat;
72+ auto& nWs = c.f_nWs;
73+ auto& blockIdx = c.f_blockIdx;
74+ auto& resource = c.f_resource;
75+ auto& xnSlotBytes = c.f_xnSlotBytes;
76+ auto& downSlotBytes = c.f_downSlotBytes;
77+ auto& c0SlotBytes = c.f_c0SlotBytes;
78+ auto& ptrXn = c.f_ptrXn;
79+ auto& ptrDown = c.f_ptrDown;
80+ auto& ptrC0Core = c.f_ptrC0Core;
81+ auto& tensorWDown = c.f_tensorWDown;
82+ auto& tensorWGate = c.f_tensorWGate;
83+ auto& tensorWUp = c.f_tensorWUp;
84+ auto& tensorUp = c.f_tensorUp;
85+ auto& localTasks = c.f_localTasks;
86+ auto& blockV0 = c.f_blockV0;
87+ auto& blockV1 = c.f_blockV1;
88+ auto& blockV2 = c.f_blockV2;
89+ 
90+ // fused 大路径按核环形槽：xn 每核 3 槽、down 每核 2 槽、
91+ // c0 每核 2 槽。c0 槽内 down 段在前、gate 段紧随，两段交错不能并成单张量，按槽各建。
92+ // wave 循环内 C0/C1 只用坐标偏移寻址，不再逐 wave 重复 SetGlobalBuffer+MakeTensor。
93+ const bool c0DpLarge = params.ptrDownPreAct != nullptr;
94+ GlobalTensor<Element> gmXnRing;
95+ gmXnRing.SetGlobalBuffer(
96+ reinterpret_cast<__gm__ Element*>(ptrXn + static_cast<int64_t>(blockIdx) * 3u * xnSlotBytes));
97+ // c0 对环形槽里的 xn 只读一次，与 v2 一样不占 L2。这段只在大 shape 路径执行。
98+ gmXnRing.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
99+ auto tensorXnRing =
100+ tla::MakeTensor(gmXnRing, tla::MakeLayout<Element, LayoutTagA>(3u * mTask, d), Arch::PositionGM{});
101+ GlobalTensor<Element> gmDownRing;
102+ gmDownRing.SetGlobalBuffer(
103+ reinterpret_cast<__gm__ Element*>(ptrDown + static_cast<int64_t>(blockIdx) * 2u * downSlotBytes));
104+ auto tensorDownRing =
105+ tla::MakeTensor(gmDownRing, tla::MakeLayout<Element, LayoutTagA>(2u * mTask, r), Arch::PositionGM{});
106+ uint32_t gateN = channels == 0 ? 1u : channels;
107+ auto layoutC0Slot = tla::MakeLayout<Element, LayoutTagC>(mTask, r);
108+ auto layoutGateSlot = tla::MakeLayout<Element, LayoutTagC>(mTask, gateN);
109+ GlobalTensor<Element> gmC0S0;
110+ gmC0S0.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrC0Core));
111+ auto tensorC0S0 = tla::MakeTensor(gmC0S0, layoutC0Slot, Arch::PositionGM{});
112+ GlobalTensor<Element> gmC0S1;
113+ gmC0S1.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrC0Core + c0SlotBytes));
114+ auto tensorC0S1 = tla::MakeTensor(gmC0S1, layoutC0Slot, Arch::PositionGM{});
115+ GlobalTensor<Element> gmGateS0;
116+ gmGateS0.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(
117+ ptrC0Core + static_cast<int64_t>(mTask) * r * static_cast<int64_t>(sizeof(Element))));
118+ auto tensorGateS0 = tla::MakeTensor(gmGateS0, layoutGateSlot, Arch::PositionGM{});
119+ GlobalTensor<Element> gmGateS1;
120+ gmGateS1.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(
121+ ptrC0Core + c0SlotBytes + static_cast<int64_t>(mTask) * r * static_cast<int64_t>(sizeof(Element))));
122+ auto tensorGateS1 = tla::MakeTensor(gmGateS1, layoutGateSlot, Arch::PositionGM{});
123+ // c0dp：down 直写 down_pre_act 全区域张量（(M, r)），按任务行 rowC0 寻址。
124+ GlobalTensor<Element> gmC0Dpa;
125+ if (c0DpLarge) {
126+ gmC0Dpa.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(params.ptrDownPreAct));
127+ }
128+ auto tensorC0Dpa = tla::MakeTensor(gmC0Dpa, tla::MakeLayout<Element, LayoutTagC>(M, r), Arch::PositionGM{});
129+ 
130+ // fused 默认 V2 delay-2：V0[t] → C0[t] → V1[t-1] → C1[t-1] → V2[t-2]
131+ // 尾部多排空 2 拍（最后 V1/C1 与最后 V2 各一拍）。
132+ for (uint32_t wave = 0; wave <= localTasks + 1; ++wave) {
133+ uint32_t task = blockIdx + wave * usedAic;
134+ if (wave < localTasks) {
135+ uint32_t row = task * mTask;
136+ uint32_t actualM = M - row < mTask ? M - row : mTask;
137+ if ASCEND_IS_AIV {
138+ int64_t offX = static_cast<int64_t>(row) * d;
139+ int64_t offXn = (static_cast<int64_t>(blockIdx) * 3u * xnSlotBytes +
140+ static_cast<int64_t>(wave % 3u) * xnSlotBytes) /
141+ static_cast<int64_t>(sizeof(Element));
142+ int64_t offInv = static_cast<int64_t>(row) * channels;
143+ blockV0(resource, offX, offXn, actualM, d, channels, hidden, eps, useCombine, false, offInv);
144+ HcAivSet(FLAG_V0);
145+ }
146+ if ASCEND_IS_AIC {
147+ HcAicWait(FLAG_V0);
148+ uint32_t slot = wave & 1;
149+ // c0dp：down 直写 down_pre_act 全区域张量（rowC0=任务行，保存原始值未除 C），
150+ // gate 仍写本核 c0 槽；两者均复用循环外一次构造的槽张量。
151+ auto const& tensorC0 = c0DpLarge ? tensorC0Dpa : (slot == 0 ? tensorC0S0 : tensorC0S1);
152+ auto const& tensorGate = slot == 0 ? tensorGateS0 : tensorGateS1;
153+ RunLargeC0NConcat(resource, tensorXnRing, tensorWDown, tensorWGate, tensorC0, tensorGate,
154+ (wave % 3u) * mTask, actualM, r, channels, d, useCombine, c0DpLarge ? row : 0u);
155+ AscendC::PipeBarrier<PIPE_ALL>();
156+ HcAicSet(slot == 0 ? FLAG_C0_0 : FLAG_C0_1);
157+ }
158+ }
159+ if (wave >= 1 && (wave - 1) < localTasks) {
160+ uint32_t prev = blockIdx + (wave - 1) * usedAic;
161+ uint32_t prevSlot = (wave - 1) & 1;
162+ uint32_t row = prev * mTask;
163+ uint32_t actualM = M - row < mTask ? M - row : mTask;
164+ if ASCEND_IS_AIV {
165+ HcAivWait(prevSlot == 0 ? FLAG_C0_0 : FLAG_C0_1);
166+ const int64_t c0SlotEls = c0SlotBytes / static_cast<int64_t>(sizeof(Element));
167+ // c0dp：V1 从 down_pre_act 读 C0 原始输出（offC0 相对 dpa 基址），
168+ // 不再单独写保存；否则从本核 c0 槽读。
169+ const bool v1C0FromDpa = params.ptrDownPreAct != nullptr;
170+ const int64_t offC0 = v1C0FromDpa ? static_cast<int64_t>(row) * r : prevSlot * c0SlotEls;
171+ const int64_t offDown = (static_cast<int64_t>(blockIdx) * 2u + prevSlot) * mTask * r;
172+ blockV1.template operator()<Arch::Resource<ArchTag>>(
173+ resource, offC0, v1C0FromDpa, offDown, static_cast<int64_t>(row) * channels, actualM, nWs,
174+ nConcat, r, channels, useCombine, false, false, 0,
175+ prevSlot * c0SlotEls + static_cast<int64_t>(mTask) * r);
176+ HcAivSet(FLAG_V1);
177+ }
178+ if ASCEND_IS_AIC {
179+ HcAicWait(FLAG_V1);
180+ GemmCoord problem{actualM, d, r};
181+ BlockSchedulerC1 scheduler(problem, MatrixCoord{L1M1, L1N1});
182+ uint32_t nTiles = scheduler.GetCoreLoops();
183+ BlockMmadC1 mmadC1(resource);
184+ for (uint32_t nTile = 0; nTile < nTiles; ++nTile) {
185+ GemmCoord coord = scheduler.GetBlockCoord(nTile);
186+ GemmCoord actual = scheduler.GetActualBlockShape(coord);
187+ uint32_t nOff = coord.n() * L1N1;
188+ auto tileA =
189+ GetTile(tensorDownRing, tla::MakeCoord(prevSlot * mTask, 0), tla::MakeShape(actual.m(), r));
190+ auto tileB = GetTile(tensorWUp, tla::MakeCoord(0, nOff), tla::MakeShape(r, actual.n()));
191+ auto tileC =
192+ GetTile(tensorUp, tla::MakeCoord(row, nOff), tla::MakeShape(actual.m(), actual.n()));
193+ mmadC1(tileA, tileB, tileC, actual);
194+ }
195+ 
196+ AscendC::PipeBarrier<PIPE_ALL>();
197+ HcAicSet(prevSlot == 0 ? FLAG_C1_0 : FLAG_C1_1);
198+ }
199+ }
200+ if (wave >= 2 && (wave - 2) < localTasks) {
201+ uint32_t prev2 = blockIdx + (wave - 2) * usedAic;
202+ uint32_t c1Slot = (wave - 2) & 1;
203+ uint32_t row = prev2 * mTask;
204+ uint32_t actualM = M - row < mTask ? M - row : mTask;
205+ if ASCEND_IS_AIV {
206+ HcAivWait(c1Slot == 0 ? FLAG_C1_0 : FLAG_C1_1);
207+ int64_t offX = static_cast<int64_t>(row) * d;
208+ int64_t offXn = (static_cast<int64_t>(blockIdx) * 3u * xnSlotBytes +
209+ static_cast<int64_t>((wave - 2) % 3u) * xnSlotBytes) /
210+ static_cast<int64_t>(sizeof(Element));
211+ blockV2(resource, offX, offXn, static_cast<int64_t>(row) * hidden, actualM, channels, hidden,
212+ false);
213+ }
214+ }
215+ }
216+ }
217+ CATLASS_DEVICE void operator()(Params const& params)
218+ {
219+ Base::template Execute<HcForwardLargeKernel>(params);
220+ }
221+};
222+} // namespace Catlass::Gemm::Kernel
Amhc/text_gated_residual/op_kernel/arch22/kernel/hc_forward_mul_sum_kernel.hpp+448-0
@@ -0,0 +1,448 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#pragma once
12+ 
13+#ifndef K_MAX_SHAPE_DIM
14+#define K_MAX_SHAPE_DIM 0
15+#endif
16+ 
17+#include "catlass/arch/arch.hpp"
18+#include "catlass/arch/resource.hpp"
19+#include "catlass/catlass.hpp"
20+#include <type_traits>
21+#include "catlass/gemm/block/block_mmad.hpp"
22+#include "catlass/gemm/block/block_swizzle.hpp"
23+#include "catlass/gemm/dispatch_policy.hpp"
24+#include "catlass/gemm/gemm_type.hpp"
25+#include "catlass/gemm/tile/tile_copy.hpp"
26+#include "catlass/layout/layout.hpp"
27+#include "catlass/gemm_coord.hpp"
28+#include "catlass/matrix_coord.hpp"
29+#include "tla/layout.hpp"
30+#include "tla/tensor.hpp"
31+ 
32+#include "../block/block_epilogue_hc_v0.hpp"
33+#include "../block/block_epilogue_hc_v1.hpp"
34+#include "../block/block_epilogue_hc_v2.hpp"
35+#include "../../catlass_hc_forward_mul_sum_tiling_data.h"
36+ 
37+namespace Catlass::Gemm::Kernel {
38+constexpr uint32_t HC_A2_VECTOR_SUBBLOCK_COUNT = 2;
39+constexpr uint32_t HC_A2_SWIZZLE_TILE_GROUP = 3;
40+constexpr uint32_t HC_A2_L1_A_STAGES = 2;
41+constexpr uint32_t HC_A2_L1_B_STAGES = 2;
42+constexpr uint32_t HC_A2_L0_A_STAGES = 2;
43+ 
44+constexpr int64_t HC_SYNC_MODE = 2;
45+constexpr int64_t FLAG_V0 = 0;
46+constexpr int64_t FLAG_C0_0 = 1;
47+constexpr int64_t FLAG_C0_1 = 2;
48+constexpr int64_t FLAG_V1 = 3;
49+constexpr int64_t FLAG_C1_0 = 4; // C1[k] Set slot=k&1；V2[k] Wait 同槽（delay-2，禁止与 C1[k+1] 串同一 flag）
50+constexpr int64_t FLAG_C1_1 = 5;
51+constexpr int64_t HC_FLAG_PAIR = 16;
52+ 
53+// modeId=4：AIV0 的 flag N 对应 AIC N，AIV1 的同一 N 由硬件映射到 AIC N+16。
54+// 对齐 example 49 fai_kernel.h 与 grouped_matmul_slice_m_gelu.hpp：
55+// AIV Set/Wait 只用基 id；AIC 连发/连等 id 与 id+16。
56+// 仅当 MIX 未拉起第二颗 AIV（GetSubBlockNum()<2）时，AIV0 空转补 +16，避免 AIC Wait(16+id) 挂死。
57+CATLASS_DEVICE
58+uint32_t HcSubBlockNum()
59+{
60+ uint32_t n = AscendC::GetSubBlockNum();
61+ return n == 0 ? 1 : n;
62+}
63+ 
64+CATLASS_DEVICE
65+bool HcNeedDummyPair()
66+{
67+ return HcSubBlockNum() < HC_A2_VECTOR_SUBBLOCK_COUNT && AscendC::GetSubBlockIdx() == 0;
68+}
69+ 
70+CATLASS_DEVICE
71+void HcAivSet(int64_t base)
72+{
73+ AscendC::PipeBarrier<PIPE_ALL>();
74+ AscendC::CrossCoreSetFlag<HC_SYNC_MODE, PIPE_MTE3>(base);
75+}
76+ 
77+CATLASS_DEVICE
78+void HcAivWait(int64_t base)
79+{
80+ AscendC::CrossCoreWaitFlag<HC_SYNC_MODE, PIPE_V>(base);
81+}
82+ 
83+CATLASS_DEVICE
84+void HcAicSet(int64_t base)
85+{
86+ AscendC::CrossCoreSetFlag<HC_SYNC_MODE, PIPE_FIX>(base);
87+}
88+ 
89+CATLASS_DEVICE
90+void HcAicWait(int64_t base)
91+{
92+ AscendC::CrossCoreWaitFlag<HC_SYNC_MODE, PIPE_MTE1>(base);
93+}
94+ 
95+// 仅 AIV 调用。硬同步占 flagId=14，禁止 AIC / SyncAll<false>。
96+CATLASS_DEVICE
97+void HcAivSyncAll()
98+{
99+ AscendC::PipeBarrier<PIPE_ALL>();
100+ AscendC::SyncAll<true>();
101+}
102+ 
103+// 每个 AIV 领一段连续行。rowBytes 不是 32 的倍数时，一段至少盖住一个 32B 块，避免两颗 AIV 改同一块 inject。
104+CATLASS_DEVICE
105+void HcAivContig(uint32_t actualM, uint32_t rowBytes, uint32_t& begin, uint32_t& end)
106+{
107+ uint32_t grain = 1;
108+ uint32_t rem = rowBytes & 31u;
109+ if (rem != 0u) {
110+ grain = (32u % rem) == 0u ? (32u / rem) : 32u;
111+ }
112+ uint32_t id = static_cast<uint32_t>(AscendC::GetBlockIdx());
113+ begin = id * grain;
114+ if (begin >= actualM) {
115+ begin = actualM;
116+ end = actualM;
117+ return;
118+ }
119+ end = begin + grain;
120+ if (end > actualM) {
121+ end = actualM;
122+ }
123+}
124+ 
125+// 一行字节数不是 32 的倍数时，相邻行共用一个 32B GM 块。A2 上两颗 AIV 同时读改写会丢更新。
126+CATLASS_DEVICE
127+void HcAivRowRange(uint32_t actualM, uint32_t rowBytes, uint32_t usedAic, uint32_t& begin, uint32_t& end)
128+{
129+ uint32_t total = usedAic * HcSubBlockNum();
130+ if (total == 0)
131+ total = 1;
132+ uint32_t grain = 1;
133+ while (((rowBytes * grain) & 31u) != 0u && grain < 32u)
134+ grain *= 2u;
135+ uint32_t groups = (actualM + grain - 1u) / grain;
136+ uint32_t id = static_cast<uint32_t>(AscendC::GetBlockIdx());
137+ begin = (groups * id / total) * grain;
138+ end = (groups * (id + 1u) / total) * grain;
139+ begin = begin < actualM ? begin : actualM;
140+ end = end < actualM ? end : actualM;
141+}
142+ 
143+template <class Element, uint32_t C1SwizzleDir, bool SmallN = false, class ArchTag_ = Arch::Ascend950,
144+ class Tiling = HcForwardMulSumTilingData>
145+struct HcForwardCommon {
146+ using ArchTag = ArchTag_;
147+ static constexpr bool kIsA2 = std::is_same_v<ArchTag, Arch::AtlasA2>;
148+ // C1 L0 32×256×32: L0B = 32*256*4*stages; stages=2 → 64KB. PreLoad unitFlag 时 L0C 必须保持 1.
149+ static constexpr uint32_t L0B_STAGES_C1 = 2;
150+ static constexpr bool ENABLE_UNIT_FLAG_C1 = true;
151+ // A2 small-N 的 M 往往小于 16。unit flag 按分形搬出时最后一行会丢，inject 对不齐。
152+ static constexpr bool kSmallNC0UnitFlag = !kIsA2;
153+ static constexpr uint32_t L0B_STAGES_C0 = 2u;
154+ static constexpr bool ENABLE_UNIT_FLAG_C0 = SmallN ? kSmallNC0UnitFlag : false;
155+ // A2 小路径 mTask=128 时 C0 有 4×6=24 块（32×64）：host 侧 usedAic 已计入 mTile，
156+ // 24 块摊到 20 个 AIC；gate-only 块排在队首由低号核算，down 峰值 2 块/核。
157+ static constexpr uint32_t L1M0 = SmallN ? ((kIsA2) ? 32u : HC_C0_L1M_SMALL) : 128;
158+ static constexpr uint32_t kA2SmallN0 = 64u;
159+ static constexpr uint32_t L1N0 = SmallN ? (kIsA2 ? kA2SmallN0 : HC_C0_L1N0_SMALL) : (kIsA2 ? 256u : HC_L1N0_LARGE);
160+ // MatmulV3 大 K：baseK=64，stepK=4，L1 上一次带 256 的 K，双缓冲即 depth=8。
161+ static constexpr uint32_t L1K0 = SmallN ? HC_L1K0_SMALL_F16 : (kIsA2 ? 256u : 128u);
162+ // C1 tile M：A2 bf16 下 K=320 全载 L1。L1M1=32 时每核每 wave 发 8 个小 m-tile 的
163+ // mmad+fixpipe（16KB/次），发射开销主导。升到 128：L1 占用 2*80K(A)+2*160K(B)=480K/512K
164+ // 仍可行，fixpipe 粒度 64KB、次数 /4。fp32 与非 A2 保持原值。
165+ static constexpr uint32_t L1M1 = (SmallN && kIsA2) ? 128u : (SmallN ? HC_L1M_SMALL : 128);
166+ static constexpr uint32_t L1N1 = HC_L1N1;
167+ // C1 的 K 是 R=320，A2 上整段放进 L1，避免每个 N 块把 A 再搬一遍。
168+ static constexpr uint32_t L1K1 = (kIsA2) ? 320u : HC_C1_L1K_F16;
169+ static constexpr uint32_t L0K0 = SmallN ? HC_L0K0_SMALL : ((kIsA2) ? 64u : 32u);
170+ static constexpr uint32_t L0K1 = (kIsA2) ? 64u : 32u;
171+ static_assert(!SmallN || kIsA2 || (L1M0 == HC_C0_L1M_SMALL && L1N0 == HC_C0_L1N0_SMALL), "C0 L1_M/L1_N knobs");
172+ static_assert(L1M0 != 0 && L1N0 != 0, "C0 tile nonempty");
173+ static_assert(L1K0 % L0K0 == 0, "C0 L1.K must be a multiple of L0.K");
174+ static_assert(L0K0 * L1N0 * static_cast<int64_t>(sizeof(Element)) * L0B_STAGES_C0 <= 64u * 1024u,
175+ "C0 L0B capacity");
176+ static_assert(L0K1 * L1N1 * static_cast<int64_t>(sizeof(Element)) * L0B_STAGES_C1 <= 64u * 1024u,
177+ "C1 L0B capacity");
178+ 
179+ using LayoutTagA = layout::RowMajor;
180+ using LayoutTagB = layout::ColumnMajor;
181+ using LayoutTagC = layout::RowMajor;
182+ 
183+ using DispatchPolicyC0 = Gemm::MmadPingpong<ArchTag, ENABLE_UNIT_FLAG_C0, false, 1, false, HC_A2_L1_A_STAGES,
184+ HC_A2_L1_B_STAGES, HC_A2_L0_A_STAGES, L0B_STAGES_C0>;
185+ // C1 的 K=R 已整段在 L1。同一行的 A 在 N 块之间驻留，只重搬 B。
186+ // 两级 L1A 留给两条 128 行。只有一条行带时第二片 N 会落到空槽，改用一级。
187+ using DispatchPolicyC1 = Gemm::MmadPingpong<ArchTag, ENABLE_UNIT_FLAG_C1, false, 1, true, HC_A2_L1_A_STAGES,
188+ HC_A2_L1_B_STAGES, HC_A2_L0_A_STAGES, L0B_STAGES_C1>;
189+ using DispatchPolicyC1A1 = Gemm::MmadPingpong<ArchTag, ENABLE_UNIT_FLAG_C1, false, 1, true, 1, HC_A2_L1_B_STAGES,
190+ HC_A2_L0_A_STAGES, L0B_STAGES_C1>;
191+ 
192+ using L1C0 = tla::Shape<tla::Int<L1M0>, tla::Int<L1N0>, tla::Int<L1K0>>;
193+ using L0C0 = tla::Shape<tla::Int<L1M0>, tla::Int<L1N0>, tla::Int<L0K0>>;
194+ using L1C1 = tla::Shape<tla::Int<L1M1>, tla::Int<L1N1>, tla::Int<L1K1>>;
195+ using L0C1 = tla::Shape<tla::Int<L1M1>, tla::Int<L1N1>, tla::Int<L0K1>>;
196+ 
197+ using TileCopyC0Packed =
198+ Gemm::Tile::PackedTileCopyTla<ArchTag, Element, LayoutTagA, Element, LayoutTagB, Element, LayoutTagC, void>;
199+ using TileCopyC0 = TileCopyC0Packed;
200+ using TileCopyC1 = TileCopyC0Packed;
201+ using BlockMmadC0 =
202+ Gemm::Block::BlockMmadTla<DispatchPolicyC0, L1C0, L0C0, Element, Element, Element, void, TileCopyC0>;
203+ using BlockMmadC1 =
204+ Gemm::Block::BlockMmadTla<DispatchPolicyC1, L1C1, L0C1, Element, Element, Element, void, TileCopyC1>;
205+ using BlockMmadC1A1 =
206+ Gemm::Block::BlockMmadTla<DispatchPolicyC1A1, L1C1, L0C1, Element, Element, Element, void, TileCopyC1>;
207+ using BlockSchedulerC1 = Gemm::Block::GemmIdentityBlockSwizzle<HC_A2_SWIZZLE_TILE_GROUP, C1SwizzleDir>;
208+ 
209+ struct Params {
210+ GM_ADDR ptrX;
211+ GM_ADDR ptrWDown;
212+ GM_ADDR ptrWUp;
213+ GM_ADDR ptrWGate;
214+ GM_ADDR ptrGamma;
215+ GM_ADDR ptrMixed;
216+ GM_ADDR ptrInject;
217+ GM_ADDR ptrInvRms;
218+ GM_ADDR ptrDownPreAct;
219+ GM_ADDR ptrUpPreAct;
220+ GM_ADDR userWs;
221+ GM_ADDR tiling;
222+ // 1：registry GET_TILING_DATA 已解包；0：直调 POD GM（.asc memcpy）。
223+ uint32_t tilingFromUnpacked;
224+ const Tiling* tilingUnpacked;
225+ };
226+ 
227+ template <class Mmad, class TensorA, class TensorB, class TensorC>
228+ CATLASS_DEVICE void RunMmad(Mmad& mmad, TensorA const& tensorA, TensorB const& tensorB, TensorC const& tensorC,
229+ uint32_t mOff, uint32_t nOffB, uint32_t nOffC, uint32_t actualM, uint32_t actualN,
230+ uint32_t kFull)
231+ {
232+ auto tileC = GetTile(tensorC, tla::MakeCoord(mOff, nOffC), tla::MakeShape(actualM, actualN));
233+ auto tileA = GetTile(tensorA, tla::MakeCoord(mOff, 0), tla::MakeShape(actualM, kFull));
234+ auto tileB = GetTile(tensorB, tla::MakeCoord(0, nOffB), tla::MakeShape(kFull, actualN));
235+ mmad(tileA, tileB, tileC, GemmCoord{actualM, actualN, kFull});
236+ }
237+ 
238+ // kTileM/kTileN 默认就是 A2 小路径的 32×64。超过 128 行改走 64×80：
239+ // down 的 N 能被 80 整除时 gate 单独占核。块数要两趟、改成 64×128 能收成一趟时才换。
240+ // 128 列时最后一块跨过 R，gate 并进这块 down，不再单独占一趟 K。fp32 的 L0B 放不下 128 列。
241+ CATLASS_DEVICE bool C0PreferN128(uint32_t actualM, uint32_t nConcat, uint32_t usedAic) const
242+ {
243+ if (actualM <= 128u || usedAic == 0u || nConcat == 0u) {
244+ return false;
245+ }
246+ uint32_t mTiles = (actualM + 63u) / 64u;
247+ uint32_t n80 = (nConcat + 79u) / 80u;
248+ uint32_t n128 = (nConcat + 127u) / 128u;
249+ uint32_t pass80 = (mTiles * n80 + usedAic - 1u) / usedAic;
250+ uint32_t pass128 = (mTiles * n128 + usedAic - 1u) / usedAic;
251+ return pass128 < pass80;
252+ }
253+ // 不超过 32 行时只有一条 M 块。32×64 对 R=320 只有 6 块，占不满 20 核。
254+ // 收到 32×32 是 10 块 down 加 1 块 gate，仍是一趟 K。再高的段收成 32 会变成两趟，不收。
255+ CATLASS_DEVICE bool C0PreferN32(uint32_t actualM, uint32_t nConcat, uint32_t usedAic) const
256+ {
257+ if (actualM == 0u || actualM > 32u || usedAic == 0u || nConcat == 0u) {
258+ return false;
259+ }
260+ uint32_t n64 = (nConcat + 63u) / 64u;
261+ uint32_t n32 = (nConcat + 31u) / 32u;
262+ if (n64 >= usedAic || n32 <= n64) {
263+ return false;
264+ }
265+ uint32_t pass32 = (n32 + usedAic - 1u) / usedAic;
266+ uint32_t pass64 = (n64 + usedAic - 1u) / usedAic;
267+ return pass32 == pass64;
268+ }
269+ CATLASS_DEVICE
270+ HcForwardCommon() {}
271+ 
272+ template <class Body>
273+ CATLASS_DEVICE void Execute(Params const& params)
274+ {
275+ using namespace AscendC;
276+ // 禁止 tdGm->field：3510 对 tiling GM 做结构体标量读会 264。
277+ // registry 入口必须 GET_TILING_DATA 解包后置 tilingFromUnpacked=1。
278+ if (params.tilingFromUnpacked == 0 || params.tilingUnpacked == nullptr) {
279+ return;
280+ }
281+ const Tiling* td = params.tilingUnpacked;
282+ uint32_t usedAic = td->usedAic;
283+ uint32_t mTask = td->mTask;
284+ uint32_t taskNum = td->taskNum;
285+ uint32_t d = td->D;
286+ uint32_t r = td->R;
287+ uint32_t hidden = td->H;
288+ uint32_t channels = td->C;
289+ uint32_t M = td->M;
290+ float eps = td->eps;
291+ bool useCombine = td->useCombine != 0;
292+ bool hasGamma = td->hasGamma != 0;
293+ uint32_t szT = td->szT;
294+ 
295+ uint32_t nConcat = td->nConcat;
296+ uint32_t nWs = td->nWs;
297+ if (nConcat == 0) {
298+ nConcat = useCombine ? (r + channels) : r;
299+ }
300+ if (nConcat > r && channels > 0) {
301+ useCombine = true;
302+ }
303+ if (nWs == 0) {
304+ nWs = ((nConcat + 15u) / 16u) * 16u;
305+ }
306+ 
307+ uint32_t blockIdx = 0;
308+ if ASCEND_IS_AIC {
309+ blockIdx = GetBlockIdx();
310+ } else {
311+ blockIdx = GetBlockIdx() / HcSubBlockNum();
312+ }
313+ if (blockIdx >= usedAic) {
314+ return;
315+ }
316+ 
317+ Arch::Resource<ArchTag> resource;
318+ 
319+ GM_ADDR ws = params.userWs;
320+ const int64_t xnSlotBytes = static_cast<int64_t>(mTask) * d * szT;
321+ const int64_t downSlotBytes = static_cast<int64_t>(mTask) * r * szT;
322+ // 正常大路径 xn 3 拍、down 2 拍，按核环形槽。小路径按全量 M。
323+ int64_t s0;
324+ int64_t s1;
325+ if constexpr (SmallN) {
326+ s0 = static_cast<int64_t>(M) * d * szT;
327+ s1 = static_cast<int64_t>(M) * r * szT;
328+ } else {
329+ s0 = static_cast<int64_t>(usedAic) * 3u * xnSlotBytes;
330+ s1 = static_cast<int64_t>(usedAic) * 2u * downSlotBytes;
331+ }
332+ int64_t s2 = 0;
333+ const uint32_t gateCols = (useCombine && channels > 0) ? channels : 0u;
334+ const uint32_t slotRows = mTask;
335+ const int64_t c0SlotBytes = static_cast<int64_t>(slotRows) * (static_cast<int64_t>(r) + gateCols) *
336+ static_cast<int64_t>(sizeof(Element));
337+ if constexpr (SmallN) {
338+ s2 = static_cast<int64_t>(2) * c0SlotBytes;
339+ } else {
340+ s2 = static_cast<int64_t>(usedAic) * c0SlotBytes * 2;
341+ }
342+ GM_ADDR ptrXn = ws;
343+ GM_ADDR ptrDown = ws + s0;
344+ GM_ADDR ptrC0Base = ws + s0 + s1;
345+ GM_ADDR ptrAfterC0 = ptrC0Base + s2;
346+ // C1 writes the backward cache directly; inference keeps private workspace.
347+ GM_ADDR ptrUp = params.ptrUpPreAct != nullptr ? params.ptrUpPreAct : ptrAfterC0;
348+ GM_ADDR ptrC0Core = SmallN ? ptrC0Base : (ptrC0Base + static_cast<int64_t>(blockIdx) * c0SlotBytes * 2);
349+ 
350+ GlobalTensor<Element> gmX;
351+ gmX.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(params.ptrX));
352+ GlobalTensor<Element> gmXn;
353+ gmXn.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrXn));
354+ GlobalTensor<Element> gmDown;
355+ gmDown.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrDown));
356+ GlobalTensor<Element> gmWDown;
357+ gmWDown.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(params.ptrWDown));
358+ GlobalTensor<Element> gmWUp;
359+ gmWUp.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(params.ptrWUp));
360+ GlobalTensor<Element> gmWGate;
361+ gmWGate.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(params.ptrWGate));
362+ GlobalTensor<Element> gmC0;
363+ gmC0.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrC0Core));
364+ GlobalTensor<Element> gmUp;
365+ gmUp.SetGlobalBuffer(reinterpret_cast<__gm__ Element*>(ptrUp));
366+ // up_pre_act 是反向专属保存：C1 写入后仅 V2 两 wave 后流式读一次，无前向复用，
367+ // 大 shape 不占 L2（与 x / down_pre_act / xn 同策略），把容量留给权重。
368+ // 仅大 shape 关闭：小 shape 的 up_pre_act 只有 KB~10MiB，对 L2 无压力，
369+ // 实测 M=32~512 全档噪声内（m128 +3.4% 偏劣），smallN 路径不设置。
370+ if (params.ptrUpPreAct != nullptr && td->smallN == 0u) {
371+ gmUp.SetL2CacheHint(CacheMode::CACHE_MODE_DISABLE);
372+ }
373+ 
374+ auto layoutXnFull = tla::MakeLayout<Element, LayoutTagA>(M, d);
375+ auto layoutDownFull = tla::MakeLayout<Element, LayoutTagA>(M, r);
376+ auto layoutWDown = tla::MakeLayout<Element, LayoutTagB>(d, r);
377+ auto layoutWUp = tla::MakeLayout<Element, LayoutTagB>(r, d);
378+ auto layoutWGate = tla::MakeLayout<Element, LayoutTagB>(d, channels);
379+ auto tensorX = tla::MakeTensor(gmX, layoutXnFull, Arch::PositionGM{});
380+ auto tensorXn = tla::MakeTensor(gmXn, layoutXnFull, Arch::PositionGM{});
381+ auto tensorDown = tla::MakeTensor(gmDown, layoutDownFull, Arch::PositionGM{});
382+ auto tensorWDown = tla::MakeTensor(gmWDown, layoutWDown, Arch::PositionGM{});
383+ auto tensorWGate = tla::MakeTensor(gmWGate, layoutWGate, Arch::PositionGM{});
384+ auto tensorWUp = tla::MakeTensor(gmWUp, layoutWUp, Arch::PositionGM{});
385+ auto layoutUpFull = tla::MakeLayout<Element, LayoutTagC>(M, d);
386+ auto tensorUp = tla::MakeTensor(gmUp, layoutUpFull, Arch::PositionGM{});
387+ 
388+ uint32_t localTasks = 0;
389+ if constexpr (SmallN) {
390+ localTasks = taskNum;
391+ } else if (usedAic > 0 && blockIdx < taskNum) {
392+ localTasks = (taskNum - blockIdx + usedAic - 1u) / usedAic;
393+ }
394+ 
395+ Epilogue::BlockEpilogueHcV0<Element> blockV0;
396+ Epilogue::BlockEpilogueHcV1<Element> blockV1;
397+ Epilogue::BlockEpilogueHcV2<Element> blockV2;
398+ blockV0.Init(params.ptrX, ptrXn, hasGamma ? params.ptrGamma : nullptr, params.ptrInvRms, !SmallN, !SmallN);
399+ blockV1.Init(ptrC0Core, ptrDown, params.ptrInject, params.ptrDownPreAct, !SmallN, channels, r);
400+ blockV2.Init(ptrUp, ptrXn, params.ptrMixed, params.ptrUpPreAct != nullptr && !SmallN, !SmallN);
401+ 
402+ const uint32_t nSliceC1 = (d + L1N1 - 1) / L1N1;
403+ 
404+ struct ExecutionContext {
405+ decltype(params)& f_params;
406+ decltype(usedAic)& f_usedAic;
407+ decltype(mTask)& f_mTask;
408+ decltype(d)& f_d;
409+ decltype(r)& f_r;
410+ decltype(hidden)& f_hidden;
411+ decltype(channels)& f_channels;
412+ decltype(M)& f_M;
413+ decltype(eps)& f_eps;
414+ decltype(useCombine)& f_useCombine;
415+ decltype(szT)& f_szT;
416+ decltype(nConcat)& f_nConcat;
417+ decltype(nWs)& f_nWs;
418+ decltype(blockIdx)& f_blockIdx;
419+ decltype(resource)& f_resource;
420+ decltype(xnSlotBytes)& f_xnSlotBytes;
421+ decltype(downSlotBytes)& f_downSlotBytes;
422+ decltype(c0SlotBytes)& f_c0SlotBytes;
423+ decltype(ptrXn)& f_ptrXn;
424+ decltype(ptrDown)& f_ptrDown;
425+ decltype(ptrC0Core)& f_ptrC0Core;
426+ decltype(tensorXn)& f_tensorXn;
427+ decltype(tensorDown)& f_tensorDown;
428+ decltype(tensorWDown)& f_tensorWDown;
429+ decltype(tensorWGate)& f_tensorWGate;
430+ decltype(tensorWUp)& f_tensorWUp;
431+ decltype(tensorUp)& f_tensorUp;
432+ decltype(localTasks)& f_localTasks;
433+ decltype(blockV0)& f_blockV0;
434+ decltype(blockV1)& f_blockV1;
435+ decltype(blockV2)& f_blockV2;
436+ decltype(nSliceC1)& f_nSliceC1;
437+ };
438+ ExecutionContext context{params, usedAic, mTask, d, r, hidden, channels,
439+ M, eps, useCombine, szT, nConcat, nWs, blockIdx,
440+ resource, xnSlotBytes, downSlotBytes, c0SlotBytes, ptrXn, ptrDown, ptrC0Core,
441+ tensorXn, tensorDown, tensorWDown, tensorWGate, tensorWUp, tensorUp, localTasks,
442+ blockV0, blockV1, blockV2, nSliceC1};
443+ Body body;
444+ body.RunBody(context);
445+ }
446+};
447+ 
448+} // namespace Catlass::Gemm::Kernel
Amhc/text_gated_residual/tests/ut/op_host/op_api/test_aclnn_text_gated_residual.cpp+40-0
Amhc/text_gated_residual/tests/ut/op_host/test_text_gated_residual_infershape.cpp+39-0
Amhc/text_gated_residual/tests/ut/op_host/test_text_gated_residual_tiling.cpp+41-0