[Architecture] VPTO IR cache 机制应迁移到 DSL 框架侧

[Zhendong404]

Background

当前 TileLib / TileLang 模板实例化路径中，PTOAS 在 ExpandTileOp pass 内维护了一份进程内缓存：

• ExpandState::specCache 以 SpecKey 为 key 缓存已经实例化出来的 func.func
• 命中时直接复用当前 module 中已克隆进去的实例函数
• 未命中时，再调用 python -m tilelang_dsl.expand_helper ...，解析输出 MLIR，克隆到 module 末尾，然后写回缓存

这套机制当前能避免“同一 module 内相同 Tile op 重复实例化”，但它的职责位置更偏 PTOAS 内部实现，而不是 DSL 框架集成边界。

从架构上看，VPTO IR / TileLib 实例缓存更适合作为 DSL 框架侧能力，而不是依赖 ExpandTileOp 的进程内 DenseMap：

• 当前 cache 生命周期只在一次 pass / 一次 ptoas 进程内有效
• cache 命中结果绑定的是 MLIR 符号克隆结果，不是一个更稳定的框架层 artifact
• 上层框架很难复用已有实例，也无法提前控制 cache key / cache invalidation / artifact 存储

Problem

把 cache 放在 PTOAS ExpandTileOp 内部，会带来一些边界和演进问题：

• 同一个模板实例跨 module / 跨进程 / 跨编译任务无法复用
• 框架侧看不到也控制不了实例缓存，无法把它纳入自己的编译缓存体系
• 当前 SpecKey、函数命名、helper 调用协议、module clone 逻辑被耦合在一起，导致 cache 语义更像 pass 内部优化，而不是稳定集成契约
• 一旦未来希望由框架直接产出或持久化 VPTO authoring IR / 实例化结果，现在的 pass 内缓存无法自然承接
• cache 命中与否影响 ExpandTileOp 的性能，但这一性能行为当前不能由 DSL 框架统一管理

Suggestion

建议把“实例化结果缓存”的主责迁移到 DSL / 上层框架侧，PTOAS 侧只保留最薄的一层消费逻辑。

1. 明确 cache 的归属边界

建议定义成：

• DSL / 框架侧负责根据稳定协议计算 specialization key
• DSL / 框架侧负责决定是否命中、如何持久化、如何失效
• PTOAS 侧负责消费实例化结果，做 IR 接入、校验、后续 lowering

2. 将当前 pass 内 cache 语义降级为实现细节或过渡机制

ExpandTileOp 内部可以保留一个很轻量的“同进程去重”机制，但不应继续作为主要 cache 语义来源。

也就是说，理想状态应该是：

• 主要实例缓存由框架侧管理
• PTOAS 接收“已经选好并实例化好的 artifact”或更明确的 DSL 产物
• pass 内缓存即使存在，也只是本地优化，而不是对外协议的一部分

3. 把 cache key 协议文档化并前移到框架侧

当前实现中已经隐含了一套 key 规则：

• Tile: dtype + shape + valid_shape + memorySpace + config
• View: 只按 dtype 进入 key，shape/strides/memorySpace 只参与约束检查
• Scalar: dtype
• 另外还有 target、op、context_attrs

建议把这套规则整理成正式协议，并明确：

• 哪些字段决定实例边界
• 哪些字段只参与校验
• key 变更时 cache invalidation 怎么做
• DSL 版本 / 模板版本是否需要进入 key

4. 明确 PTOAS 消费的 artifact 形式

建议讨论并选定一种更稳定的框架侧产物，例如：

• 已实例化的 authoring-form MLIR 函数
• 可 import 的 TileLib 实例 module
• 或者更高一层的“模板实例描述 + IR artifact”组合

核心是让 PTOAS 面向“明确产物”而不是面向“内部 helper 调用 + 进程内 cache”。

Expected outcome

• VPTO IR / TileLib 实例缓存成为 DSL 框架能力的一部分，便于跨任务、跨模块复用
• PTOAS 与 DSL 框架之间的边界更清晰：框架产出实例化 artifact，PTOAS 负责消费和 lowering
• 后续如果要做编译缓存、远程缓存、增量编译、预热实例库，会更自然
• ExpandTileOp 不再承担过多“框架缓存层”的职责

Possible acceptance criteria

• 明确一份框架侧 specialization key 协议，并与 PTOAS 当前实现对齐或完成迁移设计
• 设计并落地一种可由框架持久化的实例 artifact 形式
• PTOAS 能消费该 artifact，而不是必须自己调用 helper 才能命中缓存
• ExpandTileOp 中现有 specCache 要么退化成局部优化，要么被新的消费路径替代

Additional context

当前现状可以直接从仓内实现中看到：

• lib/PTO/Transforms/ExpandTileOp.cpp 中 ExpandState::specCache
• buildOperandSpecsJson() / buildContextAttrsJson() 负责构造 helper 输入
• invokeTilelangDSL() 负责 cache check -> 调 helper -> 解析 MLIR -> clone 到 module -> 回填 cache
• docs/designs/ptoas-tileop-expand-design.md 3.2.1 / 3.2.2 已经描述了现有 SpecKey 和缓存策略

[bingmeiyou]

---

一、Issue 分析验证

1.1 仓库核心架构

PTOAS是一个编译器，连接上层AI框架和底层NPU硬件：

TileLang DSL (Python模板库) → PTOAS (MLIR编译器) → LLVM IR → NPU二进制

两个层次的IR：

IR层级	名称	操作对象	示例
TileOp层	高层抽象	tile_buf（一个数据块）	pto.tadd ins(%a, %b) outs(%c)
Vector IR层	底层向量指令	vreg（向量寄存器）	pto.vlds, pto.vadd, pto.vsts

一条pto.tadd需要展开成几十条向量指令（循环、load、计算、store）。

1.2 模板展开机制

问题：dtype有f16/f32/bf16等，shape可以是任意大小，组合无穷多，无法手写所有实现。

方案：用Python DSL写模板，编译时按参数特化实例化。

lib/TileOps/tadd_template.py 是一个模板，当PTOAS遇到pto.tadd时：

1. 提取参数的dtype/shape/layout等信息
2. 调用Python helper生成特化的MLIR函数
3. 把函数clone进当前module，替换原TileOp为func.call

1.3 Issue指出的缓存问题

代码位置：ExpandTileOp.cpp:450-452

struct ExpandState {
  llvm::DenseMap<SpecKey, func::FuncOp, SpecKeyInfo> specCache;  // 进程内缓存
  ...
};

问题本质：

现状	限制
缓存是DenseMap	只在单次pass运行期间有效
每次pass新建ExpandState	pass结束缓存销毁
缓存存的是克隆后的FuncOp	只能在当前module内复用

实际影响：

• 同一个module内多处pto.tadd(f32, 16x64)会复用一份实例 ✅
• 不同module编译、或不同进程编译，相同参数要重新实例化 ❌
• DSL框架看不到缓存，无法提前预热或持久化 ❌

1.4 Issue建议的改进

把缓存职责从编译器pass迁移到DSL框架侧：

当前	建议
PTOAS内部DenseMap	DSL框架持久化缓存
缓存key隐含在C++代码	key协议文档化、稳定化
PTOAS调用Python helper	DSL框架产出artifact，PTOAS消费
pass内缓存是主要机制	降级为局部优化，不是对外协议

好处：

• 跨编译任务复用（比如预热的实例库）
• 框架可控制缓存失效策略
• 未来支持增量编译、远程缓存等

1.5 验证结论

Issue描述的代码行为、缓存机制、限制都是准确的。建议的架构改进也是合理的，本质是让缓存成为DSL框架的能力而非编译器内部实现细节。

---

二、设计方案

2.1 整体架构

┌─────────────────────────────────────────────────────────────────┐
│                      DSL Framework (Python)                      │
│  ┌──────────────┐  ┌──────────────┐  ┌──────────────────────┐  │
│  │ Template Lib │  │ Kernel Reg.  │  │  Instance Cache      │  │
│  │ (TileOps/*.py)│  │              │  │  - In-process cache  │  │
│  │              │  │              │  │  - Persistent cache  │  │
│  └──────────────┘  └──────────────┐  │  - Cache key protocol │  │
│                    ┌──────────────┐  └──────────────────────┘  │
│                    │ Artifact     │                              │
│                    │ Manager      │←─── 缓存命中时直接返回       │
│                    │              │                              │
│                    └──────────────┘                              │
└────────────────────────────┬────────────────────────────────────┘
                             │ Stable Protocol (JSON / MLIR text)
                             ▼
┌─────────────────────────────────────────────────────────────────┐
│                      PTOAS (C++ MLIR)                            │
│  ┌──────────────────┐  ┌──────────────────┐  ┌──────────────┐  │
│  │ ExpandTileOp     │  │ Artifact Consumer│  │ Pass-local   │  │
│  │ Pass             │→ │ (parse + clone)  │  │ Dedup Cache  │  │
│  │                  │  │                  │  │ (fallback)   │  │
│  └──────────────────┘  └──────────────────┘  └──────────────┘  │
│                                                              │
│  输入: 已实例化的 MLIR artifact (可选 cache key hint)        │
│  输出: clone 到 module 的 FuncOp                            │
└─────────────────────────────────────────────────────────────────┘

2.2 核心变更

2.2.1 DSL 框架侧新增 Instance Cache

新增 Python 模块 tilelang_dsl/instance_cache.py：

class InstanceCache:
    """框架侧模板实例缓存"""
    
    def __init__(self, 
                 cache_dir: Path | None = None,
                 enable_persistent: bool = True):
        self._in_process_cache: dict[CacheKey, InstanceArtifact] = {}
        self._persistent_cache_dir = cache_dir
        self._enable_persistent = enable_persistent
    
    def compute_cache_key(
        self,
        target: str,
        op: str,
        operand_specs: list[dict],
        context_attrs: dict | None = None,
    ) -> CacheKey:
        """按稳定协议计算 cache key"""
        ...
    
    def lookup(self, key: CacheKey) -> InstanceArtifact | None:
        """查询缓存（in-process + persistent）"""
        ...
    
    def store(self, key: CacheKey, artifact: InstanceArtifact) -> None:
        """写入缓存"""
        ...
    
    def invalidate(self, key: CacheKey | None = None) -> None:
        """失效缓存（可指定单个 key 或全部）"""
        ...

2.2.2 Cache Key 协议文档化

定义稳定的 cache key 计算规则：

@dataclass(frozen=True)
class CacheKey:
    """稳定的实例化缓存 key"""
    
    target: str                         # 架构: "a5", "a3"
    op: str                             # 算子名: "tadd", "tmul"
    operand_keys: tuple[OperandKey, ...]  # 每个操作数的 key
    context_key: frozenset[tuple[str, str]]  # context attrs 的 sorted tuple

Cache Key 规则（稳定协议）：

操作数类型	参与 Cache Key 的字段	不参与但传给模板的字段
Tile	dtype, shape, valid_shape, memory_space, blayout, slayout, fractal, pad	—
View	dtype	shape, strides, memory_space（仅约束检查）
Scalar	dtype	—

2.2.3 Instance Artifact 定义

@dataclass
class InstanceArtifact:
    """实例化产物"""
    
    cache_key: CacheKey
    mlir_text: str                     # 生成的 MLIR 文本
    function_name: str                 # 主函数名
    metadata: dict                     # 元数据
    source_descriptor: str             # 来源模板名
    
    def save(self, path: Path) -> None:
        """持久化到文件"""
        ...
    
    @classmethod
    def load(cls, path: Path) -> "InstanceArtifact":
        """从文件加载"""
        ...

2.2.4 DSL 框架 API 扩展

def instantiate_tile_op(
    target: str,
    op: str,
    operand_specs: list[dict],
    context_attrs: dict | None = None,
    *,
    cache: InstanceCache | None = None,
    use_cache: bool = True,
) -> InstanceArtifact:
    """实例化 TileOp，返回 artifact"""
    ...

def warmup_cache(
    target: str,
    ops: list[str] | None = None,
    common_specs: list[dict] | None = None,
) -> int:
    """预热缓存"""
    ...

2.2.5 PTOAS 侧消费 Artifact

修改 ExpandTileOp.cpp，支持两种模式：

模式 1：PTOAS 调用 DSL helper（向后兼容）

func::FuncOp invokeTilelangDSL(const SpecKey &key, ...);

模式 2：PTOAS 接收预实例化 Artifact（新模式）

// 新增选项: --artifact-cache-dir
struct ArtifactManifest {
  std::string cacheKeyJson;
  std::string mlirPath;
  std::string metadataJson;
};

func::FuncOp tryLoadFromArtifactCache(...);

PTOAS 侧变更摘要：

1. ExpandState 中 specCache 降级为"同 module 内去重"的 pass-local cache
2. 新增 artifactCache 用于从预实例化 artifact 目录加载
3. 优先级：artifactCache > specCache > invoke DSL helper

---

三、实现计划

3.1 DSL 框架侧文件结构

tilelang-dsl/python/tilelang_dsl/
├── instance_cache.py          ← 新增: 缓存管理
├── cache_key.py               ← 新增: cache key 协议
├── artifact.py                ← 新增: artifact 定义
├── instantiate.py             ← 新增: 实例化 API
├── expand_helper.py           ← 修改: 使用框架缓存
└── ...

3.2 PTOAS 侧改造

// lib/PTO/Transforms/ExpandTileOp.cpp

struct ExpandTileOpOptions {
  std::string artifactCacheDir;  // 新增：预实例化 artifact 目录
};

struct ExpandState {
  // pass-local cache：仅用于同一 module 内去重
  llvm::DenseMap<SpecKey, func::FuncOp, SpecKeyInfo> specCache;
  
  // 新增：artifact cache manifest
  llvm::StringMap<ArtifactManifest> artifactManifest;
  
  func::FuncOp tryLoadFromArtifactCache(...);
};

// 主流程优先级
// 1. Try artifact cache
// 2. Try pass-local cache  
// 3. Invoke DSL helper (uses framework cache internally)

3.3 Cache Key Hash 计算（C++ 与 Python 对齐）

static std::string computeCacheKeyHash(const SpecKey &key) {
  // 与 Python cache_key.py 的 to_hash_input 对齐
  std::string input;
  input += key.targetArch + "|";
  input += key.opName + "|";
  
  for (const auto &op : key.operands) {
    if (op.kind == OperandKind::Tile) {
      input += "tile:" + op.dtype + ":";
      // shape, valid_shape, memory_space, layout...
    } else if (op.kind == OperandKind::View) {
      input += "view:" + op.dtype;
    } else {
      input += "scalar:" + op.dtype;
    }
  }
  
  return llvm::SHA256Hash(input).substr(0, 16);
}

---

四、迁移路径

Phase 1：DSL 框架侧缓存（不影响 PTOAS）

1. 实现 instance_cache.py、cache_key.py、artifact.py
2. 改造 expand_helper.py 使用框架缓存
3. 添加单元测试验证缓存行为
4. PTOAS 行为不变（expand_helper 内部使用缓存）

验收标准：

• 多次调用相同参数的 expand_helper，第二次命中缓存
• 缓存命中时跳过模板实例化过程
• tilelang_dsl.get_global_cache().stats() 返回正确的命中率

Phase 2：PTOAS 侧消费 Artifact

1. 实现 PTOAS 侧 computeCacheKeyHash
2. 新增 --artifact-cache-dir 选项
3. 实现 tryLoadFromArtifactCache
4. 优先级：artifactCache > specCache > DSL helper

验收标准：

• PTOAS 可从预实例化目录加载 artifact
• 指定 artifact 目录时，优先使用而非调用 DSL
• 无 artifact 目录时，行为与当前一致

Phase 3：缓存预热与持久化

1. 实现 warmup_cache() API
2. 添加 CLI 工具预热常用实例
3. 验证跨进程缓存复用

验收标准：

• 预热后，新进程编译可命中持久化缓存
• 缓存失效 API 正常工作

Phase 4：文档与协议稳定化

1. 将 Cache Key 协议写入设计文档
2. 版本号纳入 cache key（DSL 版本、模板版本）
3. 定义缓存失效策略（版本升级、模板更新）

验收标准：

• Cache Key 协议文档完整
• 协议变更时缓存自动失效

---

五、风险与缓解

风险	缓解措施
C++ 和 Python cache key 计算不一致	单元测试验证两端 hash 一致性；提供校验工具
持久化缓存无限增长	设置 TTL 和 max_entries；提供清理 CLI
模板更新后旧缓存导致错误	DSL 版本号纳入 cache key；提供 invalidate API

---

六、验收标准汇总

阶段	验收项
Phase 1	DSL 框架缓存正常工作，命中率统计正确
Phase 2	PTOAS 可从 artifact 目录加载，优先级正确
Phase 3	跨进程缓存复用，预热 API 正常
Phase 4	协议文档完整，版本变更自动失效
整体	同一 module 内去重（pass-local cache）保持
整体	向后兼容：不指定 artifact 目录时行为不变
整体	性能：缓存命中时编译时间显著减少

---

附录：关键代码位置

文件	说明
lib/PTO/Transforms/ExpandTileOp.cpp:450-452	当前缓存实现 specCache
lib/PTO/Transforms/ExpandTileOp.cpp:575-779	invokeTilelangDSL 实例化流程
lib/PTO/Transforms/ExpandTileOp.cpp:128-164	SpecKey 定义和 hash 计算
docs/designs/ptoas-tileop-expand-design.md:468-485	现有 SpecKey 规则文档
tilelang-dsl/python/tilelang_dsl/expand_helper.py	Python helper 入口
tilelang-dsl/python/tilelang_dsl/kernel.py	VKernelDescriptor 定义
lib/TileOps/*.py	75+ 个模板文件

[Zhendong404]

补充一下我对这个方案的建议，整体方向我认同，但我倾向于再做一次收敛，把 cache 机制尽量完全隐藏在 DSL 框架内部，而不是继续扩散到 PTOAS 侧。

我比较认同下面几个点：

1. 不建议把“框架侧缓存”继续做成 expand_helper 这种一次性 CLI 内部的小优化。
   当前每次 ExpandTileOp miss 都会重新起一个 Python 子进程，重新扫描模板、导入模块、匹配 descriptor，然后输出 MLIR 后退出。这样即使 DSL 内部做了进程内 dict cache，也没有跨请求复用价值。
   所以更合适的方式是：在启动 ptoas 的早期就拉起一个 TileLang 守护进程，由 ExpandTileOp 通过一个唯一接口向它请求 MLIR 实例化结果。守护进程内部负责：

• descriptor 选择
• cache 命中判定
• 并行实例化
• 同 key 请求去重
• 版本变更后的自然失效

PTOAS 只负责发送请求和消费返回的 MLIR，不需要感知 cache 实现细节。

2. TensorView 只靠 dtype 做 cache key 不够，必须增强 key 的稳定性。
   现在 DSL 的 kernel selection / constraint evaluation 不只是看 dtype，还可能依赖：

• shape
• strides
• memory_space
• valid_shape
• context_attrs

所以 cache key 的设计原则不能只看“哪些字段影响最终代码生成”，还必须看“哪些字段会影响选中哪个 kernel / descriptor”。只要会影响选择结果，就必须进入 key，否则就可能出现缓存命中错误，拿到不应该复用的实例。

3. 不建议引入 warm_up 机制。
   这个场景下 warm_up 带来的收益有限，但会明显增加接口和维护复杂度，而且容易把“优化策略”固化成外部协议。既然守护进程可以常驻，首选还是让它在真实请求路径里自然积累缓存，而不是额外设计一套预热 API。

4. 不建议引入 artifact 机制，也尽量不要让 PTOAS 侧显式参与 cache 协议。
   我更倾向于：

• PTOAS 看不到 artifact-cache-dir
• PTOAS 不计算 hash key
• PTOAS 不区分 artifact cache / local cache / DSL helper fallback
• PTOAS 只调用一个 daemon 接口：给定 target + concrete_op + operand_specs + context_attrs，拿回 materialized MLIR module

这样缓存策略、失效策略、并发策略都能留在 DSL 框架内部演进，不会把实现细节泄漏到 C++ 边界上。

我目前更推荐的收敛方向可以概括成一句话：

ExpandTileOp 通过单一 RPC/IPC 接口向 TileLang daemon 请求已 materialize 的 MLIR module；daemon 内部负责 descriptor 选择、稳定 key 计算、single-flight、缓存与失效；PTOAS 不感知 warm-up、artifact 或任何缓存细节。

这样既能真正解决当前 Python 子进程模型下缓存收益不足的问题，也能避免把 cache 协议固化到 PTOAS 侧，后续 DSL 继续演进时边界会更清晰。

[bingmeiyou]

--

一、方案概述

核心设计

一句话总结：

ExpandTileOp 通过单一 RPC 接口向 TileLang daemon 请求 materialized MLIR module；daemon 内部负责 descriptor 选择、stable key 计算、single-flight、缓存与失效；PTOAS 不感知 warm-up、artifact 或任何缓存细节。

架构图：

┌─────────────────────────────────────────┐
│ TileLang Daemon (Python, 常驻进程)       │
│                                          │
│ ├─ Template Registry    ← 菜谱仓库       │
│ ├─ Kernel Selection     ← 找菜谱选模板   │
│ ├─ Instance Cache       ← 记忆（缓存）   │
│ ├─ Stable Key Compute   ← 算菜名        │
│ ├─ Single-flight        ← 并发去重      │
│ ├─ Version Tracking     ← 版本追踪      │
│ ├─ RPC Server           ← 接电话        │
│                                          │
│ 唯一接口: instantiate(target, op, specs, attrs) │
└─────────────────────────────────────────┘
            │ Socket (RPC/IPC)
            ▼
┌─────────────────────────────────────────┐
│ PTOAS Compiler (C++, MLIR)               │
│                                          │
│ ├─ ExpandTileOp Pass   ← 读代码          │
│ ├─ Daemon Client       ← 拨电话点菜      │
│ ├─ MLIR Consumer       ← 收菜、复制      │
│ ├─ Pass-local Dedup    ← 同module去重   │
│                                          │
│ 不感知: cache key、hash、缓存细节         │
└─────────────────────────────────────────┘

核心原则

1. 缓存完全隐藏在 DSL 框架内部：PTOAS 不感知缓存实现细节
2. 单一 RPC 接口：PTOAS 只调用 instantiate()，不管理缓存
3. 常驻进程模型：daemon 常驻运行，实现真正的跨请求缓存复用
4. 增强 stable key：View 参数（shape/strides/ms）全部进 key

---

二、对比：当前问题 vs 新方案

问题	当前状态（CLI helper）	新方案（Daemon）
进程模型	每次启动新 Python 子进程	一个进程常驻
跨请求复用	❌ 无法复用（进程退出丢失）	✅ 完全复用（常驻保留）
缓存收益	❌ 只有 pass 内小优化	✅ 真正跨请求复用
PTOAS 感知	❌ 知道细节、计算 hash	✅ 只调用接口，不感知
接口复杂度	❌ 3层 fallback	✅ 单一 RPC
View stable key	❌ 只有 dtype	✅ 增强：shape/strides/ms
并发处理	❌ 无	✅ Single-flight 去重
版本追踪	❌ 无	✅ 文件 hash 自动失效
预热机制	❌ 无	✅ 不需要，自然积累
Artifact	❌ 无	✅ 不引入，简化

---

三、关键设计细节

3.1 Daemon 生命周期

启动方式：

// PTOAS 自动启动（推荐）
void ExpandTileOpPass::runOnOperation() {
    if (!checkDaemonRunning(socketPath)) {
        // 自动启动 daemon
        system("python -m tilelang_dsl.daemon --socket /tmp/... &");
        waitForDaemonReady(socketPath);
    }
    // 开始编译
    daemonClient.instantiate(...);
}

生命周期管理：

阶段	方式
启动	PTOAS 编译时自动启动（首次编译）
运行	常驻运行，监听 socket，不自动退出
退出	手动 kill 或系统重启，不用空闲退出

理由：

• 编译任务通常批量进行，缓存有价值
• 常驻运行性价比高（启动开销 > 空闲内存）
• 内存占用可控（~10MB），不自动退出

---

3.2 Stable Key 增强

专家核心观点：

"cache key 的设计原则不能只看'哪些字段影响最终代码生成'，还必须看'哪些字段会影响选中哪个 kernel / descriptor'。只要会影响选择结果，就必须进入 key，否则可能出现缓存命中错误。"

增强规则：

操作数类型	字段	进 Key	原因
Tile	dtype, shape, valid_shape, ms, config	✅ 全部	原来就全
View	dtype	✅	dtype signature matching
View	shape	✅ 新增	可能影响 constraint evaluation
View	strides	✅ 新增	可能影响 constraint evaluation
View	memory_space	✅ 新增	可能影响 constraint evaluation
Scalar	dtype	✅	dtype signature matching

为什么增强 View：

# kernel.py 的 constraint evaluation 可能读取：
attrs[f"{spec.name}_strides"] = ...      ← strides 参与！
attrs[f"{spec.name}_memory_space"] = ... ← memory_space 参与！

# 从而影响 kernel selection：
for constraint in self.constraints:
    result = constraint(context_attrs)  ← 可能读取 strides、ms
    if not result:
        raise LookupError("kernel rejected")

结论：View 的 shape/strides/ms 可能影响选模板 → 必须进 stable key

---

3.3 Single-flight 机制

问题：多个并发请求同 key，重复实例化浪费资源

解决方案：

async def instantiate(key, ...):
    # 1. 查缓存
    if key in cache:
        return cache[key]  ← 命中
    
    # 2. Single-flight：检查有没有正在做的
    if key in pending_requests:
        return await pending_requests[key]  ← 等别人的结果
    
    # 3. 我来做
    future = asyncio.Future()
    pending_requests[key] = future
    
    mlir = await _do_instantiate(...)  ← 实例化
    cache[key] = mlir                   ← 记缓存
    
    future.set_result(mlir)             ← 通知所有等待者
    return mlir

效果：10个并发同 key请求 → 只实例化1次 → 分给10个请求

---

3.4 Version Tracking

问题：模板更新后，旧缓存导致错误

解决方案：

def _scan_template_versions():
    for py_file in template_dir.glob("*.py"):
        # 文件 hash 作为版本号
        version = hashlib.sha256(py_file.read_bytes()).hexdigest()[:8]
        template_versions[op] = version

def _check_version_change():
    current = _scan_template_versions()
    for op, version in current.items():
        if version != template_versions[op]:
            # 版本变了 → 失效缓存
            _invalidate_by_op(op)

效果：模板文件变更 → 自动失效相关缓存 → 强制重新实例化

---

3.5 缓存管理

存储位置：内存（RAM），daemon 进程的 Python dict

配置：

MAX_ENTRIES = 1000    # 缓存上限
# 内存占用：~10MB（单实例 ~10KB × 1000）
# 淘汰：FIFO（扔最早的）

淘汰机制：

def store(key, artifact):
    if len(cache) >= MAX_ENTRIES:
        oldest = next(iter(cache))  ← FIFO
        del cache[oldest]
    cache[key] = artifact

为什么不持久化：

• daemon 常驻运行，缓存在内存够用
• 持久化增加复杂度，收益不明显
• 先不实现，有需求再加

---

四、实现计划

Phase 1: TileLang Daemon（框架侧）

新增文件：

tilelang-dsl/python/tilelang_dsl/
├── daemon.py              ← 主程序、生命周期
├── daemon_core.py         ← 核心逻辑：缓存、单飞
├── stable_key.py          ← Stable key 计算（增强）
├── daemon_server.py       ← RPC server（socket）

验收标准：

• ✅ daemon 可启动并常驻
• ✅ 接 RPC 请求，返回 MLIR 文本
• ✅ 同样请求第二次命中缓存
• ✅ 同 key 并发，single-flight 去重
• ✅ 模板变更，自动失效缓存

---

Phase 2: PTOAS Client（编译器侧）

新增文件：

lib/PTO/Daemon/
├── TileLangDaemonClient.h    ← 接口定义
├── TileLangDaemonClient.cpp  ← RPC 实现

修改文件：

lib/PTO/Transforms/
├── ExpandTileOp.cpp          ← 使用 daemon client

验收标准：

• ✅ PTOAS 可连接 daemon，发送请求
• ✅ 成功返回 MLIR 文本
• ✅ 不感知 cache key、hash、缓存细节
• ✅ daemon 不可用，fallback 到 CLI

---

Phase 3: 自动启动和管理

实现：

• PTOAS 启动时检查 daemon
• 未运行则自动启动
• 等待就绪后编译

验收标准：

• ✅ 不手动启动，直接编译，daemon 自动启动
• ✅ 编译成功，缓存生效

---

Phase 4: 测试和文档

测试：

• 单次编译正确性
• 多次编译缓存命中
• 并发编译 single-flight
• 模板变更缓存失效
• fallback 正常
• 跨进程复用

文档：

• Stable key 协议（哪些字段进 key）
• Daemon 启动和管理
• Kernel selection 影响因素分析

---

五、专家建议回应

专家建议 1：不用 CLI helper，用 daemon 常驻

回应：✅ 完全采纳

当前问题：
  ├─ 每次启动 Python 子进程 → 慢（秒级）
  ├─ 进程内 dict cache → 无跨请求价值
  └─ 重复扫描、匹配、实例化
  
新方案：
  ├─ daemon 常驻 → 快（微秒级缓存命中）
  ├─ 跨请求缓存 → 真正复用
  └─ 只启动一次

---

专家建议 2：View key 增强

回应：✅ 完全采纳

旧版：
  View key = dtype  ← 不够
  
新版：
  View key = dtype + shape + strides + memory_space  ← 增强
  
原因：
  ├─ constraint evaluation 可能读取这些参数
  ├─ 影响 kernel selection → 选择哪个模板
  └─ 如果不进 key → 可能选错模板

---

专家建议 3：不用 warm-up

回应：✅ 完全采纳

理由：
  ├─ daemon 常驻，自然积累缓存
  ├─ 不需要额外预热 API
  ├─ 不增加复杂度
  └─ 用得越多，缓存越多

---

专家建议 4：不用 artifact，PTOAS 不感知缓存

回应：✅ 完全采纳

V1 方案（复杂）：
  ├─ PTOAS 知道 artifact-cache-dir
  ├─ PTOAS 计算 hash
  ├─ PTOAS 判断优先级（3层）
  └─ 边界模糊
  
V2 方案（极简）：
  ├─ PTOAS 只调用 daemon.instantiate()
  ├─ 不计算 hash、不感知缓存
  ├─ 单一 RPC 接口
  └─ 边界清晰
  
好处：
  ├─ daemon 内部改进，PTOAS 不用改
  ├─ 缓存策略变化，PTOAS 无感知
  └─ PTOAS 代码简单

---

专家建议 5：单一 RPC 接口

回应：✅ 完全采纳

# 唯一接口
def instantiate(target, op, operand_specs, context_attrs) -> str:
    """返回 materialized MLIR module"""
    
    # daemon 内部处理：
    # ├─ stable key 计算
    # ├─ 缓存查找
    # ├─ kernel selection
    # ├─ single-flight
    # ├─ 版本追踪
    # └─ 实例化
    
    return mlir_text

---

六、风险和缓解

风险	缓解措施
daemon 进程崩溃	fallback 到 CLI helper；PTOAS 自动重启
socket 连接失败	fallback；重试机制；自动启动
模板变更缓存未失效	文件 hash 版本追踪；定期检查
stable key 遗漏字段	详细分析 constraint evaluation；文档化
并发请求竞争	Single-flight；asyncio lock
内存占用过大	Cache limit（1000）；FIFO 淘汰
跨平台 socket	Unix socket + Windows named pipe

---

七、总结

核心改动

项目	改动
架构	CLI helper → TileLang daemon 常驻
缓存	pass 内 → daemon 内（跨请求）
接口	3层 fallback → 单一 RPC
PTOAS 感知	知道细节 → 不感知
View key	dtype → dtype + shape + strides + ms
并发	无 → Single-flight
预热	无 API → 不需要
Artifact	引入 → 不引入

预期收益

• ✅ 跨请求缓存复用（核心目标）
• ✅ PTOAS 边界清晰，不感知缓存细节
• ✅ View stable key 增强，正确性保证
• ✅ 简化接口，不引入 warm-up、artifact
• ✅ 并发优化，single-flight 去重
• ✅ 版本追踪，自动失效

---

附录：关键代码位置

功能	文件位置
Daemon 主程序	tilelang-dsl/daemon.py（新增）
Daemon 核心	tilelang-dsl/daemon_core.py（新增）
Stable key	tilelang-dsl/stable_key.py（新增）
RPC server	tilelang-dsl/daemon_server.py（新增）
Kernel selection	tilelang-dsl/kernel.py（复用，分析影响）
Constraint evaluation	tilelang-dsl/kernel.py:1148-1256（分析）
Daemon client	lib/PTO/Daemon/TileLangDaemonClient.cpp（新增）
ExpandTileOp	lib/PTO/Transforms/ExpandTileOp.cpp（修改）
当前 SpecKey	lib/PTO/Transforms/ExpandTileOp.cpp:128-164（参考）

[Zhendong404]

1. Daemon应该随PTOAS主进程一起退出
2. 不需要考虑模板更新缓存失效
3. 其他OK

[bingmeiyou]

一、问题分析

1.1 当前实现问题

现状：

// lib/PTO/Transforms/ExpandTileOp.cpp:450-452
struct ExpandState {
  llvm::DenseMap<SpecKey, func::FuncOp, SpecKeyInfo> specCache;
};

• 缓存是进程内 DenseMap
• 只在单次 pass 运行期间有效
• pass 结束缓存销毁

问题分析：

问题	影响
每次遇到 Tile op 启动 Python 子进程	启动开销大（秒级）
Python 子进程每次重新扫描模板	重复工作
Python 进程退出后缓存丢失	无跨 tile op 复用
缓存只在 pass 内有效	无跨 module 复用
PTOAS 内部管理缓存	框架侧无法控制

具体举例：

编译一个包含 3 个 pto.tadd 的 module：

当前流程：
  ├─ tile op1 → 启动 Python → 扫描模板 → 实例化 → 输出 MLIR → Python 退出
  ├─ tile op2 → 启动 Python → 扫描模板 → 实例化 → 输出 MLIR → Python 退出
  ├─ tile op3 → 启动 Python → 扫描模板 → 实例化 → 输出 MLIR → Python 退出
  
问题：
  ├─ 启动 Python 3 次（开销大）
  ├─ 扫描模板 3 次（重复）
  ├─ 即使 op1/op2/op3 参数相同，也要实例化 3 次
  └─ pass 结束，缓存销毁

---

1.2 Issue 建议

"将 VPTO IR 缓存机制迁移到 DSL 框架侧"

核心改动：

• 框架侧负责缓存管理
• PTOAS 只消费实例化结果
• 缓存成为框架能力，而非编译器内部细节

---

二、最终方案

2.1 核心设计

一句话总结：

TileLang daemon 作为 PTOAS 子进程运行，生命周期绑定 PTOAS；单次编译过程中提供缓存服务，避免重复实例化；PTOAS 通过单一 RPC 接口请求实例化，不感知缓存细节。

架构图：

┌─────────────────────────────────────────────────────────┐
│ PTOAS 主进程                                             │
│                                                          │
│ 启动阶段：                                                │
│   ├─ 启动 daemon 子进程                                  │
│   ├─ 等待 daemon 就绪                                    │
│   └─ 连接 daemon socket                                  │
│                                                          │
│ 编译阶段：                                                │
│   ├─ 遍历 Tile ops                                       │
│   ├─ 向 daemon 发送实例化请求                            │
│   ├─ 收到 materialized MLIR                              │
│   ├─ 解析、clone、替换 Tile op                           │
│                                                          │
│ 退出阶段：                                                │
│   ├─ daemon 作为子进程自动退出                          │
│   └─ 缓存随进程清理                                      │
│                                                          │
│ 不感知：cache key 计算、缓存实现、淘汰策略               │
└─────────────────────────────────────────────────────────┘
                            │
                            │ RPC (Socket)
                            │
                            ▼
┌─────────────────────────────────────────────────────────┐
│ TileLang daemon 子进程                                   │
│                                                          │
│ ├─ Template Registry    ← 启动时扫描模板目录             │
│ ├─ Kernel Selection     ← 选择匹配的模板                 │
│ ├─ Instance Cache       ← 进程内缓存（单次编译有效）     │
│ ├─ Stable Key Compute   ← 计算缓存 key                   │
│ ├─ Single-flight        ← 并发请求去重                   │
│ ├─ RPC Server           ← 监听 socket，处理请求          │
│                                                          │
│ 唯一接口：                                               │
│   instantiate(target, op, operand_specs, context_attrs) │
│   返回：materialized MLIR module text                    │
│                                                          │
│ 生命周期：跟随 PTOAS 主进程                             │
│   父进程退出 → 子进程自动退出                            │
└─────────────────────────────────────────────────────────┘

---

2.2 对比：方案演进

方案	核心设计	专家反馈
V0（当前）	CLI helper 每次 Python 子进程	问题：无缓存复用
V1（初稿）	框架缓存 + PTOAS 消费 artifact	专家：PTOAS 感知太多，扩散边界
V2（第一轮修订）	daemon 常驻 + 跨进程缓存 + 版本追踪	专家：跨进程不需要，版本追踪过设计
V2.1（最终）	daemon 子进程 + 进程内缓存 + 绑定生命周期	✅ 通过评审

---

2.3 关键改动点

项目	当前（V0）	最终方案（V2.1）	改动原因
进程模型	每次 tile op 启动 Python	daemon 子进程（一次编译启动一次）	减少启动开销
缓存位置	pass 内 DenseMap	daemon 进程内 dict	进程内缓存有效
缓存范围	pass 内有效	单次编译过程有效	专家：跨进程不需要
生命周期	pass 结束销毁	PTOAS 退出时销毁	专家：绑定主进程
PTOAS 感知	计算 SpecKey	不感知 cache key	专家：边界清晰
接口	CLI helper stdout	RPC socket	专家：单一接口
View key	只有 dtype	dtype + shape + strides + ms	专家：影响 kernel selection
并发	无	Single-flight	需要去重
Version Tracking	无	不引入	专家：模板不变更

---

三、详细设计

3.1 Daemon 生命周期

3.1.1 启动流程

PTOAS 编译开始：
  
  Step 1: 检查 daemon 是否运行
    ├─ 尝试连接 socket: /tmp/tilelang_daemon.sock
    ├─ 如果连接成功 → daemon 已运行 → 直接使用
    └─ 如果连接失败 → daemon 未运行 → 启动 daemon
  
  Step 2: 启动 daemon 子进程
    ├─ fork() 创建子进程
    ├─ 子进程：execvp("python", "-m tilelang_dsl.daemon", ...)
    ├─ 父进程：等待子进程就绪
    └─ 等待方式：不断尝试连接 socket（最多 5秒）
  
  Step 3: 开始编译
    ├─ 连接 daemon socket
    ├─ 遍历 Tile ops，向 daemon 发送请求
    └─ 编译完成

3.1.2 运行过程

daemon 运行阶段：
  
  ├─ 监听 socket，等待请求
  ├─ 收到请求 → 处理请求
  ├─ 处理逻辑：
    ├─ 计算 stable key
    ├─ 查缓存（命中 → 直接返回）
    ├─ 未命中 → kernel selection → specialize → materialize
    ├─ 存缓存
    └─ 返回 MLIR 文本
  ├─ 单次编译过程中缓存有效
  └─ 同 key 并发请求 → single-flight 去重

3.1.3 退出流程

PTOAS 编译结束：
  
  方式 A（推荐）：子进程自动退出
    ├─ PTOAS 主进程退出
    ├─ daemon 作为子进程
    ├─ 父进程退出 → 子进程收到 SIGTERM/SIGCHLD
    └─ 子进程自动退出（或显式处理信号）
  
  方式 B（备选）：显式通知退出
    ├─ PTOAS 发送 shutdown 信号给 daemon
    ├─ daemon 收到信号 → 清理 → 退出
    └─ PTOAS 退出

3.1.4 代码实现（PTOAS 侧）

// lib/PTO/Transforms/ExpandTileOp.cpp

void ExpandTileOpPass::runOnOperation() {
    std::string socketPath = "/tmp/tilelang_daemon_" + getProcessId() + ".sock";
    //              ↑ 每个进程独立 socket，避免冲突
    
    // Step 1: 检查并启动 daemon
    if (!checkDaemonRunning(socketPath)) {
        // 启动 daemon 子进程
        pid_t pid = fork();
        
        if (pid == 0) {
            // 子进程：运行 daemon
            std::vector<std::string> args = {
                "python3",
                "-m", "tilelang_dsl.daemon",
                "--template-dir", tilelangPath,
                "--socket", socketPath,
            };
            execvp(args[0].c_str(), args);
            exit(1);  // execvp 失败
        }
        
        // 父进程：等待就绪
        daemonPid_ = pid;  // 记录 pid
        
        int maxWaitMs = 5000;
        for (int waited = 0; waited < maxWaitMs; waited += 100) {
            usleep(100000);  // 100ms
            
            if (checkDaemonRunning(socketPath)) {
                break;  // daemon 就绪
            }
        }
        
        if (!checkDaemonRunning(socketPath)) {
            // daemon 启动失败，fallback
            useCLIFallback_ = true;
        }
    }
    
    // Step 2: 连接 daemon
    if (!useCLIFallback_) {
        daemonClient_.connect(socketPath);
    }
    
    // Step 3: 编译
    expandTileOpsInFunction(...);
}

// 退出时清理
void ExpandTileOpPass::~ExpandTileOpPass() {
    if (daemonPid_ > 0) {
        // 通知 daemon 退出（可选）
        kill(daemonPid_, SIGTERM);
        // 或等待子进程自然退出
    }
}

bool checkDaemonRunning(const std::string &socketPath) {
    try {
        Socket sock;
        sock.connect(socketPath);
        sock.close();
        return true;
    } catch (...) {
        return false;
    }
}

---

3.2 Stable Key 计算

3.2.1 为什么需要增强？

专家观点：

"cache key 的设计原则不能只看'哪些字段影响最终代码生成'，还必须看'哪些字段会影响选中哪个 kernel / descriptor'。只要会影响选择结果，就必须进入 key。"

分析：

# kernel.py 的 constraint evaluation

def _constraint_context_for_evaluation(self):
    attrs = {}
    
    for spec in self._parameter_specs:
        # 这些属性可能被 constraint 函数读取
        attrs[f"{spec.name}_shape"] = ...
        attrs[f"{spec.name}_strides"] = ...      ← strides 参与！
        attrs[f"{spec.name}_memory_space"] = ... ← memory_space 参与！
        attrs[f"{spec.name}_config"] = ...
    
    return attrs

# 约束检查
for constraint in self.constraints:
    result = constraint(context_attrs)  ← 可能读取 strides、ms
    if not result:
        raise LookupError("kernel rejected")  ← 模板不合适！

结论：View 的 shape/strides/ms 可能影响 kernel selection → 必须进 stable key

3.2.2 Stable Key 规则

操作数类型	字段	进 Key	原因
Tile	dtype	✅	dtype signature matching
Tile	shape	✅	constraint evaluation
Tile	valid_shape	✅	constraint evaluation
Tile	memory_space	✅	constraint evaluation
Tile	config (blayout, slayout, fractal, pad)	✅	constraint evaluation
View	dtype	✅	dtype signature matching
View	shape	✅ 新增	可能影响 constraint evaluation
View	strides	✅ 新增	可能影响 constraint evaluation
View	memory_space	✅ 新增	可能影响 constraint evaluation
Scalar	dtype	✅	dtype signature matching
Context	context_attrs (round_mode, cmp_mode...)	✅	直接影响 kernel selection

3.2.3 代码实现（daemon 侧）

# tilelang-dsl/stable_key.py

from dataclasses import dataclass
from typing import Any

@dataclass(frozen=True)
class StableKey:
    """稳定的实例化缓存 key"""
    
    target: str                              # 架构："a5", "a3"
    op: str                                  # 算子："tadd", "tmul"
    operand_keys: tuple[OperandStableKey, ...]  # 操作数 keys
    context_key: frozenset[tuple[str, str]]     # context attrs
    
    def to_hash(self) -> str:
        """计算 hash 用于内部索引"""
        parts = [
            self.target,
            self.op,
            *[ok.to_hash_input() for ok in self.operand_keys],
            *[f"{k}={v}" for k, v in sorted(self.context_key)],
        ]
        import hashlib
        return hashlib.sha256("|".join(parts).encode()).hexdigest()[:16]

@dataclass(frozen=True)
class OperandStableKey:
    """操作数稳定 key"""
    
    kind: str        # "tile", "view", "scalar"
    dtype: str       # "f32", "f16", "bf16", "i8", ...
    
    # === Tile 操作数 ===
    shape: tuple[int, ...] | None = None
    valid_shape: tuple[int | None, ...] | None = None  # None 表示动态
    memory_space: str | None = None
    blayout: str | None = None
    slayout: str | None = None
    fractal: int | None = None
    pad: int | None = None
    
    # === View 操作数（增强）===
    view_shape: tuple[int | None, ...] | None = None     # 新增
    view_strides: tuple[int | None, ...] | None = None   # 新增
    view_memory_space: str | None = None                 # 新增
    
    def to_hash_input(self) -> str:
        """生成 hash 输入字符串"""
        if self.kind == "tile":
            vs_str = ",".join(
                "?" if v is None else str(v) for v in self.valid_shape
            ) if self.valid_shape else ""
            s_str = ",".join(str(s) for s in self.shape) if self.shape else ""
            return f"tile:{self.dtype}:{s_str}:{vs_str}:{self.memory_space}:{self.blayout}:{self.slayout}:{self.fractal}:{self.pad}"
        
        elif self.kind == "view":
            s_str = ",".join(
                "?" if s is None else str(s) for s in self.view_shape
            ) if self.view_shape else ""
            st_str = ",".join(
                "?" if st is None else str(st) for st in self.view_strides
            ) if self.view_strides else ""
            return f"view:{self.dtype}:{s_str}:{st_str}:{self.view_memory_space}"
        
        else:  # scalar
            return f"scalar:{self.dtype}"


def compute_stable_key(
    target: str,
    op: str,
    operand_specs: list[dict],
    context_attrs: dict[str, Any] | None = None,
) -> StableKey:
    """从 operand_specs 计算 stable key"""
    
    operand_keys = []
    
    for spec in operand_specs:
        kind = spec.get("kind", "tile")
        dtype = normalize_dtype(spec.get("dtype"))
        
        if kind == "tile":
            operand_keys.append(OperandStableKey(
                kind="tile",
                dtype=dtype,
                shape=tuple(spec.get("shape", [])),
                valid_shape=tuple(spec.get("valid_shape", [])),
                memory_space=spec.get("memory_space", "ub"),
                blayout=normalize_blayout(spec.get("config", {}).get("b_layout")),
                slayout=normalize_slayout(spec.get("config", {}).get("s_layout")),
                fractal=spec.get("config", {}).get("s_fractal_size", 0),
                pad=spec.get("config", {}).get("pad_value", 0),
            ))
        
        elif kind == "view":
            # 增强：shape/strides/memory_space 都进 key
            operand_keys.append(OperandStableKey(
                kind="view",
                dtype=dtype,
                view_shape=tuple(spec.get("shape", [])),
                view_strides=tuple(spec.get("strides", [])),
                view_memory_space=spec.get("memory_space", "gm"),
            ))
        
        else:  # scalar
            operand_keys.append(OperandStableKey(
                kind="scalar",
                dtype=dtype,
            ))
    
    context_key = frozenset(
        (k, str(v)) for k, v in (context_attrs or {}).items()
    )
    
    return StableKey(
        target=target,
        op=op,
        operand_keys=tuple(operand_keys),
        context_key=context_key,
    )

---

3.3 Single-flight 机制

3.3.1 为什么需要？

场景：编译过程中可能并发请求同一个 key

示例：
  
  用户代码有 3 个相同的 pto.tadd(f32, 16x64)
  
  并发处理：
    ├─ 线程1：请求 instantiate("tadd", f32, 16x64)
    ├─ 线程2：请求 instantiate("tadd", f32, 16x64)（同时）
    ├─ 线程3：请求 instantiate("tadd", f32, 16x64)（同时）
    
  无 single-flight：
    ├─ 实例化 3 次（重复工作）
    └─ 资源浪费
    
  有 single-flight：
    ├─ 只实例化 1 次
    ├─ 结果给 3 个请求
    └─ 省资源

3.3.2 代码实现

# tilelang-dsl/daemon_core.py

import asyncio
from typing import dict

class InstanceCache:
    def __init__(self):
        self.cache: dict[StableKey, str] = {}  # key → MLIR text
        self.pending_requests: dict[StableKey, asyncio.Future] = {}  # single-flight
    
    async def instantiate(
        self,
        key: StableKey,
        target: str,
        op: str,
        operand_specs: list,
        context_attrs: dict,
    ) -> str:
        """实例化接口（带 single-flight）"""
        
        # Step 1: 查缓存
        if key in self.cache:
            return self.cache[key]  # 缓存命中
        
        # Step 2: Single-flight 检查
        if key in self.pending_requests:
            # 有其他请求正在实例化同样的 key
            # 等待它的结果（不重复实例化）
            return await self.pending_requests[key]
        
        # Step 3: 创建 future（标记"正在做"）
        future = asyncio.Future()
        self.pending_requests[key] = future
        
        try:
            # Step 4: 实例化
            mlir_text = await self._do_instantiate(
                target, op, operand_specs, context_attrs
            )
            
            # Step 5: 存缓存
            self.cache[key] = mlir_text
            
            # Step 6: 完成 future（通知所有等待者）
            future.set_result(mlir_text)
            
            return mlir_text
        
        except Exception as e:
            # 失败：通知所有等待者
            future.set_exception(e)
            raise
        
        finally:
            # 清理 pending 状态
            del self.pending_requests[key]
    
    async def _do_instantiate(self, target, op, operand_specs, context_attrs):
        """实际实例化过程"""
        
        # 1. Kernel selection
        descriptor = self.registry.select_kernel(
            target, op, operand_specs, context_attrs
        )
        
        # 2. Specialize
        tile_specializations = {}
        for param, spec in zip(descriptor.parameters, operand_specs):
            if param.kind == "tile":
                tile_specializations[param.name] = TileSpecialization(
                    shape=spec.get("shape"),
                    memory_space=spec.get("memory_space"),
                    config=spec.get("config"),
                    valid_shape=spec.get("valid_shape"),
                )
        
        specialized = descriptor.specialize(**tile_specializations)
        
        # 3. Materialize
        mlir_text = specialized.mlir_text()
        
        return mlir_text

---

3.4 缓存管理

3.4.1 缓存位置

内存（RAM）：daemon 进程的 Python dict

理由：

• 访问快（微秒级）
• daemon 作为子进程，内存占用可控
• PTOAS 退出时自动清理

3.4.2 缓存大小

配置：

MAX_ENTRIES = 1000  # 最多缓存 1000 个实例

内存估算：

单实例大小：
  ├─ StableKey: ~几百字节
  ├─ MLIR text: ~5KB ~ 50KB
  └─ 总计：~10KB（平均）

总内存：
  ├─ 1000 × 10KB = ~10MB
  └─ 服务器级内存可忽略

3.4.3 淘汰机制

FIFO（先进先出）：

def store(self, key: StableKey, mlir_text: str):
    # 检查是否满了
    if len(self.cache) >= self.max_entries:
        # FIFO：扔掉最早缓存的
        oldest_key = next(iter(self.cache))
        del self.cache[oldest_key]
    
    # 存入缓存
    self.cache[key] = mlir_text

为什么不持久化：

• 缓存只在单次编译过程中有效
• PTOAS 退出时 daemon 也退出
• 不需要跨进程持久化

---

3.5 RPC 通信

3.5.1 通信方式

Unix Socket（Linux/Mac）：

socket 路径：
  /tmp/tilelang_daemon_<pid>.sock
  
  每个 PTOAS 进程独立 socket
  避免多个编译任务冲突

Named Pipe（Windows）：

pipe 名称：
  \\.\pipe\tilelang_daemon_<pid>
  
  Windows 平台备选方案

3.5.2 通信协议

请求格式（JSON）：

{
    "target": "a5",
    "op": "tadd",
    "operand_specs": [
        {
            "kind": "tile",
            "dtype": "f32",
            "shape": [16, 64],
            "valid_shape": [16, 64],
            "memory_space": "ub",
            "config": {
                "b_layout": "row_major",
                "s_layout": "none_box",
                "s_fractal_size": 512,
                "pad_value": 0
            }
        },
        {
            "kind": "tile",
            "dtype": "f32",
            ...
        },
        {
            "kind": "tile",
            "dtype": "f32",
            ...
        }
    ],
    "context_attrs": {}
}

响应格式（JSON）：

{
    "status": "success",
    "mlir": "module {\n  func.func @__pto_tilelang_...\n    ...\n  }\n}"
}

失败响应：

{
    "status": "error",
    "error": "kernel selection failed",
    "message": "..."
}

3.5.3 代码实现

daemon 侧（Python）：

# tilelang-dsl/daemon_server.py

import asyncio
import json
import socket

class DaemonServer:
    def __init__(self, socket_path: str, cache: InstanceCache):
        self.socket_path = socket_path
        self.cache = cache
    
    async def handle_request(self, reader, writer):
        """处理单个请求"""
        
        # 读取请求
        data = await reader.read()
        request = json.loads(data.decode())
        
        # 提取参数
        target = request["target"]
        op = request["op"]
        operand_specs = request["operand_specs"]
        context_attrs = request.get("context_attrs", {})
        
        # 计算 stable key
        key = compute_stable_key(target, op, operand_specs, context_attrs)
        
        # 实例化（带 single-flight）
        try:
            mlir_text = await self.cache.instantiate(
                key, target, op, operand_specs, context_attrs
            )
            
            response = {
                "status": "success",
                "mlir": mlir_text,
            }
        
        except Exception as e:
            response = {
                "status": "error",
                "error": str(e),
            }
        
        # 发送响应
        writer.write(json.dumps(response).encode())
        await writer.drain()
        writer.close()
    
    async def run(self):
        """启动服务器"""
        server = await asyncio.start_server(
            self.handle_request,
            path=self.socket_path,
        )
        
        # 运行直到父进程退出
        async with server:
            await server.serve_forever()

PTOAS 侧（C++）：

// lib/PTO/Daemon/TileLangDaemonClient.cpp

#include <sys/socket.h>
#include <sys/un.h>
#include <json/json.h>

class TileLangDaemonClient {
public:
    bool connect(const std::string &socketPath) {
        sockfd_ = socket(AF_UNIX, SOCK_STREAM, 0);
        
        struct sockaddr_un addr;
        addr.sun_family = AF_UNIX;
        strcpy(addr.sun_path, socketPath.c_str());
        
        return ::connect(sockfd_, (struct sockaddr*)&addr, sizeof(addr)) == 0;
    }
    
    std::optional<std::string> instantiate(
        const std::string &target,
        const std::string &op,
        const std::vector<OperandSpec> &operand_specs,
        const std::map<std::string, std::string> &context_attrs
    ) {
        // 构造 JSON 请求
        Json::Value request;
        request["target"] = target;
        request["op"] = op;
        request["operand_specs"] = buildOperandSpecsJson(operand_specs);
        request["context_attrs"] = buildContextAttrsJson(context_attrs);
        
        // 发送请求
        std::string requestStr = Json::writeString(request);
        send(sockfd_, requestStr.c_str(), requestStr.size(), 0);
        
        // 接收响应
        char buffer[1024 * 1024];  // 1MB buffer
        int n = recv(sockfd_, buffer, sizeof(buffer), 0);
        
        if (n <= 0) return std::nullopt;
        
        // 解析响应
        Json::Value response;
        Json::Reader reader;
        reader.parse(buffer, buffer + n, response);
        
        if (response["status"] == "success") {
            return response["mlir"].asString();
        }
        
        return std::nullopt;
    }
    
private:
    int sockfd_;
};

---

3.6 Fallback 机制

3.6.1 为什么需要？

daemon 可能失败的情况：

• daemon 启动失败
• daemon 崩溃
• socket 连接失败
• 请求超时

3.6.2 Fallback 流程

// lib/PTO/Transforms/ExpandTileOp.cpp

func::FuncOp ExpandState::instantiateTileOp(...) {
    // 尝试 daemon
    if (!useCLIFallback_) {
        auto mlirText = daemonClient_.instantiate(...);
        
        if (mlirText) {
            // 成功：解析、clone
            return cloneMlirFunction(*mlirText, mod, ctx);
        }
        
        // 失败：标记使用 fallback
        useCLIFallback_ = true;
        llvm::errs() << "Daemon failed, fallback to CLI helper\n";
    }
    
    // Fallback：使用 CLI helper（旧方案）
    return invokeTilelangDSL_CLI(key, mod, ctx);
}

func::FuncOp invokeTilelangDSL_CLI(...) {
    // 启动 Python 子进程（每次 tile op）
    std::string cmd = "python3 -m tilelang_dsl.expand_helper ...";
    
    // 执行，读取输出
    auto mlirText = executeCommand(cmd);
    
    // 解析、clone
    return cloneMlirFunction(mlirText, mod, ctx);
}

---

四、实现计划

4.1 文件改动清单

新增文件（框架侧）

tilelang-dsl/python/tilelang_dsl/
├── daemon.py              ← daemon 主程序（生命周期管理）
├── daemon_core.py         ← 核心逻辑（缓存、单飞）
├── stable_key.py          ← stable key 计算
├── daemon_server.py       ← RPC server

新增文件（编译器侧）

lib/PTO/Daemon/
├── TileLangDaemonClient.h    ← daemon client 接口
├── TileLangDaemonClient.cpp  ← daemon client 实现
├── CMakeLists.txt            ← 编译配置

修改文件

lib/PTO/Transforms/
├── ExpandTileOp.cpp          ← 改用 daemon client

---

4.2 Phase 1: TileLang Daemon 实现

任务：

1. daemon.py：

   • 主程序入口
   • 生命周期管理（启动、运行、退出）
   • 作为子进程运行

2. daemon_core.py：

   • InstanceCache 类（缓存、单飞）
   • instantiate 接口
   • FIFO 淘汰

3. stable_key.py：

   • StableKey 类
   • OperandStableKey 类（增强 View 参数）
   • compute_stable_key 函数

4. daemon_server.py：

   • RPC server
   • socket 监听
   • 请求处理

验收标准：

• ✅ daemon 可启动（作为子进程）
• ✅ 监听 socket，接收请求
• ✅ 同样请求第二次命中缓存
• ✅ 同 key 并发，single-flight 去重
• ✅ FIFO 淘汰正常工作
• ✅ 父进程退出时自动退出

---

4.3 Phase 2: PTOAS Client 实现

任务：

1. TileLangDaemonClient.h：

   • 接口定义
   • connect 方法
   • instantiate 方法

2. TileLangDaemonClient.cpp：

   • socket 连接实现
   • JSON 请求构造
   • JSON 响应解析
   • Unix socket 支持

3. ExpandTileOp.cpp：

   • 启动 daemon 子进程
   • 使用 daemon client
   • fallback 到 CLI helper
   • pass-local dedup

验收标准：

• ✅ PTOAS 可启动 daemon 子进程
• ✅ 连接 daemon，发送请求
• ✅ 成功返回 MLIR 文本
• ✅ 解析、clone 正常
• ✅ daemon 不可用，fallback 正常
• ✅ PTOAS 不感知 cache key 计算

---

4.4 Phase 3: 测试

测试用例：

1. 单次编译正确性：

   • 输入：包含 tile ops 的 MLIR
   • 输出：展开后的向量 IR
   • 验证：与当前实现结果一致

2. 单次编译缓存命中：

   • 输入：包含 3 个相同 tile op
   • 验证：第二次、第三次命中缓存
   • 验证：daemon 统计：hits > 0

3. 并发编译 single-flight：

   • 输入：并发请求相同 tile op
   • 验证：只实例化 1 次
   • 验证：所有请求拿到相同结果

4. daemon 生命周期：

   • 验证：PTOAS 启动时 daemon 启动
   • 验证：PTOAS 退出时 daemon 退出
   • 验证：进程内缓存有效

5. Fallback：

   • 场景：daemon 启动失败
   • 验证：使用 CLI helper
   • 验证：编译成功

---

4.5 Phase 4: 文档

文档内容：

1. Stable Key 协议：

   • 哪些字段进 key
   • 为什么这样设计
   • View 参数增强原因

2. Daemon 生命周期管理：

   • 启动流程
   • 运行流程
   • 退出流程
   • 绑定 PTOAS 主进程

3. 缓存机制：

   • 进程内缓存
   • FIFO 淘汰
   • Single-flight

4. Kernel selection 影响因素分析：

   • constraint evaluation 可能读取哪些参数
   • 为什么影响 stable key

---

五、验收标准

5.1 功能验收

功能	验收标准
daemon 启动	PTOAS 自动启动 daemon 子进程
daemon 连接	PTOAS 成功连接 daemon socket
缓存命中	同 tile op 第二次命中缓存
single-flight	同 key 并发只实例化 1 次
fallback	daemon 失败时使用 CLI helper
生命周期	PTOAS 退出时 daemon 退出
stable key	View 的 shape/strides/ms 进 key
编译正确性	展开结果与当前实现一致

---

5.2 性能验收

性能指标	当前	目标
单 tile op 实例化	~2秒（启动 Python）	~1秒（首次）
缓存命中	无缓存	~0.1秒（缓存命中）
启动开销	每次 tile op	一次编译一次
内存占用	pass 内	~10MB（进程内）

---

5.3 边界验收

边界	验收标准
PTOAS 不感知	不知道 cache key 计算、缓存实现
单一接口	只调用 instantiate()
框架独立演进	daemon 内部改进，PTOAS 不改

---

六、专家反馈总结

6.1 第一轮反馈

反馈	采纳	修订
不用 CLI helper，用 daemon 常驻	✅	daemon 常驻（后修订为绑定生命周期）
View key 增强	✅	shape/strides/ms 进 key
不用 warm-up	✅	去掉预热 API
不用 artifact	✅	去掉持久化机制
PTOAS 不感知缓存	✅	单一 RPC 接口
单一 RPC 接口	✅	只调用 instantiate()

---

6.2 第二轮反馈

反馈	采纳	修订
daemon 退出绑定 PTOAS	✅	作为子进程，跟随退出
单次编译缓存够用	✅	进程内缓存，不跨进程
去掉 Version Tracking	✅	删除文件 hash 版本追踪

---

6.3 方案演进总结

V0（当前）：
  ├─ 每次 tile op 启动 Python 子进程
  ├─ pass 内 DenseMap 缓存
  └─ 无跨请求复用
  
V1（初稿）：
  ├─ 框架缓存 + PTOAS 消费 artifact
  ├─ 跨进程持久化缓存
  ├─ PTOAS 计算 hash、管理缓存
  └─ 边界扩散
  
V2（第一轮修订）：
  ├─ daemon 常驻进程
  ├─ 跨进程缓存复用
  ├─ View key 增强
  ├─ Version Tracking
  └─ PTOAS 不感知
  
V2.1（最终）：
  ├─ daemon 作为子进程，绑定生命周期 ✅
  ├─ 进程内缓存（单次编译有效） ✅
  ├─ View key 增强 ✅
  ├─ Single-flight ✅
  ├─ 去掉 Version Tracking ✅
  ├─ PTOAS 不感知 ✅
  └─ 单一 RPC 接口 ✅

---

七、附录

7.1 关键代码位置

功能	当前位置	新增/修改位置
当前缓存	ExpandTileOp.cpp:450-452	—
当前 SpecKey	ExpandTileOp.cpp:128-164	—
当前实例化	ExpandTileOp.cpp:575-779	—
daemon 主程序	—	daemon.py（新增）
daemon 核心	—	daemon_core.py（新增）
stable key	—	stable_key.py（新增）
RPC server	—	daemon_server.py（新增）
daemon client	—	TileLangDaemonClient.cpp（新增）
ExpandTileOp 改造	ExpandTileOp.cpp	ExpandTileOp.cpp（修改）
kernel selection	kernel.py	kernel.py（分析影响）
constraint evaluation	kernel.py:1148-1256	—（分析影响）

---

7.2 参考资料

资料	位置
Issue #210	#210
当前设计文档	docs/designs/ptoas-tileop-expand-design.md
kernel.py	tilelang-dsl/python/tilelang_dsl/kernel.py
expand_helper.py	tilelang-dsl/python/tilelang_dsl/expand_helper.py
模板目录	lib/TileOps/*.py

---

7.3 配置参数

参数	默认值	说明
MAX_ENTRIES	1000	缓存上限
socketPath	/tmp/tilelang_daemon_<pid>.sock	daemon socket 路径
maxWaitSeconds	5	daemon 启动等待时间
timeoutSeconds	30	RPC 请求超时时间

---

7.4 术语表

术语	说明
daemon	守护进程，这里指作为子进程运行的 TileLang 服务进程
stable key	稳定的缓存键，包含所有影响 kernel selection 的参数
single-flight	并发请求去重机制，同 key 只实例化一次
instance cache	实例化缓存，存储 materialized MLIR
kernel selection	从多个模板中选择合适的模板
constraint evaluation	约束检查，判断模板是否适合当前参数
materialized MLIR	实例化后的 MLIR 文本
View	外部数据的视图（切片），参数可能动态
Tile	数据块，参数静态已知
FIFO	先进先出淘汰策略

---

八、总结

核心设计：

• daemon 作为 PTOAS 子进程运行，生命周期绑定
• 单次编译过程中提供缓存服务
• PTOAS 通过单一 RPC 接口请求实例化，不感知缓存细节
• View 参数增强：shape/strides/ms 进 stable key
• Single-flight 去重并发请求
• 不需要 Version Tracking（模板不变更）

预期收益：

• 启动开销减少（一次编译启动一次 daemon）
• 进程内缓存有效（避免重复实例化）
• 实现简化（去掉版本追踪）
• 边界清晰（PTOAS 不感知）

实现计划：

• Phase 1: daemon 实现（框架侧）
• Phase 2: client 实现（编译器侧）
• Phase 3: 测试
• Phase 4: 文档

验收标准：

• daemon 自动启动、退出绑定 PTOAS
• 缓存命中、single-flight 正常
• fallback 可用
• 编译正确性一致
• PTOAS 不感知缓存细节