MhcPre新增M-K模板&性能优化-ops-transformer-AtomGit

Mdocs/zh/menu_aclnn_api.md+1-0文件内容审核中，请稍后刷新重试

Mdocs/zh/op_api_list.md+2-1文件内容审核中，请稍后刷新重试

	@@ -111,6 +111,13 @@ $$
111	<td>FLOAT</td>	111	<td>FLOAT</td>
112	<td>-</td>	112	<td>-</td>
113	</tr>	113	</tr>
		114	+ <tr>
		115	+ <td>op_impl_mode</td>
		116	+ <td>可选输入</td>
		117	+ <td>指定MhcPre算子的计算模式，0表示Cube使用FP32模式计算，1表示Cube使用HF32模式计算，默认值为0。</td>
		118	+ <td>INT64</td>
		119	+ <td>-</td>
		120	+ </tr>
114	<tr>	121	<tr>
115	<td>h_in</td>	122	<td>h_in</td>
116	<td>输出</td>	123	<td>输出</td>
	@@ -166,3 +173,4 @@ $$
166	\| 调用方式 \| 调用样例 \| 说明 \|	173	\| 调用方式 \| 调用样例 \| 说明 \|
167	\|--------------\|-------------------------\|--------------------------------------------------------------\|	174	\|--------------\|-------------------------\|--------------------------------------------------------------\|
168	\| aclnn调用 \| [test_aclnn_mhc_pre](examples/test_aclnn_mhc_pre.cpp) \| 通过[aclnnMhcPre](docs/aclnnMhcPre.md)接口方式调用MhcPre算子。 \|	175	\| aclnn调用 \| [test_aclnn_mhc_pre](examples/test_aclnn_mhc_pre.cpp) \| 通过[aclnnMhcPre](docs/aclnnMhcPre.md)接口方式调用MhcPre算子。 \|
		176	+\| aclnn调用 \| [test_aclnn_mhc_pre_v2](examples/test_aclnn_mhc_pre_v2.cpp) \| 通过[aclnnMhcPreV2](docs/aclnnMhcPreV2.md)接口方式调用MhcPre算子，并通过opImplMode选择Cube的FP32或HF32计算模式。 \|

Amhc/mhc_pre/docs/aclnnMhcPreV2.md+692-0

@@ -0,0 +1,692 @@
	1	+# aclnnMhcPreV2
	C C chenjiao8月5日交付件不全 L liweijian16 8月5日评论： 8月5日评论：
	2	+
	3	+[📄 查看源码](https://gitcode.com/cann/ops-transformer/tree/master/mhc/mhc_pre)
	4	+
	5	+## 产品支持情况
	6	+
	7	+<!-- npu="950" id1 -->
	8	+- <term>Ascend 950PR/Ascend 950DT</term>：支持
	9	+<!-- end id1 -->
	10	+<!-- npu="A3" id2 -->
	11	+- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：不支持
	12	+<!-- end id2 -->
	13	+<!-- npu="910b" id3 -->
	14	+- <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：不支持
	15	+<!-- end id3 -->
	16	+<!-- npu="310b" id4 -->
	17	+- <term>Atlas 200I/500 A2 推理产品</term>：不支持
	18	+<!-- end id4 -->
	19	+<!-- npu="310p" id5 -->
	20	+- <term>Atlas 推理系列产品</term>：不支持
	21	+<!-- end id5 -->
	22	+<!-- npu="910" id6 -->
	23	+- <term>Atlas 训练系列产品</term>：不支持
	24	+<!-- end id6 -->
	25	+
	26	+## 功能说明
	27	+
	28	+- 接口功能：基于一系列计算得到MHC架构中hidden层的$H^{res}$和$H^{post}$投影矩阵以及Attention或MLP层的输入矩阵$h^{in}$。
	29	+
	30	+- 相较于[aclnnMhcPre](aclnnMhcPre.md)接口，此接口新增`opImplMode`参数，用于指定MhcPre算子的计算模式：
	31	+ - `opImplMode=0`：在Cube中使用FP32模式计算。
	32	+ - `opImplMode=1`：在Cube中使用HF32模式计算。
	33	+
	34	+- 计算公式：
	35	+
	36	+ 其中，xFlat表示将x的最后两维n和D视作长度为nD的向量，gammaFlat表示将gamma视作长度为nD的向量，@表示矩阵乘法，$\odot$表示逐元素乘法。
	37	+
	38	+ $$
	39	+ \begin{aligned}
	40	+ invRms &= \left(mean(xFlat^{2}) + normEps\right)^{-\frac{1}{2}}\\
	41	+ xGamma &= \begin{cases}
	42	+ xFlat \odot gammaFlat, & gamma \ne null \\
	43	+ xFlat, & gamma = null
	44	+ \end{cases}\\
	45	+ hMix &= xGamma @ phi^{T}\\
	46	+ w &= hMix \odot invRms\\
	47	+ [pPre, pPost, pRes] &= \begin{cases}
	48	+ split(w, [n, n, n^{2}]), & alpha.shape=[3] \\
	49	+ [split(w, [n, n]), 0], & alpha.shape=[2]
	50	+ \end{cases}\\
	51	+ hPre &= \sigma(pPre \odot alpha0 + bias0) + hcEps\\
	52	+ hPost &= 2\sigma(pPost \odot alpha1 + bias1)\\
	53	+ hRes &= \begin{cases}
	54	+ pRes \odot alpha2 + bias2, & alpha.shape=[3] \\
	55	+ 0, & alpha.shape=[2]
	56	+ \end{cases}\\
	57	+ hIn_{d} &= \sum_{i=0}^{n-1} hPre_{i} x_{i,d}
	58	+ \end{aligned}
	59	+ $$
	60	+
	61	+## 函数原型
	62	+
	63	+每个算子分为[两段式接口](../../../docs/zh/context/two_phase_api.md)，必须先调用"aclnnMhcPreV2GetWorkspaceSize"接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用"aclnnMhcPreV2"接口执行计算。
	64	+
	65	+```Cpp
	66	+aclnnStatus aclnnMhcPreV2GetWorkspaceSize(
	67	+ const aclTensor *x,
	68	+ const aclTensor *phi,
	69	+ const aclTensor *alpha,
	70	+ const aclTensor *bias,
	71	+ const aclTensor *gammaOptional,
	72	+ double normEps,
	73	+ double hcEps,
	74	+ int64_t opImplMode,
	75	+ aclTensor *hIn,
	76	+ aclTensor *hPost,
	77	+ aclTensor *hRes,
	78	+ aclTensor *invRmsOptional,
	79	+ aclTensor *hMixOptional,
	80	+ aclTensor *hPreOptional,
	81	+ uint64_t *workspaceSize,
	82	+ aclOpExecutor **executor)
	83	+```
	84	+
	85	+```Cpp
	86	+aclnnStatus aclnnMhcPreV2(
	87	+ void *workspace,
	88	+ uint64_t workspaceSize,
	89	+ aclOpExecutor *executor,
	90	+ aclrtStream stream)
	91	+```
	92	+
	93	+## aclnnMhcPreV2GetWorkspaceSize
	94	+
	95	+- 参数说明：
	96	+
	97	+ <table style="undefined;table-layout: fixed; width: 1550px"><colgroup>
	98	+ <col style="width: 187px">
	99	+ <col style="width: 121px">
	100	+ <col style="width: 287px">

Amhc/mhc_pre/examples/test_aclnn_mhc_pre_v2.cpp+321-0

@@ -0,0 +1,321 @@
	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	+/*!
	12	+ * \file test_aclnn_mhc_pre_v2.cpp
	13	+ * \brief
	14	+ */
	15	+
	16	+#include <iostream>
	17	+#include <vector>
	18	+#include "acl/acl.h"
	19	+#include "aclnnop/aclnn_mhc_pre_v2.h"
	20	+
	21	+#define CHECK_RET(cond, return_expr) \
	22	+ do { \
	23	+ if (!(cond)) { \
	24	+ return_expr; \
	25	+ } \
	26	+ } while (0)
	27	+
	28	+#define LOG_PRINT(message, ...) \
	29	+ do { \
	30	+ printf(message, ##__VA_ARGS__); \
	31	+ } while (0)
	32	+
	33	+// 计算Tensor形状对应的总元素数
	34	+int64_t GetShapeSize(const std::vector<int64_t> &shape)
	35	+{
	36	+ int64_t size = 1;
	37	+ for (int64_t dim : shape) {
	38	+ size *= dim;
	39	+ }
	40	+ return size;
	41	+}
	42	+
	43	+// 将Device侧Tensor数据拷贝到Host侧并打印（float类型）
	44	+void PrintTensorDataFloat(const std::vector<int64_t> &shape, void *device_addr)
	45	+{
	46	+ int64_t size = GetShapeSize(shape);
	47	+ std::vector<float> host_data(size, 0.0f);
	48	+
	49	+ aclError ret = aclrtMemcpy(host_data.data(), size * sizeof(float), device_addr, size * sizeof(float),
	50	+ ACL_MEMCPY_DEVICE_TO_HOST);
	51	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Memcpy device to host failed, error: %d\n", ret); return);
	52	+
	53	+ LOG_PRINT("Tensor data (first 10 elements): ");
	54	+ for (int64_t i = 0; i < std::min((int64_t)10, size); ++i) {
	55	+ LOG_PRINT("%f ", host_data[i]);
	56	+ }
	57	+ LOG_PRINT("\n");
	58	+}
	59	+
	60	+// 将Device侧Tensor数据拷贝到Host侧并打印（float16类型）
	61	+void PrintTensorDataFloat16(const std::vector<int64_t> &shape, void *device_addr)
	62	+{
	63	+ int64_t size = GetShapeSize(shape);
	64	+ std::vector<aclFloat16> host_fp16(size);
	65	+ std::vector<float> host_data(size, 0.0f);
	66	+
	67	+ aclError ret = aclrtMemcpy(host_fp16.data(), size * sizeof(aclFloat16), device_addr, size * sizeof(aclFloat16),
	68	+ ACL_MEMCPY_DEVICE_TO_HOST);
	69	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Memcpy device to host failed, error: %d\n", ret); return);
	70	+
	71	+ for (int64_t i = 0; i < size; ++i) {
	72	+ host_data[i] = aclFloat16ToFloat(host_fp16[i]);
	73	+ }
	74	+
	75	+ LOG_PRINT("Tensor data (first 10 elements): ");
	76	+ for (int64_t i = 0; i < std::min((int64_t)10, size); ++i) {
	77	+ LOG_PRINT("%f ", host_data[i]);
	78	+ }
	79	+ LOG_PRINT("\n");
	80	+}
	81	+
	82	+// 初始化AscendCL环境（Device/Context/Stream）
	83	+int InitAcl(int32_t device_id, aclrtContext &context, aclrtStream &stream)
	84	+{
	85	+ aclError ret = aclInit(nullptr);
	86	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed, error: %d\n", ret); return -1);
	87	+
	88	+ ret = aclrtSetDevice(device_id);
	89	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed, error: %d\n", ret); return -1);
	90	+
	91	+ ret = aclrtCreateContext(&context, device_id);
	92	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateContext failed, error: %d\n", ret); return -1);
	93	+
	94	+ ret = aclrtSetCurrentContext(context);
	95	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetCurrentContext failed, error: %d\n", ret); return -1);
	96	+
	97	+ ret = aclrtCreateStream(&stream);
	98	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed, error: %d\n", ret); return -1);
	99	+
	100	+ return 0;
	101	+}
	102	+
	103	+// 创建FLOAT32类型Device侧aclTensor（含数据拷贝）
	104	+int CreateAclTensorFloat32(const std::vector<float> &host_data, const std::vector<int64_t> &shape, void *&device_addr,
	105	+ aclTensor *&tensor)
	106	+{
	107	+ int64_t size = GetShapeSize(shape) * sizeof(float);
	108	+
	109	+ aclError ret = aclrtMalloc(&device_addr, size, ACL_MEM_MALLOC_HUGE_FIRST);
	110	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed, error: %d\n", ret); return -1);
	111	+
	112	+ ret = aclrtMemcpy(device_addr, size, host_data.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
	113	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed, error: %d\n", ret); return -1);
	114	+
	115	+ std::vector<int64_t> strides(shape.size(), 1);
	116	+ for (int64_t i = shape.size() - 2; i >= 0; --i) {
	117	+ strides[i] = strides[i + 1] * shape[i + 1];
	118	+ }
	119	+
	120	+ tensor = aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT, strides.data(), 0, ACL_FORMAT_ND, shape.data(),
	121	+ shape.size(), device_addr);
	122	+ CHECK_RET(tensor != nullptr, LOG_PRINT("aclCreateTensor failed\n"); return -1);
	123	+
	124	+ return 0;
	125	+}
	126	+
	127	+// 创建FLOAT16类型Device侧aclTensor（含数据拷贝）
	128	+int CreateAclTensorFloat16(const std::vector<float> &host_data, const std::vector<int64_t> &shape, void *&device_addr,
	129	+ aclTensor *&tensor)
	130	+{
	131	+ int64_t size = GetShapeSize(shape);
	132	+ std::vector<aclFloat16> host_data_fp16(size);
	133	+ for (int64_t i = 0; i < size; ++i) {
	134	+ host_data_fp16[i] = aclFloat16(host_data[i]);
	135	+ }
	136	+
	137	+ int64_t byte_size = size * sizeof(aclFloat16);
	138	+
	139	+ aclError ret = aclrtMalloc(&device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST);
	140	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed, error: %d\n", ret); return -1);
	141	+
	142	+ ret = aclrtMemcpy(device_addr, byte_size, host_data_fp16.data(), byte_size, ACL_MEMCPY_HOST_TO_DEVICE);
	143	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed, error: %d\n", ret); return -1);
	144	+
	145	+ std::vector<int64_t> strides(shape.size(), 1);
	146	+ for (int64_t i = shape.size() - 2; i >= 0; --i) {
	147	+ strides[i] = strides[i + 1] * shape[i + 1];
	148	+ }
	149	+
	150	+ tensor = aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT16, strides.data(), 0, ACL_FORMAT_ND, shape.data(),
	151	+ shape.size(), device_addr);
	152	+ CHECK_RET(tensor != nullptr, LOG_PRINT("aclCreateTensor failed\n"); return -1);
	153	+
	154	+ return 0;
	155	+}
	156	+
	157	+// 创建FLOAT16类型输出aclTensor（仅申请内存）
	158	+int CreateAclTensorFloat16Output(const std::vector<int64_t> &shape, void &device_addr, aclTensor &tensor)
	159	+{
	160	+ int64_t size = GetShapeSize(shape);
	161	+ int64_t byte_size = size * sizeof(aclFloat16);
	162	+
	163	+ aclError ret = aclrtMalloc(&device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST);
	164	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed, error: %d\n", ret); return -1);
	165	+
	166	+ tensor = aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT16, nullptr, 0, ACL_FORMAT_ND, shape.data(),
	167	+ shape.size(), device_addr);
	168	+ CHECK_RET(tensor != nullptr, LOG_PRINT("aclCreateTensor failed\n"); return -1);
	169	+
	170	+ return 0;
	171	+}
	172	+
	173	+// 创建FLOAT32类型输出aclTensor（仅申请内存）
	174	+int CreateAclTensorFloat32Output(const std::vector<int64_t> &shape, void &device_addr, aclTensor &tensor)
	175	+{
	176	+ int64_t size = GetShapeSize(shape);
	177	+ int64_t byte_size = size * sizeof(float);
	178	+
	179	+ aclError ret = aclrtMalloc(&device_addr, byte_size, ACL_MEM_MALLOC_HUGE_FIRST);
	180	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed, error: %d\n", ret); return -1);
	181	+
	182	+ tensor = aclCreateTensor(shape.data(), shape.size(), ACL_FLOAT, nullptr, 0, ACL_FORMAT_ND, shape.data(),
	183	+ shape.size(), device_addr);
	184	+ CHECK_RET(tensor != nullptr, LOG_PRINT("aclCreateTensor failed\n"); return -1);
	185	+
	186	+ return 0;
	187	+}
	188	+
	189	+struct Tensors {
	190	+ void x_addr = nullptr, phi_addr = nullptr, alpha_addr = nullptr, bias_addr = nullptr, *gamma_addr = nullptr;
	191	+ void hin_addr = nullptr, h_post_addr = nullptr, h_res_addr = nullptr, inv_rms_addr = nullptr;
	192	+ void h_mix_addr = nullptr, h_pre_addr = nullptr;
	193	+ aclTensor x = nullptr, phi = nullptr, alpha = nullptr, bias = nullptr, *gamma = nullptr;
	194	+ aclTensor hin = nullptr, h_post = nullptr, h_res = nullptr, inv_rms = nullptr;
	195	+ aclTensor h_mix = nullptr, h_pre = nullptr;
	196	+};
	197	+
	198	+int CreateInputTensors(const std::vector<int64_t> &x_shape, const std::vector<int64_t> &phi_shape,
	199	+ const std::vector<int64_t> &alpha_shape, const std::vector<int64_t> &bias_shape,
	200	+ const std::vector<int64_t> &gamma_shape, Tensors &tensors)
	201	+{
	202	+ std::vector<float> x_host_data(GetShapeSize(x_shape), 1.0f);
	203	+ std::vector<float> phi_host_data(GetShapeSize(phi_shape), 1.0f);
	204	+ std::vector<float> alpha_host_data(3, 1.0f);
	205	+ std::vector<float> bias_host_data(GetShapeSize(bias_shape), 1.0f);
	206	+ std::vector<float> gamma_host_data(GetShapeSize(gamma_shape), 1.0f);
	207	+
	208	+ int ret = CreateAclTensorFloat16(x_host_data, x_shape, tensors.x_addr, tensors.x);
	209	+ CHECK_RET(ret == 0, LOG_PRINT("Create x_tensor failed\n"); return -1);
	210	+ ret = CreateAclTensorFloat32(phi_host_data, phi_shape, tensors.phi_addr, tensors.phi);
	211	+ CHECK_RET(ret == 0, LOG_PRINT("Create phi_tensor failed\n"); return -1);
	212	+ ret = CreateAclTensorFloat32(alpha_host_data, alpha_shape, tensors.alpha_addr, tensors.alpha);
	213	+ CHECK_RET(ret == 0, LOG_PRINT("Create alpha_tensor failed\n"); return -1);
	214	+ ret = CreateAclTensorFloat32(bias_host_data, bias_shape, tensors.bias_addr, tensors.bias);
	215	+ CHECK_RET(ret == 0, LOG_PRINT("Create bias_tensor failed\n"); return -1);
	216	+ ret = CreateAclTensorFloat32(gamma_host_data, gamma_shape, tensors.gamma_addr, tensors.gamma);
	217	+ CHECK_RET(ret == 0, LOG_PRINT("Create gamma_tensor failed\n"); return -1);
	218	+ return 0;
	219	+}
	220	+
	221	+int CreateOutputTensors(const std::vector<int64_t> &hin_shape, const std::vector<int64_t> &h_post_shape,
	222	+ const std::vector<int64_t> &h_res_shape, const std::vector<int64_t> &inv_rms_shape,
	223	+ const std::vector<int64_t> &h_mix_shape, const std::vector<int64_t> &h_pre_shape,
	224	+ Tensors &tensors)
	225	+{
	226	+ int ret = CreateAclTensorFloat16Output(hin_shape, tensors.hin_addr, tensors.hin);
	227	+ CHECK_RET(ret == 0, LOG_PRINT("Create hin_tensor failed\n"); return -1);
	228	+ ret = CreateAclTensorFloat32Output(h_post_shape, tensors.h_post_addr, tensors.h_post);
	229	+ CHECK_RET(ret == 0, LOG_PRINT("Create h_post_tensor failed\n"); return -1);
	230	+ ret = CreateAclTensorFloat32Output(h_res_shape, tensors.h_res_addr, tensors.h_res);
	231	+ CHECK_RET(ret == 0, LOG_PRINT("Create h_res_tensor failed\n"); return -1);
	232	+ ret = CreateAclTensorFloat32Output(inv_rms_shape, tensors.inv_rms_addr, tensors.inv_rms);
	233	+ CHECK_RET(ret == 0, LOG_PRINT("Create inv_rms_tensor failed\n"); return -1);
	234	+ ret = CreateAclTensorFloat32Output(h_mix_shape, tensors.h_mix_addr, tensors.h_mix);
	235	+ CHECK_RET(ret == 0, LOG_PRINT("Create h_mix_tensor failed\n"); return -1);
	236	+ ret = CreateAclTensorFloat32Output(h_pre_shape, tensors.h_pre_addr, tensors.h_pre);
	237	+ CHECK_RET(ret == 0, LOG_PRINT("Create h_pre_tensor failed\n"); return -1);
	238	+ return 0;
	239	+}
	240	+
	241	+void DestroyTensors(Tensors &tensors)
	242	+{
	243	+ aclDestroyTensor(tensors.x);
	244	+ aclDestroyTensor(tensors.phi);
	245	+ aclDestroyTensor(tensors.alpha);
	246	+ aclDestroyTensor(tensors.bias);
	247	+ aclDestroyTensor(tensors.gamma);
	248	+ aclDestroyTensor(tensors.hin);
	249	+ aclDestroyTensor(tensors.h_post);
	250	+ aclDestroyTensor(tensors.h_res);
	251	+ aclDestroyTensor(tensors.inv_rms);
	252	+ aclDestroyTensor(tensors.h_mix);
	253	+ aclDestroyTensor(tensors.h_pre);
	254	+}
	255	+
	256	+void FreeDeviceMemory(Tensors &tensors)
	257	+{
	258	+ aclrtFree(tensors.x_addr);
	259	+ aclrtFree(tensors.phi_addr);
	260	+ aclrtFree(tensors.alpha_addr);
	261	+ aclrtFree(tensors.bias_addr);
	262	+ aclrtFree(tensors.gamma_addr);
	263	+ aclrtFree(tensors.hin_addr);
	264	+ aclrtFree(tensors.h_post_addr);
	265	+ aclrtFree(tensors.h_res_addr);
	266	+ aclrtFree(tensors.inv_rms_addr);
	267	+ aclrtFree(tensors.h_mix_addr);
	268	+ aclrtFree(tensors.h_pre_addr);
	269	+}
	270	+
	271	+int main()
	272	+{
	273	+ int32_t device_id = 0;
	274	+ aclrtContext context = nullptr;
	275	+ aclrtStream stream = nullptr;
	276	+ Tensors tensors;
	277	+ int B = 1, S = 2048, n = 4, D = 2560;
	278	+ std::vector<int64_t> x_shape = {B * S, n, D}, phi_shape = {n * n + 2 * n, n * D}, alpha_shape = {3},
	279	+ bias_shape = {n * n + 2 * n}, gamma_shape = {n, D};
	280	+ std::vector<int64_t> hin_shape = {B * S, D}, h_post_shape = {B * S, n}, h_res_shape = {B * S, n, n},
	281	+ inv_rms_shape = {B * S}, h_mix_shape = {B * S, n * n + 2 * n}, h_pre_shape = {B * S, n};
	282	+ int ret = InitAcl(device_id, context, stream);
	283	+ CHECK_RET(ret == 0, LOG_PRINT("InitAcl failed, error: %d\n", ret); return -1);
	284	+ ret = CreateInputTensors(x_shape, phi_shape, alpha_shape, bias_shape, gamma_shape, tensors);
	285	+ CHECK_RET(ret == 0, return -1);
	286	+ ret = CreateOutputTensors(hin_shape, h_post_shape, h_res_shape, inv_rms_shape, h_mix_shape, h_pre_shape, tensors);
	287	+ CHECK_RET(ret == 0, return -1);
	288	+ uint64_t workspace_size = 0;
	289	+ aclOpExecutor *executor = nullptr;
	290	+ int64_t opImplMode = 1;
	291	+ aclnnStatus aclnn_ret = aclnnMhcPreV2GetWorkspaceSize(
	292	+ tensors.x, tensors.phi, tensors.alpha, tensors.bias, tensors.gamma, 1e-6, 1e-6, opImplMode, tensors.hin,
	293	+ tensors.h_post, tensors.h_res, tensors.inv_rms, tensors.h_mix, tensors.h_pre, &workspace_size, &executor);
	294	+ CHECK_RET(aclnn_ret == ACL_SUCCESS, LOG_PRINT("aclnnMhcPreV2GetWorkspaceSize failed, error: %d\n", aclnn_ret);
	295	+ return -1);
	296	+ void *workspace_addr = nullptr;
	297	+ if (workspace_size > 0) {
	298	+ ret = aclrtMalloc(&workspace_addr, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST);
	299	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc workspace failed, error: %d\n", ret); return -1);
	300	+ }
	301	+ aclnn_ret = aclnnMhcPreV2(workspace_addr, workspace_size, executor, stream);
	302	+ CHECK_RET(aclnn_ret == ACL_SUCCESS, LOG_PRINT("aclnnMhcPreV2 failed, error: %d\n", aclnn_ret); return -1);
	303	+ CHECK_RET(aclrtSynchronizeStream(stream) == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed\n"); return -1);
	304	+ LOG_PRINT("MhcPre compute success!\nOutput tensor data: \n");
	305	+ PrintTensorDataFloat16(hin_shape, tensors.hin_addr);
	306	+ PrintTensorDataFloat(h_post_shape, tensors.h_post_addr);
	307	+ PrintTensorDataFloat(h_res_shape, tensors.h_res_addr);
	308	+ PrintTensorDataFloat(inv_rms_shape, tensors.inv_rms_addr);
	309	+ PrintTensorDataFloat(h_mix_shape, tensors.h_mix_addr);
	310	+ PrintTensorDataFloat(h_pre_shape, tensors.h_pre_addr);
	311	+ DestroyTensors(tensors);
	312	+ FreeDeviceMemory(tensors);
	313	+ if (workspace_size > 0)
	314	+ aclrtFree(workspace_addr);
	315	+ aclrtDestroyStream(stream);
	316	+ aclrtDestroyContext(context);
	317	+ aclrtResetDevice(device_id);
	318	+ aclFinalize();
	319	+ LOG_PRINT("All resources released successfully!\n");
	320	+ return 0;
	321	+}

Mmhc/mhc_pre/op_host/mhc_pre_def.cpp+1-0

	@@ -90,6 +90,7 @@ public:
90	this->Attr("out_flag").AttrType(OPTIONAL).Int(0);	90	this->Attr("out_flag").AttrType(OPTIONAL).Int(0);
91	this->Attr("norm_eps").AttrType(OPTIONAL).Float(1e-6f);	91	this->Attr("norm_eps").AttrType(OPTIONAL).Float(1e-6f);
92	this->Attr("hc_eps").AttrType(OPTIONAL).Float(1e-6f);	92	this->Attr("hc_eps").AttrType(OPTIONAL).Float(1e-6f);
		93	+ this->Attr("op_impl_mode").AttrType(OPTIONAL).Int(0);
93	}	94	}
94	};	95	};
95		96

Mmhc/mhc_pre/op_host/mhc_pre_infershape.cpp+10-7

	@@ -29,17 +29,20 @@ using namespace ge;
29		29
30	namespace ops {	30	namespace ops {
31		31
		32	+// Input tensor indices.
32	const constexpr int64_t X_INDEX = 0;	33	const constexpr int64_t X_INDEX = 0;
33	const constexpr int64_t PHI_INDEX = 1;	34	const constexpr int64_t PHI_INDEX = 1;
34	const constexpr int64_t ALPHA_INDEX = 2;	35	const constexpr int64_t ALPHA_INDEX = 2;
35		36
		37	+// Output tensor indices.
36	const constexpr int64_t OUT_H_IN_INDEX = 0;	38	const constexpr int64_t OUT_H_IN_INDEX = 0;
37	const constexpr int64_t OUT_H_POST_INDEX = 1;	39	const constexpr int64_t OUT_H_POST_INDEX = 1;
38	const constexpr int64_t OUT_H_RES_INDEX = 2;	40	const constexpr int64_t OUT_H_RES_INDEX = 2;
39	const constexpr int64_t OUT_INV_RMS_INDEX = 3;	41	const constexpr int64_t OUT_INV_RMS_INDEX = 3;
40	-const constexpr int64_t OUT_MM_RES_INDEX = 4;	42	+const constexpr int64_t OUT_H_MIX_INDEX = 4;
41	const constexpr int64_t OUT_H_PRE_INDEX = 5;	43	const constexpr int64_t OUT_H_PRE_INDEX = 5;
42		44
		45	+// Supported input ranks and dynamic-shape marker.
43	const constexpr int64_t BSND_DIM_NUM = 4;	46	const constexpr int64_t BSND_DIM_NUM = 4;
44	const constexpr int64_t TND_DIM_NUM = 3;	47	const constexpr int64_t TND_DIM_NUM = 3;
45	const constexpr int64_t UNKNOWN_DIM_VALUE = -1LL;	48	const constexpr int64_t UNKNOWN_DIM_VALUE = -1LL;
	@@ -117,9 +120,9 @@ static ge::graphStatus InferShape4MhcPre(InferShapeContext *context)
117	// has residual	120	// has residual
118	bool hasResi = (alphaDimSize == 3);	121	bool hasResi = (alphaDimSize == 3);
119		122
120	- gert::Shape *outShapes[6] = {context->GetOutputShape(OUT_H_IN_INDEX), context->GetOutputShape(OUT_H_POST_INDEX),	123	+ gert::Shape *outShapes[6] = {context->GetOutputShape(OUT_H_IN_INDEX), context->GetOutputShape(OUT_H_POST_INDEX),
121	- context->GetOutputShape(OUT_H_RES_INDEX), context->GetOutputShape(OUT_INV_RMS_INDEX),	124	+ context->GetOutputShape(OUT_H_RES_INDEX), context->GetOutputShape(OUT_INV_RMS_INDEX),
122	- context->GetOutputShape(OUT_MM_RES_INDEX), context->GetOutputShape(OUT_H_PRE_INDEX)};	125	+ context->GetOutputShape(OUT_H_MIX_INDEX), context->GetOutputShape(OUT_H_PRE_INDEX)};
123	if (xDim == BSND_DIM_NUM) {	126	if (xDim == BSND_DIM_NUM) {
124	// b,s,n,d	127	// b,s,n,d
125	uint64_t b = xShape->GetDim(0), s = xShape->GetDim(1), n = xShape->GetDim(2), d = xShape->GetDim(3);	128	uint64_t b = xShape->GetDim(0), s = xShape->GetDim(1), n = xShape->GetDim(2), d = xShape->GetDim(3);
	@@ -133,7 +136,7 @@ static ge::graphStatus InferShape4MhcPre(InferShapeContext *context)
133	SetShape1D(outShapes[OUT_H_RES_INDEX], 0);	136	SetShape1D(outShapes[OUT_H_RES_INDEX], 0);
134	}	137	}
135	SetShape2D(outShapes[OUT_INV_RMS_INDEX], b, s);	138	SetShape2D(outShapes[OUT_INV_RMS_INDEX], b, s);
136	- SetShape3D(outShapes[OUT_MM_RES_INDEX], b, s, matN);	139	+ SetShape3D(outShapes[OUT_H_MIX_INDEX], b, s, matN);
137	SetShape3D(outShapes[OUT_H_PRE_INDEX], b, s, n);	140	SetShape3D(outShapes[OUT_H_PRE_INDEX], b, s, n);
138	} else {	141	} else {
139	// t,n,d	142	// t,n,d
	@@ -148,7 +151,7 @@ static ge::graphStatus InferShape4MhcPre(InferShapeContext *context)
148	SetShape1D(outShapes[OUT_H_RES_INDEX], 0);	151	SetShape1D(outShapes[OUT_H_RES_INDEX], 0);
149	}	152	}
150	SetShape1D(outShapes[OUT_INV_RMS_INDEX], t);	153	SetShape1D(outShapes[OUT_INV_RMS_INDEX], t);
151	- SetShape2D(outShapes[OUT_MM_RES_INDEX], t, matN);	154	+ SetShape2D(outShapes[OUT_H_MIX_INDEX], t, matN);
152	SetShape2D(outShapes[OUT_H_PRE_INDEX], t, n);	155	SetShape2D(outShapes[OUT_H_PRE_INDEX], t, n);
153	}	156	}
154		157
	@@ -163,7 +166,7 @@ static graphStatus InferDataType4MhcPre(gert::InferDataTypeContext *context)
163	context->SetOutputDataType(OUT_H_POST_INDEX, DataType::DT_FLOAT);	166	context->SetOutputDataType(OUT_H_POST_INDEX, DataType::DT_FLOAT);
164	context->SetOutputDataType(OUT_H_RES_INDEX, DataType::DT_FLOAT);	167	context->SetOutputDataType(OUT_H_RES_INDEX, DataType::DT_FLOAT);
165	context->SetOutputDataType(OUT_INV_RMS_INDEX, DataType::DT_FLOAT);	168	context->SetOutputDataType(OUT_INV_RMS_INDEX, DataType::DT_FLOAT);
166	- context->SetOutputDataType(OUT_MM_RES_INDEX, DataType::DT_FLOAT);	169	+ context->SetOutputDataType(OUT_H_MIX_INDEX, DataType::DT_FLOAT);
167	context->SetOutputDataType(OUT_H_PRE_INDEX, DataType::DT_FLOAT);	170	context->SetOutputDataType(OUT_H_PRE_INDEX, DataType::DT_FLOAT);
168	return GRAPH_SUCCESS;	171	return GRAPH_SUCCESS;
169	}	172	}

Mmhc/mhc_pre/op_host/op_api/aclnn_mhc_pre.cpp+59-14

	@@ -9,6 +9,7 @@
9	*/	9	*/
10		10
11	#include "aclnn_mhc_pre.h"	11	#include "aclnn_mhc_pre.h"
		12	+#include "aclnn_mhc_pre_v2.h"
12	#include <dlfcn.h>	13	#include <dlfcn.h>
13	#include <new>	14	#include <new>
14	#include <memory>	15	#include <memory>
	@@ -54,12 +55,13 @@ constexpr int64_t D_ALIGNMENT = 16;
54	constexpr int64_t ALPHA_DIM_SIZE_3 = 3;	55	constexpr int64_t ALPHA_DIM_SIZE_3 = 3;
55	constexpr int64_t ALPHA_DIM_SIZE_2 = 2;	56	constexpr int64_t ALPHA_DIM_SIZE_2 = 2;
56	constexpr int64_t PHI_DIM_OFFSET = 2;	57	constexpr int64_t PHI_DIM_OFFSET = 2;
		58	+constexpr int64_t MHC_PRE_USE_FP32 = 0;
		59	+constexpr int64_t MHC_PRE_USE_HF32 = 1;
57		60
58	static bool CheckAlphaShape(const aclTensor *alphaTensor);	61	static bool CheckAlphaShape(const aclTensor *alphaTensor);
59	static bool ValidateNDParams(int64_t n, int64_t d);	62	static bool ValidateNDParams(int64_t n, int64_t d);
60	static bool CheckTensorShape(const aclTensor tensor, std::initializer_list<int64_t> expectedShape, const char name);	63	static bool CheckTensorShape(const aclTensor tensor, std::initializer_list<int64_t> expectedShape, const char name);
61		64
62	-
63	static std::string TensorShapeToString(const aclTensor *tensor)	65	static std::string TensorShapeToString(const aclTensor *tensor)
64	{	66	{
65	const auto shape = tensor->GetViewShape();	67	const auto shape = tensor->GetViewShape();
	@@ -129,6 +131,7 @@ struct MhcParamsBase {
129	const aclTensor *gammaOptional = nullptr;	131	const aclTensor *gammaOptional = nullptr;
130	double normEps;	132	double normEps;
131	double hcEps;	133	double hcEps;
		134	+ int64_t opImplMode = MHC_PRE_USE_FP32;
132	aclTensor *hIn = nullptr;	135	aclTensor *hIn = nullptr;
133	aclTensor *hPost = nullptr;	136	aclTensor *hPost = nullptr;
134	aclTensor *hRes = nullptr;	137	aclTensor *hRes = nullptr;
	@@ -165,10 +168,11 @@ public:
165	return *this;	168	return *this;
166	}	169	}
167		170
168	- MhcBuilder &SetAttr(double normEps, double hcEps)	171	+ MhcBuilder &SetAttr(double normEps, double hcEps, int64_t opImplMode)
169	{	172	{
170	obj_.normEps = normEps;	173	obj_.normEps = normEps;
171	obj_.hcEps = hcEps;	174	obj_.hcEps = hcEps;
		175	+ obj_.opImplMode = opImplMode;
172	return *this;	176	return *this;
173	}	177	}
174		178
	@@ -257,7 +261,8 @@ static bool CheckOptionalOutputGroup(const MhcParamsBase &params)
257	bool hasAllOptional =	261	bool hasAllOptional =
258	params.invRmsOptional != nullptr && params.hMixOptional != nullptr && params.hPreOptional != nullptr;	262	params.invRmsOptional != nullptr && params.hMixOptional != nullptr && params.hPreOptional != nullptr;
259	if (hasAnyOptional && !hasAllOptional) {	263	if (hasAnyOptional && !hasAllOptional) {
260	- const std::string pointerStates = std::string("invRmsOptional=") +	264	+ const std::string pointerStates =
		265	+ std::string("invRmsOptional=") +
261	(params.invRmsOptional == nullptr ? "nullptr" : "non-null") + ", hMixOptional=" +	266	(params.invRmsOptional == nullptr ? "nullptr" : "non-null") + ", hMixOptional=" +
262	(params.hMixOptional == nullptr ? "nullptr" : "non-null") + ", hPreOptional=" +	267	(params.hMixOptional == nullptr ? "nullptr" : "non-null") + ", hPreOptional=" +
263	(params.hPreOptional == nullptr ? "nullptr" : "non-null");	268	(params.hPreOptional == nullptr ? "nullptr" : "non-null");
	@@ -269,7 +274,6 @@ static bool CheckOptionalOutputGroup(const MhcParamsBase &params)
269	return true;	274	return true;
270	}	275	}
271		276
272	-
273	static bool CheckOptionalOutputDims(const MhcParamsBase &params)	277	static bool CheckOptionalOutputDims(const MhcParamsBase &params)
274	{	278	{
275	if (params.invRmsOptional == nullptr) {	279	if (params.invRmsOptional == nullptr) {
	@@ -429,7 +433,6 @@ static bool ValidateNDParams(int64_t n, int64_t d)
429	return true;	433	return true;
430	}	434	}
431		435
432	-
433	static bool CheckInputDtype(const MhcParamsBase &params)	436	static bool CheckInputDtype(const MhcParamsBase &params)
434	{	437	{
435	const std::initializer_list<DataType> X_SUPPORT_DTYPE_LIST = {DataType::DT_BF16, DataType::DT_FLOAT16};	438	const std::initializer_list<DataType> X_SUPPORT_DTYPE_LIST = {DataType::DT_BF16, DataType::DT_FLOAT16};
	@@ -526,6 +529,18 @@ static bool CheckFormat(const MhcParamsBase &params)
526	{	529	{
527	return CheckInputFormat(params) && CheckOutputFormat(params) && CheckOptionalOutputFormat(params);	530	return CheckInputFormat(params) && CheckOutputFormat(params) && CheckOptionalOutputFormat(params);
528	}	531	}
		532	+
		533	+static bool CheckOpImplMode(const MhcParamsBase &params)
		534	+{
		535	+ if (params.opImplMode != MHC_PRE_USE_FP32 && params.opImplMode != MHC_PRE_USE_HF32) {
		536	+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(ACLNN_OP_NAME, "opImplMode",
		537	+ std::to_string(params.opImplMode).c_str(),
		538	+ "must be 0 (MHC_PRE_USE_FP32) or 1 (MHC_PRE_USE_HF32)");
		539	+ return false;
		540	+ }
		541	+ return true;
		542	+}
		543	+
529	static aclnnStatus CheckParams(const MhcParamsBase &params)	544	static aclnnStatus CheckParams(const MhcParamsBase &params)
530	{	545	{
531	CHECK_RET(CheckNotNull(params), ACLNN_ERR_PARAM_NULLPTR);	546	CHECK_RET(CheckNotNull(params), ACLNN_ERR_PARAM_NULLPTR);
	@@ -537,6 +552,7 @@ static aclnnStatus CheckParams(const MhcParamsBase &params)
537	CHECK_RET(CheckInputOutShape(params), ACLNN_ERR_PARAM_INVALID);	552	CHECK_RET(CheckInputOutShape(params), ACLNN_ERR_PARAM_INVALID);
538	CHECK_RET(CheckDtypeValid(params), ACLNN_ERR_PARAM_INVALID);	553	CHECK_RET(CheckDtypeValid(params), ACLNN_ERR_PARAM_INVALID);
539	CHECK_RET(CheckFormat(params), ACLNN_ERR_PARAM_INVALID);	554	CHECK_RET(CheckFormat(params), ACLNN_ERR_PARAM_INVALID);
		555	+ CHECK_RET(CheckOpImplMode(params), ACLNN_ERR_PARAM_INVALID);
540	return ACLNN_SUCCESS;	556	return ACLNN_SUCCESS;
541	}	557	}
542		558
	@@ -604,7 +620,8 @@ static aclnnStatus mHCPreCommonProcess(MhcParamsBase &params, aclOpExecutor *exe
604	int64_t outFlag = params.invRmsOptional != nullptr ? 1 : 0;	620	int64_t outFlag = params.invRmsOptional != nullptr ? 1 : 0;
605	auto outParams =	621	auto outParams =
606	l0op::MhcPre(params.xContiguous, params.phiContiguous, params.alphaContiguous, params.biasContiguous,	622	l0op::MhcPre(params.xContiguous, params.phiContiguous, params.alphaContiguous, params.biasContiguous,
607	- params.gammaOptionalContiguous, outFlag, params.normEps, params.hcEps, executor);	623	+ params.gammaOptionalContiguous, outFlag, params.normEps, params.hcEps, params.opImplMode,
		624	+ executor);
608	CHECK_RET(outParams != std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR);	625	CHECK_RET(outParams != std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr), ACLNN_ERR_INNER_NULLPTR);
609	ret = CopyRequiredOutputs(std::get<0>(outParams), std::get<1>(outParams), std::get<2>(outParams), params, executor);	626	ret = CopyRequiredOutputs(std::get<0>(outParams), std::get<1>(outParams), std::get<2>(outParams), params, executor);
610	CHECK_RET(ret == ACLNN_SUCCESS, ret);	627	CHECK_RET(ret == ACLNN_SUCCESS, ret);
	@@ -612,11 +629,11 @@ static aclnnStatus mHCPreCommonProcess(MhcParamsBase &params, aclOpExecutor *exe
612	executor);	629	executor);
613	}	630	}
614		631
615	-aclnnStatus aclnnMhcPreGetWorkspaceSize(const aclTensor x, const aclTensor phi, const aclTensor *alpha,	632	+static aclnnStatus MhcPreGetWorkspaceSizeCommon(
616	- const aclTensor bias, const aclTensor gammaOptional, double normEps,	633	+ const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,
617	- double hcEps, aclTensor hIn, aclTensor hPost, aclTensor *hRes,	634	+ const aclTensor gammaOptional, double normEps, double hcEps, int64_t opImplMode, aclTensor hIn,
618	- aclTensor invRmsOptional, aclTensor hMixOptional, aclTensor *hPreOptional,	635	+ aclTensor hPost, aclTensor hRes, aclTensor invRmsOptional, aclTensor hMixOptional, aclTensor *hPreOptional,
619	- uint64_t workspaceSize, aclOpExecutor *executor)	636	+ uint64_t workspaceSize, aclOpExecutor *executor)
620	{	637	{
621	if (workspaceSize == nullptr) {	638	if (workspaceSize == nullptr) {
622	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(ACLNN_OP_NAME, "workspaceSize", "nullptr",	639	OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(ACLNN_OP_NAME, "workspaceSize", "nullptr",
	@@ -628,13 +645,11 @@ aclnnStatus aclnnMhcPreGetWorkspaceSize(const aclTensor x, const aclTensor phi
628	"output parameter must not be nullptr");	645	"output parameter must not be nullptr");
629	return ACLNN_ERR_PARAM_NULLPTR;	646	return ACLNN_ERR_PARAM_NULLPTR;
630	}	647	}
631	- L2_DFX_PHASE_1(aclnnMhcPre, DFX_IN(x, phi, alpha, bias, gammaOptional, normEps, hcEps),
632	- DFX_OUT(hIn, hPost, hRes, invRmsOptional, hMixOptional, hPreOptional));
633	auto uniqueExecutor = CREATE_EXECUTOR();	648	auto uniqueExecutor = CREATE_EXECUTOR();
634		649
635	MhcParamsBase params = MhcBuilder::Create()	650	MhcParamsBase params = MhcBuilder::Create()
636	.SetInput(x, phi, alpha, bias, gammaOptional)	651	.SetInput(x, phi, alpha, bias, gammaOptional)
637	- .SetAttr(normEps, hcEps)	652	+ .SetAttr(normEps, hcEps, opImplMode)
638	.SetOutput(hIn, hPost, hRes)	653	.SetOutput(hIn, hPost, hRes)
639	.SetOptionalOutput(invRmsOptional, hMixOptional, hPreOptional)	654	.SetOptionalOutput(invRmsOptional, hMixOptional, hPreOptional)
640	.Build();	655	.Build();
	@@ -647,12 +662,42 @@ aclnnStatus aclnnMhcPreGetWorkspaceSize(const aclTensor x, const aclTensor phi
647	return ACLNN_SUCCESS;	662	return ACLNN_SUCCESS;
648	}	663	}
649		664
		665	+aclnnStatus aclnnMhcPreGetWorkspaceSize(const aclTensor x, const aclTensor phi, const aclTensor *alpha,
		666	+ const aclTensor bias, const aclTensor gammaOptional, double normEps,
		667	+ double hcEps, aclTensor hIn, aclTensor hPost, aclTensor *hRes,
		668	+ aclTensor invRmsOptional, aclTensor hMixOptional, aclTensor *hPreOptional,
		669	+ uint64_t workspaceSize, aclOpExecutor *executor)
		670	+{
		671	+ L2_DFX_PHASE_1(aclnnMhcPre, DFX_IN(x, phi, alpha, bias, gammaOptional, normEps, hcEps),
		672	+ DFX_OUT(hIn, hPost, hRes, invRmsOptional, hMixOptional, hPreOptional));
		673	+ return MhcPreGetWorkspaceSizeCommon(x, phi, alpha, bias, gammaOptional, normEps, hcEps, MHC_PRE_USE_FP32, hIn,
		674	+ hPost, hRes, invRmsOptional, hMixOptional, hPreOptional, workspaceSize,
		675	+ executor);
		676	+}
		677	+
		678	+aclnnStatus aclnnMhcPreV2GetWorkspaceSize(
		679	+ const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,
		680	+ const aclTensor gammaOptional, double normEps, double hcEps, int64_t opImplMode, aclTensor hIn,
		681	+ aclTensor hPost, aclTensor hRes, aclTensor invRmsOptional, aclTensor hMixOptional, aclTensor *hPreOptional,
		682	+ uint64_t workspaceSize, aclOpExecutor *executor)
		683	+{
		684	+ L2_DFX_PHASE_1(aclnnMhcPreV2, DFX_IN(x, phi, alpha, bias, gammaOptional, normEps, hcEps, opImplMode),
		685	+ DFX_OUT(hIn, hPost, hRes, invRmsOptional, hMixOptional, hPreOptional));
		686	+ return MhcPreGetWorkspaceSizeCommon(x, phi, alpha, bias, gammaOptional, normEps, hcEps, opImplMode, hIn, hPost,
		687	+ hRes, invRmsOptional, hMixOptional, hPreOptional, workspaceSize, executor);
		688	+}
650	aclnnStatus aclnnMhcPre(void workspace, uint64_t workspaceSize, aclOpExecutor executor, aclrtStream stream)	689	aclnnStatus aclnnMhcPre(void workspace, uint64_t workspaceSize, aclOpExecutor executor, aclrtStream stream)
651	{	690	{
652	L2_DFX_PHASE_2(aclnnMhcPre);	691	L2_DFX_PHASE_2(aclnnMhcPre);
653	return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);	692	return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
654	}	693	}
655		694
		695	+aclnnStatus aclnnMhcPreV2(void workspace, uint64_t workspaceSize, aclOpExecutor executor, aclrtStream stream)
		696	+{
		697	+ L2_DFX_PHASE_2(aclnnMhcPreV2);
		698	+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
		699	+}
		700	+
656	} // namespace	701	} // namespace
657	#ifdef __cplusplus	702	#ifdef __cplusplus
658	}	703	}

Amhc/mhc_pre/op_host/op_api/aclnn_mhc_pre_v2.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	+
	11	+#ifndef OP_API_INC_MHC_PRE_V2_H
	12	+#define OP_API_INC_MHC_PRE_V2_H
	13	+
	14	+#include "aclnn/aclnn_base.h"
	15	+#include "aclnn_util.h"
	16	+
	17	+#ifdef __cplusplus
	18	+extern "C" {
	19	+#endif
	20	+
	21	+/**
	22	+ * @brief The first-stage interface of aclnnMhcPreV2. Calculates the required workspace size.
	23	+ * @domain aclnn_ops_infer
	24	+ * @param [in] x: Input tensor.
	25	+ * @param [in] phi: Input projection tensor.
	26	+ * @param [in] alpha: Input scaling tensor.
	27	+ * @param [in] bias: Input bias tensor.
	28	+ * @param [in] gammaOptional: Optional RMSNorm scale tensor.
	29	+ * @param [in] normEps: RMSNorm epsilon.
	30	+ * @param [in] hcEps: Hyper-connection epsilon.
	31	+ * @param [in] opImplMode: Matrix multiplication precision mode. 0 uses FP32 and 1 uses HF32.
	32	+ * @param [out] hIn: Output tensor.
	33	+ * @param [out] hPost: Output tensor.
	34	+ * @param [out] hRes: Output tensor.
	35	+ * @param [out] invRmsOptional: Optional output tensor.
	36	+ * @param [out] hMixOptional: Optional output tensor.
	37	+ * @param [out] hPreOptional: Optional output tensor.
	38	+ * @param [out] workspaceSize: Required device workspace size.
	39	+ * @param [out] executor: Operator executor.
	40	+ * @return aclnnStatus: Status code.
	41	+ */
	42	+ACLNN_API aclnnStatus aclnnMhcPreV2GetWorkspaceSize(
	43	+ const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,
	44	+ const aclTensor gammaOptional, double normEps, double hcEps, int64_t opImplMode, aclTensor hIn,
	45	+ aclTensor hPost, aclTensor hRes, aclTensor invRmsOptional, aclTensor hMixOptional, aclTensor *hPreOptional,
	46	+ uint64_t workspaceSize, aclOpExecutor *executor);
	47	+
	48	+/**
	49	+ * @brief The second-stage interface of aclnnMhcPreV2. Executes the operator.
	50	+ * @domain aclnn_ops_infer
	51	+ * @param [in] workspace: Device workspace address.
	52	+ * @param [in] workspaceSize: Workspace size returned by aclnnMhcPreV2GetWorkspaceSize.
	53	+ * @param [in] executor: Operator executor returned by aclnnMhcPreV2GetWorkspaceSize.
	54	+ * @param [in] stream: ACL stream used to execute the operator.
	55	+ * @return aclnnStatus: Status code.
	56	+ */
	57	+ACLNN_API aclnnStatus aclnnMhcPreV2(void workspace, uint64_t workspaceSize, aclOpExecutor executor,
	58	+ aclrtStream stream);
	59	+
	60	+#ifdef __cplusplus
	61	+}
	62	+#endif
	63	+
	64	+#endif // OP_API_INC_MHC_PRE_V2_H

Mmhc/mhc_pre/op_host/op_api/mhc_pre.cpp+6-4

Mmhc/mhc_pre/op_host/op_api/mhc_pre.h+2-1

	@@ -17,7 +17,8 @@
17	namespace l0op {	17	namespace l0op {
18	const std::tuple<aclTensor , aclTensor , aclTensor , aclTensor , aclTensor , aclTensor >	18	const std::tuple<aclTensor , aclTensor , aclTensor , aclTensor , aclTensor , aclTensor >
19	MhcPre(const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,	19	MhcPre(const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,
20	- const aclTensor gammaOptional, int64_t outFlag, double normEps, double hcEps, aclOpExecutor executor);	20	+ const aclTensor *gammaOptional, int64_t outFlag, double normEps, double hcEps, int64_t opImplMode,
		21	+ aclOpExecutor *executor);
21	}	22	}
22		23
23	#endif // OP_API_MHC_PRE_H_	24	#endif // OP_API_MHC_PRE_H_

Mmhc/mhc_pre/op_host/op_tiling/arch35/mhc_pre_tiling.cpp+300-176

	@@ -26,6 +26,8 @@
26	namespace optiling {	26	namespace optiling {
27		27
28	using namespace Ops::Transformer::OpTiling;	28	using namespace Ops::Transformer::OpTiling;
		29	+
		30	+// Supported input ranks and input tensor indices.
29	const constexpr int64_t BSND_DIM_NUM = 4;	31	const constexpr int64_t BSND_DIM_NUM = 4;
30	const constexpr int64_t TND_DIM_NUM = 3;	32	const constexpr int64_t TND_DIM_NUM = 3;
31	const constexpr uint32_t X_INDEX = 0;	33	const constexpr uint32_t X_INDEX = 0;

Mmhc/mhc_pre/op_host/op_tiling/arch35/mhc_pre_tiling.h+28-7

	@@ -24,11 +24,11 @@ namespace optiling {
24		24
25	enum class TilingMode : uint8_t {	25	enum class TilingMode : uint8_t {
26	SPLIT_BS = 0,	26	SPLIT_BS = 0,
27	- SPLIT_ND = 1	27	+ SPLIT_ND = 1,
		28	+ SPLIT_M_K = 2
28	};	29	};
29		30
30	BEGIN_TILING_DATA_DEF(MhcPreTilingData)	31	BEGIN_TILING_DATA_DEF(MhcPreTilingData)
31	-TILING_DATA_FIELD_DEF_STRUCT(TCubeTiling, matmulTiling);
32	TILING_DATA_FIELD_DEF(uint32_t, coreNum);	32	TILING_DATA_FIELD_DEF(uint32_t, coreNum);
33	TILING_DATA_FIELD_DEF(uint32_t, outFlag);	33	TILING_DATA_FIELD_DEF(uint32_t, outFlag);
34	TILING_DATA_FIELD_DEF(uint32_t, hasGamma);	34	TILING_DATA_FIELD_DEF(uint32_t, hasGamma);
	@@ -43,6 +43,20 @@ TILING_DATA_FIELD_DEF(uint64_t, D);
43	TILING_DATA_FIELD_DEF(float, normEps);	43	TILING_DATA_FIELD_DEF(float, normEps);
44	TILING_DATA_FIELD_DEF(float, hcEps);	44	TILING_DATA_FIELD_DEF(float, hcEps);
45	TILING_DATA_FIELD_DEF(float, scaleMean);	45	TILING_DATA_FIELD_DEF(float, scaleMean);
		46	+TILING_DATA_FIELD_DEF(uint32_t, mL1Size);
		47	+TILING_DATA_FIELD_DEF(uint32_t, kUbSize);
		48	+TILING_DATA_FIELD_DEF(uint32_t, kL1Size);
		49	+TILING_DATA_FIELD_DEF(uint32_t, fusionAlign);
		50	+TILING_DATA_FIELD_DEF(uint32_t, cubeBlockDimM);
		51	+TILING_DATA_FIELD_DEF(uint32_t, cubeBlockDimK);
		52	+TILING_DATA_FIELD_DEF(uint32_t, multCoreSplitKSize);
		53	+TILING_DATA_FIELD_DEF(uint32_t, mkWorkspaceMmOffset);
		54	+TILING_DATA_FIELD_DEF(uint32_t, mkWorkspaceRmsOffset);
		55	+TILING_DATA_FIELD_DEF(uint32_t, mkWorkspaceFinalOffset);
		56	+TILING_DATA_FIELD_DEF(uint32_t, mkUseGmStage);
		57	+TILING_DATA_FIELD_DEF(uint32_t, stage2UsedAivNum);
		58	+TILING_DATA_FIELD_DEF(uint32_t, stage2RowsPerCore);
		59	+TILING_DATA_FIELD_DEF(uint32_t, implMode);
46	END_TILING_DATA_DEF;	60	END_TILING_DATA_DEF;
47		61
48	REGISTER_TILING_DATA_CLASS(MhcPre, MhcPreTilingData);	62	REGISTER_TILING_DATA_CLASS(MhcPre, MhcPreTilingData);
	@@ -60,7 +74,8 @@ struct MhcPreCompileInfo {
60		74
61	class MhcPreBaseTiling : public Ops::Transformer::OpTiling::TilingBaseClass {	75	class MhcPreBaseTiling : public Ops::Transformer::OpTiling::TilingBaseClass {
62	public:	76	public:
63	- explicit MhcPreBaseTiling(gert::TilingContext *context) : Ops::Transformer::OpTiling::TilingBaseClass(context) {};	77	+ explicit MhcPreBaseTiling(gert::TilingContext *context)
		78	+ : Ops::Transformer::OpTiling::TilingBaseClass(context) {};
64		79
65	~MhcPreBaseTiling() override = default;	80	~MhcPreBaseTiling() override = default;
66		81
	@@ -115,11 +130,17 @@ protected:
115	void FillTilingData();	130	void FillTilingData();
116	ge::graphStatus TilingProcess();	131	ge::graphStatus TilingProcess();
117	ge::graphStatus CheckUbBufferSize();	132	ge::graphStatus CheckUbBufferSize();
		133	+ ge::graphStatus CalculateBsTiling();
		134	+ ge::graphStatus CalculateNdTiling();
		135	+ ge::graphStatus CalculateMkTiling();
		136	+ bool CanUseMk() const;
		137	+ uint32_t SelectMkStage() const;
118	size_t GetLastRequiredInputIndex() const;	138	size_t GetLastRequiredInputIndex() const;
119		139
120	private:	140	private:
121	MhcPreTilingData tilingData_;	141	MhcPreTilingData tilingData_;
122	uint32_t blockDim_ = 32; // AIC	142	uint32_t blockDim_ = 32; // AIC
		143	+ uint32_t aivBlockDim_ = 64;
123	uint64_t totalLength_ = 0;	144	uint64_t totalLength_ = 0;
124	uint64_t m_ = 0;	145	uint64_t m_ = 0;
125	uint64_t ubSize_ = 0;	146	uint64_t ubSize_ = 0;
	@@ -136,12 +157,12 @@ private:
136	uint32_t hasGamma_ = 0;	157	uint32_t hasGamma_ = 0;
137	uint32_t chunkTSize_ = 0;	158	uint32_t chunkTSize_ = 0;
138	uint32_t v1ChunkDSize_ = 0;	159	uint32_t v1ChunkDSize_ = 0;
		160	+ uint32_t usedCoreNum_ = 0;
		161	+ float scaleMean_ = 0.0f;
		162	+ uint32_t implMode_ = 0;
139	TilingMode tilingMode_ = TilingMode::SPLIT_BS;	163	TilingMode tilingMode_ = TilingMode::SPLIT_BS;
140	uint32_t hasResi_ = 1;	164	uint32_t hasResi_ = 1;
141	-
142	-protected:
143	- matmul_tiling::MultiCoreMatmulTiling mm_;
144	};	165	};
145		166
146	} // namespace optiling	167	} // namespace optiling
147	-#endif // __OP_HOST_MHC_PRE_TILING_H__	168	+#endif // __OP_HOST_MHC_PRE_TILING_H__

Amhc/mhc_pre/op_kernel/arch35/mhc_pre_common.h+169-0

@@ -0,0 +1,169 @@
	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	+/*!
	12	+ * \file mhc_pre_common.h
	13	+ * \brief Shared kernel types, constants, and hardware helpers for MHC Pre
	14	+ */
	15	+
	16	+#ifndef MHC_PRE_COMMON_H
	17	+#define MHC_PRE_COMMON_H
	18	+
	19	+#include "kernel_operator.h"
	20	+#include "kernel_tiling/kernel_tiling.h"
	21	+#include "lib/matmul_intf.h"
	22	+
	23	+namespace MhcPre {
	24	+
	25	+using namespace AscendC;
	26	+
	27	+struct InitParams {
	28	+ GM_ADDR x;
	29	+ GM_ADDR phi;
	30	+ GM_ADDR alpha;
	31	+ GM_ADDR bias;
	32	+ GM_ADDR gamma;
	33	+ GM_ADDR hin;
	34	+ GM_ADDR h_post;
	35	+ GM_ADDR h_res;
	36	+ GM_ADDR inv_rms;
	37	+ GM_ADDR h_mix;
	38	+ GM_ADDR h_pre;
	39	+ GM_ADDR workspace;
	40	+ TPipe *tPipeIn;
	41	+ const MhcPreTilingData *tilingData;
	42	+};
	43	+
	44	+struct MatrixInfo {
	45	+ uint32_t totalLength = 0;
	46	+ uint32_t nD = 0;
	47	+ uint32_t fusionSize = 0;
	48	+ float normEps = 0.0f;
	49	+ float hcEps = 0.0f;
	50	+};
	51	+
	52	+struct VectorOffsetParams {
	53	+ uint64_t globalOffsetM = 0;
	54	+ uint64_t singleCoreM = 0;
	55	+ uint64_t offsetMStart = 0;
	56	+ uint64_t offsetMEnd = 0;
	57	+};
	58	+
	59	+struct MNConfig {
	60	+ uint64_t m = 0;
	61	+ uint64_t n = 0;
	62	+ uint64_t k = 0;
	63	+ uint64_t singleCoreM = 0;
	64	+ uint64_t singleCoreN = 0;
	65	+ uint64_t singleCoreK = 0;
	66	+ uint64_t curSingleCoreM = 0;
	67	+ uint64_t curSingleCoreN = 0;
	68	+ uint64_t curSingleCoreK = 0;
	69	+};
	70	+
	71	+// Matmul implementation modes shared by vector and Cube paths.
	72	+constexpr uint32_t MHC_PRE_IMPL_MODE_FP32 = 0U;
	73	+constexpr uint32_t MHC_PRE_IMPL_MODE_HF32 = 1U;
	74	+
	75	+// Basic API hardware granularity and alignment.
	76	+constexpr uint32_t MHC_PRE_BASIC_API_BLOCK_SIZE = 32U;
	77	+constexpr uint32_t MHC_PRE_BASIC_API_C0_SIZE = 8U;
	78	+constexpr uint32_t MHC_PRE_BASIC_API_VL_FP32 = 64U;
	79	+
	80	+// Shared L1 double-buffer layout for A and B operands.
	81	+constexpr uint64_t MHC_PRE_BASIC_API_L1_ALLOC_SIZE = 512U * 1024U;
	82	+constexpr uint64_t MHC_PRE_BASIC_API_L1_BUF_NUM = 2U;
	83	+constexpr uint64_t MHC_PRE_BASIC_API_L1_BUF_OFFSET = 128U * 256U;
	84	+
	85	+// Basic API vector cast traits.
	86	+constexpr AscendC::MicroAPI::CastTrait MHC_PRE_BASIC_API_CAST_B16_TO_B32 = {
	87	+ AscendC::MicroAPI::RegLayout::ZERO,
	88	+ AscendC::MicroAPI::SatMode::UNKNOWN,
	89	+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
	90	+ AscendC::RoundMode::UNKNOWN,
	91	+};
	92	+
	93	+// UB alignment and cross-core synchronization protocol.
	94	+constexpr uint32_t MHC_PRE_UB_ALIGN_SIZE = 32U;
	95	+constexpr uint8_t MHC_PRE_CROSS_CORE_SYNC_MODE = 4U;
	96	+constexpr uint16_t MHC_PRE_SUBBLOCK_FLAG_OFFSET = 16U;
	97	+constexpr uint16_t MHC_PRE_X_READY_FLAG = 8U;
	98	+constexpr uint16_t MHC_PRE_X_CONSUMED_FLAG = 9U;
	99	+constexpr uint16_t MHC_PRE_MM_READY_FLAG = 10U;
	100	+
	101	+// Integer shape helpers.
	102	+__aicore__ inline uint32_t BasicApiCeilDiv(uint32_t value, uint32_t div) { return (value + div - 1) / div; }
	103	+
	104	+__aicore__ inline uint32_t BasicApiAlign(uint32_t value, uint32_t align)
	105	+{
	106	+ return BasicApiCeilDiv(value, align) * align;
	107	+}
	108	+
	109	+template <typename T>
	110	+__aicore__ inline T MhcPreMin(T a, T b)
	111	+{
	112	+ return a > b ? b : a;
	113	+}
	114	+
	115	+__aicore__ inline uint64_t MhcPreCeilDiv(uint64_t value, uint64_t div)
	116	+{
	117	+ return div == 0U ? value : (value + div - 1U) / div;
	118	+}
	119	+
	120	+__aicore__ inline uint64_t MhcPreAlign(uint64_t value, uint64_t align) { return MhcPreCeilDiv(value, align) * align; }
	121	+
	122	+__aicore__ inline constexpr uint32_t MhcPreGetVRegSize()
	123	+{
	124	+#if __CCE_AICORE__ == 310
	125	+ return AscendC::VECTOR_REG_WIDTH;
	126	+#else
	127	+ return 256U;
	128	+#endif
	129	+}
	130	+
	131	+template <typename T>
	132	+__aicore__ inline uint32_t BasicApiRoundUp(uint32_t value)
	133	+{
	134	+ return BasicApiAlign(value, MHC_PRE_BASIC_API_BLOCK_SIZE / sizeof(T));
	135	+}
	136	+
	137	+// Register and tensor data-movement helpers.
	138	+__aicore__ inline void MhcPreBasicApiLoadBroadcast(AscendC::MicroAPI::RegTensor<float> &dst, __local_mem__ float *src,
	139	+ AscendC::MicroAPI::MaskReg mask, uint32_t offset)
	140	+{
	141	+ AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(dst, src + offset);
	142	+}
	143	+
	144	+template <typename T>
	145	+__aicore__ inline void BasicApiCopyToL1(const LocalTensor<T> &srcLocal, const LocalTensor<T> &dstLocal,
	146	+ const DataCopyParams &copyParams)
	147	+{
	148	+ DataCopy(dstLocal, srcLocal, copyParams);
	149	+}
	150	+
	151	+template <typename T>
	152	+__aicore__ inline void MhcPreVFCompactRows(const LocalTensor<T> &dstLocal, const LocalTensor<T> &srcLocal,
	153	+ uint16_t rowCount, uint32_t srcRowStride, uint32_t rowWidth)
	154	+{
	155	+ __ubuf__ T dst = (__ubuf__ T )dstLocal.GetPhyAddr();
	156	+ __ubuf__ T src = (__ubuf__ T )srcLocal.GetPhyAddr();
	157	+ __VEC_SCOPE__
	158	+ {
	159	+ MicroAPI::RegTensor<T> data;
	160	+ for (uint16_t row = 0; row < rowCount; ++row) {
	161	+ MicroAPI::Load<T>(data, src + static_cast<uint32_t>(row) * srcRowStride);
	162	+ MicroAPI::Store<T>(dst + static_cast<uint32_t>(row) * rowWidth, data, rowWidth);
	163	+ }
	164	+ }
	165	+}
	166	+
	167	+} // namespace MhcPre
	168	+
	169	+#endif // MHC_PRE_COMMON_H

Amhc/mhc_pre/op_kernel/arch35/mhc_pre_cube_compute.h+326-0

@@ -0,0 +1,326 @@
	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	+/*!
	12	+ * \file mhc_pre_cube_compute.h
	13	+ * \brief Basic API Cube service for MHC Pre matmul
	14	+ */
	15	+
	16	+#ifndef MHC_PRE_CUBE_COMPUTE_H
	17	+#define MHC_PRE_CUBE_COMPUTE_H
	18	+
	19	+#include "mhc_pre_common.h"
	20	+
	21	+namespace MhcPre {
	22	+
	23	+// Mmad and Fixpipe unit-flag values.
	24	+constexpr uint32_t FINAL_ACCUMULATION = 3U;
	25	+constexpr uint32_t NON_FINAL_ACCUMULATION = 2U;
	26	+
	27	+// Fixpipe row-split policy.
	28	+constexpr bool MHC_PRE_BASIC_API_SPLIT_M = true;
	29	+constexpr uint64_t MHC_PRE_BASIC_API_SPLIT_M_ALIGN = 2U;
	30	+
	31	+class MhcPreCubeCompute {
	32	+public:
	33	+ uint64_t m_{0};
	34	+ uint64_t n_{0};
	35	+ uint64_t k_{0};
	36	+ uint64_t baseM_{0};
	37	+ uint64_t baseN_{0};
	38	+ uint64_t baseK_{0};
	39	+ uint64_t kL1_{0};
	40	+ uint32_t implMode_{MHC_PRE_IMPL_MODE_FP32};
	41	+ uint8_t bL1BufferID_{0};
	42	+ uint8_t l0PingPongID_{0};
	43	+ uint8_t cl0PingPongID_{0};
	44	+
	45	+ AscendC::LocalTensor<float> aL0Ping_;
	46	+ AscendC::LocalTensor<float> aL0Pong_;
	47	+ AscendC::LocalTensor<float> bL0Ping_;
	48	+ AscendC::LocalTensor<float> bL0Pong_;
	49	+ AscendC::LocalTensor<float> cL0Ping_;
	50	+ AscendC::LocalTensor<float> cL0Pong_;
	51	+
	52	+ __aicore__ inline MhcPreCubeCompute() {}
	53	+
	54	+ __aicore__ inline uint8_t GetBL1BufferId() { return bL1BufferID_; }
	55	+
	56	+ __aicore__ inline void CopyInA1Nd2Nz(uint64_t m, uint64_t currentK, const AscendC::GlobalTensor<float> &aGlobal,
	57	+ const AscendC::LocalTensor<float> &al1Local)
	58	+ {
	59	+ CopyInA1Nd2Nz(m, currentK, currentK, aGlobal, al1Local);
	60	+ }
	61	+
	62	+ __aicore__ inline void CopyInA1Nd2Nz(uint64_t m, uint64_t currentK, uint64_t srcDValue,
	63	+ const AscendC::GlobalTensor<float> &aGlobal,
	64	+ const AscendC::LocalTensor<float> &al1Local)
	65	+ {
	66	+ AscendC::Nd2NzParams nd2nzParam;
	67	+ nd2nzParam.ndNum = 1;
	68	+ nd2nzParam.nValue = m;
	69	+ nd2nzParam.dValue = currentK;
	70	+ nd2nzParam.srcNdMatrixStride = 1;
	71	+ nd2nzParam.srcDValue = srcDValue;
	72	+ nd2nzParam.dstNzC0Stride = BasicApiAlign(m, AscendC::BLOCK_CUBE);
	73	+ nd2nzParam.dstNzNStride = 1;
	74	+ nd2nzParam.dstNzMatrixStride = 1;
	75	+ AscendC::DataCopy(al1Local, aGlobal, nd2nzParam);
	76	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(2);
	77	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(2);
	78	+ }
	79	+
	80	+ __aicore__ inline void Init(uint32_t implMode)
	81	+ {
	82	+ implMode_ = implMode;
	83	+ uint32_t aL0OneBuffer = 256 * 32;
	84	+ uint32_t bL0OneBuffer = 256 * 32;
	85	+ uint32_t cL0OneBuffer = 256 * 128;
	86	+
	87	+ aL0Ping_ = AscendC::LocalTensor<float>(AscendC::TPosition::A2, 0, aL0OneBuffer);
	88	+ aL0Pong_ = AscendC::LocalTensor<float>(AscendC::TPosition::A2, aL0OneBuffer * sizeof(float), aL0OneBuffer);
	89	+ bL0Ping_ = AscendC::LocalTensor<float>(AscendC::TPosition::B2, 0, bL0OneBuffer);
	90	+ bL0Pong_ = AscendC::LocalTensor<float>(AscendC::TPosition::B2, bL0OneBuffer * sizeof(float), bL0OneBuffer);
	91	+ cL0Ping_ = AscendC::LocalTensor<float>(AscendC::TPosition::CO1, 0, cL0OneBuffer);
	92	+ cL0Pong_ = AscendC::LocalTensor<float>(AscendC::TPosition::CO1, cL0OneBuffer * sizeof(float), cL0OneBuffer);
	93	+
	94	+ // Ping-pong L0A/L0B/L0C so MTE1, Mmad, and Fixpipe can overlap.
	95	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(0);
	96	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(1);
	97	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(3);
	98	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(4);
	99	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(0);
	100	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(1);
	101	+
	102	+ // HF32 is a Cube mode. Restrict the mode write to AIC to avoid an invalid AIV-side operation.
	103	+ if ASCEND_IS_AIC {
	104	+ if (implMode == MHC_PRE_IMPL_MODE_HF32) {
	105	+ AscendC::SetHF32Mode(1);
	106	+ AscendC::SetHF32TransMode(1);
	107	+ }
	108	+ }
	109	+ }
	110	+
	111	+ __aicore__ inline void CopyInB1Nd2Nz(uint64_t k, uint64_t currentK, uint64_t baseN,
	112	+ const AscendC::GlobalTensor<float> &bGlobal,
	113	+ const AscendC::LocalTensor<float> &bl1Local)
	114	+ {
	115	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(bL1BufferID_);
	116	+ k_ = k;
	117	+ kL1_ = currentK;
	118	+ baseN_ = baseN;
	119	+ AscendC::Nd2NzParams nd2nzParam;
	120	+ nd2nzParam.ndNum = 1;
	121	+ nd2nzParam.nValue = baseN_;
	122	+ nd2nzParam.dValue = kL1_;
	123	+ nd2nzParam.srcNdMatrixStride = 1;
	124	+ nd2nzParam.srcDValue = k_;
	125	+ nd2nzParam.dstNzC0Stride = BasicApiAlign(baseN_, AscendC::BLOCK_CUBE);
	126	+ nd2nzParam.dstNzNStride = 1;
	127	+ nd2nzParam.dstNzMatrixStride = 1;
	128	+ AscendC::DataCopy(bl1Local, bGlobal, nd2nzParam);
	129	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(bL1BufferID_);
	130	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(bL1BufferID_);
	131	+ }
	132	+
	133	+ __aicore__ inline void Process(uint64_t m, uint64_t n, uint64_t baseM, uint64_t baseK, bool isFirstKL1,
	134	+ bool isLastKL1, const AscendC::LocalTensor<float> &al1Local,
	135	+ const AscendC::LocalTensor<float> &bl1Local)
	136	+ {
	137	+ m_ = m;
	138	+ n_ = n;
	139	+ baseM_ = baseM;
	140	+ baseK_ = baseK;
	141	+ if (isFirstKL1) {
	142	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(cl0PingPongID_);
	143	+ }
	144	+ uint64_t kL1Offset = 0;
	145	+ for (uint64_t kb = 0; kb < kL1_; kb += baseK_) {
	146	+ bool isLastKL0 = (kb + baseK_) >= kL1_;
	147	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0PingPongID_ + 3);
	148	+ CopyInA2(kb, kL1Offset, al1Local);
	149	+ CopyInB2(kb, kL1Offset, bl1Local);
	150	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0PingPongID_);
	151	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0PingPongID_);
	152	+ MmadBase(kb, isFirstKL1, isLastKL1, isLastKL0);
	153	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0PingPongID_ + 3);
	154	+ l0PingPongID_ ^= 1;
	155	+ kL1Offset += baseK_;
	156	+ }
	157	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(bL1BufferID_);
	158	+ bL1BufferID_ ^= 1;
	159	+ }
	160	+
	161	+ // Own iter0 (K-L1 staging and cross-core handshakes); Process owns iter1 (K-L0 and Mmad).
	162	+ __aicore__ inline void ProcessKRange(uint64_t m, uint64_t n, uint64_t k, uint64_t baseM, uint64_t baseK,
	163	+ uint64_t kStart, uint64_t kEnd, uint64_t kL1Size, bool useSequentialPartials,
	164	+ uint64_t partialK, bool useGmStage, uint64_t stageSlotElements,
	165	+ uint64_t stageCoreOffset, const AscendC::GlobalTensor<float> &phiGlobal,
	166	+ const AscendC::GlobalTensor<float> &xStageGlobal,
	167	+ const AscendC::GlobalTensor<float> &dstGlobal, uint64_t dstGroupStride,
	168	+ const AscendC::LocalTensor<float> &al1Local,
	169	+ const AscendC::LocalTensor<float> &bl1Local)
	170	+ {
	171	+ uint8_t aL1BufferId = 0;
	172	+ for (uint64_t kOffset = kStart; kOffset < kEnd; kOffset += kL1Size) {
	173	+ uint64_t currentK = AscendC::Std::min(kL1Size, kEnd - kOffset);
	174	+ uint8_t bL1BufferId = GetBL1BufferId();
	175	+ AscendC::LocalTensor<float> currentAL1 = al1Local[aL1BufferId * MHC_PRE_BASIC_API_L1_BUF_OFFSET];
	176	+ AscendC::LocalTensor<float> currentBL1 = bl1Local[bL1BufferId * MHC_PRE_BASIC_API_L1_BUF_OFFSET];
	177	+ CopyInB1Nd2Nz(k, currentK, n, phiGlobal[kOffset], currentBL1);
	178	+
	179	+ if (useGmStage) {
	180	+ AscendC::CrossCoreWaitFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE2>(MHC_PRE_X_READY_FLAG);
	181	+ AscendC::CrossCoreWaitFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE2>(MHC_PRE_X_READY_FLAG +
	182	+ MHC_PRE_SUBBLOCK_FLAG_OFFSET);
	183	+ uint64_t stageOffset = stageCoreOffset + aL1BufferId * stageSlotElements;
	184	+ CopyInA1Nd2Nz(m, currentK, kL1Size, xStageGlobal[stageOffset], currentAL1);
	185	+ } else {
	186	+ AscendC::CrossCoreWaitFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE1>(MHC_PRE_X_READY_FLAG);
	187	+ AscendC::CrossCoreWaitFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE1>(MHC_PRE_X_READY_FLAG +
	188	+ MHC_PRE_SUBBLOCK_FLAG_OFFSET);
	189	+ }
	190	+
	191	+ bool isFirstKL1 = useSequentialPartials ? kOffset % partialK == 0U : kOffset == kStart;
	192	+ bool isLastKL1 = useSequentialPartials ?
	193	+ ((kOffset + currentK) % partialK == 0U \|\| kOffset + currentK >= kEnd) :
	194	+ kOffset + currentK >= kEnd;
	195	+ Process(m, n, baseM, baseK, isFirstKL1, isLastKL1, currentAL1, currentBL1);
	196	+ if (isLastKL1) {
	197	+ uint64_t partialIndex = useSequentialPartials ? (kOffset + currentK - 1U) / partialK : 0U;
	198	+ CopyOut(dstGlobal[partialIndex * dstGroupStride], n);
	199	+ }
	200	+ AscendC::CrossCoreSetFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE1>(MHC_PRE_X_CONSUMED_FLAG);
	201	+ AscendC::CrossCoreSetFlag<MHC_PRE_CROSS_CORE_SYNC_MODE, PIPE_MTE1>(MHC_PRE_X_CONSUMED_FLAG +
	202	+ MHC_PRE_SUBBLOCK_FLAG_OFFSET);
	203	+ aL1BufferId ^= 1U;
	204	+ }
	205	+ }
	206	+ __aicore__ inline void CopyInA2(uint64_t kOffset, uint64_t kAL1Offset, const AscendC::LocalTensor<float> &al1Local)
	207	+ {
	208	+ uint64_t mAL1 = BasicApiAlign(baseM_, AscendC::BLOCK_CUBE);
	209	+ uint64_t offsetAL1 = BasicApiAlign(kAL1Offset, MHC_PRE_BASIC_API_C0_SIZE) * mAL1;
	210	+ AscendC::LoadData2DParamsV2 loadData2dParams;
	211	+ uint64_t currM = baseM_;
	212	+ uint64_t currK = AscendC::Std::min(baseK_, kL1_ - kOffset);
	213	+ loadData2dParams.mStartPosition = 0;
	214	+ loadData2dParams.kStartPosition = 0;
	215	+ loadData2dParams.mStep = BasicApiCeilDiv(currM, AscendC::BLOCK_CUBE);
	216	+ loadData2dParams.kStep = BasicApiCeilDiv(currK, MHC_PRE_BASIC_API_C0_SIZE);
	217	+ loadData2dParams.srcStride = BasicApiCeilDiv(currM, AscendC::BLOCK_CUBE);
	218	+ loadData2dParams.dstStride = loadData2dParams.mStep;
	219	+ loadData2dParams.ifTranspose = false;
	220	+ AscendC::LoadData(l0PingPongID_ == 0 ? aL0Ping_ : aL0Pong_, al1Local[offsetAL1], loadData2dParams);
	221	+ }
	222	+
	223	+ __aicore__ inline void CopyInB2(uint64_t kOffset, uint64_t kBL1Offset, const AscendC::LocalTensor<float> &bl1Local)
	224	+ {
	225	+ uint64_t nBL1 = BasicApiAlign(baseN_, AscendC::BLOCK_CUBE);
	226	+ uint64_t offsetBL1 = BasicApiAlign(kBL1Offset, MHC_PRE_BASIC_API_C0_SIZE) * nBL1;
	227	+ AscendC::LoadData2DParamsV2 loadData2dParams;
	228	+ uint64_t currN = baseN_;
	229	+ uint64_t currK = AscendC::Std::min(baseK_, kL1_ - kOffset);
	230	+ loadData2dParams.mStartPosition = 0;
	231	+ loadData2dParams.kStartPosition = 0;
	232	+ loadData2dParams.mStep = BasicApiCeilDiv(currN, AscendC::BLOCK_CUBE);
	233	+ loadData2dParams.kStep = BasicApiCeilDiv(currK, MHC_PRE_BASIC_API_C0_SIZE);
	234	+ loadData2dParams.srcStride = BasicApiCeilDiv(currN, AscendC::BLOCK_CUBE);
	235	+ loadData2dParams.dstStride = loadData2dParams.mStep;
	236	+ loadData2dParams.ifTranspose = false;
	237	+ AscendC::LoadData(l0PingPongID_ == 0 ? bL0Ping_ : bL0Pong_, bl1Local[offsetBL1], loadData2dParams);
	238	+ }
	239	+
	240	+ __aicore__ inline void MmadBase(uint64_t kOffset, bool isFirstKL1, bool isLastKL1, bool isLastKL0)
	241	+ {
	242	+ uint32_t mmK = AscendC::Std::min(baseK_, kL1_ - kOffset);
	243	+ // FP32 keeps K=32 accumulation order for precision; HF32 consumes the full baseK to reduce loop overhead.
	244	+ uint32_t mmKStep = implMode_ == MHC_PRE_IMPL_MODE_FP32 ? 32U : mmK;
	245	+ uint64_t mAlign = BasicApiAlign(baseM_, AscendC::BLOCK_CUBE);
	246	+ uint64_t nAlign = BasicApiAlign(baseN_, AscendC::BLOCK_CUBE);
	247	+ for (uint32_t innerK = 0; innerK < mmK; innerK += mmKStep) {
	248	+ uint32_t currentK = AscendC::Std::min(mmKStep, mmK - innerK);
	249	+ bool isLastInnerK = innerK + currentK >= mmK;
	250	+ AscendC::MmadParams params;
	251	+ params.m = baseM_;
	252	+ params.n = baseN_;
	253	+ params.k = currentK;
	254	+ params.disableGemv = true;
	255	+ params.cmatrixInitVal = (isFirstKL1 && kOffset == 0 && innerK == 0);
	256	+ params.cmatrixSource = false;
	257	+ params.unitFlag = (isLastKL1 && isLastKL0 && isLastInnerK) ? FINAL_ACCUMULATION : NON_FINAL_ACCUMULATION;
	258	+ AscendC::Mmad(cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_,
	259	+ (l0PingPongID_ == 0 ? aL0Ping_ : aL0Pong_)[innerK * mAlign],
	260	+ (l0PingPongID_ == 0 ? bL0Ping_ : bL0Pong_)[innerK * nAlign], params);
	261	+ AscendC::PipeBarrier<PIPE_M>();
	262	+ }
	263	+ }
	264	+
	265	+ __aicore__ inline void CopyOut(const AscendC::LocalTensor<float> &dstLocal)
	266	+ {
	267	+ // Direct L0C-to-UB removes the hMix GM write/read round trip when the caller has aligned UB space.
	268	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(cl0PingPongID_);
	269	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(cl0PingPongID_);
	270	+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> fixpipeParams;
	271	+ uint64_t c0 = AscendC::AuxGetC0Size<float>();
	272	+ fixpipeParams.nSize = BasicApiAlign(baseN_, c0);
	273	+ fixpipeParams.mSize =
	274	+ MHC_PRE_BASIC_API_SPLIT_M ? BasicApiAlign(baseM_, MHC_PRE_BASIC_API_SPLIT_M_ALIGN) : baseM_;
	275	+ fixpipeParams.dstStride = fixpipeParams.nSize;
	276	+ fixpipeParams.srcStride = BasicApiAlign(baseM_, AscendC::BLOCK_CUBE);
	277	+ fixpipeParams.quantPre = QuantMode_t::NoQuant;
	278	+ fixpipeParams.dualDstCtl =
	279	+ MHC_PRE_BASIC_API_SPLIT_M ? static_cast<uint8_t>(AscendC::McgShfMode::DUAL_DST_SPLIT_M) : 0;
	280	+ fixpipeParams.unitFlag = FINAL_ACCUMULATION;
	281	+ fixpipeParams.params.ndNum = 1;
	282	+ fixpipeParams.params.srcNdStride = 1;
	283	+ fixpipeParams.params.dstNdStride = 1;
	284	+ AscendC::Fixpipe<float, float, AscendC::Impl::CFG_ROW_MAJOR_UB>(
	285	+ dstLocal, cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_, fixpipeParams);
	286	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(cl0PingPongID_);
	287	+ cl0PingPongID_ ^= 1;
	288	+ }
	289	+
	290	+ __aicore__ inline void CopyOut(const AscendC::GlobalTensor<float> &dstGlobal, uint32_t dstStride)
	291	+ {
	292	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(cl0PingPongID_);
	293	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(cl0PingPongID_);
	294	+ AscendC::DataCopyCO12DstParams copyParams;
	295	+ copyParams.nSize = baseN_;
	296	+ copyParams.mSize = baseM_;
	297	+ copyParams.dstStride = dstStride;
	298	+ copyParams.srcStride = BasicApiAlign(baseM_, AscendC::BLOCK_CUBE);
	299	+ copyParams.quantPre = QuantMode_t::NoQuant;
	300	+ copyParams.nz2ndEn = true;
	301	+ copyParams.unitFlag = FINAL_ACCUMULATION;
	302	+ AscendC::SetFixpipeNz2ndFlag(1, 1, 1);
	303	+ AscendC::DataCopy(dstGlobal, cl0PingPongID_ == 0 ? cL0Ping_ : cL0Pong_, copyParams);
	304	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(cl0PingPongID_);
	305	+ cl0PingPongID_ ^= 1;
	306	+ }
	307	+
	308	+ __aicore__ inline void End(uint32_t implMode)
	309	+ {
	310	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(0);
	311	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(1);
	312	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(3);
	313	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(4);
	314	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(0);
	315	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(1);
	316	+ if ASCEND_IS_AIC {
	317	+ if (implMode == MHC_PRE_IMPL_MODE_HF32) {
	318	+ AscendC::SetHF32Mode(0);
	319	+ }
	320	+ }
	321	+ }
	322	+};
	323	+
	324	+} // namespace MhcPre
	325	+
	326	+#endif // MHC_PRE_CUBE_COMPUTE_H

Amhc/mhc_pre/op_kernel/arch35/mhc_pre_m_k_split.h+990-0

Mmhc/mhc_pre/op_kernel/arch35/mhc_pre_split_bs.h+211-404

	@@ -10,436 +10,243 @@
10		10
11	/*!	11	/*!
12	* \file mhc_pre_split_bs.h	12	* \file mhc_pre_split_bs.h
13	- * \brief MHC Pre kernel for batch split mode	13	+ * \brief BS-split Basic API implementation for MHC Pre
14	*/	14	*/
15		15
16	-#ifndef MHC_PRE_SPLIT_BS_H_	16	+#ifndef MHC_PRE_SPLIT_BS_H
17	-#define MHC_PRE_SPLIT_BS_H_	17	+#define MHC_PRE_SPLIT_BS_H
18		18
19	-#include "kernel_operator.h"	19	+#include "mhc_pre_vector_compute.h"
20	-#include "lib/matmul_intf.h"	20	+#include "mhc_pre_cube_compute.h"
21	-#include "kernel_tiling/kernel_tiling.h"
22	-#include "mhc_pre_utils.h"
23	-#include "mhc_pre_kernel_base.h"
24	-#include "mhc_pre_tiling_key.h"
25		21
26	namespace MhcPre {	22	namespace MhcPre {
27		23
28	-using namespace matmul;	24	+constexpr uint32_t MHC_PRE_BS_SEQUENTIAL_PARTIAL_K = 1024U;
29	-using namespace AscendC;
30		25
31	-static constexpr uint64_t SYNC_V2C = 0x1;	26	+template <class T, class P, int8_t RESI_MODE>
32	-static constexpr uint64_t SYNC_C2V = 0x2;	27	+class MhcPreSplitBS {
33	-static constexpr uint64_t SYNC_C2V1 = 0x3;
34	-
35	-static constexpr uint32_t DEFAULT_CHUNK_SIZE = 64;
36	-static constexpr uint32_t DEFAULT_V1_CHUNK_D_SIZE = 5120;
37	-
38	-template <class T, class P, int8_t HAS_RESI>
39	-class MhcPreKernelSplitBS : public MhcPreKernelBase<T, P> {
40	public:	28	public:
41	- using Base = MhcPreKernelBase<T, P>;	29	+ __aicore__ inline MhcPreSplitBS() = default;
42	- using Base::mm;
43	- using Base::mnConfig_;
44	- using Base::coreNum_;
45	- using Base::totalLength_;
46	- using Base::N_;
47	- using Base::D_;
48	- using Base::outFlag_;
49	- using Base::hasGamma_;
50	- using Base::xGm_;
51	- using Base::phiGm_;
52	- using Base::alphaGm_;
53	- using Base::biasGm_;
54	- using Base::gammaGm_;
55	- using Base::hinGm_;
56	- using Base::hPostGm_;
57	- using Base::hResGm_;
58	- using Base::invRmsGm_;
59	- using Base::hMixGm_;
60	- using Base::hPreGm_;
61	- using Base::xFloatGm_;
62	- using Base::xInQueue_;
63	- using Base::gammaInQueue_;
64	- using Base::invRmsOutQueue_;
65	- using Base::outQueue_;
66	- using Base::biasInQue_;
67	- using Base::tmpBuff_;
68	- using Base::alphaBuf_;
69	- using Base::inputBuff_;
70	- using Base::hPreBuff_;
71	- using Base::hPostBuff_;
72	- using Base::hResBuff_;
73	- using Base::alphaInUb_;
74	- using Base::biasInUb_;
75	- using Base::matmulRes_;
76	- using Base::xLocal_;
77	- using Base::invRmsUb_;
78	- using Base::gammaUb_;
79	- using Base::preOffsetBuf_;
80	- using Base::postOffsetBuf_;
81	- using Base::resOffsetBuf_;
82	- using Base::pipe_;
83	- using Base::matrixInfo_;
84	- using Base::vectorOffset_;
85	- using Base::tiling_;
86	- using Base::chunkTSize_;
87	- using Base::v1ChunkDSize_;
88	- using Base::curSingleT_;
89	- using Base::coreIdx_;
90	- using Base::subBlockIdx_;
91	- using Base::scaleMean_;
92	- using Base::globalOffsetM_;
93	- using Base::vectorCount_;
94	- using Base::cubeCount_;
95	- using Base::mmCount_;
96	- using Base::vec1Count_;
97	- using Base::eleNumPerVf_;
98	- using Base::kDoubleBufferCount;
99	- using Base::kSingleBufferCount;
100	- using Base::kHalfSplitDivisor;
101	- using Base::ND_LENGTH;
102	- using Base::PARALLEL_NUM;
103	- using Base::hasResi_;
104	- using Base::V0_BASE_T;
105	- using Base::V1_BASE_T;
106		30
107	- static constexpr int8_t kHasResi = HAS_RESI;	31	+ __aicore__ inline void Init(InitParams initParams)
		32	+ {
		33	+ vector_.BindGlobalTensors(initParams);

Mmhc/mhc_pre/op_kernel/arch35/mhc_pre_split_nd.h+297-413

Mmhc/mhc_pre/op_kernel/arch35/mhc_pre_tiling_key.h+16-7

	@@ -20,16 +20,26 @@
20		20
21	#define MHC_PRE_SPLIT_BS 0	21	#define MHC_PRE_SPLIT_BS 0
22	#define MHC_PRE_SPLIT_ND 1	22	#define MHC_PRE_SPLIT_ND 1
		23	+#define MHC_PRE_SPLIT_M_K 2
23	#define MHC_PRE_HAS_RESI 0	24	#define MHC_PRE_HAS_RESI 0
24	#define MHC_PRE_NO_RESI 1	25	#define MHC_PRE_NO_RESI 1
25		26
26	ASCENDC_TPL_ARGS_DECL(	27	ASCENDC_TPL_ARGS_DECL(
27	MhcPre,	28	MhcPre,
28	- ASCENDC_TPL_UINT_DECL(TILING_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_BS, MHC_PRE_SPLIT_ND),	29	+ ASCENDC_TPL_UINT_DECL(
29	- ASCENDC_TPL_UINT_DECL(HAS_RESI, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, MHC_PRE_HAS_RESI, MHC_PRE_NO_RESI)	30	+ TILING_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST,
30	-);	31	+ MHC_PRE_SPLIT_BS, MHC_PRE_SPLIT_ND, MHC_PRE_SPLIT_M_K),
		32	+ ASCENDC_TPL_UINT_DECL(HAS_RESI, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, MHC_PRE_HAS_RESI, MHC_PRE_NO_RESI));
31		33
32	ASCENDC_TPL_SEL(	34	ASCENDC_TPL_SEL(
		35	+ ASCENDC_TPL_ARGS_SEL(
		36	+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
		37	+ ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_ND),
		38	+ ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_HAS_RESI)),
		39	+ ASCENDC_TPL_ARGS_SEL(
		40	+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
		41	+ ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_ND),
		42	+ ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_NO_RESI)),
33	ASCENDC_TPL_ARGS_SEL(	43	ASCENDC_TPL_ARGS_SEL(
34	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),	44	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
35	ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_BS),	45	ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_BS),
	@@ -40,12 +50,11 @@ ASCENDC_TPL_SEL(
40	ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_NO_RESI)),	50	ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_NO_RESI)),
41	ASCENDC_TPL_ARGS_SEL(	51	ASCENDC_TPL_ARGS_SEL(
42	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),	52	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
43	- ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_ND),	53	+ ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_M_K),
44	ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_HAS_RESI)),	54	ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_HAS_RESI)),
45	ASCENDC_TPL_ARGS_SEL(	55	ASCENDC_TPL_ARGS_SEL(
46	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),	56	ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_MIX_AIC_1_2),
47	- ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_ND),	57	+ ASCENDC_TPL_UINT_SEL(TILING_MODE, ASCENDC_TPL_UI_LIST, MHC_PRE_SPLIT_M_K),
48	- ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_NO_RESI))	58	+ ASCENDC_TPL_UINT_SEL(HAS_RESI, ASCENDC_TPL_UI_LIST, MHC_PRE_NO_RESI)));
49	-);
50		59
51	#endif	60	#endif

Dmhc/mhc_pre/op_kernel/arch35/mhc_pre_utils.h+0-62

	@@ -1,62 +0,0 @@
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	-/*!
12	- * \file mhc_pre_utils.h
13	- * \brief
14	- */
15	-#ifndef ASCENDC_MHC_PRE_UTILS_H
16	-#define ASCENDC_MHC_PRE_UTILS_H
17	-
18	-namespace MhcPreUtils {
19	-
20	-constexpr uint32_t UB_ALIGN_SIZE = 32;
21	-
22	-constexpr uint8_t SYNC_AIC_AIV_MODE = 4;
23	-constexpr uint16_t FLAG_ID_MAX = 16;
24	-constexpr uint16_t AIC_SYNC_AIV_FLAG = 4;
25	-constexpr uint16_t AIV_SYNC_AIC_FLAG = 6;
26	-constexpr uint16_t VEC0_FLAG_ID_OFFSET = 0;
27	-constexpr uint16_t VEC1_FLAG_ID_OFFSET = FLAG_ID_MAX;
28	-
29	-template <typename T>
30	-__aicore__ inline T Max(T a, T b) {
31	- return a > b ? a : b;
32	-}
33	-
34	-template <typename T>
35	-__aicore__ inline T Min(T a, T b) {
36	- return a > b ? b : a;
37	-}
38	-
39	-__aicore__ inline uint64_t CeilDiv(uint64_t a, uint64_t b) {
40	- if (b == 0) {
41	- return a;
42	- }
43	- return (a + b - 1) / b;
44	-}
45	-
46	-__aicore__ inline uint64_t Align(uint64_t a, uint64_t b) {
47	- return CeilDiv(a, b) * b;
48	-}
49	-/**
50	- * Get the size of vector registers in bytes
51	- */
52	-__aicore__ inline constexpr uint32_t GetVRegSize()
53	-{
54	-#if __CCE_AICORE__ == 310
55	- return AscendC::VECTOR_REG_WIDTH;
56	-#else
57	- return 256U;
58	-#endif
59	-}
60	-}
61	-
62	-#endif // ASCENDC_MHC_PRE_UTILS_H

Rmhc/mhc_pre/op_kernel/arch35/mhc_pre_kernel_base.h→mhc/mhc_pre/op_kernel/arch35/mhc_pre_vector_compute.h+473-282

Mmhc/mhc_pre/op_kernel/mhc_pre_apt.cpp+27-37

	@@ -13,16 +13,16 @@
13	* \brief	13	* \brief
14	*/	14	*/
15	#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310	15	#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310
16	-#include "arch35/mhc_pre_kernel_base.h"	16	+#include "arch35/mhc_pre_common.h"
17	#include "arch35/mhc_pre_split_bs.h"	17	#include "arch35/mhc_pre_split_bs.h"
		18	+#include "arch35/mhc_pre_m_k_split.h"
18	#include "arch35/mhc_pre_split_nd.h"	19	#include "arch35/mhc_pre_split_nd.h"
19	#include "arch35/mhc_pre_tiling_key.h"	20	#include "arch35/mhc_pre_tiling_key.h"
20		21
21	using namespace AscendC;	22	using namespace AscendC;
22	-using namespace matmul;
23	using namespace MhcPre;	23	using namespace MhcPre;
24		24
25	-template <int8_t TILING_MODE, int8_t HAS_RESI>	25	+template <int8_t TILING_MODE, int8_t RESI_MODE>
26	__global__ __aicore__ void mhc_pre(GM_ADDR x, GM_ADDR phi, GM_ADDR alpha, GM_ADDR bias, GM_ADDR gamma, GM_ADDR hin,	26	__global__ __aicore__ void mhc_pre(GM_ADDR x, GM_ADDR phi, GM_ADDR alpha, GM_ADDR bias, GM_ADDR gamma, GM_ADDR hin,
27	GM_ADDR h_post, GM_ADDR h_res, GM_ADDR inv_rms, GM_ADDR h_mix, GM_ADDR h_pre,	27	GM_ADDR h_post, GM_ADDR h_res, GM_ADDR inv_rms, GM_ADDR h_mix, GM_ADDR h_pre,
28	GM_ADDR workspaceGM, GM_ADDR tilingGM)	28	GM_ADDR workspaceGM, GM_ADDR tilingGM)
	@@ -33,41 +33,31 @@ __global__ __aicore__ void mhc_pre(GM_ADDR x, GM_ADDR phi, GM_ADDR alpha, GM_ADD
33	KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);	33	KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
34	TPipe pipe;	34	TPipe pipe;
35	if constexpr (TILING_MODE == MHC_PRE_SPLIT_BS) {	35	if constexpr (TILING_MODE == MHC_PRE_SPLIT_BS) {
36	- if constexpr (HAS_RESI == MHC_PRE_HAS_RESI) {	36	+ InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,
37	- InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,	37	+ h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};
38	- h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};	38	+ MhcPreSplitBS<DTYPE_X, float32_t, RESI_MODE> op;
39	- MT mm;	39	+ op.Init(initParams);
40	- mm.Init(&tilingData.matmulTiling, &pipe);	40	+ op.Process();
41	- MhcPreKernelSplitBS<DTYPE_X, float32_t, MHC_PRE_HAS_RESI> op(mm);	41	+ } else if constexpr (TILING_MODE == MHC_PRE_SPLIT_M_K) {
42	- op.Init(initParams);	42	+ InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,
43	- op.Process();	43	+ h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};
44	- } else {	44	+ MhcPreMKPart1<DTYPE_X, float32_t, RESI_MODE> part1;
45	- InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,	45	+ part1.Init(initParams);
46	- h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};	46	+ part1.Process();
47	- MT mm;	47	+ // Stage 1 and stage 2 use different UB layouts. Releasing stage-1 queues lets stage 2
48	- mm.Init(&tilingData.matmulTiling, &pipe);	48	+ // reuse the full UB budget for reduction and postprocessing.
49	- MhcPreKernelSplitBS<DTYPE_X, float32_t, MHC_PRE_NO_RESI> op(mm);	49	+ pipe.Destroy();
50	- op.Init(initParams);	50	+ TPipe stage2Pipe;
51	- op.Process();	51	+ initParams.tPipeIn = &stage2Pipe;
52	- }	52	+ MhcPreMKPart2<DTYPE_X, float32_t, RESI_MODE> part2;
		53	+ part2.Init(initParams);
		54	+ part2.Process();
53	} else if constexpr (TILING_MODE == MHC_PRE_SPLIT_ND) {	55	} else if constexpr (TILING_MODE == MHC_PRE_SPLIT_ND) {
54	- if constexpr (HAS_RESI == MHC_PRE_HAS_RESI) {	56	+ InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,
55	- InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,	57	+ h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};
56	- h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};	58	+ MhcPreSplitND<DTYPE_X, float32_t, RESI_MODE> op;
57	- MT mm;	59	+ op.Init(initParams);
58	- mm.Init(&tilingData.matmulTiling, &pipe);	60	+ op.Process();
59	- MhcPreKernelSplitND<DTYPE_X, float32_t, MHC_PRE_HAS_RESI> op(mm);
60	- op.Init(initParams);
61	- op.Process();
62	- } else {
63	- InitParams initParams{x, phi, alpha, bias, gamma, hin, h_post,
64	- h_res, inv_rms, h_mix, h_pre, user, &pipe, &tilingData};
65	- MT mm;
66	- mm.Init(&tilingData.matmulTiling, &pipe);
67	- MhcPreKernelSplitND<DTYPE_X, float32_t, MHC_PRE_NO_RESI> op(mm);
68	- op.Init(initParams);
69	- op.Process();
70	- }
71	}	61	}
72	}	62	}
73	#endif	63	#endif

Mmhc/mhc_pre/tests/ut/op_host/op_api/test_aclnn_mhc_pre.cpp+36-2

Mmhc/mhc_pre/tests/ut/op_host/test_mhc_pre_tiling.cpp+1137-901

	@@ -24,9 +24,10 @@ OP_TYPE_REGISTER(MhcPre);
24		24
25	const std::tuple<aclTensor , aclTensor , aclTensor , aclTensor , aclTensor , aclTensor >	25	const std::tuple<aclTensor , aclTensor , aclTensor , aclTensor , aclTensor , aclTensor >
26	MhcPre(const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,	26	MhcPre(const aclTensor x, const aclTensor phi, const aclTensor alpha, const aclTensor bias,
27	- const aclTensor gammaOptional, int64_t outFlag, double normEps, double hcEps, aclOpExecutor executor)	27	+ const aclTensor *gammaOptional, int64_t outFlag, double normEps, double hcEps, int64_t opImplMode,
		28	+ aclOpExecutor *executor)
28	{	29	{
29	- L0_DFX(MhcPre, x, phi, alpha, bias, gammaOptional, outFlag, normEps, hcEps);	30	+ L0_DFX(MhcPre, x, phi, alpha, bias, gammaOptional, outFlag, normEps, hcEps, opImplMode);
30		31
31	DataType outType = DataType::DT_FLOAT;	32	DataType outType = DataType::DT_FLOAT;
32	Format format = Format::FORMAT_ND;	33	Format format = Format::FORMAT_ND;
	@@ -40,12 +41,13 @@ MhcPre(const aclTensor x, const aclTensor phi, const aclTensor *alpha, const a
40		41
41	auto ret = INFER_SHAPE(MhcPre, OP_INPUT(x, phi, alpha, bias, gammaOptional),	42	auto ret = INFER_SHAPE(MhcPre, OP_INPUT(x, phi, alpha, bias, gammaOptional),
42	OP_OUTPUT(outHin, outHpost, outHres, outInvRms, outMmRes, outHpre),	43	OP_OUTPUT(outHin, outHpost, outHres, outInvRms, outMmRes, outHpre),
43	- OP_ATTR(outFlag, static_cast<float>(normEps), static_cast<float>(hcEps)));	44	+ OP_ATTR(outFlag, static_cast<float>(normEps), static_cast<float>(hcEps), opImplMode));
44	OP_CHECK_INFERSHAPE(ret != ACLNN_SUCCESS, return std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr),	45	OP_CHECK_INFERSHAPE(ret != ACLNN_SUCCESS, return std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr),
45	"MhcPre InferShape failed.");	46	"MhcPre InferShape failed.");
46	auto ret1 = ADD_TO_LAUNCHER_LIST_AICORE(MhcPre, OP_INPUT(x, phi, alpha, bias, gammaOptional),	47	auto ret1 = ADD_TO_LAUNCHER_LIST_AICORE(MhcPre, OP_INPUT(x, phi, alpha, bias, gammaOptional),
47	OP_OUTPUT(outHin, outHpost, outHres, outInvRms, outMmRes, outHpre),	48	OP_OUTPUT(outHin, outHpost, outHres, outInvRms, outMmRes, outHpre),
48	- OP_ATTR(outFlag, static_cast<float>(normEps), static_cast<float>(hcEps)));	49	+ OP_ATTR(outFlag, static_cast<float>(normEps), static_cast<float>(hcEps),
		50	+ opImplMode));
49	OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(ret1 != ACLNN_SUCCESS,	51	OP_CHECK_ADD_TO_LAUNCHER_LIST_AICORE(ret1 != ACLNN_SUCCESS,
50	return std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr),	52	return std::tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr),
51	"MhcPre ADD_TO_LAUNCHER_LIST_AICORE failed.");	53	"MhcPre ADD_TO_LAUNCHER_LIST_AICORE failed.");