feat(dsl): compile class-organized kernels-catlass-AtomGit

Mpython/tla_dsl/catlass/base_dsl/compiler.py+158-19

	@@ -1,10 +1,117 @@
1	from __future__ import annotations	1	from __future__ import annotations
2		2
3	-from typing import Any	3	+import ast
		CheaterAbec26 days ago PR 描述接收者规则写：@tla.kernel 作用在 staticmethod/classmethod 上会在装饰期报错并提示顺序。本 diff 文件列表与 jit_runner/_make_kernel_wrapper 路径均未见该诊断，测试里实际也没有对应用例 yuantao 26 days ago Comments： 26 days ago Comments：
		4	+import functools
		5	+import inspect
		6	+import textwrap
		7	+from typing import Any, Callable, NoReturn
4		8
5	from .jit_executor import JitCompiledFunction, JitModule	9	from .jit_executor import JitCompiledFunction, JitModule
6		10
7		11
		12	+def _check_jit_call_is_pure_forwarding(
		13	+ fn: Callable[..., Any], *, class_name: str, error_type: type[Exception]
		14	+) -> tuple[str, tuple[str, ...]]:
		15	+ """Read a one-line ``__call__`` forward as ``(method name, forwarded names)``."""
		16	+
		17	+ from .ast_preprocessor import _find_function_def
		18	+
		19	+ def reject(reason: str) -> NoReturn:
		20	+ raise error_type(
		21	+ f"{class_name}.__call__ {reason} A @tla.jit __call__ must be one line "
		22	+ "calling self.<kernel>(<its own parameters, in order>)."
		23	+ )
		24	+
		25	+ try:
		26	+ source, _ = inspect.getsourcelines(fn)
		27	+ target = _find_function_def(
		28	+ ast.parse(textwrap.dedent("".join(source))), fn.__name__
		29	+ )
		30	+ except (OSError, TypeError, SyntaxError, ValueError):
		31	+ target = None
		32	+ if target is None:
		33	+ reject("has no readable source.")
		34	+ params = [arg.arg for arg in (target.args.posonlyargs, target.args.args)]
		35	+ if not params:
		36	+ reject("takes no parameters.")
		37	+ if params[0] not in ("self", "cls"):
		38	+ reject(f"does not take a receiver first (found {params[0]!r}).")
		39	+ body = target.body[1:] if ast.get_docstring(target) else target.body
		40	+ if len(body) != 1:
		41	+ reject(f"has {len(body)} statements instead of one.")
		42	+ statement = body[0]
		43	+ call = statement.value if isinstance(statement, (ast.Return, ast.Expr)) else None
		44	+ if not isinstance(call, ast.Call):
		45	+ reject("has a single statement, but it is not a call.")
		46	+ if call.keywords:
		47	+ reject("passes keyword arguments.")
		48	+ callee = call.func
		49	+ if (
		50	+ not isinstance(callee, ast.Attribute)
		51	+ or not isinstance(callee.value, ast.Name)
		52	+ or callee.value.id != params[0]
		53	+ ):
		54	+ reject(f"calls {ast.unparse(callee)}, not {params[0]}.<kernel>(...).")
		55	+ computed = [ast.unparse(arg) for arg in call.args if not isinstance(arg, ast.Name)]
		56	+ if computed:
		57	+ reject(f"passes {', '.join(computed)} instead of its own parameter names.")
		58	+ forwarded = [arg.id for arg in call.args]
		59	+ if forwarded != params[1:]:
		60	+ reject(
		61	+ f"forwards ({', '.join(forwarded)}) instead of its own parameters "
		62	+ f"({', '.join(params[1:])}) in order."
		63	+ )
		64	+ return callee.attr, tuple(forwarded)
		65	+
		66	+
		67	+def _resolve_kernel_target(
		68	+ func: Any, *, error_type: type[Exception]
		69	+) -> tuple[Any, Any]:
		70	+ """Return ``(kernel, receiver)`` for one ``tla.compile`` entry."""
		71	+
		72	+ from ..dsl import TlaJitFunction, is_jit_callable, unwrap_jit_callable
		73	+ from .runtime.jit_arg_adapters import receiver_argument_name
		74	+
		75	+ if isinstance(func, TlaJitFunction):
		76	+ return func, func.receiver
		77	+ if isinstance(func, (type, functools.partial)) or inspect.isroutine(func):
		78	+ raise error_type(
		79	+ "tla.compile expects a @tla.kernel function, or an instance of a class "
		80	+ "that declares one."
		81	+ )
		82	+ class_name = type(func).__name__
		83	+ # The class's own ``__call__``, not the metaclass's.
		84	+ call = next(
		85	+ (vars(k)["__call__"] for k in type(func).__mro__ if "__call__" in vars(k)), None
		86	+ )
		87	+ if call is None or not is_jit_callable(call):
		CheaterAbec26 days ago 这里对「无 @tla.jit__call__」直接报错，没有任何「方法名 kernel / 类内唯一 @tla.kernel」的回落。但 PR 正文「编译入口的确定规则」仍写了这两条 fallback，仓内文档 compile_and_launch.md L101–L121 与测试 test_instance_without_jit_call_is_rejected（test_kernel_class_organization.py L291）与代码一致，建议保持统一 yuantao 26 days ago Comments： 26 days ago Comments：
		88	+ raise error_type(
		89	+ f"tla.compile() received an instance of {class_name!r}, but that class "
		90	+ "declares no @tla.jit __call__. Add one that names the kernel to compile:\n"
		91	+ " @tla.jit\n"
		92	+ " def __call__(self, ...):\n"
		93	+ " self.<kernel>(...)"
		94	+ )
		95	+ target, forwarded = _check_jit_call_is_pure_forwarding(
		96	+ unwrap_jit_callable(call), class_name=class_name, error_type=error_type
		97	+ )
		98	+ kernel = getattr(type(func), target, None)
		99	+ if not isinstance(kernel, TlaJitFunction) or kernel.kind != "kernel":
		100	+ raise error_type(
		101	+ f"{class_name}.__call__ forwards to {target!r}, which is not one of that "
		102	+ "class's @tla.kernel methods."
		103	+ )
		104	+ params = list(inspect.signature(kernel.fn).parameters.values())
		105	+ if receiver_argument_name(params) is not None:
		106	+ params = params[1:]
		107	+ if len(forwarded) != len(params):
		108	+ raise error_type(
		109	+ f"{class_name}.__call__ forwards {len(forwarded)} parameters to "
		110	+ f"{target!r}, which takes {len(params)} (self excluded)."
		111	+ )
		112	+ return kernel, func
		113	+
		114	+
8	class CompileCallable:	115	class CompileCallable:
9	"""Compile a Tla kernel and return a callable compiled executor."""	116	"""Compile a Tla kernel and return a callable compiled executor."""
10		117
	@@ -17,49 +124,81 @@ class CompileCallable:
17	path when you compile once and launch the same binary repeatedly.	124	path when you compile once and launch the same binary repeatedly.
18		125
19	Parameters:	126	Parameters:
20	- - `func` (`TlaJitFunction`): Decorated `@tla.kernel` function. Required.	127	+ - `func` (`TlaJitFunction \| Any`): Decorated `@tla.kernel` function,
		128	+ or a class instance whose `@tla.jit __call__` names its kernel. Required.
21	- `args` (`Any`): Host tensors / scalars / `@dataclass` instances used	129	- `args` (`Any`): Host tensors / scalars / `@dataclass` instances used
22	as compile type samples (e.g. `from_dlpack` or `make_fake_tensor`	130	as compile type samples (e.g. `from_dlpack` or `make_fake_tensor`
23	- results).	131	+ results); they need not be bound NPU buffers.
24	- `kwargs`: Host compile kwargs. Pass `options="--npu-arch 3510"` to	132	- `kwargs`: Host compile kwargs. Pass `options="--npu-arch 3510"` to
25	select the public chip name; see the option table above for every	133	select the public chip name; see the option table above for every
26	accepted option. Cache / IR dump / force-recompile use	134	accepted option. Cache / IR dump / force-recompile use
27	- `CATLASS_DSL_*` environment variables.	135	+ `CATLASS_DSL_*` environment variables. Launch kwargs such as
		136	+ `block_num` / `stream` belong on the returned compiled function, not
		137	+ here.
28		138
29	Constraints:	139	Constraints:
30	- - `func` must be a `@tla.kernel` `TlaJitFunction`.	140	+ - `func` must be a `@tla.kernel` `TlaJitFunction`, or an instance of a
31	- - `args` are compile-time type samples; they need not be bound NPU	141	+ class that declares one.
32	- buffers (`make_fake_tensor` is valid).	142	+ - A kernel read off an instance (`instance.kernel`) carries that instance
33	- - Pass the public chip name with `options="--npu-arch 3510"`;	143	+ as its receiver.
34	- unsupported tokens raise at compile time. Switch options such as	144	+ - Passing a class instance compiles the kernel its `@tla.jit __call__`
35	- `--cce-disable-asc-reserved-ubuf` take no value.	145	+ forwards to, so the class must declare that `__call__`. The forward must
36	- - Launch kwargs such as `block_num` / `stream` belong on the returned	146	+ be a plain `self.<kernel>(...)` call passing the receiver's own parameters
37	- compiled function, not on `tla.compile`.	147	+ in order; any other `@tla.jit __call__` body is rejected.
		148	+ - A kernel method's first `self` parameter is compile-time configuration.
		149	+ It occupies no runtime ABI slot, and `self.attr` reads fold into the
		150	+ kernel as constants while it is staged.
38		151
39	Example:	152	Example:
40	```python	153	```python
41	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")	154	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")
42	- compiled(tx, ty, block_num=1)	155	+ compiled(tx, ty, block_num=1) # launch; call again to reuse the binary
43	- compiled(tx, ty, block_num=1) # launch again with the same binary	156	+
		157	+ class VectorAddWithStaticLength:
		158	+ def __init__(self, n_ele: int, vl_ele: int = 64):
		159	+ self.n_ele, self.vl_ele = n_ele, vl_ele
		160	+
		161	+ @tla.jit
		162	+ def __call__(self, gm_a, gm_b, gm_c):
		163	+ self.kernel(gm_a, gm_b, gm_c)
		164	+
		165	+ @tla.kernel
		166	+ def kernel(self, gm_a, gm_b, gm_c):
		167	+ ... # gm_c = gm_a + gm_b; the trip count comes from self.n_ele
		168	+
		169	+ compiled = tla.compile(
		170	+ VectorAddWithStaticLength(384), gm_a, gm_b, gm_c,
		171	+ options="--npu-arch 3510",
		172	+ )
		173	+ compiled(gm_a, gm_b, gm_c, block_num=1)
44	```	174	```
45		175
46	"""	176	"""
47	from ..execution import TlaUnsupportedAbiError	177	from ..execution import TlaUnsupportedAbiError
48	- from ..dsl import TlaJitFunction, _bind_kernel_call_args	178	+ from ..dsl import _bind_kernel_call_args
49		179
50	if func is None:	180	if func is None:
51	raise TlaUnsupportedAbiError("Function is not set or invalid.")	181	raise TlaUnsupportedAbiError("Function is not set or invalid.")
52		182
53	- if not isinstance(func, TlaJitFunction):	183	+ kernel, receiver = _resolve_kernel_target(
54	- raise TlaUnsupportedAbiError("tla.compile expects a @tla.kernel function.")	184	+ func, error_type=TlaUnsupportedAbiError
		185	+ )
55	if "type_args" in kwargs:	186	if "type_args" in kwargs:
56	raise TypeError(	187	raise TypeError(
57	"tla.compile() does not accept `type_args`; "	188	"tla.compile() does not accept `type_args`; "
58	"pass compile-time sample arguments positionally."	189	"pass compile-time sample arguments positionally."
59	)	190	)
		191	+ if receiver is not None:
		CheaterAbec26 days ago instance.kernel 已通过__get__带上 receiver 时，这里会无条件插入 type_args 首位。若用户写成 tla.compile(inst.kernel, inst, lhs, ...)，会双重传入 receiver，导致 arity/绑定错误，报错未必指向「勿重复传实例」。建议：文档约束「绑定方法不要再显式传 receiver」；或在 prepend 前检测首参已是同一 receiver 并给出明确诊断 yuantao 26 days ago Comments： 26 days ago Comments：
		192	+ if args and args[0] is receiver:
		193	+ raise TlaUnsupportedAbiError(
		194	+ "tla.compile() got the receiver instance as its first argument "
		195	+ "too; the instance already rides along as the receiver, so drop "
		196	+ "the extra one."
		197	+ )
		198	+ args = (receiver, *args)
60	# Kernel args passed by name arrive mixed in with the compile options.	199	# Kernel args passed by name arrive mixed in with the compile options.
61	- args, kwargs = _bind_kernel_call_args(func.fn, args, kwargs)	200	+ args, kwargs = _bind_kernel_call_args(kernel.fn, args, kwargs)
62	- return func.compile(type_args=args if args else None, **kwargs)	201	+ return kernel.compile(type_args=args if args else None, **kwargs)
63		202
64		203
65	compile = CompileCallable()	204	compile = CompileCallable()

Mpython/tla_dsl/catlass/base_dsl/jit_executor.py+8-14

	@@ -183,9 +183,8 @@ class ExecutionArgs:
183	)	183	)
184		184
185	def filter_runtime_signature(self, sig: inspect.Signature) -> inspect.Signature:	185	def filter_runtime_signature(self, sig: inspect.Signature) -> inspect.Signature:
186	- """Drop Constexpr parameters from a signature."""	186	+ """Drop static parameters (``Constexpr`` and method receivers)."""
187	from .runtime.jit_arg_adapters import is_arg_annotation_constexpr	187	from .runtime.jit_arg_adapters import is_arg_annotation_constexpr
188	- from .typing import is_constexpr_annotation
189		188
190	filtered_params = []	189	filtered_params = []
191	for index, (name, param) in enumerate(sig.parameters.items()):	190	for index, (name, param) in enumerate(sig.parameters.items()):
	@@ -195,14 +194,10 @@ class ExecutionArgs:
195	):	194	):
196	filtered_params.append(param)	195	filtered_params.append(param)
197	continue	196	continue
198	- # Execution lowering classifies explicit kernel parameters by	197	+ # One rule classifies static parameters, so the launch signature
199	- # annotation, after callable/method binding has taken place.	198	+ # agrees with execution lowering on which slots the compiled plan
200	- is_static = (	199	+ # has: a leading ``self`` / ``cls`` receiver is static too.
201	- is_constexpr_annotation(param.annotation)	200	+ if is_arg_annotation_constexpr(param.annotation, name, index, None):
202	- if self.argument_trees is not None
203	- else is_arg_annotation_constexpr(param.annotation, name, index, None)
204	- )
205	- if is_static:
206	continue	201	continue
207	filtered_params.append(param)	202	filtered_params.append(param)
208	return sig.replace(parameters=filtered_params)	203	return sig.replace(parameters=filtered_params)
	@@ -579,8 +574,7 @@ class JitCompiledFunction:
579	"""Total Unified Buffer this kernel allocates statically, in bytes.	574	"""Total Unified Buffer this kernel allocates statically, in bytes.
580		575
581	The launch-sized region is not counted: its extent is chosen per launch,	576	The launch-sized region is not counted: its extent is chosen per launch,
582	- not compiled in. Named after CuTe DSL's ``get_kernel_smem_size``, which	577	+ not compiled in.
583	- answers the same question about shared memory.
584	"""	578	"""
585	return self.jit_module.ub_static_bytes	579	return self.jit_module.ub_static_bytes
586		580
	@@ -636,8 +630,8 @@ class JitCompiledFunction:
636	- `stream` (`Any`, via `**launch_kwargs`): Optional ACL stream handle.	630	- `stream` (`Any`, via `**launch_kwargs`): Optional ACL stream handle.
637	When omitted, uses the current stream for the executor's device.	631	When omitted, uses the current stream for the executor's device.
638	- `ub` (`int`, via `**launch_kwargs`): Bytes of UB to back the	632	- `ub` (`int`, via `**launch_kwargs`): Bytes of UB to back the
639	- kernel's `tla.arch.get_dyn_ub` region with, as CuTe DSL's	633	+ kernel's `tla.arch.get_dyn_ub` region with, i.e. the size of the
640	- `launch(smem=...)` sizes dynamic shared memory. Optional, default	634	+ dynamic shared memory the launch provides. Optional, default
641	`0`. Rejected unless the kernel declares a region, and rejected	635	`0`. Rejected unless the kernel declares a region, and rejected
642	above `artifact.max_ub_bytes`. Ask for what the kernel uses: the	636	above `artifact.max_ub_bytes`. Ask for what the kernel uses: the
643	launch hands over exactly this much, and on SIMT the UB left	637	launch hands over exactly this much, and on SIMT the UB left

Mpython/tla_dsl/catlass/base_dsl/runtime/jit_arg_adapters.py+27-0

	@@ -316,6 +316,32 @@ def is_arg_annotation_constexpr(
316	return is_constexpr_annotation(arg_annotation)	316	return is_constexpr_annotation(arg_annotation)
317		317
318		318
		319	+def receiver_argument_name(parameters: Sequence[Any]) -> str \| None:
		CheaterAbec26 days ago docstring（L322–L326）称首位 self 或 classmethod 的 cls 为 receiver。但此处固定传入 owning_func=None，而 is_arg_annotation_constexpr 对 cls 要求 _owner_marks_classmethod(owning_func)（见同文件 L313）。因此 receiver_argument_name 实际上永远不会返回 "cls"；cls 分支是死路径 yuantao 26 days ago Comments： 26 days ago Comments：
		320	+ """Name of the leading ``self`` receiver this path reserves, else ``None``.
		321	+
		322	+ A leading ``self`` carries compile-time configuration, so it takes no runtime
		323	+ ABI slot. ``is_arg_annotation_constexpr`` also reserves a classmethod's
		324	+ ``cls``, but that needs the owning function, and callers here pass none, so
		325	+ this helper never answers ``cls``. Callers that need to know which
		326	+ parameter that is — to skip it in a gate, or to drop it from a kernel's
		327	+ runtime parameter list — ask here instead of re-deriving the rule.
		328	+ """
		329	+
		330	+ if not parameters:
		331	+ return None
		332	+ first = parameters[0]
		333	+ if first.kind in (
		334	+ inspect.Parameter.VAR_POSITIONAL,
		335	+ inspect.Parameter.VAR_KEYWORD,
		336	+ ):
		337	+ return None
		338	+ if first.name not in ("self", "cls"):
		339	+ return None
		340	+ if not is_arg_annotation_constexpr(first.annotation, first.name, 0, None):
		341	+ return None
		342	+ return first.name
		343	+
		344	+
319	def is_argument_constexpr(	345	def is_argument_constexpr(
320	arg: Any,	346	arg: Any,
321	arg_annotation: Any,	347	arg_annotation: Any,
	@@ -357,5 +383,6 @@ __all__ = [
357	"_PointerLaunchArg",	383	"_PointerLaunchArg",
358	"is_arg_annotation_constexpr",	384	"is_arg_annotation_constexpr",
359	"is_argument_constexpr",	385	"is_argument_constexpr",
		386	+ "receiver_argument_name",
360	"_is_compile_time_callable",	387	"_is_compile_time_callable",
361	]	388	]

Mpython/tla_dsl/catlass/dsl.py+11-1

	@@ -4,7 +4,7 @@ from __future__ import annotations
4		4
5	import contextlib	5	import contextlib
6	import contextvars	6	import contextvars
7	-from dataclasses import dataclass	7	+from dataclasses import dataclass, replace
8	import functools	8	import functools
9	import inspect	9	import inspect
10	from typing import Any, Callable, Mapping, Sequence	10	from typing import Any, Callable, Mapping, Sequence
	@@ -161,6 +161,8 @@ class TlaJitFunction:
161	kind: str	161	kind: str
162	options: Mapping[str, Any]	162	options: Mapping[str, Any]
163	decorator_location: DSLLocation \| None = None	163	decorator_location: DSLLocation \| None = None
		164	+ #: The instance this kernel was read off, if any.
		165	+ receiver: Any \| None = None
164	_mlir: str \| None = None	166	_mlir: str \| None = None
165	_base_dsl: BaseDSL \| None = None	167	_base_dsl: BaseDSL \| None = None
166	_lowered: Any \| None = None	168	_lowered: Any \| None = None
	@@ -188,6 +190,14 @@ class TlaJitFunction:
188	self, "__annotations__", dict(annotations) if annotations else {}	190	self, "__annotations__", dict(annotations) if annotations else {}
189	)	191	)
190		192
		193	+ def __get__(self, obj: Any, objtype: type \| None = None) -> TlaJitFunction:
		194	+ """Bind the kernel to ``obj`` the way a plain method would."""
		195	+
		196	+ if obj is None:
		197	+ return self
		198	+ # The memo fields belong to the unbound kernel.
		199	+ return replace(self, receiver=obj, _mlir=None, _base_dsl=None, _lowered=None)
		200	+
191	def __call__(self, args: Any, *kwargs: Any) -> Any:	201	def __call__(self, args: Any, *kwargs: Any) -> Any:
192	raise TypeError(	202	raise TypeError(
193	"Direct @tla.kernel invocation is disabled; compile explicitly "	203	"Direct @tla.kernel invocation is disabled; compile explicitly "

Mpython/tla_dsl/catlass/execution_lowering.py+27-3

	@@ -29,13 +29,17 @@ from .base_dsl.runtime.argument_binding import (
29	_RuntimeArgumentProxy,	29	_RuntimeArgumentProxy,
30	)	30	)
31	from .base_dsl.utils.tree_utils import tree_unflatten	31	from .base_dsl.utils.tree_utils import tree_unflatten
32	-from .base_dsl.typing import Numeric, is_constexpr_annotation	32	+from .base_dsl.typing import Numeric
33	from .dsl import (	33	from .dsl import (
34	_jit_helper_inline,	34	_jit_helper_inline,
35	is_jit_callable,	35	is_jit_callable,
36	unwrap_jit_callable,	36	unwrap_jit_callable,
37	)	37	)
38	-from .base_dsl.runtime.jit_arg_adapters import _is_compile_time_callable	38	+from .base_dsl.runtime.jit_arg_adapters import (
		39	+ _is_compile_time_callable,
		40	+ is_arg_annotation_constexpr,
		41	+ receiver_argument_name,
		42	+)
39	from .frontend_diagnostics import (	43	from .frontend_diagnostics import (
40	FrontendDiagnosticError,	44	FrontendDiagnosticError,
41	SourceLocation,	45	SourceLocation,
	@@ -190,11 +194,31 @@ def lower_jit_to_tlair_module_by_execution(
190	sig = BaseDSL()._get_signature(fn)	194	sig = BaseDSL()._get_signature(fn)
191	params = list(sig.parameters.values())	195	params = list(sig.parameters.values())
192	arg_names = [p.name for p in params]	196	arg_names = [p.name for p in params]
193	- constexpr_names = {p.name for p in params if is_constexpr_annotation(p.annotation)}	197	+ # One rule decides what is compile-time: ``is_arg_annotation_constexpr``
		198	+ # reserves a leading ``self`` / ``cls`` receiver in addition to
		199	+ # ``Constexpr`` annotations. The receiver carries compile-time configuration,
		200	+ # so ``self.attr`` reads fold as the body is staged and the instance takes no
		201	+ # runtime ABI leaf.
		202	+ constexpr_names = {
		203	+ p.name
		204	+ for index, p in enumerate(params)
		205	+ if is_arg_annotation_constexpr(p.annotation, p.name, index, None)
		206	+ }
		207	+ receiver_name = receiver_argument_name(params)
194	kw_only = inspect.Parameter.KEYWORD_ONLY	208	kw_only = inspect.Parameter.KEYWORD_ONLY
195	keyword_only_names = {p.name for p in params if p.kind is kw_only}	209	keyword_only_names = {p.name for p in params if p.kind is kw_only}
196	call_args = _prepare_call_args(arg_names=arg_names, type_args=type_args)	210	call_args = _prepare_call_args(arg_names=arg_names, type_args=type_args)
197	for name, value in zip(arg_names, call_args, strict=False):	211	for name, value in zip(arg_names, call_args, strict=False):
		212	+ if name == receiver_name:
		Z Z zjw666888Sep 29 用户把真实 NPU tensor 误当 receiver 传给一个不读 self 的 kernel 方法时，编译照常成功、tensor 被静默忽略 yuantao Sep 30 Comments： Sep 30 Comments：
		213	+ # The receiver is a compile-time configuration object, not a kernel
		214	+ # value; its attributes are read while the body is staged.
		215	+ if isinstance(value, (Tensor, Numeric)):
		216	+ raise TlaLoweringError(
		217	+ f"kernel parameter {name!r} is the compile-time receiver, "
		218	+ "but a runtime value was passed for it; rename the parameter "
		219	+ "when it carries data"
		220	+ )
		221	+ continue
198	if name in constexpr_names:	222	if name in constexpr_names:
199	# Constexpr values are compile-time inputs, not runtime ABI leaves.	223	# Constexpr values are compile-time inputs, not runtime ABI leaves.
200	if (	224	if (

Mpython/tla_dsl/docs/en/api/host_api_reference.md+42-19

	@@ -367,7 +367,7 @@ The Unified Buffer is 256 KB, of which the compiler reserves 8 KB (2 KB Ascend C
367		367
368	#### `compile`	368	#### `compile`
369		369
370	-Source: [`catlass.base_dsl.compiler.CompileCallable.__call__`](../../../catlass/base_dsl/compiler.py#L11)	370	+Source: [`catlass.base_dsl.compiler.CompileCallable.__call__`](../../../catlass/base_dsl/compiler.py#L118)
371		371
372	Description:	372	Description:
373		373
	@@ -384,39 +384,62 @@ tla.compile(func: Any, args: Any, *kwargs: Any) -> JitCompiledFunction
384		384
385	Parameters:	385	Parameters:
386		386
387	-- `func` (`TlaJitFunction`): Decorated `@tla.kernel` function. Required.	387	+- `func` (`TlaJitFunction \| Any`): Decorated `@tla.kernel` function,
		388	+ or a class instance whose `@tla.jit __call__` names its kernel. Required.
388	- `args` (`Any`): Host tensors / scalars / `@dataclass` instances used	389	- `args` (`Any`): Host tensors / scalars / `@dataclass` instances used
389	as compile type samples (e.g. `from_dlpack` or `make_fake_tensor`	390	as compile type samples (e.g. `from_dlpack` or `make_fake_tensor`
390	- results).	391	+ results); they need not be bound NPU buffers.
391	- `kwargs`: Host compile kwargs. Pass `options="--npu-arch 3510"` to	392	- `kwargs`: Host compile kwargs. Pass `options="--npu-arch 3510"` to
392	select the public chip name; see the option table above for every	393	select the public chip name; see the option table above for every
393	accepted option. Cache / IR dump / force-recompile use	394	accepted option. Cache / IR dump / force-recompile use
394	- `CATLASS_DSL_*` environment variables.	395	+ `CATLASS_DSL_*` environment variables. Launch kwargs such as
		396	+ `block_num` / `stream` belong on the returned compiled function, not
		397	+ here.
395		398
396	Constraints:	399	Constraints:
397		400
398	-- `func` must be a `@tla.kernel` `TlaJitFunction`.	401	+- `func` must be a `@tla.kernel` `TlaJitFunction`, or an instance of a
399	-- `args` are compile-time type samples; they need not be bound NPU	402	+ class that declares one.
400	- buffers (`make_fake_tensor` is valid).	403	+- A kernel read off an instance (`instance.kernel`) carries that instance
401	-- Pass the public chip name with `options="--npu-arch 3510"`;	404	+ as its receiver.
402	- unsupported tokens raise at compile time. Switch options such as	405	+- Passing a class instance compiles the kernel its `@tla.jit __call__`
403	- `--cce-disable-asc-reserved-ubuf` take no value.	406	+ forwards to, so the class must declare that `__call__`. The forward must
404	-- Launch kwargs such as `block_num` / `stream` belong on the returned	407	+ be a plain `self.<kernel>(...)` call passing the receiver's own parameters
405	- compiled function, not on `tla.compile`.	408	+ in order; any other `@tla.jit __call__` body is rejected.
		409	+- A kernel method's first `self` parameter is compile-time configuration.
		410	+ It occupies no runtime ABI slot, and `self.attr` reads fold into the
		411	+ kernel as constants while it is staged.
406		412
407	Example:	413	Example:
408		414
409	```python	415	```python
410	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")	416	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")
411	-compiled(tx, ty, block_num=1)	417	+compiled(tx, ty, block_num=1) # launch; call again to reuse the binary
412	-compiled(tx, ty, block_num=1) # launch again with the same binary	418	+
		419	+class VectorAddWithStaticLength:
		420	+ def __init__(self, n_ele: int, vl_ele: int = 64):
		421	+ self.n_ele, self.vl_ele = n_ele, vl_ele
		422	+
		423	+ @tla.jit
		424	+ def __call__(self, gm_a, gm_b, gm_c):
		425	+ self.kernel(gm_a, gm_b, gm_c)
		426	+
		427	+ @tla.kernel
		428	+ def kernel(self, gm_a, gm_b, gm_c):
		429	+ ... # gm_c = gm_a + gm_b; the trip count comes from self.n_ele
		430	+
		431	+compiled = tla.compile(
		432	+ VectorAddWithStaticLength(384), gm_a, gm_b, gm_c,
		433	+ options="--npu-arch 3510",
		434	+)
		435	+compiled(gm_a, gm_b, gm_c, block_num=1)
413	```	436	```
414		437
415	---	438	---
416		439
417	#### `TlaJitFunction.compile`	440	#### `TlaJitFunction.compile`
418		441
419	-Source: [`catlass.dsl.TlaJitFunction.compile`](../../../catlass/dsl.py#L198)	442	+Source: [`catlass.dsl.TlaJitFunction.compile`](../../../catlass/dsl.py#L208)
420		443
421	Description:	444	Description:
422		445
	@@ -467,7 +490,7 @@ Run a compiled kernel on the NPU by calling the `JitCompiledFunction` returned b
467		490
468	#### `JitCompiledFunction.__call__`	491	#### `JitCompiledFunction.__call__`
469		492
470	-Source: [`catlass.base_dsl.jit_executor.JitCompiledFunction.__call__`](../../../catlass/base_dsl/jit_executor.py#L614)	493	+Source: [`catlass.base_dsl.jit_executor.JitCompiledFunction.__call__`](../../../catlass/base_dsl/jit_executor.py#L608)
471		494
472	Description:	495	Description:
473		496
	@@ -494,8 +517,8 @@ Parameters:
494	- `stream` (`Any`, via `**launch_kwargs`): Optional ACL stream handle.	517	- `stream` (`Any`, via `**launch_kwargs`): Optional ACL stream handle.
495	When omitted, uses the current stream for the executor's device.	518	When omitted, uses the current stream for the executor's device.
496	- `ub` (`int`, via `**launch_kwargs`): Bytes of UB to back the	519	- `ub` (`int`, via `**launch_kwargs`): Bytes of UB to back the
497	- kernel's `tla.arch.get_dyn_ub` region with, as CuTe DSL's	520	+ kernel's `tla.arch.get_dyn_ub` region with, i.e. the size of the
498	- `launch(smem=...)` sizes dynamic shared memory. Optional, default	521	+ dynamic shared memory the launch provides. Optional, default
499	`0`. Rejected unless the kernel declares a region, and rejected	522	`0`. Rejected unless the kernel declares a region, and rejected
500	above `artifact.max_ub_bytes`. Ask for what the kernel uses: the	523	above `artifact.max_ub_bytes`. Ask for what the kernel uses: the
501	launch hands over exactly this much, and on SIMT the UB left	524	launch hands over exactly this much, and on SIMT the UB left
	@@ -525,7 +548,7 @@ Dump frontend TLA IR without building a device binary or launching. See `TlaJitF
525		548
526	#### `TlaJitFunction.dump_mlir`	549	#### `TlaJitFunction.dump_mlir`
527		550
528	-Source: [`catlass.dsl.TlaJitFunction.dump_mlir`](../../../catlass/dsl.py#L259)	551	+Source: [`catlass.dsl.TlaJitFunction.dump_mlir`](../../../catlass/dsl.py#L269)
529		552
530	Description:	553	Description:
531		554

Mpython/tla_dsl/docs/zh/api/host_api_reference.md+27-14

	@@ -347,7 +347,7 @@ executor。缓存 / 架构 / IR dump 等非函数参数见
347		347
348	#### `compile`	348	#### `compile`
349		349
350	-源码： [`catlass.base_dsl.compiler.CompileCallable.__call__`](../../../catlass/base_dsl/compiler.py#L11)	350	+源码： [`catlass.base_dsl.compiler.CompileCallable.__call__`](../../../catlass/base_dsl/compiler.py#L118)
351		351
352	功能说明：	352	功能说明：
353		353
	@@ -363,27 +363,40 @@ tla.compile(func: Any, args: Any, *kwargs: Any) -> JitCompiledFunction
363		363
364	参数说明：	364	参数说明：
365		365
366	-- `func`（`TlaJitFunction`）：被 `@tla.kernel` 装饰的函数。必填。	366	+- `func`（`TlaJitFunction \| Any`）：被 `@tla.kernel` 装饰的函数，或由 `@tla.jit __call__` 指出其 kernel 的类实例。必填。
367	-- `args`（`Any`）：作为编译类型样本的 Host tensor / 标量 / `@dataclass` 实例	367	+- `args`（`Any`）：作为编译类型样本的 Host tensor / 标量 / `@dataclass` 实例（如 `from_dlpack` 或 `make_fake_tensor` 的返回值）；不必绑定 NPU 缓冲。
368	- （如 `from_dlpack` 或 `make_fake_tensor` 的返回值）。	368	+- `kwargs`：Host 编译参数。用 `options="--npu-arch 3510"` 指定芯片名，全部可用取值见[编译选项](#编译选项)。缓存 / IR dump / 强制重编译由 `CATLASS_DSL_*` 环境变量控制。`block_num` / `stream` 等启动参数写在返回的编译函数上，而不是这里。
369	-- `kwargs`：Host 编译参数。用 `options="--npu-arch 3510"` 指定芯片名，
370	- 全部可用取值见[编译选项](#编译选项)。
371	- 缓存 / IR dump / 强制重编译由 `CATLASS_DSL_*` 环境变量控制。
372		369
373	约束说明：	370	约束说明：
374		371
375	-- `func` 必须是 `@tla.kernel` 得到的 `TlaJitFunction`。	372	+- `func` 必须是 `@tla.kernel` 得到的 `TlaJitFunction`，或由 `@tla.jit __call__` 指出其 kernel 的类实例。
376	-- `args` 只作编译期类型样本，不必绑定 NPU 缓冲（`make_fake_tensor` 合法）。	373	+- 从实例读到的 kernel 方法（`实例.kernel`）自带该实例作为 receiver，该实例排在编译类型样本的第 0 位，落进 kernel 的 `self` 参数。
377	-- 用 `options="--npu-arch 3510"` 指定芯片名；不支持的取值在编译时报错。	374	+- 传入类实例时编译其 `@tla.jit __call__` 转发到的 kernel，因此类必须声明该 `__call__`。转发必须是 `self.<kernel>(...)` 形式的纯调用，并按自身参数顺序传参；其它形式的 `@tla.jit __call__` 函数体不支持。
378	-- `block_num` / `stream` 等启动参数写在返回的编译函数上，而不是	375	+- kernel 方法的第一个 `self` 参数是编译期配置：不占启动 ABI，`self.attr` 在编译时折叠为常量。
379	- `tla.compile`。
380		376
381	调用示例：	377	调用示例：
382		378
383	```python	379	```python
384	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")	380	compiled = tla.compile(vadd, tx, ty, options="--npu-arch 3510")
385	-compiled(tx, ty, block_num=1)	381	+compiled(tx, ty, block_num=1) # 启动；再次调用复用同一份二进制
386	-compiled(tx, ty, block_num=1) # 同一份二进制再次启动	382	+
		383	+class VectorAddWithStaticLength:
		384	+ def __init__(self, n_ele: int, vl_ele: int = 64):
		385	+ self.n_ele, self.vl_ele = n_ele, vl_ele
		386	+
		387	+ @tla.jit
		388	+ def __call__(self, gm_a, gm_b, gm_c):
		389	+ self.kernel(gm_a, gm_b, gm_c)
		390	+
		391	+ @tla.kernel
		392	+ def kernel(self, gm_a, gm_b, gm_c):
		393	+ ... # gm_c = gm_a + gm_b；循环次数由 self.n_ele 在编译期确定
		394	+
		395	+compiled = tla.compile(
		396	+ VectorAddWithStaticLength(384), gm_a, gm_b, gm_c,
		397	+ options="--npu-arch 3510",
		398	+)
		399	+compiled(gm_a, gm_b, gm_c, block_num=1)
387	```	400	```
388		401
389	---	402	---

Mpython/tla_dsl/docs/zh/kernel_development/core_concepts/compile_and_launch.md+78-0

	@@ -48,6 +48,83 @@ shape，应在编译和启动所用的 Host tensor 上正确标记动态维度
48		48
49	---	49	---
50		50
		51	+## 以类实例组织 Kernel
		52	+
		53	+编译期配置（tile size、stage 数、向量长度等）可以放在实例里，与 kernel 和编译入口封装成同一个抽象。kernel 写成实例方法，第一个参数 `self` 不进入启动 ABI，`self.attr` 在编译时折叠为常量：
		54	+
		55	+```python
		56	+from functools import cache
		57	+
		58	+import catlass.tla as tla
		59	+from catlass.tla.runtime import make_fake_tensor
		60	+
		61	+VL_ELE = 64 # f32 的 SIMD 向量长度
		62	+
		63	+
		64	+class VectorAddWithStaticLength:
		65	+ """gm_c = gm_a + gm_b，长度 n_ele 在编译期确定。"""
		66	+
		67	+ def __init__(self, n_ele: int, vl_ele: int = VL_ELE) -> None:
		68	+ self.n_ele = n_ele
		69	+ self.vl_ele = vl_ele
		70	+
		71	+ @tla.jit
		72	+ def __call__(self, gm_a, gm_b, gm_c):
		73	+ self.kernel(gm_a, gm_b, gm_c)
		74	+
		75	+ @tla.kernel
		76	+ def kernel(self, gm_a, gm_b, gm_c):
		77	+ ub_a = tla.make_tensor_like(
		78	+ tla.allocate(self.n_ele, tla.Float32, tla.AddressSpace.ub, 256),
		79	+ gm_a, tla.arch.RowMajor,
		80	+ )
		81	+ ... # 受篇幅限制，省略 kernel 主逻辑
		82	+
		83	+ with tla.vec.func(mode="simd"):
		84	+ for i in tla.range(self.n_ele // self.vl_ele):
		85	+ ... # 循环上界为编译期常量
		86	+
		87	+ @cache
		88	+ @staticmethod
		89	+ def compile(n_ele: int, vl_ele: int = VL_ELE):
		90	+ fake = make_fake_tensor(tla.Float32, (n_ele,), (1,))
		91	+ return tla.compile(
		92	+ VectorAddWithStaticLength(n_ele, vl_ele), fake, fake, fake,
		93	+ options="--npu-arch 3510",
		94	+ )
		95	+
		96	+
		97	+compiled = VectorAddWithStaticLength.compile(n_ele=384)
		98	+compiled(gm_a, gm_b, gm_c, block_num=1)
		99	+```
		100	+
		101	+### 编译入口的写法
		102	+
		103	+\| 写法 \| 入口如何确定 \| 何时使用 \|
		104	+\|------\|--------------\|----------\|
		105	+\| `tla.compile(实例, gm_a, gm_b, gm_c)` \| 取类自己的 `@tla.jit __call__`，编译它一行转发到的 kernel 方法，方法名不限 \| 默认写法 \|
		106	+\| `tla.compile(实例.kernel, gm_a, gm_b, gm_c)` \| 直接点名 kernel 方法；从实例读到它时，该实例随属性读取一并作为 receiver \| 类里有多个 kernel，需要指定其中一个 \|
		107	+
		108	+两种写法编译出的产物相同：实例作为 receiver 排在编译类型样本的第 0 位，落进 kernel 的 `self` 参数，不占启动 ABI。`n_ele` / `vl_ele` 是编译期输入，不同实例得到不同的编译产物；`gm_a`、`gm_b`、`gm_c` 仍是启动时传入的 tensor。
		109	+
		110	+完整 kernel 与可运行版本见 [vector_add_with_static_length.py](../../../../examples/end_to_end/kernel_in_class/vector_add_with_static_length.py)。
		111	+
		112	+约束：
		113	+
		114	+- kernel 的第一个参数是 `self` 时，它只承载编译期配置，不占启动 ABI。判定按位置和名字进行：位置 0 且名为 `self` 即为接收者（`cls` 需要 `classmethod` 宿主），不区分该函数是否定义在类里。
		115	+- `@tla.jit __call__` 必须是 `self.<kernel>(...)` 形式的纯转发，并按自身参数顺序传参；其它形式的 `@tla.jit __call__` 函数体不支持。
		116	+
		117	+入口不合法时的报错：
		118	+
		119	+\| 情况 \| 报错内容 \|
		120	+\|------\|----------\|
		121	+\| 类里没有 `@tla.jit __call__` \| `declares no @tla.jit __call__`，并给出正确写法 \|
		122	+\| `__call__` 不是一行纯转发 \| 指出具体原因：语句数、关键字实参、实参不是自身参数名、参数顺序不符等 \|
		123	+\| 转发目标不是本类的 `@tla.kernel` 方法 \| `not one of that class's @tla.kernel methods` \|
		124	+\| 转发参数个数与 kernel 的运行时参数个数不一致 \| `forwards N parameters to 'kernel', which takes M (self excluded)` \|
		125	+
		126	+---
		127	+
51	## 用户需要认识的对象	128	## 用户需要认识的对象
52		129
53	编译和启动过程中会遇到以下对象：	130	编译和启动过程中会遇到以下对象：
	@@ -176,3 +253,4 @@ executor。
176	- 不同 stream：复用同一个 `compiled`，在启动前切换 stream 上下文。	253	- 不同 stream：复用同一个 `compiled`，在启动前切换 stream 上下文。
177	- 仅需类型样本：用 `make_fake_tensor` 编译，启动时换成真实绑定 tensor。	254	- 仅需类型样本：用 `make_fake_tensor` 编译，启动时换成真实绑定 tensor。
178	- 跨 shape 复用：正确标记动态 layout；不要依赖缓存忽略静态 shape 差异。	255	- 跨 shape 复用：正确标记动态 layout；不要依赖缓存忽略静态 shape 差异。
		256	+- 配置封装：把编译期配置放进实例属性，用类实例作为 `tla.compile` 的第一个参数。

Apython/tla_dsl/examples/end_to_end/kernel_in_class/README.md+9-0

@@ -0,0 +1,9 @@
	1	+# 以类实例组织 kernel
	2	+
	3	+这些示例是 [基础 VADD 端到端示例](../basic_vadd/README.md)的类版本：编译期配置、kernel 与编译入口封装在同一个类里，再交给 `tla.compile`。
	4	+
	5	+\| 示例 \| 内容 \|
	6	+\|---\|---\|
	7	+\| `vector_add_with_static_length.py` \| 向量长度 `n_ele` 作为实例属性，在 kernel 编排时折叠为常量；不同 `n_ele` 各得一份编译产物 \|
	8	+
	9	+写法与约束见[编译与启动](../../../docs/zh/kernel_development/core_concepts/compile_and_launch.md)。

Apython/tla_dsl/examples/end_to_end/kernel_in_class/vector_add_with_static_length.py+147-0

@@ -0,0 +1,147 @@
	1	+# -----------------------------------------------------------------------------------------------------------
	2	+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
	4	+# CANN Open Software License Agreement Version 2.0 (the "License").
	5	+# Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+# See LICENSE in the root of the software repository for the full text of the License.
	9	+# -----------------------------------------------------------------------------------------------------------
	10	+
	11	+"""gm_c = gm_a + gm_b with the length fixed at compile time.
	12	+
	13	+Kernel body, compile-time configuration and the compile entry point live in one
	14	+class. ``n_ele`` is an instance attribute, so it folds to a constant while the
	15	+kernel is staged: each length gets its own binary and its own cache key.
	16	+
	17	+Runnable on Ascend 950 (``--npu-arch 3510``):
	18	+ python examples/end_to_end/kernel_in_class/vector_add_with_static_length.py
	19	+"""
	20	+
	21	+from __future__ import annotations
	22	+
	23	+import sys
	24	+from pathlib import Path
	25	+
	26	+_DSL_EXAMPLE_PATH = str((Path(__file__).resolve().parent / "..").resolve())
	27	+
	28	+if _DSL_EXAMPLE_PATH not in sys.path:
	29	+ sys.path.append(_DSL_EXAMPLE_PATH)
	30	+
	31	+import catlass.tla as tla
	32	+from catlass.tla.runtime import from_dlpack
	33	+
	34	+VL_ELE = 64
	35	+
	36	+
	37	+class VectorAddWithStaticLength:
	38	+ """gm_c = gm_a + gm_b; ``n_ele`` is fixed when the instance is built."""
	39	+
	40	+ def __init__(self, n_ele: int, vl_ele: int = VL_ELE) -> None:
	41	+ if n_ele <= 0 or n_ele % vl_ele:
	42	+ raise ValueError(
	43	+ f"n_ele={n_ele} must be a positive multiple of vl_ele={vl_ele}"
	44	+ )
	45	+ self.n_ele = n_ele
	46	+ self.vl_ele = vl_ele
	47	+
	48	+ @tla.jit
	49	+ def __call__(self, gm_a, gm_b, gm_c):
	50	+ self.kernel(gm_a, gm_b, gm_c)
	51	+
	52	+ @tla.kernel
	53	+ def kernel(self, gm_a, gm_b, gm_c):
	54	+ ub_loaded = tla.flag("ub_loaded", tla.arch.MTE2, tla.arch.VECTOR)
	55	+ vec_done = tla.flag("vec_done", tla.arch.VECTOR, tla.arch.MTE3)
	56	+
	57	+ ub_a = tla.make_tensor_like(
	58	+ tla.allocate(self.n_ele, tla.Float32, tla.AddressSpace.ub, 256),
	59	+ gm_a,
	60	+ tla.arch.RowMajor,
	61	+ )
	62	+ ub_b = tla.make_tensor_like(
	63	+ tla.allocate(self.n_ele, tla.Float32, tla.AddressSpace.ub, 256),
	64	+ gm_b,
	65	+ tla.arch.RowMajor,
	66	+ )
	67	+ ub_c = tla.make_tensor_like(
	68	+ tla.allocate(self.n_ele, tla.Float32, tla.AddressSpace.ub, 256),
	69	+ gm_c,
	70	+ tla.arch.RowMajor,
	71	+ )
	72	+
	73	+ with tla.vector():
	74	+ tla.copy(ub_a, gm_a)
	75	+ tla.copy(ub_b, gm_b)
	76	+ tla.set_flag(ub_loaded)
	77	+ tla.wait_flag(ub_loaded)
	78	+
	79	+ with tla.vec.func(mode="simd"):
	80	+ if tla.const_expr(self.n_ele == self.vl_ele):
	81	+ vl = tla.make_shape(self.vl_ele)
	82	+ at = tla.make_coord(0)
	83	+ tla.tile_view(ub_c, vl, at).store(
	84	+ tla.add(
	85	+ tla.tile_view(ub_a, vl, at).load(),
	86	+ tla.tile_view(ub_b, vl, at).load(),
	87	+ )
	88	+ )
	89	+ else:
	90	+ for i in tla.range(self.n_ele // self.vl_ele):
	91	+ vl = tla.make_shape(self.vl_ele)
	92	+ at = tla.make_coord(i)
	93	+ tla.tile_view(ub_c, vl, at).store(
	94	+ tla.add(
	95	+ tla.tile_view(ub_a, vl, at).load(),
	96	+ tla.tile_view(ub_b, vl, at).load(),
	97	+ )
	98	+ )
	99	+
	100	+ tla.set_flag(vec_done)
	101	+ tla.wait_flag(vec_done)
	102	+
	103	+ tla.copy(gm_c, ub_c)
	104	+ tla.pipe_barrier(tla.pipes.ALL)
	105	+
	106	+
	107	+def main() -> int:
	108	+ import torch
	109	+ import torch_npu
	110	+
	111	+ from common import compare, get_block_num
	112	+
	113	+ torch.npu.set_device(0)
	114	+ torch.npu.manual_seed(0)
	115	+ block_num = get_block_num(-1, 0, kind="vector")
	116	+
	117	+ def bind(buf):
	118	+ return from_dlpack(
	119	+ buf.contiguous(), layout_tag=tla.arch.RowMajor
	120	+ ).mark_compact_shape_dynamic(0)
	121	+
	122	+ ok = True
	123	+ for n_ele in (64, 384):
	124	+ a = torch.rand(n_ele, dtype=torch.float32, device="npu") * 10.0 - 5.0
	125	+ b = torch.rand(n_ele, dtype=torch.float32, device="npu") * 10.0 - 5.0
	126	+ c = torch.full((n_ele,), -7.0, dtype=torch.float32, device="npu")
	127	+ tla_a, tla_b, tla_c = bind(a), bind(b), bind(c)
	128	+
	129	+ compiled = tla.compile(
	130	+ VectorAddWithStaticLength(n_ele),
	131	+ tla_a,
	132	+ tla_b,
	133	+ tla_c,
	134	+ options="--npu-arch 3510",
	135	+ )
	136	+ compiled(tla_a, tla_b, tla_c, block_num=block_num)
	137	+ torch.npu.synchronize()
	138	+
	139	+ passed = compare(c, a + b, rtol=0.0, atol=1e-4)
	140	+ ok = ok and passed
	141	+ print(f"n_ele={n_ele} passed={passed} cache_key={compiled.cache_key}")
	142	+
	143	+ return 0 if ok else 1
	144	+
	145	+
	146	+if __name__ == "__main__":
	147	+ raise SystemExit(main())

Mpython/tla_dsl/examples/end_to_end/structured_arguments/static_tensor_interfaces.py+9-4

	@@ -31,9 +31,14 @@ def copy_with_offset(aux, output: tla.Tensor, get_tensor: tla.Constexpr):
31	output[0] = src[0] + src.element_type(1)	31	output[0] = src[0] + src.element_type(1)
32		32
33		33
		34	+# The first parameter is deliberately not named ``self``: a leading ``self`` is
		35	+# the compile-time receiver reserved by the kernel-class organization rule, so it
		36	+# would leave the runtime ABI.
34	@tla.kernel	37	@tla.kernel
35	-def copy_from_self(self: tla.Tensor, output: tla.Tensor, offset: tla.Constexpr):	38	+def copy_from_tensor(
36	- output[0] = self[0] + self.element_type(offset)	39	+ input_tensor: tla.Tensor, output: tla.Tensor, offset: tla.Constexpr
		40	+):
		41	+ output[0] = input_tensor[0] + input_tensor.element_type(offset)
37		42
38		43
39	def main():	44	def main():
	@@ -48,11 +53,11 @@ def main():
48	src = torch.tensor([3.0], dtype=dtype, device=f"npu:{args.device}")	53	src = torch.tensor([3.0], dtype=dtype, device=f"npu:{args.device}")
49	out = torch.zeros_like(src)	54	out = torch.zeros_like(src)
50	ts, to = [tla.from_dlpack(t, layout_tag=tla.arch.RowMajor) for t in (src, out)]	55	ts, to = [tla.from_dlpack(t, layout_tag=tla.arch.RowMajor) for t in (src, out)]
51	- artifact = tla.compile(copy_from_self, ts, to, 1, options="--npu-arch 3510")	56	+ artifact = tla.compile(copy_from_tensor, ts, to, 1, options="--npu-arch 3510")
52	artifact(ts, to, block_num=1)	57	artifact(ts, to, block_num=1)
53	torch.npu.synchronize()	58	torch.npu.synchronize()
54	assert out.item() == 4.0	59	assert out.item() == 4.0
55	- print(f"PASS {dtype} self: runtime parameter and Constexpr launch signature")	60	+ print(f"PASS {dtype} tensor: runtime parameter and Constexpr launch signature")
56	cases = (	61	cases = (
57	("direct", ts, lambda x: x),	62	("direct", ts, lambda x: x),
58	("tuple", (ts,), lambda x: x[0]),	63	("tuple", (ts,), lambda x: x[0]),

Mpython/tla_dsl/pyproject.toml+1-0

	@@ -94,6 +94,7 @@ exclude = ["catlass/_version.py"]
94	"catlass/tla/__init__.py" = ["F401"]	94	"catlass/tla/__init__.py" = ["F401"]
95	"examples/*/.py" = ["E402", "F401", "F722", "F841"]	95	"examples/*/.py" = ["E402", "F401", "F722", "F841"]
96	"tests/lit/lit.cfg.py" = ["F821"]	96	"tests/lit/lit.cfg.py" = ["F821"]
		97	+"tests/frontend/test_kernel_class_organization.py" = ["F821"]
97	"tests/test_execution_lowering.py" = ["F841"]	98	"tests/test_execution_lowering.py" = ["F841"]
98	"tests/test_generated_ops_surface.py" = ["F841"]	99	"tests/test_generated_ops_surface.py" = ["F841"]
99	"tests/test_frontend_lowering.py" = ["F722", "F821"]	100	"tests/test_frontend_lowering.py" = ["F722", "F821"]

Apython/tla_dsl/tests/frontend/test_kernel_class_organization.py+548-0

@@ -0,0 +1,548 @@
	1	+"""Class-organized kernels: compile-time configuration carried by an instance.
	2	+
	3	+These cases are about how ``tla.compile`` picks the entry and where the instance
	4	+lands; the kernel bodies are arbitrary and carry no computation of their own.
	5	+"""
	6	+
	7	+from __future__ import annotations
	8	+
	9	+from pathlib import Path
	10	+from typing import Any
	11	+
	12	+import pytest
	13	+
	14	+pytest.importorskip("catlass", exc_type=ImportError)
	15	+
	16	+import catlass.tla as tla
	17	+from catlass.base_dsl import BaseDSL
	18	+from catlass.execution import TlaUnsupportedAbiError
	19	+from catlass.execution_lowering import TlaLoweringError
	20	+from catlass.tla.runtime import make_fake_tensor
	21	+
	22	+
	23	+def _op(value: int) -> None:
	24	+ tla.make_shape(value)
	25	+
	26	+
	27	+class _ForwardingKernel:
	28	+ """The issue's shape: a ``@tla.jit __call__`` naming the kernel to compile."""
	29	+
	30	+ def __init__(self, vector_ele: int) -> None:
	31	+ self.vector_ele = vector_ele
	32	+
	33	+ @tla.jit
	34	+ def __call__(self, lhs, rhs, output) -> None:
	35	+ self.kernel(lhs, rhs, output)
	36	+
	37	+ @tla.kernel
	38	+ def kernel(self, lhs, rhs, output) -> None:
	39	+ _op(self.vector_ele)
	40	+
	41	+
	42	+class _NoJitCallKernel:
	43	+ """A kernel method but no ``@tla.jit __call__``: nothing names the entry."""
	44	+
	45	+ @tla.kernel
	46	+ def kernel(self, lhs, rhs) -> None:
	47	+ _op(8)
	48	+
	49	+
	50	+class _UndecoratedCallKernel:
	51	+ """A plain ``__call__`` does not name the entry either."""
	52	+
	53	+ def __call__(self, lhs, rhs) -> None:
	54	+ pass
	55	+
	56	+ @tla.kernel
	57	+ def kernel(self, lhs, rhs) -> None:
	58	+ _op(8)
	59	+
	60	+
	61	+class _RunKernel:
	62	+ """The forward, not the method name, picks the kernel to compile."""
	63	+
	64	+ @tla.jit
	65	+ def __call__(self, lhs, rhs) -> None:
	66	+ self.run(lhs, rhs)
	67	+
	68	+ @tla.kernel
	69	+ def run(self, lhs, rhs) -> None:
	70	+ _op(8)
	71	+
	72	+
	73	+class _NoKernelClass:
	74	+ """An instance whose class declares no kernel at all."""
	75	+
	76	+ @tla.jit
	77	+ def __call__(self, lhs, rhs) -> None:
	78	+ pass
	79	+
	80	+
	81	+class _ReorderedCallKernel:
	82	+ """``__call__`` forwards its parameters out of order."""
	83	+
	84	+ @tla.jit
	85	+ def __call__(self, lhs, rhs, output) -> None:
	86	+ self.kernel(rhs, lhs, output)
	87	+
	88	+ @tla.kernel
	89	+ def kernel(self, lhs, rhs, output) -> None:
	90	+ _op(8)
	91	+
	92	+
	93	+class _ShortCallKernel:
	94	+ """``__call__`` forwards fewer parameters than the kernel takes."""
	95	+
	96	+ @tla.jit
	97	+ def __call__(self, lhs, rhs) -> None:
	98	+ self.kernel(lhs, rhs)
	99	+
	100	+ @tla.kernel

Mpython/tla_dsl/tests/test_execution_args.py+28-12

	@@ -112,10 +112,19 @@ def test_launch_rejects_unbound_tensor_in_each_container(
112	assert checks == [tensor, tensor, tensor]	112	assert checks == [tensor, tensor, tensor]
113		113
114		114
115	-@pytest.mark.parametrize("name", ["value", "self", "cls"])	115	+@pytest.mark.parametrize(
116	-def test_tree_launch_signature_uses_annotations_not_parameter_names(name):	116	+ "name,runtime_names",
117	- from catlass.base_dsl.runtime.argument_tree import build_runtime_argument_tree	117	+ [
118	-	118	+ ("value", ("value",)),
		119	+ # ``is_arg_annotation_constexpr`` reserves a leading ``self`` by position
		120	+ # and name, so it leaves the runtime signature like a ``Constexpr`` does.
		121	+ ("self", ()),
		122	+ # ``cls`` is reserved only when the owner is a classmethod, and this
		123	+ # synthetic signature has no owner, so it stays a runtime parameter.
		124	+ ("cls", ("cls",)),
		125	+ ],
		126	+)
		127	+def test_tree_launch_signature_follows_ported_constexpr_rule(name, runtime_names):
119	sig = inspect.Signature(	128	sig = inspect.Signature(
120	[	129	[
121	inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD),	130	inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD),
	@@ -126,21 +135,28 @@ def test_tree_launch_signature_uses_annotations_not_parameter_names(name):
126	),	135	),
127	]	136	]
128	)	137	)
		138	+ runtime_args = (1,) if runtime_names else ()
129	binder = ExecutionArgs(	139	binder = ExecutionArgs(
130	- original_signature=sig, argument_trees=(build_runtime_argument_tree(1),)	140	+ original_signature=sig,
		141	+ argument_trees=tuple(
		142	+ build_runtime_argument_tree(1) for _ in runtime_names
		143	+ ),
131	)	144	)
132	- assert tuple(binder.signature.parameters) == (name,)	145	+ assert tuple(binder.signature.parameters) == runtime_names
133	- assert binder.get_rectified_args_from_original_args((1, 128)) == (1,)	146	+ assert binder.get_rectified_args_from_original_args((1, 128)) == runtime_args
134	- assert binder.get_rectified_args_from_original_args((1,)) == (1,)	147	+ assert binder.get_rectified_args_from_original_args((1,)) == runtime_args
135		148
136		149
137	def test_lowered_self_parameter_and_launch_signature_agree():	150	def test_lowered_self_parameter_and_launch_signature_agree():
138	from catlass.base_dsl import BaseDSL	151	from catlass.base_dsl import BaseDSL
139		152
140	- def kernel(self, tile: tla.Constexpr):	153	+ def kernel(self, value, tile: tla.Constexpr):
141	- _ = self + tile	154	+ _ = self + value + tile
142		155
143	- lowered = BaseDSL()._lower(kernel, kind="kernel", options={}, type_args=(1, 128))	156	+ lowered = BaseDSL()._lower(
		157	+ kernel, kind="kernel", options={}, type_args=(1, 3, 128)
		158	+ )
		159	+ # ``self`` and ``tile`` are compile-time; only ``value`` reaches the ABI.
144	assert len(lowered.argument_trees) == 1	160	assert len(lowered.argument_trees) == 1
145	binder = ExecutionArgs(	161	binder = ExecutionArgs(
146	original_signature=inspect.signature(kernel),	162	original_signature=inspect.signature(kernel),
	@@ -149,7 +165,7 @@ def test_lowered_self_parameter_and_launch_signature_agree():
149	lowered.module	165	lowered.module
150	).kernel_abi,	166	).kernel_abi,
151	)	167	)
152	- assert tuple(binder.signature.parameters) == ("self",)	168	+ assert tuple(binder.signature.parameters) == ("value",)
153	assert binder.generate_launch_payload((3,))[:4] == struct.pack("i", 3)	169	assert binder.generate_launch_payload((3,))[:4] == struct.pack("i", 3)
154		170
155		171