feat(language): integrate hand-written C++ kernels via external_source

docs/en/dev/language/01-external-kernels.md

@@ -0,0 +1,115 @@
+# Integrating Hand-Written C++ Kernels
+
+PyPTO can call an **existing hand-written C++ InCore kernel** from a
+PyPTO-authored orchestration, without going through PyPTO's tile codegen. You
+declare a *signature-only function header* in the DSL; the orchestration calls
+it like any InCore kernel, but the compiler skips codegen for its body and
+compiles the referenced `.cpp` as the kernel instead.
+
+Use this when you already have a tuned AICore kernel (e.g. a bespoke attention
+kernel) and want to drive it from a PyPTO orchestration, reusing PyPTO's
+task scheduling, dependency analysis, and runtime dispatch.
+
+## Contract
+
+An external kernel is a normal InCore kernel from the runtime's point of view,
+so the hand-written source must meet the same ABI PyPTO-generated kernels do:
+
+- Export a single `extern "C" void kernel_entry(__gm__ int64_t* args)` entry
+  (one kernel per `.cpp`; the runtime dispatches by `func_id`, the symbol is
+  fixed). This is the same entry PyPTO's own generated kernels export.
+- The declared **parameter order and directions** (`pl.Out` / `pl.InOut`) must
+  match how the kernel reads its arguments — the orchestration builds the task
+  payload (`add_input` / `add_inout` / `add_output`) from the declaration.
+
+The declaration carries only the signature; its body is a bare `...`.
+
+## `@pl.program` route
+
+Declare the kernel with `@pl.function(type=AIC/AIV, external_source=...)` and an
+empty body. A mixed **AIC + AIV** kernel (one `MixedKernels` submit) is a
+`pl.FunctionType.Group` of one AIC member and one AIV member:
+
+```python
+from pathlib import Path
+import pypto.language as pl
+
+KDIR = Path(__file__).parent / "kernels"
+
+@pl.program
+class PagedAttention:
+    @pl.function(type=pl.FunctionType.AIC, external_source=KDIR / "aic/pa.cpp")
+    def PA_AIC(self, query: pl.Tensor[[B, H, D], pl.FP16], ...,
+               out: pl.Out[pl.Tensor[[B, H, D], pl.FP16]], ...
+               ) -> pl.Tensor[[B, H, D], pl.FP16]:
+        ...                                   # body lives in aic/pa.cpp
+
+    @pl.function(type=pl.FunctionType.AIV, external_source=KDIR / "aic/pa.cpp")
+    def PA_AIV(self, ...same signature...) -> ...:
+        ...
+
+    @pl.function(type=pl.FunctionType.Group)
+    def PA(self, ...same signature...) -> ...:
+        r = self.PA_AIC(...)                  # defines the group members
+        self.PA_AIV(...)
+        return r
+
+    @pl.function(type=pl.FunctionType.Orchestration)
+    def entry(self, query, ..., out):
+        # build tiling / workspace tensors here (host-side), then dispatch:
+        out = self.PA(query, ..., out)        # -> MixedKernels submit
+        return out
+```
+
+`external_source` accepts an absolute path or one relative to the file that
+defines the program. The AIC and AIV members may point at the **same** source
+(compiled once per core) or different files.
+
+For a single-core kernel, declare just one `AIC` or `AIV` function and call it
+directly from the orchestration (no group).
+
+## `@pl.jit.extern` route
+
+Under `@pl.jit`, declare the kernel with `@pl.jit.extern`. A `core_type="mixed"`
+kernel auto-expands to the AIC + AIV + Group form above:
+
+```python
+@pl.jit.extern(core_type="mixed",
+               aic_source="kernels/aic/pa.cpp",
+               aiv_source="kernels/aic/pa.cpp")
+def pa(query: pl.Tensor[[B, H, D], pl.FP16], ...,
+       out: pl.Out[pl.Tensor[[B, H, D], pl.FP16]], ...
+       ) -> pl.Tensor[[B, H, D], pl.FP16]: ...
+
+@pl.jit
+def decode(query: pl.Tensor, ..., out: pl.Out[pl.Tensor]):
+    out = pa(query, ..., out)                 # dep discovered automatically
+    return out
+```
+
+Single-core form: `@pl.jit.extern(core_type="aic"|"aiv", source="k.cpp")`.
+
+Paths resolve relative to the file defining the kernel. Editing the referenced
+`.cpp` changes the JIT cache key, so a kernel change triggers recompilation even
+though the Python stub is unchanged.
+
+## What the compiler does
+
+- **Passes**: the header-only function survives the pass pipeline unchanged
+  (its empty body is a no-op for tile passes); it is exempt only from the
+  `ReturnParamsExplicit` property, which requires a `ReturnStmt` a header cannot
+  have.
+- **Orchestration codegen**: assigns the kernel a `func_id` and emits the submit
+  exactly as for a DSL kernel — a single AIC/AIV `rt_submit_*_task`, or a
+  `MixedKernels{aic_id, aiv_id, ...}` + `rt_submit_task` for a group.
+- **Backend**: skips ptoas for the external kernel and copies its `.cpp` to
+  `kernels/<aic|aiv>/<name>.cpp`, so the generated `kernel_config.py` manifest
+  lists it exactly like a generated kernel (`func_id`, `source`, `core_type`,
+  `signature`, `arg_index`).
+
+## Restrictions
+
+- `external_source` is only valid on `FunctionType.AIC` / `FunctionType.AIV`.
+- The body must be a bare `...` (signature only).
+- A `Group` must be all-external or all-DSL — mixing external and DSL members in
+  one group is rejected.

docs/zh-cn/dev/language/01-external-kernels.md

@@ -0,0 +1,108 @@
+# 集成手写 C++ Kernel
+
+PyPTO 可以从 PyPTO 编写的 orchestration（编排）中直接调用一个**已有的手写 C++
+InCore kernel（核内核函数）**，而不经过 PyPTO 的 tile 代码生成。你在 DSL 里声明一个
+*仅有签名的函数头*；orchestration 像调用普通 InCore kernel 那样调用它，但编译器会
+跳过其函数体的代码生成，转而把所引用的 `.cpp` 编译为该 kernel。
+
+适用场景：你已有一个调优过的 AICore kernel（例如定制的 attention kernel），希望用
+PyPTO orchestration 驱动它，复用 PyPTO 的任务调度、依赖分析与运行时派发。
+
+## 契约（Contract）
+
+从运行时的角度看，外部 kernel 就是一个普通 InCore kernel，因此手写源码必须满足与
+PyPTO 生成 kernel 相同的 ABI：
+
+- 导出唯一的 `extern "C" void kernel_entry(__gm__ int64_t* args)` 入口（一个 `.cpp`
+  一个 kernel；运行时按 `func_id` 派发，符号名固定）。这与 PyPTO 自身生成 kernel 导出
+  的入口完全一致。
+- 声明的**参数顺序与方向**（`pl.Out` / `pl.InOut`）必须与 kernel 读取参数的方式一致
+  —— orchestration 会依据声明构建任务负载（`add_input` / `add_inout` /
+  `add_output`）。
+
+声明只携带签名，函数体为一个裸的 `...`。
+
+## `@pl.program` 写法
+
+用 `@pl.function(type=AIC/AIV, external_source=...)` + 空体声明 kernel。混合的
+**AIC + AIV** kernel（一次 `MixedKernels` 派发）表示为一个 `pl.FunctionType.Group`，
+其成员为一个 AIC 与一个 AIV：
+
+```python
+from pathlib import Path
+import pypto.language as pl
+
+KDIR = Path(__file__).parent / "kernels"
+
+@pl.program
+class PagedAttention:
+    @pl.function(type=pl.FunctionType.AIC, external_source=KDIR / "aic/pa.cpp")
+    def PA_AIC(self, query: pl.Tensor[[B, H, D], pl.FP16], ...,
+               out: pl.Out[pl.Tensor[[B, H, D], pl.FP16]], ...
+               ) -> pl.Tensor[[B, H, D], pl.FP16]:
+        ...                                   # 函数体在 aic/pa.cpp 中
+
+    @pl.function(type=pl.FunctionType.AIV, external_source=KDIR / "aic/pa.cpp")
+    def PA_AIV(self, ...same signature...) -> ...:
+        ...
+
+    @pl.function(type=pl.FunctionType.Group)
+    def PA(self, ...same signature...) -> ...:
+        r = self.PA_AIC(...)                  # 定义 group 成员
+        self.PA_AIV(...)
+        return r
+
+    @pl.function(type=pl.FunctionType.Orchestration)
+    def entry(self, query, ..., out):
+        # 在此构建 tiling / workspace 张量（host 侧），然后派发：
+        out = self.PA(query, ..., out)        # -> MixedKernels 派发
+        return out
+```
+
+`external_source` 接受绝对路径，或相对于定义该 program 的文件的相对路径。AIC 与 AIV
+成员可指向**同一**源文件（按核各编译一次），也可指向不同文件。
+
+单核 kernel：只声明一个 `AIC` 或 `AIV` 函数，并在 orchestration 中直接调用（无需
+group）。
+
+## `@pl.jit.extern` 写法
+
+在 `@pl.jit` 下，用 `@pl.jit.extern` 声明 kernel。`core_type="mixed"` 会自动展开为上面
+的 AIC + AIV + Group 形式：
+
+```python
+@pl.jit.extern(core_type="mixed",
+               aic_source="kernels/aic/pa.cpp",
+               aiv_source="kernels/aic/pa.cpp")
+def pa(query: pl.Tensor[[B, H, D], pl.FP16], ...,
+       out: pl.Out[pl.Tensor[[B, H, D], pl.FP16]], ...
+       ) -> pl.Tensor[[B, H, D], pl.FP16]: ...
+
+@pl.jit
+def decode(query: pl.Tensor, ..., out: pl.Out[pl.Tensor]):
+    out = pa(query, ..., out)                 # 依赖自动发现
+    return out
+```
+
+单核写法：`@pl.jit.extern(core_type="aic"|"aiv", source="k.cpp")`。
+
+路径相对于定义该 kernel 的文件解析。修改所引用的 `.cpp` 会改变 JIT 缓存键，因此即使
+Python stub 未变，kernel 变更也会触发重新编译。
+
+## 编译器行为
+
+- **Passes**：仅有函数头的函数原样穿过 pass 流水线（空体对 tile passes 是 no-op）；它
+  仅豁免 `ReturnParamsExplicit` 属性 —— 该属性要求存在 `ReturnStmt`，而函数头本就没有。
+- **Orchestration 代码生成**：为该 kernel 分配 `func_id`，并像 DSL kernel 一样生成派发
+  —— 单核 AIC/AIV 的 `rt_submit_*_task`，或 group 的 `MixedKernels{aic_id, aiv_id, ...}`
+  + `rt_submit_task`。
+- **Backend**：对外部 kernel 跳过 ptoas，并把其 `.cpp` 复制到
+  `kernels/<aic|aiv>/<name>.cpp`，因此生成的 `kernel_config.py` manifest 会像列出生成
+  kernel 一样列出它（`func_id`、`source`、`core_type`、`signature`、`arg_index`）。
+
+## 约束
+
+- `external_source` 仅在 `FunctionType.AIC` / `FunctionType.AIV` 上有效。
+- 函数体必须是裸的 `...`（仅签名）。
+- 一个 `Group` 必须全为外部 kernel 或全为 DSL kernel —— 不允许在同一 group 中混用外部
+  与 DSL 成员。

python/pypto/backend/pto_backend.py

@@ -965,6 +965,36 @@ def _get_kernel_output_path(
     return os.path.join("kernels", ct_str, f"{func.name}.{suffix}")
 
 
+def _external_source_of(func: _ir_core.Function) -> str | None:
+    """Return the ``external_source`` path of a header-only external kernel, else None.
+
+    External kernels (declared via ``@pl.function(type=AIC/AIV,
+    external_source=...)``) carry the path to a hand-written C++ ``.cpp`` in
+    their ``external_source`` attr and have an empty ``...`` DSL body. The
+    backend copies that source into ``kernels/<ct>/<name>.cpp`` instead of
+    running PyPTO codegen + ptoas.
+    """
+    return dict(func.attrs).get("external_source")
+
+
+def _emit_external_kernel(func: _ir_core.Function, result_files: dict[str, str]) -> None:
+    """Copy a hand-written external kernel source into the per-core kernel path.
+
+    The orchestration codegen still assigns the function a ``func_id`` and emits
+    its submit; only the kernel body is taken from the referenced ``.cpp``
+    instead of being generated. The manifest (``kernel_config.py``) picks up the
+    file at the standard ``kernels/<ct>/<name>.cpp`` location, so no manifest
+    change is needed.
+    """
+    source_path = _external_source_of(func)
+    if source_path is None:
+        raise RuntimeError(f"Internal error: '{func.name}' has no external_source attr")
+    with open(source_path) as f:
+        source = f.read()
+    out_path = _get_kernel_output_path(func, "cpp")
+    result_files[out_path] = source
+
+
 def _compile_pto_module(
     pto_code: str,
     unit_name: str,
@@ -1524,6 +1554,21 @@ def _generate_single_chip(
     # Grouped functions: one MLIR module per group
     for group_name, members in groups.items():
         try:
+            # External kernels: copy the hand-written .cpp per member and skip
+            # PyPTO codegen. A group must be all-external or all-DSL — mixing is
+            # rejected (the DSL members would need cross-core protocol wiring the
+            # external source can't participate in).
+            ext_members = [m for m in members if _external_source_of(m) is not None]
+            if ext_members:
+                if len(ext_members) != len(members):
+                    dsl = ", ".join(m.name for m in members if _external_source_of(m) is None)
+                    raise RuntimeError(
+                        f"Group '{group_name}' mixes external and DSL kernels "
+                        f"(DSL members: {dsl}). A group must be all-external or all-DSL."
+                    )
+                for member in members:
+                    _emit_external_kernel(member, result_files)
+                continue
             grouped_program = _ir_core.Program(members, group_name, transformed_program.span)
             stage = StageRecord(name=f"kernel_codegen:{group_name}", start=time.perf_counter())
             ir_record = StageRecord(name="ir_to_mlir", start=time.perf_counter())
@@ -1538,6 +1583,10 @@ def _generate_single_chip(
 
     for func in ungrouped:
         try:
+            # External kernel: copy the hand-written .cpp and skip PyPTO codegen.
+            if _external_source_of(func) is not None:
+                _emit_external_kernel(func, result_files)
+                continue
             peer_names = _extract_peer_function_names(func)
             peer_funcs: list[_ir_core.Function] = []
             for name in peer_names: