Star75
18
代码
代码
Issues1
Pull Requests3
项目讨论
Wiki
分析
Star75
18
已合并
[feature] 提升msKL tiling/kernel快速调用接口使用体验，支持liboptiling.so和kernel .o自动匹配 #73
gong-siwei创建于 16 天前
[feature] 提升msKL tiling/kernel快速调用接口使用体验，支持liboptiling.so和kernel .o自动匹配 #73
已合并
gong-siwei创建于 16 天前
从gong-siwei/mskl_3797dev-0817合入到Ascend/msklmaster
共 9 个文件变更+987-76
Mdocs/en/api_reference/mskl_api_reference.md+23-12
@@ -65,14 +65,15 @@ The APIs provided by the msKL tool can call the tiling function in the msOpGen p
65 65 
66Calls the user's tiling function.66Calls the user's tiling function.
67 67 
68-> [!NOTE] Note 68+> [!NOTE] Note
69> tiling_func does not support calling the GetCompileInfo interface in [Basic Data Structure and Interface Reference](https://www.hiascend.com/document/detail/en/canncommercial/83RC1/API/basicdataapi/atlasopapi_07_00001.html).69> tiling_func does not support calling the GetCompileInfo interface in [Basic Data Structure and Interface Reference](https://www.hiascend.com/document/detail/en/canncommercial/83RC1/API/basicdataapi/atlasopapi_07_00001.html).
70 70 
71**Prototype**71**Prototype**
72 72 
73```python73```python
74-def tiling_func(op_type: str, inputs: list, outputs: list, lib_path: str,74+def tiling_func(op_type: str, inputs: list = None, outputs: list = None, lib_path: str = None,
75- inputs_info: list = None, outputs_info: list = None, attr=None, soc_version: str = None) -> TilingOutput75+ inputs_info: list = None, outputs_info: list = None, attr=None,
76+ soc_version: str = None, workspace: str = None) -> TilingOutput
76```77```
77 78 
78**Parameters**79**Parameters**
@@ -201,10 +202,20 @@ def tiling_func(op_type: str, inputs: list, outputs: list, lib_path: str,
201</td>202</td>
202<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p14829330143515"><a name="p14829330143515"></a><a name="p14829330143515"></a>Optional parameter.</p>203<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p14829330143515"><a name="p14829330143515"></a><a name="p14829330143515"></a>Optional parameter.</p>
203</td>204</td>
204-<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p16858115495219"><a name="p16858115495219"></a><a name="p16858115495219"></a>Path to the liboptiling.so file generated by compiling the msOpGen project. It can be found in the project directory using <strong id="b8137203510342"><a name="b8137203510342"></a><a name="b8137203510342"></a>find . -name 'liboptiling.so'</strong>. The msKL tool retrieves the user's tiling function in the order of deployed operators and then searches for <strong id="b1026950296"><a name="b1026950296"></a><a name="b1026950296"></a>.so</strong> files.</p>205+<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p16858115495219"><a name="p16858115495219"></a><a name="p16858115495219"></a>Path to the liboptiling.so file generated by compiling the msOpGen project. If this parameter is not specified, the msKL tool automatically searches for the liboptiling.so file recursively under the <strong id="b1026950296"><a name="b1026950296"></a><a name="b1026950296"></a>workspace</strong> directory (if multiple files are found, the package under the <strong id="b8137203510342"><a name="b8137203510342"></a><a name="b8137203510342"></a>_CPack_Packages</strong> or <strong id="b8137203510343"><a name="b8137203510343"></a><a name="b8137203510343"></a>op_tiling</strong> directory is preferred). If no such file is found, it falls back to the operators deployed in the CANN environment. If this parameter is explicitly configured, the path specified by this parameter is used and automatic search is skipped.</p>
205<p id="p2026182761114"><a name="p2026182761114"></a><a name="p2026182761114"></a>Data Type: str.</p>206<p id="p2026182761114"><a name="p2026182761114"></a><a name="p2026182761114"></a>Data Type: str.</p>
206</td>207</td>
207</tr>208</tr>
209+<tr id="row198901511"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p198901511"><a name="p198901511"></a><a name="p198901511"></a>workspace</p>
210+</td>
211+<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p198901512"><a name="p198901512"></a><a name="p198901512"></a>Input</p>
212+</td>
213+<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p198901513"><a name="p198901513"></a><a name="p198901513"></a>Optional parameter.</p>
214+</td>
215+<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p198901514"><a name="p198901514"></a><a name="p198901514"></a>Root directory of the operator project, for example, the CustomOp directory of an msOpGen project. If this parameter is not specified, the current directory is used by default. The msKL tool automatically searches for the liboptiling.so file and the kernel .o file recursively under this directory (the <strong id="b198901515"><a name="b198901515"></a><a name="b198901515"></a>liboptiling.so</strong> file and the kernel .o file whose path contains the <strong id="b198901516"><a name="b198901516"></a><a name="b198901516"></a>kernel</strong> directory) for use by tiling_func and get_kernel_from_binary.</p>
216+<p id="p198901516"><a name="p198901516"></a><a name="p198901516"></a>Data Type: str.</p>
217+</td>
218+</tr>
208<tr id="row331912112322"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p18319161123219"><a name="p18319161123219"></a><a name="p18319161123219"></a>soc_version</p>219<tr id="row331912112322"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p18319161123219"><a name="p18319161123219"></a><a name="p18319161123219"></a>soc_version</p>
209</td>220</td>
210<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p0700152183213"><a name="p0700152183213"></a><a name="p0700152183213"></a>Input</p>221<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p0700152183213"><a name="p0700152183213"></a><a name="p0700152183213"></a>Input</p>
@@ -278,7 +289,7 @@ output = np.zeros([M, N]).astype(np.float32)
278tiling_output = mskl.tiling_func(289tiling_output = mskl.tiling_func(
279 op_type="MatmulLeakyreluCustom",290 op_type="MatmulLeakyreluCustom",
280 inputs=[input_a, input_b, input_bias], outputs=[output],291 inputs=[input_a, input_b, input_bias], outputs=[output],
281- lib_path="liboptiling.so", # Tiling code compilation artifact. 292+ lib_path="liboptiling.so", # Tiling code compilation artifact.
282)293)
283```294```
284 295 
@@ -291,7 +302,7 @@ Generates an instance that can call the user's Kernel function.
291**Prototype**302**Prototype**
292 303 
293```python304```python
294-def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, tiling_key: int = None) -> CompiledKernel305+def get_kernel_from_binary(kernel_binary_file: str = None, kernel_type: str = None, tiling_key: int = None) -> CompiledKernel
295```306```
296 307 
297**Parameters**308**Parameters**
@@ -311,9 +322,9 @@ def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, til
311</td>322</td>
312<td class="cellrowborder" valign="top" width="10.07%" headers="mcps1.1.5.1.2 "><p id="p1239865692418"><a name="p1239865692418"></a><a name="p1239865692418"></a>Input</p>323<td class="cellrowborder" valign="top" width="10.07%" headers="mcps1.1.5.1.2 "><p id="p1239865692418"><a name="p1239865692418"></a><a name="p1239865692418"></a>Input</p>
313</td>324</td>
314-<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p224816370358"><a name="p224816370358"></a><a name="p224816370358"></a>Required parameter.</p>325+<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p224816370358"><a name="p224816370358"></a><a name="p224816370358"></a>Optional parameter.</p>
315</td>326</td>
316-<td class="cellrowborder" valign="top" width="65.25%" headers="mcps1.1.5.1.4 "><p id="p858655320395"><a name="p858655320395"></a><a name="p858655320395"></a>Path to the operator's kernel.o file. It can be found by executing the <strong id="b23211229105614"><a name="b23211229105614"></a><a name="b23211229105614"></a>find . -name '*.o'</strong> command in the project directory.</p>327+<td class="cellrowborder" valign="top" width="65.25%" headers="mcps1.1.5.1.4 "><p id="p858655320395"><a name="p858655320395"></a><a name="p858655320395"></a>Path to the operator's kernel.o file. If this parameter is not specified, the msKL tool automatically searches for <strong id="b23211229105614"><a name="b23211229105614"></a><a name="b23211229105614"></a>*.o</strong> files whose path contains the <strong id="b23211229105616"><a name="b23211229105616"></a><a name="b23211229105616"></a>kernel</strong> directory recursively under the tiling_func <strong id="b23211229105617"><a name="b23211229105617"></a><a name="b23211229105617"></a>workspace</strong> directory. If multiple .o files are found, the target kernel .o file is automatically selected based on <strong id="b23211229105618"><a name="b23211229105618"></a><a name="b23211229105618"></a>op_type → device soc → the same-directory JSON (dtype/format/shape)</strong>.</p>
317<p id="p960565114599"><a name="p960565114599"></a><a name="p960565114599"></a>Data Type: str.</p>328<p id="p960565114599"><a name="p960565114599"></a><a name="p960565114599"></a>Data Type: str.</p>
318</td>329</td>
319</tr>330</tr>
@@ -399,7 +410,7 @@ An executable Kernel object.
399</tbody>410</tbody>
400</table>411</table>
401 412 
402-> [!NOTE] 413+> [!NOTE]
403> The Kernel object type is CompiledKernel, which supports invoking the Kernel as follows: kernel[blockdim](arg1, arg2, ..., timeout=-1, device_id=0, repeat=1). During actual invocation, ensure that the input arguments of the CompiledKernel function are consistent with those when invoking the Kernel.414> The Kernel object type is CompiledKernel, which supports invoking the Kernel as follows: kernel[blockdim](arg1, arg2, ..., timeout=-1, device_id=0, repeat=1). During actual invocation, ensure that the input arguments of the CompiledKernel function are consistent with those when invoking the Kernel.
404 415 
405**Sample**416**Sample**
@@ -557,8 +568,8 @@ The File Path of the generated code.
557**Sample**568**Sample**
558 569 
559```py570```py
560-config = mskl.KernelInvokeConfig(kernel_file, kernel_name) 571+config = mskl.KernelInvokeConfig(kernel_file, kernel_name)
561-gen_file = mskl.Launcher(config).code_gen() 572+gen_file = mskl.Launcher(config).code_gen()
562```573```
563 574 
564**Related Class/Structure Definitions**575**Related Class/Structure Definitions**
@@ -572,7 +583,7 @@ class KernelInvokeConfig:
572 pass583 pass
573# The user can only pass a parameter of type KernelInvokeConfig.584# The user can only pass a parameter of type KernelInvokeConfig.
574class Launcher:585class Launcher:
575- def __init__(self, config: KernelInvokeConfig): 586+ def __init__(self, config: KernelInvokeConfig):
576 ...587 ...
577 a class that generates launch source code for a kernel588 a class that generates launch source code for a kernel
578 589 
Mdocs/zh/api_reference/mskl_api_reference.md+23-12
@@ -65,14 +65,15 @@ msKL工具提供的接口可以调用msOpGen工程中的tiling函数以及用户
65 65 
66调用用户的tiling函数。66调用用户的tiling函数。
67 67 
68-> [!NOTE] 68+> [!NOTE]
69> tiling_func不支持调用《[基础数据结构和接口参考](https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/API/basicdataapi/atlasopapi_07_00001.html)》中的GetCompileInfo接口。69> tiling_func不支持调用《[基础数据结构和接口参考](https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/API/basicdataapi/atlasopapi_07_00001.html)》中的GetCompileInfo接口。
70 70 
71**函数原型**71**函数原型**
72 72 
73```python73```python
74-def tiling_func(op_type: str, inputs: list, outputs: list, lib_path: str,74+def tiling_func(op_type: str, inputs: list = None, outputs: list = None, lib_path: str = None,
75- inputs_info: list = None, outputs_info: list = None, attr=None, soc_version: str = None) -> TilingOutput75+ inputs_info: list = None, outputs_info: list = None, attr=None,
76+ soc_version: str = None, workspace: str = None) -> TilingOutput
76```77```
77 78 
78**参数说明**79**参数说明**
@@ -201,10 +202,20 @@ def tiling_func(op_type: str, inputs: list, outputs: list, lib_path: str,
201</td>202</td>
202<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p14829330143515"><a name="p14829330143515"></a><a name="p14829330143515"></a>可选参数。</p>203<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p14829330143515"><a name="p14829330143515"></a><a name="p14829330143515"></a>可选参数。</p>
203</td>204</td>
204-<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p16858115495219"><a name="p16858115495219"></a><a name="p16858115495219"></a>msOpGen工程编译生成的liboptiling.so文件的路径，可在工程目录下通过<strong id="b8137203510342"><a name="b8137203510342"></a><a name="b8137203510342"></a>find . -name 'liboptiling.so'</strong>进行查找。msKL工具会按已部署算子、<strong id="b1026950296"><a name="b1026950296"></a><a name="b1026950296"></a>.so</strong>文件的查找顺序获取用户tiling函数。</p>205+<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p16858115495219"><a name="p16858115495219"></a><a name="p16858115495219"></a>msOpGen工程编译生成的liboptiling.so文件的路径。若不配置该参数，msKL工具会自动在<strong id="b8137203510342"><a name="b8137203510342"></a><a name="b8137203510342"></a>workspace</strong>目录下递归搜索liboptiling.so文件（存在多个时优先选择<strong id="b1026950296"><a name="b1026950296"></a><a name="b1026950296"></a>_CPack_Packages</strong>或<strong id="b1026950297"><a name="b1026950297"></a><a name="b1026950297"></a>op_tiling</strong>目录下的部署包）；若未搜索到，再回退使用已部署在CANN环境中的算子。若显式配置该参数，则以该参数指定的路径为准，不再执行自动搜索。</p>
205<p id="p2026182761114"><a name="p2026182761114"></a><a name="p2026182761114"></a>数据类型：str。</p>206<p id="p2026182761114"><a name="p2026182761114"></a><a name="p2026182761114"></a>数据类型：str。</p>
206</td>207</td>
207</tr>208</tr>
209+<tr id="row198901511"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p198901511"><a name="p198901511"></a><a name="p198901511"></a>workspace</p>
210+</td>
211+<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p198901512"><a name="p198901512"></a><a name="p198901512"></a>输入</p>
212+</td>
213+<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p198901513"><a name="p198901513"></a><a name="p198901513"></a>可选参数。</p>
214+</td>
215+<td class="cellrowborder" valign="top" width="65.27%" headers="mcps1.1.5.1.4 "><p id="p198901514"><a name="p198901514"></a><a name="p198901514"></a>算子工程根目录路径，例如msOpGen工程的CustomOp目录。若不配置该参数，默认使用当前目录。msKL工具会在该目录下递归搜索<strong id="b198901515"><a name="b198901515"></a><a name="b198901515"></a>liboptiling.so</strong>文件和路径包含<strong id="b198901516"><a name="b198901516"></a><a name="b198901516"></a>kernel</strong>目录的kernel .o文件，供tiling_func和get_kernel_from_binary使用。</p>
216+<p id="p198901516"><a name="p198901516"></a><a name="p198901516"></a>数据类型：str。</p>
217+</td>
218+</tr>
208<tr id="row331912112322"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p18319161123219"><a name="p18319161123219"></a><a name="p18319161123219"></a>soc_version</p>219<tr id="row331912112322"><td class="cellrowborder" valign="top" width="17.16%" headers="mcps1.1.5.1.1 "><p id="p18319161123219"><a name="p18319161123219"></a><a name="p18319161123219"></a>soc_version</p>
209</td>220</td>
210<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p0700152183213"><a name="p0700152183213"></a><a name="p0700152183213"></a>输入</p>221<td class="cellrowborder" valign="top" width="10.05%" headers="mcps1.1.5.1.2 "><p id="p0700152183213"><a name="p0700152183213"></a><a name="p0700152183213"></a>输入</p>
@@ -277,7 +288,7 @@ output = np.zeros([M, N]).astype(np.float32)
277tiling_output = mskl.tiling_func(288tiling_output = mskl.tiling_func(
278 op_type="MatmulLeakyreluCustom",289 op_type="MatmulLeakyreluCustom",
279 inputs=[input_a, input_b, input_bias], outputs=[output],290 inputs=[input_a, input_b, input_bias], outputs=[output],
280- lib_path="liboptiling.so", # tiling函数编译产物 291+ lib_path="liboptiling.so", # tiling函数编译产物
281)292)
282```293```
283 294 
@@ -290,7 +301,7 @@ tiling_output = mskl.tiling_func(
290**函数原型**301**函数原型**
291 302 
292```python303```python
293-def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, tiling_key: int = None) -> CompiledKernel304+def get_kernel_from_binary(kernel_binary_file: str = None, kernel_type: str = None, tiling_key: int = None) -> CompiledKernel
294```305```
295 306 
296**参数说明**307**参数说明**
@@ -310,9 +321,9 @@ def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, til
310</td>321</td>
311<td class="cellrowborder" valign="top" width="10.07%" headers="mcps1.1.5.1.2 "><p id="p1239865692418"><a name="p1239865692418"></a><a name="p1239865692418"></a>输入</p>322<td class="cellrowborder" valign="top" width="10.07%" headers="mcps1.1.5.1.2 "><p id="p1239865692418"><a name="p1239865692418"></a><a name="p1239865692418"></a>输入</p>
312</td>323</td>
313-<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p224816370358"><a name="p224816370358"></a><a name="p224816370358"></a>必选参数。</p>324+<td class="cellrowborder" valign="top" width="7.5200000000000005%" headers="mcps1.1.5.1.3 "><p id="p224816370358"><a name="p224816370358"></a><a name="p224816370358"></a>可选参数。</p>
314</td>325</td>
315-<td class="cellrowborder" valign="top" width="65.25%" headers="mcps1.1.5.1.4 "><p id="p858655320395"><a name="p858655320395"></a><a name="p858655320395"></a>算子kernel.o路径，可以在工程目录下执行<strong id="b23211229105614"><a name="b23211229105614"></a><a name="b23211229105614"></a>find . -name '*.o'</strong>命令进行查找。</p>326+<td class="cellrowborder" valign="top" width="65.25%" headers="mcps1.1.5.1.4 "><p id="p858655320395"><a name="p858655320395"></a><a name="p858655320395"></a>算子kernel.o路径。若不配置该参数，msKL工具会自动在tiling_func的<strong id="b23211229105614"><a name="b23211229105614"></a><a name="b23211229105614"></a>workspace</strong>目录下递归搜索路径包含<strong id="b23211229105616"><a name="b23211229105616"></a><a name="b23211229105616"></a>kernel</strong>目录的<strong id="b23211229105617"><a name="b23211229105617"></a><a name="b23211229105617"></a>*.o</strong>文件；若搜索到多个，则按<strong id="b23211229105618"><a name="b23211229105618"></a><a name="b23211229105618"></a>算子名(op_type) → 设备soc → 同目录json(dtype/format/shape)</strong>自动选取目标kernel .o文件。</p>
316<p id="p960565114599"><a name="p960565114599"></a><a name="p960565114599"></a>数据类型：str。</p>327<p id="p960565114599"><a name="p960565114599"></a><a name="p960565114599"></a>数据类型：str。</p>
317</td>328</td>
318</tr>329</tr>
@@ -399,7 +410,7 @@ def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, til
399</table>410</table>
400 411 
401> [!NOTE]412> [!NOTE]
402-> 413+>
403> Kernel对象类型为CompiledKernel，支持如下方式调用Kernel：kernel[blockdim](arg1, arg2, ..., timeout=-1, device_id=0, repeat=1)，实际调用时，需保证CompiledKernel函数的入参和调用Kernel时的入参一致。414> Kernel对象类型为CompiledKernel，支持如下方式调用Kernel：kernel[blockdim](arg1, arg2, ..., timeout=-1, device_id=0, repeat=1)，实际调用时，需保证CompiledKernel函数的入参和调用Kernel时的入参一致。
404 415 
405**调用示例**416**调用示例**
@@ -557,8 +568,8 @@ def code_gen(self, gen_file)
557**调用示例**568**调用示例**
558 569 
559```py570```py
560-config = mskl.KernelInvokeConfig(kernel_file, kernel_name) 571+config = mskl.KernelInvokeConfig(kernel_file, kernel_name)
561-mskl.Launcher(config).code_gen(gen_file) 572+mskl.Launcher(config).code_gen(gen_file)
562```573```
563 574 
564**相关类/结构体定义**575**相关类/结构体定义**
@@ -572,7 +583,7 @@ class KernelInvokeConfig:
572 pass583 pass
573# 用户仅能传KernelInvokeConfig类型584# 用户仅能传KernelInvokeConfig类型
574class Launcher:585class Launcher:
575- def __init__(self, config: KernelInvokeConfig): 586+ def __init__(self, config: KernelInvokeConfig):
576 ...587 ...
577 a class that generates launch source code for a kernel588 a class that generates launch source code for a kernel
578 589 
Mmskl/launcher/config.py+14-13
@@ -23,7 +23,7 @@ import numpy as np
23from ..utils import safe_check, logger23from ..utils import safe_check, logger
24from ..utils.safe_check import FileChecker24from ..utils.safe_check import FileChecker
25from .context import context25from .context import context
26-from ..utils.launcher_utils import get_cann_path26+from ..utils.launcher_utils import get_cann_path, search_tiling_so, get_workspace_path
27from ..utils.autotune_utils import (27from ..utils.autotune_utils import (
28 is_torch_tensor_instance,28 is_torch_tensor_instance,
29 safe_prod,29 safe_prod,
@@ -217,6 +217,7 @@ class TilingConfig:
217 outputs_info: list = None,217 outputs_info: list = None,
218 attr=None,218 attr=None,
219 soc_version: str = None,219 soc_version: str = None,
220+ workspace: str = None,
220 ):221 ):
221 # op_type 做透传处理222 # op_type 做透传处理
222 if not self._is_valid_key_str(op_type):223 if not self._is_valid_key_str(op_type):
@@ -224,7 +225,7 @@ class TilingConfig:
224 self.op_type = op_type225 self.op_type = op_type
225 self._parse_io_params(inputs, outputs, inputs_info, outputs_info)226 self._parse_io_params(inputs, outputs, inputs_info, outputs_info)
226 self._parse_attr(attr)227 self._parse_attr(attr)
227- self._parse_lib_path(lib_path)228+ self._parse_lib_path(lib_path, workspace)
228 self._parse_soc_version(soc_version)229 self._parse_soc_version(soc_version)
229 230 
230 @staticmethod231 @staticmethod
@@ -361,9 +362,14 @@ class TilingConfig:
361 elif (t is not None) and (not is_torch_or_numpy_tensor(t)):362 elif (t is not None) and (not is_torch_or_numpy_tensor(t)):
362 raise ValueError(f'Type of {para}[{idx}] should be torch.Tensor, numpy.ndarray or the list of above')363 raise ValueError(f'Type of {para}[{idx}] should be torch.Tensor, numpy.ndarray or the list of above')
363 364 
364- def _parse_lib_path(self, lib_path: str):365+ def _parse_lib_path(self, lib_path: str, workspace: str = None):
365 if lib_path is None:366 if lib_path is None:
366- # load liboptiling.so in cann as default367+ # 未显式传入lib_path时，在整个workspace目录下递归搜索liboptiling.so
368+ # （存在多个时优先打包产物_CPack_Packages/op_tiling下的部署包），搜索失败则回退到cann默认的liboptiling.so
369+ tiling_so = search_tiling_so(get_workspace_path(workspace))
370+ if tiling_so is not None:
371+ self.lib_path = _escape_cpp_string(os.path.realpath(tiling_so))
372+ return
367 cann_path = get_cann_path()373 cann_path = get_cann_path()
368 tiling_so = os.path.join(cann_path, "lib64/liboptiling.so")374 tiling_so = os.path.join(cann_path, "lib64/liboptiling.so")
369 self.lib_path = _escape_cpp_string(tiling_so) if os.path.exists(tiling_so) else ''375 self.lib_path = _escape_cpp_string(tiling_so) if os.path.exists(tiling_so) else ''
@@ -688,7 +694,7 @@ class TilingConfig:
688 checker = FileChecker(data_path, 'file')694 checker = FileChecker(data_path, 'file')
689 if not checker.check_input_file():695 if not checker.check_input_file():
690 raise PermissionError(f'{param_name}["data_path"] check permission failed')696 raise PermissionError(f'{param_name}["data_path"] check permission failed')
691- info['data_path'] = repr(os.path.abspath(data_path))[1:-1]697+ info['data_path'] = _escape_cpp_string(os.path.abspath(data_path))
692 698 
693 ch = ChainHandler(699 ch = ChainHandler(
694 [verify_fmt, verify_ori_format, verify_shape, verify_ori_shape, verify_dtype, verify_data_path]700 [verify_fmt, verify_ori_format, verify_shape, verify_ori_shape, verify_dtype, verify_data_path]
@@ -742,14 +748,9 @@ class TilingConfig:
742 )748 )
743 last_param = ''749 last_param = ''
744 if tensor['data_path'] != '':750 if tensor['data_path'] != '':
745- escaped_path = (751+ # data_path 已由 _escape_cpp_string 统一完成转义（反斜杠/双引号/换行/制表符等），
746- tensor['data_path']752+ # 直接嵌入C++字符串字面量，避免二次转义导致路径错误
747- .replace('\\', '\\\\')753+ escaped_path = tensor['data_path']
748- .replace('"', '\\"')
749- .replace('\n', '\\n')
750- .replace('\t', '\\t')
751- .replace('\r', '\\r')
752- )
753 last_param = f", string{{\"{escaped_path}\"}}"754 last_param = f", string{{\"{escaped_path}\"}}"
754 elif tensor['dtype'] in supported_types and tensor['addr'] != 0:755 elif tensor['dtype'] in supported_types and tensor['addr'] != 0:
755 if not isinstance(tensor['addr'], int) or tensor['addr'] < 0:756 if not isinstance(tensor['addr'], int) or tensor['addr'] < 0:
Mmskl/launcher/context.py+39-4
@@ -16,6 +16,7 @@
16# See the Mulan PSL v2 for more details.16# See the Mulan PSL v2 for more details.
17# -------------------------------------------------------------------------17# -------------------------------------------------------------------------
18 18 
19+ 
19class Context:20class Context:
20 """21 """
21 Maintain context of modules(code_generator, compile, etc.)22 Maintain context of modules(code_generator, compile, etc.)
@@ -28,8 +29,10 @@ class Context:
28 self._launch_src_file = None29 self._launch_src_file = None
29 self._build_script = None30 self._build_script = None
30 self._blockdim = None31 self._blockdim = None
31- self._tiling_output = None # mskl.launcher.opgen_workflow.TilingOutput32+ self._tiling_output = None
32- self._op_type = None # str like AddCustom33+ self._op_type = None
34+ self._workspace = None
35+ self._io_info = None
33 self._autotune_in_progress = False36 self._autotune_in_progress = False
34 self._prelaunch_flag = False37 self._prelaunch_flag = False
35 38 
@@ -142,8 +145,40 @@ class Context:
142 def op_type(self, value):145 def op_type(self, value):
143 self._op_type = value146 self._op_type = value
144 147 
148+ @property
149+ def workspace(self):
150+ return self._workspace
151+ 
152+ @workspace.setter
153+ def workspace(self, value):
154+ self._workspace = value
155+ 
156+ @property
157+ def io_info(self):
158+ # {'inputs': [{'dtype','format','shape'},...], 'outputs': [...]}，用于kernel .o自动匹配
159+ return self._io_info
160+ 
161+ @io_info.setter
162+ def io_info(self, value):
163+ self._io_info = value
164+ 
145 def reset(self):165 def reset(self):
146- self.__init__()166+ # 重置实例状态，避免直接调用 __init__（pylint C2801）
167+ self._config = None
168+ self._kernel_name = None
169+ self._kernel_src_file = None
170+ self._launch_src_file = None
171+ self._build_script = None
172+ self._blockdim = None
173+ self._tiling_output = None
174+ self._op_type = None
175+ self._workspace = None
176+ self._io_info = None
177+ self._autotune_in_progress = False
178+ self._prelaunch_flag = False
179+ self._decl_args = None
180+ self._template_args = None
181+ self._kernel_args = None
147 182 
148 183 
149-context = Context()184+context = Context()
Mmskl/launcher/opgen_workflow.py+139-8
@@ -23,6 +23,14 @@ from .code_generator import Launcher
23from .compiler import compile_tiling, compile_kernel_binary, CompiledKernel23from .compiler import compile_tiling, compile_kernel_binary, CompiledKernel
24from .context import context24from .context import context
25from ..utils.safe_check import FileChecker, DATA_DIRECTORY_AUTHORITY25from ..utils.safe_check import FileChecker, DATA_DIRECTORY_AUTHORITY
26+from ..utils.launcher_utils import (
27+ get_workspace_path,
28+ search_kernel_binary_files,
29+ match_kernel_binary_by_op_type,
30+ match_kernel_binary_by_soc,
31+ match_kernel_binary_by_json,
32+ get_device_soc_name,
33+)
26from ..utils import logger34from ..utils import logger
27 35 
28TILING_FUNC_CNT = 036TILING_FUNC_CNT = 0
@@ -50,12 +58,84 @@ class TilingOutput:
50 self.tiling_key = tiling_output["tiling_key"]58 self.tiling_key = tiling_output["tiling_key"]
51 59 
52 60 
53-def tiling_func(op_type: str, inputs: list = None, outputs: list = None, lib_path: str = None,61+def _tensor_dtype(t) -> str:
54- inputs_info: list = None, outputs_info: list = None, attr=None,62+ """获取tensor的dtype字符串（float16/float32等），供kernel .o的json匹配使用。"""
55- soc_version: str = None) -> TilingOutput:63+ if t is None:
64+ return ''
65+ if isinstance(t, np.ndarray):
66+ return t.dtype.name
67+ if hasattr(t, 'dtype'):
68+ return str(t.dtype).replace('torch.', '')
69+ return ''
70+ 
71+ 
72+def _build_io_info(inputs, outputs, inputs_info, outputs_info) -> dict:
73+ """从tiling_func入参提取inputs/outputs的dtype/format/shape信息，用于kernel .o自动匹配。"""
74+ 
75+ def _io_of(t, info):
76+ if info is None:
77+ info = {}
78+ dtype = _tensor_dtype(t) or info.get('dtype', '')
79+ fmt = info.get('format', '') or 'ND'
80+ shape = list(t.shape) if (t is not None and hasattr(t, 'shape')) else list(info.get('shape', []))
81+ return {'dtype': dtype, 'format': fmt, 'shape': shape}
82+ 
83+ def _flatten(tensors, infos):
84+ tensors = list(tensors) if tensors else []
85+ infos = list(infos) if infos else []
86+ length = max(len(tensors), len(infos))
87+ res = []
88+ for i in range(length):
89+ t = tensors[i] if i < len(tensors) else None
90+ info = infos[i] if i < len(infos) else None
91+ if isinstance(t, list) or isinstance(info, list):
92+ t_list = t if isinstance(t, list) else [None]
93+ info_list = info if isinstance(info, list) else [info]
94+ for j in range(max(len(t_list), len(info_list))):
95+ res.append(
96+ _io_of(t_list[j] if j < len(t_list) else None, info_list[j] if j < len(info_list) else None)
97+ )
98+ else:
99+ res.append(_io_of(t, info))
100+ return res
101+ 
102+ return {
103+ 'inputs': _flatten(inputs, inputs_info),
104+ 'outputs': _flatten(outputs, outputs_info),
105+ }
106+ 
107+ 
108+def tiling_func(
109+ op_type: str,
110+ inputs: list = None,
111+ outputs: list = None,
112+ lib_path: str = None,
113+ inputs_info: list = None,
114+ outputs_info: list = None,
115+ attr=None,
116+ soc_version: str = None,
117+ workspace: str = None,
118+) -> TilingOutput:
119+ """
120+ :param op_type: op type, e.g. AddCustom
121+ :param inputs: tensors of kernel inputs
122+ :param outputs: tensors of kernel outputs
123+ :param lib_path: optional, path of liboptiling.so. If not set, it will be searched recursively in
124+ the <workspace> directory (the package under _CPack_Packages/op_tiling is preferred)
125+ :param inputs_info: info of kernel inputs
126+ :param outputs_info: info of kernel outputs
127+ :param attr: attrs of the op
128+ :param soc_version: soc version
129+ :param workspace: optional, root directory of the op project, default is current directory.
130+ Used to search liboptiling.so and kernel .o automatically
131+ :return: TilingOutput
132+ """
56 global TILING_FUNC_CNT133 global TILING_FUNC_CNT
57 TILING_FUNC_CNT += 1134 TILING_FUNC_CNT += 1
58- config = TilingConfig(op_type, inputs, outputs, lib_path, inputs_info, outputs_info, attr, soc_version)135+ workspace = get_workspace_path(workspace)
136+ context.workspace = workspace
137+ context.io_info = _build_io_info(inputs, outputs, inputs_info, outputs_info)
138+ config = TilingConfig(op_type, inputs, outputs, lib_path, inputs_info, outputs_info, attr, soc_version, workspace)
59 tmp_path = init_tmp_folder()139 tmp_path = init_tmp_folder()
60 cpp_path = os.path.join(tmp_path, f'_mskl_gen_tiling.{TILING_FUNC_CNT}.cpp')140 cpp_path = os.path.join(tmp_path, f'_mskl_gen_tiling.{TILING_FUNC_CNT}.cpp')
61 Launcher(config).code_gen(cpp_path)141 Launcher(config).code_gen(cpp_path)
@@ -71,15 +151,65 @@ def tiling_func(op_type: str, inputs: list = None, outputs: list = None, lib_pat
71 return output151 return output
72 152 
73 153 
74-def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, tiling_key: int = None) -> CompiledKernel:154+def _auto_search_kernel_binary() -> str:
155+ """根据tiling_func传入的workspace自动搜索kernel .o文件。
156+ 选取顺序：算子名(op_type) -> 设备soc目录 -> 同目录json(dtype/format/shape)。
75 """157 """
76- :param kernel_binary_file: path of kernel.o158+ workspace = context.workspace if context.workspace is not None else os.getcwd()
159+ o_files = search_kernel_binary_files(workspace)
160+ if not o_files:
161+ raise FileNotFoundError(
162+ f'Cannot find any kernel .o file under {workspace}, please input [kernel_binary_file] or check [workspace]'
163+ )
164+ # 1. 按算子名过滤：kernel .o 文件名以 <op_type>_ 开头
165+ op_type = context.op_type
166+ if op_type:
167+ o_files = match_kernel_binary_by_op_type(o_files, op_type)
168+ if not o_files:
169+ raise FileNotFoundError(
170+ f'Cannot find kernel .o matching op_type {op_type} under {workspace}, '
171+ f'please input [kernel_binary_file] or check [op_type]'
172+ )
173+ if len(o_files) == 1:
174+ return o_files[0]
175+ # 2. 优先选择与当前设备soc匹配的kernel .o
176+ soc_files = match_kernel_binary_by_soc(o_files, get_device_soc_name())
177+ candidates = soc_files if soc_files else o_files
178+ if len(candidates) == 1:
179+ return candidates[0]
180+ # 3. 根据kernel .o同目录json中的inputs/outputs信息（dtype/format/shape）区分
181+ json_files = match_kernel_binary_by_json(candidates, context.io_info)
182+ if len(json_files) == 1:
183+ return json_files[0]
184+ if len(json_files) > 1:
185+ logger.warning(
186+ f'Multiple kernel .o files match op_type {op_type}, soc and io: '
187+ f'{json_files}, use the first one {json_files[0]}'
188+ )
189+ return json_files[0]
190+ logger.warning(f'Multiple kernel .o files match op_type {op_type}: {candidates}, use the first one {candidates[0]}')
191+ return candidates[0]
192+ 
193+ 
194+def get_kernel_from_binary(
195+ kernel_binary_file: str = None, kernel_type: str = None, tiling_key: int = None
196+) -> CompiledKernel:
197+ """
198+ :param kernel_binary_file: optional, path of kernel.o. If not set, it will be searched recursively in
199+ the <workspace> directory and matched by op_type/soc/json
77 :param kernel_type: ['mix', 'cube', 'vec']200 :param kernel_type: ['mix', 'cube', 'vec']
78 :param tiling_key: None will use tiling_func()'s return value201 :param tiling_key: None will use tiling_func()'s return value
79 :return: CompiledKernel202 :return: CompiledKernel
80 """203 """
81 global GET_KERNEL_FROM_BINARY_CNT204 global GET_KERNEL_FROM_BINARY_CNT
82 GET_KERNEL_FROM_BINARY_CNT += 1205 GET_KERNEL_FROM_BINARY_CNT += 1
206+ if kernel_binary_file is None:
207+ if tiling_key is None:
208+ if context.tiling_output is None:
209+ raise Exception('Please call mskl.tiling_func or assign [tiling_key]')
210+ tiling_key = context.tiling_output.tiling_key
211+ kernel_binary_file = _auto_search_kernel_binary()
212+ logger.debug(f'Automatically searched kernel binary file is {kernel_binary_file}')
83 config = KernelBinaryInvokeConfig(kernel_binary_file, kernel_type, tiling_key)213 config = KernelBinaryInvokeConfig(kernel_binary_file, kernel_type, tiling_key)
84 tmp_path = init_tmp_folder()214 tmp_path = init_tmp_folder()
85 context.opgen_tmp_dir_path = tmp_path215 context.opgen_tmp_dir_path = tmp_path
@@ -87,6 +217,7 @@ def get_kernel_from_binary(kernel_binary_file: str, kernel_type: str = None, til
87 Launcher(config).code_gen(cpp_path)217 Launcher(config).code_gen(cpp_path)
88 so_path = os.path.join(tmp_path, f'_mskl_gen_binary_module.{GET_KERNEL_FROM_BINARY_CNT}.so')218 so_path = os.path.join(tmp_path, f'_mskl_gen_binary_module.{GET_KERNEL_FROM_BINARY_CNT}.so')
89 kernel = compile_kernel_binary(cpp_path, so_path)219 kernel = compile_kernel_binary(cpp_path, so_path)
90- logger.debug(f'Call get_kernel_from_binary {GET_KERNEL_FROM_BINARY_CNT} success, '220+ logger.debug(
91- f'kernel path is {kernel_binary_file}')221+ f'Call get_kernel_from_binary {GET_KERNEL_FROM_BINARY_CNT} success, kernel path is {kernel_binary_file}'
222+ )
92 return kernel223 return kernel
Mmskl/utils/launcher_utils.py+204-0
@@ -16,9 +16,213 @@
16# See the Mulan PSL v2 for more details.16# See the Mulan PSL v2 for more details.
17# -------------------------------------------------------------------------17# -------------------------------------------------------------------------
18 18 
19+import json
19import os20import os
21+from mskl.utils import logger
20from mskl.utils.safe_check import FileChecker22from mskl.utils.safe_check import FileChecker
21 23 
24+TILING_SO_NAME = 'liboptiling.so'
25+KERNEL_DIR_NAME = 'kernel'
26+# 递归搜索workspace时跳过的目录，避免遍历无关内容
27+_SKIP_DIRS = {'.git', '__pycache__', 'node_modules', '.venv', 'third_party'}
28+ 
29+ 
30+def get_workspace_path(workspace: str = None) -> str:
31+ """将workspace参数归一化为绝对路径，未指定时使用当前目录。"""
32+ if workspace is None:
33+ workspace = os.getcwd()
34+ if not isinstance(workspace, str) or not workspace:
35+ raise ValueError('workspace should be a non-empty str')
36+ workspace = os.path.abspath(workspace)
37+ checker = FileChecker(workspace, "dir")
38+ if not checker.check_input_file():
39+ if not os.path.exists(workspace):
40+ raise FileNotFoundError(f'workspace {workspace} not exist, please verify path')
41+ raise PermissionError(f'workspace {workspace} check permission failed, please verify path and permissions')
42+ return workspace
43+ 
44+ 
45+def search_tiling_so(workspace: str = None):
46+ """在整个workspace目录下递归搜索 liboptiling.so。
47+ 存在多个时，优先选择打包产物（路径含 _CPack_Packages 或 op_tiling）中的文件。
48+ """
49+ workspace = get_workspace_path(workspace)
50+ found = []
51+ for dirpath, dirnames, filenames in os.walk(workspace):
52+ dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]
53+ for name in filenames:
54+ if name == TILING_SO_NAME:
55+ found.append(os.path.join(dirpath, name))
56+ if not found:
57+ return None
58+ for path in found:
59+ parts = os.path.normpath(path).split(os.sep)
60+ if '_CPack_Packages' in parts or 'op_tiling' in parts:
61+ return path
62+ return found[0]
63+ 
64+ 
65+def search_kernel_binary_files(workspace: str = None):
66+ """在整个workspace目录下递归搜索路径含 kernel 目录的 *.o 文件。"""
67+ workspace = get_workspace_path(workspace)
68+ o_files = []
69+ for dirpath, dirnames, filenames in os.walk(workspace):
70+ dirnames[:] = [d for d in dirnames if d not in _SKIP_DIRS]
71+ if KERNEL_DIR_NAME not in dirpath.split(os.sep):
72+ continue
73+ for name in filenames:
74+ if name.endswith('.o'):
75+ o_files.append(os.path.join(dirpath, name))
76+ return o_files
77+ 
78+ 
79+def match_kernel_binary_by_op_type(o_files, op_type: str):
80+ """从kernel .o文件中，根据算子名称（文件名以 <op_type>_ 开头）筛选出匹配的文件。"""
81+ if not op_type:
82+ return o_files
83+ prefix = op_type + '_'
84+ return [o_file for o_file in o_files if os.path.basename(o_file).startswith(prefix)]
85+ 
86+ 
87+def get_device_soc_name() -> str:
88+ """获取当前设备soc名称（如 Ascend950PR_9579），获取失败时返回空串。"""
89+ cann_path = get_cann_path()
90+ try:
91+ import ctypes
92+ 
93+ acl = ctypes.CDLL(os.path.join(cann_path, 'lib64/libascendcl.so'), mode=ctypes.RTLD_GLOBAL)
94+ acl.aclrtGetSocName.restype = ctypes.c_char_p
95+ name = acl.aclrtGetSocName()
96+ if name:
97+ return name.decode()
98+ except Exception as e: # 获取失败时回退到rtGetSocVersion
99+ logger.debug(f'Get soc name by aclrtGetSocName failed: {e}')
100+ try:
101+ import ctypes
102+ 
103+ rt = ctypes.CDLL(os.path.join(cann_path, 'lib64/libruntime.so'), mode=ctypes.RTLD_GLOBAL)
104+ rt.rtGetSocVersion.argtypes = [ctypes.c_char_p, ctypes.c_uint64]
105+ rt.rtGetSocVersion.restype = ctypes.c_int
106+ buf = ctypes.create_string_buffer(64)
107+ if rt.rtGetSocVersion(buf, 64) == 0 and buf.value:
108+ return buf.value.decode()
109+ except Exception as e: # 获取失败时返回空串，后续回退到默认匹配
110+ logger.debug(f'Get soc version by rtGetSocVersion failed: {e}')
111+ return ''
112+ 
113+ 
114+def _get_soc_dir(o_file: str) -> str:
115+ """从kernel .o路径中提取soc目录名，路径形如 .../kernel/<soc>/<op>/<file>.o。"""
116+ parts = os.path.normpath(o_file).split(os.sep)
117+ try:
118+ idx = parts.index(KERNEL_DIR_NAME)
119+ if idx + 1 < len(parts):
120+ return parts[idx + 1].lower()
121+ except ValueError:
122+ pass
123+ return ''
124+ 
125+ 
126+def match_kernel_binary_by_soc(o_files, soc_name: str):
127+ """按设备soc（如 Ascend950PR_9579）筛选kernel .o，优先匹配路径中soc目录与设备soc一致的。"""
128+ if not soc_name:
129+ return []
130+ soc_lower = soc_name.lower()
131+ matched = []
132+ for o_file in o_files:
133+ soc_dir = _get_soc_dir(o_file)
134+ # ascend950 匹配 Ascend950PR_9579 / ascend950pr；ascend910b 匹配 Ascend910B3 等
135+ if soc_dir and (soc_lower.startswith(soc_dir) or soc_dir.startswith(soc_lower)):
136+ matched.append(o_file)
137+ return matched
138+ 
139+ 
140+def _normalize_dtype(dtype: str) -> str:
141+ """归一化dtype名称，兼容 float/float32、half/float16 等别名。"""
142+ if not dtype:
143+ return ''
144+ aliases = {
145+ 'half': 'float16',
146+ 'float': 'float32',
147+ 'fp16': 'float16',
148+ 'fp32': 'float32',
149+ 'fp64': 'float64',
150+ 'double': 'float64',
151+ 'int': 'int32',
152+ }
153+ d = str(dtype).lower()
154+ return aliases.get(d, d)
155+ 
156+ 
157+def _match_json_dtype(json_dtype: str, actual_dtype: str) -> bool:
158+ if not json_dtype or not actual_dtype:
159+ return True # 信息缺失时不作为排除条件
160+ return _normalize_dtype(json_dtype) == _normalize_dtype(actual_dtype)
161+ 
162+ 
163+def _match_json_format(json_format: str, actual_format: str, match_mode: str = '') -> bool:
164+ if not json_format or not actual_format:
165+ return True
166+ if 'agnostic' in str(match_mode).lower():
167+ return True # FormatAgnostic：任意format均可
168+ return str(json_format).lower() == str(actual_format).lower()
169+ 
170+ 
171+def _match_json_shape(json_shape, actual_shape) -> bool:
172+ if not json_shape:
173+ return True
174+ # 全部为负值（如 [-2]）表示动态shape，匹配任意shape
175+ if all(dim < 0 for dim in json_shape):
176+ return True
177+ if len(json_shape) != len(actual_shape):
178+ return False
179+ for json_dim, actual_dim in zip(json_shape, actual_shape):
180+ if json_dim < 0: # 单个维度为负值表示该维度动态
181+ continue
182+ if json_dim != actual_dim:
183+ return False
184+ return True
185+ 
186+ 
187+def _json_supports_io(meta: dict, io_info: dict) -> bool:
188+ """判断kernel meta json中记录的inputs/outputs是否与本次调用的io信息匹配。"""
189+ support = meta.get('supportInfo', {})
190+ for io_key, actual_list in (('inputs', io_info.get('inputs', [])), ('outputs', io_info.get('outputs', []))):
191+ json_list = support.get(io_key, [])
192+ if len(actual_list) > len(json_list):
193+ return False
194+ for idx, actual in enumerate(actual_list):
195+ spec = json_list[idx] or {}
196+ if not _match_json_dtype(spec.get('dtype', ''), actual.get('dtype', '')):
197+ return False
198+ if not _match_json_format(
199+ spec.get('format', ''), actual.get('format', ''), spec.get('format_match_mode', '')
200+ ):
201+ return False
202+ if not _match_json_shape(spec.get('shape'), actual.get('shape', [])):
203+ return False
204+ return True
205+ 
206+ 
207+def match_kernel_binary_by_json(o_files, io_info: dict):
208+ """根据kernel .o同目录json中的inputs/outputs信息（dtype/format/shape）筛选kernel .o。"""
209+ if not o_files or not io_info:
210+ return []
211+ matched = []
212+ for o_file in o_files:
213+ json_path = o_file[:-1] + 'json' # .o -> .json
214+ if not os.path.exists(json_path):
215+ continue
216+ try:
217+ with open(json_path, 'r', encoding='utf-8') as f:
218+ meta = json.load(f)
219+ except Exception as e: # json解析失败时跳过该候选
220+ logger.warning(f'Read kernel meta json {json_path} failed, skip it. error: {e}')
221+ continue
222+ if _json_supports_io(meta, io_info):
223+ matched.append(o_file)
224+ return matched
225+ 
22 226 
23def get_cann_path() -> str:227def get_cann_path() -> str:
24 cann_path = os.getenv('ASCEND_HOME_PATH')228 cann_path = os.getenv('ASCEND_HOME_PATH')
Mtest/launcher/test_config.py+71-1
@@ -21,10 +21,11 @@ import unittest
21import pathlib21import pathlib
22import stat22import stat
23import tempfile23import tempfile
24+import shutil
24import numpy as np25import numpy as np
25 26 
26from test.utils.test_base import TestBase # pylint: disable=E061127from test.utils.test_base import TestBase # pylint: disable=E0611
27-from mskl.launcher.config import KernelInvokeConfig, TilingConfig, KernelBinaryInvokeConfig28+from mskl.launcher.config import KernelInvokeConfig, TilingConfig, KernelBinaryInvokeConfig, _escape_cpp_string
28 29 
29 30 
30class MockTilingOutput:31class MockTilingOutput:
@@ -102,6 +103,25 @@ class TestTilingConfig(unittest.TestCase):
102 )103 )
103 self.assertEqual(len(config.outputs_list), 1)104 self.assertEqual(len(config.outputs_list), 1)
104 105 
106+ def test_tiling_config_data_path_escaped_once(self):
107+ # data_path 含反斜杠/双引号，应统一由 _escape_cpp_string 转义，并直接嵌入生成的tds中（不二次转义）
108+ tmp = tempfile.mkdtemp()
109+ try:
110+ data_path = os.path.join(tmp, 'data "quote"\\dir.bin')
111+ with open(data_path, 'w', encoding='utf-8') as f:
112+ f.write('')
113+ config = TilingConfig(
114+ op_type="TestOp",
115+ inputs=[self.test_input],
116+ outputs=[self.test_output],
117+ inputs_info=[{"shape": [128, 128], "dtype": "float32", "format": "nd", "data_path": data_path}],
118+ )
119+ escaped = _escape_cpp_string(os.path.abspath(data_path))
120+ self.assertEqual(config.inputs_list[0][0]['data_path'], escaped)
121+ self.assertIn(f'string{{"{escaped}"}}', config.tds)
122+ finally:
123+ shutil.rmtree(tmp, ignore_errors=True)
124+ 
105 def test_tiling_config_with_soc_version(self):125 def test_tiling_config_with_soc_version(self):
106 config = TilingConfig(126 config = TilingConfig(
107 op_type="TestOp", inputs=[self.test_input], outputs=[self.test_output], soc_version="Ascend910B"127 op_type="TestOp", inputs=[self.test_input], outputs=[self.test_output], soc_version="Ascend910B"
@@ -135,6 +155,56 @@ class TestTilingConfig(unittest.TestCase):
135 def test_tiling_config_empty_lib_path(self):155 def test_tiling_config_empty_lib_path(self):
136 self.assertRaises(ValueError, TilingConfig, "TestOp", [self.test_input], [self.test_output], lib_path="")156 self.assertRaises(ValueError, TilingConfig, "TestOp", [self.test_input], [self.test_output], lib_path="")
137 157 
158+ def test_tiling_config_auto_search_liboptiling_by_workspace(self):
159+ ws = './_test_mskl_cfg_ws'
160+ so_path = os.path.join(
161+ ws,
162+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
163+ 'op_impl/ai_core/tbe/op_tiling/liboptiling.so',
164+ )
165+ try:
166+ if os.path.exists(ws):
167+ shutil.rmtree(ws, ignore_errors=True)
168+ os.makedirs(os.path.dirname(so_path), exist_ok=True)
169+ with open(so_path, 'w', encoding='utf-8') as f:
170+ f.write('test so')
171+ config = TilingConfig(op_type="TestOp", inputs=[self.test_input], outputs=[self.test_output], workspace=ws)
172+ self.assertIn(os.path.realpath(so_path), config.lib_path)
173+ finally:
174+ if os.path.exists(ws):
175+ shutil.rmtree(ws, ignore_errors=True)
176+ 
177+ def test_tiling_config_explicit_lib_path_priority_over_workspace(self):
178+ ws = './_test_mskl_cfg_ws'
179+ with tempfile.NamedTemporaryFile(suffix='.so', delete=False) as f:
180+ lib_path = f.name
181+ try:
182+ if os.path.exists(ws):
183+ shutil.rmtree(ws, ignore_errors=True)
184+ os.makedirs(ws, exist_ok=True)
185+ config = TilingConfig(
186+ op_type="TestOp",
187+ inputs=[self.test_input],
188+ outputs=[self.test_output],
189+ lib_path=lib_path,
190+ workspace=ws,
191+ )
192+ self.assertEqual(config.lib_path, lib_path.replace('\\', '\\\\'))
193+ finally:
194+ os.remove(lib_path)
195+ if os.path.exists(ws):
196+ shutil.rmtree(ws, ignore_errors=True)
197+ 
198+ def test_tiling_config_invalid_workspace(self):
199+ self.assertRaises(
200+ Exception,
201+ TilingConfig,
202+ "TestOp",
203+ [self.test_input],
204+ [self.test_output],
205+ workspace='./_not_exist_mskl_ws_dir',
206+ )
207+ 
138 def test_tiling_config_invalid_op_type(self):208 def test_tiling_config_invalid_op_type(self):
139 # 只测试真正会触发的情况209 # 只测试真正会触发的情况
140 invalid_ops = ["", "invalid op", "op@name"]210 invalid_ops = ["", "invalid op", "op@name"]
Atest/launcher/test_launcher_utils.py+202-0
@@ -0,0 +1,202 @@
1+#!/usr/bin/python
2+# -*- coding: UTF-8 -*-
3+# -------------------------------------------------------------------------
4+# This file is part of the MindStudio project.
5+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
6+#
7+# MindStudio is licensed under Mulan PSL v2.
8+# You can use this software according to the terms and conditions of the Mulan PSL v2.
9+# You may obtain a copy of Mulan PSL v2 at:
10+#
11+# http://license.coscl.org.cn/MulanPSL2
12+#
13+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
14+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
15+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
16+# See the Mulan PSL v2 for more details.
17+# -------------------------------------------------------------------------
18+ 
19+import os
20+import json
21+import unittest
22+ 
23+from mskl.utils.launcher_utils import (
24+ get_workspace_path,
25+ search_tiling_so,
26+ search_kernel_binary_files,
27+ match_kernel_binary_by_op_type,
28+ match_kernel_binary_by_soc,
29+ match_kernel_binary_by_json,
30+ _match_json_shape,
31+)
32+ 
33+ 
34+def _touch(path):
35+ os.makedirs(os.path.dirname(path), exist_ok=True)
36+ with open(path, 'w', encoding='utf-8') as f:
37+ f.write('test file')
38+ 
39+ 
40+class TestLauncherUtils(unittest.TestCase):
41+ def setUp(self):
42+ self.workspace = './_test_mskl_workspace'
43+ self._cleanup()
44+ os.makedirs(self.workspace, exist_ok=True)
45+ 
46+ def tearDown(self):
47+ self._cleanup()
48+ 
49+ def _cleanup(self):
50+ if os.path.exists(self.workspace):
51+ import shutil
52+ 
53+ shutil.rmtree(self.workspace, ignore_errors=True)
54+ 
55+ def test_get_workspace_path_default_to_cwd(self):
56+ self.assertEqual(get_workspace_path(), os.getcwd())
57+ 
58+ def test_get_workspace_path_with_relative_path(self):
59+ os.makedirs(self.workspace, exist_ok=True)
60+ self.assertEqual(get_workspace_path(self.workspace), os.path.abspath(self.workspace))
61+ 
62+ def test_get_workspace_path_invalid(self):
63+ for invalid in ['', None]:
64+ if invalid is None:
65+ continue # None表示使用当前目录
66+ self.assertRaises(ValueError, get_workspace_path, invalid)
67+ 
68+ def test_get_workspace_path_not_exist(self):
69+ nonexist = os.path.join(self.workspace, 'not_exist_dir')
70+ self.assertRaises(FileNotFoundError, get_workspace_path, nonexist)
71+ 
72+ def test_search_tiling_so_found(self):
73+ so_path = os.path.join(
74+ self.workspace,
75+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
76+ 'op_impl/ai_core/tbe/op_tiling/liboptiling.so',
77+ )
78+ _touch(so_path)
79+ self.assertEqual(search_tiling_so(self.workspace), os.path.abspath(so_path))
80+ 
81+ def test_search_tiling_so_not_found(self):
82+ self.assertIsNone(search_tiling_so(self.workspace))
83+ 
84+ def test_search_tiling_so_prefer_cpack(self):
85+ # 同时存在时优先选择_CPack_Packages/op_tiling下的liboptiling.so
86+ cpack_so = os.path.join(
87+ self.workspace,
88+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
89+ 'op_impl/ai_core/tbe/op_tiling/liboptiling.so',
90+ )
91+ other_so = os.path.join(self.workspace, 'op_host/liboptiling.so')
92+ _touch(cpack_so)
93+ _touch(other_so)
94+ self.assertEqual(search_tiling_so(self.workspace), os.path.abspath(cpack_so))
95+ 
96+ def test_search_kernel_binary_files_only_in_kernel_dir(self):
97+ kernel_o = os.path.join(
98+ self.workspace,
99+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
100+ 'op_impl/ai_core/tbe/kernel/ascend910b/custom_op/OpCustom_1.o',
101+ )
102+ other_o = os.path.join(
103+ self.workspace,
104+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
105+ 'op_impl/ai_core/tbe/op_tiling/other.o',
106+ )
107+ _touch(kernel_o)
108+ _touch(other_o)
109+ result = search_kernel_binary_files(self.workspace)
110+ self.assertEqual(result, [os.path.abspath(kernel_o)])
111+ 
112+ def test_match_kernel_binary_by_op_type(self):
113+ o_files = [
114+ '/p/kernel/ascend950/add_custom/AddCustom_hash_0.o',
115+ '/p/kernel/ascend950/leaky_relu_custom/LeakyReluCustom_hash_0.o',
116+ '/p/kernel/ascend950/add_custom_template/AddCustomTemplate_hash_1.o',
117+ ]
118+ self.assertEqual(
119+ match_kernel_binary_by_op_type(o_files, 'AddCustom'),
120+ ['/p/kernel/ascend950/add_custom/AddCustom_hash_0.o'],
121+ )
122+ # 未传op_type时不进行过滤
123+ self.assertEqual(match_kernel_binary_by_op_type(o_files, None), o_files)
124+ 
125+ def test_match_kernel_binary_by_soc(self):
126+ o_files = [
127+ '/p/kernel/ascend910_93/add_custom/AddCustom_hash_0.o',
128+ '/p/kernel/ascend950/add_custom/AddCustom_hash_0.o',
129+ '/p/kernel/ascend910b/add_custom/AddCustom_hash_0.o',
130+ ]
131+ self.assertEqual(
132+ match_kernel_binary_by_soc(o_files, 'Ascend950PR_9579'),
133+ ['/p/kernel/ascend950/add_custom/AddCustom_hash_0.o'],
134+ )
135+ # 空soc名称时不匹配
136+ self.assertEqual(match_kernel_binary_by_soc(o_files, ''), [])
137+ 
138+ def test_match_json_shape(self):
139+ # 全为负值（[-2]）表示动态shape
140+ self.assertTrue(_match_json_shape([-2], [8, 2048]))
141+ # -1 表示该维度动态
142+ self.assertTrue(_match_json_shape([-1, 8], [4, 8]))
143+ self.assertTrue(_match_json_shape([-1, 8], [16, 8]))
144+ # 固定维度需要精确匹配
145+ self.assertFalse(_match_json_shape([8, 2048], [4, 8]))
146+ # rank不匹配
147+ self.assertFalse(_match_json_shape([8], [8, 2048]))
148+ 
149+ def test_match_kernel_binary_by_json(self):
150+ kernel_a = os.path.join(self.workspace, 'kernel/ascend950/add_custom_template/AddCustomTemplate_hashA.o')
151+ kernel_b = os.path.join(self.workspace, 'kernel/ascend950/add_custom_template/AddCustomTemplate_hashB.o')
152+ _touch(kernel_a)
153+ _touch(kernel_b)
154+ for kernel, dtype in ((kernel_a, 'float16'), (kernel_b, 'float32')):
155+ json_path = kernel[:-1] + 'json'
156+ meta = {
157+ 'supportInfo': {
158+ 'inputs': [
159+ {
160+ 'name': 'x',
161+ 'dtype': dtype,
162+ 'format': 'ND',
163+ 'shape': [-2],
164+ 'format_match_mode': 'FormatAgnostic',
165+ },
166+ {
167+ 'name': 'y',
168+ 'dtype': dtype,
169+ 'format': 'ND',
170+ 'shape': [-2],
171+ 'format_match_mode': 'FormatAgnostic',
172+ },
173+ ],
174+ 'outputs': [
175+ {
176+ 'name': 'z',
177+ 'dtype': dtype,
178+ 'format': 'ND',
179+ 'shape': [-2],
180+ 'format_match_mode': 'FormatAgnostic',
181+ },
182+ ],
183+ }
184+ }
185+ with open(json_path, 'w', encoding='utf-8') as f:
186+ json.dump(meta, f)
187+ 
188+ io_info = {
189+ 'inputs': [
190+ {'dtype': 'float16', 'format': 'ND', 'shape': [8, 2048]},
191+ {'dtype': 'float16', 'format': 'ND', 'shape': [8, 2048]},
192+ ],
193+ 'outputs': [{'dtype': 'float16', 'format': 'ND', 'shape': [8, 2048]}],
194+ }
195+ matched = match_kernel_binary_by_json([kernel_a, kernel_b], io_info)
196+ self.assertEqual(matched, [kernel_a])
197+ # json不存在时不参与匹配
198+ self.assertEqual(match_kernel_binary_by_json(['/p/kernel/x/no_meta.o'], io_info), [])
199+ 
200+ 
201+if __name__ == '__main__':
202+ unittest.main()
Mtest/launcher/test_opgen_workflow.py+272-26
@@ -19,16 +19,23 @@ import os
19import shutil19import shutil
20import numpy as np20import numpy as np
21import unittest21import unittest
22-from unittest.mock import patch, Mock22+from unittest.mock import patch
23 23 
24-from mskl.launcher.opgen_workflow import TilingOutput, tiling_func, get_kernel_from_binary, TMP_FOLDER24+from mskl.launcher.opgen_workflow import (
25-from mskl.utils import safe_check25+ TilingOutput,
26+ tiling_func,
27+ get_kernel_from_binary,
28+ TMP_FOLDER,
29+ _build_io_info,
30+)
31+from mskl.launcher import opgen_workflow
26from mskl.launcher.context import context32from mskl.launcher.context import context
33+from mskl.utils import safe_check
27 34 
28tiling_dict = {35tiling_dict = {
29 "blockdim": 8,36 "blockdim": 8,
30 "workspace_size": 64,37 "workspace_size": 64,
31- "tiling_data": [1,2,3,4],38+ "tiling_data": [1, 2, 3, 4],
32 "tiling_key": 1111,39 "tiling_key": 1111,
33}40}
34 41 
@@ -37,6 +44,27 @@ def mock_tiling_func():
37 return tiling_dict44 return tiling_dict
38 45 
39 46 
47+def _touch(path):
48+ os.makedirs(os.path.dirname(path), exist_ok=True)
49+ with open(path, 'w', encoding='utf-8') as f:
50+ f.write('test file')
51+ 
52+ 
53+def _build_workspace(with_tiling_so=True, kernel_o_files=(), kernel_soc='ascend910b'):
54+ ws = os.path.join(os.getcwd(), '_test_mskl_ws')
55+ if os.path.exists(ws):
56+ shutil.rmtree(ws, ignore_errors=True)
57+ base = os.path.join(
58+ ws,
59+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/op_impl/ai_core/tbe',
60+ )
61+ if with_tiling_so:
62+ _touch(os.path.join(base, 'op_tiling/liboptiling.so'))
63+ for o in kernel_o_files:
64+ _touch(os.path.join(base, 'kernel', kernel_soc, 'custom_op', o))
65+ return ws
66+ 
67+ 
40class TestOpgenWorkflow(unittest.TestCase):68class TestOpgenWorkflow(unittest.TestCase):
41 LIB = 'test_tiling.so'69 LIB = 'test_tiling.so'
42 CPP = '_mskl_gen_tiling.cpp'70 CPP = '_mskl_gen_tiling.cpp'
@@ -47,16 +75,15 @@ class TestOpgenWorkflow(unittest.TestCase):
47 75 
48 @classmethod76 @classmethod
49 def setUpClass(cls):77 def setUpClass(cls):
50- with os.fdopen(os.open(TestOpgenWorkflow.LIB,78+ with os.fdopen(
51- safe_check.OPEN_FLAGS,79+ os.open(TestOpgenWorkflow.LIB, safe_check.OPEN_FLAGS, safe_check.SAVE_DATA_FILE_AUTHORITY), 'w'
52- safe_check.SAVE_DATA_FILE_AUTHORITY),80+ ) as f:
53- 'w') as f:
54 f.truncate()81 f.truncate()
55 f.write('this is a test so')82 f.write('this is a test so')
56- with os.fdopen(os.open(TestOpgenWorkflow.KERNEL_BINARY_PATH,83+ with os.fdopen(
57- safe_check.OPEN_FLAGS,84+ os.open(TestOpgenWorkflow.KERNEL_BINARY_PATH, safe_check.OPEN_FLAGS, safe_check.SAVE_DATA_FILE_AUTHORITY),
58- safe_check.SAVE_DATA_FILE_AUTHORITY),85+ 'w',
59- 'w') as f:86+ ) as f:
60 f.truncate()87 f.truncate()
61 f.write('this is a test kernel binary file')88 f.write('this is a test kernel binary file')
62 89 
@@ -85,13 +112,17 @@ class TestOpgenWorkflow(unittest.TestCase):
85 a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)112 a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
86 b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)113 b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
87 c = np.zeros([2, 32]).astype(np.float16)114 c = np.zeros([2, 32]).astype(np.float16)
88- inputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"},115+ inputs_info = [
89- {"shape": [2, 32], "dtype": "float16", "format": "ND"}]116+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
117+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
118+ ]
90 outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]119 outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]
91 output = tiling_func(120 output = tiling_func(
92 op_type="AddCustom",121 op_type="AddCustom",
93- inputs_info=inputs_info, outputs_info=outputs_info,122+ inputs_info=inputs_info,
94- inputs=[a, b], outputs=[c],123+ outputs_info=outputs_info,
124+ inputs=[a, b],
125+ outputs=[c],
95 attr={126 attr={
96 "a1": 1,127 "a1": 1,
97 "a2": False,128 "a2": False,
@@ -101,9 +132,9 @@ class TestOpgenWorkflow(unittest.TestCase):
101 "a6": [True, False],132 "a6": [True, False],
102 "a7": ["asdf", "zxcv"],133 "a7": ["asdf", "zxcv"],
103 "a8": [[1, 2, 3, 4], [5, 6, 7, 8], [5646, 2345]],134 "a8": [[1, 2, 3, 4], [5, 6, 7, 8], [5646, 2345]],
104- "a9": [111, 222, 333]135+ "a9": [111, 222, 333],
105 },136 },
106- lib_path=TestOpgenWorkflow.LIB137+ lib_path=TestOpgenWorkflow.LIB,
107 )138 )
108 mock_compile_tiling.assert_called_once()139 mock_compile_tiling.assert_called_once()
109 self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)140 self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)
@@ -116,16 +147,21 @@ class TestOpgenWorkflow(unittest.TestCase):
116 a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)147 a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
117 b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)148 b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
118 c = np.zeros([2, 32]).astype(np.float16)149 c = np.zeros([2, 32]).astype(np.float16)
119- inputs_info = [[{"shape": [2, 32], "dtype": "float16", "format": "ND"},150+ inputs_info = [
120- {"shape": [2, 32], "dtype": "float16", "format": "ND"},151+ [
121- {"shape": [2, 32], "dtype": "float16", "format": "ND"}],152+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
122- {"shape": [2, 32], "dtype": "float16", "format": "ND"},153+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
123- ]154+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
155+ ],
156+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
157+ ]
124 outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]158 outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]
125 output = tiling_func(159 output = tiling_func(
126 op_type="AddCustom",160 op_type="AddCustom",
127- inputs_info=inputs_info, outputs_info=outputs_info,161+ inputs_info=inputs_info,
128- inputs=[[a, a, b], b], outputs=[c],162+ outputs_info=outputs_info,
163+ inputs=[[a, a, b], b],
164+ outputs=[c],
129 attr=[165 attr=[
130 {"name": "a1", "dtype": "int", "value": 1},166 {"name": "a1", "dtype": "int", "value": 1},
131 {"name": "a2", "dtype": "bool", "value": False},167 {"name": "a2", "dtype": "bool", "value": False},
@@ -142,7 +178,7 @@ class TestOpgenWorkflow(unittest.TestCase):
142 {"name": "a11", "dtype": "int64", "value": 2},178 {"name": "a11", "dtype": "int64", "value": 2},
143 {"name": "a12", "dtype": "float32", "value": 1.3},179 {"name": "a12", "dtype": "float32", "value": 1.3},
144 ],180 ],
145- lib_path=TestOpgenWorkflow.LIB181+ lib_path=TestOpgenWorkflow.LIB,
146 )182 )
147 mock_compile_tiling.assert_called_once()183 mock_compile_tiling.assert_called_once()
148 self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)184 self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)
@@ -158,3 +194,213 @@ class TestOpgenWorkflow(unittest.TestCase):
158 kernel = get_kernel_from_binary(TestOpgenWorkflow.KERNEL_BINARY_PATH, 'mix')194 kernel = get_kernel_from_binary(TestOpgenWorkflow.KERNEL_BINARY_PATH, 'mix')
159 mock_compile_kernel_binary.assert_called_once()195 mock_compile_kernel_binary.assert_called_once()
160 self.assertEqual(kernel, None)196 self.assertEqual(kernel, None)
197+ 
198+ @patch("mskl.launcher.opgen_workflow.compile_tiling")
199+ def test_tiling_func_auto_search_liboptiling_by_workspace(self, mock_compile_tiling):
200+ mock_compile_tiling.return_value = mock_tiling_func
201+ ws = _build_workspace(with_tiling_so=True)
202+ 
203+ a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
204+ b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
205+ c = np.zeros([2, 32]).astype(np.float16)
206+ inputs_info = [
207+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
208+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
209+ ]
210+ outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]
211+ try:
212+ output = tiling_func(
213+ op_type="AddCustom",
214+ inputs_info=inputs_info,
215+ outputs_info=outputs_info,
216+ inputs=[a, b],
217+ outputs=[c],
218+ workspace=ws,
219+ )
220+ mock_compile_tiling.assert_called_once()
221+ self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)
222+ self.assertEqual(context.workspace, os.path.abspath(ws))
223+ expected_so = os.path.join(
224+ ws,
225+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
226+ 'op_impl/ai_core/tbe/op_tiling/liboptiling.so',
227+ )
228+ gen_cpp = os.path.join(TMP_FOLDER, f'_mskl_gen_tiling.{opgen_workflow.TILING_FUNC_CNT}.cpp')
229+ with open(gen_cpp, 'r', encoding='utf-8') as f:
230+ self.assertIn(os.path.realpath(expected_so), f.read())
231+ finally:
232+ context.workspace = None
233+ if os.path.exists(ws):
234+ shutil.rmtree(ws, ignore_errors=True)
235+ 
236+ @patch("mskl.launcher.opgen_workflow.compile_tiling")
237+ def test_tiling_func_explicit_lib_path_priority_over_workspace(self, mock_compile_tiling):
238+ mock_compile_tiling.return_value = mock_tiling_func
239+ ws = _build_workspace(with_tiling_so=True)
240+ with open(TestOpgenWorkflow.LIB, 'w', encoding='utf-8') as f:
241+ f.write('this is a test so')
242+ 
243+ a = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
244+ b = np.random.uniform(1, 100, [2, 32]).astype(np.float16)
245+ c = np.zeros([2, 32]).astype(np.float16)
246+ inputs_info = [
247+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
248+ {"shape": [2, 32], "dtype": "float16", "format": "ND"},
249+ ]
250+ outputs_info = [{"shape": [2, 32], "dtype": "float16", "format": "ND"}]
251+ try:
252+ output = tiling_func(
253+ op_type="AddCustom",
254+ inputs_info=inputs_info,
255+ outputs_info=outputs_info,
256+ inputs=[a, b],
257+ outputs=[c],
258+ lib_path=TestOpgenWorkflow.LIB,
259+ workspace=ws,
260+ )
261+ mock_compile_tiling.assert_called_once()
262+ self.assertEqual(output.blockdim, TilingOutput(tiling_dict).blockdim)
263+ gen_cpp = os.path.join(TMP_FOLDER, f'_mskl_gen_tiling.{opgen_workflow.TILING_FUNC_CNT}.cpp')
264+ with open(gen_cpp, 'r', encoding='utf-8') as f:
265+ gen_src = f.read()
266+ self.assertIn(os.path.realpath(TestOpgenWorkflow.LIB), gen_src)
267+ self.assertNotIn('liboptiling.so', gen_src)
268+ finally:
269+ context.workspace = None
270+ if os.path.exists(ws):
271+ shutil.rmtree(ws, ignore_errors=True)
272+ 
273+ @patch("mskl.launcher.opgen_workflow.compile_kernel_binary")
274+ def test_get_kernel_from_binary_auto_search_single_o(self, mock_compile_kernel_binary):
275+ with patch('mskl.utils.launcher_utils.check_runtime_impl', return_value=True):
276+ with patch('mskl.utils.launcher_utils.get_cann_path', return_value=""):
277+ mock_compile_kernel_binary.return_value = None
278+ ws = _build_workspace(with_tiling_so=False, kernel_o_files=['TestOpType_1.o'])
279+ context.tiling_output = TilingOutput(tiling_dict)
280+ context.op_type = 'TestOpType'
281+ context.workspace = os.path.abspath(ws)
282+ try:
283+ kernel = get_kernel_from_binary(None, 'mix', 1)
284+ mock_compile_kernel_binary.assert_called_once()
285+ self.assertEqual(kernel, None)
286+ gen_cpp = os.path.join(
287+ TMP_FOLDER, f'_mskl_gen_binary_launch.{opgen_workflow.GET_KERNEL_FROM_BINARY_CNT}.cpp'
288+ )
289+ with open(gen_cpp, 'r', encoding='utf-8') as f:
290+ self.assertIn('TestOpType_1.o', f.read())
291+ finally:
292+ context.workspace = None
293+ if os.path.exists(ws):
294+ shutil.rmtree(ws, ignore_errors=True)
295+ 
296+ def test_get_kernel_from_binary_auto_search_select_by_soc(self):
297+ # 多个soc目录下存在同名同tiling_key的kernel，应优先选择与设备soc匹配的
298+ with patch('mskl.launcher.opgen_workflow.compile_kernel_binary') as mock_compile_kernel_binary:
299+ with patch('mskl.utils.launcher_utils.check_runtime_impl', return_value=True):
300+ with patch('mskl.utils.launcher_utils.get_cann_path', return_value=""):
301+ with patch('mskl.launcher.opgen_workflow.get_device_soc_name', return_value='Ascend950PR_9579'):
302+ mock_compile_kernel_binary.return_value = None
303+ ws = _build_workspace(
304+ with_tiling_so=False, kernel_o_files=['TestOpType_1.o'], kernel_soc='ascend910_93'
305+ )
306+ _build_workspace_extra = os.path.join(
307+ os.getcwd(),
308+ '_test_mskl_ws',
309+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
310+ 'op_impl/ai_core/tbe/kernel',
311+ )
312+ for soc in ('ascend950', 'ascend910b'):
313+ os.makedirs(os.path.join(_build_workspace_extra, soc, 'custom_op'), exist_ok=True)
314+ _touch(os.path.join(_build_workspace_extra, soc, 'custom_op', 'TestOpType_1.o'))
315+ context.tiling_output = TilingOutput(tiling_dict)
316+ context.op_type = 'TestOpType'
317+ context.workspace = os.path.abspath(ws)
318+ try:
319+ kernel = get_kernel_from_binary(None, 'mix', 1)
320+ mock_compile_kernel_binary.assert_called_once()
321+ self.assertEqual(kernel, None)
322+ gen_cpp = os.path.join(
323+ TMP_FOLDER,
324+ f'_mskl_gen_binary_launch.{opgen_workflow.GET_KERNEL_FROM_BINARY_CNT}.cpp',
325+ )
326+ with open(gen_cpp, 'r', encoding='utf-8') as f:
327+ self.assertIn('ascend950/custom_op/TestOpType_1.o', f.read())
328+ self.assertNotIn('ascend910_93', f.read())
329+ finally:
330+ context.workspace = None
331+ if os.path.exists(ws):
332+ shutil.rmtree(ws, ignore_errors=True)
333+ 
334+ def test_build_io_info_from_tensors(self):
335+ x = np.zeros([8, 2048], dtype=np.float16)
336+ y = np.zeros([8, 2048], dtype=np.float16)
337+ z = np.zeros([8, 2048], dtype=np.float16)
338+ info = _build_io_info([x, y], [z], None, None)
339+ self.assertEqual(info['inputs'][0]['dtype'], 'float16')
340+ self.assertEqual(info['inputs'][0]['format'], 'ND')
341+ self.assertEqual(info['inputs'][0]['shape'], [8, 2048])
342+ self.assertEqual(len(info['inputs']), 2)
343+ self.assertEqual(info['outputs'][0]['dtype'], 'float16')
344+ 
345+ @patch("mskl.launcher.opgen_workflow.compile_kernel_binary")
346+ def test_get_kernel_from_binary_auto_search_select_by_json(self, mock_compile_kernel_binary):
347+ # 多个soc相同、同算子名的kernel，根据同目录json中的dtype区分
348+ with patch('mskl.utils.launcher_utils.check_runtime_impl', return_value=True):
349+ with patch('mskl.utils.launcher_utils.get_cann_path', return_value=""):
350+ with patch('mskl.launcher.opgen_workflow.get_device_soc_name', return_value='Ascend950PR_9579'):
351+ mock_compile_kernel_binary.return_value = None
352+ ws = _build_workspace(
353+ with_tiling_so=False,
354+ kernel_soc='ascend950',
355+ kernel_o_files=['TestOpType_f16.o', 'TestOpType_f32.o'],
356+ )
357+ # 为两个 .o 生成描述不同dtype的json
358+ base = os.path.join(
359+ ws,
360+ 'build_out/_CPack_Packages/Linux/External/custom.run/packages/vendors/customize/'
361+ 'op_impl/ai_core/tbe/kernel/ascend950/custom_op',
362+ )
363+ for o_name, dtype in (('TestOpType_f16.o', 'float16'), ('TestOpType_f32.o', 'float32')):
364+ meta = {
365+ 'supportInfo': {
366+ 'inputs': [{'dtype': dtype, 'format': 'ND', 'shape': [-2]}],
367+ 'outputs': [{'dtype': dtype, 'format': 'ND', 'shape': [-2]}],
368+ }
369+ }
370+ with open(os.path.join(base, o_name[:-1] + 'json'), 'w', encoding='utf-8') as f:
371+ import json
372+ 
373+ json.dump(meta, f)
374+ context.tiling_output = TilingOutput(tiling_dict)
375+ context.op_type = 'TestOpType'
376+ context.workspace = os.path.abspath(ws)
377+ context.io_info = {
378+ 'inputs': [{'dtype': 'float16', 'format': 'ND', 'shape': [8, 2048]}],
379+ 'outputs': [{'dtype': 'float16', 'format': 'ND', 'shape': [8, 2048]}],
380+ }
381+ try:
382+ kernel = get_kernel_from_binary(None, 'mix', 1)
383+ mock_compile_kernel_binary.assert_called_once()
384+ self.assertEqual(kernel, None)
385+ gen_cpp = os.path.join(
386+ TMP_FOLDER, f'_mskl_gen_binary_launch.{opgen_workflow.GET_KERNEL_FROM_BINARY_CNT}.cpp'
387+ )
388+ with open(gen_cpp, 'r', encoding='utf-8') as f:
389+ self.assertIn('TestOpType_f16.o', f.read())
390+ finally:
391+ context.workspace = None
392+ context.io_info = None
393+ if os.path.exists(ws):
394+ shutil.rmtree(ws, ignore_errors=True)
395+ 
396+ def test_get_kernel_from_binary_auto_search_no_o_found(self):
397+ ws = _build_workspace(with_tiling_so=False)
398+ context.tiling_output = TilingOutput(tiling_dict)
399+ context.op_type = 'TestOpType'
400+ context.workspace = os.path.abspath(ws)
401+ try:
402+ self.assertRaises(FileNotFoundError, get_kernel_from_binary, None, 'mix', 1)
403+ finally:
404+ context.workspace = None
405+ if os.path.exists(ws):
406+ shutil.rmtree(ws, ignore_errors=True)