已合并
bmm支持小m,k轴且n轴=1模板 #6190
zhang-junming21创建于 6月17日
bmm支持小m,k轴且n轴=1模板 #6190
已合并
zhang-junming21创建于 6月17日
10 个文件变更+1204-19
@@ -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+}
@@ -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 
27using namespace optiling::batch_mat_mul_v3;28using namespace optiling::batch_mat_mul_v3;
28using Ops::NN::MathUtil;29using Ops::NN::MathUtil;
@@ -57,6 +58,9 @@ constexpr uint64_t L2_SIZE_2 = 192UL * 1024UL * 1024UL;
57constexpr uint64_t MAX_TRANS_CONFLICT = 6;58constexpr uint64_t MAX_TRANS_CONFLICT = 6;
58constexpr double TAIL_CONFLICT_RATIO = 0.5;59constexpr double TAIL_CONFLICT_RATIO = 0.5;
59constexpr uint64_t MAX_INT32_VALUE = 2147483647uL;60constexpr 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 
61static inline uint64_t LastPower2(uint64_t n)65static 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+ 
997ge::graphStatus BatchMatmulV3BaseTiling::PostTiling()1222ge::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) ) &&
@@ -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+ 
140protected:182protected:
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};
149protected:191protected:
150 TilingEnable tilingEnable_;192 TilingEnable tilingEnable_;
193+ bool isVectorMode_ = false;
151};194};
152}195}
153}196}
@@ -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 
18using namespace AscendC;19using namespace AscendC;
19using namespace matmul;20using 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+ 
140template<int MULTIBATCHL1FULLLOAD, int MULTIBATCH, int LOADMODE, int ISMULTIBATCHOUT, int MIXND2NZ>152template<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_FLOAT216#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 &&
@@ -53,11 +53,31 @@ struct alignas(8) MultiBatchInfo{
53};53};
54#pragma pack(pop)54#pragma pack(pop)
55 55 
56+#pragma pack(push, 8)
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+ 
56#pragma pack(push, 8)75#pragma pack(push, 8)
57// 8 means 8 bytes aligned76// 8 means 8 bytes aligned
58struct alignas(8) BatchMatmulTilingData{77struct alignas(8) BatchMatmulTilingData{
59 MatmulTilingData matmulTiling;78 MatmulTilingData matmulTiling;
60 MultiBatchInfo multiBatchInfo;79 MultiBatchInfo multiBatchInfo;
80+ VectorTilingInfo vectorTilingInfo;
61};81};
62#pragma pack(pop)82#pragma pack(pop)
63 83 
@@ -27,6 +27,7 @@
27#define BATCH_MAT_MUL_V3_BASE_FULLLOAD 027#define BATCH_MAT_MUL_V3_BASE_FULLLOAD 0
28#define BATCH_MAT_MUL_V3_AL1_FULLLOAD 128#define BATCH_MAT_MUL_V3_AL1_FULLLOAD 1
29#define BATCH_MAT_MUL_V3_BL1_FULLLOAD 229#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 032#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_FALSE 0
32#define BATCH_MAT_MUL_V3_ISMULTIBATCHOUT_TRUE 133#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_FLOAT137#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(
@@ -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
@@ -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+}
@@ -36,8 +36,9 @@ string get_map_string(const std::map<string, string>& map, const string& key) {
36}36}
37bool IsDisplayTilingdata(const string& case_name, size_t index, uint64_t tilingKey)37bool IsDisplayTilingdata(const string& case_name, size_t index, uint64_t tilingKey)
38{38{
39- // 0-18 22-27 30-32 48之后 表示bmm实际用到的tilingdata39+ // 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校验全部的tilingdata44 // 基础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 
@@ -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}