feat(dsl): support user-provided Ascend C functions-catlass-AtomGit

Mpython/tla_dsl/catlass/_mlir_bindings/tla_ops_gen.py+38-0

	@@ -403,6 +403,44 @@ class BlockNumOp(_ods_ir.OpView):
403	def arch_block_num(num, *, loc=None, ip=None) -> _ods_ir.Value:	403	def arch_block_num(num, *, loc=None, ip=None) -> _ods_ir.Value:
404	return _get_op_result_or_op_results(BlockNumOp(num=num, loc=loc, ip=ip))	404	return _get_op_result_or_op_results(BlockNumOp(num=num, loc=loc, ip=ip))
405		405
		406	+@_ods_cext.register_operation(_Dialect)
		407	+class CallExternOp(_ods_ir.OpView):
		408	+ OPERATION_NAME = "tla.call_extern"
		409	+
		410	+ _ODS_REGIONS = (0, True)
		411	+
		412	+ def __init__(self, callee, operands_, *, loc=None, ip=None):
		413	+ operands = []
		414	+ results = []
		415	+ attributes = {}
		416	+ regions = None
		417	+ operands.extend(_get_op_results_or_values(operands_))
		418	+ _ods_context = _ods_get_default_loc_context(loc)
		419	+ attributes["callee"] = (callee if (
		420	+ isinstance(callee, _ods_ir.Attribute) or
		421	+ not _ods_ir.AttrBuilder.contains('FlatSymbolRefAttr')) else
		422	+ _ods_ir.AttrBuilder.get('FlatSymbolRefAttr')(callee, context=_ods_context))
		423	+ _ods_successors = None
		424	+ super().__init__(self.build_generic(attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip))
		425	+
		426	+ @builtins.property
		427	+ def operands_(self):
		428	+ _ods_variadic_group_length = len(self.operation.operands) - 1 + 1
		429	+ return self.operation.operands[0:0 + _ods_variadic_group_length]
		430	+
		431	+ @builtins.property
		432	+ def callee(self):
		433	+ return self.operation.attributes["callee"]
		434	+
		435	+ @callee.setter
		436	+ def callee(self, value):
		437	+ if value is None:
		438	+ raise ValueError("'None' not allowed as value for mandatory attributes")
		439	+ self.operation.attributes["callee"] = value
		440	+
		441	+def call_extern(callee, operands_, *, loc=None, ip=None) -> _ods_ir.Operation:
		442	+ return _get_op_result_or_op_results(CallExternOp(callee=callee, operands_=operands_, loc=loc, ip=ip))
		443	+
406	@_ods_cext.register_operation(_Dialect)	444	@_ods_cext.register_operation(_Dialect)
407	class CastOp(_ods_ir.OpView):	445	class CastOp(_ods_ir.OpView):
408	OPERATION_NAME = "tla.cast"	446	OPERATION_NAME = "tla.cast"

Mpython/tla_dsl/catlass/base_dsl/op.py+4-0

	@@ -38,6 +38,10 @@ class _FrontendEmitState:
38	active_regions: list[str] = field(default_factory=list)	38	active_regions: list[str] = field(default_factory=list)
39	#: Stack of ``mode`` values for the enclosing ``tla.vec.func`` regions.	39	#: Stack of ``mode`` values for the enclosing ``tla.vec.func`` regions.
40	vec_func_modes: list[str] = field(default_factory=list)	40	vec_func_modes: list[str] = field(default_factory=list)
		41	+ #: The one external function used by this kernel in the v1 implementation.
		42	+ extern_function: Any \| None = None
		43	+ #: Core types from call sites of ``extern_function``.
		44	+ extern_core_types: set[str] = field(default_factory=set)
41		45
42		46
43	_FRONTEND_EMIT_STATE: contextvars.ContextVar[_FrontendEmitState \| None] = (	47	_FRONTEND_EMIT_STATE: contextvars.ContextVar[_FrontendEmitState \| None] = (

Mpython/tla_dsl/catlass/base_dsl/typing.py+45-0

	@@ -6,6 +6,7 @@ import ctypes
6	import operator	6	import operator
7	import struct	7	import struct
8	from abc import ABC, abstractmethod	8	from abc import ABC, abstractmethod
		9	+from dataclasses import dataclass
9	from typing import (	10	from typing import (
10	Any,	11	Any,
11	Callable,	12	Callable,
	@@ -21,6 +22,7 @@ from typing import (
21	import numpy as np	22	import numpy as np
22	from mlir import ir as mlir_ir # type: ignore[assignment]	23	from mlir import ir as mlir_ir # type: ignore[assignment]
23		24
		25	+from ..address_space import AddressSpace
24	from .op import (	26	from .op import (
25	dsl_user_op,	27	dsl_user_op,
26	_bind_frontend_category,	28	_bind_frontend_category,
	@@ -1101,9 +1103,51 @@ class JitArgument(Protocol):
1101	raise NotImplementedError	1103	raise NotImplementedError
1102		1104
1103		1105
		1106	+@dataclass(frozen=True)
		1107	+class TypedPointer:
		1108	+ """Type descriptor for a pointer element type and memory space.
		1109	+
		1110	+ ``Pointer[dtype, space]`` returns a ``TypedPointer`` for APIs that need a
		1111	+ compile-time pointer type, such as an external function ABI. It is not a
		1112	+ runtime or SSA pointer value.
		1113	+
		1114	+ Args:
		1115	+ dtype: Concrete :class:`Numeric` element type.
		1116	+ space: Target memory space as a
		1117	+ :class:`~catlass.address_space.AddressSpace` member.
		1118	+ """
		1119	+
		1120	+ dtype: type[Numeric]
		1121	+ space: AddressSpace
		1122	+
		1123	+ def __post_init__(self) -> None:
		1124	+ # Validate `dtype` and `space`
		1125	+ if (
		1126	+ not isinstance(self.dtype, type)
		1127	+ or not issubclass(self.dtype, Numeric)
		1128	+ or not self.dtype.dtype
		1129	+ ):
		1130	+ raise TypeError(
		1131	+ "TypedPointer dtype must be a concrete Numeric type, "
		1132	+ f"got {self.dtype!r}"
		1133	+ )
		1134	+ if not isinstance(self.space, AddressSpace):
		1135	+ raise TypeError(
		1136	+ f"TypedPointer memory space must be AddressSpace, got {self.space!r}"
		1137	+ )
		1138	+
		1139	+ def __repr__(self) -> str:
		1140	+ return f"TypedPointer[{self.dtype}, {self.space}]"
		1141	+
		1142	+
1104	class Pointer(ABC):	1143	class Pointer(ABC):
1105	"""Abstract JIT pointer (typed ``Pointer`` protocol)."""	1144	"""Abstract JIT pointer (typed ``Pointer`` protocol)."""
1106		1145
		1146	+ def __class_getitem__(cls, args: Any) -> TypedPointer:
		1147	+ if not isinstance(args, tuple) or len(args) != 2:
		1148	+ raise TypeError("Pointer[...] expects (dtype, memory_space)")
		1149	+ return TypedPointer(*args)
		1150	+
1107	@property	1151	@property
1108	@abstractmethod	1152	@abstractmethod
1109	def dtype(self) -> Any: ...	1153	def dtype(self) -> Any: ...
	@@ -1145,6 +1189,7 @@ __all__ = [
1145	"BFloat16",	1189	"BFloat16",
1146	"Float32",	1190	"Float32",
1147	"Pointer",	1191	"Pointer",
		1192	+ "TypedPointer",
1148	"Constexpr",	1193	"Constexpr",
1149	"JitArgument",	1194	"JitArgument",
1150	]	1195	]

Mpython/tla_dsl/catlass/core_api.py+120-2

	@@ -8,7 +8,7 @@ import math
8	import sys	8	import sys
9	from enum import Enum	9	from enum import Enum
10	from itertools import chain	10	from itertools import chain
11	-from typing import Any, Callable, Iterable, NoReturn, Sequence, TypeAlias	11	+from typing import TYPE_CHECKING, Any, Callable, Iterable, NoReturn, Sequence, TypeAlias
12		12
13	from mlir import ir as mlir_ir # type: ignore[assignment]	13	from mlir import ir as mlir_ir # type: ignore[assignment]
14	from mlir._mlir_libs._mlir import ( # type: ignore[import-not-found]	14	from mlir._mlir_libs._mlir import ( # type: ignore[import-not-found]
	@@ -26,7 +26,7 @@ from ._mlir_bindings import tla_ops_gen as _tla_ops_gen
26	from .base_dsl import ast_helpers as _ast_helpers	26	from .base_dsl import ast_helpers as _ast_helpers
27	from .base_dsl.op import dsl_user_op, _capture_user_loc	27	from .base_dsl.op import dsl_user_op, _capture_user_loc
28	from .base_dsl.typing import Bool, Float32, Int8, Int32, Numeric, as_numeric	28	from .base_dsl.typing import Bool, Float32, Int8, Int32, Numeric, as_numeric
29	-from .base_dsl.typing import Pointer	29	+from .base_dsl.typing import Pointer, TypedPointer
30	from .tla.tensor import normalize_tile_view_coord	30	from .tla.tensor import normalize_tile_view_coord
31	from .tla.typing import Tensor	31	from .tla.typing import Tensor
32	from . import runtime as _runtime	32	from . import runtime as _runtime
	@@ -7443,6 +7443,124 @@ def update_mask(
7443	return MaskSSA(mask_value), as_numeric(new_true_shape)	7443	return MaskSSA(mask_value), as_numeric(new_true_shape)
7444		7444
7445		7445
		7446	+# Keep this type-only: tla.ffi lazily calls back into core_api, so a
		7447	+# runtime import here would add an unnecessary reverse dependency.
		7448	+if TYPE_CHECKING:
		7449	+ from .tla.ffi import ExternFunction
		7450	+
		7451	+
		7452	+def _emit_extern_call(
		7453	+ extern_function: ExternFunction,
		7454	+ args: tuple[object, ...],
		7455	+ *,
		7456	+ loc: mlir_ir.Location \| None = None,
		7457	+) -> None:
		7458	+ """Validate and emit a call to a declared external function."""
		7459	+
		7460	+ _require_frontend_state("extern")
		7461	+ in_vector = _runtime._has_enclosing_region("vector")
		7462	+ in_cube = _runtime._has_enclosing_region("cube")
		7463	+ # Validate: scope, args count.
		7464	+ if in_vector == in_cube:
		7465	+ _op_error(
		7466	+ "extern",
		7467	+ "call must be nested inside exactly one of tla.vector() or tla.cube()",
		7468	+ )
		7469	+ if in_vector and _runtime._has_enclosing_region("vec.func"):
		7470	+ _op_error("extern", "call must be outside tla.vec.func()")
		7471	+ core_type = "aiv" if in_vector else "aic"
		7472	+ if len(args) != len(extern_function.arg_types):
		7473	+ _op_error(
		7474	+ "extern",
		7475	+ f"{extern_function.symbol} expects {len(extern_function.arg_types)} arguments, got {len(args)}",
		7476	+ )
		7477	+
		7478	+ state = _runtime._current_frontend_state()
		7479	+ # Validate: only 1 external function per kernel; same function can be called multiple times.
		7480	+ assert state is not None
		7481	+ if state.extern_function is None:
		7482	+ state.extern_function = extern_function
		7483	+ elif state.extern_function is not extern_function:
		7484	+ _op_error(
		7485	+ "extern",
		7486	+ "v1 supports at most one external function per kernel; "
		7487	+ "the same function may be called multiple times",
		7488	+ )
		7489	+ state.extern_core_types.add(core_type)
		7490	+
		7491	+ operands: list[mlir_ir.Value] = []
		7492	+ for position, (arg, expected) in enumerate(
		7493	+ zip(args, extern_function.arg_types, strict=True)
		7494	+ ):
		7495	+ if isinstance(expected, TypedPointer):
		7496	+ # Pointer case: get and validate ptr type, pointee type and address space.
		7497	+ resolved = _resolve_bound_value(arg)
		7498	+ if isinstance(arg, Tensor) or (
		7499	+ isinstance(resolved, mlir_ir.Value)
		7500	+ and _tla_type_bridge.type_is_tensor(resolved.type)
		7501	+ ):
		7502	+ _op_error(
		7503	+ "extern",
		7504	+ f"argument {position} of {extern_function.symbol} expects a Pointer, "
		7505	+ "but received a Tensor; pass tensor.ptr explicitly",
		7506	+ )
		7507	+ value = _as_value(arg)
		7508	+ ptr_type = PtrType.try_cast(value.type)
		7509	+ if ptr_type is None:
		7510	+ _op_error(
		7511	+ "extern",
		7512	+ f"argument {position} of {extern_function.symbol} must be "
		7513	+ f"!tla.ptr<{expected.dtype.dtype}, {expected.space.name}, ...>; "
		7514	+ f"got {value.type}",
		7515	+ )
		7516	+ actual_dtype = Numeric.from_mlir_type(ptr_type.pointee)
		7517	+ actual_space = AddressSpace.from_mlir_token(ptr_type.addrspace)
		7518	+ if actual_dtype is not expected.dtype or actual_space is not expected.space:
		7519	+ _op_error(
		7520	+ "extern",
		7521	+ f"argument {position} of {extern_function.symbol} expects pointer "
		7522	+ f"to {expected.dtype.__name__} in {expected.space.name}, "
		7523	+ f"got {actual_dtype.__name__} in {actual_space.name}",
		7524	+ )
		7525	+ else:
		7526	+ # Numeric case: get and validate numeric type.
		7527	+ resolved = _resolve_bound_value(arg)
		7528	+ if isinstance(resolved, Numeric):
		7529	+ value = resolved.ir_value(loc=loc)
		7530	+ elif isinstance(resolved, mlir_ir.Value):
		7531	+ value = resolved
		7532	+ elif isinstance(resolved, (bool, int, float)):
		7533	+ value = expected(resolved).ir_value(loc=loc)
		7534	+ else:
		7535	+ _op_error(
		7536	+ "extern",
		7537	+ f"argument {position} of {extern_function.symbol} must be "
		7538	+ f"{expected.__name__}, got {_type_name(arg)}",
		7539	+ )
		7540	+ if isinstance(value.type, mlir_ir.IndexType):
		7541	+ _op_error(
		7542	+ "extern",
		7543	+ f"argument {position} of {extern_function.symbol} must not be index; "
		7544	+ f"declare and pass Int32 or Int64 explicitly",
		7545	+ )
		7546	+ try:
		7547	+ actual_dtype = Numeric.from_mlir_type(value.type)
		7548	+ except TypeError:
		7549	+ _op_error(
		7550	+ "extern",
		7551	+ f"argument {position} of {extern_function.symbol} must be {expected.__name__}, got {value.type}",
		7552	+ )
		7553	+ if actual_dtype is not expected:
		7554	+ _op_error(
		7555	+ "extern",
		7556	+ f"argument {position} of {extern_function.symbol} expects "
		7557	+ f"{expected.__name__}, got {actual_dtype.__name__}",
		7558	+ )
		7559	+ operands.append(value)
		7560	+
		7561	+ _tla_ops_gen.call_extern(extern_function.symbol, operands, loc=loc)
		7562	+
		7563	+
7446	_mask_namespace = _Namespace()	7564	_mask_namespace = _Namespace()
7447	for _mask_pattern_token in _MASK_PATTERN_TOKENS:	7565	for _mask_pattern_token in _MASK_PATTERN_TOKENS:
7448	_mask_namespace._set(_mask_pattern_token, _MaskPattern(_mask_pattern_token))	7566	_mask_namespace._set(_mask_pattern_token, _MaskPattern(_mask_pattern_token))

Mpython/tla_dsl/catlass/execution.py+196-5

	@@ -16,7 +16,7 @@ import sys
16	import tempfile	16	import tempfile
17	import threading	17	import threading
18	from pathlib import Path	18	from pathlib import Path
19	-from typing import Any, Callable, Iterable, Mapping, Sequence	19	+from typing import TYPE_CHECKING, Any, Callable, Iterable, Mapping, Sequence
20		20
21	from .base_dsl.arch import (	21	from .base_dsl.arch import (
22	DEFAULT_NPU_ARCH,	22	DEFAULT_NPU_ARCH,
	@@ -44,6 +44,9 @@ from .compiler_bridge import (
44	)	44	)
45	from .types import dtype_size_bytes	45	from .types import dtype_size_bytes
46		46
		47	+if TYPE_CHECKING:
		48	+ from .tla.ffi import ExternFunction
		49	+
47	# CATLASS_DSL_KEEP tokens: ir / ir-debug / kernel.	50	# CATLASS_DSL_KEEP tokens: ir / ir-debug / kernel.
48	_KEEP_ALL_TOKENS: frozenset[str] = frozenset({"ir", "ir-debug", "kernel"})	51	_KEEP_ALL_TOKENS: frozenset[str] = frozenset({"ir", "ir-debug", "kernel"})
49	_POINTER_ABI_SIZE = 8	52	_POINTER_ABI_SIZE = 8
	@@ -324,11 +327,18 @@ def compile_kernel(
324	type_args=type_args,	327	type_args=type_args,
325	location=decorator_location,	328	location=decorator_location,
326	)	329	)
		330	+ extern_function = lowered.extern_function
		331	+ extern_core_types = lowered.extern_core_types
327	tlair_mlir = lowered.asm(generic=True)	332	tlair_mlir = lowered.asm(generic=True)
328	entrypoint = _extract_entrypoint(tlair_mlir)	333	entrypoint = _extract_entrypoint(tlair_mlir)
329	compiler_bridge_path = resolve_bridge_extension_path()	334	compiler_bridge_path = resolve_bridge_extension_path()
330	hivmc = _resolve_hivmc_a5()	335	hivmc = _resolve_hivmc_a5()
331	target = _resolve_kernel_target(runtime)	336	target = _resolve_kernel_target(runtime)
		337	+ extern_targets = _resolve_extern_targets(
		338	+ extern_function,
		339	+ extern_core_types=extern_core_types,
		340	+ target_arch=target.target_arch,
		341	+ )
332	cache_dir = runtime.cache_dir or _default_cache_dir()	342	cache_dir = runtime.cache_dir or _default_cache_dir()
333	cache_key = _cache_key(	343	cache_key = _cache_key(
334	tlair_mlir=tlair_mlir,	344	tlair_mlir=tlair_mlir,
	@@ -337,6 +347,8 @@ def compile_kernel(
337	compiler_bridge_path=compiler_bridge_path,	347	compiler_bridge_path=compiler_bridge_path,
338	hivmc=hivmc,	348	hivmc=hivmc,
339	target=target,	349	target=target,
		350	+ extern_function=extern_function,
		351	+ extern_targets=extern_targets,
340	)	352	)
341	artifact_dir = cache_dir / cache_key	353	artifact_dir = cache_dir / cache_key
342	manifest = artifact_dir / "manifest.json"	354	manifest = artifact_dir / "manifest.json"
	@@ -445,6 +457,17 @@ def compile_kernel(
445	hivmc_mlir_path, template_bitcode = _create_stamped_hivmc_input(	457	hivmc_mlir_path, template_bitcode = _create_stamped_hivmc_input(
446	mlir_path, runtime_for_hivmc	458	mlir_path, runtime_for_hivmc
447	)	459	)
		460	+ if extern_function is not None:
		461	+ user_bitcodes = _compile_ascendc_extern_function(
		462	+ extern_function,
		463	+ artifact_dir=artifact_dir,
		464	+ targets=extern_targets,
		465	+ )
		466	+ if template_bitcode is None:
		467	+ template_bitcode = _resolve_hivm_template_bitcode(runtime_for_hivmc)
		468	+ template_bitcode = ",".join(
		469	+ (template_bitcode, *(str(path) for path in user_bitcodes))
		470	+ )
448	try:	471	try:
449	_run_checked(	472	_run_checked(
450	_build_hivmc_a5_command(	473	_build_hivmc_a5_command(
	@@ -1648,7 +1671,12 @@ def _cache_key(
1648	compiler_bridge_path: Path \| None,	1671	compiler_bridge_path: Path \| None,
1649	hivmc: Path,	1672	hivmc: Path,
1650	target: TlaKernelTarget,	1673	target: TlaKernelTarget,
		1674	+ extern_function: ExternFunction \| None = None,
		1675	+ extern_targets: Sequence[TlaKernelTarget] = (),
1651	) -> str:	1676	) -> str:
		1677	+ extern_compile = (
		1678	+ None if extern_function is None else _ascendc_extern_compile_identity()
		1679	+ )
1652	key_payload = {	1680	key_payload = {
1653	"debug_print_workspace_abi_revision": _DEBUG_PRINT_WORKSPACE_ABI_REVISION,	1681	"debug_print_workspace_abi_revision": _DEBUG_PRINT_WORKSPACE_ABI_REVISION,
1654	"print_tensor_workspace_abi_revision": _PRINT_TENSOR_WORKSPACE_ABI_REVISION,	1682	"print_tensor_workspace_abi_revision": _PRINT_TENSOR_WORKSPACE_ABI_REVISION,
	@@ -1664,6 +1692,13 @@ def _cache_key(
1664	"hivmc_fingerprint": _tool_fingerprint(hivmc),	1692	"hivmc_fingerprint": _tool_fingerprint(hivmc),
1665	"mlir": tlair_mlir,	1693	"mlir": tlair_mlir,
1666	"print_ir": runtime.print_ir,	1694	"print_ir": runtime.print_ir,
		1695	+ "extern_source_sha256": (
		1696	+ None
		1697	+ if extern_function is None
		1698	+ else hashlib.sha256(extern_function.source.encode("utf-8")).hexdigest()
		1699	+ ),
		1700	+ "extern_targets": [target.arch_scope for target in extern_targets],
		1701	+ "extern_compile": extern_compile,
1667	}	1702	}
1668	return hashlib.sha256(	1703	return hashlib.sha256(
1669	json.dumps(key_payload, sort_keys=True).encode("utf-8")	1704	json.dumps(key_payload, sort_keys=True).encode("utf-8")
	@@ -1818,6 +1853,162 @@ def _resolve_hivmc_a5() -> Path:
1818	)	1853	)
1819		1854
1820		1855
		1856	+def _resolve_ccec() -> Path:
		1857	+ """Resolve ``ccec`` from PATH / ``ASCEND_HOME_PATH`` after ``set_env.sh``."""
		1858	+ which = shutil.which("ccec")
		1859	+ if which:
		1860	+ return Path(which).resolve()
		1861	+ ascend_home = os.getenv("ASCEND_HOME_PATH")
		1862	+ if ascend_home:
		1863	+ candidate = Path(ascend_home).expanduser().resolve() / "bin" / "ccec"
		1864	+ if candidate.exists():
		1865	+ return candidate.resolve()
		1866	+ raise TlaBackendCompilerNotFoundError(
		1867	+ "ccec not found on PATH. Source the CANN toolkit set_env.sh so user "
		1868	+ "Ascend C external functions can be compiled."
		1869	+ )
		1870	+
		1871	+
		1872	+def _ascendc_include_dirs(ascend_home: Path) -> list[Path]:
		1873	+ asc_root = (ascend_home / "asc").resolve()
		1874	+ highlevel_api = asc_root.parent / "ascendc" / "include" / "highlevel_api"
		1875	+ return [
		1876	+ path.resolve()
		1877	+ for path in (
		1878	+ asc_root,
		1879	+ asc_root / "impl" / "adv_api",
		1880	+ asc_root / "impl" / "basic_api",
		1881	+ asc_root / "impl" / "basic_api" / "reg_compute",
		1882	+ asc_root / "impl" / "c_api",
		1883	+ asc_root / "impl" / "micro_api",
		1884	+ asc_root / "impl" / "simt_api",
		1885	+ asc_root / "impl" / "utils",
		1886	+ asc_root / "include",
		1887	+ asc_root / "include" / "adv_api",
		1888	+ asc_root / "include" / "aicpu_api",
		1889	+ asc_root / "include" / "basic_api",
		1890	+ asc_root / "include" / "basic_api" / "reg_compute",
		1891	+ asc_root / "include" / "c_api",
		1892	+ asc_root / "include" / "interface",
		1893	+ asc_root / "include" / "micro_api",
		1894	+ asc_root / "include" / "simt_api",
		1895	+ asc_root / "include" / "tiling",
		1896	+ asc_root / "include" / "utils",
		1897	+ highlevel_api,
		1898	+ Path(__file__).resolve().parents[3] / "include",
		1899	+ )
		1900	+ if path.is_dir()
		1901	+ ]
		1902	+
		1903	+
		1904	+def _ascendc_compiler_inputs() -> tuple[Path, list[Path]]:
		1905	+ compiler = _resolve_ccec()
		1906	+ ascend_home_env = os.getenv("ASCEND_HOME_PATH")
		1907	+ if not ascend_home_env:
		1908	+ raise TlaBackendCompilerNotFoundError(
		1909	+ "ASCEND_HOME_PATH is not set. Source the CANN toolkit set_env.sh so "
		1910	+ "Ascend C headers can be found."
		1911	+ )
		1912	+ ascend_home = Path(ascend_home_env).expanduser().resolve()
		1913	+ return compiler, _ascendc_include_dirs(ascend_home)
		1914	+
		1915	+
		1916	+def _resolve_extern_targets(
		1917	+ extern_function: ExternFunction \| None,
		1918	+ *,
		1919	+ extern_core_types: Iterable[str],
		1920	+ target_arch: str,
		1921	+) -> tuple[TlaKernelTarget, ...]:
		1922	+ if extern_function is None:
		1923	+ return ()
		1924	+
		1925	+ core_types = frozenset(extern_core_types)
		1926	+ if not core_types:
		1927	+ raise TlaKernelCompileError(
		1928	+ f"external function {extern_function.symbol!r} has no call target"
		1929	+ )
		1930	+ unsupported = core_types.difference(("aic", "aiv"))
		1931	+ if unsupported:
		1932	+ raise TlaKernelCompileError(
		1933	+ f"unsupported external function core types: {sorted(unsupported)}"
		1934	+ )
		1935	+
		1936	+ targets = tuple(
		1937	+ _get_kernel_target(target_arch=target_arch, core_type=core_type)
		1938	+ for core_type in ("aic", "aiv")
		1939	+ if core_type in core_types
		1940	+ )
		1941	+ return targets
		1942	+
		1943	+
		1944	+def _ascendc_extern_compile_identity() -> dict[str, object]:
		1945	+ """Return the resolved compiler configuration tracked by the kernel cache."""
		1946	+
		1947	+ compiler, include_dirs = _ascendc_compiler_inputs()
		1948	+ return {
		1949	+ "ccec": str(compiler),
		1950	+ "ccec_version": _tool_version(compiler),
		1951	+ "ccec_fingerprint": _tool_fingerprint(compiler),
		1952	+ "include_dirs": [str(path) for path in include_dirs],
		1953	+ }
		1954	+
		1955	+
		1956	+def _compile_ascendc_extern_function(
		1957	+ extern_function: ExternFunction,
		1958	+ *,
		1959	+ artifact_dir: Path,
		1960	+ targets: Sequence[TlaKernelTarget],
		1961	+) -> tuple[Path, ...]:
		1962	+ # ``targets`` is produced by ``_resolve_extern_targets`` in the compile path.
		1963	+ targets = tuple(targets)
		1964	+ assert 1 <= len(targets) <= 2
		1965	+ source = artifact_dir / "extern.cpp"
		1966	+ source.write_text(extern_function.source, encoding="utf-8")
		1967	+ compiler, include_dirs = _ascendc_compiler_inputs()
		1968	+ command = [
		1969	+ str(compiler),
		1970	+ "-O2",
		1971	+ "-x",
		1972	+ "cce",
		1973	+ "--cce-auto-sync=off",
		1974	+ "--cce-aicore-only",
		1975	+ "--cce-generic-addrspace=off",
		1976	+ str(source),
		1977	+ "-emit-llvm",
		1978	+ "-c",
		1979	+ "-mllvm",
		1980	+ "-disable-llvm-optzns",
		1981	+ "-DCATLASS_ARCH=3510",
		1982	+ "-DTILING_KEY_VAR",
		1983	+ "-Wno-ignored-attributes",
		1984	+ "-std=c++17",
		1985	+ ]
		1986	+ for include_dir in include_dirs:
		1987	+ command.extend(["-I", str(include_dir)])
		1988	+
		1989	+ outputs = tuple(
		1990	+ artifact_dir / f"extern.{target.arch_scope}.bc" for target in targets
		1991	+ )
		1992	+ for target, output in zip(targets, outputs, strict=True):
		1993	+ target_command = [
		1994	+ *command,
		1995	+ f"--cce-aicore-arch={target.cce_arch}",
		1996	+ "-o",
		1997	+ str(output),
		1998	+ ]
		1999	+ _run_checked(
		2000	+ target_command,
		2001	+ label=f"ccec external function {extern_function.symbol}",
		2002	+ cwd=artifact_dir,
		2003	+ )
		2004	+ if not output.exists():
		2005	+ raise TlaKernelCompileError(
		2006	+ "ccec completed but external function bitcode was not created at "
		2007	+ f"{output}"
		2008	+ )
		2009	+ return tuple(output.resolve() for output in outputs)
		2010	+
		2011	+
1821	def _build_hivmc_a5_command(	2012	def _build_hivmc_a5_command(
1822	*,	2013	*,
1823	compiler: Path,	2014	compiler: Path,
	@@ -1860,10 +2051,10 @@ def _create_stamped_hivmc_input(
1860	) -> tuple[Path, str \| None]:	2051	) -> tuple[Path, str \| None]:
1861	"""Stamp a private HIVMC input only when debug-print helpers are present.	2052	"""Stamp a private HIVMC input only when debug-print helpers are present.
1862		2053
1863	- Ordinary kernels rely on ``--link-aicore-bitcode`` alone. Debug /	2054	+ Ordinary and external-function kernels rely on ``--link-aicore-bitcode``.
1864	- ``print_tensor`` helpers also need module attrs ``hivm.aiv_bitcode`` /	2055	+ Debug / ``print_tensor`` helpers also need module attrs
1865	- ``hivm.aic_bitcode`` (and optionally helper bitcode), so copy+stamp a	2056	+ ``hivm.aiv_bitcode`` / ``hivm.aic_bitcode`` (and optionally helper bitcode),
1866	- private ``*.hivmc-input.mlir`` in those cases only.	2057	+ so copy and stamp a private ``*.hivmc-input.mlir`` only for those kernels.
1867	"""	2058	"""
1868	compiler_text = mlir_path.read_text()	2059	compiler_text = mlir_path.read_text()
1869	if (	2060	if (

Mpython/tla_dsl/catlass/execution_lowering.py+19-5

	@@ -6,7 +6,7 @@ import dataclasses
6	import inspect	6	import inspect
7	import linecache	7	import linecache
8	from dataclasses import dataclass	8	from dataclasses import dataclass
9	-from typing import Any, Mapping, Sequence	9	+from typing import TYPE_CHECKING, Any, Mapping, Sequence
10		10
11	from mlir import ir as mlir_ir # type: ignore[assignment]	11	from mlir import ir as mlir_ir # type: ignore[assignment]
12		12
	@@ -23,6 +23,9 @@ from .base_dsl.typing import Numeric, is_constexpr_annotation
23	from .dsl import _jit_helper_transformer	23	from .dsl import _jit_helper_transformer
24	from .tla.typing import Tensor	24	from .tla.typing import Tensor
25		25
		26	+if TYPE_CHECKING:
		27	+ from .tla.ffi import ExternFunction
		28	+
26		29
27	class TlaLoweringError(RuntimeError):	30	class TlaLoweringError(RuntimeError):
28	"""Raised when Tla DSL lowering fails."""	31	"""Raised when Tla DSL lowering fails."""
	@@ -72,6 +75,8 @@ class LoweredTlaIR:
72	module: mlir_ir.Module	75	module: mlir_ir.Module
73	generic: bool = False	76	generic: bool = False
74	_asm: str \| None = None	77	_asm: str \| None = None
		78	+ extern_function: ExternFunction \| None = None
		79	+ extern_core_types: frozenset[str] = frozenset()
75		80
76	def asm(self, *, generic: bool \| None = None) -> str:	81	def asm(self, *, generic: bool \| None = None) -> str:
77	emit_generic = self.generic if generic is None else bool(generic)	82	emit_generic = self.generic if generic is None else bool(generic)
	@@ -164,7 +169,7 @@ def lower_jit_to_tlair_module_by_execution(
164	module = mlir_ir.Module.create()	169	module = mlir_ir.Module.create()
165	with mlir_ir.InsertionPoint(module.body):	170	with mlir_ir.InsertionPoint(module.body):
166	fn_loc = _coerce_location(ctx, location)	171	fn_loc = _coerce_location(ctx, location)
167	- _build_tla_func(	172	+ extern_function, extern_core_types = _build_tla_func(
168	fn=fn,	173	fn=fn,
169	module=module,	174	module=module,
170	fn_name=fn.__name__,	175	fn_name=fn.__name__,
	@@ -177,7 +182,13 @@ def lower_jit_to_tlair_module_by_execution(
177	fn_loc=fn_loc,	182	fn_loc=fn_loc,
178	auto_sync=auto_sync,	183	auto_sync=auto_sync,
179	)	184	)
180	- lowered = LoweredTlaIR(context=ctx, module=module, generic=bool(generic))	185	+ lowered = LoweredTlaIR(
		186	+ context=ctx,
		187	+ module=module,
		188	+ generic=bool(generic),
		189	+ extern_function=extern_function,
		190	+ extern_core_types=extern_core_types,
		191	+ )
181	lowered._asm = module.operation.get_asm(	192	lowered._asm = module.operation.get_asm(
182	print_generic_op_form=bool(generic),	193	print_generic_op_form=bool(generic),
183	assume_verified=False,	194	assume_verified=False,
	@@ -236,7 +247,7 @@ def _build_tla_func(
236	ctx: mlir_ir.Context,	247	ctx: mlir_ir.Context,
237	fn_loc: mlir_ir.Location,	248	fn_loc: mlir_ir.Location,
238	auto_sync: str \| None,	249	auto_sync: str \| None,
239	-) -> None:	250	+) -> tuple[ExternFunction \| None, frozenset[str]]:
240	runtime_arg_names = [name for name in arg_names if name not in constexpr_names]	251	runtime_arg_names = [name for name in arg_names if name not in constexpr_names]
241		252
242	# Dynamic GM host tensors enter as unified GM memref + originShape0/1 index args.	253	# Dynamic GM host tensors enter as unified GM memref + originShape0/1 index args.
	@@ -492,7 +503,7 @@ def _build_tla_func(
492	category_bindings=category_bindings,	503	category_bindings=category_bindings,
493	tensor_host_by_value=tensor_host_by_value,	504	tensor_host_by_value=tensor_host_by_value,
494	module=module,	505	module=module,
495	- ):	506	+ ) as frontend_state:
496	from .core_api import (	507	from .core_api import (
497	_register_tla_tensor_metadata,	508	_register_tla_tensor_metadata,
498	_register_tla_tensor_type,	509	_register_tla_tensor_type,
	@@ -599,7 +610,10 @@ def _build_tla_func(
599	if message is None:	610	if message is None:
600	message = f"Execution-mode lowering failed while running `{fn.__name__}`: {exc}"	611	message = f"Execution-mode lowering failed while running `{fn.__name__}`: {exc}"
601	raise UnsupportedExecutionLowering(message) from exc	612	raise UnsupportedExecutionLowering(message) from exc
		613	+ extern_function = frontend_state.extern_function
		614	+ extern_core_types = frozenset(frontend_state.extern_core_types)
602	mlir_ir.Operation.create("tla.return", loc=fn_loc)	615	mlir_ir.Operation.create("tla.return", loc=fn_loc)
		616	+ return extern_function, extern_core_types
603		617
604		618
605	def _coerce_location(	619	def _coerce_location(

Mpython/tla_dsl/catlass/tla/__init__.py+8-0

	@@ -57,8 +57,10 @@ _PARENT_EXPORTS = (
57	"const_expr",	57	"const_expr",
58	"Constexpr",	58	"Constexpr",
59	"Pointer",	59	"Pointer",
		60	+ "TypedPointer",
60	"JitArgument",	61	"JitArgument",
61	"AddressSpace",	62	"AddressSpace",
		63	+ "extern",
62	"Numeric",	64	"Numeric",
63	"Integer",	65	"Integer",
64	"Float",	66	"Float",
	@@ -121,6 +123,7 @@ def __getattr__(name: str) -> Any:
121		123
122	if name == "Constexpr" or name in (	124	if name == "Constexpr" or name in (
123	"Pointer",	125	"Pointer",
		126	+ "TypedPointer",
124	"JitArgument",	127	"JitArgument",
125	"as_numeric",	128	"as_numeric",
126	"cast",	129	"cast",
	@@ -138,6 +141,11 @@ def __getattr__(name: str) -> Any:
138		141
139	return AddressSpace	142	return AddressSpace
140		143
		144	+ if name == "extern":
		145	+ from .ffi import extern
		146	+
		147	+ return extern
		148	+
141	if name in (	149	if name in (
142	"Numeric",	150	"Numeric",
143	"Integer",	151	"Integer",

Apython/tla_dsl/catlass/tla/ffi.py+133-0

@@ -0,0 +1,133 @@
	1	+"""Declaration API for user-provided TLA device functions."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import inspect
	6	+import re
	7	+from dataclasses import dataclass
	8	+from typing import Callable, TypeAlias, get_type_hints
	9	+
	10	+from mlir import ir as mlir_ir # type: ignore[assignment]
	11	+
	12	+from ..base_dsl.op import dsl_user_op
	13	+from ..base_dsl.typing import Numeric, TypedPointer
	14	+
	15	+
	16	+_SYMBOL_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
	17	+
	18	+
	19	+ExternArgType: TypeAlias = TypedPointer \| type[Numeric]
	20	+
	21	+
	22	+@dataclass(frozen=True)
	23	+class ExternFunction:
	24	+ """A C ABI entry point supplied by user source code."""
	25	+
	26	+ source: str
	27	+ symbol: str
	28	+ arg_types: tuple[ExternArgType, ...]
	29	+
	30	+ @dsl_user_op
	31	+ def __call__(
	32	+ self,
	33	+ *args: object,
	34	+ loc: mlir_ir.Location \| None = None,
	35	+ ) -> None:
	36	+ """Emit a call to this external function during frontend lowering."""
	37	+
	38	+ # Importing here keeps declaration independent from the frontend's
	39	+ # comparatively heavy core_api module.
	40	+ from ..core_api import _emit_extern_call
	41	+
	42	+ _emit_extern_call(self, args, loc=loc)
	43	+
	44	+
	45	+def _validate_symbol(symbol: object) -> str:
	46	+ if not isinstance(symbol, str) or _SYMBOL_RE.fullmatch(symbol) is None:
	47	+ raise ValueError(f"tla.extern name must be a C identifier, got {symbol!r}")
	48	+ return symbol
	49	+
	50	+
	51	+def _validate_arg_type(annotation: object, *, parameter: str) -> ExternArgType:
	52	+ if isinstance(annotation, TypedPointer):
	53	+ return annotation
	54	+ if (
	55	+ isinstance(annotation, type)
	56	+ and issubclass(annotation, Numeric)
	57	+ and annotation.dtype
	58	+ ):
	59	+ return annotation
	60	+ raise TypeError(
	61	+ "tla.extern parameter annotations must be Pointer[dtype, memory_space] "
	62	+ f"or a concrete Numeric type; parameter {parameter!r} has {annotation!r}"
	63	+ )
	64	+
	65	+
	66	+def extern(
	67	+ *,
	68	+ source: str,
	69	+ name: str \| None = None,
	70	+) -> Callable[[Callable[..., None]], ExternFunction]:
	71	+ """Declare one C ABI entry point supplied by inline Ascend C source."""
	72	+
	73	+ # Validity check
	74	+ if not isinstance(source, str):
	75	+ raise TypeError(f"tla.extern source must be str, got {type(source).__name__}")
	76	+ if not source.strip():
	77	+ raise ValueError("tla.extern source must not be empty")
	78	+ explicit_symbol = None if name is None else _validate_symbol(name)
	79	+
	80	+ def decorator(fn: Callable[..., None]) -> ExternFunction:
	81	+ signature = inspect.signature(fn)
	82	+ hints = get_type_hints(fn)
	83	+
	84	+ # Validity check
	85	+ return_annotation = hints.get("return", signature.return_annotation)
	86	+ if return_annotation not in (None, type(None)):
	87	+ if return_annotation is inspect.Signature.empty:
	88	+ detail = "is missing"
	89	+ else:
	90	+ detail = f"must be None, got {return_annotation!r}"
	91	+ raise TypeError(f"tla.extern return annotation for {fn.__name__} {detail}")
	92	+
	93	+ arg_types: list[ExternArgType] = []
	94	+ for parameter in signature.parameters.values():
	95	+ # Validity check
	96	+ if parameter.kind not in (
	97	+ inspect.Parameter.POSITIONAL_ONLY,
	98	+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
	99	+ ):
	100	+ raise TypeError(
	101	+ "tla.extern parameters must be positional and fixed; "
	102	+ f"parameter {parameter.name!r} has kind {parameter.kind.description}"
	103	+ )
	104	+ if parameter.default is not inspect.Parameter.empty:
	105	+ raise TypeError(
	106	+ "tla.extern parameters must not have default values; "
	107	+ f"parameter {parameter.name!r} has default {parameter.default!r}"
	108	+ )
	109	+ annotation = hints.get(parameter.name, parameter.annotation)
	110	+ if annotation is inspect.Parameter.empty:
	111	+ raise TypeError(
	112	+ f"tla.extern parameter {parameter.name!r} is missing an annotation"
	113	+ )
	114	+ # Add to arg_types after validation
	115	+ arg_types.append(_validate_arg_type(annotation, parameter=parameter.name))
	116	+
	117	+ symbol = (
	118	+ explicit_symbol
	119	+ if explicit_symbol is not None
	120	+ else _validate_symbol(fn.__name__)
	121	+ )
	122	+ return ExternFunction(
	123	+ source=source,
	124	+ symbol=symbol,
	125	+ arg_types=tuple(arg_types),
	126	+ )
	127	+
	128	+ return decorator
	129	+
	130	+
	131	+__all__ = [
	132	+ "extern",
	133	+]

Mpython/tla_dsl/csrc/mlir/include/Dialect/Tla/IR/Tla.td+15-0

	@@ -935,6 +935,21 @@ def Tla_DebugPrintOp : Tla_Op<"debug_print", [MemoryEffects<[MemWrite]>]> {
935	let hasVerifier = 1;	935	let hasVerifier = 1;
936	}	936	}
937		937
		938	+def Tla_CallExternOp : Tla_Op<"call_extern", [MemoryEffects<[MemRead, MemWrite]>]> {
		939	+ let summary = "Call a user-provided device function";
		940	+ let description = [{
		941	+ Calls one C ABI device function supplied by user LLVM bitcode. Pointer
		942	+ operands remain !tla.ptr until tla-lower-ptr converts them to i64 byte
		943	+ addresses. v1 has no results and conservatively models unknown memory
		944	+ reads and writes.
		945	+ }];
		946	+ let arguments = (ins FlatSymbolRefAttr:$callee, Variadic<AnyType>:$operands);
		947	+ let assemblyFormat = [{
		948	+ $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)
		949	+ }];
		950	+ let hasVerifier = 1;
		951	+}
		952	+
938	def Tla_PrintTensorOp : Tla_Op<"print_tensor", [MemoryEffects<[MemWrite]>]> {	953	def Tla_PrintTensorOp : Tla_Op<"print_tensor", [MemoryEffects<[MemWrite]>]> {
939	let summary = "Print a prefix of one contiguous supported GM or UB tensor";	954	let summary = "Print a prefix of one contiguous supported GM or UB tensor";
940	let description = [{	955	let description = [{

Mpython/tla_dsl/csrc/mlir/include/Passes.h+1-0

	@@ -7,6 +7,7 @@ namespace tla {
7		7
8	std::unique_ptr<mlir::Pass> createTlaLowerDebugPrintPass();	8	std::unique_ptr<mlir::Pass> createTlaLowerDebugPrintPass();
9	std::unique_ptr<mlir::Pass> createTlaLowerPtrPass();	9	std::unique_ptr<mlir::Pass> createTlaLowerPtrPass();
		10	+std::unique_ptr<mlir::Pass> createTlaLowerExternCallPass();
10	std::unique_ptr<mlir::Pass> createTlaCubeRegionPass();	11	std::unique_ptr<mlir::Pass> createTlaCubeRegionPass();
11	std::unique_ptr<mlir::Pass> createTlaFinalizeMemrefPass();	12	std::unique_ptr<mlir::Pass> createTlaFinalizeMemrefPass();
12	std::unique_ptr<mlir::Pass> createTlaLowerFlagBarrierToHivmPass();	13	std::unique_ptr<mlir::Pass> createTlaLowerFlagBarrierToHivmPass();

Mpython/tla_dsl/csrc/mlir/lib/Dialect/Tla/IR/TlaOps.cpp+16-0

	@@ -302,6 +302,22 @@ mlir::LogicalResult MmadOp::verify()
302	return mlir::success();	302	return mlir::success();
303	}	303	}
304		304
		305	+mlir::LogicalResult CallExternOp::verify()
		306	+{
		307	+ bool inVector = hasEnclosing<VectorOp>(getOperation());
		308	+ bool inCube = hasEnclosing<CubeOp>(getOperation());
		309	+ if (inVector == inCube)
		310	+ return emitOpError("must be nested inside exactly one tla.vector or tla.cube region");
		311	+ if (inVector && hasEnclosing<VecFuncOp>(getOperation()))
		312	+ return emitOpError("must be outside tla.vec.func");
		313	+ for (mlir::Type type : getOperandTypes()) {
		314	+ if (mlir::isa<PtrType, mlir::IntegerType, mlir::FloatType>(type))
		315	+ continue;
		316	+ return emitOpError("operands must be !tla.ptr or scalar integer/float (not index), got ") << type;
		317	+ }
		318	+ return mlir::success();
		319	+}
		320	+
305	// One thread block may hold at most this many threads on the supported targets.	321	// One thread block may hold at most this many threads on the supported targets.
306	// The lowering packs the product into hivm_regbaseintrins::SIMT_EntryAttr, whose	322	// The lowering packs the product into hivm_regbaseintrins::SIMT_EntryAttr, whose
307	// value is a uint32_t, so an unchecked product would also truncate.	323	// value is a uint32_t, so an unchecked product would also truncate.

Mpython/tla_dsl/csrc/mlir/lib/Passes/CMakeLists.txt+1-0

	@@ -11,6 +11,7 @@ add_library(TlaPasses
11	TlaLowerBlockIdxPass.cpp	11	TlaLowerBlockIdxPass.cpp
12	TlaScratchAllocation.cpp	12	TlaScratchAllocation.cpp
13	TlaLowerPtrPass.cpp	13	TlaLowerPtrPass.cpp
		14	+ TlaLowerExternCallPass.cpp
14	TlaLowerMutexToStdPass.cpp	15	TlaLowerMutexToStdPass.cpp
15	TlaFinalizeMemrefPass.cpp	16	TlaFinalizeMemrefPass.cpp
16	TlaCubeRegionPass.cpp	17	TlaCubeRegionPass.cpp

Mpython/tla_dsl/csrc/mlir/lib/Passes/PassRegistry.cpp+2-0

Mpython/tla_dsl/csrc/mlir/lib/Passes/PassesInternal.h+1-0

	@@ -23,6 +23,7 @@ void registerTlaLowerBlockIdxPass();
23	void registerTlaVectorRegionPass();	23	void registerTlaVectorRegionPass();
24	void registerTlaLowerFlagBarrierToHivmPass();	24	void registerTlaLowerFlagBarrierToHivmPass();
25	void registerTlaLowerPtrPass();	25	void registerTlaLowerPtrPass();
		26	+void registerTlaLowerExternCallPass();
26	void registerTlaLowerMutexToStdPass();	27	void registerTlaLowerMutexToStdPass();
27	void registerTlaCubeRegionPass();	28	void registerTlaCubeRegionPass();
28	void registerTlaFinalizeMemrefPass();	29	void registerTlaFinalizeMemrefPass();

Mpython/tla_dsl/csrc/mlir/lib/Passes/TlaInsertAutoMutexPass.cpp+19-1

	@@ -351,6 +351,23 @@ static LogicalResult validateNoManualLocalSync(func::FuncOp func)
351	"or local flag synchronization; cross_core_* remains explicit");	351	"or local flag synchronization; cross_core_* remains explicit");
352	}	352	}
353		353
		354	+// In TLA DSL extern op support v1, the user is responsible for ensuring that
		355	+// external calls are properly synchronized. we do not attempt to automatically
		356	+// synchronize them.
		357	+static LogicalResult validateNoExternCalls(func::FuncOp func)
		358	+{
		359	+ ::tla::CallExternOp invalid;
		360	+ func.walk([&](::tla::CallExternOp op) {
		361	+ if (!invalid)
		362	+ invalid = op;
		363	+ });
		364	+ if (!invalid)
		365	+ return success();
		366	+ return invalid.emitError(
		367	+ "auto_sync='v0' cannot be combined with tla.call_extern; "
		368	+ "external calls require explicit synchronization in v1");
		369	+}
		370	+
354	static LogicalResult validateCopyUnitFlags(func::FuncOp func)	371	static LogicalResult validateCopyUnitFlags(func::FuncOp func)
355	{	372	{
356	LogicalResult result = success();	373	LogicalResult result = success();
	@@ -929,7 +946,8 @@ public:
929	func->removeAttr(kAutoSyncAttrName);	946	func->removeAttr(kAutoSyncAttrName);
930	continue;	947	continue;
931	}	948	}
932	- if (failed(validateNoManualLocalSync(func)) \|\| failed(validateCopyUnitFlags(func))) {	949	+ if (failed(validateNoExternCalls(func)) \|\| failed(validateNoManualLocalSync(func)) \|\|
		950	+ failed(validateCopyUnitFlags(func))) {
933	signalPassFailure();	951	signalPassFailure();
934	return;	952	return;
935	}	953	}

Apython/tla_dsl/csrc/mlir/lib/Passes/TlaLowerExternCallPass.cpp+99-0

@@ -0,0 +1,99 @@
	1	+#include "PassesCommon.h"
	2	+#include "PassesInternal.h"
	3	+
	4	+#include "mlir/Dialect/Func/IR/FuncOps.h"
	5	+#include "mlir/IR/SymbolTable.h"
	6	+
	7	+namespace tla {
	8	+namespace {
	9	+
	10	+static bool isVectorCall(::tla::CallExternOp op)
	11	+{
	12	+ return op->getParentOfType<::tla::VectorOp>() != nullptr;
	13	+}
	14	+
	15	+static void updateCallCoreType(func::FuncOp callee, bool isVector)
	16	+{
	17	+ MLIRContext* ctx = callee.getContext();
	18	+ hivm::TFuncCoreType callerCoreType = isVector ? hivm::TFuncCoreType::AIV : hivm::TFuncCoreType::AIC;
	19	+ auto calleeCoreType = callee->getAttrOfType<hivm::TFuncCoreTypeAttr>(hivm::TFuncCoreTypeAttr::name);
	20	+ if (!calleeCoreType) {
	21	+ callee->setAttr(hivm::TFuncCoreTypeAttr::name, hivm::TFuncCoreTypeAttr::get(ctx, callerCoreType));
	22	+ return;
	23	+ }
	24	+ if (calleeCoreType.getFuncCoreType() != callerCoreType) {
	25	+ callee->setAttr(
	26	+ hivm::TFuncCoreTypeAttr::name, hivm::TFuncCoreTypeAttr::get(ctx, hivm::TFuncCoreType::AIC_OR_AIV));
	27	+ }
	28	+}
	29	+
	30	+// Lower tla.call_extern to a private func.func declaration + func.call.
	31	+class TlaLowerExternCallPass : public PassWrapper<TlaLowerExternCallPass, OperationPass<ModuleOp>> {
	32	+public:
	33	+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TlaLowerExternCallPass)
	34	+
	35	+ StringRef getArgument() const override
	36	+ {
	37	+ return "tla-lower-extern-call";
	38	+ }
	39	+ StringRef getName() const override
	40	+ {
	41	+ return "TlaLowerExternCallPass";
	42	+ }
	43	+ StringRef getDescription() const override
	44	+ {
	45	+ return "Lower tla.call_extern to a private func.call declaration";
	46	+ }
	47	+ void getDependentDialects(DialectRegistry& registry) const override
	48	+ {
	49	+ registry.insert<func::FuncDialect, ::tla::TlaDialect>();
	50	+ }
	51	+
	52	+ void runOnOperation() override
	53	+ {
	54	+ ModuleOp module = getOperation();
	55	+ SmallVector<::tla::CallExternOp, 4> calls;
	56	+ module.walk([&](::tla::CallExternOp op) { calls.push_back(op); });
	57	+
	58	+ for (::tla::CallExternOp op : calls) {
	59	+ StringRef symbol = op.getCallee();
	60	+ bool isVector = isVectorCall(op);
	61	+ SmallVector<Type, 4> operandTypes(op.getOperandTypes());
	62	+ auto functionType = FunctionType::get(module.getContext(), operandTypes, TypeRange{});
	63	+ func::FuncOp callee = module.lookupSymbol<func::FuncOp>(symbol);
	64	+ if (!callee) {
	65	+ OpBuilder moduleBuilder(module.getBodyRegion());
	66	+ moduleBuilder.setInsertionPointToStart(module.getBody());
	67	+ callee = moduleBuilder.create<func::FuncOp>(op.getLoc(), symbol, functionType);
	68	+ callee.setPrivate();
	69	+ } else if (!callee.isDeclaration()) {
	70	+ op.emitOpError() << "external symbol @" << symbol << " conflicts with a defined function";
	71	+ signalPassFailure();
	72	+ return;
	73	+ } else if (callee.getFunctionType() != functionType) {
	74	+ op.emitOpError() << "external symbol @" << symbol << " was called with incompatible function types";
	75	+ signalPassFailure();
	76	+ return;
	77	+ }
	78	+ updateCallCoreType(callee, isVector);
	79	+
	80	+ OpBuilder builder(op);
	81	+ builder.create<func::CallOp>(op.getLoc(), callee, op.getOperands());
	82	+ op.erase();
	83	+ }
	84	+ }
	85	+};
	86	+
	87	+} // namespace
	88	+
	89	+std::unique_ptr<Pass> createTlaLowerExternCallPass()
	90	+{
	91	+ return std::make_unique<TlaLowerExternCallPass>();
	92	+}
	93	+
	94	+void registerTlaLowerExternCallPass()
	95	+{
	96	+ PassRegistration<TlaLowerExternCallPass>();
	97	+}
	98	+
	99	+} // namespace tla

Apython/tla_dsl/examples/end_to_end/extern_op/README.md+165-0

@@ -0,0 +1,165 @@
	1	+# 外部算子端到端示例
	2	+
	3	+本目录下样例演示如何使用 `@tla.extern` 将用户提供的 Ascend C 函数嵌入 CATLASS DSL Kernel：`extern_vecadd.py` 展示外部函数与 TLA DSL 原语协同完成向量加法，`extern_dual_core.py` 展示同一份源码中的同一个外部函数同时被 AIC 和 AIV 调用。
	4	+
	5	+## VecAdd 功能说明
	6	+
	7	+向量加法算子实现两个一维向量的逐元素加法，计算公式为：
	8	+
	9	+$$
	10	+\begin{aligned}
	11	+C &= A + B
	12	+\end{aligned}
	13	+$$
	14	+
	15	+本样例仅使用外部 Ascend C 函数 `tla_user_gm_to_ub_f32` 替换输入 A、B 的 GM（Global Memory）到 UB（Unified Buffer）数据搬运，向量加法和结果写回仍由 TLA DSL 完成。整体流程如下：
	16	+
	17	+1. 通过 `@tla.extern` 声明外部函数的 C ABI，并以内联字符串形式提供 Ascend C 源码；
	18	+2. 在 `tla.vector()` 区域内调用外部函数，将 A、B 从 GM 搬运到 UB；
	19	+3. 使用 TLA DSL 的 `load`、`tla.add` 和 `store` 完成分块向量加法；
	20	+4. 使用 `tla.copy` 将结果从 UB 写回 GM。
	21	+
	22	+## AIC/AIV 共享外部函数
	23	+
	24	+`extern_dual_core.py` 只提供一份 `OP_SOURCE_CODES` 和一个 `tla_user_store_i32` 外部函数声明。TLA Kernel 在 `tla.cube()` 与 `tla.vector()` 区域中复用这个声明。输出包含三个 64-byte cache line：AIC 写入索引 0，两个 AIV sub-block 分别写入索引 16、32，使三个执行单元写入不同的 cache line。对应位置的期望值为 `[101, 202, 202]`，其余元素为 0。
	25	+
	26	+## 代码组织
	27	+
	28	+本目录组织结构如下所示：
	29	+
	30	+```plain
	31	+./extern_op
	32	+├── extern_dual_core.py
	33	+├── extern_vecadd.py
	34	+└── README.md
	35	+```
	36	+
	37	+\| 文件 \| 概述 \|
	38	+\|------\|------\|
	39	+\| [`extern_dual_core.py`](extern_dual_core.py) \| 单份 Ascend C 源码、单个 extern op 同时在 AIC 与 AIV 区域调用，并进行上板结果校验。 \|
	40	+\| [`extern_vecadd.py`](extern_vecadd.py) \| 外部 Ascend C 函数与 TLA DSL 混合编程示例，包含外部函数源码、ABI 声明、VecAdd Kernel 和上板精度校验。 \|
	41	+
	42	+## 约束说明
	43	+
	44	+- 本样例的向量长度固定为 256，输入和输出数据类型固定为 `float32`，启动核数固定为 1，编译目标为 `--npu-arch 3510`。
	45	+- `@tla.extern` 的 `source` 必须是非空的 Ascend C 源码字符串；`name` 必须是合法的 C 标识符，省略时使用被装饰的 Python 函数名。
	46	+- 外部函数声明只支持位置固定、无默认值的参数。参数类型必须标注为 `tla.Pointer[dtype, address_space]` 或具体的 TLA 数值类型（如 `tla.Int32`），返回类型必须标注为 `None`。
	47	+- Kernel 调用外部函数时，需要传入显式指针（例如 `tensor.ptr`），且参数个数、数据类型和指针地址空间必须与声明完全一致。
	48	+- 外部函数必须在一个 `tla.vector()` 或 `tla.cube()` 区域内调用，不能在 `tla.vec.func(...)` 中调用。
	49	+- 单个 Kernel 当前最多依赖一个不同的外部函数，但可以像本样例一样多次调用该函数。
	50	+- 外部调用当前需要显式编写流水同步，不能与 `@tla.kernel(auto_sync="v0")` 组合使用。本样例通过 `flag` 的 `set` / `wait` 保证外部 GM→UB 搬运完成后再进行向量计算。
	51	+
	52	+## 使用示例
	53	+
	54	+要运行本路径下的样例，请参考[环境配置](../../../docs/zh/dev_guide/00_environment_setup.md)完成部署。
	55	+
	56	+### 命令行参数
	57	+
	58	+```text
	59	+extern_vecadd.py [-h] [--device DEVICE]
	60	+extern_dual_core.py [-h] [--device DEVICE]
	61	+```
	62	+
	63	+上述命令行参数具体说明如下：
	64	+
	65	+\| 参数 \| 默认值 \| 说明 \|
	66	+\|------\|--------\|------\|
	67	+\| `--device` \| `0` \| 上板执行使用的 NPU 设备号。 \|
	68	+
	69	+### 执行示例
	70	+
	71	+在 `python/tla_dsl` 目录下执行：
	72	+
	73	+```bash
	74	+cd python/tla_dsl
	75	+
	76	+# 使用 0 号 NPU 执行
	77	+python examples/end_to_end/extern_op/extern_vecadd.py --device 0
	78	+
	79	+# 同一个 extern op 分别在 AIC 和 AIV 中调用
	80	+python examples/end_to_end/extern_op/extern_dual_core.py --device 0
	81	+
	82	+# 忽略进程内缓存和磁盘缓存，强制重新编译后执行
	83	+CATLASS_DSL_FORCE_RECOMPILE=1 \
	84	+ python examples/end_to_end/extern_op/extern_vecadd.py --device 0
	85	+```
	86	+
	87	+执行测试后，预期输出：
	88	+
	89	+```plain
	90	+passed; kernel=<cache_dir>/<cache_key>/kernel.o
	91	+```
	92	+
	93	+两个程序都使用 `torch.testing.assert_close` 进行结果校验；校验通过后输出 `passed` 和编译产物路径，校验失败则抛出异常。`cache_dir` 是编译缓存目录，`cache_key` 是编译缓存的哈希值。
	94	+
	95	+---
	96	+
	97	+## 特性介绍
	98	+
	99	+### 声明外部 Ascend C 函数
	100	+
	101	+外部函数源码保存在 `OP_SOURCE_CODES` 字符串中，并使用 `extern "C"` 导出与 `@tla.extern` 声明一致的符号。示例中的 Ascend C 函数通过 `AscendC::DataCopy` 完成 GM 到 UB 的数据搬运：
	102	+
	103	+```cpp
	104	+extern "C" {
	105	+
	106	+[aicore] __attribute__((always_inline)) void tla_user_gm_to_ub_f32(
	107	+ uint64_t src_gm_addr, uint64_t dst_ub_addr, int32_t count) {
	108	+ AscendC::GlobalTensor<float> src;
	109	+ src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(src_gm_addr),
	110	+ static_cast<uint32_t>(count));
	111	+ AscendC::LocalTensor<float> dst(AscendC::TPosition::VECCALC,
	112	+ static_cast<uint32_t>(dst_ub_addr),
	113	+ static_cast<uint32_t>(count));
	114	+ AscendC::DataCopy(dst, src, static_cast<uint32_t>(count));
	115	+}
	116	+
	117	+} // extern "C"
	118	+```
	119	+
	120	+`@tla.extern` 根据 Python 函数注解描述 C ABI。声明函数的 Python 函数体不会执行：
	121	+
	122	+```python
	123	+@tla.extern(
	124	+ name="tla_user_gm_to_ub_f32",
	125	+ source=OP_SOURCE_CODES,
	126	+)
	127	+def tla_user_gm_to_ub_f32(
	128	+ gm_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.gm],
	129	+ ub_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.ub],
	130	+ ele_num: tla.Int32,
	131	+) -> None: ...
	132	+```
	133	+
	134	+### 在 TLA Kernel 中调用外部函数
	135	+
	136	+调用被 `@tla.extern` 装饰的函数会在 TLA IR 中生成 `tla.call_extern`。本样例在 AIV 区域内传入 GM/UB 指针和元素个数，连续完成两路输入搬运：
	137	+
	138	+```python
	139	+with tla.vector():
	140	+ tla_user_gm_to_ub_f32(gm_a.ptr, ub_ptr_a, TILE_ELE)
	141	+ tla_user_gm_to_ub_f32(gm_b.ptr, ub_ptr_b, TILE_ELE)
	142	+ tla.set_flag(ub_loaded)
	143	+ tla.wait_flag(ub_loaded)
	144	+
	145	+ with tla.vec.func(mode="simd"):
	146	+ # 使用 TLA DSL 完成分块 load、add 和 store
	147	+ # ...
	148	+```
	149	+
	150	+`tla.compile()` 会根据外部函数的实际调用区域选择目标。本样例仅在 AIV 区域调用，因此内联源码会由 Ascend C 编译器编译为 `extern.aiv.c310.bc`，再加入 `hivmc-a5 --link-aicore-bitcode` 的链接输入。外部源码内容和编译器信息也会参与 Kernel 缓存键计算；修改 `OP_SOURCE_CODES` 后会生成新的缓存项。
	151	+
	152	+### 在 AIC 和 AIV 中调用同一个符号
	153	+
	154	+`extern_dual_core.py` 的两个区域在 Python/TLA 层复用同一个 extern 声明：
	155	+
	156	+```python
	157	+with tla.cube():
	158	+ tla_user_store_i32(result.ptr, 0, AIC_VALUE)
	159	+
	160	+with tla.vector():
	161	+ index = (1 + tla.arch.sub_block_idx()) * ELEMENTS_PER_CACHE_LINE
	162	+ tla_user_store_i32(result.ptr, index, AIV_VALUE)
	163	+```
	164	+
	165	+Lower Extern Call 阶段保留原始符号 `tla_user_store_i32`。由于该符号同时从 AIC 和 AIV 区域调用，其声明会被标记为 `AIC_OR_AIV`，拆分后的两个 mixed kernel 入口仍然调用同一个符号。CCEC 分别使用 AIC 和 AIV target 编译同一份 `extern.cpp`，生成各自定义原始符号的 bitcode，再由 HIVMC 链接到对应的 mixed kernel 入口。

Apython/tla_dsl/examples/end_to_end/extern_op/extern_dual_core.py+93-0

@@ -0,0 +1,93 @@
	1	+"""One Ascend C function called from both AIC and AIV regions."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import argparse
	6	+
	7	+import catlass.tla as tla
	8	+from catlass.tla.runtime import from_dlpack
	9	+
	10	+
	11	+AIC_VALUE = 101
	12	+AIV_VALUE = 202
	13	+CACHE_LINE_BYTES = 64
	14	+INT32_BYTES = 4
	15	+ELEMENTS_PER_CACHE_LINE = CACHE_LINE_BYTES // INT32_BYTES
	16	+RESULT_SIZE = 3 * ELEMENTS_PER_CACHE_LINE
	17	+
	18	+OP_SOURCE_CODES = r"""
	19	+#include <cstdint>
	20	+#include "kernel_operator.h"
	21	+
	22	+extern "C" {
	23	+
	24	+[aicore] __attribute__((always_inline)) void tla_user_store_i32(
	25	+ uint64_t dst_gm_addr, int32_t index, int32_t value) {
	26	+ auto dst = reinterpret_cast<__gm__ int32_t *>(dst_gm_addr);
	27	+ dst[index] = value;
	28	+}
	29	+
	30	+} // extern "C"
	31	+"""
	32	+
	33	+
	34	+@tla.extern(
	35	+ name="tla_user_store_i32",
	36	+ source=OP_SOURCE_CODES,
	37	+)
	38	+def tla_user_store_i32(
	39	+ dst: tla.Pointer[tla.Int32, tla.AddressSpace.gm],
	40	+ index: tla.Int32,
	41	+ value: tla.Int32,
	42	+) -> None: ...
	43	+
	44	+
	45	+@tla.kernel
	46	+def extern_dual_core(result: tla.Tensor) -> None:
	47	+ with tla.cube():
	48	+ tla_user_store_i32(result.ptr, 0, AIC_VALUE)
	49	+ tla.pipe_barrier(tla.pipes.ALL)
	50	+
	51	+ with tla.vector():
	52	+ # Ensure each sub-block writes to a different cache line in the result
	53	+ # tensor. This avoids write conflicts.
	54	+ index = (1 + tla.arch.sub_block_idx()) * ELEMENTS_PER_CACHE_LINE
	55	+ tla_user_store_i32(result.ptr, index, AIV_VALUE)
	56	+ tla.pipe_barrier(tla.pipes.ALL)
	57	+
	58	+
	59	+def run(device: int = 0) -> None:
	60	+ import torch
	61	+ import torch_npu # noqa: F401
	62	+
	63	+ torch.npu.set_device(device)
	64	+ result = torch.zeros(RESULT_SIZE, dtype=torch.int32, device="npu")
	65	+ tla_result = from_dlpack(result, layout_tag=tla.arch.RowMajor)
	66	+
	67	+ executor = tla.compile(
	68	+ extern_dual_core,
	69	+ tla_result,
	70	+ options="--npu-arch 3510",
	71	+ )
	72	+ executor(tla_result, block_num=1)
	73	+ torch.npu.synchronize()
	74	+
	75	+ expected = torch.zeros(RESULT_SIZE, dtype=torch.int32, device="npu")
	76	+ expected[0] = AIC_VALUE
	77	+ expected[ELEMENTS_PER_CACHE_LINE] = AIV_VALUE
	78	+ expected[2 * ELEMENTS_PER_CACHE_LINE] = AIV_VALUE
	79	+ torch.testing.assert_close(result, expected, rtol=0.0, atol=0.0)
	80	+ print(
	81	+ f"passed; result={result.cpu().tolist()}; kernel={executor.kernel_binary_path}"
	82	+ )
	83	+
	84	+
	85	+def main() -> int:
	86	+ parser = argparse.ArgumentParser()
	87	+ parser.add_argument("--device", type=int, default=0)
	88	+ run(parser.parse_args().device)
	89	+ return 0
	90	+
	91	+
	92	+if __name__ == "__main__":
	93	+ raise SystemExit(main())

Apython/tla_dsl/examples/end_to_end/extern_op/extern_vecadd.py+114-0

@@ -0,0 +1,114 @@
	1	+"""Vecadd whose GM-to-UB copies are supplied by a user Ascend C function."""
	2	+
	3	+from __future__ import annotations
	4	+
	5	+import argparse
	6	+
	7	+import catlass.tla as tla
	8	+from catlass.tla.runtime import from_dlpack
	9	+
	10	+
	11	+TILE_ELE = 256
	12	+VL_ELE = 64
	13	+
	14	+OP_SOURCE_CODES = r"""
	15	+#include <cstdint>
	16	+#include "kernel_operator.h"
	17	+
	18	+extern "C" {
	19	+
	20	+[aicore] __attribute__((always_inline)) void tla_user_gm_to_ub_f32(
	21	+ uint64_t src_gm_addr, uint64_t dst_ub_addr, int32_t count) {
	22	+ AscendC::GlobalTensor<float> src;
	23	+ src.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(src_gm_addr),
	24	+ static_cast<uint32_t>(count));
	25	+ AscendC::LocalTensor<float> dst(AscendC::TPosition::VECCALC,
	26	+ static_cast<uint32_t>(dst_ub_addr),
	27	+ static_cast<uint32_t>(count));
	28	+ AscendC::DataCopy(dst, src, static_cast<uint32_t>(count));
	29	+}
	30	+
	31	+} // extern "C"
	32	+"""
	33	+
	34	+
	35	+@tla.extern(
	36	+ name="tla_user_gm_to_ub_f32",
	37	+ source=OP_SOURCE_CODES,
	38	+)
	39	+def tla_user_gm_to_ub_f32(
	40	+ gm_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.gm],
	41	+ ub_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.ub],
	42	+ ele_num: tla.Int32,
	43	+) -> None: ...
	44	+
	45	+
	46	+@tla.kernel
	47	+def extern_vecadd(
	48	+ gm_a: tla.Tensor,
	49	+ gm_b: tla.Tensor,
	50	+ gm_c: tla.Tensor,
	51	+) -> None:
	52	+ ub_loaded = tla.flag("ub_loaded", tla.arch.MTE2, tla.arch.VECTOR)
	53	+ vec_done = tla.flag("vec_done", tla.arch.VECTOR, tla.arch.MTE3)
	54	+
	55	+ ub_ptr_a = tla.allocate(TILE_ELE, tla.Float32, tla.AddressSpace.ub, 256)
	56	+ ub_ptr_b = tla.allocate(TILE_ELE, tla.Float32, tla.AddressSpace.ub, 256)
	57	+ ub_ptr_c = tla.allocate(TILE_ELE, tla.Float32, tla.AddressSpace.ub, 256)
	58	+ ub_a = tla.make_tensor_like(ub_ptr_a, gm_a, tla.arch.RowMajor)
	59	+ ub_b = tla.make_tensor_like(ub_ptr_b, gm_b, tla.arch.RowMajor)
	60	+ ub_c = tla.make_tensor_like(ub_ptr_c, gm_c, tla.arch.RowMajor)
	61	+
	62	+ with tla.vector():
	63	+ tla_user_gm_to_ub_f32(gm_a.ptr, ub_ptr_a, TILE_ELE)
	64	+ tla_user_gm_to_ub_f32(gm_b.ptr, ub_ptr_b, TILE_ELE)
	65	+ tla.set_flag(ub_loaded)
	66	+ tla.wait_flag(ub_loaded)
	67	+
	68	+ with tla.vec.func(mode="simd"):
	69	+ for i in tla.range(TILE_ELE // VL_ELE):
	70	+ a = tla.tile_view(ub_a, tla.make_shape(VL_ELE), tla.make_coord(i))
	71	+ b = tla.tile_view(ub_b, tla.make_shape(VL_ELE), tla.make_coord(i))
	72	+ c = tla.tile_view(ub_c, tla.make_shape(VL_ELE), tla.make_coord(i))
	73	+ c.store(tla.add(a.load(), b.load()))
	74	+
	75	+ tla.set_flag(vec_done)
	76	+ tla.wait_flag(vec_done)
	77	+ tla.copy(gm_c, ub_c)
	78	+ tla.pipe_barrier(tla.pipes.ALL)
	79	+
	80	+
	81	+def run(device: int = 0) -> None:
	82	+ import torch
	83	+ import torch_npu # noqa: F401
	84	+
	85	+ torch.npu.set_device(device)
	86	+ a = torch.rand(TILE_ELE, dtype=torch.float32, device="npu")
	87	+ b = torch.rand(TILE_ELE, dtype=torch.float32, device="npu")
	88	+ c = torch.empty_like(a)
	89	+ tla_a = from_dlpack(a, layout_tag=tla.arch.RowMajor)
	90	+ tla_b = from_dlpack(b, layout_tag=tla.arch.RowMajor)
	91	+ tla_c = from_dlpack(c, layout_tag=tla.arch.RowMajor)
	92	+
	93	+ executor = tla.compile(
	94	+ extern_vecadd,
	95	+ tla_a,
	96	+ tla_b,
	97	+ tla_c,
	98	+ options="--npu-arch 3510",
	99	+ )
	100	+ executor(tla_a, tla_b, tla_c, block_num=1)
	101	+ torch.npu.synchronize()
	102	+ torch.testing.assert_close(c, a + b, rtol=0.0, atol=1e-4)
	103	+ print(f"passed; kernel={executor.kernel_binary_path}")
	104	+
	105	+
	106	+def main() -> int:
	107	+ parser = argparse.ArgumentParser()
	108	+ parser.add_argument("--device", type=int, default=0)
	109	+ run(parser.parse_args().device)
	110	+ return 0
	111	+
	112	+
	113	+if __name__ == "__main__":
	114	+ raise SystemExit(main())

Apython/tla_dsl/tests/lit/tla-compile/auto-mutex-extern-error.mlir+15-0

@@ -0,0 +1,15 @@
	1	+// RUN: not %tla_compile %s -o - 2>&1 \| %filecheck %s
	2	+
	3	+module {
	4	+ tla.func @auto_mutex_extern(%gm: !tla.ptr<f32, gm, 4>) attributes {tla.auto_sync = "v0"} {
	5	+ %ub = tla.alloc_ptr{size_bytes = 1024} -> !tla.ptr<f32, ub, 256>
	6	+ %count = arith.constant 256 : i32
	7	+ "tla.vector"() ({
	8	+ tla.call_extern @tla_user_gm_to_ub_f32(%gm, %ub, %count) :
	9	+ (!tla.ptr<f32, gm, 4>, !tla.ptr<f32, ub, 256>, i32) -> ()
	10	+ }) : () -> ()
	11	+ tla.return
	12	+ }
	13	+}
	14	+
	15	+// CHECK: error: auto_sync='v0' cannot be combined with tla.call_extern; external calls require explicit synchronization in v1

Apython/tla_dsl/tests/lit/tla-compile/extern-call-symbol-conflict.mlir+13-0

@@ -0,0 +1,13 @@
	1	+// RUN: not %tla_compile %s -o - 2>&1 \| %filecheck %s
	2	+
	3	+module {
	4	+ tla.func @conflicting_symbol() {
	5	+ %count = arith.constant 1 : i32
	6	+ "tla.vector"() ({
	7	+ tla.call_extern @conflicting_symbol(%count) : (i32) -> ()
	8	+ }) : () -> ()
	9	+ tla.return
	10	+ }
	11	+}
	12	+
	13	+// CHECK: error: 'tla.call_extern' op external symbol @conflicting_symbol conflicts with a defined function

Apython/tla_dsl/tests/lit/tla-compile/extern-call.mlir+47-0

@@ -0,0 +1,47 @@
	1	+// RUN: %tla_compile %s -o - \| %filecheck %s
	2	+
	3	+module {
	4	+ tla.func @extern_call_kernel() {
	5	+ %addr = arith.constant 0 : i64
	6	+ %gm = tla.inttoptr %addr : i64 -> !tla.ptr<f32, gm, 4>
	7	+ %ub = tla.alloc_ptr{size_bytes = 1024} -> !tla.ptr<f32, ub, 256>
	8	+ %count = arith.constant 256 : i32
	9	+ "tla.vector"() ({
	10	+ tla.call_extern @tla_user_gm_to_ub_f32(%gm, %ub, %count) :
	11	+ (!tla.ptr<f32, gm, 4>, !tla.ptr<f32, ub, 256>, i32) -> ()
	12	+ }) : () -> ()
	13	+ tla.return
	14	+ }
	15	+
	16	+ tla.func @extern_call_cube_kernel() {
	17	+ %value = arith.constant 1 : i32
	18	+ "tla.cube"() ({
	19	+ tla.call_extern @tla_user_cube_only(%value) : (i32) -> ()
	20	+ }) : () -> ()
	21	+ tla.return
	22	+ }
	23	+
	24	+ tla.func @extern_call_mix_kernel() {
	25	+ %value = arith.constant 1 : i32
	26	+ "tla.cube"() ({
	27	+ tla.call_extern @tla_user_shared(%value) : (i32) -> ()
	28	+ }) : () -> ()
	29	+ "tla.vector"() ({
	30	+ tla.call_extern @tla_user_shared(%value) : (i32) -> ()
	31	+ }) : () -> ()
	32	+ tla.return
	33	+ }
	34	+}
	35	+
	36	+// CHECK-DAG: func.func private @tla_user_gm_to_ub_f32(i64, i64, i32) attributes {hivm.func_core_type = #hivm.func_core_type<AIV>}
	37	+// CHECK-DAG: func.func private @tla_user_cube_only(i32) attributes {hivm.func_core_type = #hivm.func_core_type<AIC>}
	38	+// CHECK-DAG: func.func private @tla_user_shared(i32) attributes {hivm.func_core_type = #hivm.func_core_type<AIC_OR_AIV>}
	39	+// CHECK-LABEL: func.func @extern_call_kernel
	40	+// CHECK: call @tla_user_gm_to_ub_f32({{.*}}) : (i64, i64, i32) -> ()
	41	+// CHECK-LABEL: func.func @extern_call_cube_kernel
	42	+// CHECK: call @tla_user_cube_only({{.*}}) : (i32) -> ()
	43	+// CHECK-LABEL: func.func @extern_call_mix_kernel_mix_aic
	44	+// CHECK: call @tla_user_shared({{.*}}) : (i32) -> ()
	45	+// CHECK-LABEL: func.func @extern_call_mix_kernel_mix_aiv
	46	+// CHECK: call @tla_user_shared({{.*}}) : (i32) -> ()
	47	+// CHECK-NOT: tla.call_extern

Mpython/tla_dsl/tests/test_compiler_kernel_abi.py+2-0

	@@ -490,6 +490,8 @@ def test_kernel_abi_from_dict_rejects_incoherent_structured_scalar(
490		490
491	class _FakeLowered:	491	class _FakeLowered:
492	module = object()	492	module = object()
		493	+ extern_function = None
		494	+ extern_core_types = frozenset()
493		495
494	def asm(self, *, generic: bool = False) -> str:	496	def asm(self, *, generic: bool = False) -> str:
495	del generic	497	del generic

Mpython/tla_dsl/tests/test_execution_hivm_runtime.py+2-0

	@@ -631,6 +631,8 @@ def test_debug_print_aic_output_uses_scalar_frame(
631	class _FakeLowered:	631	class _FakeLowered:
632	def __init__(self, text: str, module: object \| None = None) -> None:	632	def __init__(self, text: str, module: object \| None = None) -> None:
633	self.module = module	633	self.module = module
		634	+ self.extern_function = None
		635	+ self.extern_core_types = frozenset()
634	self._text = text	636	self._text = text
635		637
636	def asm(self, *, generic: bool = False) -> str:	638	def asm(self, *, generic: bool = False) -> str:

Apython/tla_dsl/tests/test_extern_op.py+473-0

@@ -0,0 +1,473 @@
	1	+from __future__ import annotations
	2	+
	3	+import inspect
	4	+from pathlib import Path
	5	+import importlib.util
	6	+
	7	+import pytest
	8	+
	9	+import catlass.tla as tla
	10	+from catlass.base_dsl import BaseDSL
	11	+from catlass.runtime import TlaCoreAPIError
	12	+from catlass import execution
	13	+from catlass.base_dsl.arch import get_kernel_target
	14	+
	15	+
	16	+EXAMPLE_PATH = (
	17	+ Path(__file__).parents[1]
	18	+ / "examples"
	19	+ / "end_to_end"
	20	+ / "extern_op"
	21	+ / "extern_vecadd.py"
	22	+)
	23	+DUAL_CORE_EXAMPLE_PATH = EXAMPLE_PATH.with_name("extern_dual_core.py")
	24	+EXTERN_SOURCE_CODE = 'extern "C" void tla_user_gm_to_ub_f32() {}\n'
	25	+DUAL_CORE_SOURCE_CODE = r"""
	26	+#include <cstdint>
	27	+
	28	+extern "C" {
	29	+[aicore] __attribute__((noinline)) void tla_user_dual_core(int32_t value) {
	30	+ (void)value;
	31	+}
	32	+}
	33	+"""
	34	+
	35	+
	36	+def _gm_to_ub(source_code: str = EXTERN_SOURCE_CODE):
	37	+ @tla.extern(
	38	+ source=source_code,
	39	+ name="tla_user_gm_to_ub_f32",
	40	+ )
	41	+ def gm_to_ub(
	42	+ gm_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.gm],
	43	+ ub_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.ub],
	44	+ ele_num: tla.Int32,
	45	+ ) -> None: ...
	46	+
	47	+ return gm_to_ub
	48	+
	49	+
	50	+GM_TO_UB = _gm_to_ub()
	51	+
	52	+
	53	+@tla.extern(source=EXTERN_SOURCE_CODE, name="tla_user_other_op")
	54	+def OTHER_EXTERN(
	55	+ gm_ptr: tla.Pointer[tla.Float32, tla.AddressSpace.gm],
	56	+) -> None: ...
	57	+
	58	+
	59	+@tla.extern(source=DUAL_CORE_SOURCE_CODE)
	60	+def tla_user_dual_core(value: tla.Int32) -> None: ...
	61	+
	62	+
	63	+@tla.kernel
	64	+def extern_load_kernel(gm: tla.Tensor) -> None:
	65	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	66	+ with tla.vector():
	67	+ GM_TO_UB(gm.ptr, ub, 256)
	68	+
	69	+
	70	+@tla.kernel
	71	+def two_extern_ops_kernel(gm: tla.Tensor) -> None:
	72	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	73	+ with tla.vector():
	74	+ GM_TO_UB(gm.ptr, ub, 256)
	75	+ OTHER_EXTERN(gm.ptr)
	76	+
	77	+
	78	+@tla.kernel
	79	+def wrong_address_space_kernel(gm: tla.Tensor) -> None:
	80	+ with tla.vector():
	81	+ GM_TO_UB(gm.ptr, gm.ptr, 256)
	82	+
	83	+
	84	+@tla.kernel
	85	+def tensor_argument_kernel(gm: tla.Tensor) -> None:
	86	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	87	+ with tla.vector():
	88	+ GM_TO_UB(gm, ub, 256)
	89	+
	90	+
	91	+@tla.kernel
	92	+def outside_vector_kernel(gm: tla.Tensor) -> None:
	93	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	94	+ GM_TO_UB(gm.ptr, ub, 256)
	95	+
	96	+
	97	+@tla.kernel
	98	+def extern_cube_kernel(value: tla.Int32) -> None:
	99	+ with tla.cube():
	100	+ tla_user_dual_core(value)
	101	+
	102	+
	103	+@tla.kernel
	104	+def extern_mix_kernel(value: tla.Int32) -> None:
	105	+ with tla.cube():
	106	+ tla_user_dual_core(value)
	107	+ with tla.vector():
	108	+ tla_user_dual_core(value)
	109	+
	110	+
	111	+@tla.kernel
	112	+def inside_vec_func_kernel(gm: tla.Tensor) -> None:
	113	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	114	+ with tla.vector():
	115	+ with tla.vec.func(mode="simd"):
	116	+ GM_TO_UB(gm.ptr, ub, 256)
	117	+
	118	+
	119	+@tla.kernel
	120	+def wrong_scalar_type_kernel(gm: tla.Tensor) -> None:
	121	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	122	+ with tla.vector():
	123	+ GM_TO_UB(gm.ptr, ub, tla.Int64(256))
	124	+
	125	+
	126	+@tla.kernel
	127	+def wrong_argument_count_kernel(gm: tla.Tensor) -> None:
	128	+ ub = tla.allocate(256, tla.Float32, tla.AddressSpace.ub, 256)
	129	+ with tla.vector():
	130	+ GM_TO_UB(gm.ptr, ub)
	131	+
	132	+
	133	+@tla.kernel
	134	+def unused_extern_kernel(gm: tla.Tensor) -> None:
	135	+ pass
	136	+
	137	+
	138	+def _fake_gm_tensor():
	139	+ return tla.make_fake_tensor(
	140	+ tla.Float32,
	141	+ (256,),
	142	+ (1,),
	143	+ origin_shape=(256,),
	144	+ coord=(0,),
	145	+ layout_tag=tla.arch.RowMajor,
	146	+ )
	147	+
	148	+
	149	+def _lower_kernel(kernel, *type_args):
	150	+ return BaseDSL()._lower(
	151	+ kernel.fn,
	152	+ kind="kernel",
	153	+ options={},
	154	+ type_args=type_args,
	155	+ location=kernel.decorator_location,
	156	+ )
	157	+
	158	+
	159	+def test_extern_decorator_describes_inline_source_without_compiling(
	160	+ monkeypatch,
	161	+) -> None:
	162	+ source_code = 'extern "C" void op() {}\n'
	163	+ calls = []
	164	+ monkeypatch.setattr("subprocess.run", lambda args, *kwargs: calls.append(args))
	165	+
	166	+ @tla.extern(source=source_code)
	167	+ def op(value: tla.Int32) -> None:
	168	+ raise AssertionError("the declaration body must not execute")
	169	+
	170	+ assert op.source == source_code
	171	+ assert op.symbol == "op"
	172	+ assert op.arg_types == (tla.Int32,)
	173	+ assert callable(op)
	174	+ assert calls == []
	175	+
	176	+
	177	+def test_pointer_subscription_returns_public_typed_pointer() -> None:
	178	+ pointer_type = tla.Pointer[tla.Float32, tla.AddressSpace.gm]
	179	+
	180	+ assert isinstance(pointer_type, tla.TypedPointer)
	181	+ assert pointer_type.dtype is tla.Float32
	182	+ assert pointer_type.space is tla.AddressSpace.gm
	183	+
	184	+
	185	+def test_pointer_subscription_requires_dtype_and_memory_space() -> None:
	186	+ with pytest.raises(TypeError, match="expects \$dtype, memory_space\$"):
	187	+ tla.Pointer[tla.Float32]
	188	+
	189	+
	190	+def test_extern_rejects_non_string_or_empty_source(tmp_path) -> None:
	191	+ with pytest.raises(TypeError, match="source must be str"):
	192	+ tla.extern(source=tmp_path / "op.cpp")
	193	+ with pytest.raises(ValueError, match="source must not be empty"):
	194	+ tla.extern(source=" \n")
	195	+
	196	+
	197	+def test_extern_rejects_invalid_name() -> None:
	198	+ with pytest.raises(ValueError, match="name must be a C identifier"):
	199	+ tla.extern(source=EXTERN_SOURCE_CODE, name="op-name")
	200	+
	201	+
	202	+def test_extern_rejects_invalid_function_signature() -> None:
	203	+ def missing_return(value: tla.Int32): ...
	204	+
	205	+ with pytest.raises(TypeError, match="return annotation.*is missing"):
	206	+ tla.extern(source=EXTERN_SOURCE_CODE)(missing_return)
	207	+
	208	+ def variadic(*values: tla.Int32) -> None: ...
	209	+
	210	+ with pytest.raises(TypeError, match="positional and fixed"):
	211	+ tla.extern(source=EXTERN_SOURCE_CODE)(variadic)
	212	+
	213	+ def missing_parameter_annotation(value) -> None: ...
	214	+
	215	+ with pytest.raises(TypeError, match="missing an annotation"):
	216	+ tla.extern(source=EXTERN_SOURCE_CODE)(missing_parameter_annotation)
	217	+
	218	+
	219	+def test_extern_frontend_emits_call_and_tracks_dependency() -> None:
	220	+ lowered = _lower_kernel(extern_load_kernel, _fake_gm_tensor())
	221	+ mlir = lowered.asm()
	222	+ assert "tla.call_extern" in mlir
	223	+ assert "@tla_user_gm_to_ub_f32" in mlir
	224	+ assert "!tla.ptr<f32, gm" in mlir
	225	+ assert "!tla.ptr<f32, ub" in mlir
	226	+ assert lowered.extern_function is GM_TO_UB
	227	+ assert lowered.extern_core_types == frozenset({"aiv"})
	228	+
	229	+
	230	+def test_extern_call_location_points_to_user_call() -> None:
	231	+ source_lines, first_line = inspect.getsourcelines(extern_load_kernel.fn)
	232	+ call_line = first_line + next(
	233	+ index for index, line in enumerate(source_lines) if "GM_TO_UB(" in line
	234	+ )
	235	+
	236	+ lowered = _lower_kernel(extern_load_kernel, _fake_gm_tensor())
	237	+ with lowered.context:
	238	+ mlir = lowered.module.operation.get_asm(
	239	+ enable_debug_info=True,
	240	+ assume_verified=False,
	241	+ )
	242	+
	243	+ assert f'test_extern_op.py":{call_line}:' in mlir
	244	+
	245	+
	246	+def test_extern_rejects_second_op_in_one_kernel() -> None:
	247	+ with pytest.raises(TlaCoreAPIError, match="at most one external function"):
	248	+ two_extern_ops_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	249	+
	250	+
	251	+def test_extern_checks_pointer_address_space() -> None:
	252	+ with pytest.raises(TlaCoreAPIError, match="expects pointer.*ub"):
	253	+ wrong_address_space_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	254	+
	255	+
	256	+def test_extern_rejects_tensor_and_requests_explicit_pointer() -> None:
	257	+ with pytest.raises(TlaCoreAPIError, match=r"Tensor; pass tensor\.ptr explicitly"):
	258	+ tensor_argument_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	259	+
	260	+
	261	+def test_extern_checks_call_region() -> None:
	262	+ with pytest.raises(TlaCoreAPIError, match="exactly one.vector.cube"):
	263	+ outside_vector_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	264	+ with pytest.raises(TlaCoreAPIError, match="outside tla.vec.func"):
	265	+ inside_vec_func_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	266	+
	267	+
	268	+def test_extern_infers_aic_and_mixed_call_targets() -> None:
	269	+ cube = _lower_kernel(extern_cube_kernel, tla.Int32(0))
	270	+ mixed = _lower_kernel(extern_mix_kernel, tla.Int32(0))
	271	+
	272	+ assert cube.extern_core_types == frozenset({"aic"})
	273	+ assert mixed.extern_core_types == frozenset({"aic", "aiv"})
	274	+ assert mixed.asm().count("tla_user_dual_core") == 2
	275	+
	276	+
	277	+def test_extern_checks_argument_count_and_scalar_type() -> None:
	278	+ with pytest.raises(TlaCoreAPIError, match="expects 3 arguments, got 2"):
	279	+ wrong_argument_count_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	280	+ with pytest.raises(TlaCoreAPIError, match="expects Int32, got Int64"):
	281	+ wrong_scalar_type_kernel.dump_mlir(type_args=(_fake_gm_tensor(),))
	282	+
	283	+
	284	+def test_unused_extern_declaration_is_not_a_kernel_dependency() -> None:
	285	+ lowered = _lower_kernel(unused_extern_kernel, _fake_gm_tensor())
	286	+
	287	+ assert lowered.extern_function is None
	288	+ assert lowered.extern_core_types == frozenset()
	289	+
	290	+
	291	+def test_ascendc_compile_command_reuses_bitcode_backend(tmp_path, monkeypatch) -> None:
	292	+ compiler = tmp_path / "ccec"
	293	+ compiler.write_text("")
	294	+ commands = []
	295	+
	296	+ def fake_run(command, **kwargs):
	297	+ del kwargs
	298	+ commands.append(command)
	299	+ Path(command[command.index("-o") + 1]).write_bytes(b"BC")
	300	+
	301	+ monkeypatch.setattr(execution, "_resolve_ccec", lambda: compiler)
	302	+ monkeypatch.setattr(execution, "_run_checked", fake_run)
	303	+ monkeypatch.setenv("ASCEND_HOME_PATH", str(tmp_path))
	304	+
	305	+ result = execution._compile_ascendc_extern_function(
	306	+ GM_TO_UB,
	307	+ artifact_dir=tmp_path,
	308	+ targets=(get_kernel_target(target_arch="c310", core_type="aiv"),),
	309	+ )
	310	+
	311	+ assert [path.name for path in result] == ["extern.aiv.c310.bc"]
	312	+ command = commands[0]
	313	+ assert "-emit-llvm" in command
	314	+ assert "--cce-aicore-arch=dav-c310-vec" in command
	315	+ source = tmp_path / "extern.cpp"
	316	+ assert str(source) in command
	317	+ assert source.read_text() == EXTERN_SOURCE_CODE
	318	+
	319	+
	320	+def test_ascendc_multi_target_compile_uses_one_ccec_invocation_per_core(
	321	+ tmp_path, monkeypatch
	322	+) -> None:
	323	+ compiler = tmp_path / "ccec"
	324	+ compiler.write_text("")
	325	+ commands = []
	326	+
	327	+ def fake_run(command, **kwargs):
	328	+ del kwargs
	329	+ commands.append(command)
	330	+ Path(command[command.index("-o") + 1]).write_bytes(b"BC")
	331	+
	332	+ monkeypatch.setattr(execution, "_resolve_ccec", lambda: compiler)
	333	+ monkeypatch.setattr(execution, "_run_checked", fake_run)
	334	+ monkeypatch.setenv("ASCEND_HOME_PATH", str(tmp_path))
	335	+ targets = execution._resolve_extern_targets(
	336	+ tla_user_dual_core,
	337	+ extern_core_types={"aic", "aiv"},
	338	+ target_arch="c310",
	339	+ )
	340	+
	341	+ result = execution._compile_ascendc_extern_function(
	342	+ tla_user_dual_core,
	343	+ artifact_dir=tmp_path,
	344	+ targets=targets,
	345	+ )
	346	+
	347	+ assert [path.name for path in result] == [
	348	+ "extern.aic.c310.bc",
	349	+ "extern.aiv.c310.bc",
	350	+ ]
	351	+ assert len(commands) == 2
	352	+ assert "--cce-aicore-arch=dav-c310-cube" in commands[0]
	353	+ assert "--cce-aicore-arch=dav-c310-vec" in commands[1]
	354	+ assert commands[0][commands[0].index("-o") + 1].endswith(
	355	+ "extern.aic.c310.bc"
	356	+ )
	357	+ assert commands[1][commands[1].index("-o") + 1].endswith(
	358	+ "extern.aiv.c310.bc"
	359	+ )
	360	+
	361	+
	362	+def _cache_key_kwargs(tmp_path, target):
	363	+ return {
	364	+ "tlair_mlir": "module { func.func @kernel() }",
	365	+ "entrypoint": "kernel",
	366	+ "runtime": execution.TlaRuntimeOptions(arch_scope="aiv.c310"),
	367	+ "compiler_bridge_path": None,
	368	+ "hivmc": tmp_path / "hivmc-a5",
	369	+ "target": target,
	370	+ }
	371	+
	372	+
	373	+def test_extern_compile_identity_participates_in_kernel_cache_key(
	374	+ tmp_path, monkeypatch
	375	+) -> None:
	376	+ compiler = (tmp_path / "toolchain" / "bin" / "ccec").resolve()
	377	+ compiler.parent.mkdir(parents=True)
	378	+ compiler.write_bytes(b"ccec")
	379	+ compiler_version = ["first"]
	380	+ target = get_kernel_target(target_arch="c310", core_type="aiv")
	381	+ monkeypatch.setenv("ASCEND_HOME_PATH", str(tmp_path))
	382	+ monkeypatch.setattr(execution, "_resolve_ccec", lambda: compiler)
	383	+ monkeypatch.setattr(
	384	+ execution,
	385	+ "_tool_version",
	386	+ lambda path: compiler_version[0] if path == compiler else "hivmc-version",
	387	+ )
	388	+ monkeypatch.setattr(execution, "_tool_fingerprint", lambda _path: "fingerprint")
	389	+ monkeypatch.setattr(
	390	+ execution,
	391	+ "_ascendc_include_dirs",
	392	+ lambda _ascend_home: [tmp_path / "include"],
	393	+ )
	394	+ kwargs = _cache_key_kwargs(tmp_path, target)
	395	+ kwargs.update(extern_function=GM_TO_UB, extern_targets=(target,))
	396	+ first = execution._cache_key(**kwargs)
	397	+
	398	+ compiler_version[0] = "second"
	399	+
	400	+ assert execution._cache_key(**kwargs) != first
	401	+
	402	+
	403	+def test_extern_source_participates_in_kernel_cache_key(tmp_path, monkeypatch) -> None:
	404	+ target = get_kernel_target(target_arch="c310", core_type="aiv")
	405	+ monkeypatch.setattr(execution, "_tool_version", lambda _path: "version")
	406	+ monkeypatch.setattr(execution, "_tool_fingerprint", lambda _path: "fingerprint")
	407	+ monkeypatch.setattr(
	408	+ execution,
	409	+ "_ascendc_extern_compile_identity",
	410	+ lambda: {"revision": "same"},
	411	+ )
	412	+ kwargs = _cache_key_kwargs(tmp_path, target)
	413	+
	414	+ first = execution._cache_key(
	415	+ **kwargs,
	416	+ extern_function=_gm_to_ub('extern "C" void op() {}\n'),
	417	+ extern_targets=(target,),
	418	+ )
	419	+ second = execution._cache_key(
	420	+ **kwargs,
	421	+ extern_function=_gm_to_ub('extern "C" void op() { /* changed */ }\n'),
	422	+ extern_targets=(target,),
	423	+ )
	424	+
	425	+ assert second != first
	426	+
	427	+
	428	+def _load_vecadd_example():
	429	+ spec = importlib.util.spec_from_file_location("extern_vecadd_example", EXAMPLE_PATH)
	430	+ assert spec and spec.loader
	431	+ module = importlib.util.module_from_spec(spec)
	432	+ spec.loader.exec_module(module)
	433	+ return module
	434	+
	435	+
	436	+def _load_dual_core_example():
	437	+ spec = importlib.util.spec_from_file_location(
	438	+ "extern_dual_core_example", DUAL_CORE_EXAMPLE_PATH
	439	+ )
	440	+ assert spec and spec.loader
	441	+ module = importlib.util.module_from_spec(spec)
	442	+ spec.loader.exec_module(module)
	443	+ return module
	444	+
	445	+
	446	+def test_extern_vecadd_frontend_uses_custom_load_and_tla_compute() -> None:
	447	+ example = _load_vecadd_example()
	448	+ tensor = _fake_gm_tensor()
	449	+ lowered = _lower_kernel(example.extern_vecadd, tensor, tensor, tensor)
	450	+ mlir = lowered.asm()
	451	+ assert mlir.count("tla.call_extern") == 2
	452	+ assert "tla.add" in mlir
	453	+ assert "tla.copy" in mlir
	454	+ assert lowered.extern_function is example.tla_user_gm_to_ub_f32
	455	+
	456	+
	457	+def test_extern_dual_core_example_calls_one_op_from_aic_and_aiv() -> None:
	458	+ example = _load_dual_core_example()
	459	+ tensor = tla.make_fake_tensor(
	460	+ tla.Int32,
	461	+ (example.RESULT_SIZE,),
	462	+ (1,),
	463	+ origin_shape=(example.RESULT_SIZE,),
	464	+ coord=(0,),
	465	+ layout_tag=tla.arch.RowMajor,
	466	+ )
	467	+ lowered = _lower_kernel(example.extern_dual_core, tensor)
	468	+ mlir = lowered.asm()
	469	+
	470	+ assert mlir.count("tla.call_extern") == 2
	471	+ assert mlir.count("@tla_user_store_i32") == 2
	472	+ assert lowered.extern_function is example.tla_user_store_i32
	473	+ assert lowered.extern_core_types == frozenset({"aic", "aiv"})

Mtests/dsl_battery/test_battery.py+10-0

	@@ -150,6 +150,16 @@ def _cases(device: int) -> Iterator[tuple[str, list[list[str]]]]:
150	]],	150	]],
151	)	151	)
152		152
		153	+ # --- extern_op: user-provided Ascend C functions called from TLA kernels ---
		154	+ yield (
		155	+ "extern-vecadd",
		156	+ [["extern_op/extern_vecadd.py", *dev]],
		157	+ )
		158	+ yield (
		159	+ "extern-dual-core",
		160	+ [["extern_op/extern_dual_core.py", *dev]],
		161	+ )
		162	+
153	# --- basic_mixed ---	163	# --- basic_mixed ---
154	# One test on purpose: the second invocation exercises cache reuse and is	164	# One test on purpose: the second invocation exercises cache reuse and is
155	# only meaningful straight after the one that populates the cache.	165	# only meaningful straight after the one that populates the cache.

Mtests/run_dsl_test.sh+12-0

	@@ -12,6 +12,7 @@
12	# End-to-end validation for python/tla_dsl/examples/end_to_end/basic_mmad (basic_matmul*.py, basic_mmad_ptr.py,	12	# End-to-end validation for python/tla_dsl/examples/end_to_end/basic_mmad (basic_matmul*.py, basic_mmad_ptr.py,
13	# basic_matmul_l0c2l1.py),	13	# basic_matmul_l0c2l1.py),
14	# python/tla_dsl/examples/end_to_end/basic_vadd (basic_vadd.py),	14	# python/tla_dsl/examples/end_to_end/basic_vadd (basic_vadd.py),
		15	+# python/tla_dsl/examples/end_to_end/extern_op (extern_vecadd.py, extern_dual_core.py),
15	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed.py), python/tla_dsl/examples/end_to_end/basic_mixed_mutex (basic_mixed_mutex.py) and	16	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed.py), python/tla_dsl/examples/end_to_end/basic_mixed_mutex (basic_mixed_mutex.py) and
16	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed_ub2l1.py, basic_mixed_store_zN.py,	17	# python/tla_dsl/examples/end_to_end/basic_mixed (basic_mixed_ub2l1.py, basic_mixed_store_zN.py,
17	# basic_mixed_store_zNUnAlign.py, basic_mixed_fixpipe_nz2dn.py).	18	# basic_mixed_store_zNUnAlign.py, basic_mixed_fixpipe_nz2dn.py).
	@@ -99,6 +100,8 @@ BASIC_MMAD_AUTO_SYNC_REL="examples/end_to_end/basic_mmad/basic_matmul_auto_sync.
99	BASIC_MMAD_PTR_REL="examples/end_to_end/basic_mmad/basic_mmad_ptr.py"	100	BASIC_MMAD_PTR_REL="examples/end_to_end/basic_mmad/basic_mmad_ptr.py"
100	BASIC_MMAD_L0C2L1_REL="examples/end_to_end/basic_mmad/basic_matmul_l0c2l1.py"	101	BASIC_MMAD_L0C2L1_REL="examples/end_to_end/basic_mmad/basic_matmul_l0c2l1.py"
101	BASIC_VADD_REL="examples/end_to_end/basic_vadd/basic_vadd.py"	102	BASIC_VADD_REL="examples/end_to_end/basic_vadd/basic_vadd.py"
		103	+EXTERN_VECADD_REL="examples/end_to_end/extern_op/extern_vecadd.py"
		104	+EXTERN_DUAL_CORE_REL="examples/end_to_end/extern_op/extern_dual_core.py"
102	BASIC_MIXED_REL="examples/end_to_end/basic_mixed/basic_mixed.py"	105	BASIC_MIXED_REL="examples/end_to_end/basic_mixed/basic_mixed.py"
103	BASIC_MIXED_MUTEX_REL="examples/end_to_end/basic_mixed/basic_mixed_mutex.py"	106	BASIC_MIXED_MUTEX_REL="examples/end_to_end/basic_mixed/basic_mixed_mutex.py"
104	BASIC_MIXED_UB2L1_REL="examples/end_to_end/basic_mixed/basic_mixed_ub2l1.py"	107	BASIC_MIXED_UB2L1_REL="examples/end_to_end/basic_mixed/basic_mixed_ub2l1.py"
	@@ -168,6 +171,7 @@ Run end-to-end validation for:
168	- basic_mmad_ptr (basic_mmad_ptr.py)	171	- basic_mmad_ptr (basic_mmad_ptr.py)
169	- basic_mmad_l0c2l1 (basic_matmul_l0c2l1.py: cube-only E=(A@B)@D with L0C->L1 staging)	172	- basic_mmad_l0c2l1 (basic_matmul_l0c2l1.py: cube-only E=(A@B)@D with L0C->L1 staging)
170	- basic_vadd (basic_vadd.py with per-dtype CLI invocations, plus mutex variants)	173	- basic_vadd (basic_vadd.py with per-dtype CLI invocations, plus mutex variants)
		174	+ - extern_op (extern_vecadd.py and extern_dual_core.py)
171	- basic_mixed (basic_mixed.py with dynamic GM mnk list, including --use-mutex; basic_mixed_ub2l1.py,	175	- basic_mixed (basic_mixed.py with dynamic GM mnk list, including --use-mutex; basic_mixed_ub2l1.py,
172	basic_mixed_store_zN.py, basic_mixed_store_zNUnAlign.py for m=64/m=50)	176	basic_mixed_store_zN.py, basic_mixed_store_zNUnAlign.py for m=64/m=50)
173	- binary_op (binary_op.py <op> --all-dtypes for add/sub/mul/div/max/min/add_unalign/add_brc_b32)	177	- binary_op (binary_op.py <op> --all-dtypes for add/sub/mul/div/max/min/add_unalign/add_brc_b32)
	@@ -412,6 +416,14 @@ if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_VADD_REL}" ]]; then
412	echo "error: missing ${BASIC_VADD_REL} under ${CATLASS_DSL_DIR}" >&2	416	echo "error: missing ${BASIC_VADD_REL} under ${CATLASS_DSL_DIR}" >&2
413	exit 1	417	exit 1
414	fi	418	fi
		419	+if [[ ! -f "${CATLASS_DSL_DIR}/${EXTERN_VECADD_REL}" ]]; then
		420	+ echo "error: missing ${EXTERN_VECADD_REL} under ${CATLASS_DSL_DIR}" >&2
		421	+ exit 1
		422	+fi
		423	+if [[ ! -f "${CATLASS_DSL_DIR}/${EXTERN_DUAL_CORE_REL}" ]]; then
		424	+ echo "error: missing ${EXTERN_DUAL_CORE_REL} under ${CATLASS_DSL_DIR}" >&2
		425	+ exit 1
		426	+fi
415	if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_MIXED_REL}" ]]; then	427	if [[ ! -f "${CATLASS_DSL_DIR}/${BASIC_MIXED_REL}" ]]; then
416	echo "error: missing ${BASIC_MIXED_REL} under ${CATLASS_DSL_DIR}" >&2	428	echo "error: missing ${BASIC_MIXED_REL} under ${CATLASS_DSL_DIR}" >&2
417	exit 1	429	exit 1

	@@ -30,6 +30,7 @@ void registerTlaPasses()
30	registerTlaVectorRegionPass();	30	registerTlaVectorRegionPass();
31	registerTlaLowerFlagBarrierToHivmPass();	31	registerTlaLowerFlagBarrierToHivmPass();
32	registerTlaLowerPtrPass();	32	registerTlaLowerPtrPass();
		33	+ registerTlaLowerExternCallPass();
33	registerTlaLowerMutexToStdPass();	34	registerTlaLowerMutexToStdPass();
34	registerTlaCubeRegionPass();	35	registerTlaCubeRegionPass();
35	registerTlaFinalizeMemrefPass();	36	registerTlaFinalizeMemrefPass();
	@@ -47,6 +48,7 @@ void buildTlaPipeline(OpPassManager& pm)
47	// HACC machinery and the mixed-func split both consume.	48	// HACC machinery and the mixed-func split both consume.
48	pm.addPass(createTlaLowerFuncPass());	49	pm.addPass(createTlaLowerFuncPass());
49	pm.addPass(createTlaInsertAutoMutexPass());	50	pm.addPass(createTlaInsertAutoMutexPass());
		51	+ pm.addPass(createTlaLowerExternCallPass());
50	pm.addPass(createTlaLowerPtrPass());	52	pm.addPass(createTlaLowerPtrPass());
51	pm.addPass(createTlaSplitMixedFuncPass());	53	pm.addPass(createTlaSplitMixedFuncPass());
52	// Materialize tensor-view producer chains as tla.tensor_desc before region	54	// Materialize tensor-view producer chains as tla.tensor_desc before region