[WIP]Draft: perf(dsl) 单算子 RUN 包支持 Native AOT 并行-ops-transformer-zyh-AtomGit

Mattention/mixed_quant_sparse_flash_mla/op_dsl/CMakeLists.txt+1-1

	@@ -38,5 +38,5 @@ if find_spec("cannbotdsl") is not None and find_spec("opkit") is not None:
38	endif()	38	endif()
39		39
40	if(BUILD_OPEN_PROJECT)	40	if(BUILD_OPEN_PROJECT)
41	- enable_cannbotdsl_kernel(mixed_quant_sparse_flash_mla "${CMAKE_CURRENT_LIST_DIR}")	41	+ enable_cannbotdsl_kernel(mixed_quant_sparse_flash_mla "${CMAKE_CURRENT_LIST_DIR}" PARALLEL)
42	endif()	42	endif()

Mattention/mixed_quant_sparse_flash_mla/op_dsl/mixed_quant_sparse_flash_mla.py+52-69

 import math
 import os
 import threading
--from itertools import product
++from itertools import islice, product
 from types import SimpleNamespace
 import cannbotdsl
 import torch
     return out, lse
++def _aot_profiles():
++    for tnd_padding, has_sinks, head_bucket in product((False, True), (True, False), (64, 32, 63)):
++        for has_cmp in (True, False):
++            cmp_lengths = (True, False) if has_cmp else (False,)
++            for has_ori_length, has_cmp_length, has_lse, addr_mode in product(
++                (True, False), cmp_lengths, (True, False), ("pre", "off", "arena")
++            ):
++                for pa_fast in ((False,) if addr_mode == "off" else (False, True)):
++                    yield dict(
++                        has_cmp=has_cmp,
++                        has_ori_length=has_ori_length,
++                        has_cmp_length=has_cmp_length,
++                        has_lse=has_lse,
++                        addr_mode=addr_mode,
++                        layout_kv="PA_BBND",
++                        pa_fast=pa_fast,
++                        head_bucket=head_bucket,
++                        has_sinks=has_sinks,
++                        tnd_padding=tnd_padding,
++                    )
++        for has_cmp in (True, False):
++            cmp_lengths = (True, False) if has_cmp else (False,)
++            for has_ori_length, has_cmp_length, has_lse, layout_kv in product(
++                (True, False), cmp_lengths, (True, False), ("TND", "BSND")
++            ):
++                yield dict(
++                    has_cmp=has_cmp,
++                    has_ori_length=has_ori_length,
++                    has_cmp_length=has_cmp_length,
++                    has_lse=has_lse,
++                    addr_mode="off",
++                    layout_kv=layout_kv,
++                    head_bucket=head_bucket,
++                    has_sinks=has_sinks,
++                    tnd_padding=tnd_padding,
++                )
++
++
++def _export_mqsmla_shard(index, count):
++    clear_caches()
++    try:
++        for profile in islice(_aot_profiles(), index, None, count):
++            _get_compiled_variant(**profile)
++    finally:
++        clear_caches()
++
++
 @aot.export("mixed_quant_sparse_flash_mla_dsl")
 def export_mqsmla():
     """Collect all optional-input/layout/addressing variants with dynamic shapes."""
--    clear_caches()
++    _export_mqsmla_shard(0, 1)
--    try:
++
--        for tnd_padding in (False, True):
++
--            for has_sinks in (True, False):
++export_mqsmla.__transformer_aot_shard__ = _export_mqsmla_shard
--                for head_bucket in (64, 32, 63):
--                    for has_cmp in (True, False):
--                        cmp_lengths = (True, False) if has_cmp else (False,)
--                        for (
--                            has_ori_length,
--                            has_cmp_length,
--                            has_lse,
--                            addr_mode,
--                        ) in product(
--                            (True, False),
--                            cmp_lengths,
--                            (True, False),
--                            ("pre", "off", "arena"),
--                        ):
--                            _get_compiled_variant(
--                                has_cmp,
--                                has_ori_length,
--                                has_cmp_length,
--                                has_lse,
--                                addr_mode,
--                                layout_kv="PA_BBND",
--                                head_bucket=head_bucket,
--                                has_sinks=has_sinks,
--                                tnd_padding=tnd_padding,
--                            )
--                            if addr_mode != "off":
--                                _get_compiled_variant(
--                                    has_cmp,
--                                    has_ori_length,
--                                    has_cmp_length,
--                                    has_lse,
--                                    addr_mode,
--                                    layout_kv="PA_BBND",
--                                    head_bucket=head_bucket,
--                                    has_sinks=has_sinks,
--                                    tnd_padding=tnd_padding,
--                                    pa_fast=True,
--                                )
--                    for has_cmp in (True, False):
--                        cmp_lengths = (True, False) if has_cmp else (False,)
--                        for (
--                            has_ori_length,
--                            has_cmp_length,
--                            has_lse,
--                            layout_kv,
--                        ) in product(
--                            (True, False),
--                            cmp_lengths,
--                            (True, False),
--                            ("TND", "BSND"),
--                        ):
--                            _get_compiled_variant(
--                                has_cmp,
--                                has_ori_length,
--                                has_cmp_length,
--                                has_lse,
--                                "off",
--                                layout_kv,
--                                head_bucket=head_bucket,
--                                has_sinks=has_sinks,
--                                tnd_padding=tnd_padding,
--                            )
--    finally:
--        clear_caches()

Mbuild.sh+4-0

	@@ -2221,6 +2221,10 @@ if [ -n "${INPUT_JOB_NUM}" ]; then
2221	JOB_NUM="-j${INPUT_JOB_NUM}"	2221	JOB_NUM="-j${INPUT_JOB_NUM}"
2222	fi	2222	fi
2223		2223
		2224	+if [[ "$ENABLE_BUILD_PKG" == "TRUE" && "${JOB_NUM}" =~ ^-j([0-9]+)$ ]]; then
		2225	+ CUSTOM_OPTION="${CUSTOM_OPTION} -DCANNBOTDSL_BUILD_JOBS=${BASH_REMATCH[1]}"
		2226	+fi
		2227	+
2224	# 非打包命令调用，打包模式会打进同一个包里	2228	# 非打包命令调用，打包模式会打进同一个包里
2225	function set_compute_unit_option() {	2229	function set_compute_unit_option() {
2226	# 使用共享校验库逐项校验 SoC 列表(支持逗号/分号分隔的多个 soc), 替换原有的整串精确匹配	2230	# 使用共享校验库逐项校验 SoC 列表(支持逗号/分号分隔的多个 soc), 替换原有的整串精确匹配

Mcmake/func.cmake+8-2

	@@ -804,6 +804,12 @@ function(enable_cannbotdsl_kernel op_name source_dir)
804	file(STRINGS "${_root}/inputs.list" _dependencies ENCODING UTF-8)	804	file(STRINGS "${_root}/inputs.list" _dependencies ENCODING UTF-8)
805	set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${_dependencies}")	805	set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${_dependencies}")
806	include("${_root}/inputs.cmake")	806	include("${_root}/inputs.cmake")
		807	+ set(_native_build_args)
		808	+ # Only a single selected operator owns this build's DSL worker budget.
		809	+ if(ENABLE_BUILD_PKG AND _dsl_native AND NOT _dsl_aclnn AND "${ASCEND_OP_NAME}" STREQUAL "${op_name}" AND
		810	+ "${ARGV2}" STREQUAL "PARALLEL" AND CANNBOTDSL_BUILD_JOBS GREATER 1)
		811	+ set(_native_build_args --jobs "${CANNBOTDSL_BUILD_JOBS}")
		812	+ endif()
807	if(NOT _dsl_native AND NOT _dsl_aclnn)	813	if(NOT _dsl_native AND NOT _dsl_aclnn)
808	message(FATAL_ERROR "No Native or ACLNN export found in ${_source}")	814	message(FATAL_ERROR "No Native or ACLNN export found in ${_source}")
809	endif()	815	endif()
	@@ -826,7 +832,7 @@ function(enable_cannbotdsl_kernel op_name source_dir)
826	if(ENABLE_OPS_HOST AND _dsl_native)	832	if(ENABLE_OPS_HOST AND _dsl_native)
827	# Imported archive targets and runtime sources must exist at configure time.	833	# Imported archive targets and runtime sources must exist at configure time.
828	execute_process(COMMAND "${CMAKE_COMMAND}" -E env ${_env}	834	execute_process(COMMAND "${CMAKE_COMMAND}" -E env ${_env}
829	- "${HI_PYTHON}" "${_builder}" ${_target_args} --mode host --source "${_source}" --output "${_root}/host"	835	+ "${HI_PYTHON}" "${_builder}" ${_target_args} ${_native_build_args} --mode host --source "${_source}" --output "${_root}/host"
830	--inputs "${_inputs}"	836	--inputs "${_inputs}"
831	RESULT_VARIABLE _result)	837	RESULT_VARIABLE _result)
832	if(NOT _result EQUAL 0)	838	if(NOT _result EQUAL 0)
	@@ -876,7 +882,7 @@ function(enable_cannbotdsl_kernel op_name source_dir)
876	set(_stamp "${_root}/kernel.built")	882	set(_stamp "${_root}/kernel.built")
877	add_custom_command(OUTPUT "${_stamp}"	883	add_custom_command(OUTPUT "${_stamp}"
878	COMMAND "${CMAKE_COMMAND}" -E env ${_env}	884	COMMAND "${CMAKE_COMMAND}" -E env ${_env}
879	- "${HI_PYTHON}" "${_builder}" ${_target_args} --mode kernel --source "${_source}"	885	+ "${HI_PYTHON}" "${_builder}" ${_target_args} ${_native_build_args} --mode kernel --source "${_source}"
880	--output "${_root}/device" --kernel-output "${ASCEND_BINARY_OUT_DIR}/${CANNBOTDSL_COMPUTE_UNIT}/bin"	886	--output "${_root}/device" --kernel-output "${ASCEND_BINARY_OUT_DIR}/${CANNBOTDSL_COMPUTE_UNIT}/bin"
881	--inputs "${_inputs}"	887	--inputs "${_inputs}"
882	COMMAND "${CMAKE_COMMAND}" -E touch "${_stamp}"	888	COMMAND "${CMAKE_COMMAND}" -E touch "${_stamp}"

Mcmake/scripts/build_cannbotdsl.py+97-15

	@@ -13,16 +13,19 @@ import argparse
13	import fcntl	13	import fcntl
14	import hashlib	14	import hashlib
15	import json	15	import json
		16	+import multiprocessing
16	import os	17	import os
17	import re	18	import re
18	import shutil	19	import shutil
19	import subprocess	20	import subprocess
20	import sys	21	import sys
21	import tempfile	22	import tempfile
		23	+from concurrent.futures import ProcessPoolExecutor
22	from contextlib import contextmanager	24	from contextlib import contextmanager
23	from importlib import import_module	25	from importlib import import_module
24	from importlib.metadata import PackageNotFoundError, distribution	26	from importlib.metadata import PackageNotFoundError, distribution
25	from pathlib import Path	27	from pathlib import Path
		28	+from types import SimpleNamespace
26		29
27		30
28	def digest(path):	31	def digest(path):
	@@ -316,6 +319,77 @@ def validate_aclnn_host(root, aclnn_type, target="dav-3510"):
316	raise ValueError(f"Missing ACLNN module sources: {name}")	319	raise ValueError(f"Missing ACLNN module sources: {name}")
317		320
318		321
		322	+def compile_native_shard(source, target, mode, directory, index, count):
		323	+ from cannbotdsl.aot.collection import collect, discover
		324	+ from cannbotdsl.core.compiler.builder import lower_to_source
		325	+ from opkit.api import _compilation_symbols
		326	+ from opkit.device import DevicePublication
		327	+
		328	+ device = DevicePublication(directory, target)
		329	+ device._native_host_abi = mode == "host"
		330	+ try:
		331	+ with _compilation_symbols(), discover(collectors=source.parent, pattern=source.name) as tasks:
		332	+ with lower_to_source(), collect(device):
		333	+ for name, callback in tasks.items():
		334	+ device.task = name
		335	+ callback.__transformer_aot_shard__(index, count)
		336	+ if mode == "kernel":
		337	+ for key, record in device.records.items():
		338	+ device.modules[key] = device._compile(key, record["operator"])
		339	+ return device.records, device.modules
		340	+ except Exception as exc:
		341	+ raise RuntimeError(f"{source.name} {mode} shard {index + 1}/{count} failed") from exc
		342	+ finally:
		343	+ device.close()
		344	+
		345	+
		346	+def build_native_parallel(source, output, target, mode, jobs):
		347	+ # Use independent processes because CANNBotDSL collection owns its process
		348	+ # and does not support concurrent threads or forked compiler state.
		349	+ from opkit.device import DevicePublication
		350	+ from opkit.package import Package, _publish_native
		351	+
		352	+ jobs = min(jobs, 32)
		353	+ print(f"CANNBotDSL {source.stem} {mode}: compiling with {jobs} workers", flush=True)
		354	+ with tempfile.TemporaryDirectory(prefix=".shards-", dir=output.parent) as scratch:
		355	+ scratch = Path(scratch)
		356	+ device = DevicePublication(scratch / "device", target)
		357	+ device.root.mkdir()
		358	+ try:
		359	+ with ProcessPoolExecutor(max_workers=jobs, mp_context=multiprocessing.get_context("spawn")) as pool:
		360	+ futures = [
		361	+ pool.submit(compile_native_shard, source, target, mode, scratch / str(index), index, jobs)
		362	+ for index in range(jobs)
		363	+ ]
		364	+ for index, future in enumerate(futures):
		365	+ records, modules = future.result()
		366	+ for key, record in records.items():
		367	+ device.task = record["operator"]
		368	+ for origin in record["sources"]:
		369	+ function = SimpleNamespace(__qualname__=origin["function"], __module__=origin["module"])
		370	+ device.record(record["artifact"], function)
		371	+ if key in modules and key not in device.modules:
		372	+ module = modules[key]
		373	+ for field in ("binary", "metadata"):
		374	+ relative = module[field]
		375	+ destination = device.root / relative
		376	+ destination.parent.mkdir(parents=True, exist_ok=True)
		377	+ shutil.copy2(scratch / str(index) / relative, destination)
		378	+ device.modules[key] = module
		379	+ if not device.records:
		380	+ raise ValueError(f"{source.name} collected no DSL compile requests")
		381	+ output.mkdir()
		382	+ if mode == "host":
		383	+ _publish_native(output, device, mode)
		384	+ else:
		385	+ package = Package(output)
		386	+ device.publish_all(output)
		387	+ package.finish(dict(target=target, tasks=sorted(device.tasks), binary_format="device", mode=mode, aclnn={}))
		388	+ print(f"CANNBotDSL {source.stem} {mode}: published {len(device.records)} profiles", flush=True)
		389	+ finally:
		390	+ device.close()
		391	+
		392	+
319	def build(	393	def build(
320	source,	394	source,
321	output,	395	output,
	@@ -325,6 +399,7 @@ def build(
325	inputs=None,	399	inputs=None,
326	aclnn_type=None,	400	aclnn_type=None,
327	target="dav-3510",	401	target="dav-3510",
		402	+ jobs=1,
328	):	403	):
329	import opkit	404	import opkit
330		405
	@@ -350,21 +425,24 @@ def build(
350	l2=aclnn_type == "aclnn",	425	l2=aclnn_type == "aclnn",
351	)	426	)
352	)	427	)
353	- opkit.build(	428	+ if jobs > 1 and aclnn_type is None:
354	- collectors=source.parent,	429	+ build_native_parallel(source, stage, target, mode, jobs)
355	- # Only ACLNN Host publication collects the L2 declaration.	430	+ else:
356	- # Native Host and device profiles come from the implementation.	431	+ opkit.build(
357	- pattern=(	432	+ collectors=source.parent,
358	- "*.py"	433	+ # Only ACLNN Host publication collects the L2 declaration.
359	- if source.with_name("aclnn_" + source.name).is_file()	434	+ # Native Host and device profiles come from the implementation.
360	- and aclnn_type is not None	435	+ pattern=(
361	- else source.name	436	+ "*.py"
362	- ),	437	+ if source.with_name("aclnn_" + source.name).is_file()
363	- target=target,	438	+ and aclnn_type is not None
364	- output=stage,	439	+ else source.name
365	- mode=mode,	440	+ ),
366	- **options,	441	+ target=target,
367	- )	442	+ output=stage,
		443	+ mode=mode,
		444	+ **options,
		445	+ )
368	if aclnn_type is None and mode == "host":	446	if aclnn_type is None and mode == "host":
369	extract_host_objects(stage)	447	extract_host_objects(stage)
370	if aclnn_type is not None and mode == "host":	448	if aclnn_type is not None and mode == "host":
	@@ -387,6 +465,7 @@ def main():
387	parser.add_argument("--output", type=Path)	465	parser.add_argument("--output", type=Path)
388	parser.add_argument("--soc", default="ascend950")	466	parser.add_argument("--soc", default="ascend950")
389	parser.add_argument("--target", default="dav-3510")	467	parser.add_argument("--target", default="dav-3510")
		468	+ parser.add_argument("--jobs", type=int, default=1)
390	parser.add_argument("--kernel-output", type=Path)	469	parser.add_argument("--kernel-output", type=Path)
391	parser.add_argument("--packages", nargs="*", type=Path, default=[])	470	parser.add_argument("--packages", nargs="*", type=Path, default=[])
392	parser.add_argument("--inputs", type=Path)	471	parser.add_argument("--inputs", type=Path)
	@@ -395,6 +474,8 @@ def main():
395	"--aclnn-type", choices=("aclnn", "aclnn_inner", "aclnn_exclude")	474	"--aclnn-type", choices=("aclnn", "aclnn_inner", "aclnn_exclude")
396	)	475	)
397	args = parser.parse_args()	476	args = parser.parse_args()
		477	+ if args.jobs < 1:
		478	+ parser.error("--jobs must be positive")
398	if args.output is None:	479	if args.output is None:
399	parser.error("--output is required for publication modes")	480	parser.error("--output is required for publication modes")
400	if args.mode == "inputs":	481	if args.mode == "inputs":
	@@ -421,6 +502,7 @@ def main():
421	args.inputs,	502	args.inputs,
422	args.aclnn_type,	503	args.aclnn_type,
423	args.target,	504	args.target,
		505	+ args.jobs,
424	)	506	)
425		507
426		508