diff --git a/aikg/tools/gen_docs/DSL_DOCS_GENERATION_TASK.md b/aikg/tools/gen_docs/DSL_DOCS_GENERATION_TASK.md new file mode 100644 index 0000000000000000000000000000000000000000..f0a079079623fd41d2306680c967bbf41400b4b4 --- /dev/null +++ b/aikg/tools/gen_docs/DSL_DOCS_GENERATION_TASK.md @@ -0,0 +1,288 @@ +# DSL文档生成任务 + +--- +## 【填写以下信息】 + +**DSL名称**: _填写DSL名_ (例如: tilelang, taichi, mojo) + +**目标后端硬件**: _填写硬件_ (例如: CUDA GPU, Ascend NPU, CPU) + +**DSL代码仓路径**: _填写路径_ (例如: /path/to/your_dsl, https://github.com/xxx/your_dsl) + +--- + +## 你的任务 + +为上述DSL生成符合AIKG规范的完整文档集。 + +## 第一步:理解AIKG的文档规范 + +先阅读以下内容,理解AIKG对DSL文档的要求: + +### 1. 阅读DDI规范 +打开 `aikg/docs/CN/DocDrivenIntegration.md`,了解文档驱动式接入的基本要求。 + +### 2. 重点学习triton_docs(最重要) +打开 `aikg/python/ai_kernel_generator/resources/docs/triton_docs/`,这是最完整的参考示例: +- 阅读 `basic_docs.md` - 理解基础教程的写法 +- 阅读 `api/api.md` - 理解API参考的组织方式 +- 阅读 `suggestion_docs.md` - 理解优化建议的内容 +- 查看 `examples/` 目录 - 理解示例代码的命名和格式 + +### 3. 对比其他DSL文档 +快速浏览 `cuda_docs/`, `cpu_docs/`, `swft_docs/`,了解不同DSL的文档差异。 + +### 4. 查看配置使用方式 +打开 `aikg/python/ai_kernel_generator/config/` 目录下的配置文件(如 `default_triton_config.yaml`),看看 `docs_dir` 字段如何配置。 + +## 第二步:探索目标DSL仓库并发现 API + +现在去探索上面填写的DSL代码仓,找到以下关键信息: + +### 1. 核心概念和设计目标 +- 打开 `README.md` 和主要文档,提取DSL的核心概念 +- 找到关键术语(如内核定义方式、执行模型、内存层次等) + +### 2. 发现和列举 API(重点!) + +**目标**:找到 DSL 的所有公开 API,确保不遗漏 + +#### 2.1 找到 API 定义的位置 + +根据 DSL 语言类型,采用不同策略: + +**Python DSL**: +- 找主入口:`__init__.py` 的 `__all__` 列表或 import 语句 +- 找核心模块:`language.py`, `lang.py`, `ops.py`, `api.py` 等 +- 看目录结构:`kernel/`, `ops/`, `primitives/` 等目录 + +**C++/C DSL**: +- 找头文件:`include/` 目录下的 `.h` 或 `.hpp` 文件 +- 看 namespace 和 class 定义 +- 找宏定义:`#define` 开头的 API + +**其他语言**: +- 查看主模块的导出函数 +- 看官方文档的 API 章节 + +#### 2.2 系统化列举 API + +**方法 1:从代码入口追溯** +``` +1. 打开主入口文件(如 `__init__.py`) +2. 列出所有 from xxx import yyy +3. 逐个打开被 import 的模块 +4. 列出所有不以 _ 开头的函数/类 +``` + +**方法 2:从示例代码反推** +``` +1. 找到典型示例(如 vector_add, matmul) +2. 列出示例中用到的所有 API +3. 按功能分类(内存分配、数据传输、计算等) +4. 检查是否有遗漏的类别 +``` + +**方法 3:从文档中提取** +``` +1. 查找是否有 API Reference 章节 +2. 提取所有函数名和类名 +3. 交叉验证代码中是否都能找到 +``` + +#### 2.3 API 分类和组织 + +将发现的 API 按功能分类(建议分类): +- **内核装饰器/定义**:如 `@jit`, `@kernel`, `prim_func` +- **内存分配**:如 `alloc_shared`, `alloc_fragment`, `alloc_local` +- **数据传输**:如 `copy`, `load`, `store`, `fill`, `clear` +- **计算操作**:如 `gemm`, `dot`, `matmul`, `reduce_*` +- **执行控制**:如 `Kernel`, `Parallel`, `Pipelined`, `Grid` +- **数学函数**:如 `exp`, `log`, `sqrt`, `max`, `min` +- **原子操作**:如 `atomic_add`, `atomic_max` +- **其他**:无法归类的 + +#### 2.4 记录 API 清单 + +对每个 API 记录: +- **名称**:如 `alloc_shared` +- **签名**:如 `alloc_shared(shape, dtype)` +- **所属模块**:如 `tilelang.language` +- **功能简述**:1 句话说明(可从 docstring 提取) +- **是否核心**:装饰器、alloc、copy、gemm 等是核心 + +**输出示例**: +``` +内存分配类 (6个): +- alloc_shared(shape, dtype) - 分配共享内存缓冲区 +- alloc_fragment(shape, dtype) - 分配寄存器缓冲区 +- alloc_local(shape, dtype) - 分配局部内存 +- alloc_reducer(shape, dtype) - 分配归约器 +- alloc_barrier() - 分配同步屏障 +- alloc_var(dtype, scope) - 分配变量 + +数据传输类 (5个): +- copy(src, dst) - 数据拷贝 +- load(ptr, mask) - 加载数据 +- store(ptr, value, mask) - 存储数据 +- fill(buffer, value) - 填充值 +- clear(buffer) - 清零 + +... 更多 +``` + +### 3. 示例代码 +- 在 `examples/`, `tutorials/`, `tests/` 等目录中找示例 +- 重点关注:vector_add, matmul, softmax, layer_norm 等典型算子 +- 理解示例代码的结构和模式 + +### 4. 优化建议和最佳实践 +- 查找性能优化相关的文档或注释 +- 找到常见错误、注意事项、调试技巧 + +## 第三步:生成AIKG格式的文档 + +在 `aikg/python/ai_kernel_generator/resources/docs/{dsl_name}_docs/` 目录下创建以下文档: + +### 必需:`basic_docs.md` +参考 `triton_docs/basic_docs.md` 的结构,编写包含: +- 核心概念(内核、网格、内存层次等) +- 标准内核结构(5步模式) +- 编程模式(向量操作、归约、矩阵运算) +- 边界处理示例 + +### 强烈推荐:`api/api.md` + +**生成要求**: +- **覆盖率目标**:至少覆盖 **80%** 的公开 API +- **必须覆盖所有核心 API**(装饰器、内存分配、数据传输、核心计算) + +参考 `triton_docs/api/api.md` 的结构,按你在第二步整理的分类组织: + +```markdown +# {DSL名称} API 参考手册 + +## 1. 内核装饰器 +[列出所有装饰器 API,如 @jit, @kernel] + +## 2. 内存分配 API +[列出所有 alloc_* API] + +## 3. 数据传输 API +[列出 copy, load, store 等] + +## 4. 计算操作 API +[列出 gemm, reduce_* 等] + +## 5. 数学函数 API +[列出 exp, log, sqrt 等] + +## 6. 其他 API... +``` + +**对每个 API,包含**: +- 函数签名(从清单复制) +- 参数说明 +- 返回值 +- 简短用途说明(1-2句) +- 简单示例(可选,核心 API 必须有) + +**压缩策略**: +- **核心 API**(~30个):详细说明 + 代码示例 +- **常用 API**(~50个):签名 + 简短说明 +- **其他 API**(~50+个):仅列举签名,简述用途 + +### 强烈推荐:`suggestion_docs.md` +参考 `triton_docs/suggestion_docs.md`,提供: +- 性能优化策略(块大小选择、内存访问优化) +- 数值稳定性技巧 +- API使用限制和替代方案 +- 常见错误速查表 + +### 强烈推荐:`examples/` 目录 +创建示例代码,文件命名格式:`{framework}_*.py` +- `torch_vector_add.py` - 向量加法 +- `torch_matmul.py` - 矩阵乘法 +- `torch_softmax.py` - Softmax +- `numpy_*.py` - NumPy版本(如需要) + +示例代码必须包含完整的kernel定义和启动函数。 + +## 第四步:校验文档(结构 + API 覆盖率) + +### 4.1 结构校验 + +运行校验脚本: +```bash +cd aikg/tools/gen_docs +python validate_dsl_docs.py ../../python/ai_kernel_generator/resources/docs/{dsl_name}_docs +``` + +确保: +- ✅ `basic_docs.md` 存在 +- ⚠️ `api/api.md` 存在(强烈推荐) +- ⚠️ `suggestion_docs.md` 存在(强烈推荐) +- ⚠️ `examples/` 目录存在(强烈推荐) + +### 4.2 API 覆盖率自查(重要!) + +**验证方法**:对照第二步整理的 API 清单,逐一检查是否在文档中被提及 + +#### 步骤 1:从文档中提取提到的 API + +在生成的 `api/api.md` 中查找: +- 统计有多少个 API 函数/类被列出 +- 按分类统计:每个类别覆盖了多少个 + +#### 步骤 2:与 API 清单对比 + +对照第二步的清单,检查: + +``` +你的清单中有 X 个 API,分为 Y 个类别 + +文档中覆盖了: +- 内存分配类:6/6 个 ✅ +- 数据传输类:4/5 个 ⚠️ 缺失:xxx +- 计算操作类:8/10 个 ⚠️ 缺失:xxx, yyy +- 数学函数类:15/30 个 ❌ 覆盖率过低 + +总覆盖率:XX/YY = ZZ% +``` + +#### 步骤 3:识别关键遗漏 + +检查是否遗漏了关键 API: +- [ ] 所有装饰器(@jit, @kernel 等)都有文档? +- [ ] 所有 alloc_* 函数都有文档? +- [ ] copy, load, store 等核心传输 API 都有文档? +- [ ] gemm, matmul, reduce 等核心计算 API 都有文档? + +**如果覆盖率 < 80% 或关键 API 有遗漏**: +1. 回到第三步,补充缺失的 API +2. 重新验证 + +#### 步骤 4:检查是否有拼写错误 + +检查文档中提到的 API 名称: +- 是否都能在代码仓库中找到? +- 是否存在打错的名字(如把 `alloc_shared` 写成 `alloc_share`)? + +**自查清单**: +``` +✅ basic_docs.md 存在且内容完整 +✅ api/api.md 存在 +✅ API 覆盖率 ≥ 80% +✅ 所有核心 API 都有文档 +✅ 所有 API 名称拼写正确 +✅ suggestion_docs.md 存在 +✅ examples/ 目录存在,至少 3-4 个示例 +``` + +## 质量要求 + +- 文档结构必须与 `triton_docs` 保持一致 +- 代码示例必须可运行且符合命名规范 +- 必须涵盖DSL的核心特性和常用模式 +- 语言清晰、示例完整 + diff --git a/aikg/tools/gen_docs/README.md b/aikg/tools/gen_docs/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ce3cdad315a63b6ac19efedb66851c3fca63ec5d --- /dev/null +++ b/aikg/tools/gen_docs/README.md @@ -0,0 +1,104 @@ +# DSL文档生成工具 + +## 这个工具是做什么的 + +为AIKG生成新DSL的标准文档,让LLM agent自动完成文档接入工作。 + +## 目录文件 + +- **`DSL_DOCS_GENERATION_TASK.md`** - 任务文档(给LLM agent执行) +- **`validate_dsl_docs.py`** - 校验脚本(检查生成的文档结构) + +## 怎么使用 + +### 第1步:填写DSL信息 + +打开 `DSL_DOCS_GENERATION_TASK.md`,在开头填写3个信息: + +```markdown +**DSL名称**: tilelang +**目标后端硬件**: CUDA GPU +**DSL代码仓路径**: /path/to/tilelang +``` + +### 第2步:让LLM agent执行任务 + +把 `DSL_DOCS_GENERATION_TASK.md` 给你的LLM agent(Claude、Cursor等)读。 + +它会按照文档指引: +1. 学习AIKG的文档规范(triton_docs) +2. **系统化发现和列举 DSL 的所有 API** +3. 生成符合格式的文档(**要求 API 覆盖率 ≥ 80%**) +4. **自我验证**:检查是否遗漏关键 API + +### 第3步:校验结果 + +```bash +cd aikg/tools/gen_docs +python validate_dsl_docs.py ../../python/ai_kernel_generator/resources/docs/{dsl_name}_docs +``` + +检查: +- ✅ `basic_docs.md` 是否存在(必需) +- ⚠️ `api/api.md` 是否存在(推荐) +- ⚠️ `suggestion_docs.md` 是否存在(推荐) +- ⚠️ `examples/` 是否存在(推荐) + +### 第4步:在AIKG中使用 + +打开 `aikg/python/ai_kernel_generator/config/your_config.yaml`,添加: + +```yaml +docs_dir: + designer: "resources/docs/{dsl_name}_docs" + coder: "resources/docs/{dsl_name}_docs" +``` + +## 完整示例(以tilelang为例) + +```bash +cd aikg/tools/gen_docs + +# 1. 填写任务信息 +vim DSL_DOCS_GENERATION_TASK.md +# 填写 DSL 名称、硬件、路径 + +# 2. 给 LLM 读任务文档 +# LLM 会: +# - 发现和列举所有 API +# - 生成文档(覆盖率 ≥ 80%) +# - 自我验证是否有遗漏 + +# 3. 校验结构 +python validate_dsl_docs.py ../../python/ai_kernel_generator/resources/docs/tilelang_docs + +# 4. 配置使用 +# 编辑 config/your_config.yaml,加上 docs_dir 配置 +``` + +## 核心改进 + +相比之前的版本,新增了关于 **API 发现和覆盖率验证** 的详细指引: + +### API 发现方法 +- **Python DSL**:从 `__init__.py` 入口追溯 +- **C++ DSL**:查找 `include/` 头文件 +- **示例反推**:从典型示例中提取使用的 API + +### API 文档要求 +- **覆盖率目标**:≥ 80% 的公开 API +- **核心 API 必须 100% 覆盖**(装饰器、内存、传输、计算) +- **分层描述**:核心详细、常用简洁、其他列举 + +### 自我验证方法 +- 对照清单统计覆盖率 +- 检查关键 API 是否遗漏 +- 验证 API 名称拼写 + +## 注意事项 + +- 文档必须放在 `aikg/python/ai_kernel_generator/resources/docs/` 下 +- 目录命名格式:`{dsl_name}_docs`(如 `triton_docs`) +- 示例代码命名格式:`{framework}_*.py`(如 `torch_vector_add.py`) +- **API 覆盖率要求:≥ 80%** +- 参考 `triton_docs` 是最完整的示例 diff --git a/aikg/tools/gen_docs/validate_dsl_docs.py b/aikg/tools/gen_docs/validate_dsl_docs.py new file mode 100755 index 0000000000000000000000000000000000000000..2d4e01784298d64fac5ff652c59d5333c479ef57 --- /dev/null +++ b/aikg/tools/gen_docs/validate_dsl_docs.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +DSL文档结构校验脚本 + +功能:检查DSL文档目录是否包含必需的文件 +""" + +import os +import sys +from pathlib import Path +import argparse + + +class DSLDocsValidator: + """DSL文档校验器""" + + # 必需文件 + REQUIRED_FILES = [ + "basic_docs.md", + ] + + # 推荐文件/目录 + RECOMMENDED_ITEMS = [ + "api/api.md", + "suggestion_docs.md", + "examples/", + ] + + def __init__(self, docs_dir: str): + self.docs_dir = Path(docs_dir) + self.errors = [] + self.warnings = [] + + def validate(self) -> bool: + """执行校验""" + if not self.docs_dir.exists(): + self.errors.append(f"文档目录不存在: {self.docs_dir}") + return False + + if not self.docs_dir.is_dir(): + self.errors.append(f"路径不是目录: {self.docs_dir}") + return False + + # 检查必需文件 + self._check_required_files() + + # 检查推荐文件 + self._check_recommended_items() + + return len(self.errors) == 0 + + def _check_required_files(self): + """检查必需文件""" + for file_path in self.REQUIRED_FILES: + full_path = self.docs_dir / file_path + if not full_path.exists(): + self.errors.append(f"缺少必需文件: {file_path}") + + def _check_recommended_items(self): + """检查推荐文件/目录""" + for item_path in self.RECOMMENDED_ITEMS: + full_path = self.docs_dir / item_path + if not full_path.exists(): + self.warnings.append(f"缺少推荐项: {item_path}") + + def print_report(self): + """打印校验报告""" + print(f"\n{'='*60}") + print(f"DSL文档校验报告") + print(f"{'='*60}") + print(f"文档目录: {self.docs_dir}") + print(f"{'-'*60}") + + if not self.errors and not self.warnings: + print("✅ 校验通过!文档结构完整。") + else: + if self.errors: + print(f"\n❌ 错误 ({len(self.errors)}):") + for error in self.errors: + print(f" - {error}") + + if self.warnings: + print(f"\n⚠️ 警告 ({len(self.warnings)}):") + for warning in self.warnings: + print(f" - {warning}") + + print(f"{'='*60}\n") + + +def main(): + parser = argparse.ArgumentParser( + description="校验DSL文档目录结构", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +示例: + python validate_dsl_docs.py ../python/ai_kernel_generator/resources/docs/triton_docs + python validate_dsl_docs.py --docs-dir ./my_new_dsl_docs + """ + ) + + parser.add_argument( + 'docs_dir', + nargs='?', + help='DSL文档目录路径' + ) + + parser.add_argument( + '--docs-dir', + dest='docs_dir_alt', + help='DSL文档目录路径(替代位置参数)' + ) + + args = parser.parse_args() + + # 获取文档目录路径 + docs_dir = args.docs_dir or args.docs_dir_alt + + if not docs_dir: + parser.print_help() + sys.exit(1) + + # 执行校验 + validator = DSLDocsValidator(docs_dir) + success = validator.validate() + validator.print_report() + + # 返回退出码 + sys.exit(0 if success else 1) + + +if __name__ == "__main__": + main() +