bmm支持小m,k轴且n轴=1模板-ops-nn-AtomGit

已合并

bmm支持小m,k轴且n轴=1模板 #6103

zhang-junming21创建于 6月15日从hid54649903/ops-nn： bmm_opt 合入到 cann/ops-nn： master

bmm支持小m,k轴且n轴=1模板 #6103

已合并

zhang-junming21创建于 6月15日

从hid54649903/ops-nn bmm_opt合入到cann/ops-nn master

共 10 个文件变更+1204-20

Amatmul/batch_mat_mul_v3/examples/test_aclnn_batchmatmul_vector.cpp+550-0

@@ -0,0 +1,550 @@
	1	+/**
	2	+ * Copyright (c) 2025-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_batchmatmul_v3_vector.cpp
	13	+ * @brief 测试 aclnnBatchMatMul 的 vector 计算路径。
	14	+ * 输入条件严格参考 CheckVectorComputationCondition():
	15	+ * 1. A、B 维度范围 [3, 6]
	16	+ * 2. A 和 B 维数相同
	17	+ * 3. B 最后一维 == 1 (vector kernel 的核心特征)
	18	+ * 4. A最后一维(K轴) == B倒数第二维(K轴)，满足矩阵乘法约束
	19	+ * 5. 数据类型全部 fp32
	20	+ * 6. batch 维度（前 dim-2 维）全部相等（广播维度必须匹配）
	21	+ */
	22	+
	23	+#include <iostream>
	24	+#include <memory>
	25	+#include <vector>
	26	+#include "acl/acl.h"
	27	+#include "aclnnop/aclnn_batch_matmul.h"
	28	+
	29	+#define CHECK_RET(cond, return_expr) \
	30	+ do { \
	31	+ if (!(cond)) { \
	32	+ return_expr; \
	33	+ } \
	34	+ } while (0)
	35	+
	36	+#define LOG_PRINT(message, ...) \
	37	+ do { \
	38	+ printf(message, ##__VA_ARGS__); \
	39	+ } while (0)
	40	+
	41	+int64_t GetShapeSize(const std::vector<int64_t>& shape)
	42	+{
	43	+ int64_t shapeSize = 1;
	44	+ for (auto i : shape) {
	45	+ shapeSize *= i;
	46	+ }
	47	+ return shapeSize;
	48	+}
	49	+
	50	+int Init(int32_t deviceId, aclrtStream* stream)
	51	+{
	52	+ // 固定写法，资源初始化
	53	+ auto ret = aclInit(nullptr);
	54	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
	55	+ ret = aclrtSetDevice(deviceId);
	56	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
	57	+ ret = aclrtCreateStream(stream);
	58	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
	59	+ return 0;
	60	+}
	61	+
	62	+template <typename T>
	63	+int CreateAclTensor(
	64	+ const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr, aclDataType dataType,
	65	+ aclTensor** tensor)
	66	+{
	67	+ auto size = GetShapeSize(shape) * sizeof(T);
	68	+ // 调用aclrtMalloc申请Device侧内存
	69	+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
	70	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
	71	+
	72	+ // 调用aclrtMemcpy将Host侧数据拷贝到Device侧内存上
	73	+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
	74	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
	75	+
	76	+ // 计算连续tensor的strides
	77	+ std::vector<int64_t> strides(shape.size(), 1);
	78	+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
	79	+ strides[i] = shape[i + 1] * strides[i + 1];
	80	+ }
	81	+
	82	+ // 调用aclCreateTensor接口创建aclTensor
	83	+ *tensor = aclCreateTensor(
	84	+ shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(),
	85	+ *deviceAddr);
	86	+ return 0;
	87	+}
	88	+
	89	+/**
	90	+ * @brief 封装 aclnnBatchMatMul 调用：GetWorkspaceSize → 申请workspace → 执行 → 同步 → 拷贝结果
	91	+ * @return ACL_SUCCESS 表示成功，其他值表示失败
	92	+ */
	93	+int RunBatchMatMul(
	94	+ aclTensor* self, aclTensor* mat2, aclTensor* out, int8_t cubeMathType,
	95	+ aclrtStream stream, void* outDeviceAddr, std::vector<float>& resultData)
	96	+{
	97	+ uint64_t workspaceSize = 0;
	98	+ aclOpExecutor* executor = nullptr;
	99	+
	100	+ // aclnnBatchMatMul接口调用示例
	101	+ // 3. 调用CANN算子库API，需要修改为具体的API名称
	102	+ // 调用aclnnBatchMatMul第一段接口
	103	+ auto ret = aclnnBatchMatMulGetWorkspaceSize(self, mat2, out, cubeMathType, &workspaceSize, &executor);
	104	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBatchMatMulGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
	105	+ // 根据第一段接口计算出的workspaceSize申请device内存
	106	+ void* workspaceAddr = nullptr;
	107	+ if (workspaceSize > 0UL) {
	108	+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
	109	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
	110	+ }
	111	+ // 调用aclnnBatchMatMul第二段接口
	112	+ ret = aclnnBatchMatMul(workspaceAddr, workspaceSize, executor, stream);
	113	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnBatchMatMul failed. ERROR: %d\n", ret); return ret);
	114	+
	115	+ // 4. （固定写法）同步等待任务执行结束
	116	+ ret = aclrtSynchronizeStream(stream);
	117	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
	118	+
	119	+ // 5. 获取输出的值，将Device侧内存上的结果拷贝至Host侧，需要根据具体API的接口定义修改
	120	+ ret = aclrtMemcpy(
	121	+ resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr,
	122	+ resultData.size() * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
	123	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
	124	+
	125	+ // 释放workspace
	126	+ if (workspaceAddr != nullptr) {
	127	+ aclrtFree(workspaceAddr);
	128	+ }
	129	+
	130	+ return 0;
	131	+}
	132	+
	133	+int main()
	134	+{
	135	+ // 1. （固定写法）device/stream初始化，参考acl API手册
	136	+ // 根据自己的实际device填写deviceId
	137	+ int32_t deviceId = 0;
	138	+ aclrtStream stream;
	139	+ auto ret = Init(deviceId, &stream);
	140	+ // check根据自己的需要处理
	141	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
	142	+
	143	+ // =========================================================================
	144	+ // Test Case 1: A[150, 4, 4] x B[150, 4, 1] = C[150, 4, 1]
	145	+ // 满足条件: dim=3, B最后维=1, K轴匹配(4==4), fp32, batch=150
	146	+ // =========================================================================
	147	+ {
	148	+ LOG_PRINT("\n===== Test Case 1: [150,4,4] x [150,4,1] =====\n");
	149	+
	150	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	151	+ std::vector<int64_t> selfShape = {150, 4, 4};
	152	+ std::vector<int64_t> mat2Shape = {150, 4, 1};
	153	+ std::vector<int64_t> outShape = {150, 4, 1};
	154	+ void* selfDeviceAddr = nullptr;
	155	+ void* mat2DeviceAddr = nullptr;
	156	+ void* outDeviceAddr = nullptr;
	157	+ aclTensor* self = nullptr;
	158	+ aclTensor* mat2 = nullptr;
	159	+ aclTensor* out = nullptr;
	160	+ // A[150,4,4]: 全1.0, 共2400个元素
	161	+ // B[150,4,1]: 全1.0, 共600个元素
	162	+ // 期望C[150,4,1]: 每元素 = 4*1.0 = 4.0
	163	+ std::vector<float> selfHostData(150 * 4 * 4, 1.0f);
	164	+ std::vector<float> mat2HostData(150 * 4 * 1, 1.0f);
	165	+ std::vector<float> outHostData(150 * 4 * 1, 0);
	166	+ int8_t cubeMathType = 1;
	167	+ // 创建self aclTensor
	168	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	169	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	170	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	171	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	172	+ // 创建mat2 aclTensor
	173	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	174	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	175	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	176	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	177	+ // 创建out aclTensor
	178	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	179	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	180	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	181	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	182	+
	183	+ std::vector<float> resultData(150 * 4 * 1, 0);
	184	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	185	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	186	+
	187	+ float expectedVal = 4.0f;
	188	+ for (int64_t i = 0; i < 150 * 4; i++) {
	189	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	190	+ }
	191	+ }
	192	+
	193	+ // =========================================================================
	194	+ // Test Case 2: A[150, 3, 3] x B[150, 3, 1] = C[150, 3, 1]
	195	+ // 满足条件: dim=3, B最后维=1, K轴匹配(3==3), fp32, batch=150
	196	+ // =========================================================================
	197	+ {
	198	+ LOG_PRINT("\n===== Test Case 2: [150,3,3] x [150,3,1] =====\n");
	199	+
	200	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	201	+ std::vector<int64_t> selfShape = {150, 3, 3};
	202	+ std::vector<int64_t> mat2Shape = {150, 3, 1};
	203	+ std::vector<int64_t> outShape = {150, 3, 1};
	204	+ void* selfDeviceAddr = nullptr;
	205	+ void* mat2DeviceAddr = nullptr;
	206	+ void* outDeviceAddr = nullptr;
	207	+ aclTensor* self = nullptr;
	208	+ aclTensor* mat2 = nullptr;
	209	+ aclTensor* out = nullptr;
	210	+ // A[150,3,3]: 全1.0, 共1350个元素
	211	+ // B[150,3,1]: 全1.0, 共450个元素
	212	+ // 期望C[150,3,1]: 每元素 = 3*1.0 = 3.0
	213	+ std::vector<float> selfHostData(150 * 3 * 3, 1.0f);
	214	+ std::vector<float> mat2HostData(150 * 3 * 1, 1.0f);
	215	+ std::vector<float> outHostData(150 * 3 * 1, 0);
	216	+ int8_t cubeMathType = 1;
	217	+ // 创建self aclTensor
	218	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	219	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	220	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	221	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	222	+ // 创建mat2 aclTensor
	223	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	224	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	225	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	226	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	227	+ // 创建out aclTensor
	228	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	229	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	230	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	231	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	232	+
	233	+ std::vector<float> resultData(150 * 3 * 1, 0);
	234	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	235	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	236	+
	237	+ float expectedVal = 3.0f;
	238	+ for (int64_t i = 0; i < 150 * 3; i++) {
	239	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	240	+ }
	241	+ }
	242	+
	243	+ // =========================================================================
	244	+ // Test Case 3: A[40001, 4, 4] x B[40001, 4, 1] = C[40001, 4, 1]
	245	+ // 满足条件: dim=3, B最后维=1, K轴匹配(4==4), fp32, batch=40001
	246	+ // =========================================================================
	247	+ {
	248	+ LOG_PRINT("\n===== Test Case 3: [40001,4,4] x [40001,4,1] =====\n");
	249	+
	250	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	251	+ std::vector<int64_t> selfShape = {40001, 4, 4};
	252	+ std::vector<int64_t> mat2Shape = {40001, 4, 1};
	253	+ std::vector<int64_t> outShape = {40001, 4, 1};
	254	+ void* selfDeviceAddr = nullptr;
	255	+ void* mat2DeviceAddr = nullptr;
	256	+ void* outDeviceAddr = nullptr;
	257	+ aclTensor* self = nullptr;
	258	+ aclTensor* mat2 = nullptr;
	259	+ aclTensor* out = nullptr;
	260	+ // A[40001,4,4]: 全1.0
	261	+ // B[40001,4,1]: 全1.0
	262	+ // 期望C[40001,4,1]: 每元素 = 4*1.0 = 4.0
	263	+ std::vector<float> selfHostData(40001 * 4 * 4, 1.0f);
	264	+ std::vector<float> mat2HostData(40001 * 4 * 1, 1.0f);
	265	+ std::vector<float> outHostData(40001 * 4 * 1, 0);
	266	+ int8_t cubeMathType = 1;
	267	+ // 创建self aclTensor
	268	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	269	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	270	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	271	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	272	+ // 创建mat2 aclTensor
	273	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	274	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	275	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	276	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	277	+ // 创建out aclTensor
	278	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	279	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	280	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	281	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	282	+
	283	+ std::vector<float> resultData(40001 * 4 * 1, 0);
	284	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	285	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	286	+
	287	+ float expectedVal = 4.0f;
	288	+ for (int64_t i = 40001 * 4 - 200; i < 40001 * 4; i++) {
	289	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	290	+ }
	291	+ }
	292	+
	293	+ // =========================================================================
	294	+ // Test Case 4: A[40001, 3, 3] x B[40001, 3, 1] = C[40001, 3, 1]
	295	+ // 满足条件: dim=3, B最后维=1, K轴匹配(3==3), fp32, batch=40001
	296	+ // =========================================================================
	297	+ {
	298	+ LOG_PRINT("\n===== Test Case 4: [40001,3,3] x [40001,3,1] =====\n");
	299	+
	300	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	301	+ std::vector<int64_t> selfShape = {40001, 3, 3};
	302	+ std::vector<int64_t> mat2Shape = {40001, 3, 1};
	303	+ std::vector<int64_t> outShape = {40001, 3, 1};
	304	+ void* selfDeviceAddr = nullptr;
	305	+ void* mat2DeviceAddr = nullptr;
	306	+ void* outDeviceAddr = nullptr;
	307	+ aclTensor* self = nullptr;
	308	+ aclTensor* mat2 = nullptr;
	309	+ aclTensor* out = nullptr;
	310	+ // A[40001,3,3]: 全1.0, 共40123个元素
	311	+ // B[40001,3,1]: 全1.0, 共120个元素
	312	+ // 期望C[40001,3,1]: 每元素 = 3*1.0 = 3.0
	313	+ std::vector<float> selfHostData(40001 * 3 * 3, 1.0f);
	314	+ std::vector<float> mat2HostData(40001 * 3 * 1, 1.0f);
	315	+ std::vector<float> outHostData(40001 * 3 * 1, 0);
	316	+ int8_t cubeMathType = 1;
	317	+ // 创建self aclTensor
	318	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	319	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	320	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	321	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	322	+ // 创建mat2 aclTensor
	323	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	324	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	325	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	326	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	327	+ // 创建out aclTensor
	328	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	329	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	330	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	331	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	332	+
	333	+ std::vector<float> resultData(40001 * 3 * 1, 0);
	334	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	335	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	336	+
	337	+ float expectedVal = 3.0f;
	338	+ for (int64_t i = 40001 * 3 - 200; i < 40001 * 3; i++) {
	339	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	340	+ }
	341	+ }
	342	+
	343	+ // =========================================================================
	344	+ // Test Case 5: A[150, 7, 5] x B[150, 5, 1] = C[150, 7, 1]
	345	+ // 满足条件: dim=3, B最后维=1, K轴匹配(5==5), fp32, batch=150
	346	+ // =========================================================================
	347	+ {
	348	+ LOG_PRINT("\n===== Test Case 5: [150,7,5] x [150,5,1] =====\n");
	349	+
	350	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	351	+ std::vector<int64_t> selfShape = {150, 7, 5};
	352	+ std::vector<int64_t> mat2Shape = {150, 5, 1};
	353	+ std::vector<int64_t> outShape = {150, 7, 1};
	354	+ void* selfDeviceAddr = nullptr;
	355	+ void* mat2DeviceAddr = nullptr;
	356	+ void* outDeviceAddr = nullptr;
	357	+ aclTensor* self = nullptr;
	358	+ aclTensor* mat2 = nullptr;
	359	+ aclTensor* out = nullptr;
	360	+ // A[150,7,5]: 全1.0
	361	+ // B[150,5,1]: 全1.0
	362	+ // 期望C[150,7,1]: 每元素 = 5*1.0 = 5.0
	363	+ std::vector<float> selfHostData(150 * 7 * 5, 1.0f);
	364	+ std::vector<float> mat2HostData(150 * 5 * 1, 1.0f);
	365	+ std::vector<float> outHostData(150 * 7 * 1, 0);
	366	+ int8_t cubeMathType = 1;
	367	+ // 创建self aclTensor
	368	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	369	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	370	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	371	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	372	+ // 创建mat2 aclTensor
	373	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	374	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	375	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	376	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	377	+ // 创建out aclTensor
	378	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	379	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	380	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	381	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	382	+
	383	+ std::vector<float> resultData(150 * 7 * 1, 0);
	384	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	385	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	386	+
	387	+ float expectedVal = 5.0f;
	388	+ for (int64_t i = 0; i < 150 * 7 * 1; i++) {
	389	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	390	+ }
	391	+ }
	392	+
	393	+ // =========================================================================
	394	+ // Test Case 6: A[40001, 7, 5] x B[40001, 5, 1] = C[40001, 7, 1]
	395	+ // 满足条件: dim=3, B最后维=1, K轴匹配(5==5), fp32, batch=40001
	396	+ // =========================================================================
	397	+ {
	398	+ LOG_PRINT("\n===== Test Case 6: [40001,7,5] x [40001,5,1] =====\n");
	399	+
	400	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	401	+ std::vector<int64_t> selfShape = {40001, 7, 5};
	402	+ std::vector<int64_t> mat2Shape = {40001, 5, 1};
	403	+ std::vector<int64_t> outShape = {40001, 7, 1};
	404	+ void* selfDeviceAddr = nullptr;
	405	+ void* mat2DeviceAddr = nullptr;
	406	+ void* outDeviceAddr = nullptr;
	407	+ aclTensor* self = nullptr;
	408	+ aclTensor* mat2 = nullptr;
	409	+ aclTensor* out = nullptr;
	410	+ // A[40001, 7, 5]: 全1.0
	411	+ // B[40001,5,1]: 全1.0
	412	+ // 期望C[40001,7,1]: 每元素 = 5*1.0 = 5.0
	413	+ std::vector<float> selfHostData(40001 * 7 * 5, 1.0f);
	414	+ std::vector<float> mat2HostData(40001 * 5 * 1, 1.0f);
	415	+ std::vector<float> outHostData(40001 * 7 * 1, 0);
	416	+ int8_t cubeMathType = 1;
	417	+ // 创建self aclTensor
	418	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	419	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	420	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	421	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	422	+ // 创建mat2 aclTensor
	423	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	424	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	425	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	426	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	427	+ // 创建out aclTensor
	428	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	429	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	430	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	431	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	432	+
	433	+ std::vector<float> resultData(40001 * 7 * 1, 0);
	434	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	435	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	436	+
	437	+ float expectedVal = 5.0f;
	438	+ for (int64_t i = 40001 * 7 * 1 - 200; i < 40001 * 7 * 1; i++) {
	439	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	440	+ }
	441	+ }
	442	+
	443	+ // =========================================================================
	444	+ // Test Case 7: A[150, 6, 8] x B[150, 8, 1] = C[150, 6, 1]
	445	+ // 满足条件: dim=3, B最后维=1, K轴匹配(8==8), fp32, batch=150
	446	+ // =========================================================================
	447	+ {
	448	+ LOG_PRINT("\n===== Test Case 7: [150,6,8] x [150,8,1] =====\n");
	449	+
	450	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	451	+ std::vector<int64_t> selfShape = {150, 6, 8};
	452	+ std::vector<int64_t> mat2Shape = {150, 8, 1};
	453	+ std::vector<int64_t> outShape = {150, 6, 1};
	454	+ void* selfDeviceAddr = nullptr;
	455	+ void* mat2DeviceAddr = nullptr;
	456	+ void* outDeviceAddr = nullptr;
	457	+ aclTensor* self = nullptr;
	458	+ aclTensor* mat2 = nullptr;
	459	+ aclTensor* out = nullptr;
	460	+ // A[150,6,8]: 全1.0
	461	+ // B[150,8,1]: 全1.0
	462	+ // 期望C[150,6,1]: 每元素 = 8*1.0 = 8.0
	463	+ std::vector<float> selfHostData(150 * 6 * 8, 1.0f);
	464	+ std::vector<float> mat2HostData(150 * 8 * 1, 1.0f);
	465	+ std::vector<float> outHostData(150 * 6 * 1, 0);
	466	+ int8_t cubeMathType = 1;
	467	+ // 创建self aclTensor
	468	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	469	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	470	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	471	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	472	+ // 创建mat2 aclTensor
	473	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	474	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	475	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	476	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	477	+ // 创建out aclTensor
	478	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	479	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	480	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	481	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	482	+
	483	+ std::vector<float> resultData(150 * 6 * 1, 0);
	484	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	485	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	486	+
	487	+ float expectedVal = 8.0f;
	488	+ for (int64_t i = 0; i < 150 * 6 * 1; i++) {
	489	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	490	+ }
	491	+ }
	492	+
	493	+ // =========================================================================
	494	+ // Test Case 8: A[40001, 6, 8] x B[40001, 8, 1] = C[40001, 6, 1]
	495	+ // 满足条件: dim=3, B最后维=1, K轴匹配(8==8), fp32, batch=40001
	496	+ // =========================================================================
	497	+ {
	498	+ LOG_PRINT("\n===== Test Case 8: [40001,6,8] x [40001,8,1] =====\n");
	499	+
	500	+ // 2. 构造输入与输出，需要根据API的接口自定义构造
	501	+ std::vector<int64_t> selfShape = {40001, 6, 8};
	502	+ std::vector<int64_t> mat2Shape = {40001, 8, 1};
	503	+ std::vector<int64_t> outShape = {40001, 6, 1};
	504	+ void* selfDeviceAddr = nullptr;
	505	+ void* mat2DeviceAddr = nullptr;
	506	+ void* outDeviceAddr = nullptr;
	507	+ aclTensor* self = nullptr;
	508	+ aclTensor* mat2 = nullptr;
	509	+ aclTensor* out = nullptr;
	510	+ // A[40001,6,8]: 全1.0
	511	+ // B[40001,8,1]: 全1.0
	512	+ // 期望C[40001,6,1]: 每元素 = 8*1.0 = 8.0
	513	+ std::vector<float> selfHostData(40001 * 6 * 8, 1.0f);
	514	+ std::vector<float> mat2HostData(40001 * 8 * 1, 1.0f);
	515	+ std::vector<float> outHostData(40001 * 6 * 1, 0);
	516	+ int8_t cubeMathType = 1;
	517	+ // 创建self aclTensor
	518	+ ret = CreateAclTensor(selfHostData, selfShape, &selfDeviceAddr, aclDataType::ACL_FLOAT, &self);
	519	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> selfTensorPtr(self, aclDestroyTensor);
	520	+ std::unique_ptr<void, aclError ()(void)> selfDeviceAddrPtr(selfDeviceAddr, aclrtFree);
	521	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	522	+ // 创建mat2 aclTensor
	523	+ ret = CreateAclTensor(mat2HostData, mat2Shape, &mat2DeviceAddr, aclDataType::ACL_FLOAT, &mat2);
	524	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> mat2TensorPtr(mat2, aclDestroyTensor);
	525	+ std::unique_ptr<void, aclError ()(void)> mat2DeviceAddrPtr(mat2DeviceAddr, aclrtFree);
	526	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	527	+ // 创建out aclTensor
	528	+ ret = CreateAclTensor(outHostData, outShape, &outDeviceAddr, aclDataType::ACL_FLOAT, &out);
	529	+ std::unique_ptr<aclTensor, aclnnStatus ()(const aclTensor)> outTensorPtr(out, aclDestroyTensor);
	530	+ std::unique_ptr<void, aclError ()(void)> outdeviceAddrPtr(outDeviceAddr, aclrtFree);
	531	+ CHECK_RET(ret == ACL_SUCCESS, return ret);
	532	+
	533	+ std::vector<float> resultData(40001 * 6 * 1, 0);
	534	+ ret = RunBatchMatMul(self, mat2, out, cubeMathType, stream, outDeviceAddr, resultData);
	535	+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("RunBatchMatMul failed. ERROR: %d\n", ret); return ret);
	536	+
	537	+ float expectedVal = 8.0f;
	538	+ for (int64_t i = 40001 * 6 * 1 - 200; i < 40001 * 6 * 1; i++) {
	539	+ LOG_PRINT("result[%ld] is: %f (expected: %f)\n", i, resultData[i], expectedVal);
	540	+ }
	541	+ }
	542	+
	543	+ LOG_PRINT("\n===== All vector kernel tests completed =====\n");
	544	+
	545	+ // 6. 释放device资源，需要根据具体API的接口定义修改
	546	+ aclrtDestroyStream(stream);
	547	+ aclrtResetDevice(deviceId);
	548	+ aclFinalize();
	549	+ return 0;
	550	+}

Mmatmul/batch_mat_mul_v3/op_host/op_tiling/batch_mat_mul_v3_base_tiling.cpp+243-14

	@@ -23,6 +23,7 @@
23	#include "matmul/common/op_host/math_util.h"	23	#include "matmul/common/op_host/math_util.h"
24	#include "matmul/common/op_host/op_tiling/debug_tiling.h"	24	#include "matmul/common/op_host/op_tiling/debug_tiling.h"
25	#include "platform/platform_infos_def.h"	25	#include "platform/platform_infos_def.h"
		26	+#include "tiling/platform/platform_ascendc.h"
26		27
27	using namespace optiling::batch_mat_mul_v3;	28	using namespace optiling::batch_mat_mul_v3;
28	using Ops::NN::MathUtil;	29	using Ops::NN::MathUtil;
	@@ -57,6 +58,9 @@ constexpr uint64_t L2_SIZE_2 = 192UL * 1024UL * 1024UL;
57	constexpr uint64_t MAX_TRANS_CONFLICT = 6;	58	constexpr uint64_t MAX_TRANS_CONFLICT = 6;
58	constexpr double TAIL_CONFLICT_RATIO = 0.5;	59	constexpr double TAIL_CONFLICT_RATIO = 0.5;
59	constexpr uint64_t MAX_INT32_VALUE = 2147483647uL;	60	constexpr uint64_t MAX_INT32_VALUE = 2147483647uL;
		61	+constexpr uint64_t VECTOR_DIM_THRESHOLD = 3;
		62	+constexpr uint64_t VECTOR_ALIGN_NUM = 64;
		63	+constexpr uint64_t VECTOR_BUFFER_MULTIPLIER = 24 * 4;
60		64
61	static inline uint64_t LastPower2(uint64_t n)	65	static inline uint64_t LastPower2(uint64_t n)
62	{	66	{
	@@ -289,9 +293,8 @@ void BatchMatmulV3BaseTiling::DoL2CacheAndCalOrderTiling()
289	}	293	}
290	}	294	}
291		295
292	-ge::graphStatus BatchMatmulV3BaseTiling::DoLibApiTiling()	296	+void BatchMatmulV3BaseTiling::SetBatchDimInfo()
293	{	297	{
294	- auto ret = MatmulV3BaseTiling::DoLibApiTiling();
295	bmmTilingData_.multiBatchInfo.batchUsedCoreNum = bmmTilingData_.matmulTiling.matmulTiling.usedCoreNum;	298	bmmTilingData_.multiBatchInfo.batchUsedCoreNum = bmmTilingData_.matmulTiling.matmulTiling.usedCoreNum;
296	bmmTilingData_.multiBatchInfo.aBatchDim3 = static_cast<uint32_t>(batchInfo_.batchA3);	299	bmmTilingData_.multiBatchInfo.aBatchDim3 = static_cast<uint32_t>(batchInfo_.batchA3);
297	bmmTilingData_.multiBatchInfo.aBatchDim2 = static_cast<uint32_t>(batchInfo_.batchA2);	300	bmmTilingData_.multiBatchInfo.aBatchDim2 = static_cast<uint32_t>(batchInfo_.batchA2);
	@@ -305,7 +308,10 @@ ge::graphStatus BatchMatmulV3BaseTiling::DoLibApiTiling()
305	bmmTilingData_.multiBatchInfo.cBatchDim2 = static_cast<uint32_t>(batchInfo_.batchC2);	308	bmmTilingData_.multiBatchInfo.cBatchDim2 = static_cast<uint32_t>(batchInfo_.batchC2);
306	bmmTilingData_.multiBatchInfo.cBatchDim1 = static_cast<uint32_t>(batchInfo_.batchC1);	309	bmmTilingData_.multiBatchInfo.cBatchDim1 = static_cast<uint32_t>(batchInfo_.batchC1);
307	bmmTilingData_.multiBatchInfo.cBatchDim0 = static_cast<uint32_t>(batchInfo_.batchC0);	310	bmmTilingData_.multiBatchInfo.cBatchDim0 = static_cast<uint32_t>(batchInfo_.batchC0);
		311	+}
308		312
		313	+void BatchMatmulV3BaseTiling::CalcBatchDimAll()
		314	+{
309	aBatchDimAll_ = batchInfo_.batchA0 * batchInfo_.batchA1 * batchInfo_.batchA2 * batchInfo_.batchA3;	315	aBatchDimAll_ = batchInfo_.batchA0 * batchInfo_.batchA1 * batchInfo_.batchA2 * batchInfo_.batchA3;
310	bBatchDimAll_ = batchInfo_.batchB0 * batchInfo_.batchB1 * batchInfo_.batchB2 * batchInfo_.batchB3;	316	bBatchDimAll_ = batchInfo_.batchB0 * batchInfo_.batchB1 * batchInfo_.batchB2 * batchInfo_.batchB3;
311	cBatchDimAll_ = batchInfo_.batchC0 * batchInfo_.batchC1 * batchInfo_.batchC2 * batchInfo_.batchC3;	317	cBatchDimAll_ = batchInfo_.batchC0 * batchInfo_.batchC1 * batchInfo_.batchC2 * batchInfo_.batchC3;
	@@ -313,31 +319,56 @@ ge::graphStatus BatchMatmulV3BaseTiling::DoLibApiTiling()
313	bmmTilingData_.multiBatchInfo.bBatchDimAll = static_cast<uint32_t>(bBatchDimAll_);	319	bmmTilingData_.multiBatchInfo.bBatchDimAll = static_cast<uint32_t>(bBatchDimAll_);
314	bmmTilingData_.multiBatchInfo.cBatchDimAll = static_cast<uint32_t>(cBatchDimAll_);	320	bmmTilingData_.multiBatchInfo.cBatchDimAll = static_cast<uint32_t>(cBatchDimAll_);
315	bmmTilingData_.multiBatchInfo.batchTileBlock = static_cast<uint32_t>(cBatchDimAll_);	321	bmmTilingData_.multiBatchInfo.batchTileBlock = static_cast<uint32_t>(cBatchDimAll_);
		322	+}
		323	+
		324	+bool BatchMatmulV3BaseTiling::CheckNd2NzOnTheFlyLimit()
		325	+{
		326	+ uint64_t innerSizeA = args_.isATrans ? args_.mValue : args_.kValue;
		327	+ uint64_t innerSizeB = args_.isBTrans ? args_.kValue : args_.nValue;
		328	+ return innerSizeA > ND2NZ_ON_THE_FLY_LIMIT \|\| innerSizeB > ND2NZ_ON_THE_FLY_LIMIT;
		329	+}
		330	+
		331	+void BatchMatmulV3BaseTiling::DoMultiBatchAndL1FullLoadTiling()
		332	+{
		333	+ if (compileInfo_.supportL0c2out && tilingSelect_ != TilingCalcSelect::COMMON &&
		334	+ std::string(context_->GetNodeType()) != "TransposeBatchMatMul" &&
		335	+ args_.bFormat != ge::FORMAT_FRACTAL_NZ) {
		336	+ OP_LOGD("Enter DoMultiBatchTiling");
		337	+ DoMultiBatchTiling();
		338	+ if (IsMultiBatchAL1FullLoad()) {
		339	+ OP_LOGD("Enter DoMultiBatchL1FullLoadTiling");
		340	+ DoMultiBatchL1FullLoadTiling();
		341	+ }
		342	+ }
		343	+}
		344	+
		345	+ge::graphStatus BatchMatmulV3BaseTiling::DoLibApiTiling()
		346	+{
		347	+ auto ret = MatmulV3BaseTiling::DoLibApiTiling();
		348	+ SetBatchDimInfo();
		349	+ CalcBatchDimAll();
316	if (CheckBMMTilingDataIsVaild()) {	350	if (CheckBMMTilingDataIsVaild()) {
317	return ge::GRAPH_FAILED;	351	return ge::GRAPH_FAILED;
318	}	352	}
319	bmmTilingData_.multiBatchInfo.biasWithBatch = static_cast<uint32_t>(batchInfo_.biasWithBatch);	353	bmmTilingData_.multiBatchInfo.biasWithBatch = static_cast<uint32_t>(batchInfo_.biasWithBatch);
320	bmmTilingData_.multiBatchInfo.mOri = static_cast<uint32_t>(args_.mOriValue);	354	bmmTilingData_.multiBatchInfo.mOri = static_cast<uint32_t>(args_.mOriValue);
321		355
322	- uint64_t innerSizeA = args_.isATrans ? args_.mValue : args_.kValue;	356	+ if (CheckNd2NzOnTheFlyLimit()) {
323	- uint64_t innerSizeB = args_.isBTrans ? args_.kValue : args_.nValue;
324	- if (innerSizeA > ND2NZ_ON_THE_FLY_LIMIT \|\| innerSizeB > ND2NZ_ON_THE_FLY_LIMIT) {
325	DoUnAlignCommonTiling();	357	DoUnAlignCommonTiling();
326	DoTilingKeyCustom();	358	DoTilingKeyCustom();
327	return ret;	359	return ret;
328	}	360	}
329		361
330	DoCommonTiling();	362	DoCommonTiling();
		363	+ if (CheckVectorComputationCondition()) {
		364	+ DoVectorTiling();
		365	+ DoTilingKeyCustom();
		366	+ return ret;
		367	+ }
		368	+
331	DoL1FullLoadTiling();	369	DoL1FullLoadTiling();
332	DoL2CacheAndCalOrderTiling();	370	DoL2CacheAndCalOrderTiling();
333	- if (compileInfo_.supportL0c2out && tilingSelect_ != TilingCalcSelect::COMMON &&	371	+ DoMultiBatchAndL1FullLoadTiling();
334	- std::string(context_->GetNodeType()) != "TransposeBatchMatMul" &&
335	- args_.bFormat != ge::FORMAT_FRACTAL_NZ) {
336	- DoMultiBatchTiling();
337	- if (IsMultiBatchAL1FullLoad()) { // 多batch AL1全载
338	- DoMultiBatchL1FullLoadTiling();
339	- }
340	- }
341	DoTilingKeyCustom();	372	DoTilingKeyCustom();
342	return ret;	373	return ret;
343	}	374	}
	@@ -994,6 +1025,200 @@ void BatchMatmulV3BaseTiling::DoL1FullLoadTiling()
994	}	1025	}
995	}	1026	}
996		1027
		1028	+bool BatchMatmulV3BaseTiling::CalcVectorShapeInfo(VectorShapeInfo &shapeInfo)
		1029	+{
		1030	+ auto platformInfoptr = context_->GetPlatformInfo();
		1031	+ if (platformInfoptr == nullptr) { return false; }
		1032	+ auto ascendplatformInfo = platform_ascendc::PlatformAscendC(platformInfoptr);
		1033	+ shapeInfo.coreNumber = ascendplatformInfo.GetCoreNumAiv();
		1034	+
		1035	+ auto aStorageShape = context_->GetInputShape(0)->GetStorageShape();
		1036	+ auto bStorageShape = context_->GetInputShape(1)->GetStorageShape();
		1037	+ const auto dimNum = aStorageShape.GetDimNum();
		1038	+ if (dimNum < VECTOR_DIM_THRESHOLD) { return false; }
		1039	+
		1040	+ shapeInfo.aTotalSize = 1;
		1041	+ for (size_t i = 0; i < aStorageShape.GetDimNum(); i++) {
		1042	+ shapeInfo.aTotalSize *= aStorageShape.GetDim(i);
		1043	+ }
		1044	+ shapeInfo.dimSizeSecondLast = aStorageShape.GetDim(dimNum - 2);
		1045	+ shapeInfo.dimSizeLast = aStorageShape.GetDim(dimNum - 1);
		1046	+
		1047	+ shapeInfo.bTotalSize = 1;
		1048	+ for (size_t i = 0; i < bStorageShape.GetDimNum(); i++) {
		1049	+ shapeInfo.bTotalSize *= bStorageShape.GetDim(i);
		1050	+ }
		1051	+ if (shapeInfo.dimSizeLast == 0) { return false; }
		1052	+
		1053	+ shapeInfo.batchSize = shapeInfo.aTotalSize / shapeInfo.dimSizeLast;
		1054	+ const uint64_t bDimNum = bStorageShape.GetDimNum();
		1055	+ shapeInfo.bRowsPerBatch = bStorageShape.GetDim(bDimNum - 1);
		1056	+ return true;
		1057	+}
		1058	+
		1059	+bool BatchMatmulV3BaseTiling::CalcVectorCoreParams(const VectorShapeInfo &shapeInfo, VectorCoreParams &coreParams)
		1060	+{
		1061	+ auto platformInfoptr = context_->GetPlatformInfo();
		1062	+ auto ascendplatformInfo = platform_ascendc::PlatformAscendC(platformInfoptr);
		1063	+
		1064	+ coreParams.coreData = ops::CeilDiv(shapeInfo.batchSize, shapeInfo.coreNumber);
		1065	+ coreParams.coreData = ops::CeilAlign(coreParams.coreData, VECTOR_ALIGN_NUM);
		1066	+
		1067	+ uint64_t totalUbBytes;
		1068	+ ascendplatformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, totalUbBytes);
		1069	+ coreParams.alignedDimSizeLast = ops::CeilAlign(shapeInfo.dimSizeLast * aDtypeSize_, (uint64_t)32) / aDtypeSize_;
		1070	+ const uint64_t perRowUbCost = 4 * coreParams.alignedDimSizeLast * aDtypeSize_ + aDtypeSize_;
		1071	+ coreParams.rowsPerCore = totalUbBytes / perRowUbCost;
		1072	+ coreParams.rowsPerCore = ops::FloorAlign(coreParams.rowsPerCore, VECTOR_ALIGN_NUM);
		1073	+
		1074	+ if (coreParams.coreData == 0 \|\| coreParams.rowsPerCore == 0) { return false; }
		1075	+ return true;
		1076	+}
		1077	+
		1078	+void BatchMatmulV3BaseTiling::SetVectorTilingParams(const VectorShapeInfo &shapeInfo,
		1079	+ const VectorCoreParams &coreParams)
		1080	+{
		1081	+ uint64_t usedCoreNum = ops::CeilDiv(shapeInfo.batchSize, coreParams.coreData);
		1082	+ usedCoreNum = std::min(usedCoreNum, shapeInfo.coreNumber);
		1083	+
		1084	+ bmmTilingData_.vectorTilingInfo.coreNumber = usedCoreNum;
		1085	+ bmmTilingData_.vectorTilingInfo.coreData = coreParams.coreData;
		1086	+ bmmTilingData_.vectorTilingInfo.rowsPerCore = coreParams.rowsPerCore;
		1087	+ bmmTilingData_.vectorTilingInfo.aTotalSize = shapeInfo.aTotalSize;
		1088	+ bmmTilingData_.vectorTilingInfo.bTotalSize = shapeInfo.bTotalSize;
		1089	+ bmmTilingData_.vectorTilingInfo.dimSizeSecondLast = shapeInfo.dimSizeSecondLast;
		1090	+ bmmTilingData_.vectorTilingInfo.dimSizeLast = shapeInfo.dimSizeLast;
		1091	+ bmmTilingData_.vectorTilingInfo.alignedDimSizeLast = coreParams.alignedDimSizeLast;
		1092	+ bmmTilingData_.vectorTilingInfo.bRowsPerBatch = shapeInfo.bRowsPerBatch;
		1093	+ bmmTilingData_.vectorTilingInfo.cTotalSize = shapeInfo.batchSize;
		1094	+
		1095	+ isVectorMode_ = true;
		1096	+ // 注意：vector kernel被实现为特殊的VECTOR_FULL_LOAD模式，这是为了架构一致性
		1097	+ // vector计算不使用L1全载，而是使用UB全载进行数据处理
		1098	+ tilingEnable_.tilingEnableMultiBatchL1FullLoad = TilingEnableMultiBatchL1FullLoad::IS_FALSE;
		1099	+ tilingEnable_.tilingEnableMultiBatch = TilingEnableMultiBatch::IS_TRUE;
		1100	+ tilingEnable_.tilingEnableLoadMode = TilingEnableLoadMode::VECTOR_FULL_LOAD;
		1101	+ tilingEnable_.tilingEnableMultiBatchOut = TilingEnableMultiBatchOut::IS_FALSE;
		1102	+ tilingEnable_.tilingEnableMixNd2Nz = TilingEnableMixNd2Nz::IS_FALSE;
		1103	+}
		1104	+
		1105	+void BatchMatmulV3BaseTiling::DoVectorTiling()
		1106	+{
		1107	+ VectorShapeInfo shapeInfo{};
		1108	+ if (!CalcVectorShapeInfo(shapeInfo)) { return; }
		1109	+
		1110	+ VectorCoreParams coreParams{};
		1111	+ if (!CalcVectorCoreParams(shapeInfo, coreParams)) { return; }
		1112	+
		1113	+ SetVectorTilingParams(shapeInfo, coreParams);
		1114	+}
		1115	+
		1116	+bool BatchMatmulV3BaseTiling::CheckVectorNpuArch()
		1117	+{
		1118	+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());
		1119	+ auto npuArch = ascendcPlatform.GetCurNpuArch();
		1120	+ if (npuArch != NpuArch::DAV_2201) {
		1121	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A2/A3 is required. "
		1122	+ "Bmm vector opt version not supported.");
		1123	+ return false;
		1124	+ }
		1125	+ return true;
		1126	+}
		1127	+
		1128	+bool BatchMatmulV3BaseTiling::CheckVectorShapeDims()
		1129	+{
		1130	+ auto aShape = context_->GetInputShape(0)->GetOriginShape();
		1131	+ auto bShape = context_->GetInputShape(1)->GetOriginShape();
		1132	+ size_t aDims = aShape.GetDimNum();
		1133	+ size_t bDims = bShape.GetDimNum();
		1134	+
		1135	+ // 检查是否都为3-6维
		1136	+ const size_t MIN_DIM = 3;
		1137	+ const size_t MAX_DIM = 6;
		1138	+ if (aDims < MIN_DIM \|\| aDims > MAX_DIM \|\| bDims < MIN_DIM \|\| bDims > MAX_DIM) {
		1139	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A and B must have between %d and %d dimensions. "
		1140	+ "Bmm vector opt version not supported.", MIN_DIM, MAX_DIM);
		1141	+ return false;
		1142	+ }
		1143	+
		1144	+ // 检查B的最后一维是否为1（内轴，因为当n=1时B默认转置）
		1145	+ if (bShape.GetDim(bDims - 2) != 1) {
		1146	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: shape does not meet requirements. "
		1147	+ "Bmm vector opt version not supported.");
		1148	+ return false;
		1149	+ }
		1150	+
		1151	+ // 检查m和k是否至多为8
		1152	+ int64_t MAX_MK_DIM = 8;
		1153	+ if (aShape.GetDim(aDims - 1) > MAX_MK_DIM \|\| aShape.GetDim(aDims - 2) > MAX_MK_DIM) {
		1154	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A and B must have at most 8 on M and K axis. "
		1155	+ "Bmm vector opt version not supported.");
		1156	+ return false;
		1157	+ }
		1158	+
		1159	+ // 检查两个输入张量的维数是否相同
		1160	+ if (aDims != bDims) {
		1161	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A and B must have the same number of dimensions. "
		1162	+ "Bmm vector opt version not supported.");
		1163	+ return false;
		1164	+ }
		1165	+ return true;
		1166	+}
		1167	+
		1168	+bool BatchMatmulV3BaseTiling::CheckVectorDtypeAndKAxis()
		1169	+{
		1170	+ if (!(args_.aType == ge::DT_FLOAT && args_.bType == ge::DT_FLOAT && args_.cType == ge::DT_FLOAT)) {
		1171	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: input A, B, and output C must be fp32 data type. "
		1172	+ "Bmm vector opt version not supported.");
		1173	+ return false;
		1174	+ }
		1175	+
		1176	+ if (args_.aFormat != ge::FORMAT_ND \|\| args_.bFormat != ge::FORMAT_ND) {
		1177	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: input A and B must be ND format. "
		1178	+ "Bmm vector opt version not supported.");
		1179	+ return false;
		1180	+ }
		1181	+
		1182	+ // 检查A矩阵的K轴与B矩阵的K轴长度是否相同
		1183	+ auto aShape = context_->GetInputShape(0)->GetOriginShape();
		1184	+ auto bShape = context_->GetInputShape(1)->GetOriginShape();
		1185	+ size_t aDims = aShape.GetDimNum();
		1186	+ size_t bDims = bShape.GetDimNum();
		1187	+ if (aShape.GetDim(aDims - 1) != bShape.GetDim(bDims - 1)) {
		1188	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A and B must have the same length on K axis. "
		1189	+ "Bmm vector opt version not supported.");
		1190	+ return false;
		1191	+ }
		1192	+ return true;
		1193	+}
		1194	+
		1195	+bool BatchMatmulV3BaseTiling::CheckVectorBatchBroadcast()
		1196	+{
		1197	+ auto aShape = context_->GetInputShape(0)->GetOriginShape();
		1198	+ auto bShape = context_->GetInputShape(1)->GetOriginShape();
		1199	+ size_t aDims = aShape.GetDimNum();
		1200	+ for (size_t i = 0; i < aDims - 2; i++) {
		1201	+ if (aShape.GetDim(i) != bShape.GetDim(i)) {
		1202	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: A and B must have the same batch size. "
		1203	+ "Bmm vector opt version not supported.");
		1204	+ return false;
		1205	+ }
		1206	+ }
		1207	+ return true;
		1208	+}
		1209	+
		1210	+bool BatchMatmulV3BaseTiling::CheckVectorComputationCondition()
		1211	+{
		1212	+ if (!CheckVectorNpuArch()) { return false; }
		1213	+ if (!CheckVectorShapeDims()) { return false; }
		1214	+ if (!CheckVectorDtypeAndKAxis()) { return false; }
		1215	+ if (!CheckVectorBatchBroadcast()) { return false; }
		1216	+
		1217	+ OP_LOGD(args_.opName, "BatchMatmulV3BaseTiling: vector tiling condition check passed, "
		1218	+ "enter bmm vector opt version");
		1219	+ return true;
		1220	+}
		1221	+
997	ge::graphStatus BatchMatmulV3BaseTiling::PostTiling()	1222	ge::graphStatus BatchMatmulV3BaseTiling::PostTiling()
998	{	1223	{
999	size_t tilingDataSize = sizeof(BatchMatmulTilingData);	1224	size_t tilingDataSize = sizeof(BatchMatmulTilingData);
	@@ -1006,7 +1231,11 @@ ge::graphStatus BatchMatmulV3BaseTiling::PostTiling()
1006	return ge::GRAPH_FAILED;	1231	return ge::GRAPH_FAILED;
1007	}	1232	}
1008	context_->GetRawTilingData()->SetDataSize(tilingDataSize);	1233	context_->GetRawTilingData()->SetDataSize(tilingDataSize);
1009	- context_->SetBlockDim(compileInfo_.aicNum);	1234	+ if (isVectorMode_) {
		1235	+ context_->SetBlockDim(bmmTilingData_.vectorTilingInfo.coreNumber);
		1236	+ } else {
		1237	+ context_->SetBlockDim(compileInfo_.aicNum);
		1238	+ }
1010	auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());	1239	auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->GetPlatformInfo());
1011	auto npuArch = ascendcPlatform.GetCurNpuArch();	1240	auto npuArch = ascendcPlatform.GetCurNpuArch();
1012	if (( (npuArch == NpuArch::DAV_2201) \|\| (npuArch == NpuArch::DAV_3003) ) &&	1241	if (( (npuArch == NpuArch::DAV_2201) \|\| (npuArch == NpuArch::DAV_3003) ) &&

Mmatmul/batch_mat_mul_v3/op_host/op_tiling/batch_mat_mul_v3_base_tiling.h+43-0

	@@ -66,6 +66,7 @@ enum class TilingEnableLoadMode : int32_t // 互斥flag, 对应不同全载模
66	BASE = 0,	66	BASE = 0,
67	AL1_FULL_LOAD = 1,	67	AL1_FULL_LOAD = 1,
68	BL1_FULL_LOAD = 2,	68	BL1_FULL_LOAD = 2,
		69	+ VECTOR_FULL_LOAD = 3,
69	MAX = 10 //模板类别不能超过10个	70	MAX = 10 //模板类别不能超过10个
70	};	71	};
71		72
	@@ -121,6 +122,10 @@ protected:
121	bool GetBatchInfo();	122	bool GetBatchInfo();
122	bool GetBiasWithBatchInfo();	123	bool GetBiasWithBatchInfo();
123	void MergeBatchAndMAxis();	124	void MergeBatchAndMAxis();
		125	+ void SetBatchDimInfo();
		126	+ void CalcBatchDimAll();
		127	+ bool CheckNd2NzOnTheFlyLimit();
		128	+ void DoMultiBatchAndL1FullLoadTiling();
124	bool CheckBMMTilingDataIsVaild() const;	129	bool CheckBMMTilingDataIsVaild() const;
125	void DoTilingKeyCustom();	130	void DoTilingKeyCustom();
126	void DoUnAlignCommonTiling();	131	void DoUnAlignCommonTiling();
	@@ -137,6 +142,43 @@ protected:
137	void CalculateNd2nzWorkspaceSize();	142	void CalculateNd2nzWorkspaceSize();
138	void CheckandSetDiagonalConflict(uint64_t mCnt, uint64_t nCnt, uint64_t batch, uint64_t usedCoreNum, uint64_t transConflict, uint64_t newMcnt);	143	void CheckandSetDiagonalConflict(uint64_t mCnt, uint64_t nCnt, uint64_t batch, uint64_t usedCoreNum, uint64_t transConflict, uint64_t newMcnt);
139	void DoL2CacheAndCalOrderTiling();	144	void DoL2CacheAndCalOrderTiling();
		145	+ // 计算vector tiling参数并填充到vectorTilingInfo
		146	+ // vector计算使用AIV核心和UB空间，与cube计算的tiling逻辑不同
		147	+ void DoVectorTiling();
		148	+ // 检查是否满足vector计算条件
		149	+ // vector计算适用于：
		150	+ // 1. x1（A矩阵）：形状为4-6维，且可以与x2互相广播
		151	+ // 2. x2（B矩阵）：形状为4-6维，最后一维为1，且可以与x1互相广播
		152	+ // 3. 两个输入张量的维数相同，且各维度可以互相广播
		153	+ // 4. 输入和输出的数据类型均为fp32
		154	+ // 5. A矩阵的最后一维（K轴）与B矩阵的倒数第二维（K轴）长度相同，满足矩阵乘法约束
		155	+ bool CheckVectorComputationCondition();
		156	+
		157	+private:
		158	+ // CheckVectorComputationCondition的子检查函数
		159	+ bool CheckVectorNpuArch();
		160	+ bool CheckVectorShapeDims();
		161	+ bool CheckVectorDtypeAndKAxis();
		162	+ bool CheckVectorBatchBroadcast();
		163	+ // DoVectorTiling的子计算函数
		164	+ struct VectorShapeInfo {
		165	+ uint64_t coreNumber;
		166	+ uint64_t aTotalSize;
		167	+ uint64_t bTotalSize;
		168	+ uint64_t dimSizeSecondLast;
		169	+ uint64_t dimSizeLast;
		170	+ uint64_t batchSize;
		171	+ uint64_t bRowsPerBatch;
		172	+ };
		173	+ bool CalcVectorShapeInfo(VectorShapeInfo &shapeInfo);
		174	+ struct VectorCoreParams {
		175	+ uint64_t coreData;
		176	+ uint64_t rowsPerCore;
		177	+ uint64_t alignedDimSizeLast;
		178	+ };
		179	+ bool CalcVectorCoreParams(const VectorShapeInfo &shapeInfo, VectorCoreParams &coreParams);
		180	+ void SetVectorTilingParams(const VectorShapeInfo &shapeInfo, const VectorCoreParams &coreParams);
		181	+
140	protected:	182	protected:
141	BatchShapeInfo batchInfo_;	183	BatchShapeInfo batchInfo_;
142	BatchMatmulTilingData &bmmTilingData_;	184	BatchMatmulTilingData &bmmTilingData_;
	@@ -148,6 +190,7 @@ private:
148	uint64_t cBatchDimAll_{1};	190	uint64_t cBatchDimAll_{1};
149	protected:	191	protected:
150	TilingEnable tilingEnable_;	192	TilingEnable tilingEnable_;
		193	+ bool isVectorMode_ = false;
151	};	194	};
152	}	195	}
153	}	196	}

Mmatmul/batch_mat_mul_v3/op_kernel/batch_mat_mul_v3.cpp+19-0

	@@ -14,6 +14,7 @@
14	*/	14	*/
15	#include "batch_mat_mul_v3.h"	15	#include "batch_mat_mul_v3.h"
16	#include "batch_mat_mul_v3_tiling_key.h"	16	#include "batch_mat_mul_v3_tiling_key.h"
		17	+#include "batch_mat_mul_v3_vector.h"
17		18
18	using namespace AscendC;	19	using namespace AscendC;
19	using namespace matmul;	20	using namespace matmul;
	@@ -137,6 +138,17 @@ constexpr CubeFormat format_y = CubeFormat::ND;
137	op.Process(); \	138	op.Process(); \
138	} while (0)	139	} while (0)
139		140
		141	+#define BMMV3_IMPL_VECTOR_CLASS(templateClass) \
		142	+ do { \
		143	+ using aType = MatmulType<AscendC::TPosition::GM, format_x1, float, false, LayoutMode::NORMAL>; \
		144	+ using bType = MatmulType<AscendC::TPosition::GM, format_x2, float, false, LayoutMode::NORMAL>; \
		145	+ using cType = MatmulType<AscendC::TPosition::GM, format_y, float, false, LayoutMode::NORMAL>; \
		146	+ TPipe pipe; \
		147	+ templateClass<aType, bType, cType> op; \
		148	+ op.Init(aGM, bGM, cGM, biasGM, offsetWGM, workspaceGM, &tilingData, &pipe); \
		149	+ op.Process(); \
		150	+ } while (0)
		151	+
140	template<int MULTIBATCHL1FULLLOAD, int MULTIBATCH, int LOADMODE, int ISMULTIBATCHOUT, int MIXND2NZ>	152	template<int MULTIBATCHL1FULLLOAD, int MULTIBATCH, int LOADMODE, int ISMULTIBATCHOUT, int MIXND2NZ>
141	__global__ __aicore__ void batch_mat_mul_v3(	153	__global__ __aicore__ void batch_mat_mul_v3(
142	GM_ADDR aGM, GM_ADDR bGM, GM_ADDR biasGM, GM_ADDR offsetWGM, GM_ADDR cGM, GM_ADDR workspaceGM, GM_ADDR tilingGM)	154	GM_ADDR aGM, GM_ADDR bGM, GM_ADDR biasGM, GM_ADDR offsetWGM, GM_ADDR cGM, GM_ADDR workspaceGM, GM_ADDR tilingGM)
	@@ -194,6 +206,13 @@ __global__ __aicore__ void batch_mat_mul_v3(
194	MULTIBATCH == BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE && LOADMODE == BATCH_MAT_MUL_V3_BASE_FULLLOAD &&	206	MULTIBATCH == BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE && LOADMODE == BATCH_MAT_MUL_V3_BASE_FULLLOAD &&
195	ISMULTIBATCHOUT == BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE && MIXND2NZ == BATCH_MAT_MUL_V3_MIXND2NZ_TRUE) {	207	ISMULTIBATCHOUT == BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE && MIXND2NZ == BATCH_MAT_MUL_V3_MIXND2NZ_TRUE) {
196	BMMV3_IMPL_CLASS(BatchMatMulUnalignedMultiBatchKernel, BatchMatMulUnalignedMultiBatchBaseBlock, MM_CFG_MULTI_BATCH_OUT);	208	BMMV3_IMPL_CLASS(BatchMatMulUnalignedMultiBatchKernel, BatchMatMulUnalignedMultiBatchBaseBlock, MM_CFG_MULTI_BATCH_OUT);
		209	+ } else if constexpr (
		210	+ MULTIBATCHL1FULLLOAD == BATCH_MAT_MUL_V3_MULTI_BATCH_L1_FULLLOAD_FALSE &&
		211	+ MULTIBATCH == BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE && LOADMODE == BATCH_MAT_MUL_V3_VECTOR_FULLLOAD &&
		212	+ ISMULTIBATCHOUT == BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE && MIXND2NZ == BATCH_MAT_MUL_V3_MIXND2NZ_FALSE &&
		213	+ format_x1 == CubeFormat::ND && format_x2 == CubeFormat::ND &&
		214	+ std::is_same_v<DTYPE_X1, float> && std::is_same_v<DTYPE_X2, float> && std::is_same_v<DTYPE_Y, float>) {
		215	+ BMMV3_IMPL_VECTOR_CLASS(BatchMatmulVectorKernel);
197	#if defined(ORIG_DTYPE_X1) && ORIG_DTYPE_X1 == DT_FLOAT	216	#if defined(ORIG_DTYPE_X1) && ORIG_DTYPE_X1 == DT_FLOAT
198	} else if constexpr (	217	} else if constexpr (
199	MULTIBATCHL1FULLLOAD == BATCH_MAT_MUL_V3_MULTI_BATCH_L1_FULLLOAD_TRUE &&	218	MULTIBATCHL1FULLLOAD == BATCH_MAT_MUL_V3_MULTI_BATCH_L1_FULLLOAD_TRUE &&

Mmatmul/batch_mat_mul_v3/op_kernel/batch_mat_mul_v3_tiling_data.h+20-1

	@@ -54,10 +54,29 @@ struct alignas(8) MultiBatchInfo{
54	#pragma pack(pop)	54	#pragma pack(pop)
55		55
56	#pragma pack(push, 8)	56	#pragma pack(push, 8)
57	-// 8 means 8 bytes aligned	57	+struct alignas(8) VectorTilingInfo {
		58	+ uint64_t coreNumber;
		59	+ uint64_t coreData;
		60	+ uint64_t copyLoop;
		61	+ uint64_t copyTail;
		62	+ uint64_t lastCopyLoop;
		63	+ uint64_t lastCopyTail;
		64	+ uint64_t rowsPerCore;
		65	+ uint64_t aTotalSize;
		66	+ uint64_t bTotalSize;
		67	+ uint64_t cTotalSize;
		68	+ uint64_t dimSizeSecondLast;
		69	+ uint64_t dimSizeLast;
		70	+ uint64_t alignedDimSizeLast;
		71	+ uint64_t bRowsPerBatch;
		72	+};
		73	+#pragma pack(pop)
		74	+
		75	+#pragma pack(push, 8)
58	struct alignas(8) BatchMatmulTilingData{	76	struct alignas(8) BatchMatmulTilingData{
59	MatmulTilingData matmulTiling;	77	MatmulTilingData matmulTiling;
60	MultiBatchInfo multiBatchInfo;	78	MultiBatchInfo multiBatchInfo;
		79	+ VectorTilingInfo vectorTilingInfo;
61	};	80	};
62	#pragma pack(pop)	81	#pragma pack(pop)
63		82

Mmatmul/batch_mat_mul_v3/op_kernel/batch_mat_mul_v3_tiling_key.h+10-1

	@@ -27,6 +27,7 @@
27	#define BATCH_MAT_MUL_V3_BASE_FULLLOAD 0	27	#define BATCH_MAT_MUL_V3_BASE_FULLLOAD 0
28	#define BATCH_MAT_MUL_V3_AL1_FULLLOAD 1	28	#define BATCH_MAT_MUL_V3_AL1_FULLLOAD 1
29	#define BATCH_MAT_MUL_V3_BL1_FULLLOAD 2	29	#define BATCH_MAT_MUL_V3_BL1_FULLLOAD 2
		30	+#define BATCH_MAT_MUL_V3_VECTOR_FULLLOAD 3
30		31
31	#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE 0	32	#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE 0
32	#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE 1	33	#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE 1
	@@ -44,7 +45,7 @@ ASCENDC_TPL_ARGS_DECL(
44	BATCH_MAT_MUL_V3_MULTI_BATCH_FALSE, BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE),	45	BATCH_MAT_MUL_V3_MULTI_BATCH_FALSE, BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE),
45	ASCENDC_TPL_UINT_DECL(	46	ASCENDC_TPL_UINT_DECL(
46	LOADMODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST,	47	LOADMODE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST,
47	- BATCH_MAT_MUL_V3_BASE_FULLLOAD, BATCH_MAT_MUL_V3_AL1_FULLLOAD, BATCH_MAT_MUL_V3_BL1_FULLLOAD),	48	+ BATCH_MAT_MUL_V3_BASE_FULLLOAD, BATCH_MAT_MUL_V3_AL1_FULLLOAD, BATCH_MAT_MUL_V3_BL1_FULLLOAD, BATCH_MAT_MUL_V3_VECTOR_FULLLOAD),
48	ASCENDC_TPL_UINT_DECL(	49	ASCENDC_TPL_UINT_DECL(
49	ISMULTIBATCHOUT, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST,	50	ISMULTIBATCHOUT, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST,
50	BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE, BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE),	51	BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE, BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE),
	@@ -125,6 +126,14 @@ ASCENDC_TPL_SEL(
125	ASCENDC_TPL_UINT_SEL(LOADMODE, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_BASE_FULLLOAD),	126	ASCENDC_TPL_UINT_SEL(LOADMODE, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_BASE_FULLLOAD),
126	ASCENDC_TPL_UINT_SEL(ISMULTIBATCHOUT, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE),	127	ASCENDC_TPL_UINT_SEL(ISMULTIBATCHOUT, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE),
127	ASCENDC_TPL_UINT_SEL(MIXND2NZ, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_MIXND2NZ_FALSE), ),	128	ASCENDC_TPL_UINT_SEL(MIXND2NZ, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_MIXND2NZ_FALSE), ),
		129	+ /* VECTOR_FULLLOAD */
		130	+ ASCENDC_TPL_ARGS_SEL(
		131	+ ASCENDC_TPL_KERNEL_TYPE_SEL(ASCENDC_TPL_AIV_ONLY),
		132	+ ASCENDC_TPL_UINT_SEL(MULTIBATCHL1FULLLOAD, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_MULTI_BATCH_L1_FULLLOAD_FALSE),
		133	+ ASCENDC_TPL_UINT_SEL(MULTIBATCH, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_MULTI_BATCH_TRUE),
		134	+ ASCENDC_TPL_UINT_SEL(LOADMODE, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_VECTOR_FULLLOAD),
		135	+ ASCENDC_TPL_UINT_SEL(ISMULTIBATCHOUT, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE),
		136	+ ASCENDC_TPL_UINT_SEL(MIXND2NZ, ASCENDC_TPL_UI_LIST, BATCH_MAT_MUL_V3_MIXND2NZ_FALSE), ),
128	#if defined(ORIG_DTYPE_X1) && ORIG_DTYPE_X1 == DT_FLOAT	137	#if defined(ORIG_DTYPE_X1) && ORIG_DTYPE_X1 == DT_FLOAT
129	/* MultiBatchL1_FULLLOAD */	138	/* MultiBatchL1_FULLLOAD */
130	ASCENDC_TPL_ARGS_SEL(	139	ASCENDC_TPL_ARGS_SEL(

Amatmul/batch_mat_mul_v3/op_kernel/batch_mat_mul_v3_vector.h+213-0

@@ -0,0 +1,213 @@
	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 batch_mat_mul_v3_vector.h
	13	+ * \brief vector kernel for batch_mat_mul_v3, used when B matrix last dim is 1, actual shape [1, k]
	14	+ * forcefully transposed before entering this kernel
	15	+ */
	16	+#ifndef BATCH_MAT_MUL_V3_VECTOR_H
	17	+#define BATCH_MAT_MUL_V3_VECTOR_H
	18	+
	19	+#include "kernel_operator.h"
	20	+#include "lib/matmul_intf.h"
	21	+
	22	+using namespace AscendC;
	23	+
	24	+template <typename A_TYPE, typename B_TYPE, typename C_TYPE>
	25	+class BatchMatmulVectorKernel {
	26	+public:
	27	+ __aicore__ inline BatchMatmulVectorKernel() {}
	28	+ using A_T = typename A_TYPE::T;
	29	+ using B_T = typename B_TYPE::T;
	30	+ using C_T = typename C_TYPE::T;
	31	+
	32	+ TBuf<TPosition::VECIN> aBuf, bBuf;
	33	+ TBuf<TPosition::VECOUT> cBuf;
	34	+ TBuf<TPosition::VECCALC> sharedTmpBuf, mulBuf;
	35	+ GlobalTensor<A_T> aGm;
	36	+ GlobalTensor<B_T> bGm;
	37	+ GlobalTensor<C_T> cGm;
	38	+ uint64_t coreNumber;
	39	+ uint64_t coreData;
	40	+ uint64_t rowsPerCore;
	41	+ uint64_t aTotalSize;
	42	+ uint64_t bTotalSize;
	43	+ uint64_t dimSizeSecondLast;
	44	+ uint64_t dimSizeLast;
	45	+ uint64_t alignedDimSizeLast;
	46	+ uint64_t bRowsPerBatch;
	47	+ uint64_t cTotalSize;
	48	+ LocalTensor<A_T> aUb;
	49	+ LocalTensor<B_T> bUb;
	50	+ LocalTensor<C_T> mulUb;
	51	+ LocalTensor<C_T> cUb;
	52	+ LocalTensor<C_T> sharedTmpUb;
	53	+
	54	+ DataCopyPadExtParams<A_T> padParamsA;
	55	+ DataCopyPadExtParams<B_T> padParamsB;
	56	+
	57	+ __aicore__ inline void Init(GM_ADDR aGM,
	58	+ GM_ADDR bGM,
	59	+ GM_ADDR cGM,
	60	+ GM_ADDR biasGM,
	61	+ GM_ADDR offsetWGM,
	62	+ GM_ADDR workspaceGM,
	63	+ BatchMatmulTilingData* tilingData, TPipe* pipe)
	64	+ {
	65	+ ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
	66	+ this->coreNumber = tilingData->vectorTilingInfo.coreNumber;
	67	+ this->coreData = tilingData->vectorTilingInfo.coreData;
	68	+ this->rowsPerCore = tilingData->vectorTilingInfo.rowsPerCore;
	69	+ this->aTotalSize = tilingData->vectorTilingInfo.aTotalSize;
	70	+ this->bTotalSize = tilingData->vectorTilingInfo.bTotalSize;
	71	+ this->cTotalSize = tilingData->vectorTilingInfo.cTotalSize;
	72	+ this->dimSizeSecondLast = tilingData->vectorTilingInfo.dimSizeSecondLast;
	73	+ this->dimSizeLast = tilingData->vectorTilingInfo.dimSizeLast;
	74	+ this->alignedDimSizeLast = tilingData->vectorTilingInfo.alignedDimSizeLast;
	75	+ this->bRowsPerBatch = tilingData->vectorTilingInfo.bRowsPerBatch;
	76	+
	77	+ uint8_t rightPad = static_cast<uint8_t>(this->alignedDimSizeLast - this->dimSizeLast);
	78	+ this->padParamsA = DataCopyPadExtParams<A_T>{true, 0, rightPad, static_cast<A_T>(0)};
	79	+ this->padParamsB = DataCopyPadExtParams<B_T>{true, 0, rightPad, static_cast<B_T>(0)};
	80	+
	81	+ aGm.SetGlobalBuffer((__gm__ A_T*)aGM, this->aTotalSize);
	82	+ bGm.SetGlobalBuffer((__gm__ B_T*)bGM, this->bTotalSize);
	83	+ cGm.SetGlobalBuffer((__gm__ C_T*)cGM, this->cTotalSize);
	84	+
	85	+ pipe->InitBuffer(aBuf, this->rowsPerCore * alignedDimSizeLast * sizeof(A_T));
	86	+ pipe->InitBuffer(bBuf, this->rowsPerCore * alignedDimSizeLast * sizeof(B_T));
	87	+ pipe->InitBuffer(mulBuf, this->rowsPerCore * alignedDimSizeLast * sizeof(C_T));
	88	+ pipe->InitBuffer(sharedTmpBuf, this->rowsPerCore * alignedDimSizeLast * sizeof(C_T));
	89	+ pipe->InitBuffer(cBuf, this->rowsPerCore * sizeof(C_T));
	90	+ }
	91	+
	92	+ __aicore__ inline void Process()
	93	+ {
	94	+ uint32_t coreId = GetBlockIdx();
	95	+ if (coreId >= this->coreNumber) {
	96	+ return;
	97	+ }
	98	+ uint64_t startAddress = coreId * this->coreData;
	99	+ if (startAddress >= this->cTotalSize) {
	100	+ return;
	101	+ }
	102	+ uint64_t totalRows = this->coreData;
	103	+ // 最后一个核处理剩余行
	104	+ if (coreId == this->coreNumber - 1) {
	105	+ totalRows = this->cTotalSize - startAddress;
	106	+ }
	107	+
	108	+ // 按照rowsPerCore分块处理，每次迭代计算rowsPerCore行，即ub可操作的最大行数
	109	+ uint64_t fullIters = totalRows / this->rowsPerCore;
	110	+ uint64_t tailRows = totalRows % this->rowsPerCore;
	111	+ for (uint64_t i = 0; i < fullIters; i++) {
	112	+ uint64_t address = startAddress + i * this->rowsPerCore;
	113	+ indicesCompute(this->rowsPerCore, address);
	114	+ }
	115	+ // 处理尾块计算
	116	+ if (tailRows != 0) {
	117	+ uint64_t address = startAddress + fullIters * this->rowsPerCore;
	118	+ indicesCompute(tailRows, address);
	119	+ }
	120	+ }
	121	+
	122	+private:
	123	+ __aicore__ inline void indicesCompute(int32_t tensorSize, uint64_t address)
	124	+ {
	125	+ aUb = aBuf.Get<A_T>();
	126	+ bUb = bBuf.Get<B_T>();
	127	+ mulUb = mulBuf.Get<C_T>();
	128	+ cUb = cBuf.Get<C_T>();
	129	+ sharedTmpUb = sharedTmpBuf.Get<C_T>();
	130	+
	131	+ // A: 一条DataCopyPad拷贝所有行
	132	+ // rightPadding已将每个block扩展到alignedDimSizeLast，dstStride=0即可
	133	+ DataCopyExtParams copyParamsA{1, (uint32_t)(dimSizeLast * sizeof(A_T)), 0, 0, 0};
	134	+ copyParamsA.blockCount = tensorSize;
	135	+ DataCopyPad(aUb, aGm[address * dimSizeLast], copyParamsA, padParamsA);
	136	+
	137	+ CopyAndBroadcastB(tensorSize, address);
	138	+ BmmByVec(tensorSize, address);
	139	+ }
	140	+
	141	+ /**
	142	+ * @brief 复制并广播B向量: (1, dimSizeLast) -> (tensorSize, dimSizeLast)
	143	+ *
	144	+ * @param tensorSize 单核内矩阵行数
	145	+ * @param address 矩阵起始地址
	146	+ */
	147	+ __aicore__ inline void CopyAndBroadcastB(int32_t tensorSize, uint64_t address)
	148	+ {
	149	+ // 计算需要读取的B矩阵batch范围
	150	+ uint64_t firstBatchIdx = address / dimSizeSecondLast;
	151	+ uint64_t lastBatchIdx = (address + tensorSize - 1) / dimSizeSecondLast;
	152	+ int32_t numBatches = static_cast<int32_t>(lastBatchIdx - firstBatchIdx + 1);
	153	+
	154	+ // Step 1: 一次性从GM拷贝所有B行到mulUb临时空间
	155	+ // B数据在GM中连续存储，每个batch一行(dimSizeLast个元素)
	156	+ DataCopyExtParams copyParamsB{static_cast<uint16_t>(numBatches),
	157	+ (uint32_t)(dimSizeLast * sizeof(B_T)), 0, 0, 0};
	158	+ auto tempUb = mulUb.template ReinterpretCast<B_T>();
	159	+ DataCopyPad(tempUb, bGm[firstBatchIdx * bRowsPerBatch], copyParamsB, padParamsB);
	160	+ PipeBarrier<PIPE_ALL>();
	161	+
	162	+ // Step 2: 从临时空间广播到bUb目标位置
	163	+ // 负责广播每个batch的单行到bUb目标位置
	164	+ // BLOCK_BYTE_SIZE = 32
	165	+ constexpr uint32_t elementsPerBlock = BLOCK_BYTE_SIZE / sizeof(B_T);
	166	+ for (int32_t i = 0; i < tensorSize; ) {
	167	+ uint64_t curAddress = address + i;
	168	+ uint64_t bBatchIdx = curAddress / dimSizeSecondLast;
	169	+ uint64_t batchEndRow = (bBatchIdx + 1) * dimSizeSecondLast;
	170	+ int32_t batchCount = (tensorSize - i < static_cast<int32_t>(batchEndRow - address - i))
	171	+ ? (tensorSize - i) : static_cast<int32_t>(batchEndRow - address - i);
	172	+ int32_t localBatchIdx = static_cast<int32_t>(bBatchIdx - firstBatchIdx);
	173	+
	174	+ // 广播: 从tempUb[localBatchIdx行] 到 bUb[i行..i+batchCount-1行]
	175	+ uint64_t mask = alignedDimSizeLast;
	176	+ CopyRepeatParams copyRepeatParams;
	177	+ copyRepeatParams.srcStride = 0;
	178	+ copyRepeatParams.dstStride = 0;
	179	+ copyRepeatParams.srcRepeatSize = 0; // 源不前进
	180	+ copyRepeatParams.dstRepeatSize = alignedDimSizeLast / elementsPerBlock;
	181	+ Copy(bUb[i * alignedDimSizeLast], tempUb[localBatchIdx * alignedDimSizeLast],
	182	+ mask, static_cast<uint8_t>(batchCount), copyRepeatParams);
	183	+ i += batchCount;
	184	+ }
	185	+ PipeBarrier<PIPE_ALL>();
	186	+ }
	187	+
	188	+ /**
	189	+ * @brief 矩阵乘法: vector核内进行逐元素相乘，然后沿dimSizeLast维度进行归一化计算reduce sum
	190	+ *
	191	+ * @param tensorSize 单核内矩阵行数
	192	+ * @param address 矩阵起始地址
	193	+ */
	194	+ __aicore__ inline void BmmByVec(int32_t tensorSize, uint64_t address)
	195	+ {
	196	+ // 1. 逐元素相乘
	197	+ Mul(mulUb, aUb, bUb, tensorSize * alignedDimSizeLast);
	198	+ PipeBarrier<PIPE_ALL>();
	199	+ uint32_t srcShape[2] = {static_cast<uint32_t>(tensorSize),
	200	+ static_cast<uint32_t>(alignedDimSizeLast)};
	201	+
	202	+ // 2. 沿dimSizeLast维度进行归一化计算reduce sum
	203	+ ReduceSum<C_T, AscendC::Pattern::Reduce::AR, true>(cUb, mulUb,
	204	+ sharedTmpUb.template ReinterpretCast<uint8_t>(), srcShape, true);
	205	+ PipeBarrier<PIPE_ALL>();
	206	+
	207	+ DataCopyExtParams copyParamsC{1, (uint32_t)(tensorSize * sizeof(C_T)), 0, 0, 0};
	208	+ DataCopyPad(cGm[address], cUb, copyParamsC);
	209	+ PipeBarrier<PIPE_ALL>();
	210	+ }
	211	+};
	212	+
	213	+#endif // BATCH_MAT_MUL_V3_VECTOR_H

Amatmul/batch_mat_mul_v3/tests/ut/op_host/test_aclnn_batch_matmul_vector_compute.cpp+88-0

@@ -0,0 +1,88 @@
	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 <array>
	12	+#include <vector>
	13	+#include "gtest/gtest.h"
	14	+
	15	+#include "../../../op_host/op_api/aclnn_batch_matmul.h"
	16	+#include "opdev/platform.h"
	17	+#include "op_api_ut_common/op_api_ut.h"
	18	+#include "op_api_ut_common/scalar_desc.h"
	19	+#include "op_api_ut_common/tensor_desc.h"
	20	+
	21	+using namespace std;
	22	+using namespace op;
	23	+
	24	+class l2_batch_matmul_vector_compute_test : public testing::Test
	25	+{
	26	+protected:
	27	+ static void SetUpTestCase()
	28	+ {
	29	+ cout << "batch_matmul_vector_compute_test SetUp" << endl;
	30	+ }
	31	+
	32	+ static void TearDownTestCase()
	33	+ {
	34	+ cout << "batch_matmul_vector_compute_test TearDown" << endl;
	35	+ }
	36	+};
	37	+
	38	+TEST_F(l2_batch_matmul_vector_compute_test, ascend910b_vector_case_1)
	39	+{
	40	+ auto self_desc = TensorDesc({12200, 4, 4}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	41	+ auto mat2_desc = TensorDesc({12200, 4, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	42	+ auto out_desc = TensorDesc({12200, 4, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2).Precision(0.005, 0.005);
	43	+ int8_t cube_math_type = 1;
	44	+ auto ut = OP_API_UT(aclnnBatchMatMul, INPUT(self_desc, mat2_desc), OUTPUT(out_desc), cube_math_type);
	45	+
	46	+ uint64_t workspace_size = 0;
	47	+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
	48	+ EXPECT_EQ(aclRet, ACL_SUCCESS);
	49	+}
	50	+
	51	+TEST_F(l2_batch_matmul_vector_compute_test, ascend910b_vector_case_2)
	52	+{
	53	+ auto self_desc = TensorDesc({220, 4, 4}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	54	+ auto mat2_desc = TensorDesc({220, 4, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	55	+ auto out_desc = TensorDesc({220, 4, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2).Precision(0.005, 0.005);
	56	+ int8_t cube_math_type = 1;
	57	+ auto ut = OP_API_UT(aclnnBatchMatMul, INPUT(self_desc, mat2_desc), OUTPUT(out_desc), cube_math_type);
	58	+
	59	+ uint64_t workspace_size = 0;
	60	+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
	61	+ EXPECT_EQ(aclRet, ACL_SUCCESS);
	62	+}
	63	+
	64	+TEST_F(l2_batch_matmul_vector_compute_test, ascend910b_vector_case_3)
	65	+{
	66	+ auto self_desc = TensorDesc({12200, 3, 3}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	67	+ auto mat2_desc = TensorDesc({12200, 3, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	68	+ auto out_desc = TensorDesc({12200, 3, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2).Precision(0.005, 0.005);
	69	+ int8_t cube_math_type = 1;
	70	+ auto ut = OP_API_UT(aclnnBatchMatMul, INPUT(self_desc, mat2_desc), OUTPUT(out_desc), cube_math_type);
	71	+
	72	+ uint64_t workspace_size = 0;
	73	+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
	74	+ EXPECT_EQ(aclRet, ACL_SUCCESS);
	75	+}
	76	+
	77	+TEST_F(l2_batch_matmul_vector_compute_test, ascend910b_vector_case_4)
	78	+{
	79	+ auto self_desc = TensorDesc({220, 3, 3}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	80	+ auto mat2_desc = TensorDesc({220, 3, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2);
	81	+ auto out_desc = TensorDesc({220, 3, 1}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-2, 2).Precision(0.005, 0.005);
	82	+ int8_t cube_math_type = 1;
	83	+ auto ut = OP_API_UT(aclnnBatchMatMul, INPUT(self_desc, mat2_desc), OUTPUT(out_desc), cube_math_type);
	84	+
	85	+ uint64_t workspace_size = 0;
	86	+ aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspace_size);
	87	+ EXPECT_EQ(aclRet, ACL_SUCCESS);
	88	+}

Mmatmul/batch_mat_mul_v3/tests/ut/op_host/test_batch_mat_mul_v3_tiling.cpp+5-4

	@@ -36,8 +36,9 @@ string get_map_string(const std::map<string, string>& map, const string& key) {
36	}	36	}
37	bool IsDisplayTilingdata(const string& case_name, size_t index, uint64_t tilingKey)	37	bool IsDisplayTilingdata(const string& case_name, size_t index, uint64_t tilingKey)
38	{	38	{
39	- // 0-18 22-27 30-32 48之后表示bmm实际用到的tilingdata	39	+ // 0-18 22-27 30-32 48-91 表示bmm实际用到的tilingdata(不含VectorTilingInfo)
40	- if (index < 18 \|\| (index >= 22 && index <= 27) \|\| (index >= 30 && index <= 32) \|\| index >= 48) {	40	+ // VectorTilingInfo从index 92开始(sizeof(MatmulTilingData)+sizeof(MultiBatchInfo)=368字节=92个int32)
		41	+ if (index < 18 \|\| (index >= 22 && index <= 27) \|\| (index >= 30 && index <= 32) \|\| (index >= 48 && index < 92)) {
41	return true;	42	return true;
42	}	43	}
43	// 基础API校验全部的tilingdata	44	// 基础API校验全部的tilingdata
	@@ -379,7 +380,7 @@ static TilingTestParam ascend910B_cases_params[] = {
379	"hardware_info": {"BT_SIZE": 1024, "load3d_constraints": "unknown", "Intrinsic_fix_pipe_l0c2out": true, "Intrinsic_data_move_l12ub": false, "Intrinsic_data_move_l0c2ub": false, "Intrinsic_data_move_out2l1_nd2nz": true, "UB_SIZE": 196608, "L2_SIZE": 201326592, "L1_SIZE": 524288, "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, "CORE_NUM": 24, "vector_core_cnt": 48, "socVersion": "Ascend910B" },	380	"hardware_info": {"BT_SIZE": 1024, "load3d_constraints": "unknown", "Intrinsic_fix_pipe_l0c2out": true, "Intrinsic_data_move_l12ub": false, "Intrinsic_data_move_l0c2ub": false, "Intrinsic_data_move_out2l1_nd2nz": true, "UB_SIZE": 196608, "L2_SIZE": 201326592, "L1_SIZE": 524288, "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, "CORE_NUM": 24, "vector_core_cnt": 48, "socVersion": "Ascend910B" },
380	"format_a":"ND","format_b":"ND","repo_range":{},"repo_seeds":{}})",	381	"format_a":"ND","format_b":"ND","repo_range":{},"repo_seeds":{}})",
381	ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, false, false, 0, false, {1500, 1, 512}, {1500, 512, 128}, {1500, 1, 128}, false, 0, 0, 24, 65537,	382	ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, false, false, 0, false, {1500, 1, 512}, {1500, 512, 128}, {1500, 1, 128}, false, 0, 0, 24, 65537,
382	- "24 1 128 512 512 1 128 512 16 128 64 64 8 1 1 0 0 0 0 65536 1024 0 1 1 1 1 32 4 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 24 1500 1500 1500 1 1 1 1 1 1 1 1 1 1500 1500 1500 0 0 1 1500 7 1 "	383	+ "24 1 128 512 512 1 128 512 16 128 64 64 8 1 1 0 0 0 0 65536 1024 0 1 1 1 1 32 4 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 24 1500 1500 1500 1 1 1 1 1 1 1 1 1 1500 1500 1500 0 0 1 1500 7 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 "
383	},	384	},
384	{"BatchMatMulV3_MultiBatch_AL1FullLoad_general_test_02",	385	{"BatchMatMulV3_MultiBatch_AL1FullLoad_general_test_02",
385	"BatchMatMulV3",	386	"BatchMatMulV3",
	@@ -389,7 +390,7 @@ static TilingTestParam ascend910B_cases_params[] = {
389	"hardware_info": {"BT_SIZE": 1024, "load3d_constraints": "unknown", "Intrinsic_fix_pipe_l0c2out": true, "Intrinsic_data_move_l12ub": false, "Intrinsic_data_move_l0c2ub": false, "Intrinsic_data_move_out2l1_nd2nz": true, "UB_SIZE": 196608, "L2_SIZE": 201326592, "L1_SIZE": 524288, "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, "CORE_NUM": 24, "vector_core_cnt": 48, "socVersion": "Ascend910B" },	390	"hardware_info": {"BT_SIZE": 1024, "load3d_constraints": "unknown", "Intrinsic_fix_pipe_l0c2out": true, "Intrinsic_data_move_l12ub": false, "Intrinsic_data_move_l0c2ub": false, "Intrinsic_data_move_out2l1_nd2nz": true, "UB_SIZE": 196608, "L2_SIZE": 201326592, "L1_SIZE": 524288, "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, "CORE_NUM": 24, "vector_core_cnt": 48, "socVersion": "Ascend910B" },
390	"format_a":"ND","format_b":"ND","repo_range":{},"repo_seeds":{}})",	391	"format_a":"ND","format_b":"ND","repo_range":{},"repo_seeds":{}})",
391	ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, false, true, 0, false, {1500, 1, 128}, {1500, 512, 128}, {1500, 1, 512}, false, 0, 0, 24, 65537,	392	ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, false, true, 0, false, {1500, 1, 128}, {1500, 512, 128}, {1500, 1, 512}, false, 0, 0, 24, 65537,
392	- "24 1 512 128 128 1 512 128 16 64 128 16 8 1 1 0 0 0 0 24576 2048 0 1 1 1 1 8 4 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 1 16 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 24 1500 1500 1500 1 1 1 1 1 1 1 1 1 1500 1500 1500 0 0 1 1500 31 1 "	393	+ "24 1 512 128 128 1 512 128 16 64 128 16 8 1 1 0 0 0 0 24576 2048 0 1 1 1 1 8 4 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 1 16 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 24 1500 1500 1500 1 1 1 1 1 1 1 1 1 1500 1500 1500 0 0 1 1500 31 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 "
393	}	394	}
394	};	395	};
395		396

Mmatmul/common/op_host/op_api/batch_matmul_util.cpp+13-0

	@@ -241,6 +241,19 @@ static bool CheckAscendCScenario(
241	return false;	241	return false;
242	}	242	}
243		243
		244	+ // vector kernel scenario: effective N dim of x2 is 1
		245	+ // 当nDim = 1时，触发vector kernel场景
		246	+ size_t x2DimNum = x2->GetViewShape().GetDimNum();
		247	+ int64_t nDim = x2->GetViewShape().GetDim(x2DimNum - 2);
		248	+ size_t x1DimNum = x1->GetViewShape().GetDimNum();
		249	+ int64_t mDim = x1->GetViewShape().GetDim(x1DimNum - 2);
		250	+ int64_t kDim = x1->GetViewShape().GetDim(x1DimNum - 1);
		251	+ int64_t MAX_MK_DIM = 8;
		252	+ if (npuArch == NpuArch::DAV_2201 && adjX2 == 1 && nDim == 1 && mDim <= MAX_MK_DIM && kDim <= MAX_MK_DIM) {
		253	+ OP_LOGI("Hit batch_mat_mul_v3 vector kernel scenario: effective N dimension is 1.");
		254	+ return true;
		255	+ }
		256	+
244	return Ops::NN::BmmCheckHitV3Shape(x1, x2, bias, adjX1, adjX2, mmOpInfo.support_info.self_format,	257	return Ops::NN::BmmCheckHitV3Shape(x1, x2, bias, adjX1, adjX2, mmOpInfo.support_info.self_format,
245	mmOpInfo.support_info.mat2_format, mmOpInfo.enableFp16Bf16InFp32Out);	258	mmOpInfo.support_info.mat2_format, mmOpInfo.enableFp16Bf16InFp32Out);
246	}	259	}