Add SSA compiler pipeline and unified backends

Author: whjthu

.gitignore

@@ -160,3 +160,12 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Local/generated compiler validation artifacts.
+reports/
+outputs/
+tmp_*/
+remote_*/
+remote_scripts/
+compiler-skills/
+ninetoothed.skill/

docs/design/README.md

@@ -0,0 +1,27 @@
+# NineToothed Compiler Design Notes
+
+This directory contains the design documents that should accompany the compiler implementation changes. Keep generated benchmark logs, large backend artifacts, and exploratory reports out of the main repository history.
+
+## Recommended Reading Order
+
+1. `ir_design_zh.md`
+   - Overall IR motivation, expressiveness, extensibility, and current scope.
+2. `ssa_ir_lowering_pipeline_zh.md`
+   - Source-to-SSA lowering layers and target-annotated SSA flow.
+3. `ssa_pass_registry_pipeline_zh.md`
+   - Pass registry, default pipeline, backend-specific optimization pass organization.
+4. `multi_backend_ir.md`
+   - Earlier multi-backend design context and compatibility notes.
+
+## Repository Boundary
+
+The main `ninetoothed` repository should keep:
+
+- compiler source code under `src/`
+- tests under `tests/`
+- reusable validation scripts under `scripts/`
+- concise design documents under `docs/design/`
+
+Generated reports and backend artifacts should remain local validation output unless a specific report is intentionally selected for a release note or PR discussion.
+
+AI/Codex workflow skills are maintained separately in the `ninetoothed.skill` repository.

docs/design/ir_design_zh.md

@@ -0,0 +1,77 @@
+# NineToothed Multi-Backend IR Design
+
+## Objectives
+
+NineToothed currently exposes a tensor-oriented metaprogramming DSL and lowers it to Triton. The target backend programming languages are now:
+
+- Triton: keep the current production path and use it as the correctness baseline.
+- CUDA: add native CUDA source generation while preserving the current AOT CUDA launcher path.
+- TileLang: add a TileLang path for tile-level scheduling experiments.
+- TVM: add a TensorIR/TVMScript path for a widely used tensor compiler stack.
+
+SGLang is intentionally excluded from this backend list because it is an inference engine, not a kernel programming language target.
+
+## IR Stack
+
+The design should remain multi-level instead of forcing all targets through one flat IR:
+
+- TOM Graph IR: tensor arrangement history, shapes, strides, jagged metadata, dtypes, constexpr values, and symbolic constraints.
+- Program IR: loads, stores, pointer/index expressions, masks, arithmetic, dot, reductions, math intrinsics, helper calls, and function boundaries.
+- Schedule IR: grid/block decomposition, thread/warp mapping, memory scopes, vectorization, unrolling, pipelining, and target capabilities.
+- Target IR: backend-specific syntax, ABI metadata, runtime dispatch metadata, and compiler flags.
+
+## First Implementation Slice
+
+The current implementation slice adds a compact `KernelIR` bridge:
+
+- `TensorTypeIR` records public tensor parameter metadata.
+- `LaunchIR` records launch function name, launch args, and grid text.
+- `KernelIR` records canonical source, entrypoint, launch metadata, compiler options, and extensible metadata.
+- Backend lowerers consume `KernelIR` and return `BackendArtifact`.
+
+This bridge does not replace the Triton generator yet. It creates a stable backend contract while the production Triton path remains untouched.
+
+## Backend Status
+
+| Backend | Current Status | Execution Status |
+| --- | --- | --- |
+| Triton | Existing production path plus artifact lowerer | Executable |
+| CUDA | Emits `.cu` manifest and ABI shell; existing `caller="cuda"` AOT remains the only executable CUDA path | Native source shell is not executable |
+| TileLang | Emits TileLang Python module shell and metadata | Not executable |
+| TVM | Emits TVMScript module shell and metadata | Not executable |
+
+## API Direction
+
+```python
+artifact = ninetoothed.lower(arrangement, application, tensors, backend="triton")
+cuda_artifact = ninetoothed.lower(arrangement, application, tensors, backend="cuda")
+tilelang_artifact = ninetoothed.lower(arrangement, application, tensors, backend="tilelang")
+tvm_artifact = ninetoothed.lower(arrangement, application, tensors, backend="tvm")
+```
+
+Existing `make`, `jit`, and `aot` calls should keep working. New backend selection should remain explicit until native CUDA, TileLang, and TVM lowering are feature complete enough to pass the full correctness gates.
+
+## Lowering Roadmap
+
+1. Keep the bridge IR and backend registry in place.
+2. Extract TOM Graph IR from `Tensor` arrangement history instead of reflecting only generated source.
+3. Split semantic lowering from Triton rendering.
+4. Add Program IR statement and expression nodes for loads, stores, masks, arithmetic, dot, reductions, and intrinsics.
+5. Re-render Triton from Program IR and prove parity with the existing generator.
+6. Implement native CUDA rendering for elementwise, reductions, and matmul-style tiled dot.
+7. Implement TileLang rendering for elementwise, reductions, and matmul-style tiled dot.
+8. Implement TVM TensorIR rendering for elementwise, reductions, and matmul-style tiled dot.
+9. Add target capability descriptions for AscendC, BangC, and similar vendor languages.
+
+## Validation Plan
+
+The acceptance gate is intentionally strict:
+
+- All existing NineToothed tests pass.
+- All `ntops` tests pass.
+- All runnable `ntops.lab` operators pass.
+- Backend artifact audit passes for Triton, CUDA, TileLang, and TVM.
+- Executable backend artifacts compile and run.
+- Outputs match existing PyTorch/reference checks.
+
+Until CUDA, TileLang, and TVM emit real executable target code, backend compatibility must be reported as incomplete.

docs/design/multi_backend_ir.md

@@ -0,0 +1,110 @@
+# 九齿 SSA IR Lowering Pipeline 设计
+
+## 目标
+
+九齿的核心 IR 不应把一个完整算子表示成一个粗粒度节点。前端看到的是 Python/NineToothed 计算过程，IR 应抽象其中的控制流、数据流、访存、规约、矩阵原语和数学函数，然后通过 pass 逐步加入 schedule、memory scope 和后端 intrinsic 信息，最后生成 Triton、CUDA、TileLang、TVM 等后端语言。
+
+因此 canonical IR 的基本单元是：
+
+- `arith.*`：标量/张量表达式。
+- `math.*`：`exp/log/sqrt/rsqrt` 等数学函数。
+- `reduce.*`：`sum/max/min` 等规约。
+- `linalg.*`：`dot/matmul/transpose` 等细粒度线性代数原语。
+- `scf.*`：循环、条件、yield 等控制流 region。
+- `mem.*`：load/store、memory effect。
+- `index.*`、`shape.*`：形状和索引计算。
+
+`AttentionOpIR`、`FlashAttentionOpIR` 这类节点只能作为历史兼容层或 pattern 名称出现，不能成为 canonical SSA 的表达粒度。以 `tests/test_attention.py` 中的参考实现为例，canonical SSA 是 `scf.for + linalg.dot + reduce.max/reduce.sum + math.exp2 + select.where + mem.store`，不是一个 attention 节点。
+
+## IR 分层
+
+1. Frontend AST/Tensor Semantics
+
+   输入是 NineToothed 的 arrangement/application 或已有结构化 ProgramIR。此阶段只负责理解用户代码与张量元信息。
+
+2. Generic SSA
+
+   语义层 IR，保持后端无关。它描述计算发生了什么，不描述 CUDA block、Triton program、TileLang fragment 或 TVM thread binding。
+
+3. Canonical SSA
+
+   通过 `ssa.canonicalize` 规范化 opcode、metadata 和 region 结构，为后续分析 pass 提供稳定输入。
+
+4. Analyzed SSA
+
+   `ssa.analyze_effects` 统计 store、reduction、loop、dot、online-softmax-shape 等事实。这里仍不做后端决策。
+
+5. Scheduled SSA
+
+   `ssa.select_schedule` 根据后端和分析事实加入 schedule annotation，例如 elementwise grid、parallel reduction、blocked linalg、online-softmax-dot。此阶段开始出现 tile size、vector width、threads、num_warps/num_stages 等策略，但仍保持为 IR metadata。
+
+6. Target-Annotated SSA
+
+   `ssa.lower_memory_scopes` 和 `ssa.lower_backend_intrinsics` 将抽象 schedule 映射到目标后端：
+
+   - Triton：`tl.program_id`、`tl.load/tl.store`、`tl.dot`、Triton tensor/register 表达。
+   - CUDA：`blockIdx/threadIdx`、thread-local register、`__shared__`、`__expf/expf`、后续 `mma.sync` 候选。
+   - TileLang：`T.Kernel`、`T.get_thread_binding`、`local.fragment/shared/global`、`T.gemm/T.dot` 候选。
+   - TVM：`T.thread_binding`、`T.alloc_buffer(scope="local/shared")`、tensorize intrinsic 候选。
+
+7. Backend Source
+
+   后端 emitter 读取 target-annotated SSA 和历史 ProgramIR 信息，生成目标语言源码。当前 CUDA/TileLang/TVM 的结构化算子 emitter 已能接收带 pass trace 的 SSA；Triton 普通算子仍可复用历史 Triton source path，但同样生成 SSA 审计信息。
+
+## 可表达性
+
+SSA 的表达能力来自三个方面：
+
+- 数据流：每个 operation 显式列出 operands/results，支持 CSE、DCE、fusion、layout rewrite。
+- 控制流：`scf.for/scf.if` 通过 region 表达循环和条件，loop-carried state 通过 block args 和 `scf.yield` 表示，能覆盖 online softmax、rowwise normalize、scan-like 更新。
+- memory effect：`mem.store` 与 future `mem.load/alloc` 明确建模 side effect，便于在 schedule pass 中决定 coalescing、shared memory staging、local scratch placement。
+
+这足以覆盖当前已验证的：
+
+- elementwise/fill/copy。
+- 1D reduction、axis reduction、rowwise softmax/layernorm 类融合。
+- transpose、matmul。
+- `tests/test_attention.py` 形式的 flash attention reference lowering。
+
+## 可扩展性
+
+新增后端时不需要改前端 IR，只需增加：
+
+1. 目标 backend enum/registry。
+2. backend-specific schedule policy。
+3. memory scope mapping。
+4. intrinsic mapping。
+5. source emitter。
+
+新增优化 pass 时应满足：
+
+- 输入输出仍是 SSAProgramIR。
+- pass 只处理细粒度 opcode，不按完整算子名改写。
+- 若识别 pattern，也只能作为 schedule 选择依据，而不是替换成粗粒度语义节点。
+
+## 当前实现状态
+
+已新增 `ninetoothed.ssa_passes`：
+
+```text
+ssa.canonicalize
+ssa.analyze_effects
+ssa.select_schedule
+ssa.lower_memory_scopes
+ssa.lower_backend_intrinsics
+```
+
+公共 `ninetoothed.lowering.lower` 和直接 `ninetoothed.backends.lower(KernelIR, backend=...)` 都会运行该 pipeline。`BackendArtifact.metadata` 会带出：
+
+- `ssa_pass_trace`
+- `ssa_schedule`
+- `ssa_metadata`
+- `kernel_metadata`
+
+## 后续优化方向
+
+- Triton 普通算子 emitter 从历史 source passthrough 逐步迁移到 SSA source emission。
+- CUDA matmul/attention 引入 shared memory tile、warp-level reduction、`mma.sync` tensorization pass。
+- TileLang 引入 `T.gemm`/block fragment schedule，而不是 correctness-first scalar loops。
+- TVM 引入 schedule rule/tensorize rule，生成标准 host wrapper 的环境兼容路径。
+- 为 pass 增加 cost model/autotune metadata，并把测得性能反馈回 schedule selection。

docs/design/ssa_ir_lowering_pipeline_zh.md

@@ -0,0 +1,169 @@
+# 九齿 SSA Pass 注册器与优化遍执行流程
+
+本文档描述当前九齿 SSA lowering 中的标准编译器式 pass 执行逻辑。目标是把优化遍从固定函数调用整理为可注册、可分类、可配置、可自动选择的 pipeline，为后续更多优化和自动调优留出接口。
+
+## 设计目标
+
+1. Pass 粒度保持在 SSA 计算、访存、控制和 schedule annotation 层，不引入 `AttentionOpIR` 之类粗粒度算子节点。
+2. 优化遍按硬件无关和硬件相关分层执行。
+3. 硬件相关优化可以按后端拆分，例如 `ssa.triton.optimize_schedule`、`ssa.cuda.optimize_schedule`。
+4. 默认 pipeline 可以覆盖当前 Triton/CUDA/TileLang/TVM 后端。
+5. 用户或 autotune 系统可以通过 declarative spec 替换 pass 序列和 pass 参数。
+
+## Pass 分类
+
+当前 pass registry 中的分类如下：
+
+| 分类 | 作用 | 当前 pass |
+| --- | --- | --- |
+| `hardware_independent` | 不依赖后端，负责规范化和语义分析 | `ssa.canonicalize`、`ssa.analyze_effects` |
+| `hardware_dependent` | 依赖后端能力，但多个后端共享同一 pass 框架 | `ssa.select_schedule`、`ssa.lower_memory_scopes`、`ssa.lower_backend_intrinsics` |
+| `backend_specific` | 后端专属优化策略 | `ssa.triton.optimize_schedule`、`ssa.cuda.optimize_schedule`、`ssa.tilelang.optimize_schedule`、`ssa.tvm.optimize_schedule` |
+
+默认 pipeline 形态为：
+
+```text
+ssa.canonicalize
+-> ssa.analyze_effects
+-> ssa.select_schedule
+-> ssa.<backend>.optimize_schedule
+-> ssa.lower_memory_scopes
+-> ssa.lower_backend_intrinsics
+```
+
+## 执行流程
+
+编译入口在 backend lowering 前调用 `lower_ssa_for_backend`。整体流程为：
+
+1. 如果 `KernelIR` 只有 `ProgramIR`，先转换为 generic `SSAProgramIR`。
+2. 根据目标后端、`compiler_options`、`kernel_metadata` 和显式参数选择 pipeline。
+3. 通过 `SSAPassRegistry` 把 pass 名称解析为 descriptor 和 pass 实例。
+4. 逐个执行 pass，并在 SSA metadata 中记录：
+   - `pass_trace`
+   - `pipeline_selection`
+   - `schedule`
+   - `optimization`
+   - `memory_scope`
+   - `backend_intrinsics`
+5. 后端代码生成器读取这些 metadata 和 operation annotation，生成目标代码。
+
+## 注册器接口
+
+核心对象：
+
+```python
+from ninetoothed.ssa_passes import (
+    DEFAULT_SSA_PASS_REGISTRY,
+    registered_ssa_passes,
+)
+```
+
+查询 pass：
+
+```python
+registered_ssa_passes(category="hardware_independent")
+registered_ssa_passes(category="backend_specific", backend="triton")
+```
+
+新增 pass 的基本方式：
+
+```python
+from ninetoothed.ssa_passes import SSAPass, BACKEND_SPECIFIC
+from ninetoothed.backends.base import BackendName
+
+
+class MyTritonPass(SSAPass):
+    name = "ssa.triton.my_pass"
+    category = BACKEND_SPECIFIC
+    phase = "optimization"
+    supported_backends = (BackendName.TRITON,)
+
+    def run(self, program, context):
+        return program
+
+
+DEFAULT_SSA_PASS_REGISTRY.register(MyTritonPass)
+```
+
+## 自定义 Pipeline
+
+可以直接传入 `SSAPipelineSpec`：
+
+```python
+from ninetoothed.ssa_passes import SSAPipelineSpec, lower_ssa_for_backend
+
+spec = SSAPipelineSpec(
+    passes=(
+        "ssa.canonicalize",
+        "ssa.analyze_effects",
+        "ssa.select_schedule",
+        "ssa.triton.optimize_schedule",
+        "ssa.lower_memory_scopes",
+        "ssa.lower_backend_intrinsics",
+    ),
+    mode="custom",
+    pass_options={
+        "ssa.triton.optimize_schedule": {
+            "tile": {"block_m": 64, "block_n": 32, "block_k": 32},
+            "num_warps": 4,
+        }
+    },
+)
+
+target_ssa = lower_ssa_for_backend(generic_ssa, backend="triton", pass_pipeline=spec)
+```
+
+也可以通过 `KernelIR.compiler_options` 或 `KernelIR.metadata` 传入：
+
+```python
+compiler_options={
+    "ssa_pass_pipeline": {
+        "mode": "custom",
+        "passes": (...),
+        "pass_options": {...},
+    }
+}
+```
+
+## 自动调优入口
+
+当前提供 policy-based autotune 入口：
+
+```python
+target_ssa = lower_ssa_for_backend(generic_ssa, backend="triton", autotune=True)
+```
+
+它会根据 SSA opcodes、program kind、目标后端选择候选 pipeline，并把选择结果记录到：
+
+```python
+target_ssa.metadata["pipeline_selection"]
+```
+
+当前 autotune 还不是运行时测量型，而是静态策略型。它已经具备三个关键扩展点：
+
+1. `candidate_pipelines`：记录候选 pass 序列。
+2. `pass_options`：记录候选配置，例如 tile、num_warps、num_stages。
+3. `reason`：记录选择原因，便于报告和调试。
+
+后续可以在这个接口上接入真实运行时 benchmark：
+
+1. 生成多个 `SSAPipelineSpec`。
+2. 对每个 spec 生成后端代码并运行 microbenchmark。
+3. 缓存 shape/operator/backend 对应的最优 spec。
+4. 以后遇到同类算子直接复用缓存 spec。
+
+## 当前后端消费情况
+
+Triton backend 已经实际消费 `optimization` metadata：
+
+- elementwise/fill/copy：读取 `block_size`，使用更大的 coalesced vector block。
+- matmul：读取 tile、num_warps、num_stages、input_precision、小问题阈值。
+- 小矩阵使用 micro-kernel，中大矩阵使用二维 block + `tl.dot`。
+
+CUDA/TileLang/TVM 目前也能收到后端专属 optimize pass 产生的 metadata，后续可以逐步把更多 schedule 决策接入各自代码生成器。
+
+## 正确性约束
+
+- pass registry 只控制 SSA pass 执行，不改变前端语义。
+- backend-specific pass 只能添加或转换细粒度 SSA annotation，不允许把整个算子包成粗粒度节点。
+- 所有 pass 执行结果都必须可序列化进 SSA metadata，便于审计、报告和自动调优缓存。