{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# 核函数调用与验证\n",
"\n",
"## 概述\n",
"\n",
"本节介绍如何从Host侧调用pyasc核函数（Launch），以及如何编写完整的算子验证程序。同时将pyasc的Host侧调用方式与Ascend C进行对比，帮助已有Ascend C经验的开发者快速迁移。\n",
"\n",
"### 学习目标\n",
"\n",
"1. 掌握Launch函数的实现方法；\n",
"2. 掌握内核调用符`kernel[核数, stream]()`的使用；\n",
"3. 能编写完整的算子验证程序（数据生成→核函数调用→结果校验）；\n",
"4. 了解手动版与框架版Add样例的区别。\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 环境初始化\n",
"!mkdir -p Sources/02.06\n",
"\n",
"import os, subprocess\n",
"env = subprocess.check_output(\"bash -l -c 'source $ASCEND_TOOLKIT_HOME/set_env.sh && env'\", shell=True, text=True)\n",
"for line in env.splitlines():\n",
" if \"=\" in line: os.environ.__setitem__(*line.split(\"=\", 1))\n",
"print(\"环境初始化完成\")"
]
},
{
"cell_type": "markdown",
"id": "a979cee4",
"metadata": {},
"source": [
"---\n",
"# 1. Launch函数\n",
"\n",
"Launch函数是Host侧调用核函数的入口，负责数据准备、核函数调用和结果返回。以Add算子为例：\n",
"\n",
"```python\n",
"import asc.lib.runtime as rt\n",
"\n",
"def vadd_launch(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:\n",
" z = torch.zeros_like(x)\n",
" total_length = z.numel()\n",
" block_length = total_length // USE_CORE_NUM\n",
" vadd_kernel[USE_CORE_NUM, rt.current_stream()](x, y, z, block_length)\n",
" return z\n",
"```\n",
"\n",
"<table style=\"margin-left: 0; margin-right: auto; width: auto; border: 1px solid #ddd;\">\n",
" <thead>\n",
" <tr>\n",
" <th align=\"left\">步骤</th>\n",
" <th align=\"left\">操作</th>\n",
" <th align=\"left\">关键代码</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr><td>1</td><td>创建输出Tensor</td><td><code>z = torch.zeros_like(x)</code></td></tr>\n",
" <tr><td>2</td><td>计算数据分片</td><td><code>block_length = total_length // USE_CORE_NUM</code></td></tr>\n",
" <tr><td>3</td><td>调用核函数</td><td><code>vadd_kernel[核数, stream](args)</code></td></tr>\n",
" <tr><td>4</td><td>返回结果</td><td><code>return z</code></td></tr>\n",
" </tbody>\n",
"</table>\n",
"\n",
"<img src=\"./images/kernel_call_flow.png\" alt=\"核函数调用流程\" width=\"700px\">\n"
]
},
{
"cell_type": "markdown",
"id": "7a5dffdb",
"metadata": {},
"source": [
"---\n",
"# 2. 内核调用符\n",
"\n",
"pyasc使用中括号`[]`语法调用核函数，称为**内核调用符（Kernel Call Operator）**：\n",
"\n",
"```python\n",
"vadd_kernel[USE_CORE_NUM, rt.current_stream()](x, y, z, block_length)\n",
"```\n",
"\n",
"<table style=\"margin-left: 0; margin-right: auto; width: auto; border: 1px solid #ddd;\">\n",
" <thead>\n",
" <tr>\n",
" <th align=\"left\">位置</th>\n",
" <th align=\"left\">参数</th>\n",
" <th align=\"left\">说明</th>\n",
" <th align=\"left\">Ascend C对应</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr><td>中括号第1个</td><td><code>USE_CORE_NUM</code></td><td>启用的核数</td><td><code><<<blockDim>>></code></td></tr>\n",
" <tr><td>中括号第2个</td><td><code>rt.current_stream()</code></td><td>执行流</td><td><code>stream</code></td></tr>\n",
" <tr><td>小括号</td><td><code>x, y, z, block_length</code></td><td>核函数参数</td><td><code>(args)</code></td></tr>\n",
" </tbody>\n",
"</table>\n"
]
},
{
"cell_type": "markdown",
"id": "a3540aed",
"metadata": {},
"source": [
"---\n",
"# 3. 完整验证程序\n",
"\n",
"一个完整的pyasc算子验证程序包含以下六部分：\n",
"\n",
"<table style=\"margin-left: 0; margin-right: auto; width: auto; border: 1px solid #ddd;\">\n",
" <thead>\n",
" <tr>\n",
" <th align=\"left\">部分</th>\n",
" <th align=\"left\">内容</th>\n",
" <th align=\"left\">关键代码</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr><td>1. 导入依赖</td><td>asc, torch, config等</td><td><code>import asc; import torch</code></td></tr>\n",
" <tr><td>2. 常量定义</td><td>核数、缓冲区数等</td><td><code>USE_CORE_NUM = 8</code></td></tr>\n",
" <tr><td>3. 核函数</td><td>@asc.jit修饰的kernel</td><td><code>@asc.jit def vadd_kernel(...)</code></td></tr>\n",
" <tr><td>4. Launch函数</td><td>Host侧调用入口</td><td><code>def vadd_launch(...)</code></td></tr>\n",
" <tr><td>5. 验证函数</td><td>生成数据+验证结果</td><td><code>def vadd_custom(...)</code></td></tr>\n",
" <tr><td>6. 主函数</td><td>解析参数+运行</td><td><code>if __name__ == ...</code></td></tr>\n",
" </tbody>\n",
"</table>\n",
"\n",
"下面查看Add样例的完整代码：\n",
"\n",
"### 3.1 结果验证：torch.allclose\n",
"\n",
"验证程序的核心是 `assert torch.allclose(z, x + y)`，其中：\n",
"\n",
"- `torch.allclose(a, b)`：逐元素比较两个Tensor是否在容忍范围内相等\n",
"- 默认参数：`rtol=1e-5`（相对容忍），`atol=1e-8`（绝对容忍）\n",
"- 对于float32类型的算子验证，通常足够\n",
"\n",
"如果算子实现有误，`allclose`返回False，`assert`抛出AssertionError，程序报错停止；如果验证通过，输出`Sample add run success.`。\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13532443",
"metadata": {},
"outputs": [],
"source": [
"# 查看Add样例的完整验证程序\n",
"!cat ./src/add.py"
]
},
{
"cell_type": "markdown",
"id": "10b7a9e0",
"metadata": {},
"source": [
"---\n",
"# 4. 运行验证\n",
"\n",
"运行Add样例的完整验证程序："
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fd2811f6",
"metadata": {},
"outputs": [],
"source": [
"!python3 ./src/add.py -r NPU"
]
},
{
"cell_type": "markdown",
"id": "0e15592f",
"metadata": {},
"source": [
"运行成功后输出`Sample add run success.`，表示算子计算结果正确（`torch.allclose`验证通过）。\n"
]
},
{
"cell_type": "markdown",
"id": "bafeca39",
"metadata": {},
"source": [
"---\n",
"# 5. 框架版样例对比\n",
"\n",
"除了本节的`add.py`（手动同步版），pyasc项目还提供了`add_framework.py`（框架版）。两者实现相同的Add算子，但采用了不同的编程模型：\n",
"\n",
"<table style=\"margin-left: 0; margin-right: auto; width: auto; border: 1px solid #ddd;\">\n",
" <thead>\n",
" <tr>\n",
" <th align=\"left\">对比项</th>\n",
" <th align=\"left\">手动版（01_add）</th>\n",
" <th align=\"left\">框架版（02_add_framework）</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr><td>同步管理</td><td>手动set_flag/wait_flag</td><td>TPipe/TQue框架自动管理</td></tr>\n",
" <tr><td>Tensor管理</td><td>手动创建LocalTensor</td><td>TQue的alloc/deque管理</td></tr>\n",
" <tr><td>代码结构</td><td>单一大核函数</td><td>核函数+Device函数拆分</td></tr>\n",
" <tr><td>代码复杂度</td><td>高</td><td>低</td></tr>\n",
" <tr><td>灵活性</td><td>高</td><td>中</td></tr>\n",
" </tbody>\n",
"</table>\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6d481163",
"metadata": {},
"outputs": [],
"source": [
"# 查看框架版Add样例\n",
"!cat ./src/add_framework.py"
]
},
{
"cell_type": "markdown",
"id": "4877a5e6",
"metadata": {},
"source": [
"框架版使用`asc.TPipe()`和`asc.TQue`自动管理同步和内存，核函数逻辑更简洁。这些将在第3章详细讲解。\n",
"\n",
"---\n",
"# 6. 小结\n",
"\n",
"- **Launch函数**：Host侧调用核函数的入口，负责数据准备、核函数调用和结果返回\n",
"- **内核调用符**：`kernel[核数, stream](args)`，对应Ascend C的`<<<blockDim, stream>>>(args)`\n",
"- **验证程序**：六部分结构（导入→常量→核函数→Launch→验证→主函数）\n",
"- **框架版优势**：TPipe/TQue自动管理同步和内存，代码更简洁（第3章详解）\n",
"\n",
"\n",
"### 5.1 关键代码逐行对比\n",
"\n",
"**核函数定义对比：**\n",
"\n",
"| 手动版（add.py） | 框架版（add_framework.py） |\n",
"|---|---|\n",
"| `x_local = asc.LocalTensor(data_type, asc.TPosition.VECIN, 0, ...)` | `x_local = in_queue_x.alloc_tensor(x_gm.dtype)` |\n",
"| `asc.data_copy(x_local[...], x_gm[...], tile_length)` | `asc.data_copy(x_local, x_gm[...], tile_length)` |\n",
"| `asc.set_flag(asc.HardEvent.MTE2_V, buf_id)` | `in_queue_x.enque(x_local)` |\n",
"| `asc.wait_flag(asc.HardEvent.MTE2_V, buf_id)` | `x_local = in_queue_x.deque(z_gm.dtype)` |\n",
"| `asc.add(z_local[...], x_local[...], y_local[...], ...)` | `asc.add(z_local, x_local, y_local, tile_length)` |\n",
"| `asc.set_flag(asc.HardEvent.V_MTE3, buf_id)` | `out_queue_z.enque(z_local)` |\n",
"| `asc.wait_flag(asc.HardEvent.V_MTE3, buf_id)` | `z_local = out_queue_z.deque(z_gm.dtype)` |\n",
"| `asc.data_copy(z_gm[...], z_local[...], tile_length)` | `asc.data_copy(z_gm[...], z_local, tile_length)` |\n",
"| `asc.set_flag(asc.HardEvent.MTE3_MTE2, buf_id)` | `out_queue_z.free_tensor(z_local)` |\n",
"| `asc.wait_flag(asc.HardEvent.MTE3_MTE2, buf_id)` | （框架自动处理） |\n",
"\n",
"核心差异：手动版需要10行代码（含6行同步），框架版只需9行代码（0行手动同步），同步逻辑由TQue的enque/deque自动管理。\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"---\n",
"\n",
"## 课后实践\n",
"\n",
"**实践任务：** 修改Add样例的参数，观察运行结果变化。\n",
"\n",
"1. 尝试修改`USE_CORE_NUM`为4，观察是否能正常运行？\n",
"2. 尝试修改`size = 8 * 2048`为`size = 4 * 2048`，观察运行结果。\n",
"3. 思考：`block_length = total_length // USE_CORE_NUM`，如果`total_length`不能被`USE_CORE_NUM`整除会怎样？\n",
"\n",
"<details>\n",
"<summary>点击查看参考答案</summary>\n",
"\n",
"1. 修改`USE_CORE_NUM`为4后，每个核处理`8*2048/4=4096`个元素。只要`block_length`能被`TILE_NUM * BUFFER_NUM`整除即可正常运行。\n",
"2. 修改size后，每个核处理`4*2048/8=1024`个元素，正常运行。\n",
"3. 如果不能整除，`block_length`向下取整，部分数据不被处理。框架版使用向上取整`(total + N - 1) // N`处理此情况。\n",
"\n",
"</details>"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 5
}