| @@ -4,7 +4,7 @@ | |||
| 4 | "cell_type": "markdown", | 4 | "cell_type": "markdown", |
| 5 | "metadata": {}, | 5 | "metadata": {}, |
| 6 | "source": [ | 6 | "source": [ |
| 7 | - "# CANN、芯片以及 PyPTO 介绍\n", | 7 | + "# 芯片、CANN以及 PyPTO 介绍\n", |
| 8 | "\n", | 8 | "\n", |
| 9 | "本章是 PyPTO 算子开发学习路径的起点。正式编写算子之前,需要先建立三层基础认知:昇腾芯片提供什么样的硬件能力,CANN 在软件栈中承担什么职责,PyPTO 又如何把开发者的算法表达落到芯片执行。\n", | 9 | "本章是 PyPTO 算子开发学习路径的起点。正式编写算子之前,需要先建立三层基础认知:昇腾芯片提供什么样的硬件能力,CANN 在软件栈中承担什么职责,PyPTO 又如何把开发者的算法表达落到芯片执行。\n", |
| 10 | "\n", | 10 | "\n", |
| @@ -21,26 +21,9 @@ | |||
| 21 | "\n", | 21 | "\n", |
| 22 | "学习 PyPTO 时,经常会同时遇到 CANN、昇腾芯片和 PyPTO 这三个概念。它们不是并列关系,而是从硬件到软件再到开发接口的分层关系:\n", | 22 | "学习 PyPTO 时,经常会同时遇到 CANN、昇腾芯片和 PyPTO 这三个概念。它们不是并列关系,而是从硬件到软件再到开发接口的分层关系:\n", |
| 23 | "\n", | 23 | "\n", |
| 24 | - "- **昇腾芯片(NPU)**:硬件层,提供 AI 计算所需的物理算力,包括 AI Core、多级片上存储、Global Memory 以及多精度计算能力。\n", | 24 | + "- `昇腾芯片(NPU)`:硬件层,提供 AI 计算所需的物理算力,核心计算单元是 AI Core。\n", |
| 25 | - "- **CANN(Compute Architecture for Neural Networks)**:软件栈层,提供驱动固件、Toolkit、Ops 包、编译优化工具链、运行时调度和性能分析等能力,是上层框架访问昇腾硬件的基础。\n", | 25 | + "- `CANN(Compute Architecture for Neural Networks)`:软件栈层,包含Toolkit包、算子包(ops)、加速库(NNAL),对上支持多种AI框架,对下服务AI处理器与编程,发挥承上启下的关键作用,是提升昇腾AI处理器计算效率的关键平台。\n", |
| 26 | - "- **PyPTO**:编程框架层,基于 CANN 的能力,为开发者提供 Python 友好的算子开发接口,通过 PTO 编程范式和 Tile 编程模型描述计算逻辑。\n", | 26 | + "- `PyPTO`:编程框架层,CANN推出的一款面向AI加速器的高效编程框架,旨在简化算子开发流程,同时保持高性能计算能力。PyPTO提供PyPTO Tensor与PyPTO Pro两种编程方式,本课程主要介绍Tensor编程方式。" |
| 27 | - "\n", | ||
| 28 | - "它们的关系可以概括为:\n", | ||
| 29 | - "\n", | ||
| 30 | - "```text\n", | ||
| 31 | - "开发者\n", | ||
| 32 | - " |\n", | ||
| 33 | - " v\n", | ||
| 34 | - "PyPTO:用 Python 和 Tensor/Tile 抽象描述“算什么”\n", | ||
| 35 | - " |\n", | ||
| 36 | - " v\n", | ||
| 37 | - "CANN:完成编译优化、代码生成、运行时调度和工具链支持\n", | ||
| 38 | - " |\n", | ||
| 39 | - " v\n", | ||
| 40 | - "昇腾芯片:在 AI Core 和内存层次上真正执行计算\n", | ||
| 41 | - "```\n", | ||
| 42 | - "\n", | ||
| 43 | - "也就是说,开发者通常通过 PyPTO 表达计算,CANN 负责把这种表达逐步 lowering、优化并调度执行,最终由昇腾芯片完成实际计算。" | ||
| 44 | ] | 27 | ] |
| 45 | }, | 28 | }, |
| 46 | { | 29 | { |
| @@ -88,12 +71,10 @@ | |||
| 88 | "| 小节 | 主题 | 关注重点 |\n", | 71 | "| 小节 | 主题 | 关注重点 |\n", |
| 89 | "| --- | --- | --- |\n", | 72 | "| --- | --- | --- |\n", |
| 90 | "| 1.1 | 章节介绍 | 建立 CANN、昇腾芯片和 PyPTO 的整体关系 |\n", | 73 | "| 1.1 | 章节介绍 | 建立 CANN、昇腾芯片和 PyPTO 的整体关系 |\n", |
| 91 | - "| 1.2 | CANN | 软件栈定位、核心组成、功能特性和安装方式 |\n", | 74 | + "| 1.2 | 昇腾芯片 | AI Core、内存层次、精度支持、MPMD 执行模式和应用场景 |\n", |
| 92 | - "| 1.3 | 昇腾芯片 | AI Core、内存层次、精度支持、MPMD 执行模式和应用场景 |\n", | 75 | + "| 1.3 | CANN | 软件栈定位、核心组成、功能特性 |\n", |
| 93 | "| 1.4 | PyPTO | PTO 编程范式、Tile 编程模型、多层级计算图、核心特性和适用场景 |\n", | 76 | "| 1.4 | PyPTO | PTO 编程范式、Tile 编程模型、多层级计算图、核心特性和适用场景 |\n", |
| 94 | - "| 1.5 | 章节实践 | 综合检查 CANN、昇腾芯片和 PyPTO 的层级关系与关键概念 |\n", | 77 | + "| 1.5 | 章节实践 | 综合检查 昇腾芯片、CANN和 PyPTO 的层级关系与关键概念 |\n" |
| 95 | - "\n", | ||
| 96 | - "阅读顺序很重要。先理解 CANN,才能知道 PyPTO 生成的内容如何进入编译和运行流程;先理解芯片的计算与内存结构,才能理解 Tile 为什么是 PyPTO 的关键抽象;最后再看 PyPTO,并通过章节实践完成自测,才能把编程模型与底层执行机制对应起来。\n" | ||
| 97 | ] | 78 | ] |
| 98 | }, | 79 | }, |
| 99 | { | 80 | { |
| @@ -104,8 +85,8 @@ | |||
| 104 | "\n", | 85 | "\n", |
| 105 | "本章偏概念导入,阅读时可以重点思考以下问题:\n", | 86 | "本章偏概念导入,阅读时可以重点思考以下问题:\n", |
| 106 | "\n", | 87 | "\n", |
| 107 | - "1. CANN 在芯片和上层开发框架之间承担了哪些职责?\n", | 88 | + "1. 昇腾芯片的内存层次和 AI Core 并行能力,为什么会影响算子的编程模型?\n", |
| 108 | - "2. 昇腾芯片的内存层次和 AI Core 并行能力,为什么会影响算子的编程模型?\n", | 89 | + "2. CANN 在芯片和上层开发框架之间承担了哪些职责?\n", |
| 109 | "3. PyPTO 的 Tile 编程模型解决了什么问题?它和硬件内存搬运、并行计算有什么关系?\n", | 90 | "3. PyPTO 的 Tile 编程模型解决了什么问题?它和硬件内存搬运、并行计算有什么关系?\n", |
| 110 | "4. Tensor Graph、Tile Graph、Block Graph、Execute Graph 分别更接近算法表达、硬件感知优化还是最终执行?\n", | 91 | "4. Tensor Graph、Tile Graph、Block Graph、Execute Graph 分别更接近算法表达、硬件感知优化还是最终执行?\n", |
| 111 | "5. 作为算子开发者,为什么通常希望使用 PyPTO,而不是直接面对底层硬件指令?\n", | 92 | "5. 作为算子开发者,为什么通常希望使用 PyPTO,而不是直接面对底层硬件指令?\n", |
| @@ -163,7 +144,7 @@ | |||
| 163 | "\n", | 144 | "\n", |
| 164 | "本节完成了本章导读:昇腾芯片是硬件算力基础,CANN 是连接硬件与上层框架的软件栈,PyPTO 是面向高性能算子开发的 Python 编程框架。\n", | 145 | "本节完成了本章导读:昇腾芯片是硬件算力基础,CANN 是连接硬件与上层框架的软件栈,PyPTO 是面向高性能算子开发的 Python 编程框架。\n", |
| 165 | "\n", | 146 | "\n", |
| 166 | - "接下来将分别介绍 CANN、昇腾芯片和 PyPTO,并在章节实践中完成综合自测。学完这几部分后,读者应能从全栈角度理解 PyPTO 算子开发:代码如何被描述,如何被 CANN 编译优化,又如何最终运行在昇腾芯片上。\n" | 147 | + "接下来将分别介绍 昇腾芯片、CANN和 PyPTO,并在章节实践中完成综合自测。学完这几部分后,读者应能从全栈角度理解 PyPTO 算子开发:代码如何被描述,如何被 CANN 编译优化,又如何最终运行在昇腾芯片上。\n" |
| 167 | ] | 148 | ] |
| 168 | } | 149 | } |
| 169 | ], | 150 | ], |
| @@ -0,0 +1,157 @@ | |||
| 1 | +{ | ||
| 2 | + "cells": [ | ||
| 3 | + { | ||
| 4 | + "cell_type": "markdown", | ||
| 5 | + "metadata": {}, | ||
| 6 | + "source": [ | ||
| 7 | + "# 昇腾芯片\n", | ||
| 8 | + "\n", | ||
| 9 | + "昇腾(Ascend)是华为面向人工智能计算设计的 AI 处理器系列,也常被称为 NPU(Neural Processing Unit)。在 AI 计算场景中,NPU 会针对神经网络中的矩阵计算、向量计算和数据搬运等典型负载进行专门优化,以提升训练和推理任务的执行效率。\n", | ||
| 10 | + "\n", | ||
| 11 | + "在 PyPTO 学习中,理解昇腾芯片并不是为了直接编写硬件指令,而是为了理解为什么 PyPTO 会强调 Tile、内存层次、并行粒度和数据类型。这些概念都和芯片的执行方式密切相关。" | ||
| 12 | + ] | ||
| 13 | + }, | ||
| 14 | + { | ||
| 15 | + "cell_type": "markdown", | ||
| 16 | + "metadata": {}, | ||
| 17 | + "source": [ | ||
| 18 | + "在异构计算架构中,专用处理器承担密集型计算任务,实现算力多元化调度与效能最大化。下文以Atlas A2训练系列产品为例,拆解昇腾设备的异构计算实现逻辑,清晰呈现NPU的硬件架构与计算机制。\n", | ||
| 19 | + "\n", | ||
| 20 | + "## 1. Host和Device\n", | ||
| 21 | + "典型应用场景中,将Atlas A2加速卡插入服务器(或个人PC)后,程序的整体逻辑控制、任务调度均在CPU侧执行;当触发大规模数据密集型计算任务时,CPU侧会将待计算数据传输至NPU侧内存,由NPU完成专用并行计算,最终将结果回传CPU侧,形成“控制-计算”的异构协同闭环。\n", | ||
| 22 | + "\n", | ||
| 23 | + "<img src=\"./images/host_and_device.png\" alt=\"host_and_device\" width=\"700px\" >\n", | ||
| 24 | + "\n", | ||
| 25 | + "- **Host(主机侧)**:即服务器CPU及配套内存所在侧,核心负责任务发起、逻辑判断、数据传输调度及结果汇总。\n", | ||
| 26 | + "\n", | ||
| 27 | + "- **Device(设备侧)**:即昇腾加速卡的NPU及专属内存所在侧,核心负责承接密集型计算任务,通过专用硬件单元实现高效并行运算。\n", | ||
| 28 | + "\n", | ||
| 29 | + "## 2. NPU内部细节\n", | ||
| 30 | + "明确Host与Device的协同关系后,进一步拆解NPU内部单元组成,理解其专用计算能力的硬件支撑。\n", | ||
| 31 | + "\n", | ||
| 32 | + "<img src=\"./images/npu_processor.png\" alt=\"npu_processor\" width=\"700px\" >\n", | ||
| 33 | + "\n", | ||
| 34 | + "- **内存**:NPU专属内存,用于存储CPU传输的待计算数据、计算过程中的中间结果及最终输出数据,为高速运算提供数据缓存支撑。\n", | ||
| 35 | + "\n", | ||
| 36 | + "- **AI Core**:昇腾NPU的核心计算单元,专为矩阵、向量、标量等密集型计算任务设计,是算子加速执行的核心载体。\n", | ||
| 37 | + "\n", | ||
| 38 | + "- **AI CPU**:负责处理不适合在AI Core上执行的任务,如轻量级逻辑运算、非并行化处理的辅助算子,补充AI Core的计算场景覆盖。\n", | ||
| 39 | + "\n", | ||
| 40 | + "- **控制CPU**:专注于NPU整体运行控制,协调内部各单元的工作时序,保障计算流程有序推进。\n", | ||
| 41 | + "\n", | ||
| 42 | + "- **任务调度器**:基于任务类型与硬件资源状态,实现计算任务在AI Core、AI CPU间的高效分配与动态调度,最大化硬件利用率。\n", | ||
| 43 | + "\n", | ||
| 44 | + "- **数字视觉预处理模块**:专用图像硬件处理单元,负责图像解码、编码、格式转换等预处理任务,减少AI Core的非核心计算负载。\n", | ||
| 45 | + "\n", | ||
| 46 | + "## 3. AI Core内部细节\n", | ||
| 47 | + "AI Core作为NPU的计算核心,绝大多数算子的加速执行均在此完成。其架构延续传统芯片“计算-存储-控制”的三大核心模块,通过专用化设计实现极致并行效能,下文逐一解析各模块功能。\n", | ||
| 48 | + "\n", | ||
| 49 | + "<img src=\"./images/ai_core_architecture.png\" alt=\"ai_core_architecture\" width=\"700px\" >\n", | ||
| 50 | + "\n", | ||
| 51 | + "### 3.1 计算单元\n", | ||
| 52 | + "AI Core内置三类专用计算单元,分别适配矩阵、向量、标量不同维度的计算需求,实现分工协作与并行提速。\n", | ||
| 53 | + "\n", | ||
| 54 | + "<img src=\"./images/computing_unit.png\" alt=\"computing_unit\" width=\"700px\" >\n", | ||
| 55 | + "\n", | ||
| 56 | + "1. **矩阵计算单元(Cube Unit)** \n", | ||
| 57 | + " 核心负责矩阵乘加运算,搭配累加器实现高效数据累加。硬件层面支持高精度并行计算:FP16精度下,单时钟周期可完成16×16与16×16矩阵乘(4096次乘加运算);INT8精度下,单时钟周期可完成16×32与32×16矩阵乘(8192次乘加运算)。累加器可将当前矩阵乘结果与历史中间结果叠加,天然适配卷积运算中偏置(bias)添加等场景。\n", | ||
| 58 | + "\n", | ||
| 59 | + "2. **向量计算单元(Vector Unit)** \n", | ||
| 60 | + " 专注于向量级运算,支持FP16、FP32、Int32、Int8等多数据类型,覆盖基本算术运算与定制化向量操作。运算效能表现为:单时钟周期可完成两组128长度FP16向量的加/乘运算,或64个FP32/Int32向量的加/乘运算,适配激活函数、数据归一化等向量密集型任务。\n", | ||
| 61 | + "\n", | ||
| 62 | + "3. **标量计算单元(Scalar Unit)** \n", | ||
| 63 | + " 承担标量运算与AI Core整体控制职责,相当于微型CPU。核心功能包括:循环控制、分支判断、地址计算与参数配置(为Cube/Vector单元提供数据地址及运算参数),同时支持基础算术运算,保障各计算单元的协同有序运行。\n", | ||
| 64 | + "\n", | ||
| 65 | + "### 3.2 存储系统 \n", | ||
| 66 | + "由片上存储单元与数据通路组成,通过分层存储设计减少外部总线访问频次,降低延迟、提升带宽,为高速计算提供数据支撑。\n", | ||
| 67 | + "\n", | ||
| 68 | + "<img src=\"./images/storage_system.png\" alt=\"storage_system\" width=\"700px\" >\n", | ||
| 69 | + "\n", | ||
| 70 | + "1. **存储转换引擎** \n", | ||
| 71 | + " 负责AI Core内部不同缓冲区的数据读写管理,同时支持多种数据格式转换操作,如Padding(填充)、Transpose(转置)、Img2Col(3D图像转2D矩阵)等预处理/后处理操作。此外,可通过总线接口直接访问AI Core外部的低层级缓存,拓展数据访问范围。\n", | ||
| 72 | + "\n", | ||
| 73 | + "2. **缓冲区** \n", | ||
| 74 | + " 包含L1缓冲区、L0A/L0B缓冲区、L0C缓冲区、统一缓冲区及标量缓冲区,核心作用是缓存高频复用数据与中间结果:一方面,将频繁访问的数据暂存片上,避免反复从外部读取,减少总线拥堵与功耗消耗;另一方面,存储神经网络各层计算的中间结果,为下一层运算快速提供数据,相较总线访问大幅降低延迟、提升运算效率。\n", | ||
| 75 | + "\n", | ||
| 76 | + "3. **寄存器** \n", | ||
| 77 | + " 主要为标量计算单元服务,用于暂存标量数据、运算指令及控制参数,保障标量运算的高速执行。\n", | ||
| 78 | + "\n", | ||
| 79 | + "### 3.3 控制单元\n", | ||
| 80 | + "作为AI Core的“指挥中枢”,负责全流程指令控制与时序协调,确保各单元并行运算的有序性与数据一致性。核心组成及功能如下:\n", | ||
| 81 | + "\n", | ||
| 82 | + "<img src=\"./images/control_unit.png\" alt=\"control_unit\" width=\"700px\" >\n", | ||
| 83 | + "\n", | ||
| 84 | + "- **系统控制模块**:管控任务块(AI Core最小计算任务粒度)的执行进程,任务块完成后执行中断处理与状态上报;若运算过程中出现错误,及时向任务调度器反馈错误状态。\n", | ||
| 85 | + "\n", | ||
| 86 | + "- **指令缓存**:提前预取后续待执行指令,一次性读取多条指令缓存,避免指令逐条读取的延迟,提升指令执行效率。\n", | ||
| 87 | + "\n", | ||
| 88 | + "- **标量指令处理队列**:指令解码后导入该队列,完成地址解码与运算控制,覆盖矩阵、向量、存储转换等各类指令。\n", | ||
| 89 | + "\n", | ||
| 90 | + "- **指令发射模块**:读取标量指令处理队列中的指令地址与参数,解码后按指令类型分发至对应执行队列,标量指令则留存于该队列中执行。\n", | ||
| 91 | + "\n", | ||
| 92 | + "- **指令执行队列**:分为矩阵运算队列、向量运算队列、存储转换队列,不同类型指令按顺序在对应队列中执行,实现并行流水线运算。\n", | ||
| 93 | + "\n", | ||
| 94 | + "- **事件同步模块**:实时监控各指令流水线的执行状态,分析不同流水线的依赖关系,解决数据依赖与时序同步问题(如矩阵乘完成后再执行向量加法),保障运算结果正确性。\n" | ||
| 95 | + ] | ||
| 96 | + }, | ||
| 97 | + { | ||
| 98 | + "cell_type": "markdown", | ||
| 99 | + "metadata": {}, | ||
| 100 | + "source": [ | ||
| 101 | + "## 4. 课后练习\n", | ||
| 102 | + "\n", | ||
| 103 | + "本节练习用于检查昇腾芯片、AI Core、内存层次和 Tile 思想。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | ||
| 104 | + "\n", | ||
| 105 | + "1. (单选题)在昇腾NPU中,Host侧的核心职责是什么? \n", | ||
| 106 | + " A. 执行密集型计算 \n", | ||
| 107 | + " B. 任务发起、逻辑判断、数据传输调度 \n", | ||
| 108 | + " C. 专用并行运算 \n", | ||
| 109 | + " D. 图像预处理 \n", | ||
| 110 | + "\n", | ||
| 111 | + "2. (单选题)AI Core的计算单元中,负责矩阵乘加运算的是? \n", | ||
| 112 | + " A. Vector Unit \n", | ||
| 113 | + " B. Scalar Unit \n", | ||
| 114 | + " C. Cube Unit \n", | ||
| 115 | + " D. 存储转换引擎 \n", | ||
| 116 | + "\n", | ||
| 117 | + "\n", | ||
| 118 | + "**执行以下代码获取答案。**\n" | ||
| 119 | + ] | ||
| 120 | + }, | ||
| 121 | + { | ||
| 122 | + "cell_type": "code", | ||
| 123 | + "execution_count": null, | ||
| 124 | + "metadata": {}, | ||
| 125 | + "outputs": [], | ||
| 126 | + "source": [ | ||
| 127 | + "!cat ./answer/01.02_answer.txt\n" | ||
| 128 | + ] | ||
| 129 | + }, | ||
| 130 | + { | ||
| 131 | + "cell_type": "markdown", | ||
| 132 | + "metadata": {}, | ||
| 133 | + "source": [ | ||
| 134 | + "## 5. 本节小结\n", | ||
| 135 | + "\n", | ||
| 136 | + "昇腾NPU的计算加速流程可概括为:开发者通过API编写的计算逻辑,经架构转换为一条条硬件可执行指令,下发至NPU后,由控制单元按类型分发至对应指令队列;AI Core内的矩阵、向量、存储转换等单元从队列中并行取指执行;针对存在先后依赖的任务(如矩阵乘后紧跟向量加法),开发者调用同步API,由事件同步模块管控流水线时序,确保依赖关系满足。\n", | ||
| 137 | + "\n", | ||
| 138 | + "需说明的是,本文介绍的为AI Core基础架构,不同代际的昇腾NPU(如310、910B、910C)在硬件排布、数据通路设计上可能存在差异。\n", | ||
| 139 | + "\n", | ||
| 140 | + "理解了芯片的这些基本特征后,就更容易明白为什么 PyPTO 要围绕 Tensor、Tile、多层级计算图和自动代码生成来设计。" | ||
| 141 | + ] | ||
| 142 | + } | ||
| 143 | + ], | ||
| 144 | + "metadata": { | ||
| 145 | + "kernelspec": { | ||
| 146 | + "display_name": "Python 3", | ||
| 147 | + "language": "python", | ||
| 148 | + "name": "python3" | ||
| 149 | + }, | ||
| 150 | + "language_info": { | ||
| 151 | + "name": "python", | ||
| 152 | + "version": "3.10.0" | ||
| 153 | + } | ||
| 154 | + }, | ||
| 155 | + "nbformat": 4, | ||
| 156 | + "nbformat_minor": 4 | ||
| 157 | +} | ||
| @@ -6,9 +6,9 @@ | |||
| 6 | "source": [ | 6 | "source": [ |
| 7 | "# CANN 是什么\n", | 7 | "# CANN 是什么\n", |
| 8 | "\n", | 8 | "\n", |
| 9 | - "CANN(Compute Architecture for Neural Networks)是昇腾 AI 处理器的软件栈。它位于上层 AI 框架和底层昇腾硬件之间,负责提供开发、编译、调试、运行和性能分析所需的一整套基础能力。\n", | 9 | + "CANN(Compute Architecture for Neural Networks)是华为针对AI场景推出的异构计算架构,对上支持多种AI框架,对下服务AI处理器与编程,发挥承上启下的关键作用,是提升昇腾AI处理器计算效率的关键平台。 \n", |
| 10 | "\n", | 10 | "\n", |
| 11 | - "如果把昇腾芯片理解为提供算力的硬件基础,那么 CANN 就是让这些算力能够被开发者和框架使用起来的软件底座。无论是使用 PyTorch、TensorFlow,还是使用 PyPTO 开发自定义算子,最终都需要依赖 CANN 提供的驱动、编译工具链、运行时和算子库等能力。" | 11 | + "简单来说,CANN就像AI芯片与上层应用之间的“翻译官+调度员”,把开发者写的AI算法代码,转换成芯片能高效执行的指令,同时优化算力分配,最大化AI芯片的性能。" |
| 12 | ] | 12 | ] |
| 13 | }, | 13 | }, |
| 14 | { | 14 | { |
| @@ -80,17 +80,58 @@ | |||
| 80 | "cell_type": "markdown", | 80 | "cell_type": "markdown", |
| 81 | "metadata": {}, | 81 | "metadata": {}, |
| 82 | "source": [ | 82 | "source": [ |
| 83 | - "## 4. 安装与使用方式\n", | 83 | + "## 4. 架构介绍\n", |
| 84 | "\n", | 84 | "\n", |
| 85 | - "CANN 支持多种安装和使用方式,选择哪一种主要取决于是否拥有本地昇腾设备,以及希望以什么方式搭建开发环境。\n", | 85 | + "### 4.1 什么是异构计算架构\n", |
| 86 | + "异构计算架构是“使能硬件异构并行计算的软件栈”,最简单的结构就是通用CPU和专用处理器的并行计算组合,其最大的好处是使能多元算力,化解算力瓶颈,从而实现算力最大化。\n", | ||
| 86 | "\n", | 87 | "\n", |
| 87 | - "| 方式 | 说明 | 适用场景 |\n", | 88 | + "<img src=\"./images/heterogeneous_computing_architecture.png\" alt=\"heterogeneous_computing_architecture\" width=\"350px\" >\n", |
| 88 | - "| --- | --- | --- |\n", | ||
| 89 | - "| CANNLab | 云端开发平台,提供一站式环境 | 没有本地昇腾设备,希望快速体验或学习 |\n", | ||
| 90 | - "| 主机安装 | 在宿主机上安装驱动、固件、Toolkit 和相关包 | 拥有昇腾设备,需要本地开发和调优 |\n", | ||
| 91 | - "| Docker | 使用预集成镜像快速部署环境 | 希望快速复现实验环境或减少本地配置成本 |\n", | ||
| 92 | "\n", | 89 | "\n", |
| 93 | - "初学阶段更重要的是理解 CANN 的定位和组成。真正进入实践时,再根据设备条件和项目要求选择具体安装方式。" | 90 | + "### 4.2 为什么要用异构计算架构\n", |
| 91 | + "传统CPU以标量计算为核心,仅能逐元素执行单一运算。而专用硬件可直接完成向量、矩阵级的并行计算,以16×16矩阵乘这一AI场景的核心运算为例,不同计算架构的效率差距极为显著。\n", | ||
| 92 | + "\n", | ||
| 93 | + "<img src=\"./images/matrix_multiplication_example.png\" alt=\"matrix_multiplication_example\" width=\"700px\" >\n", | ||
| 94 | + "\n", | ||
| 95 | + "在CPU上执行该计算时,每个时钟周期仅能完成一次标量乘加运算,核心计算逻辑如下:\n", | ||
| 96 | + "```\n", | ||
| 97 | + "for (int i=0; i<16; i++) {\n", | ||
| 98 | + " for (int j=0; j<16; j++) {\n", | ||
| 99 | + " for (int k=0; k<16; k++) {\n", | ||
| 100 | + " // 乘、加操作各占1个时钟周期,单次共需2个cycle\n", | ||
| 101 | + " c[i][j] += a[i][k] * b[k][j];\n", | ||
| 102 | + " }\n", | ||
| 103 | + " }\n", | ||
| 104 | + "}\n", | ||
| 105 | + "```\n", | ||
| 106 | + "**总耗时周期**:Cycle = 16×16×16×2 = 8192。\n", | ||
| 107 | + "\n", | ||
| 108 | + "基于 Vector 矢量硬件单元的专用架构,每个时钟周期可完成一组行与列的矢量乘加运算,核心计算逻辑如下:\n", | ||
| 109 | + "```\n", | ||
| 110 | + "for (int i=0; i<16; i++) {\n", | ||
| 111 | + " for (int j=0; j<16; j++) {\n", | ||
| 112 | + " // 一行与一列的所有元素同时完成乘加运算\n", | ||
| 113 | + " c[i][j] = a[i][:] *+ b[:][j];\n", | ||
| 114 | + " }\n", | ||
| 115 | + "}\n", | ||
| 116 | + "```\n", | ||
| 117 | + "**总耗时周期**:Cycle = 16×16 = 256。\n", | ||
| 118 | + "\n", | ||
| 119 | + "如果是集成度更高的Cube矩阵乘硬件单元,单个时钟周期即可完成整份矩阵的乘加运算,核心计算逻辑如下:\n", | ||
| 120 | + "```\n", | ||
| 121 | + "// 两个16×16矩阵一次性完成并行乘加\n", | ||
| 122 | + "c[:][:] = a[:][:] * b[:][:];\n", | ||
| 123 | + "```\n", | ||
| 124 | + "**总耗时周期**:Cycle = 1。\n", | ||
| 125 | + "\n", | ||
| 126 | + "由此可见,面对 AI 场景的密集型计算需求,让专用硬件单元承接核心计算任务,而 CPU 专注于逻辑判断、指令下发等通用任务,可实现 “专人干专事” 的算力最优分配,大幅提升整体计算效率。\n", | ||
| 127 | + "\n", | ||
| 128 | + "### 4.3 CANN架构介绍\n", | ||
| 129 | + "CANN正是针对昇腾NPU的**异构计算架构**,它采用分层解耦的设计,上层组件体现CANN的内部能力,下层对接硬件原子能力。通过开放的接口帮助开发者快速调用底层算力,完成计算加速。\n", | ||
| 130 | + "- 提供高性能算子及通信算法,帮助开发者进行大模型并行加速,释放芯片澎湃算力\n", | ||
| 131 | + "- 提供多种算子开发方式,支持开发者进行高效开发与迁移\n", | ||
| 132 | + "- 全面开源,给开发者提供丰富参考实践,让开发者具备自主创新能力\n", | ||
| 133 | + "\n", | ||
| 134 | + "<img src=\"./images/cann_software_architecture.png\" alt=\"cann_software_architecture\" width=\"900px\" >" | ||
| 94 | ] | 135 | ] |
| 95 | }, | 136 | }, |
| 96 | { | 137 | { |
| @@ -132,7 +173,7 @@ | |||
| 132 | "metadata": {}, | 173 | "metadata": {}, |
| 133 | "outputs": [], | 174 | "outputs": [], |
| 134 | "source": [ | 175 | "source": [ |
| 135 | - "!cat ./answer/01.02_answer.txt\n" | 176 | + "!cat ./answer/01.03_answer.txt\n" |
| 136 | ] | 177 | ] |
| 137 | }, | 178 | }, |
| 138 | { | 179 | { |
| @@ -141,9 +182,7 @@ | |||
| 141 | "source": [ | 182 | "source": [ |
| 142 | "## 6. 本节小结\n", | 183 | "## 6. 本节小结\n", |
| 143 | "\n", | 184 | "\n", |
| 144 | - "CANN 是昇腾 AI 开发的软件底座,向下连接芯片,向上支撑框架和开发工具。它通过驱动固件、Toolkit、Ops 包、编译工具链和运行时环境,把开发者描述的计算变成可以在昇腾芯片上执行的任务。\n", | 185 | + "CANN 是昇腾 AI 开发的软件底座,向下连接芯片,向上支撑框架和开发工具。它通过驱动固件、Toolkit、Ops 包、编译工具链和运行时环境,把开发者描述的计算变成可以在昇腾芯片上执行的任务。" |
| 145 | - "\n", | ||
| 146 | - "下一节将继续向下看硬件层,了解昇腾芯片的 AI Core、内存层次、精度支持和执行模式。这些硬件特性也是理解 PyPTO Tile 编程模型的重要基础。" | ||
| 147 | ] | 186 | ] |
| 148 | } | 187 | } |
| 149 | ], | 188 | ], |
| @@ -160,4 +199,4 @@ | |||
| 160 | }, | 199 | }, |
| 161 | "nbformat": 4, | 200 | "nbformat": 4, |
| 162 | "nbformat_minor": 4 | 201 | "nbformat_minor": 4 |
| 163 | -} | 202 | +} |
| @@ -1,165 +0,0 @@ | |||
| 1 | -{ | ||
| 2 | - "cells": [ | ||
| 3 | - { | ||
| 4 | - "cell_type": "markdown", | ||
| 5 | - "metadata": {}, | ||
| 6 | - "source": [ | ||
| 7 | - "# 昇腾芯片\n", | ||
| 8 | - "\n", | ||
| 9 | - "昇腾(Ascend)是华为面向人工智能计算设计的 AI 处理器系列,也常被称为 NPU(Neural Processing Unit)。在 AI 计算场景中,NPU 会针对神经网络中的矩阵计算、向量计算和数据搬运等典型负载进行专门优化,以提升训练和推理任务的执行效率。\n", | ||
| 10 | - "\n", | ||
| 11 | - "在 PyPTO 学习中,理解昇腾芯片并不是为了直接编写硬件指令,而是为了理解为什么 PyPTO 会强调 Tile、内存层次、并行粒度和数据类型。这些概念都和芯片的执行方式密切相关。" | ||
| 12 | - ] | ||
| 13 | - }, | ||
| 14 | - { | ||
| 15 | - "cell_type": "markdown", | ||
| 16 | - "metadata": {}, | ||
| 17 | - "source": [ | ||
| 18 | - "## 1. 面向 AI 计算的硬件特点\n", | ||
| 19 | - "\n", | ||
| 20 | - "昇腾芯片围绕 AI 计算场景进行了专门设计,典型特点包括:\n", | ||
| 21 | - "\n", | ||
| 22 | - "- **专用计算单元**:通过 AI Core 等计算单元加速矩阵、向量和张量相关计算。\n", | ||
| 23 | - "- **多核并行能力**:多个 AI Core 可以并行处理不同任务或不同数据块,提高整体吞吐。\n", | ||
| 24 | - "- **多级内存层次**:通过片上缓存和片外内存配合,降低数据访问成本。\n", | ||
| 25 | - "- **多精度支持**:支持训练、推理和量化场景中常见的多种数据类型。\n", | ||
| 26 | - "\n", | ||
| 27 | - "这些硬件特点决定了高性能算子开发不能只关注数学公式,还需要关注数据如何切分、如何搬运、如何复用,以及如何映射到多个计算核心上。" | ||
| 28 | - ] | ||
| 29 | - }, | ||
| 30 | - { | ||
| 31 | - "cell_type": "markdown", | ||
| 32 | - "metadata": {}, | ||
| 33 | - "source": [ | ||
| 34 | - "## 2. AI Core 与并行执行\n", | ||
| 35 | - "\n", | ||
| 36 | - "AI Core 是昇腾芯片执行 AI 计算的重要单元。实际运行时,一个算子通常会被拆分成多个可以并行执行的任务,分配到不同 AI Core 上处理。\n", | ||
| 37 | - "\n", | ||
| 38 | - "在这种模式下,开发者需要理解两个层面的并行:\n", | ||
| 39 | - "\n", | ||
| 40 | - "1. **计算并行**:多个 AI Core 同时执行计算,提高整体吞吐。\n", | ||
| 41 | - "2. **数据并行**:大 Tensor 被切分成多个 Tile,不同核心处理不同数据块。\n", | ||
| 42 | - "\n", | ||
| 43 | - "昇腾 AI 处理器支持 MPMD(Multiple Program Multiple Data)执行模式,即不同核心可以执行不同程序,并处理不同数据。对于复杂融合算子或大模型组件来说,这种模式为更灵活的调度和并行优化提供了基础。" | ||
| 44 | - ] | ||
| 45 | - }, | ||
| 46 | - { | ||
| 47 | - "cell_type": "markdown", | ||
| 48 | - "metadata": {}, | ||
| 49 | - "source": [ | ||
| 50 | - "## 3. 内存层次与 Tile 思想\n", | ||
| 51 | - "\n", | ||
| 52 | - "AI 计算的性能不仅取决于算力,也取决于数据是否能高效地到达计算单元。昇腾芯片采用多级内存层次,例如 L0 Buffer、L1 Buffer、Unified Buffer、Global Memory 以及片外 HBM 等。\n", | ||
| 53 | - "\n", | ||
| 54 | - "不同层级的内存容量、带宽和访问延迟不同。通常越靠近计算单元,访问越快但容量越小;越远离计算单元,容量越大但访问成本越高。因此,高性能算子需要尽量让数据在靠近计算单元的缓存中复用,减少频繁访问外部内存。\n", | ||
| 55 | - "\n", | ||
| 56 | - "这也是 PyPTO 强调 Tile 编程模型的重要原因:\n", | ||
| 57 | - "\n", | ||
| 58 | - "- 大 Tensor 会被切分成适合片上缓存和并行计算的数据块。\n", | ||
| 59 | - "- 每个 Tile 可以在有限的片上内存中完成搬运、计算和复用。\n", | ||
| 60 | - "- 合理的 Tile 形状会影响并行粒度、内存访问效率和整体性能。\n", | ||
| 61 | - "\n", | ||
| 62 | - "理解内存层次,有助于后续理解 PyPTO 中 TileShape、数据搬运和算子融合等概念。" | ||
| 63 | - ] | ||
| 64 | - }, | ||
| 65 | - { | ||
| 66 | - "cell_type": "markdown", | ||
| 67 | - "metadata": {}, | ||
| 68 | - "source": [ | ||
| 69 | - "## 4. 精度支持\n", | ||
| 70 | - "\n", | ||
| 71 | - "昇腾芯片支持多种数据类型,以适配不同训练和推理场景。常见类型包括 FP32、FP16、BF16、INT8、INT4 等,也包括部分面向低精度计算的格式,如 FP8E4M3、FP8E5M2、FP8E8M0。\n", | ||
| 72 | - "\n", | ||
| 73 | - "不同精度通常对应不同的性能、显存占用和数值稳定性:\n", | ||
| 74 | - "\n", | ||
| 75 | - "- 高精度类型更有利于数值稳定,但计算和存储成本较高。\n", | ||
| 76 | - "- 低精度类型可以降低内存占用、提升吞吐,常用于推理或量化场景。\n", | ||
| 77 | - "- 大模型训练和推理中,常会结合多种精度来平衡性能与精度。\n", | ||
| 78 | - "\n", | ||
| 79 | - "因此,在 PyPTO 算子开发中,数据类型不仅是接口参数,也会影响计算路径、内存占用和性能优化策略。" | ||
| 80 | - ] | ||
| 81 | - }, | ||
| 82 | - { | ||
| 83 | - "cell_type": "markdown", | ||
| 84 | - "metadata": {}, | ||
| 85 | - "source": [ | ||
| 86 | - "## 5. 软件生态与开发支持\n", | ||
| 87 | - "\n", | ||
| 88 | - "昇腾芯片并不是孤立使用的硬件设备,而是配套完整软件生态共同工作:\n", | ||
| 89 | - "\n", | ||
| 90 | - "1. **CANN 软件栈**:提供驱动、固件、Toolkit、Ops 包、编译链和运行时环境。\n", | ||
| 91 | - "2. **PyPTO 框架**:面向昇腾高性能算子开发,提供 Tensor 和 Tile 级别的编程抽象。\n", | ||
| 92 | - "3. **框架适配层**:支持 PyTorch、TensorFlow 等主流深度学习框架在昇腾硬件上运行,例如 PyTorch 可通过 `torch_npu` 扩展接入。\n", | ||
| 93 | - "4. **底层开发工具**:包括 Ascend C、性能分析工具、精度验证工具等。\n", | ||
| 94 | - "\n", | ||
| 95 | - "从应用角度看,昇腾芯片可用于大模型训练、推理服务、边缘计算和云端 AI 服务等场景。对于本教程而言,重点是理解这些硬件能力如何影响 PyPTO 的算子开发方式。" | ||
| 96 | - ] | ||
| 97 | - }, | ||
| 98 | - { | ||
| 99 | - "cell_type": "markdown", | ||
| 100 | - "metadata": {}, | ||
| 101 | - "source": [ | ||
| 102 | - "## 6. 课后练习\n", | ||
| 103 | - "\n", | ||
| 104 | - "本节练习用于检查昇腾芯片、AI Core、内存层次和 Tile 思想。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | ||
| 105 | - "\n", | ||
| 106 | - "1. (选择题)昇腾芯片在 PyPTO 算子开发链路中主要提供什么? \n", | ||
| 107 | - " A. Markdown 渲染能力 \n", | ||
| 108 | - " B. AI 计算所需的硬件算力 \n", | ||
| 109 | - " C. Python 包管理能力 \n", | ||
| 110 | - " D. 数据可视化能力\n", | ||
| 111 | - "\n", | ||
| 112 | - "2. (填空题)昇腾芯片面向 AI 计算的典型特点包括\\_\\_\\_、\\_\\_\\_、\\_\\_\\_和\\_\\_\\_?\n", | ||
| 113 | - "\n", | ||
| 114 | - "3. (选择题)AI Core 与 Tile 切分的关系是什么? \n", | ||
| 115 | - " A. 大 Tensor 可切成多个 Tile,不同 AI Core 可以并行处理不同数据块 \n", | ||
| 116 | - " B. Tile 会删除所有输入数据 \n", | ||
| 117 | - " C. AI Core 只能处理字符串 \n", | ||
| 118 | - " D. Tile 与并行执行无关\n", | ||
| 119 | - "\n", | ||
| 120 | - "4. (填空题)昇腾芯片的多级内存层次中,教程提到的示例包括\\_\\_\\_、\\_\\_\\_、\\_\\_\\_、\\_\\_\\_和\\_\\_\\_?\n", | ||
| 121 | - "\n", | ||
| 122 | - "5. (选择题)理解内存层次对学习 PyPTO 有什么帮助? \n", | ||
| 123 | - " A. 有助于理解 TileShape、数据搬运和算子融合等概念 \n", | ||
| 124 | - " B. 可以完全不验证数值结果 \n", | ||
| 125 | - " C. 可以让所有算子不需要输入 \n", | ||
| 126 | - " D. 可以把 NPU 当成普通硬盘使用\n", | ||
| 127 | - "\n", | ||
| 128 | - "**执行以下代码获取答案。**\n" | ||
| 129 | - ] | ||
| 130 | - }, | ||
| 131 | - { | ||
| 132 | - "cell_type": "code", | ||
| 133 | - "execution_count": null, | ||
| 134 | - "metadata": {}, | ||
| 135 | - "outputs": [], | ||
| 136 | - "source": [ | ||
| 137 | - "!cat ./answer/01.03_answer.txt\n" | ||
| 138 | - ] | ||
| 139 | - }, | ||
| 140 | - { | ||
| 141 | - "cell_type": "markdown", | ||
| 142 | - "metadata": {}, | ||
| 143 | - "source": [ | ||
| 144 | - "## 7. 本节小结\n", | ||
| 145 | - "\n", | ||
| 146 | - "昇腾芯片提供 AI 计算的硬件基础。AI Core、多级内存、多精度支持和 MPMD 执行模式共同决定了算子开发需要关注并行、数据搬运、内存复用和数据类型选择。\n", | ||
| 147 | - "\n", | ||
| 148 | - "下一节将介绍 PyPTO。理解了芯片的这些基本特征后,就更容易明白为什么 PyPTO 要围绕 Tensor、Tile、多层级计算图和自动代码生成来设计。" | ||
| 149 | - ] | ||
| 150 | - } | ||
| 151 | - ], | ||
| 152 | - "metadata": { | ||
| 153 | - "kernelspec": { | ||
| 154 | - "display_name": "Python 3", | ||
| 155 | - "language": "python", | ||
| 156 | - "name": "python3" | ||
| 157 | - }, | ||
| 158 | - "language_info": { | ||
| 159 | - "name": "python", | ||
| 160 | - "version": "3.10.0" | ||
| 161 | - } | ||
| 162 | - }, | ||
| 163 | - "nbformat": 4, | ||
| 164 | - "nbformat_minor": 4 | ||
| 165 | -} | ||
感觉就是全是文字堆叠,没有任何突出,就一点一点的堆砌有点干瘪,而且全是一条一条的


| @@ -6,97 +6,177 @@ | |||
| 6 | "source": [ | 6 | "source": [ |
| 7 | "# PyPTO 是什么\n", | 7 | "# PyPTO 是什么\n", |
| 8 | "\n", | 8 | "\n", |
| 9 | - "PyPTO(发音:pai p-t-o)是 CANN 推出的高性能编程框架,主要面向昇腾 AI 加速器上的算子开发和复杂模型组件优化。它的目标是在保留高性能优化空间的同时,降低开发者直接面对底层硬件指令、内存搬运和调度细节的成本。\n", | 9 | + "PyPTO是CANN推出的一款面向AI加速器的高效编程框架,旨在简化算子开发流程,同时保持高性能计算能力。该框架采用创新的PTO(Parallel Tensor/Tile Operation)编程范式,以基于Tile的编程模型为核心设计理念,通过多层次的计算图表达,将用户通过API构建的AI模型从高层次的Tensor计算图逐步编译成硬件指令,最终生成可在目标平台上高效执行的代码,并由设备侧以MPMD(Multiple Program Multiple Data)方式调度执行。" |
| 10 | - "\n", | ||
| 11 | - "从定位上看,PyPTO 不是一个通用深度学习训练框架,而是更靠近“高性能算子开发工具”。开发者可以使用 Python 风格的接口描述 Tensor 级计算,再由 PyPTO 和 CANN 逐步完成图优化、Tile 化、代码生成和运行时调度。" | ||
| 12 | ] | 10 | ] |
| 13 | }, | 11 | }, |
| 14 | { | 12 | { |
| 15 | "cell_type": "markdown", | 13 | "cell_type": "markdown", |
| 16 | "metadata": {}, | 14 | "metadata": {}, |
| 17 | "source": [ | 15 | "source": [ |
| 18 | - "## 1. 为什么需要 PyPTO\n", | 16 | + "## 1. 核心架构\n", |
| 19 | "\n", | 17 | "\n", |
| 20 | - "在实际模型开发中,常见算子通常可以直接使用框架或算子库提供的实现。但当开发者需要实现自定义算子、融合多个计算步骤、优化大模型组件,或者处理动态 Shape 时,通用算子库往往不够灵活。\n", | 18 | + "PyPTO框架采用分层架构设计,从用户API到底层硬件执行,共分为以下几个层次:\n", |
| 21 | "\n", | 19 | "\n", |
| 22 | - "直接使用底层语言开发算子虽然灵活,但需要处理大量硬件相关细节,例如数据如何切分、如何搬运到片上内存、如何映射到多个 AI Core 上执行。PyPTO 的价值就在于提供更高层的 Tensor 和 Tile 抽象,让开发者可以先表达“算什么”,再由框架和 CANN 共同完成“怎么高效执行”。" | 20 | + "### 1.1 用户接口层(User Interface Layer)\n", |
| 21 | + "PyPTO框架与开发者交互的接口层,提供Python友好的编程接口,使开发者能够以直观的方式表达计算逻辑,而无需深入了解底层硬件实现细节。\n", | ||
| 22 | + "\n", | ||
| 23 | + "### 1.2 计算图编译层(Compute Graph Compile Layer) \n", | ||
| 24 | + "PyPTO采用多层级计算图表达,支持从高到低多个抽象层次的计算图的优化和转换。 \n", | ||
| 25 | + " Tensor Graph:高层次的Tensor操作,贴近算法设计者的数学表达式。 \n", | ||
| 26 | + " Tile Graph:硬件感知的Tile操作,充分利用硬件并行性和内存层次结构。 \n", | ||
| 27 | + " Block Graph:子图分区,支持并行执行和资源管理。 \n", | ||
| 28 | + " Execute Graph:执行图,包含依赖关系和调度信息。 \n", | ||
| 29 | + "\n", | ||
| 30 | + "编译过程通过模块化的Pass实现,每个阶段由多个Pass组成,负责特定阶段优化或转换任务。 \n", | ||
| 31 | + " Tensor Graph阶段:实现和硬件无关的图优化,包括冗余操作消除、类型转换、内存冲突推断等。 \n", | ||
| 32 | + " Tile Graph阶段:根据TileShape进行Tile展开,实现Tile级别的优化,包括内存类型分配、移动操作生成、子图切分等。 \n", | ||
| 33 | + " Block Graph阶段:切分生成计算子图,进行Block级别的优化,包括乱序调度、内存重用、同步点插入等。 \n", | ||
| 34 | + " Execute Graph阶段:整合计算子图信息,编排生成最终的执行图。 \n", | ||
| 35 | + "\n", | ||
| 36 | + "### 1.3 代码生成层(Code Generation Layer)\n", | ||
| 37 | + "将优化后的计算图转换为目标平台的可执行代码。\n", | ||
| 38 | + " 虚拟指令生成:从Execute Graph生成PTO虚拟指令代码(PTO Virtual Instructions)。\n", | ||
| 39 | + " 目标平台编译:将虚拟指令编译为目标平台代码。\n", | ||
| 40 | + "\n", | ||
| 41 | + "### 1.4 调度执行层(Scheduling & Execution Layer)\n", | ||
| 42 | + "负责将可执行代码在设备上调度执行。\n", | ||
| 43 | + " MPMD调度:可执行代码在设备上通过MPMD方式调度到设备处理器核。\n", | ||
| 44 | + " 控制流执行:管理任务依赖关系,执行控制流逻辑。\n", | ||
| 45 | + "\n" | ||
| 23 | ] | 46 | ] |
| 24 | }, | 47 | }, |
| 25 | { | 48 | { |
| 26 | "cell_type": "markdown", | 49 | "cell_type": "markdown", |
| 27 | "metadata": {}, | 50 | "metadata": {}, |
| 28 | "source": [ | 51 | "source": [ |
| 29 | - "## 2. 核心思想:Tensor 表达,Tile 执行\n", | 52 | + "## 2. 核心特性\n", |
| 30 | "\n", | 53 | "\n", |
| 31 | - "PyPTO 采用 PTO(Parallel Tensor/Tile Operation)编程范式。可以简单理解为:在高层用 Tensor 表达算法逻辑,在底层用 Tile 适配硬件执行。\n", | 54 | + "基于Tile的编程模型:计算基于Tile(硬件感知的数据块)进行,充分利用硬件的并行计算能力和内存层次结构。\n", |
| 32 | "\n", | 55 | "\n", |
| 33 | - "Tensor 层更接近算法开发者熟悉的数学表达,关注输入输出、计算关系和操作组合。Tile 层则更接近硬件执行,关注大 Tensor 如何被切分成数据块,数据块如何在片上内存中搬运和复用,以及如何分配给多个 AI Core 并行执行。\n", | 56 | + "多层级计算图表达和优化:通过计算图编译层将Tensor Graph转换为Tile Graph、Block Graph和Execute Graph,每一步都包含一系列Pass优化流程。\n", |
| 34 | "\n", | 57 | "\n", |
| 35 | - "因此,Tile 不是普通意义上的数据切片,而是连接算法表达和硬件执行的重要抽象。前面介绍的芯片内存层次、AI Core 并行和 MPMD 执行模式,都会影响 Tile 的组织方式。" | 58 | + "自动化代码生成:编译结果通过代码生成层生成PTO虚拟指令代码,然后通过编译器编译为目标平台的可执行代码。\n", |
| 59 | + "\n", | ||
| 60 | + "MPMD执行调度:可执行代码被加载到设备侧,通过MPMD方式调度到设备处理器核,实现高效并行执行。\n", | ||
| 61 | + "\n", | ||
| 62 | + "完整的工具链支持:全流程的编译中间产物和运行时性能数据可以通过IDE集成的工具链进行可视化,以便识别性能瓶颈。开发者还可以通过工具链控制编译和调度行为。\n", | ||
| 63 | + "\n", | ||
| 64 | + "Python友好API:提供直观的Tensor级别抽象,贴近算法开发者的思维模式,支持动态Shape和符号化编程。\n", | ||
| 65 | + "\n", | ||
| 66 | + "分层抽象设计:对不同开发者暴露不同的抽象层次,算法开发者使用Tensor层次,性能专家使用Tile层次,系统开发者使用Block层次。" | ||
| 36 | ] | 67 | ] |
| 37 | }, | 68 | }, |
| 38 | { | 69 | { |
| 39 | "cell_type": "markdown", | 70 | "cell_type": "markdown", |
| 40 | "metadata": {}, | 71 | "metadata": {}, |
| 41 | "source": [ | 72 | "source": [ |
| 42 | - "## 3. 从代码到芯片执行\n", | 73 | + "## 3. 总览和适用场景\n", |
| 43 | "\n", | 74 | "\n", |
| 44 | - "PyPTO 会把开发者编写的高层计算表达逐步转换为硬件可执行任务。这个过程可以概括为以下链路:\n", | 75 | + "<img src=\"./images/pto_overview.png\" alt=\"pto_overview\" width=\"700px\" >\n", |
| 45 | "\n", | 76 | "\n", |
| 46 | - "```text\n", | 77 | + "使用场景\n", |
| 47 | - "Python API / Tensor 表达\n", | 78 | + " - 深度学习算子开发:快速实现各种神经网络算子。\n", |
| 48 | - " -> Tensor Graph:描述高层计算逻辑\n", | 79 | + " - 大模型开发:支持Attention、MoE、FFN等大模型组件。\n", |
| 49 | - " -> Tile Graph:加入 Tile 切分和硬件感知信息\n", | 80 | + " - 动态Shape处理:支持动态Batch Size等动态Shape场景。" |
| 50 | - " -> Block Graph:组织可并行执行的任务块\n", | ||
| 51 | - " -> Execute Graph:形成最终调度和执行计划\n", | ||
| 52 | - " -> CANN 编译与运行时:生成并调度 NPU 可执行任务\n", | ||
| 53 | - "```\n", | ||
| 54 | - "\n", | ||
| 55 | - "这个分层过程的意义在于:开发者不需要一开始就编写底层指令,但框架仍然可以在编译过程中逐步加入硬件信息,完成内存优化、并行调度和代码生成。" | ||
| 56 | ] | 81 | ] |
| 57 | }, | 82 | }, |
| 58 | { | 83 | { |
| 59 | "cell_type": "markdown", | 84 | "cell_type": "markdown", |
| 60 | "metadata": {}, | 85 | "metadata": {}, |
| 61 | "source": [ | 86 | "source": [ |
| 62 | - "## 4. 适用场景\n", | 87 | + "## 4. 设计理念\n", |
| 63 | "\n", | 88 | "\n", |
| 64 | - "PyPTO 适合需要兼顾开发效率和硬件性能的场景。典型例子包括自定义算子开发、融合算子开发、大模型中的 Attention、MoE、FFN 等组件优化,以及动态 Batch Size、动态序列长度等动态 Shape 场景。\n", | 89 | + "传统的模型开发通常分为算法开发人员和算子开发人员。这种分工的根源在于高性能算子开发的复杂性:算子开发人员不仅需要理解算子的数学计算属性,还必须考虑如何将其转换为对硬件友好的执行方式。这类似于早期CPU时代,在乱序执行和编译器技术尚未成熟时,程序员需要手动安排流水线指令。\n", |
| 65 | "\n", | 90 | "\n", |
| 66 | - "这些场景通常既需要比通用算子库更灵活的表达能力,又需要比纯底层开发更高的开发效率。PyPTO 正是处在这个中间位置:它让开发者用较高层的方式描述计算,同时保留面向昇腾硬件的优化路径。" | 91 | + "为了降低这种复杂性,PyPTO提出了一种新的编程框架设计理念,旨在简化算子开发流程,同时保留高性能计算的潜力。\n", |
| 92 | + "\n", | ||
| 93 | + "### 4.1 计算层设计\n", | ||
| 94 | + "\n", | ||
| 95 | + "计算层的设计理念是尽可能贴近算法设计者的数学表达式,使用Tensor而非单个元素来描述计算过程。用户通过API构建的AI模型通过Tensor Graph表达,这种设计保留了最大化的优化潜力,包括:\n", | ||
| 96 | + "\n", | ||
| 97 | + "内存布局优化:自动优化数据在内存中的排布方式\n", | ||
| 98 | + "\n", | ||
| 99 | + "数据搬运优化:最小化数据在不同内存层次间的传输\n", | ||
| 100 | + "\n", | ||
| 101 | + "多算子联合优化:识别并融合可优化的算子组合\n", | ||
| 102 | + "\n", | ||
| 103 | + "通过Tensor作为基本数据单位,计算层能够更自然地表达复杂的数学运算,同时为后续的编译优化提供丰富的信息。\n", | ||
| 104 | + "\n", | ||
| 105 | + "### 4.2 编译层设计\n", | ||
| 106 | + "\n", | ||
| 107 | + "编译层是连接计算层和执行层的关键环节,负责将Tensor Graph转换为硬件友好的执行形式。编译过程通过多阶段的Lowering Pipeline实现:\n", | ||
| 108 | + "\n", | ||
| 109 | + "Tensor Graph到Tile Graph:通过编译Pass将Tensor操作转换为Tile操作,选择Tiling策略,进行布局变换、Tile融合、Tile重排序等。\n", | ||
| 110 | + "\n", | ||
| 111 | + "Tile Graph到Block Graph:将Tile图分区为子图,检测同构子图,规范化Block Graph,追踪依赖关系。\n", | ||
| 112 | + "\n", | ||
| 113 | + "Block Graph到Execute Graph:构建执行图,分析Block Graph之间的依赖关系,规划全局资源,生成调度提示。\n", | ||
| 114 | + "\n", | ||
| 115 | + "每个阶段都包含多个优化Pass,通过模块化的图变换和优化流程,将计算层保留的优化空间转化为实际性能提升。\n", | ||
| 116 | + "\n", | ||
| 117 | + "编译层提供了以下核心能力:\n", | ||
| 118 | + "\n", | ||
| 119 | + "快速可用:保证第一时间生成可运行的结果,满足快速开发的需求。\n", | ||
| 120 | + "\n", | ||
| 121 | + "灵活调优:支持性能敏感的配置调整,方便开发者根据实际需求进行优化。\n", | ||
| 122 | + "\n", | ||
| 123 | + "深度优化:允许高级用户深度定制编译流程,追求极致性能。\n", | ||
| 124 | + "\n", | ||
| 125 | + "### 4.3 执行层设计\n", | ||
| 126 | + "\n", | ||
| 127 | + "执行层负责将编译后的代码转换为硬件友好的指令并执行。执行过程包括:\n", | ||
| 128 | + "\n", | ||
| 129 | + "代码生成:编译结果通过CodeGen生成底层PTO虚拟指令代码。\n", | ||
| 130 | + "\n", | ||
| 131 | + "目标平台编译:通过编译器将虚拟指令代码编译成目标NPU平台的可执行代码。\n", | ||
| 132 | + "\n", | ||
| 133 | + "MPMD调度:可执行代码被加载到设备侧,通过MPMD方式调度到设备上的处理器核。\n", | ||
| 134 | + "\n", | ||
| 135 | + "通过自动化代码生成技术,执行层能够根据硬件特性自动生成最优的执行指令,充分释放硬件算力。这种设计避免了传统算子开发中手动调整硬件指令的复杂性,同时确保了高性能计算的实现。\n", | ||
| 136 | + "\n", | ||
| 137 | + "### 4.4 工具链设计\n", | ||
| 138 | + "\n", | ||
| 139 | + "PyPTO提供了完整的工具链支持,包括:\n", | ||
| 140 | + "\n", | ||
| 141 | + "编译中间产物可视化:支持在编译的不同阶段(如Tensor Graph、Tile Graph、Block Graph、Execute Graph等)保存中间产物(计算图),便于调试和分析。\n", | ||
| 142 | + "\n", | ||
| 143 | + "运行时性能分析:收集并可视化运行时性能数据(泳道图),帮助识别性能瓶颈。\n", | ||
| 144 | + "\n", | ||
| 145 | + "编译和调度控制:开发者可以通过工具链控制编译Pass的执行和调度行为,实现深度定制。\n", | ||
| 146 | + "\n", | ||
| 147 | + "通过上述设计理念,PyPTO实现了算法开发与算子开发的高效协同,显著降低了算子开发的复杂性,同时保留了高性能计算的能力。" | ||
| 67 | ] | 148 | ] |
| 68 | }, | 149 | }, |
| 69 | { | 150 | { |
| 70 | "cell_type": "markdown", | 151 | "cell_type": "markdown", |
| 71 | "metadata": {}, | 152 | "metadata": {}, |
| 72 | "source": [ | 153 | "source": [ |
| 73 | - "## 5. 课后练习\n", | 154 | + "## 5. 课后练习\n" |
| 155 | + ] | ||
| 156 | + }, | ||
| 157 | + { | ||
| 158 | + "cell_type": "markdown", | ||
| 159 | + "metadata": {}, | ||
| 160 | + "source": [ | ||
| 161 | + "1. (选择题)PyPTO 框架的分层架构从用户 API 到底层硬件执行,以下哪个排列顺序是正确的?\n", | ||
| 162 | + "A. 用户接口层 → 调度执行层 → 计算图编译层 → 代码生成层\n", | ||
| 163 | + "B. 用户接口层 → 计算图编译层 → 代码生成层 → 调度执行层\n", | ||
| 164 | + "C. 计算图编译层 → 用户接口层 → 代码生成层 → 调度执行层\n", | ||
| 165 | + "D. 代码生成层 → 计算图编译层 → 用户接口层 → 调度执行层\n", | ||
| 74 | "\n", | 166 | "\n", |
| 75 | - "本节练习用于检查 PyPTO 的定位、PTO 编程范式和从代码到芯片执行的转换链路。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | 167 | + "2. 关于 PyPTO 的分层抽象设计,以下描述正确的是?\n", |
| 168 | + "A. 所有开发者都必须使用 Tile 层次进行开发\n", | ||
| 169 | + "B. 算法开发者使用 Tensor 层次,性能专家使用 Tile 层次,系统开发者使用 Block 层次\n", | ||
| 170 | + "C. 算法开发者使用 Block 层次,性能专家使用 Tensor 层次\n", | ||
| 171 | + "D. 分层抽象仅用于编译优化,与开发者角色无关\n", | ||
| 76 | "\n", | 172 | "\n", |
| 77 | - "1. (选择题)PyPTO 的定位更接近哪一种? \n", | 173 | + "3. (选择题)关于 PyPTO 编译层设计中 Block Graph 阶段负责的优化任务,以下描述正确的是?\n", |
| 78 | - " A. 通用网页开发框架 \n", | 174 | + "A. 根据 TileShape 进行 Tile 展开和内存类型分配\n", |
| 79 | - " B. 面向昇腾 AI 加速器的高性能算子开发框架 \n", | 175 | + "B. 实现冗余操作消除、类型转换和内存冲突推断\n", |
| 80 | - " C. 数据库查询语言 \n", | 176 | + "C. 检测同构子图、规范化 Block Graph 并追踪依赖关系\n", |
| 81 | - " D. 图片压缩工具\n", | 177 | + "D. 整合计算子图信息,编排生成最终的执行图\n", |
| 82 | "\n", | 178 | "\n", |
| 83 | - "2. (填空题)PyPTO 的核心思想可以概括为:用________表达算法逻辑,用________连接硬件执行。\n", | 179 | + "**执行以下代码获取答案。**" |
| 84 | - "\n", | ||
| 85 | - "3. (选择题)为什么需要 PyPTO? \n", | ||
| 86 | - " A. 在保留高性能优化空间的同时,降低开发者直接处理底层硬件细节的成本 \n", | ||
| 87 | - " B. 为了替代所有 Python 语法 \n", | ||
| 88 | - " C. 为了让算子只能在 CPU 上运行 \n", | ||
| 89 | - " D. 为了删除 CANN 软件栈\n", | ||
| 90 | - "\n", | ||
| 91 | - "4. (填空题)PyPTO 从代码到芯片执行的链路中,教程提到的图层级包括\\_\\_\\_、\\_\\_\\_、\\_\\_\\_和\\_\\_\\_。\n", | ||
| 92 | - "\n", | ||
| 93 | - "5. (选择题)Tile 在 PyPTO 中的作用是什么? \n", | ||
| 94 | - " A. 连接算法表达和硬件执行的重要抽象 \n", | ||
| 95 | - " B. 只用于给变量改名 \n", | ||
| 96 | - " C. 只用于打印日志 \n", | ||
| 97 | - " D. 只用于删除 Tensor 维度\n", | ||
| 98 | - "\n", | ||
| 99 | - "**执行以下代码获取答案。**\n" | ||
| 100 | ] | 180 | ] |
| 101 | }, | 181 | }, |
| 102 | { | 182 | { |
| @@ -105,7 +185,7 @@ | |||
| 105 | "metadata": {}, | 185 | "metadata": {}, |
| 106 | "outputs": [], | 186 | "outputs": [], |
| 107 | "source": [ | 187 | "source": [ |
| 108 | - "!cat ./answer/01.04_answer.txt\n" | 188 | + "!cat ./answer/01.04_answer.txt" |
| 109 | ] | 189 | ] |
| 110 | }, | 190 | }, |
| 111 | { | 191 | { |
| @@ -114,9 +194,7 @@ | |||
| 114 | "source": [ | 194 | "source": [ |
| 115 | "## 6. 本节小结\n", | 195 | "## 6. 本节小结\n", |
| 116 | "\n", | 196 | "\n", |
| 117 | - "PyPTO 是面向昇腾 AI 加速器的高性能算子开发框架。它用 Tensor 表达算法逻辑,用 Tile 连接硬件执行,并通过多层级计算图、自动代码生成和 CANN 软件栈完成从高层描述到芯片执行的转换。\n", | 197 | + "PyPTO 的核心设计目标是打破传统算法开发与算子开发的分工壁垒——让算法开发者用 Tensor 级别的数学表达式描述计算,由框架的编译层自动完成从 Tensor Graph 到硬件指令的全链路优化和调度,同时通过工具链提供可视化调试与性能分析能力。后续章节将进入实际的 API 使用和算子开发实践。" |
| 118 | - "\n", | ||
| 119 | - "到这里,本章已经从三个层次建立了基础认知:昇腾芯片提供硬件算力,CANN 提供软件栈支撑,PyPTO 提供面向开发者的编程模型。后续学习算子开发时,可以始终沿着这条链路理解每一个概念的位置和作用。" | ||
| 120 | ] | 198 | ] |
| 121 | } | 199 | } |
| 122 | ], | 200 | ], |
| @@ -4,7 +4,7 @@ | |||
| 4 | "cell_type": "markdown", | 4 | "cell_type": "markdown", |
| 5 | "metadata": {}, | 5 | "metadata": {}, |
| 6 | "source": [ | 6 | "source": [ |
| 7 | - "# 01.05 章节实践\n", | 7 | + "# 章节实践\n", |
| 8 | "\n", | 8 | "\n", |
| 9 | "本节用于完成第一章的综合自测。完成练习时,可以把重点放在三件事上:CANN、昇腾芯片和 PyPTO 分别处在什么层级;它们在算子开发链路中如何协作;为什么 PyPTO 要围绕 Tensor、Tile 和多层级计算图组织计算。\n" | 9 | "本节用于完成第一章的综合自测。完成练习时,可以把重点放在三件事上:CANN、昇腾芯片和 PyPTO 分别处在什么层级;它们在算子开发链路中如何协作;为什么 PyPTO 要围绕 Tensor、Tile 和多层级计算图组织计算。\n" |
| 10 | ] | 10 | ] |
| @@ -1,5 +1,2 @@ | |||
| 1 | -1. B | 1 | +1.B |
| 2 | -2. 驱动固件;Toolkit;Ops 包。 | 2 | +2.C |
| 3 | -3. B | ||
| 4 | -4. 编译优化;运行时调度。 | ||
| 5 | -5. A | ||
| @@ -1,5 +1,5 @@ | |||
| 1 | 1. B | 1 | 1. B |
| 2 | -2. 专用计算单元;多核并行能力;多级内存层次;多精度支持。 | 2 | +2. 驱动固件;Toolkit;Ops 包。 |
| 3 | -3. A | 3 | +3. B |
| 4 | -4. L0 Buffer;L1 Buffer;Unified Buffer;Global Memory;片外 HBM。 | 4 | +4. 编译优化;运行时调度。 |
| 5 | -5. A | 5 | +5. A |
| @@ -1,5 +1,3 @@ | |||
| 1 | 1. B | 1 | 1. B |
| 2 | -2. Tensor;Tile。 | 2 | +2. B |
| 3 | -3. A | 3 | +3. C |
| 4 | -4. Tensor Graph;Tile Graph;Block Graph;Execute Graph。 | ||
| 5 | -5. A | ||
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "b5680a57", | 5 | "id": "b5680a57", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 02.01 PyPTO 算子开发基础知识\n", | 8 | + "# PyPTO 算子开发基础知识\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "本章建立 PyPTO 算子开发的基础认知。PyPTO 算子可以先理解为“用 Python 语法描述一段 Tensor 计算,再由框架把这段计算交给设备执行”的程序单元。后续初级教程会继续展开逐元素算子、矩阵乘法、规约和复杂算子封装;本章先从最小闭环开始,说明一个算子如何定义、如何准备输入输出、如何运行、如何验证。\n", | 10 | "本章建立 PyPTO 算子开发的基础认知。PyPTO 算子可以先理解为“用 Python 语法描述一段 Tensor 计算,再由框架把这段计算交给设备执行”的程序单元。后续初级教程会继续展开逐元素算子、矩阵乘法、规约和复杂算子封装;本章先从最小闭环开始,说明一个算子如何定义、如何准备输入输出、如何运行、如何验证。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -42,10 +42,10 @@ | |||
| 42 | "| --- | --- | --- |\n", | 42 | "| --- | --- | --- |\n", |
| 43 | "| 明确计算目标 | 写清楚 Tensor 之间的数学关系 | `out = x + y` |\n", | 43 | "| 明确计算目标 | 写清楚 Tensor 之间的数学关系 | `out = x + y` |\n", |
| 44 | "| 定义 kernel | 标记一段需要被 PyPTO 处理的函数 | `@pypto.frontend.jit(...)` |\n", | 44 | "| 定义 kernel | 标记一段需要被 PyPTO 处理的函数 | `@pypto.frontend.jit(...)` |\n", |
| 45 | - "| 描述输入输出 | 说明参数的 dtype 和 shape 特征 | `pypto.Tensor([...], pypto.DT_FP32)` |\n", | 45 | + "| 描述输入输出 | 说明参数的 dtype 和 shape 特征 | `pypto.Tensor()` |\n", |
| 46 | - "| 设置 Tile Shape | 描述设备侧分块组织方式 | `pypto.set_vec_tile_shapes(1, 4, 1, 64)` |\n", | 46 | + "| 设置 Tile Shape | 描述设备侧分块组织方式 | `pypto.set_vec_tile_shapes()` |\n", |
| 47 | "| 准备数据 | 创建真实输入和输出 Tensor | `torch.rand`、`torch.empty` |\n", | 47 | "| 准备数据 | 创建真实输入和输出 Tensor | `torch.rand`、`torch.empty` |\n", |
| 48 | - "| 调用并验证 | 执行 kernel 并比较参考结果 | `assert_allclose` |\n", | 48 | + "| 调用并验证 | 执行 kernel 并比较参考结果 | `assert_close` |\n", |
| 49 | "\n", | 49 | "\n", |
| 50 | "后续更复杂的算子仍然沿着这条主线展开,只是计算表达式、shape 关系、dtype 选择和 Tile 设置会更加丰富。\n" | 50 | "后续更复杂的算子仍然沿着这条主线展开,只是计算表达式、shape 关系、dtype 选择和 Tile 设置会更加丰富。\n" |
| 51 | ] | 51 | ] |
| @@ -55,16 +55,16 @@ | |||
| 55 | "id": "fcc941e6", | 55 | "id": "fcc941e6", |
| 56 | "metadata": {}, | 56 | "metadata": {}, |
| 57 | "source": [ | 57 | "source": [ |
| 58 | - "## 2. 本章 Notebook 关系\n", | 58 | + "## 2. 各章节介绍\n", |
| 59 | "\n", | 59 | "\n", |
| 60 | "本章各节承担不同角色:\n", | 60 | "本章各节承担不同角色:\n", |
| 61 | "\n", | 61 | "\n", |
| 62 | "| Notebook | 作用 |\n", | 62 | "| Notebook | 作用 |\n", |
| 63 | "| --- | --- |\n", | 63 | "| --- | --- |\n", |
| 64 | - "| `02.01_chapter_intro.ipynb` | 建立章节入口、开发主线和最小运行时检查 |\n", | 64 | + "| `02.01_chapter_intro.ipynb` | 建立章节入口、开发主线 |\n", |
| 65 | "| `02.02_run_hello_world.ipynb` | 完整展开 Hello World 加法算子 |\n", | 65 | "| `02.02_run_hello_world.ipynb` | 完整展开 Hello World 加法算子 |\n", |
| 66 | "| `02.03_programming_paradigm_mpmd.ipynb` | 解释 Host/Kernel 划分、PyPTO API、计算图和 MPMD |\n", | 66 | "| `02.03_programming_paradigm_mpmd.ipynb` | 解释 Host/Kernel 划分、PyPTO API、计算图和 MPMD |\n", |
| 67 | - "| `02.04_api_and_compute_graph.ipynb` | 汇总常用 API 分类,并把计算图层次和开发闭环串联起来 |\n", | 67 | + "| `02.04_graph_and_swimlane.ipynb` | 介绍如何生成和查看计算图和泳道图 |\n", |
| 68 | "| `02.05_chapter_test.ipynb` | 通过选择、填空和综合编程检查本章掌握情况 |\n", | 68 | "| `02.05_chapter_test.ipynb` | 通过选择、填空和综合编程检查本章掌握情况 |\n", |
| 69 | "\n", | 69 | "\n", |
| 70 | "这个顺序体现了从实践到概念、再回到整体框架的学习路径:先看到一个完整算子,再解释它背后的编程范式,最后建立 API 和计算图的索引,并通过章节实践确认自己能独立完成一个基础算子的改写与验证。\n" | 70 | "这个顺序体现了从实践到概念、再回到整体框架的学习路径:先看到一个完整算子,再解释它背后的编程范式,最后建立 API 和计算图的索引,并通过章节实践确认自己能独立完成一个基础算子的改写与验证。\n" |
| @@ -72,67 +72,9 @@ | |||
| 72 | }, | 72 | }, |
| 73 | { | 73 | { |
| 74 | "cell_type": "markdown", | 74 | "cell_type": "markdown", |
| 75 | - "id": "66f7072a", | ||
| 76 | "metadata": {}, | 75 | "metadata": {}, |
| 77 | "source": [ | 76 | "source": [ |
| 78 | - "## 3. 最小运行时检查\n", | 77 | + "## 3. 课后练习\n", |
| 79 | - "\n", | ||
| 80 | - "本章入口只需要确认 Python 环境能够导入 `torch` 、 `pypto`和`torch_npu` 扩展。当前在线环境如果没有pypto,可以使用以下命令安装。如果在执行多个章节代码时报错`maybe others are using current device`,可以使用命令`npu-smi info`查看当前npu上的进程,使用kill命令移除掉之前的进程。\n" | ||
| 81 | - ] | ||
| 82 | - }, | ||
| 83 | - { | ||
| 84 | - "cell_type": "code", | ||
| 85 | - "execution_count": null, | ||
| 86 | - "id": "98684011", | ||
| 87 | - "metadata": {}, | ||
| 88 | - "outputs": [], | ||
| 89 | - "source": [ | ||
| 90 | - "!pip install pypto\n" | ||
| 91 | - ] | ||
| 92 | - }, | ||
| 93 | - { | ||
| 94 | - "cell_type": "code", | ||
| 95 | - "execution_count": null, | ||
| 96 | - "id": "fb67172c", | ||
| 97 | - "metadata": {}, | ||
| 98 | - "outputs": [], | ||
| 99 | - "source": [ | ||
| 100 | - "import os\n", | ||
| 101 | - "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | ||
| 102 | - "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 103 | - "import torch\n", | ||
| 104 | - "import torch_npu\n", | ||
| 105 | - "import pypto\n", | ||
| 106 | - "\n", | ||
| 107 | - "def describe_runtime():\n", | ||
| 108 | - " print(\"python runtime ready\")\n", | ||
| 109 | - " print(\"torch:\", torch.__version__)\n", | ||
| 110 | - " print(\"pypto:\", pypto.__file__)\n", | ||
| 111 | - " print(\"TILE_FWK_DEVICE_ID:\", os.environ.get(\"TILE_FWK_DEVICE_ID\", \"<未设置,默认 0>\"))\n", | ||
| 112 | - " device_id = int(os.environ.get(\"TILE_FWK_DEVICE_ID\", \"0\"))\n", | ||
| 113 | - " device = f\"npu:{device_id}\"\n", | ||
| 114 | - " print(\"default device:\", device)\n", | ||
| 115 | - " return device\n", | ||
| 116 | - "\n", | ||
| 117 | - "\n", | ||
| 118 | - "device = describe_runtime()\n" | ||
| 119 | - ] | ||
| 120 | - }, | ||
| 121 | - { | ||
| 122 | - "cell_type": "markdown", | ||
| 123 | - "id": "e4d5cdeb", | ||
| 124 | - "metadata": {}, | ||
| 125 | - "source": [ | ||
| 126 | - "这段检查只回答三个问题:`torch` 是否可用,`pypto` 是否可用,当前环境是否带有 `torch_npu`。如果 `torch_npu` 存在,设备编号默认来自 `TILE_FWK_DEVICE_ID`,未设置时使用进程内的 `0`。\n", | ||
| 127 | - "\n", | ||
| 128 | - "需要注意的是,`TILE_FWK_DEVICE_ID=0` 表示当前进程可见设备中的第 0 张,并不一定等同于物理机器上的第 0 张卡。可见设备范围通常由外部环境决定。\n" | ||
| 129 | - ] | ||
| 130 | - }, | ||
| 131 | - { | ||
| 132 | - "cell_type": "markdown", | ||
| 133 | - "metadata": {}, | ||
| 134 | - "source": [ | ||
| 135 | - "## 4. 课后练习\n", | ||
| 136 | "\n", | 78 | "\n", |
| 137 | "本节练习用于检查 PyPTO 算子开发基础章节的学习主线。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | 79 | "本节练习用于检查 PyPTO 算子开发基础章节的学习主线。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", |
| 138 | "\n", | 80 | "\n", |
| @@ -142,7 +84,7 @@ | |||
| 142 | " C. 只创建输出 Tensor,不需要输入 \n", | 84 | " C. 只创建输出 Tensor,不需要输入 \n", |
| 143 | " D. 只运行 `pip install` 就完成算子开发\n", | 85 | " D. 只运行 `pip install` 就完成算子开发\n", |
| 144 | "\n", | 86 | "\n", |
| 145 | - "2. (填空题)本章知识路径中,计算目标之后的三个关键步骤是________、________和________。\n", | 87 | + "2. (填空题)本章知识路径中,计算目标之后的三个关键步骤是\\_\\_\\_\\_\\_\\_、\\_\\_\\_\\_\\_\\_和\\_\\_\\_\\_\\_\\_。\n", |
| 146 | "\n", | 88 | "\n", |
| 147 | "3. (选择题)Hello World 加法算子的数学目标是什么? \n", | 89 | "3. (选择题)Hello World 加法算子的数学目标是什么? \n", |
| 148 | " A. `out = x + y` \n", | 90 | " A. `out = x + y` \n", |
| @@ -175,9 +117,9 @@ | |||
| 175 | "id": "a76768d5", | 117 | "id": "a76768d5", |
| 176 | "metadata": {}, | 118 | "metadata": {}, |
| 177 | "source": [ | 119 | "source": [ |
| 178 | - "## 5. 本节小结\n", | 120 | + "## 4. 本节小结\n", |
| 179 | "\n", | 121 | "\n", |
| 180 | - "本节建立了 PyPTO 算子开发的主线:先定义计算目标,再用 JIT kernel 描述 Tensor 运算,随后由 Host 侧准备数据、调用 kernel 并完成验证。下一节的 Hello World 加法算子会把这条主线落到一段完整代码中;再往后,编程范式、MPMD、API 分类和计算图层次会解释这段代码为什么这样组织。完成 02.05 章节实践后,读者应能独立说清并改写一个基础 PyPTO 算子的开发闭环。\n" | 122 | + "本节建立了 PyPTO 算子开发的主线:先定义计算目标,再用 JIT kernel 描述 Tensor 运算,随后由 Host 侧准备数据、调用 kernel 并完成验证。下一节的 Hello World 加法算子会把这条主线落到一段完整代码中;再往后,编程范式、MPMD、计算图层次会解释这段代码为什么这样组织。完成 02.05 章节实践后,读者应能独立说清并改写一个基础 PyPTO 算子的开发闭环。\n" |
| 181 | ] | 123 | ] |
| 182 | } | 124 | } |
| 183 | ], | 125 | ], |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "1708bb6b", | 5 | "id": "1708bb6b", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 02.02 PyPTO Hello World 加法算子\n", | 8 | + "# PyPTO Hello World 加法算子\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "Hello World 示例用于建立 PyPTO 算子开发的最小闭环。这个闭环包含四个动作:准备输入 Tensor,定义计算规则,执行计算,再用 PyTorch 结果进行验证。\n", | 10 | "Hello World 示例用于建立 PyPTO 算子开发的最小闭环。这个闭环包含四个动作:准备输入 Tensor,定义计算规则,执行计算,再用 PyTorch 结果进行验证。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -15,7 +15,10 @@ | |||
| 15 | "out = x + y\n", | 15 | "out = x + y\n", |
| 16 | "```\n", | 16 | "```\n", |
| 17 | "\n", | 17 | "\n", |
| 18 | - "`x` 和 `y` 是形状相同的输入 Tensor,`out` 是输出 Tensor。逐元素加法表示两个输入在相同位置上的数字相加,例如 `[1, 2, 3] + [10, 20, 30]` 得到 `[11, 22, 33]`。PyPTO kernel 负责描述这个计算,PyTorch 负责构造输入数据和参考结果。\n" | 18 | + "`x` 和 `y` 是形状相同的输入 Tensor,`out` 是输出 Tensor。逐元素加法表示两个输入在相同位置上的数字相加,例如 `[1, 2, 3] + [10, 20, 30]` 得到 `[11, 22, 33]`。PyPTO kernel 负责描述这个计算,PyTorch 负责构造输入数据和参考结果。\n", |
| 19 | + "\n", | ||
| 20 | + "\n", | ||
| 21 | + "\n" | ||
| 19 | ] | 22 | ] |
| 20 | }, | 23 | }, |
| 21 | { | 24 | { |
| @@ -23,11 +26,11 @@ | |||
| 23 | "id": "cloud-env", | 26 | "id": "cloud-env", |
| 24 | "metadata": {}, | 27 | "metadata": {}, |
| 25 | "source": [ | 28 | "source": [ |
| 26 | - "## 1. 安装 PyPTO\n", | 29 | + "## 1. 环境准备\n", |
| 27 | "\n", | 30 | "\n", |
| 28 | - "开始运行 PyPTO 示例前,先执行下面的代码单元安装 PyPTO。Notebook 中的命令需要保留开头的 `!`,这样命令会在当前 Python 环境对应的 shell 中执行。\n", | 31 | + "本教程可以在CANNLab上运行,创建自己的环境后,使用Web IDE方式打开。\n", |
| 29 | "\n", | 32 | "\n", |
| 30 | - "如果在终端中安装,则使用 `pip install pypto`,不需要添加 `!`。安装完成后,再继续执行下面的 Hello World 加法算子代码。\n" | 33 | + "如果提示没有pytpo模块之前的版本,需要手动安装pypto。\n" |
| 31 | ] | 34 | ] |
| 32 | }, | 35 | }, |
| 33 | { | 36 | { |
| @@ -42,20 +45,28 @@ | |||
| 42 | }, | 45 | }, |
| 43 | { | 46 | { |
| 44 | "cell_type": "markdown", | 47 | "cell_type": "markdown", |
| 45 | - "id": "aca458a4", | 48 | + "id": "4ac7cdca", |
| 46 | "metadata": {}, | 49 | "metadata": {}, |
| 47 | "source": [ | 50 | "source": [ |
| 48 | - "## 2. 输入输出规格\n", | 51 | + "配置环境变量" |
| 49 | - "\n", | 52 | + ] |
| 50 | - "在实现算子之前,需要先明确输入、输出、shape 和 dtype。shape 表示 Tensor 每个维度的长度,dtype 表示每个元素的数据类型。\n", | 53 | + }, |
| 51 | - "\n", | 54 | + { |
| 52 | - "| 名称 | 含义 | shape | dtype |\n", | 55 | + "cell_type": "code", |
| 53 | - "| --- | --- | --- | --- |\n", | 56 | + "execution_count": null, |
| 54 | - "| `input_data0` | 第一个输入 Tensor | `(1, 4, 1, 64)` | FP32 |\n", | 57 | + "id": "1f5ac374", |
| 55 | - "| `input_data1` | 第二个输入 Tensor | `(1, 4, 1, 64)` | FP32 |\n", | 58 | + "metadata": {}, |
| 56 | - "| `output_data` | 加法结果输出 Tensor | `(1, 4, 1, 64)` | FP32 |\n", | 59 | + "outputs": [], |
| 57 | - "\n", | 60 | + "source": [ |
| 58 | - "三个 Tensor 的 shape 保持一致,因为逐元素加法要求每个位置都能一一对应。`(1, 4, 1, 64)` 一共包含 `1 * 4 * 1 * 64 = 256` 个元素,输出 Tensor 的 256 个位置分别保存两个输入对应位置的和。\n" | 61 | + "!source /home/developer/Ascend/ascend-toolkit/set_env.sh" |
🟡 Medium Priority changed line: affected behavior: 在 Jupyter/IPython 中, failure mode: 如果用户依赖此步骤来配置 CANN 环境变量(如 NPU 驱动库路径),后续 注意:虽然同一 notebook 中后续的代码单元手动设置了 ![]() ![]() 不准确? | |||
| 62 | + ] | ||
| 63 | + }, | ||
| 64 | + { | ||
| 65 | + "cell_type": "markdown", | ||
| 66 | + "id": "a232e9e1", | ||
| 67 | + "metadata": {}, | ||
| 68 | + "source": [ | ||
| 69 | + "如果在执行多个章节代码时报错`maybe others are using current device`,可以使用命令`npu-smi info`查看当前npu上的进程,使用kill命令移除掉之前的进程,或者关闭多余的ipynb文件,这也会让进程退出。" | ||
| 59 | ] | 70 | ] |
| 60 | }, | 71 | }, |
| 61 | { | 72 | { |
| @@ -63,7 +74,7 @@ | |||
| 63 | "id": "7127caae", | 74 | "id": "7127caae", |
| 64 | "metadata": {}, | 75 | "metadata": {}, |
| 65 | "source": [ | 76 | "source": [ |
| 66 | - "## 3. 完整代码\n", | 77 | + "## 2. 完整代码\n", |
| 67 | "\n", | 78 | "\n", |
| 68 | "下面的代码单元给出完整的 Hello World 加法算子流程。代码中包含依赖导入、设备选择、PyPTO kernel 定义、输入输出构造、运行调用和误差验证。\n" | 79 | "下面的代码单元给出完整的 Hello World 加法算子流程。代码中包含依赖导入、设备选择、PyPTO kernel 定义、输入输出构造、运行调用和误差验证。\n" |
| 69 | ] | 80 | ] |
| @@ -113,173 +124,10 @@ | |||
| 113 | "id": "e05669b4", | 124 | "id": "e05669b4", |
| 114 | "metadata": {}, | 125 | "metadata": {}, |
| 115 | "source": [ | 126 | "source": [ |
| 116 | - "## 4. 运行结果说明\n", | 127 | + "## 3. 运行结果说明\n", |
| 117 | "\n", | 128 | "\n", |
| 118 | - "成功执行时,输出会包含运行模式、设备、输入输出 shape 和最大误差。典型输出形式如下:\n", | 129 | + "通过比较 PyPTO 输出和 PyTorch 参考结果之间的差值,如果数值在误差阈值内时,`assert_close` 校验通过。\n", |
| 119 | - "\n", | 130 | + "\n" |
| 120 | - "```text\n", | ||
| 121 | - "✓ Test add_kernel completed successfully\n", | ||
| 122 | - "```\n", | ||
| 123 | - "\n", | ||
| 124 | - "`Running on device: npu:0` 表示 Tensor 创建在当前进程可见的第 0 张 NPU 上。\n", | ||
| 125 | - "`Max difference` 是 PyPTO 输出和 PyTorch 参考结果之间的最大差值;数值在误差阈值内时,`assert_close` 校验通过。\n", | ||
| 126 | - "最后的通过信息表示加法算子的运行和验证闭环完成。\n" | ||
| 127 | - ] | ||
| 128 | - }, | ||
| 129 | - { | ||
| 130 | - "cell_type": "markdown", | ||
| 131 | - "id": "84e7dce8", | ||
| 132 | - "metadata": {}, | ||
| 133 | - "source": [ | ||
| 134 | - "## 5. 源码结构\n", | ||
| 135 | - "\n", | ||
| 136 | - "这个示例可以分为 Host 侧代码和 Kernel 侧代码。\n", | ||
| 137 | - "\n", | ||
| 138 | - "| 层次 | 作用 | 对应函数 |\n", | ||
| 139 | - "| --- | --- | --- |\n", | ||
| 140 | - "| Host 侧 | 读取设备、创建 Tensor、调用 kernel、校验结果 | `get_device`、`test_add_direct`、`main` |\n", | ||
| 141 | - "| Kernel 侧 | 描述要执行的 Tensor 计算 | `create_add_kernel` 内部的 `add_kernel` |\n", | ||
| 142 | - "\n", | ||
| 143 | - "Host 侧代码由 Python 和 PyTorch 执行,负责把输入、输出和运行参数准备好。Kernel 侧代码由 `@pypto.frontend.jit` 标记,描述需要交给 PyPTO 编译和执行的计算。\n" | ||
| 144 | - ] | ||
| 145 | - }, | ||
| 146 | - { | ||
| 147 | - "cell_type": "markdown", | ||
| 148 | - "id": "853a2530", | ||
| 149 | - "metadata": {}, | ||
| 150 | - "source": [ | ||
| 151 | - "## 6. 设备编号和设备字符串\n", | ||
| 152 | - "\n", | ||
| 153 | - "`device_id` 是设备编号,`device` 是创建 Tensor 时使用的设备字符串。二者的关系如下:\n", | ||
| 154 | - "\n", | ||
| 155 | - "```python\n", | ||
| 156 | - "device_id = 0\n", | ||
| 157 | - "device = \"npu:0\"\n", | ||
| 158 | - "```\n", | ||
| 159 | - "\n", | ||
| 160 | - "当运行模式是 `npu` 且 `device_id` 有效时,输入、输出 Tensor 创建在 `npu:0` 上;当运行模式是 `sim` 或没有可用 NPU 时,Tensor 创建在 `cpu` 上。\n", | ||
| 161 | - "\n", | ||
| 162 | - "在设置可见设备列表的场景下,进程内部的 NPU 编号可能从 0 重新开始。例如某个物理设备被设置为当前进程唯一可见设备后,代码中仍然使用 `npu:0` 表示这张可见设备。切换物理卡需要在 `torch_npu` 初始化前完成,因此修改可见设备配置后需要重启 kernel。\n" | ||
| 163 | - ] | ||
| 164 | - }, | ||
| 165 | - { | ||
| 166 | - "cell_type": "markdown", | ||
| 167 | - "id": "2a4f6138", | ||
| 168 | - "metadata": {}, | ||
| 169 | - "source": [ | ||
| 170 | - "## 7. `create_add_kernel`:定义计算规则\n", | ||
| 171 | - "\n", | ||
| 172 | - "`create_add_kernel` 根据运行模式选择 PyPTO 的执行后端,并返回一个 JIT kernel。JIT 可以理解为“在运行时记录并编译计算描述”的机制。\n", | ||
| 173 | - "\n", | ||
| 174 | - "```python\n", | ||
| 175 | - "@pypto.frontend.jit(runtime_options={\"run_mode\": mode})\n", | ||
| 176 | - "def add_kernel(...):\n", | ||
| 177 | - " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | ||
| 178 | - " out[:] = x + y\n", | ||
| 179 | - "```\n", | ||
| 180 | - "\n", | ||
| 181 | - "`pypto.Tensor([...], pypto.DT_FP32)` 描述 kernel 的输入输出参数是 FP32 Tensor。这里的 `pypto.Tensor` 不是创建真实数据,而是描述进入 PyPTO kernel 的参数类型。真实数据由 Host 侧的 PyTorch 创建,再传入 PyPTO kernel。\n" | ||
| 182 | - ] | ||
| 183 | - }, | ||
| 184 | - { | ||
| 185 | - "cell_type": "markdown", | ||
| 186 | - "id": "16ab5426", | ||
| 187 | - "metadata": {}, | ||
| 188 | - "source": [ | ||
| 189 | - "## 8. Tile Shape 和输出写回\n", | ||
| 190 | - "\n", | ||
| 191 | - "Kernel 内部的核心语句是:\n", | ||
| 192 | - "\n", | ||
| 193 | - "```python\n", | ||
| 194 | - "pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | ||
| 195 | - "out[:] = x + y\n", | ||
| 196 | - "```\n", | ||
| 197 | - "\n", | ||
| 198 | - "Tile Shape 可以先理解为底层执行时使用的数据分块形状。它不改变数学结果,只影响计算如何组织。本例是向量类逐元素计算,因此使用 `set_vec_tile_shapes`。这里的 Tile Shape 与输入 shape 一致,表示按 `(1, 4, 1, 64)` 这一组维度组织计算。\n", | ||
| 199 | - "\n", | ||
| 200 | - "`out[:] = x + y` 表示将加法结果写回输出 Tensor 的全部位置。`out` 是调用者提前分配好的输出缓冲区,kernel 不重新创建输出对象,而是把结果放入这个缓冲区。\n" | ||
| 201 | - ] | ||
| 202 | - }, | ||
| 203 | - { | ||
| 204 | - "cell_type": "markdown", | ||
| 205 | - "id": "74776a6a", | ||
| 206 | - "metadata": {}, | ||
| 207 | - "source": [ | ||
| 208 | - "## 9. `test_add_direct`:构造数据并执行 kernel\n", | ||
| 209 | - "\n", | ||
| 210 | - "`test_add_direct` 将完整验证流程串起来。首先用 PyTorch 构造两个随机输入:\n", | ||
| 211 | - "\n", | ||
| 212 | - "```python\n", | ||
| 213 | - "input_data0 = torch.rand(shape, dtype=torch.float).to(device)\n", | ||
| 214 | - "input_data1 = torch.rand(shape, dtype=torch.float).to(device)\n", | ||
| 215 | - "```\n", | ||
| 216 | - "\n", | ||
| 217 | - "然后分配输出:\n", | ||
| 218 | - "\n", | ||
| 219 | - "```python\n", | ||
| 220 | - "output_data = torch.empty(shape, dtype=torch.float32, device=device)\n", | ||
| 221 | - "```\n", | ||
| 222 | - "\n", | ||
| 223 | - "随机输入先在 CPU 上生成,再通过 `.to(device)` 搬到目标设备。这样可以避免依赖设备侧随机数算子,使示例重点集中在 PyPTO 加法 kernel 的执行和验证上。输出 Tensor 仍然直接在目标设备上分配,随后调用 `add_kernel(input_data0, input_data1, output_data)`,加法结果被写入 `output_data`。\n" | ||
| 224 | - ] | ||
| 225 | - }, | ||
| 226 | - { | ||
| 227 | - "cell_type": "markdown", | ||
| 228 | - "id": "72a17450", | ||
| 229 | - "metadata": {}, | ||
| 230 | - "source": [ | ||
| 231 | - "## 10. PyTorch 参考结果和误差验证\n", | ||
| 232 | - "\n", | ||
| 233 | - "验证闭环使用 PyTorch 生成参考结果:\n", | ||
| 234 | - "\n", | ||
| 235 | - "```python\n", | ||
| 236 | - "golden = torch.add(input_data0, input_data1)\n", | ||
| 237 | - "```\n", | ||
| 238 | - "\n", | ||
| 239 | - "`golden` 的数学含义与 kernel 中的 `out[:] = x + y` 一致,因此可以作为参考答案。随后将 PyPTO 输出和 PyTorch 输出转到 CPU 并转换为 NumPy 数组:\n", | ||
| 240 | - "\n", | ||
| 241 | - "```python\n", | ||
| 242 | - "output_np = output_data.cpu().numpy()\n", | ||
| 243 | - "golden_np = golden.cpu().numpy()\n", | ||
| 244 | - "```\n", | ||
| 245 | - "\n", | ||
| 246 | - "误差验证使用:\n", | ||
| 247 | - "\n", | ||
| 248 | - "```python\n", | ||
| 249 | - "assert_allclose(output_np, golden_np, rtol=3e-3, atol=3e-3)\n", | ||
| 250 | - "```\n", | ||
| 251 | - "\n", | ||
| 252 | - "`rtol` 是相对误差阈值,`atol` 是绝对误差阈值。FP32 加法在这个示例中通常可以得到非常小的误差;保留阈值是为了形成通用的验证写法。\n" | ||
| 253 | - ] | ||
| 254 | - }, | ||
| 255 | - { | ||
| 256 | - "cell_type": "markdown", | ||
| 257 | - "id": "3411b14c", | ||
| 258 | - "metadata": {}, | ||
| 259 | - "source": [ | ||
| 260 | - "## 11. `main`:组织示例入口\n", | ||
| 261 | - "\n", | ||
| 262 | - "`main` 负责组织示例入口,而不是描述算子计算本身。它完成三件事:\n", | ||
| 263 | - "\n", | ||
| 264 | - "1. 使用 `argparse` 解析 `--run_mode`、`--list` 和示例 ID。\n", | ||
| 265 | - "2. 使用 `examples` 字典登记可运行的示例函数。\n", | ||
| 266 | - "3. 根据运行模式准备设备,并调用 `test_add_direct`。\n", | ||
| 267 | - "\n", | ||
| 268 | - "这种组织方式适合后续扩展多个示例。新增示例时,可以在 `examples` 字典中增加新的 ID、名称、说明和函数入口,而不需要改变已有加法 kernel 的实现。\n" | ||
| 269 | - ] | ||
| 270 | - }, | ||
| 271 | - { | ||
| 272 | - "cell_type": "markdown", | ||
| 273 | - "id": "44386a9b", | ||
| 274 | - "metadata": {}, | ||
| 275 | - "source": [ | ||
| 276 | - "## 12. 补充说明\n", | ||
| 277 | - "\n", | ||
| 278 | - "`torch` 版本中出现 `+cpu` 并不必然表示无法使用 NPU。NPU 能力由 `torch_npu` 扩展提供;当 `torch_npu` 可用时,PyTorch 可以通过扩展创建 NPU Tensor。\n", | ||
| 279 | - "\n", | ||
| 280 | - "SIM 模式用于在模拟路径中验证代码结构和编译路径,NPU 模式用于在真实设备上执行 kernel。二者使用同一套 kernel 描述,区别在于 `pypto.RunMode` 的取值不同。\n", | ||
| 281 | - "\n", | ||
| 282 | - "设备侧执行通常是异步的。如果设备执行发生异常,错误可能在后续 `.cpu()` 或同步操作中暴露。因此,验证阶段的同步和结果拷回也是检查 kernel 是否真正执行成功的重要步骤。\n" | ||
| 283 | ] | 131 | ] |
| 284 | }, | 132 | }, |
| 285 | { | 133 | { |
| @@ -287,11 +135,11 @@ | |||
| 287 | "id": "d07b7555", | 135 | "id": "d07b7555", |
| 288 | "metadata": {}, | 136 | "metadata": {}, |
| 289 | "source": [ | 137 | "source": [ |
| 290 | - "## 13. 课后实践\n", | 138 | + "## 4. 课后实践\n", |
| 291 | "1. (选择题)修改 kernel 数学表达式后,PyTorch 参考结果是否也需要同步修改? \n", | 139 | "1. (选择题)修改 kernel 数学表达式后,PyTorch 参考结果是否也需要同步修改? \n", |
| 292 | " A. 需要 \n", | 140 | " A. 需要 \n", |
| 293 | " B. 不需要\n", | 141 | " B. 不需要\n", |
| 294 | - "2. (选择题)`assert_allclose` 的作用是什么? \n", | 142 | + "2. (选择题)`assert_close` 的作用是什么? \n", |
| 295 | " A. 比较 PyPTO 输出和 PyTorch 参考结果是否在误差阈值内 \n", | 143 | " A. 比较 PyPTO 输出和 PyTorch 参考结果是否在误差阈值内 \n", |
| 296 | " B. 创建 NPU 设备 \n", | 144 | " B. 创建 NPU 设备 \n", |
| 297 | " C. 修改 Tensor shape \n", | 145 | " C. 修改 Tensor shape \n", |
| @@ -315,11 +163,11 @@ | |||
| 315 | "id": "dea68ea7", | 163 | "id": "dea68ea7", |
| 316 | "metadata": {}, | 164 | "metadata": {}, |
| 317 | "source": [ | 165 | "source": [ |
| 318 | - "## 14. 本节小结\n", | 166 | + "## 5. 本节小结\n", |
| 319 | "\n", | 167 | "\n", |
| 320 | "Hello World 加法算子展示了 PyPTO 开发的基本闭环:明确输入输出规格,定义 JIT kernel,设置 Tile Shape,写回输出 Tensor,构造 PyTorch 参考结果,并完成误差验证。\n", | 168 | "Hello World 加法算子展示了 PyPTO 开发的基本闭环:明确输入输出规格,定义 JIT kernel,设置 Tile Shape,写回输出 Tensor,构造 PyTorch 参考结果,并完成误差验证。\n", |
| 321 | "\n", | 169 | "\n", |
| 322 | - "在这个闭环中,`create_add_kernel` 描述计算规则,`test_add_direct` 负责运行和验证,`main` 负责组织示例入口。后续更复杂的算子仍然沿用这条主线,只是计算表达式、shape 关系和 Tile 设置会更加丰富。\n" | 170 | + "在这个闭环中,`add_kernel` 描述计算规则,`main` 负责组织示例入口。后续更复杂的算子仍然沿用这条主线,只是计算表达式、shape 关系和 Tile 设置会更加丰富。\n" |
| 323 | ] | 171 | ] |
| 324 | } | 172 | } |
| 325 | ], | 173 | ], |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "444273af", | 5 | "id": "444273af", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 02.03 PyPTO 编程范式与 MPMD 执行模型\n", | 8 | + "# PyPTO 编程范式与 MPMD 执行模型\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "PyPTO 算子开发的关键不在于把普通 Python 语句逐行搬到设备上,而在于用 Python 函数描述 Tensor 计算,再由 PyPTO 记录、编译并调度执行。本节重点解释这种编程范式:哪些代码在 Host 侧运行,哪些代码进入 Kernel 描述,以及 MPMD 执行模型为什么适合组织不同类型的设备任务。\n", | 10 | "PyPTO 算子开发的关键不在于把普通 Python 语句逐行搬到设备上,而在于用 Python 函数描述 Tensor 计算,再由 PyPTO 记录、编译并调度执行。本节重点解释这种编程范式:哪些代码在 Host 侧运行,哪些代码进入 Kernel 描述,以及 MPMD 执行模型为什么适合组织不同类型的设备任务。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -16,36 +16,12 @@ | |||
| 16 | "3. MPMD 如何把不同类型的计算任务按依赖关系组织起来。\n" | 16 | "3. MPMD 如何把不同类型的计算任务按依赖关系组织起来。\n" |
| 17 | ] | 17 | ] |
| 18 | }, | 18 | }, |
| 19 | - { | ||
| 20 | - "cell_type": "markdown", | ||
| 21 | - "id": "a22a144c", | ||
| 22 | - "metadata": {}, | ||
| 23 | - "source": [ | ||
| 24 | - "## 1. 最小运行时检查\n", | ||
| 25 | - "\n", | ||
| 26 | - "本节主要讲解编程范式和执行模型,只需要确认 `torch` 与 `pypto` 可以导入。这里不重复完整环境初始化,也不设置具体 NPU 设备。输出只用于确认运行时依赖可用。后续内容重点在代码结构和执行模型,即使不执行 NPU kernel,也可以理解本节概念。" | ||
| 27 | - ] | ||
| 28 | - }, | ||
| 29 | - { | ||
| 30 | - "cell_type": "code", | ||
| 31 | - "execution_count": null, | ||
| 32 | - "id": "8adc6caa", | ||
| 33 | - "metadata": {}, | ||
| 34 | - "outputs": [], | ||
| 35 | - "source": [ | ||
| 36 | - "import torch\n", | ||
| 37 | - "import pypto\n", | ||
| 38 | - "\n", | ||
| 39 | - "print(\"torch:\", torch.__version__)\n", | ||
| 40 | - "print(\"pypto:\", pypto.__file__)" | ||
| 41 | - ] | ||
| 42 | - }, | ||
| 43 | { | 19 | { |
| 44 | "cell_type": "markdown", | 20 | "cell_type": "markdown", |
| 45 | "id": "b7e18e73", | 21 | "id": "b7e18e73", |
| 46 | "metadata": {}, | 22 | "metadata": {}, |
| 47 | "source": [ | 23 | "source": [ |
| 48 | - "## 2. 从普通 Python 到 PyPTO JIT\n", | 24 | + "## 1. 从普通 Python 到 PyPTO JIT\n", |
| 49 | "\n", | 25 | "\n", |
| 50 | "普通 Python 函数通常按语句顺序立即执行。例如:\n", | 26 | "普通 Python 函数通常按语句顺序立即执行。例如:\n", |
| 51 | "\n", | 27 | "\n", |
| @@ -70,7 +46,7 @@ | |||
| 70 | "id": "a9ec7481", | 46 | "id": "a9ec7481", |
| 71 | "metadata": {}, | 47 | "metadata": {}, |
| 72 | "source": [ | 48 | "source": [ |
| 73 | - "## 3. Host 侧和 Kernel 侧\n", | 49 | + "## 2. Host 侧和 Kernel 侧\n", |
| 74 | "\n", | 50 | "\n", |
| 75 | "一个 PyPTO 示例通常包含 Host 侧代码和 Kernel 侧代码。Host 侧代码运行在普通 Python 环境中,负责准备数据、选择设备、调用 kernel 和验证结果。Kernel 侧代码由 `@pypto.frontend.jit` 标记,描述需要交给 PyPTO 编译和执行的 Tensor 运算。\n", | 51 | "一个 PyPTO 示例通常包含 Host 侧代码和 Kernel 侧代码。Host 侧代码运行在普通 Python 环境中,负责准备数据、选择设备、调用 kernel 和验证结果。Kernel 侧代码由 `@pypto.frontend.jit` 标记,描述需要交给 PyPTO 编译和执行的 Tensor 运算。\n", |
| 76 | "\n", | 52 | "\n", |
| @@ -87,7 +63,7 @@ | |||
| 87 | "id": "b5bcaa47", | 63 | "id": "b5bcaa47", |
| 88 | "metadata": {}, | 64 | "metadata": {}, |
| 89 | "source": [ | 65 | "source": [ |
| 90 | - "## 4. Kernel 函数签名的含义\n", | 66 | + "## 3. Kernel 函数签名的含义\n", |
| 91 | "\n", | 67 | "\n", |
| 92 | "Kernel 函数的参数通常写成:\n", | 68 | "Kernel 函数的参数通常写成:\n", |
| 93 | "\n", | 69 | "\n", |
| @@ -110,14 +86,13 @@ | |||
| 110 | "id": "5dd616a9", | 86 | "id": "5dd616a9", |
| 111 | "metadata": {}, | 87 | "metadata": {}, |
| 112 | "source": [ | 88 | "source": [ |
| 113 | - "## 5. Hello World 中体现的最小 API 骨架\n", | 89 | + "## 4. Hello World 中体现的最小 API 骨架\n", |
| 114 | "\n", | 90 | "\n", |
| 115 | - "本节只关注 Hello World 为了表达编程范式必须出现的 API。更完整的 API 分类放在 `00.04` 中统一整理。\n", | ||
| 116 | "\n", | 91 | "\n", |
| 117 | "| API 或写法 | 在编程范式中的作用 |\n", | 92 | "| API 或写法 | 在编程范式中的作用 |\n", |
| 118 | "| --- | --- |\n", | 93 | "| --- | --- |\n", |
| 119 | "| `@pypto.frontend.jit(...)` | 标记 Kernel 侧计算描述入口 |\n", | 94 | "| `@pypto.frontend.jit(...)` | 标记 Kernel 侧计算描述入口 |\n", |
| 120 | - "| `pypto.RunMode.NPU` / `pypto.RunMode.SIM` | 指定 kernel 的运行后端 |\n", | 95 | + "| `pypto.RunMode.NPU` | 指定 kernel 的运行后端 |\n", |
| 121 | "| `pypto.Tensor([...], pypto.DT_FP32)` | 描述 kernel 输入输出参数 |\n", | 96 | "| `pypto.Tensor([...], pypto.DT_FP32)` | 描述 kernel 输入输出参数 |\n", |
| 122 | "| `pypto.set_vec_tile_shapes(...)` | 描述设备侧分块执行组织 |\n", | 97 | "| `pypto.set_vec_tile_shapes(...)` | 描述设备侧分块执行组织 |\n", |
| 123 | "| `out[:] = x + y` | 描述 Tensor 计算并写回输出 |\n", | 98 | "| `out[:] = x + y` | 描述 Tensor 计算并写回输出 |\n", |
| @@ -130,7 +105,7 @@ | |||
| 130 | "id": "1bd5c818", | 105 | "id": "1bd5c818", |
| 131 | "metadata": {}, | 106 | "metadata": {}, |
| 132 | "source": [ | 107 | "source": [ |
| 133 | - "## 6. Tensor 表达在编程范式中的位置\n", | 108 | + "## 5. Tensor 表达在编程范式中的位置\n", |
| 134 | "\n", | 109 | "\n", |
| 135 | "加法 kernel 的核心表达式是:\n", | 110 | "加法 kernel 的核心表达式是:\n", |
| 136 | "\n", | 111 | "\n", |
| @@ -148,44 +123,18 @@ | |||
| 148 | "id": "60f08d42", | 123 | "id": "60f08d42", |
| 149 | "metadata": {}, | 124 | "metadata": {}, |
| 150 | "source": [ | 125 | "source": [ |
| 151 | - "## 7. Tile Shape 与执行组织\n", | 126 | + "## 6. Tile Shape 与执行组织\n", |
| 152 | "\n", | 127 | "\n", |
| 153 | "Kernel 中常见的 Tile 设置语句是:\n", | 128 | "Kernel 中常见的 Tile 设置语句是:\n", |
| 154 | - "\n", | ||
| 155 | "```python\n", | 129 | "```python\n", |
| 156 | - "pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | 130 | + "pypto.set_vec_tile_shapes()\n", |
| 157 | "```\n", | 131 | "```\n", |
| 158 | "\n", | 132 | "\n", |
| 159 | "Tile Shape 可以理解为设备执行时组织数据的小块形状。它不改变 `out = x + y` 的数学含义,只影响底层如何分块执行。\n", | 133 | "Tile Shape 可以理解为设备执行时组织数据的小块形状。它不改变 `out = x + y` 的数学含义,只影响底层如何分块执行。\n", |
| 160 | "\n", | 134 | "\n", |
| 161 | - "对于 shape 为 `(1, 4, 1, 64)` 的逐元素加法,输入和输出共有 256 个元素。`set_vec_tile_shapes(1, 4, 1, 64)` 表示这个向量类计算可以按这组维度组织。后续遇到矩阵乘法时,会出现面向 Cube 类计算的 Tile 设置;逐元素计算和矩阵计算的执行组织方式不同,因此 Tile API 也不同。\n" | 135 | + "<img src=\"./images/tiling_partition_schematic.png\" alt=\"tiling_partition_schematic\" />\n", |
| 162 | - ] | ||
| 163 | - }, | ||
| 164 | - { | ||
| 165 | - "cell_type": "markdown", | ||
| 166 | - "id": "d9772ec8", | ||
| 167 | - "metadata": {}, | ||
| 168 | - "source": [ | ||
| 169 | - "## 8. MPMD 的直观含义\n", | ||
| 170 | "\n", | 136 | "\n", |
| 171 | - "MPMD 是 Multiple Program Multiple Data 的缩写,可以理解为“多个程序片段处理多份数据”。它强调不同任务可以使用不同的程序片段,而不是所有计算单元都执行完全相同的代码。\n", | 137 | + "后续遇到矩阵乘法时,会出现面向 Cube 类计算的 Tile 设置;逐元素计算和矩阵计算的执行组织方式不同,因此 Tile API 也不同。\n" |
| 172 | - "\n", | ||
| 173 | - "| 模型 | 直观含义 |\n", | ||
| 174 | - "| --- | --- |\n", | ||
| 175 | - "| SPMD | 同一段程序在多个处理单元上处理不同数据 |\n", | ||
| 176 | - "| MPMD | 不同程序片段按依赖关系处理不同数据 |\n", | ||
| 177 | - "\n", | ||
| 178 | - "在 PyPTO 项目中,不同算子会形成不同类型的任务。逐元素加法主要是向量类计算;矩阵乘法更接近 Cube 类计算;Softmax 会组合最大值、减法、指数、求和和除法;复杂网络片段还可能同时包含矩阵计算、逐元素计算、规约和数据搬运。MPMD 模型的价值在于让这些不同任务按照依赖关系组合起来。\n", | ||
| 179 | - "\n", | ||
| 180 | - "| 算子或计算片段 | 可能包含的任务类型 |\n", | ||
| 181 | - "| --- | --- |\n", | ||
| 182 | - "| Hello World 加法 | 逐元素加法 |\n", | ||
| 183 | - "| MatMul + Bias | 矩阵乘法、逐元素加法 |\n", | ||
| 184 | - "| Row Sum | 规约求和 |\n", | ||
| 185 | - "| Softmax | 最大值规约、减法、指数、求和、除法 |\n", | ||
| 186 | - "| Transformer Block 片段 | MatMul、归一化、激活、残差加法、数据重排 |\n", | ||
| 187 | - "\n", | ||
| 188 | - "这些任务的执行方式并不相同。MPMD 模型使框架能够根据任务类型和依赖关系组织执行,而不是把所有设备核都强行套入同一段程序。\n" | ||
| 189 | ] | 138 | ] |
| 190 | }, | 139 | }, |
| 191 | { | 140 | { |
| @@ -193,20 +142,18 @@ | |||
| 193 | "id": "12021c83", | 142 | "id": "12021c83", |
| 194 | "metadata": {}, | 143 | "metadata": {}, |
| 195 | "source": [ | 144 | "source": [ |
| 196 | - "## 9. 从 Kernel 描述到 MPMD 执行\n", | 145 | + "## 7. MPMD 执行模型\n", |
| 197 | "\n", | 146 | "\n", |
| 198 | - "PyPTO 的执行可以概括为下面的流程:\n", | 147 | + "PyPTO基于MPMD(Multiple Program Multiple Data)执行模型,与传统的SPMD(Single Program Multiple Data)模型相比: \n", |
| 199 | "\n", | 148 | "\n", |
| 200 | - "```text\n", | 149 | + " SPMD:用户需编写单一内核逻辑并实例化到多个处理器核上运行,带来同步开销和性能瓶颈 \n", |
| 201 | - "Kernel 中的 Tensor 表达\n", | 150 | + " MPMD:计算被抽象为一组异构任务,任务之间通过依赖关系组织。运行时调度器根据依赖关系将任务分配到合适的执行单元,避免了全局同步限制,提升了整体利用率与效率 \n", |
| 202 | - " -> 框架记录计算关系\n", | ||
| 203 | - " -> 生成设备侧任务\n", | ||
| 204 | - " -> 按依赖关系调度执行\n", | ||
| 205 | - "```\n", | ||
| 206 | "\n", | 151 | "\n", |
| 207 | - "在 Hello World 加法算子中,任务很小,主要就是一次逐元素加法。虽然示例简单,但它已经包含完整链路:Host 侧创建输入输出,Kernel 侧描述计算,PyPTO 记录计算关系,再把任务交给设备执行。\n", | 152 | + "MPMD执行模型的优势包括: \n", |
| 208 | - "\n", | 153 | + "1)、灵活的调度:不同任务可以分配到不同的处理器核,避免全局同步 \n", |
| 209 | - "复杂算子的差别在于任务类型更多。例如一个矩阵乘法加激活函数的算子,可能同时包含矩阵计算、逐元素加法和逐元素激活。MPMD 的作用就是在这些不同任务之间建立有序执行关系。\n" | 154 | + "2)、更好的资源利用:根据任务特性选择合适的执行单元 \n", |
| 155 | + "3)、细粒度并行:计算负载既可在细粒度上并行切分,又能在任务级别灵活调度 \n", | ||
| 156 | + "4)、适配多核架构:更好地适配NPU的多核架构 \n" | ||
| 210 | ] | 157 | ] |
| 211 | }, | 158 | }, |
| 212 | { | 159 | { |
| @@ -214,7 +161,7 @@ | |||
| 214 | "id": "837f402b", | 161 | "id": "837f402b", |
| 215 | "metadata": {}, | 162 | "metadata": {}, |
| 216 | "source": [ | 163 | "source": [ |
| 217 | - "## 10. 哪些语句属于 Kernel 描述\n", | 164 | + "## 8. 哪些语句属于 Kernel 描述\n", |
| 218 | "\n", | 165 | "\n", |
| 219 | "判断一段代码是否属于 PyPTO kernel 描述,可以先看它是否位于 `@pypto.frontend.jit` 修饰的函数内部,以及它是否描述 Tensor 计算或执行组织。\n", | 166 | "判断一段代码是否属于 PyPTO kernel 描述,可以先看它是否位于 `@pypto.frontend.jit` 修饰的函数内部,以及它是否描述 Tensor 计算或执行组织。\n", |
| 220 | "\n", | 167 | "\n", |
| @@ -234,7 +181,7 @@ | |||
| 234 | "id": "a84212ee", | 181 | "id": "a84212ee", |
| 235 | "metadata": {}, | 182 | "metadata": {}, |
| 236 | "source": [ | 183 | "source": [ |
| 237 | - "## 11. 编程范式小结\n", | 184 | + "## 9. 本节小结\n", |
| 238 | "\n", | 185 | "\n", |
| 239 | "PyPTO 的基础编程范式可以概括为三句话:\n", | 186 | "PyPTO 的基础编程范式可以概括为三句话:\n", |
| 240 | "\n", | 187 | "\n", |
| @@ -249,7 +196,7 @@ | |||
| 249 | "cell_type": "markdown", | 196 | "cell_type": "markdown", |
| 250 | "metadata": {}, | 197 | "metadata": {}, |
| 251 | "source": [ | 198 | "source": [ |
| 252 | - "## 12. 课后练习\n", | 199 | + "## 10. 课后练习\n", |
| 253 | "\n", | 200 | "\n", |
| 254 | "本节练习用于检查 Host/Kernel 分工、JIT 参数描述和 MPMD 执行模型。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | 201 | "本节练习用于检查 Host/Kernel 分工、JIT 参数描述和 MPMD 执行模型。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", |
| 255 | "\n", | 202 | "\n", |
| @@ -1,264 +0,0 @@ | |||
| 1 | -{ | ||
| 2 | - "cells": [ | ||
| 3 | - { | ||
| 4 | - "cell_type": "markdown", | ||
| 5 | - "id": "9c4f03d5", | ||
| 6 | - "metadata": {}, | ||
| 7 | - "source": [ | ||
| 8 | - "# 02.04 PyPTO API 与计算图索引\n", | ||
| 9 | - "\n", | ||
| 10 | - "本节作为本章的收束章节,集中整理 PyPTO API、Operation 和计算图层次。上一节已经说明 Host/Kernel 分工和 MPMD 的直观含义;本节进一步回答:常见 API 分别描述什么,Tensor 表达如何对应到 Operation,框架为什么需要多层计算图。\n", | ||
| 11 | - "\n", | ||
| 12 | - "本节的定位是“索引”和“归纳”,不是重新讲一遍 Hello World。后续进入逐元素算子、矩阵乘法、规约和 Softmax 时,可以用这里的分类方式快速判断一行代码处在 PyPTO 开发链路的哪个位置。\n" | ||
| 13 | - ] | ||
| 14 | - }, | ||
| 15 | - { | ||
| 16 | - "cell_type": "markdown", | ||
| 17 | - "id": "17386f5f", | ||
| 18 | - "metadata": {}, | ||
| 19 | - "source": [ | ||
| 20 | - "## 1. 最小依赖检查\n", | ||
| 21 | - "\n", | ||
| 22 | - "本节只需要导入 `torch` 和 `pypto`。代码单元用于确认运行时依赖,不依赖任何源码目录或文档目录。\n" | ||
| 23 | - ] | ||
| 24 | - }, | ||
| 25 | - { | ||
| 26 | - "cell_type": "code", | ||
| 27 | - "execution_count": null, | ||
| 28 | - "id": "9a6a3d00", | ||
| 29 | - "metadata": {}, | ||
| 30 | - "outputs": [], | ||
| 31 | - "source": [ | ||
| 32 | - "import torch\n", | ||
| 33 | - "import pypto\n", | ||
| 34 | - "\n", | ||
| 35 | - "print(\"torch:\", torch.__version__)\n", | ||
| 36 | - "print(\"pypto:\", pypto.__file__)\n" | ||
| 37 | - ] | ||
| 38 | - }, | ||
| 39 | - { | ||
| 40 | - "cell_type": "markdown", | ||
| 41 | - "id": "c99e4204", | ||
| 42 | - "metadata": {}, | ||
| 43 | - "source": [ | ||
| 44 | - "输出中的 `torch` 版本表示 Host 侧张量库可用,`pypto` 路径表示 PyPTO 包可以被当前 Python 环境导入。本节后续内容以概念和代码结构为主,不需要执行 NPU kernel。\n" | ||
| 45 | - ] | ||
| 46 | - }, | ||
| 47 | - { | ||
| 48 | - "cell_type": "markdown", | ||
| 49 | - "id": "f36b0bef", | ||
| 50 | - "metadata": {}, | ||
| 51 | - "source": [ | ||
| 52 | - "## 2. PyPTO API 的系统分类\n", | ||
| 53 | - "\n", | ||
| 54 | - "PyPTO API 可以先按“在开发链路中解决什么问题”分类。这个分类比单纯记 API 名称更重要,因为同一类 API 往往会在不同算子中反复出现。\n", | ||
| 55 | - "\n", | ||
| 56 | - "| 类别 | 典型 API 或写法 | 解决的问题 | 常见位置 |\n", | ||
| 57 | - "| --- | --- | --- | --- |\n", | ||
| 58 | - "| JIT 与运行配置 | `pypto.frontend.jit`、`pypto.RunMode.NPU`、`pypto.RunMode.SIM` | 标记 kernel,指定运行模式 | Kernel 定义处 |\n", | ||
| 59 | - "| Tensor 参数描述 | `pypto.Tensor([...], pypto.DT_FP32)` | 描述 kernel 输入输出的 shape 特征和 dtype | Kernel 函数签名 |\n", | ||
| 60 | - "| 数据类型 | `pypto.DT_FP32`、`pypto.DT_FP16`、`pypto.DT_BF16` | 表示 Tensor 元素类型 | 参数描述、算子输出 dtype |\n", | ||
| 61 | - "| Tile 配置 | `pypto.set_vec_tile_shapes`、`pypto.set_cube_tile_shapes` | 描述向量类或矩阵类计算的分块执行方式 | Kernel 函数内部 |\n", | ||
| 62 | - "| Tensor 运算 | `x + y`、`pypto.add`、`pypto.matmul`、`pypto.sum`、`pypto.exp` | 形成计算图中的 Operation 节点 | Kernel 函数内部 |\n", | ||
| 63 | - "| 输出写回 | `out[:] = ...`、`out.move(...)` | 把计算结果写入调用者提供的输出 Tensor | Kernel 函数内部 |\n", | ||
| 64 | - "| 控制流 | `pypto.loop`、`pypto.cond` | 描述需要进入计算图的循环或条件逻辑 | 复杂 Kernel 内部 |\n", | ||
| 65 | - "| 配置与调试 | `set_*_options`、`get_*_options` | 控制编译、验证、调试等行为 | 工程化调优阶段 |\n", | ||
| 66 | - "\n", | ||
| 67 | - "这个表格中的 API 共同服务于一个目标:把 Tensor 级计算清楚地描述给 PyPTO,让框架能够记录依赖、生成图结构并组织执行。\n" | ||
| 68 | - ] | ||
| 69 | - }, | ||
| 70 | - { | ||
| 71 | - "cell_type": "markdown", | ||
| 72 | - "id": "d735b3ee", | ||
| 73 | - "metadata": {}, | ||
| 74 | - "source": [ | ||
| 75 | - "## 3. 从 Tensor 表达到 Operation\n", | ||
| 76 | - "\n", | ||
| 77 | - "在 PyPTO kernel 内部,Tensor 运算会形成计算图中的 Operation。以加法为例:\n", | ||
| 78 | - "\n", | ||
| 79 | - "```python\n", | ||
| 80 | - "out[:] = x + y\n", | ||
| 81 | - "```\n", | ||
| 82 | - "\n", | ||
| 83 | - "这行代码可以拆成两个层次理解:\n", | ||
| 84 | - "\n", | ||
| 85 | - "| 片段 | 图中的含义 |\n", | ||
| 86 | - "| --- | --- |\n", | ||
| 87 | - "| `x + y` | 形成逐元素加法 Operation |\n", | ||
| 88 | - "| `out[:] = ...` | 将 Operation 结果连接到输出 Tensor |\n", | ||
| 89 | - "\n", | ||
| 90 | - "如果使用显式 API,含义接近:\n", | ||
| 91 | - "\n", | ||
| 92 | - "```python\n", | ||
| 93 | - "tmp = pypto.add(x, y)\n", | ||
| 94 | - "out.move(tmp)\n", | ||
| 95 | - "```\n", | ||
| 96 | - "\n", | ||
| 97 | - "语法糖和显式 API 的共同点是:Tensor 运算描述“算什么”,输出写回描述“结果放到哪里”。在矩阵乘法、规约、Softmax 等算子中,图里的 Operation 会更多,但判断方式相同。\n" | ||
| 98 | - ] | ||
| 99 | - }, | ||
| 100 | - { | ||
| 101 | - "cell_type": "markdown", | ||
| 102 | - "id": "47e18117", | ||
| 103 | - "metadata": {}, | ||
| 104 | - "source": [ | ||
| 105 | - "## 4. 计算图的四个层次\n", | ||
| 106 | - "\n", | ||
| 107 | - "PyPTO 的计算图可以先分成四个层次理解。层次越靠前,越接近 Python 中写出的 Tensor 表达;层次越靠后,越接近设备侧执行计划。\n", | ||
| 108 | - "\n", | ||
| 109 | - "| 图层次 | 主要含义 | 直观理解 | 典型来源 |\n", | ||
| 110 | - "| --- | --- | --- | --- |\n", | ||
| 111 | - "| Tensor Graph | Tensor 级计算关系 | 数学表达式层 | `x + y`、`matmul`、`sum` |\n", | ||
| 112 | - "| Tile Graph | 按 Tile 展开的计算关系 | 分块执行层 | `set_vec_tile_shapes`、`set_cube_tile_shapes` |\n", | ||
| 113 | - "| Block Graph | 可划分、可调度的子图结构 | 任务组织层 | 子图划分、依赖分析 |\n", | ||
| 114 | - "| Execute Graph | 带依赖和调度信息的执行图 | 设备执行层 | 任务提交与调度 |\n", | ||
| 115 | - "\n", | ||
| 116 | - "Hello World 的 Tensor Graph 很小,只包含输入、加法和输出。复杂算子会让图中出现更多 Operation,例如矩阵乘法、规约、指数、除法、广播和数据搬运。图变大以后,依赖关系和执行组织就会变得更重要。\n" | ||
| 117 | - ] | ||
| 118 | - }, | ||
| 119 | - { | ||
| 120 | - "cell_type": "markdown", | ||
| 121 | - "id": "2ae4c1a7", | ||
| 122 | - "metadata": {}, | ||
| 123 | - "source": [ | ||
| 124 | - "## 5. API 到计算图的对应关系\n", | ||
| 125 | - "\n", | ||
| 126 | - "可以用下面的表格把 API 和计算图层次对应起来:\n", | ||
| 127 | - "\n", | ||
| 128 | - "| API 或写法 | 主要影响的层次 | 说明 |\n", | ||
| 129 | - "| --- | --- | --- |\n", | ||
| 130 | - "| `pypto.Tensor([...], dtype)` | Tensor Graph 入口 | 声明图的输入输出参数 |\n", | ||
| 131 | - "| `x + y`、`pypto.add(x, y)` | Tensor Graph | 生成逐元素加法 Operation |\n", | ||
| 132 | - "| `pypto.matmul(a, b, out_dtype=...)` | Tensor Graph | 生成矩阵乘法 Operation |\n", | ||
| 133 | - "| `pypto.sum(x, dim=..., keepdim=...)` | Tensor Graph | 生成规约 Operation |\n", | ||
| 134 | - "| `pypto.set_vec_tile_shapes(...)` | Tile Graph | 描述向量类计算的 Tile 组织 |\n", | ||
| 135 | - "| `pypto.set_cube_tile_shapes(...)` | Tile Graph | 描述矩阵类计算的 Tile 组织 |\n", | ||
| 136 | - "| `pypto.loop(...)`、`pypto.cond(...)` | Tensor Graph / Block Graph | 描述动态循环或条件结构 |\n", | ||
| 137 | - "| `out[:] = ...`、`out.move(...)` | Tensor Graph 输出 | 连接计算结果和输出 Tensor |\n", | ||
| 138 | - "\n", | ||
| 139 | - "这个对应关系说明:API 不是孤立调用,而是在构造图。参数描述给出图的边界,Operation 给出图的节点,Tile API 给出执行组织线索,输出写回确定结果落点。\n" | ||
| 140 | - ] | ||
| 141 | - }, | ||
| 142 | - { | ||
| 143 | - "cell_type": "markdown", | ||
| 144 | - "id": "88bbc038", | ||
| 145 | - "metadata": {}, | ||
| 146 | - "source": [ | ||
| 147 | - "## 6. 从计算图到 MPMD 的关系\n", | ||
| 148 | - "\n", | ||
| 149 | - "MPMD 在本节中只作为计算图执行阶段的视角出现。可以把完整链路概括为:\n", | ||
| 150 | - "\n", | ||
| 151 | - "```text\n", | ||
| 152 | - "API 描述\n", | ||
| 153 | - " -> Tensor Graph\n", | ||
| 154 | - " -> Tile Graph\n", | ||
| 155 | - " -> Block Graph\n", | ||
| 156 | - " -> Execute Graph\n", | ||
| 157 | - " -> MPMD 调度执行\n", | ||
| 158 | - "```\n", | ||
| 159 | - "\n", | ||
| 160 | - "对于 Hello World,加法 Operation 很少,MPMD 的复杂性并不明显。对于 Softmax、LayerNorm、Attention 片段,同一张图中会同时出现规约、逐元素、矩阵计算和数据搬运。不同任务需要不同程序片段配合,MPMD 执行模型的作用就会变得更清楚。\n" | ||
| 161 | - ] | ||
| 162 | - }, | ||
| 163 | - { | ||
| 164 | - "cell_type": "markdown", | ||
| 165 | - "id": "de9fd909", | ||
| 166 | - "metadata": {}, | ||
| 167 | - "source": [ | ||
| 168 | - "## 7. 常见混淆点\n", | ||
| 169 | - "\n", | ||
| 170 | - "| 容易混淆的点 | 正确理解 |\n", | ||
| 171 | - "| --- | --- |\n", | ||
| 172 | - "| `pypto.Tensor(...)` 是否创建真实数据 | 它描述 kernel 参数,不负责创建真实 Tensor 数据 |\n", | ||
| 173 | - "| `torch.rand(...)` 是否进入 PyPTO 计算图 | 它属于 Host 侧数据构造,不进入 PyPTO 计算图 |\n", | ||
| 174 | - "| Tile Shape 是否改变数学结果 | 不改变数学语义,只影响设备侧分块组织方式 |\n", | ||
| 175 | - "| `out[:] = ...` 是否只是 Python 赋值 | 在 JIT kernel 中,它描述输出写回 |\n", | ||
| 176 | - "| `npu:0` 是否一定是物理 0 卡 | 它表示当前进程可见设备中的第 0 个逻辑设备 |\n", | ||
| 177 | - "| `SIM` 是否等同于真实 NPU 性能 | SIM 用于理解流程和功能验证,不代表真实硬件性能 |\n" | ||
| 178 | - ] | ||
| 179 | - }, | ||
| 180 | - { | ||
| 181 | - "cell_type": "markdown", | ||
| 182 | - "id": "2935f544", | ||
| 183 | - "metadata": {}, | ||
| 184 | - "source": [ | ||
| 185 | - "## 8. 本章闭环\n", | ||
| 186 | - "\n", | ||
| 187 | - "本章可以收束为下面这张表:\n", | ||
| 188 | - "\n", | ||
| 189 | - "| 知识点 | 对应能力 |\n", | ||
| 190 | - "| --- | --- |\n", | ||
| 191 | - "| Hello World 加法算子 | 看到一个完整 PyPTO kernel 的定义、调用和验证 |\n", | ||
| 192 | - "| Host/Kernel 划分 | 判断一行代码是在准备数据,还是在描述设备侧计算 |\n", | ||
| 193 | - "| PyPTO API 分类 | 知道 JIT、Tensor、dtype、Tile、Operation、控制流分别解决什么问题 |\n", | ||
| 194 | - "| 计算图层次 | 理解 Tensor 表达如何逐步变成执行图 |\n", | ||
| 195 | - "| MPMD 执行模型 | 理解不同类型任务如何按依赖关系组织执行 |\n", | ||
| 196 | - "\n", | ||
| 197 | - "后续进入初级计算算子章节时,每个新算子都可以放回这条主线中理解:先看数学目标,再看输入输出规格,然后看 kernel 中的 Tensor 表达、Tile 设置和输出写回,最后用 PyTorch reference 建立验证闭环。\n" | ||
| 198 | - ] | ||
| 199 | - }, | ||
| 200 | - { | ||
| 201 | - "cell_type": "markdown", | ||
| 202 | - "metadata": {}, | ||
| 203 | - "source": [ | ||
| 204 | - "## 9. 课后练习\n", | ||
| 205 | - "\n", | ||
| 206 | - "本节练习用于检查 PyPTO API 分类、Operation 和计算图层次。题型包含选择题和填空题,完成后可执行下一单元查看参考答案。\n", | ||
| 207 | - "\n", | ||
| 208 | - "1. (选择题)`x + y` 或 `pypto.add(x, y)` 在 PyPTO kernel 中通常会形成什么? \n", | ||
| 209 | - " A. Host 侧随机输入 \n", | ||
| 210 | - " B. Tensor Graph 中的逐元素加法 Operation \n", | ||
| 211 | - " C. Notebook 标题 \n", | ||
| 212 | - " D. pip 安装命令\n", | ||
| 213 | - "\n", | ||
| 214 | - "2. (填空题)PyPTO 计算图层次可以概括为\\_\\_\\_、\\_\\_\\_、\\_\\_\\_和\\_\\_\\_。\n", | ||
| 215 | - "\n", | ||
| 216 | - "3. (选择题)`pypto.Tensor(...)` 是否负责创建真实 Tensor 数据? \n", | ||
| 217 | - " A. 否,它描述 kernel 参数 \n", | ||
| 218 | - " B. 是,它直接创建真实 NPU Tensor \n", | ||
| 219 | - " C. 是,它负责生成 PyTorch reference \n", | ||
| 220 | - " D. 是,它负责安装依赖\n", | ||
| 221 | - "\n", | ||
| 222 | - "4. (填空题)`pypto.set_vec_tile_shapes(...)` 主要影响________层。\n", | ||
| 223 | - "\n", | ||
| 224 | - "5. (选择题)`torch.add(input0, input1)` 是否进入 PyPTO 计算图? \n", | ||
| 225 | - " A. 进入 \n", | ||
| 226 | - " B. 不进入\n", | ||
| 227 | - "\n", | ||
| 228 | - "6. (填空题)`assert_allclose(...)` 属于________,它________ PyPTO 计算图。\n", | ||
| 229 | - "\n", | ||
| 230 | - "**执行以下代码获取答案。**\n" | ||
| 231 | - ] | ||
| 232 | - }, | ||
| 233 | - { | ||
| 234 | - "cell_type": "code", | ||
| 235 | - "execution_count": null, | ||
| 236 | - "metadata": {}, | ||
| 237 | - "outputs": [], | ||
| 238 | - "source": [ | ||
| 239 | - "!cat ./answer/02.04_answer.txt\n" | ||
| 240 | - ] | ||
| 241 | - } | ||
| 242 | - ], | ||
| 243 | - "metadata": { | ||
| 244 | - "kernelspec": { | ||
| 245 | - "display_name": "Python 3 (ipykernel)", | ||
| 246 | - "language": "python", | ||
| 247 | - "name": "python3" | ||
| 248 | - }, | ||
| 249 | - "language_info": { | ||
| 250 | - "codemirror_mode": { | ||
| 251 | - "name": "ipython", | ||
| 252 | - "version": 3 | ||
| 253 | - }, | ||
| 254 | - "file_extension": ".py", | ||
| 255 | - "mimetype": "text/x-python", | ||
| 256 | - "name": "python", | ||
| 257 | - "nbconvert_exporter": "python", | ||
| 258 | - "pygments_lexer": "ipython3", | ||
| 259 | - "version": "3.12.3" | ||
| 260 | - } | ||
| 261 | - }, | ||
| 262 | - "nbformat": 4, | ||
| 263 | - "nbformat_minor": 5 | ||
| 264 | -} | ||
| @@ -0,0 +1,315 @@ | |||||||
| 1 | +{ | ||||||
| 2 | + "cells": [ | ||||||
| 3 | + { | ||||||
| 4 | + "cell_type": "markdown", | ||||||
| 5 | + "id": "title", | ||||||
| 6 | + "metadata": {}, | ||||||
| 7 | + "source": [ | ||||||
| 8 | + "# 计算图和泳道图\n", | ||||||
| 9 | + "\n", | ||||||
| 10 | + "这一节把视角从“怎么写 kernel”推进到“怎么分析和复用 kernel 的执行”。\n" | ||||||
| 11 | + ] | ||||||
| 12 | + }, | ||||||
| 13 | + { | ||||||
| 14 | + "cell_type": "markdown", | ||||||
| 15 | + "id": "env-title", | ||||||
| 16 | + "metadata": {}, | ||||||
| 17 | + "source": [ | ||||||
| 18 | + "\n", | ||||||
| 19 | + "## 1. 环境准备\n" | ||||||
| 20 | + ] | ||||||
| 21 | + }, | ||||||
| 22 | + { | ||||||
| 23 | + "cell_type": "code", | ||||||
| 24 | + "execution_count": null, | ||||||
| 25 | + "id": "env-code", | ||||||
| 26 | + "metadata": {}, | ||||||
| 27 | + "outputs": [], | ||||||
| 28 | + "source": [ | ||||||
| 29 | + "import os\n", | ||||||
| 30 | + "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | ||||||
| 31 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||||||
| 32 | + "import torch\n", | ||||||
| 33 | + "import pypto\n", | ||||||
| 34 | + "import numpy as np\n", | ||||||
| 35 | + "from numpy.testing import assert_allclose\n", | ||||||
| 36 | + "import torch_npu\n", | ||||||
| 37 | + "\n", | ||||||
| 38 | + "\n", | ||||||
| 39 | + "def get_device():\n", | ||||||
| 40 | + " device_id = int(os.environ.get(\"TILE_FWK_DEVICE_ID\", \"0\"))\n", | ||||||
| 41 | + " return f\"npu:{device_id}\"\n", | ||||||
| 42 | + "\n", | ||||||
| 43 | + "\n", | ||||||
| 44 | + "device = get_device()\n", | ||||||
| 45 | + "RUN_MODE = pypto.RunMode.NPU\n", | ||||||
| 46 | + "\n", | ||||||
| 47 | + "\n", | ||||||
| 48 | + "print(\"TILE_FWK_DEVICE_ID:\", os.environ.get(\"TILE_FWK_DEVICE_ID\", \"<not set>\"))\n", | ||||||
| 49 | + "print(\"device:\", device)\n", | ||||||
| 50 | + "print(\"run_mode:\", RUN_MODE)\n", | ||||||
| 51 | + "print(\"pypto:\", pypto.__file__)" | ||||||
| 52 | + ] | ||||||
| 53 | + }, | ||||||
| 54 | + { | ||||||
| 55 | + "cell_type": "markdown", | ||||||
| 56 | + "id": "goal-title", | ||||||
| 57 | + "metadata": {}, | ||||||
| 58 | + "source": [ | ||||||
| 59 | + "\n", | ||||||
| 60 | + "## 2. 代码展示\n", | ||||||
| 61 | + "\n", | ||||||
| 62 | + "使用PyPTO框架实现一个简单的算子,并通过测试用例验证其正确性。本节可以先初步了解计算图和泳道图,随着课程学完,在后续的开发工作中真正实践。\n", | ||||||
| 63 | + "\n", | ||||||
| 64 | + "可以暂时不了解softmax_core的数学含义,通过本节的学习,了解如何使用PyPTO的API来构建自定义算子。并且在程序完成后,通过`PyPTO Toolkit`可视化工具查看计算图结构,并观测算子的各项性能数据。\n" | ||||||
| 65 | + ] | ||||||
| 66 | + }, | ||||||
| 67 | + { | ||||||
| 68 | + "cell_type": "code", | ||||||
| 69 | + "execution_count": null, | ||||||
| 70 | + "id": "a012a4c2", | ||||||
| 71 | + "metadata": {}, | ||||||
| 72 | + "outputs": [], | ||||||
| 73 | + "source": [ | ||||||
| 74 | + "def softmax_core(x: pypto.Tensor) -> pypto.Tensor:\n", | ||||||
| 75 | + " row_max = pypto.amax(x, dim=-1, keepdim=True)\n", | ||||||
| 76 | + " sub = x - row_max\n", | ||||||
| 77 | + " exp = pypto.exp(sub)\n", | ||||||
| 78 | + " esum = pypto.sum(exp, dim=-1, keepdim=True)\n", | ||||||
| 79 | + " return exp / esum\n", | ||||||
| 80 | + "\n", | ||||||
| 81 | + "\n", | ||||||
| 82 | + "@pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", | ||||||
| 83 | + "def softmax_kernel(\n", | ||||||
| 84 | + " input_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),\n", | ||||||
| 85 | + " output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32),\n", | ||||||
| 86 | + "):\n", | ||||||
| 87 | + " bs, seqlen, head, dim = input_tensor.shape\n", | ||||||
| 88 | + " tile_b = 1\n", | ||||||
| 89 | + " b_loop = bs // tile_b\n", | ||||||
| 90 | + "\n", | ||||||
| 91 | + " # Tiling shape setting for efficient execution\n", | ||||||
| 92 | + " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | ||||||
| 93 | + "\n", | ||||||
| 94 | + " for idx in pypto.loop(0, b_loop, 1, name=\"LOOP_L0_bIdx\", idx_name=\"idx\"):\n", | ||||||
| 95 | + " b_offset = idx * tile_b\n", | ||||||
| 96 | + " b_offset_end = (idx + 1) * tile_b\n", | ||||||
| 97 | + " input_view = input_tensor[b_offset:b_offset_end, :seqlen, :head, :dim]\n", | ||||||
| 98 | + " softmax_out = softmax_core(input_view)\n", | ||||||
| 99 | + " output_tensor[b_offset:b_offset_end, ...] = softmax_out\n", | ||||||
| 100 | + "\n", | ||||||
🟠 High Priority changed line: affected behavior: 在 softmax_kernel 的 batch 循环中,每轮迭代计算一个 batch 元素( failure mode:
注意:变量 suggested fix: 将 建议:将输出切片的结束位置从 改动建议
![]() ![]() 不准确? | |||||||
| 101 | + "\n", | ||||||
| 102 | + "def test_softmax(device_id: int = None, dynamic: bool = True) -> None:\n", | ||||||
| 103 | + " device = get_device()\n", | ||||||
| 104 | + " shape = (32, 32, 1, 256)\n", | ||||||
| 105 | + " x = torch.rand(shape, dtype=torch.float, device=device)\n", | ||||||
| 106 | + " y = torch.zeros(shape, dtype=torch.float, device=device)\n", | ||||||
| 107 | + " softmax_kernel(x, y)\n", | ||||||
| 108 | + " golden = torch.softmax(x, dim=-1).cpu()\n", | ||||||
| 109 | + " y = y.cpu()\n", | ||||||
| 110 | + " max_diff = np.abs(y.numpy() - golden.numpy()).max()\n", | ||||||
| 111 | + " print(f\"Input shape: {x.shape}\")\n", | ||||||
| 112 | + " print(f\"Output shape: {y.shape}\")\n", | ||||||
| 113 | + " print(f\"Max difference: {max_diff:.6f}\")\n", | ||||||
| 114 | + " assert_allclose(np.array(y), np.array(golden), rtol=3e-3, atol=3e-3)\n", | ||||||
| 115 | + " print(\"✓ Softmax test passed\")\n", | ||||||
| 116 | + " print()\n", | ||||||
| 117 | + "\n", | ||||||
| 118 | + "\n", | ||||||
| 119 | + "test_softmax()" | ||||||
| 120 | + ] | ||||||
| 121 | + }, | ||||||
| 122 | + { | ||||||
| 123 | + "cell_type": "markdown", | ||||||
| 124 | + "id": "shared-title", | ||||||
| 125 | + "metadata": {}, | ||||||
| 126 | + "source": [ | ||||||
| 127 | + "\n", | ||||||
| 128 | + "## 3. 编译与执行\n" | ||||||
| 129 | + ] | ||||||
| 130 | + }, | ||||||
| 131 | + { | ||||||
| 132 | + "cell_type": "markdown", | ||||||
| 133 | + "id": "30689e9d", | ||||||
| 134 | + "metadata": {}, | ||||||
| 135 | + "source": [ | ||||||
| 136 | + "### 3.1 查看计算图" | ||||||
| 137 | + ] | ||||||
| 138 | + }, | ||||||
| 139 | + { | ||||||
| 140 | + "cell_type": "markdown", | ||||||
| 141 | + "id": "80487f94", | ||||||
| 142 | + "metadata": {}, | ||||||
| 143 | + "source": [ | ||||||
| 144 | + "开启图编译阶段调试模式开关。 " | ||||||
| 145 | + ] | ||||||
| 146 | + }, | ||||||
| 147 | + { | ||||||
| 148 | + "cell_type": "markdown", | ||||||
| 149 | + "id": "05aec023", | ||||||
| 150 | + "metadata": {}, | ||||||
| 151 | + "source": [ | ||||||
| 152 | + "@pypto.frontend.jit(\n", | ||||||
| 153 | + " debug_options={\"runtime_debug_mode\": 1}\n", | ||||||
| 154 | + ")" | ||||||
| 155 | + ] | ||||||
| 156 | + }, | ||||||
| 157 | + { | ||||||
| 158 | + "cell_type": "markdown", | ||||||
| 159 | + "id": "09d14f6f", | ||||||
| 160 | + "metadata": {}, | ||||||
| 161 | + "source": [ | ||||||
| 162 | + "PyPTO程序在编译过程中,会自动生成由Tensor和Operation组合而成的图结构,即计算图。该计算图经过PyPTO编译优化流程,完成从原始计算图到可执行图的编译过程,最终生成可在昇腾硬件环境中运行的可执行代码,以实现实际的计算任务。用户可借助PyPTO Toolkit可视化工具可以直观地查看计算图结构,了解计算图的节点信息,更便捷地进行算子功能调试。\n", | ||||||
| 163 | + "\n", | ||||||
| 164 | + "右键单击${work_path}/output/output_*/program.json文件,在弹出的菜单中选择“使用PyPTO Toolkit打开”。\n", | ||||||
| 165 | + "\n", | ||||||
| 166 | + "program.json文件包含了Execute Graph和Block Graph的汇总信息,图中的关键信息为:左右两边的卡片为Tensor节点(代表输入/输出数据)、中间卡片为调用节点(带有fx标识,单击可以实现信息钻取)。\n", | ||||||
| 167 | + "\n" | ||||||
| 168 | + ] | ||||||
| 169 | + }, | ||||||
| 170 | + { | ||||||
| 171 | + "cell_type": "markdown", | ||||||
| 172 | + "id": "cost-helper-code", | ||||||
| 173 | + "metadata": {}, | ||||||
| 174 | + "source": [ | ||||||
| 175 | + "<img src=\"./images/tensor_graph.png\" alt=\"tensor_graph\" width=\"900px\" >" | ||||||
| 176 | + ] | ||||||
| 177 | + }, | ||||||
| 178 | + { | ||||||
| 179 | + "cell_type": "markdown", | ||||||
| 180 | + "id": "d67e51c1", | ||||||
| 181 | + "metadata": {}, | ||||||
| 182 | + "source": [ | ||||||
| 183 | + "双击中间卡片逐层钻取到下图所示的Execute Graph。" | ||||||
| 184 | + ] | ||||||
| 185 | + }, | ||||||
| 186 | + { | ||||||
| 187 | + "cell_type": "markdown", | ||||||
| 188 | + "id": "cost-kernel-code", | ||||||
| 189 | + "metadata": {}, | ||||||
| 190 | + "source": [ | ||||||
| 191 | + "<img src=\"./images/execute_graph.png\" alt=\"execute_graph\" width=\"900px\" >" | ||||||
| 192 | + ] | ||||||
| 193 | + }, | ||||||
| 194 | + { | ||||||
| 195 | + "cell_type": "markdown", | ||||||
| 196 | + "id": "cost-note", | ||||||
| 197 | + "metadata": {}, | ||||||
| 198 | + "source": [ | ||||||
| 199 | + "双击上图的调用节点,可以看到Block Graph子图信息,标识着任务的具体执行过程。放大后可以看到图中具体的Tensor和Operation节点信息和连接关系,例如图中可以看到SUB和EXP操作。" | ||||||
| 200 | + ] | ||||||
| 201 | + }, | ||||||
| 202 | + { | ||||||
| 203 | + "cell_type": "markdown", | ||||||
| 204 | + "id": "cost-run", | ||||||
| 205 | + "metadata": {}, | ||||||
| 206 | + "source": [ | ||||||
| 207 | + "<img src=\"./images/block_graph.png\" alt=\"block_graph\" width=\"900px\" >" | ||||||
| 208 | + ] | ||||||
| 209 | + }, | ||||||
| 210 | + { | ||||||
| 211 | + "cell_type": "markdown", | ||||||
| 212 | + "id": "acl-title", | ||||||
| 213 | + "metadata": {}, | ||||||
| 214 | + "source": [ | ||||||
| 215 | + "\n", | ||||||
| 216 | + "### 3.2 查看泳道图\n", | ||||||
| 217 | + "\n", | ||||||
| 218 | + "泳道图用于直观展示计算图的实际调度与执行过程,清晰呈现任务的执行顺序和耗时信息,帮助开发者分析算子性能瓶颈。下面将介绍如何采集泳道图数据,并通过PyPTO Toolkit查看泳道图。\n", | ||||||
| 219 | + "\n", | ||||||
| 220 | + "通过给@pypto.frontend.jit装饰器的入参debug_options配置图执行阶段调试开关启动性能数据采集功能。重新执行代码可以生成泳道图\n" | ||||||
| 221 | + ] | ||||||
| 222 | + }, | ||||||
| 223 | + { | ||||||
| 224 | + "cell_type": "markdown", | ||||||
| 225 | + "id": "acl-kernel-code", | ||||||
| 226 | + "metadata": {}, | ||||||
| 227 | + "source": [ | ||||||
| 228 | + "@pypto.frontend.jit(debug_options={\"runtime_debug_mode\": 1})" | ||||||
| 229 | + ] | ||||||
| 230 | + }, | ||||||
| 231 | + { | ||||||
| 232 | + "cell_type": "markdown", | ||||||
| 233 | + "id": "21888c8d", | ||||||
| 234 | + "metadata": {}, | ||||||
| 235 | + "source": [ | ||||||
| 236 | + "在${work_path}/output/output_*/目录(*代表时间戳)下生成泳道图数据文件,文件名为:merged_swimlane.json。右键单击merged_swimlane.json,在弹出的菜单中选择“使用PyPTO Toolkit打开”,如下图所示。" | ||||||
| 237 | + ] | ||||||
| 238 | + }, | ||||||
| 239 | + { | ||||||
| 240 | + "cell_type": "markdown", | ||||||
| 241 | + "id": "3d07899c", | ||||||
| 242 | + "metadata": {}, | ||||||
| 243 | + "source": [ | ||||||
| 244 | + "<img src=\"./images/swimlane.png\" alt=\"swimlane\" width=\"900px\" >" | ||||||
| 245 | + ] | ||||||
| 246 | + }, | ||||||
| 247 | + { | ||||||
| 248 | + "cell_type": "markdown", | ||||||
| 249 | + "id": "b3121364", | ||||||
| 250 | + "metadata": {}, | ||||||
| 251 | + "source": [ | ||||||
| 252 | + "上图中带有色块的部分即为泳道,展示了每个AIC/AIV上的任务执行情况。泳道条目的长度对应任务的耗时,能够直观地反映计算的密集程度。用户可以通过观察相邻泳道之间的空闲间隔以及耗时较长的泳道条目,来分析可能存在的性能瓶颈问题。" | ||||||
| 253 | + ] | ||||||
| 254 | + }, | ||||||
| 255 | + { | ||||||
| 256 | + "cell_type": "markdown", | ||||||
| 257 | + "metadata": {}, | ||||||
| 258 | + "source": [ | ||||||
| 259 | + "## 4. 课后练习\n", | ||||||
| 260 | + "\n", | ||||||
| 261 | + "本节练习用于复盘计算图和泳道图的作用。请完成以下题目。\n", | ||||||
| 262 | + "\n", | ||||||
| 263 | + "1. (选择题)关于 PyPTO 计算图的查看方式,以下说法正确的是? \n", | ||||||
| 264 | + "A. 计算图数据存储在 merged_swimlane.json 文件中,需要用文本编辑器打开 \n", | ||||||
| 265 | + "B. program.json 文件包含了 Execute Graph 和 Block Graph 的汇总信息,通过右键选择\"使用 PyPTO Toolkit 打开\"查看 \n", | ||||||
| 266 | + "C. 计算图只能通过命令行工具 pypto-cli 导出为 SVG 图片查看 \n", | ||||||
| 267 | + "D. 计算图在程序运行时直接在终端以文本形式输出,无需额外文件 \n", | ||||||
| 268 | + "\n", | ||||||
| 269 | + "2. (选择题)要采集泳道图性能数据,需要在 @pypto.frontend.jit 装饰器中做怎样的配置? \n", | ||||||
| 270 | + "A. 在 runtime_options 中设置 {\"run_mode\": pypto.RunMode.SIM} \n", | ||||||
| 271 | + "B. 在 pass_options 中设置 {\"enable_swimlane\": True} \n", | ||||||
| 272 | + "C. 在 debug_options 中设置 {\"runtime_debug_mode\": 1} \n", | ||||||
| 273 | + "D. 在 codegen_options 中设置 {\"output_swimlane\": True} \n", | ||||||
| 274 | + "\n", | ||||||
| 275 | + "3. (选择题)关于泳道图的作用与特征,以下描述正确的是? \n", | ||||||
| 276 | + "A. 泳道图用于展示计算图的编译优化过程,每个泳道代表一个编译 Pass \n", | ||||||
| 277 | + "B. 泳道图中的色块代表 AIC/AIV 上的任务执行情况,泳道条目长度对应任务耗时,可用于分析性能瓶颈 \n", | ||||||
| 278 | + "C. 泳道图和计算图查看的是同一个 program.json 文件,只是展示视角不同 \n", | ||||||
| 279 | + "D. 泳道图只能在 NPU 硬件上生成,SIM 模式下无法生成泳道数据 \n", | ||||||
| 280 | + " \n", | ||||||
| 281 | + "**执行以下代码获取答案。**\n" | ||||||
| 282 | + ] | ||||||
| 283 | + }, | ||||||
| 284 | + { | ||||||
| 285 | + "cell_type": "code", | ||||||
| 286 | + "execution_count": null, | ||||||
| 287 | + "metadata": {}, | ||||||
| 288 | + "outputs": [], | ||||||
| 289 | + "source": [ | ||||||
| 290 | + "!cat ./answer/02.04_answer.txt\n" | ||||||
| 291 | + ] | ||||||
| 292 | + } | ||||||
| 293 | + ], | ||||||
| 294 | + "metadata": { | ||||||
| 295 | + "kernelspec": { | ||||||
| 296 | + "display_name": "Python 3 (ipykernel)", | ||||||
| 297 | + "language": "python", | ||||||
| 298 | + "name": "python3" | ||||||
| 299 | + }, | ||||||
| 300 | + "language_info": { | ||||||
| 301 | + "codemirror_mode": { | ||||||
| 302 | + "name": "ipython3", | ||||||
| 303 | + "version": 3 | ||||||
| 304 | + }, | ||||||
| 305 | + "file_extension": ".py", | ||||||
| 306 | + "mimetype": "text/x-python", | ||||||
| 307 | + "name": "python", | ||||||
| 308 | + "nbconvert_exporter": "python", | ||||||
| 309 | + "pygments_lexer": "ipython3", | ||||||
| 310 | + "version": "3.12.3" | ||||||
| 311 | + } | ||||||
| 312 | + }, | ||||||
| 313 | + "nbformat": 4, | ||||||
| 314 | + "nbformat_minor": 5 | ||||||
| 315 | +} | ||||||
| @@ -4,7 +4,7 @@ | |||
| 4 | "cell_type": "markdown", | 4 | "cell_type": "markdown", |
| 5 | "metadata": {}, | 5 | "metadata": {}, |
| 6 | "source": [ | 6 | "source": [ |
| 7 | - "# 02.05 章节实践\n", | 7 | + "# 章节实践\n", |
| 8 | "\n", | 8 | "\n", |
| 9 | "本节用于检查第二章的掌握情况。练习会围绕一个最小 PyPTO 算子的开发闭环展开:明确计算目标、编写 Kernel、准备 Host 侧输入输出、设置运行模式并用 PyTorch 结果完成验证。\n" | 9 | "本节用于检查第二章的掌握情况。练习会围绕一个最小 PyPTO 算子的开发闭环展开:明确计算目标、编写 Kernel、准备 Host 侧输入输出、设置运行模式并用 PyTorch 结果完成验证。\n" |
| 10 | ] | 10 | ] |
| @@ -55,11 +55,30 @@ | |||
| 55 | "\n", | 55 | "\n", |
| 56 | "6. (填空题)PyPTO 计算图可以帮助理解从 Tensor 表达逐步进入更低层执行组织的过程,其中本章提到的层次包括 Tensor Graph、Tile Graph、Block Graph 和________。\n", | 56 | "6. (填空题)PyPTO 计算图可以帮助理解从 Tensor 表达逐步进入更低层执行组织的过程,其中本章提到的层次包括 Tensor Graph、Tile Graph、Block Graph 和________。\n", |
| 57 | "\n", | 57 | "\n", |
| 58 | - "7. (编程题)在 Hello World 加法算子的基础上完成综合改写:把计算目标从 `out = x + y` 改为 `out = (x + y) * 0.5`。要求: \n", | 58 | + "7. (编程题)实现half_add_kernel方法,实现 `out = (x + y) * 0.5`。要求Kernel 侧完成加法和乘常数,Host 侧代码已给出,只需编写kernel侧代码。\n", |
| 59 | - " - Kernel 侧完成加法和乘常数。 \n", | 59 | + " \n" |
| 60 | - " - Host 侧仍创建两个输入 Tensor 和一个输出 Tensor。 \n", | 60 | + ] |
| 61 | - " - PyTorch 参考结果同步改为 `(input_data0 + input_data1) * 0.5`。 \n", | 61 | + }, |
| 62 | - " - 使用 `assert_allclose` 验证输出。\n" | 62 | + { |
| 63 | + "cell_type": "code", | ||
| 64 | + "execution_count": null, | ||
| 65 | + "metadata": {}, | ||
| 66 | + "outputs": [], | ||
| 67 | + "source": [ | ||
| 68 | + "# TODO 此处添加kernel函数的实现\n", | ||
| 69 | + "\n", | ||
| 70 | + "\n", | ||
| 71 | + "def test_half_add_kernel():\n", | ||
| 72 | + " shape = (64, 64)\n", | ||
| 73 | + " device = get_device()\n", | ||
| 74 | + " x = torch.randn(shape, dtype=torch.float, device=device)\n", | ||
| 75 | + " y = torch.randn(shape, dtype=torch.float, device=device)\n", | ||
| 76 | + " out = torch.empty(shape, dtype=torch.float, device=device)\n", | ||
| 77 | + " half_add_kernel(x, y, out)\n", | ||
| 78 | + " torch.testing.assert_close((x + y) * 0.5, out, atol=1e-3, rtol=1e-3)\n", | ||
| 79 | + " print(\"✓ Test completed successfully\")\n", | ||
| 80 | + "\n", | ||
| 81 | + "test_half_add_kernel()" | ||
| 63 | ] | 82 | ] |
| 64 | }, | 83 | }, |
| 65 | { | 84 | { |
| @@ -3,4 +3,4 @@ | |||
| 3 | 3. B | 3 | 3. B |
| 4 | 4. 参数类型。 | 4 | 4. 参数类型。 |
| 5 | 5. A | 5 | 5. A |
| 6 | -6. 逐元素加法;矩阵乘法、逐元素加法;最大值规约、减法、指数、求和、除法。 | 6 | +6. 最大值规约、减法、指数、求和、除法。 |
| @@ -1,6 +1,3 @@ | |||
| 1 | 1. B | 1 | 1. B |
| 2 | -2. Tensor Graph;Tile Graph;Block Graph;Execute Graph。 | 2 | +2. C |
| 3 | -3. A | 3 | +3. B |
| 4 | -4. Tile Graph。 | ||
| 5 | -5. B | ||
| 6 | -6. Host 侧参考实现;不进入 PyPTO 计算图。 | ||
| @@ -9,14 +9,4 @@ | |||
| 9 | 9 | ||
| 10 | def half_add_kernel(x: pypto.Tensor[...], y: pypto.Tensor[...], out: pypto.Tensor[...]): | 10 | def half_add_kernel(x: pypto.Tensor[...], y: pypto.Tensor[...], out: pypto.Tensor[...]): |
| 11 | pypto.set_vec_tile_shapes(32, 32) | 11 | pypto.set_vec_tile_shapes(32, 32) |
| 12 | - out[:] = pypto.mul(pypto.add(x, y), 0.5) | 12 | + out[:] = pypto.mul(pypto.add(x, y), 0.5) |
| 13 | - | ||
| 14 | -def test_half_add_kernel(): | ||
| 15 | - shape = (64, 64) | ||
| 16 | - device = get_device() | ||
| 17 | - x = torch.randn(shape, dtype=torch.float, device=device) | ||
| 18 | - y = torch.randn(shape, dtype=torch.float, device=device) | ||
| 19 | - out = torch.empty(shape, dtype=torch.float, device=device) | ||
| 20 | - half_add_kernel(x, y, out) | ||
| 21 | - torch.testing.assert_close((x + y) * 0.5, out, atol=1e-3, rtol=1e-3) | ||
| 22 | - print("✓ Test completed successfully") | ||
| @@ -5,9 +5,9 @@ | |||
| 5 | "id": "a0907a15", | 5 | "id": "a0907a15", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.01 章节介绍\n", | 8 | + "# 章节介绍\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | - "本节是 PyPTO 初级教程中“计算算子样例”章节的导学内容。后续章节会深入学习逐元素算子及其 API 全量覆盖、矩阵乘法、规约算子及其 API 全量覆盖、Tiling 及其 API 全量覆盖、形状变换和 Softmax。本节先介绍 Tensor 描述对象,然后用一个最小 kernel 建立共同的阅读方法。\n", | 10 | + "本节是 PyPTO 初级教程中“计算算子样例”章节的导学内容。后续章节会深入学习逐元素算子、矩阵乘法、规约算子、Tiling、形状变换和 Softmax。本节先介绍 Tensor 描述对象,然后用一个最小 kernel 建立共同的阅读方法。\n", |
| 11 | "\n", | 11 | "\n", |
| 12 | "一个 PyPTO 示例通常同时包含两种代码:kernel 外部的 PyTorch 代码,以及 kernel 内部的 PyPTO 计算描述。前者负责创建真实数据、选择设备、分配输出和验证结果;后者负责描述要编译执行的 Tensor 计算。初学时只要把这两层分清,很多代码就不再混乱。\n", | 12 | "一个 PyPTO 示例通常同时包含两种代码:kernel 外部的 PyTorch 代码,以及 kernel 内部的 PyPTO 计算描述。前者负责创建真实数据、选择设备、分配输出和验证结果;后者负责描述要编译执行的 Tensor 计算。初学时只要把这两层分清,很多代码就不再混乱。\n", |
| 13 | "\n", | 13 | "\n", |
| @@ -41,6 +41,7 @@ | |||
| 41 | "source": [ | 41 | "source": [ |
| 42 | "import os\n", | 42 | "import os\n", |
| 43 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 43 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 44 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 44 | "import torch\n", | 45 | "import torch\n", |
| 45 | "import pypto\n", | 46 | "import pypto\n", |
| 46 | "import torch_npu\n", | 47 | "import torch_npu\n", |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "8f8c4167", | 5 | "id": "8f8c4167", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.02 逐元素算子与 set_vec_tile_shapes\n", | 8 | + "# 逐元素算子与 set_vec_tile_shapes\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "上一节用一个最小 kernel 建立了 PyPTO 算子的基本闭环:host 侧准备真实 Tensor,kernel 侧描述计算,结果写入 `out`,最后用 PyTorch 参考实现验证。本节进入第一类具体计算:逐元素算子。\n", | 10 | "上一节用一个最小 kernel 建立了 PyPTO 算子的基本闭环:host 侧准备真实 Tensor,kernel 侧描述计算,结果写入 `out`,最后用 PyTorch 参考实现验证。本节进入第一类具体计算:逐元素算子。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -32,6 +32,7 @@ | |||
| 32 | "source": [ | 32 | "source": [ |
| 33 | "import os\n", | 33 | "import os\n", |
| 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 35 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 35 | "import torch\n", | 36 | "import torch\n", |
| 36 | "import pypto\n", | 37 | "import pypto\n", |
| 37 | "import torch_npu\n", | 38 | "import torch_npu\n", |
| @@ -55,27 +56,35 @@ | |||
| 55 | "source": [ | 56 | "source": [ |
| 56 | "## 1. 什么是逐元素计算\n", | 57 | "## 1. 什么是逐元素计算\n", |
| 57 | "\n", | 58 | "\n", |
| 58 | - "逐元素计算可以直观理解为“每个位置各算各的”。例如输入 `x` 和 `scale` 的形状都是 `[8, 8]`,表达式:\n", | 59 | + "逐元素计算可以直观理解为“每个位置各算各的”。例如输入 `A` 和 `B` 的形状都是 `[8, 8]`,表达式:\n", |
| 59 | "\n", | 60 | "\n", |
| 60 | "```python\n", | 61 | "```python\n", |
| 61 | - "y = x * scale\n", | 62 | + "C = A * B\n", |
| 62 | "```\n", | 63 | "```\n", |
| 63 | "\n", | 64 | "\n", |
| 64 | - "表示输出中 `(i, j)` 位置的值只依赖 `x[i, j]` 和 `scale[i, j]`:\n", | 65 | + "表示输出中 `(i, j)` 位置的值只依赖 `A[i, j]` 和 `B[i, j]`:\n", |
| 65 | "\n", | 66 | "\n", |
| 66 | "```text\n", | 67 | "```text\n", |
| 67 | - "y[i, j] = x[i, j] * scale[i, j]\n", | 68 | + "C[i, j] = A[i, j] * B[i, j]\n", |
| 68 | "```\n", | 69 | "```\n", |
| 69 | "\n", | 70 | "\n", |
| 70 | "不同位置之间没有相互依赖,也不会把一组元素合成一个元素。这个特点和后面要学习的规约不同:逐元素计算通常保持 shape,规约通常会压缩某个维度。" | 71 | "不同位置之间没有相互依赖,也不会把一组元素合成一个元素。这个特点和后面要学习的规约不同:逐元素计算通常保持 shape,规约通常会压缩某个维度。" |
| 71 | ] | 72 | ] |
| 72 | }, | 73 | }, |
| 74 | + { | ||
| 75 | + "cell_type": "markdown", | ||
| 76 | + "id": "86d93dab", | ||
| 77 | + "metadata": {}, | ||
| 78 | + "source": [ | ||
| 79 | + "<img src=\"./images/elemul.png\" alt=\"乘法\" width=\"700px\" >" | ||
| 80 | + ] | ||
| 81 | + }, | ||
| 73 | { | 82 | { |
| 74 | "cell_type": "markdown", | 83 | "cell_type": "markdown", |
| 75 | "id": "a1fed2e0", | 84 | "id": "a1fed2e0", |
| 76 | "metadata": {}, | 85 | "metadata": {}, |
| 77 | "source": [ | 86 | "source": [ |
| 78 | - "## 1.1 逐元素算子的常见模式\n", | 87 | + "## 2 逐元素算子的常见模式\n", |
| 79 | "\n", | 88 | "\n", |
| 80 | "逐元素算子不只包含最基础的加法。实际使用时,它通常会以几种固定模式出现:两个 Tensor 对应位置计算、Tensor 和标量计算、带缩放参数的计算、不同 shape 之间的广播计算,以及对每个元素独立应用数学函数。\n", | 89 | "逐元素算子不只包含最基础的加法。实际使用时,它通常会以几种固定模式出现:两个 Tensor 对应位置计算、Tensor 和标量计算、带缩放参数的计算、不同 shape 之间的广播计算,以及对每个元素独立应用数学函数。\n", |
| 81 | "\n", | 90 | "\n", |
| @@ -185,7 +194,7 @@ | |||
| 185 | "\n", | 194 | "\n", |
| 186 | "1. **函数签名**:`a`、`b`、`out` 都被标注为 `pypto.Tensor([], pypto.DT_FP32)`,表示它们是进入 PyPTO 编译流程的 Tensor 参数;真实数据会在调用 kernel 时传入。\n", | 195 | "1. **函数签名**:`a`、`b`、`out` 都被标注为 `pypto.Tensor([], pypto.DT_FP32)`,表示它们是进入 PyPTO 编译流程的 Tensor 参数;真实数据会在调用 kernel 时传入。\n", |
| 187 | "2. **Tile 设置**:`set_vec_tile_shapes(2, 8)` 表示这个二维向量计算按 `2 x 8` Tile 组织。它不改变 `a + b` 的数学结果。\n", | 196 | "2. **Tile 设置**:`set_vec_tile_shapes(2, 8)` 表示这个二维向量计算按 `2 x 8` Tile 组织。它不改变 `a + b` 的数学结果。\n", |
| 188 | - "3. **输出写回**:`out.move(pypto.add(a, b))` 将加法结果写回输出 Tensor。没有这一步,调用者提前分配的 `out` 不会得到最终结果。\n", | 197 | + "3. **输出写回**:`out[:] = pypto.add(a, b)` 将加法结果写回输出 Tensor。没有这一步,调用者提前分配的 `out` 不会得到最终结果。\n", |
| 189 | "\n", | 198 | "\n", |
| 190 | "用自然语言描述就是:接收两个 Tensor,按 vec tile 组织逐元素加法,然后把结果写入 `out`。" | 199 | "用自然语言描述就是:接收两个 Tensor,按 vec tile 组织逐元素加法,然后把结果写入 `out`。" |
| 191 | ] | 200 | ] |
| @@ -195,7 +204,7 @@ | |||
| 195 | "id": "16e21fbb", | 204 | "id": "16e21fbb", |
| 196 | "metadata": {}, | 205 | "metadata": {}, |
| 197 | "source": [ | 206 | "source": [ |
| 198 | - "这段代码中,`pypto.add(x, y)` 表达逐元素加法,`out.move(...)` 将结果写回输出 Tensor。为了验证它是否正确,我们使用 PyTorch 的 `x + y` 作为参考实现。\n" | 207 | + "这段代码中,`pypto.add(a, b)` 表达逐元素加法,`out[:] = ...` 将结果写回输出 Tensor。为了验证它是否正确,我们使用 PyTorch 的 `x + y` 作为参考实现。\n" |
| 199 | ] | 208 | ] |
| 200 | }, | 209 | }, |
| 201 | { | 210 | { |
| @@ -500,15 +509,43 @@ | |||
| 500 | "out = y * y\n", | 509 | "out = y * y\n", |
| 501 | "```\n", | 510 | "```\n", |
| 502 | "\n", | 511 | "\n", |
| 503 | - "这个练习仍然是逐元素计算。新增的平方操作不会改变 shape,只是在 ReLU 结果的每个位置上再乘一次自身。\n", | 512 | + "这个练习仍然是逐元素计算。新增的平方操作不会改变 shape,只是在 ReLU 结果的每个位置上再乘一次自身。Host 侧代码已给出,只需编写kernel侧 fused_square_practice_kernel 方法的代码。\n" |
| 513 | + ] | ||
| 514 | + }, | ||
| 515 | + { | ||
| 516 | + "cell_type": "code", | ||
| 517 | + "execution_count": null, | ||
| 518 | + "id": "d5f688c2", | ||
| 519 | + "metadata": {}, | ||
| 520 | + "outputs": [], | ||
| 521 | + "source": [ | ||
| 522 | + "# TODO 此处添加kernel函数的实现\n", | ||
| 504 | "\n", | 523 | "\n", |
| 505 | - "建议步骤:\n", | 524 | + "def main_fused_square_practice(device_id: int = None):\n", |
| 525 | + " device_local = get_device()\n", | ||
| 526 | + " x = torch.randn(8, 8, dtype=torch.float16, device=device_local)\n", | ||
| 527 | + " scale = torch.full((8, 8), 1.5, dtype=torch.float16, device=device_local)\n", | ||
| 528 | + " bias = torch.full((8, 8), -0.1, dtype=torch.float16, device=device_local)\n", | ||
| 529 | + " out = torch.empty_like(x)\n", | ||
| 506 | "\n", | 530 | "\n", |
| 507 | - "1. 复制 `fused_mul_add_relu_kernel` 的数据流。\n", | 531 | + " fused_square_practice_kernel(x, scale, bias, out)\n", |
| 508 | - "2. 在 `maximum` 后增加一次逐元素乘法。\n", | ||
| 509 | - "3. 使用 PyTorch 参考实现验证。\n", | ||
| 510 | "\n", | 532 | "\n", |
| 511 | - "**执行以下代码获取答案。**\n" | 533 | + " ref = torch.maximum(x * scale + bias, torch.zeros_like(x))\n", |
| 534 | + " ref = ref * ref\n", | ||
| 535 | + " \n", | ||
| 536 | + " torch.testing.assert_close(out, ref, rtol=3e-3, atol=3e-3)\n", | ||
| 537 | + " max_diff = (out - ref).abs().max().item()\n", | ||
| 538 | + " print(f\"最大误差: {max_diff:.6f}\")\n", | ||
| 539 | + "\n", | ||
| 540 | + "main_fused_square_practice()" | ||
| 541 | + ] | ||
| 542 | + }, | ||
| 543 | + { | ||
| 544 | + "cell_type": "markdown", | ||
| 545 | + "id": "665dc084", | ||
| 546 | + "metadata": {}, | ||
| 547 | + "source": [ | ||
| 548 | + "**执行以下代码获取答案。**" | ||
| 512 | ] | 549 | ] |
| 513 | }, | 550 | }, |
| 514 | { | 551 | { |
| @@ -528,7 +565,7 @@ | |||
| 528 | "source": [ | 565 | "source": [ |
| 529 | "## 6. 本节小结\n", | 566 | "## 6. 本节小结\n", |
| 530 | "\n", | 567 | "\n", |
| 531 | - "本节从 shape 最直观的逐元素计算开始,依次解释了标量、广播、vec tile、输出写回和验证闭环。补充部分覆盖了 kernel 内 Tensor 创建(arange、full)和逐元素 API 全量清单。\n", | 568 | + "本节从 shape 最直观的逐元素计算开始,依次解释了标量、广播、vec tile、输出写回和验证闭环。\n", |
| 532 | "\n", | 569 | "\n", |
| 533 | "需要重点记住:逐元素算子的核心是各位置独立计算;广播解决 shape 对齐问题;`set_vec_tile_shapes` 组织执行但不改变数学结果;PyTorch 参考实现用于证明 PyPTO kernel 的输出正确。" | 570 | "需要重点记住:逐元素算子的核心是各位置独立计算;广播解决 shape 对齐问题;`set_vec_tile_shapes` 组织执行但不改变数学结果;PyTorch 参考实现用于证明 PyPTO kernel 的输出正确。" |
| 534 | ] | 571 | ] |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "a3d23a83", | 5 | "id": "a3d23a83", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.03 矩阵乘法与 set_cube_tile_shapes\n", | 8 | + "# 矩阵乘法与 set_cube_tile_shapes\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "上一节学习的逐元素算子主要依赖向量计算,每个输出位置通常只依赖对应位置的输入。本节进入另一类非常重要的计算:矩阵乘法。矩阵乘法是深度学习模型中的核心计算,线性层、MLP、Attention Score 都会大量使用它。\n", | 10 | "上一节学习的逐元素算子主要依赖向量计算,每个输出位置通常只依赖对应位置的输入。本节进入另一类非常重要的计算:矩阵乘法。矩阵乘法是深度学习模型中的核心计算,线性层、MLP、Attention Score 都会大量使用它。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -32,6 +32,7 @@ | |||
| 32 | "source": [ | 32 | "source": [ |
| 33 | "import os\n", | 33 | "import os\n", |
| 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 35 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 35 | "import torch\n", | 36 | "import torch\n", |
| 36 | "import pypto\n", | 37 | "import pypto\n", |
| 37 | "import torch_npu\n", | 38 | "import torch_npu\n", |
| @@ -97,6 +98,14 @@ | |||
| 97 | "理解这三个维度非常重要,因为后续的 Cube Tile 设置、矩阵切分和 Attention Score 都离不开 `M、K、N`。\n" | 98 | "理解这三个维度非常重要,因为后续的 Cube Tile 设置、矩阵切分和 Attention Score 都离不开 `M、K、N`。\n" |
| 98 | ] | 99 | ] |
| 99 | }, | 100 | }, |
| 101 | + { | ||
| 102 | + "cell_type": "markdown", | ||
| 103 | + "id": "44c5f67b", | ||
| 104 | + "metadata": {}, | ||
| 105 | + "source": [ | ||
| 106 | + "<img src=\"./images/matrix_multiplication.png\" alt=\"矩阵乘法\" width=\"700px\" >" | ||
| 107 | + ] | ||
| 108 | + }, | ||
| 100 | { | 109 | { |
| 101 | "cell_type": "markdown", | 110 | "cell_type": "markdown", |
| 102 | "id": "44699a94", | 111 | "id": "44699a94", |
| @@ -245,10 +254,8 @@ | |||
| 245 | "### 代码细节解释\n", | 254 | "### 代码细节解释\n", |
| 246 | "\n", | 255 | "\n", |
| 247 | "- `set_cube_tile_shapes([32, 32], [64, 64], [64, 64])`:为矩阵乘法设置 Cube Tile。它面向矩阵乘法的分块执行,不等同于输入 Tensor 的完整 shape。\n", | 256 | "- `set_cube_tile_shapes([32, 32], [64, 64], [64, 64])`:为矩阵乘法设置 Cube Tile。它面向矩阵乘法的分块执行,不等同于输入 Tensor 的完整 shape。\n", |
| 248 | - "- `pypto.matmul(a, b, out_dtype=out.dtype)`:表达矩阵乘法,并将输出 dtype 与 `out` 对齐。\n", | 257 | + "- `pypto.matmul(a, b, pypto.DT_FP32)`:表达矩阵乘法。\n", |
| 249 | - "- `out.move(...)`:把矩阵乘法表达式结果写回输出 Tensor。\n", | 258 | + "- `out.move(...)`:把矩阵乘法表达式结果写回输出 Tensor。\n" |
| 250 | - "\n", | ||
| 251 | - "由于 BF16 的数值精度低于 FP32,验证时误差阈值通常需要比 FP32 更宽松。\n" | ||
| 252 | ] | 259 | ] |
| 253 | }, | 260 | }, |
| 254 | { | 261 | { |
| @@ -267,7 +274,6 @@ | |||
| 267 | "- `a` 和 `b` 是输入 Tensor。\n", | 274 | "- `a` 和 `b` 是输入 Tensor。\n", |
| 268 | "- `out` 是输出 Tensor。\n", | 275 | "- `out` 是输出 Tensor。\n", |
| 269 | "- `pypto.matmul(...)` 是真正表达矩阵乘法的地方。\n", | 276 | "- `pypto.matmul(...)` 是真正表达矩阵乘法的地方。\n", |
| 270 | - "- `out_dtype=out.dtype` 表示输出数字类型跟 `out` 保持一致。\n", | ||
| 271 | "- `set_cube_tile_shapes(...)` 是告诉 PyPTO 这是一类适合按矩阵块组织的计算。\n" | 277 | "- `set_cube_tile_shapes(...)` 是告诉 PyPTO 这是一类适合按矩阵块组织的计算。\n" |
| 272 | ] | 278 | ] |
| 273 | }, | 279 | }, |
| @@ -278,10 +284,8 @@ | |||
| 278 | "source": [ | 284 | "source": [ |
| 279 | "这里有两个点需要注意:\n", | 285 | "这里有两个点需要注意:\n", |
| 280 | "\n", | 286 | "\n", |
| 281 | - "1. `pypto.matmul(a, b, out_dtype=out.dtype)` 明确指定矩阵乘法输出 dtype。\n", | 287 | + "1. `pypto.matmul(a, b, pypto.DT_FP32)` 明确指定矩阵乘法输出 dtype。\n", |
| 282 | - "2. 输出写回使用 `out.move(...)`,表示将矩阵乘法结果写入调用者传入的输出 Tensor。\n", | 288 | + "2. 输出写回使用 `out.move(...)`,表示将矩阵乘法结果写入调用者传入的输出 Tensor。\n" |
| 283 | - "\n", | ||
| 284 | - "下面的代码模块会构造 `[64, 128] @ [128, 64] -> [64, 64]` 的输入,调用 kernel,并用 `torch.matmul` 做参考验证。" | ||
| 285 | ] | 289 | ] |
| 286 | }, | 290 | }, |
| 287 | { | 291 | { |
| @@ -672,14 +676,6 @@ | |||
| 672 | "source": [ | 676 | "source": [ |
| 673 | "## 9. 课后实践\n", | 677 | "## 9. 课后实践\n", |
| 674 | "\n", | 678 | "\n", |
| 675 | - "练习目标是在 `matmul_bias_kernel` 基础上增加 ReLU,完成:\n", | ||
| 676 | - "\n", | ||
| 677 | - "```python\n", | ||
| 678 | - "out = maximum(a @ b + bias, 0)\n", | ||
| 679 | - "```\n", | ||
| 680 | - "\n", | ||
| 681 | - "思考题:\n", | ||
| 682 | - "\n", | ||
| 683 | "1. (选择题)矩阵乘法部分应该使用向量 Tile 还是 Cube Tile? \n", | 679 | "1. (选择题)矩阵乘法部分应该使用向量 Tile 还是 Cube Tile? \n", |
| 684 | " A. 向量 Tile \n", | 680 | " A. 向量 Tile \n", |
| 685 | " B. Cube Tile\n", | 681 | " B. Cube Tile\n", |
| @@ -692,9 +688,44 @@ | |||
| 692 | " C. 避免创建输出 Tensor \n", | 688 | " C. 避免创建输出 Tensor \n", |
| 693 | " D. 跳过 PyTorch 参考验证\n", | 689 | " D. 跳过 PyTorch 参考验证\n", |
| 694 | "\n", | 690 | "\n", |
| 695 | - "编程题:补全 `matmul_bias_relu_kernel`,先完成 `a @ b + bias`,再对结果做 ReLU,并写回输出 Tensor。\n", | 691 | + "4. (编程题)实现带bias的矩阵乘法,再对结果做 ReLU。结果要和PyaTorch进行比较。Host 侧代码已给出,只需编写kernel侧 matmul_bias_relu_kernel 方法的代码。" |
| 692 | + ] | ||
| 693 | + }, | ||
| 694 | + { | ||
| 695 | + "cell_type": "code", | ||
| 696 | + "execution_count": null, | ||
| 697 | + "id": "ce6acf79", | ||
| 698 | + "metadata": {}, | ||
| 699 | + "outputs": [], | ||
| 700 | + "source": [ | ||
| 701 | + "# TODO 此处添加kernel函数的实现\n", | ||
| 696 | "\n", | 702 | "\n", |
| 697 | - "**执行以下代码获取答案。**\n" | 703 | + "\n", |
| 704 | + "def test():\n", | ||
| 705 | + " a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, device=device)\n", | ||
| 706 | + " b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32, device=device)\n", | ||
| 707 | + " bias = torch.tensor([[1, 2]], dtype=torch.float32, device=device)\n", | ||
| 708 | + " out = torch.empty((2, 2), dtype=torch.float32, device=device)\n", | ||
| 709 | + " matmul_bias_relu_kernel(a, b, bias, out)\n", | ||
| 710 | + " ref = torch.maximum(torch.matmul(a, b) + bias, torch.tensor(0.0, dtype=torch.float32, device=device))\n", | ||
| 711 | + " max_diff = (out - ref).abs().max().item()\n", | ||
| 712 | + " torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3)\n", | ||
| 713 | + " print(\"matmul_bias_relu_kernel 验证通过\")\n", | ||
| 714 | + " print(\"device:\", device, \"run_mode:\", RUN_MODE)\n", | ||
| 715 | + " print(\"a shape:\", tuple(a.shape), \"b shape:\", tuple(b.shape), \"bias shape:\", tuple(bias.shape))\n", | ||
| 716 | + " print(\"输出:\", out.cpu())\n", | ||
| 717 | + " print(\"参考:\", ref.cpu())\n", | ||
| 718 | + " print(\"最大误差:\", max_diff)\n", | ||
| 719 | + "\n", | ||
| 720 | + "test()\n" | ||
| 721 | + ] | ||
| 722 | + }, | ||
| 723 | + { | ||
| 724 | + "cell_type": "markdown", | ||
| 725 | + "id": "3da86fb3", | ||
| 726 | + "metadata": {}, | ||
| 727 | + "source": [ | ||
| 728 | + "**执行以下代码获取答案。**" | ||
| 698 | ] | 729 | ] |
| 699 | }, | 730 | }, |
| 700 | { | 731 | { |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "ba301c08", | 5 | "id": "ba301c08", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.04 规约算子\n", | 8 | + "# 规约算子\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "前两节分别学习了逐元素计算和矩阵乘法。本节学习第三类基础计算:规约。逐元素计算通常是“一个位置对应一个位置”,矩阵乘法是“行和列做乘加”,而规约是“沿某个维度把一组元素合成更少的元素”。\n", | 10 | "前两节分别学习了逐元素计算和矩阵乘法。本节学习第三类基础计算:规约。逐元素计算通常是“一个位置对应一个位置”,矩阵乘法是“行和列做乘加”,而规约是“沿某个维度把一组元素合成更少的元素”。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -32,10 +32,11 @@ | |||
| 32 | "source": [ | 32 | "source": [ |
| 33 | "import os\n", | 33 | "import os\n", |
| 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 34 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 35 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 35 | "import torch\n", | 36 | "import torch\n", |
| 36 | "import pypto\n", | 37 | "import pypto\n", |
| 37 | "import torch_npu\n", | 38 | "import torch_npu\n", |
| 38 | - "\n", | 39 | + "import numpy as np\n", |
| 39 | "\n", | 40 | "\n", |
| 40 | "def get_device():\n", | 41 | "def get_device():\n", |
| 41 | " device_id = int(os.environ.get(\"TILE_FWK_DEVICE_ID\", \"0\"))\n", | 42 | " device_id = int(os.environ.get(\"TILE_FWK_DEVICE_ID\", \"0\"))\n", |
| @@ -196,7 +197,7 @@ | |||
| 196 | "pypto.sum(x, dim=-1, keepdim=False)\n", | 197 | "pypto.sum(x, dim=-1, keepdim=False)\n", |
| 197 | "```\n", | 198 | "```\n", |
| 198 | "\n", | 199 | "\n", |
| 199 | - "输出形状是 `[2]`。它表示每一行得到一个数字,但这个结果已经不再是二维列向量。\n", | 200 | + "输出形状是 `[2]`。它表示每一行得到一个数字,但这个结果已经不再是二维Tensor。\n", |
| 200 | "\n", | 201 | "\n", |
| 201 | "而:\n", | 202 | "而:\n", |
| 202 | "\n", | 203 | "\n", |
| @@ -340,7 +341,7 @@ | |||
| 340 | "\n", | 341 | "\n", |
| 341 | "- `pypto.sum(x, dim=-1, keepdim=False)` 沿最后一个维度求和。\n", | 342 | "- `pypto.sum(x, dim=-1, keepdim=False)` 沿最后一个维度求和。\n", |
| 342 | "- 输入 shape `[2, 3]` 会变成输出 shape `[2]`。\n", | 343 | "- 输入 shape `[2, 3]` 会变成输出 shape `[2]`。\n", |
| 343 | - "- 这里使用 `set_vec_tile_shapes(2, 8)`,不是因为输入有 8 列,而是因为规约最后一维的 Tile 需要满足 32Byte 对齐。\n", | 344 | + "- 这里使用 `set_vec_tile_shapes(8, 8)`,不是因为输入有 8 列,而是因为规约最后一维的 Tile 需要满足 32Byte 对齐。\n", |
| 344 | "- 当前 dtype 是 FP32,每个元素 4Byte,因此最后一维 Tile 至少按 8 个元素对齐,`8 * 4Byte = 32Byte`。\n", | 345 | "- 当前 dtype 是 FP32,每个元素 4Byte,因此最后一维 Tile 至少按 8 个元素对齐,`8 * 4Byte = 32Byte`。\n", |
| 345 | "- 输入实际 shape 仍然是 `[2, 3]`,多出来的 Tile 宽度是执行组织上的对齐要求,不改变数学输入和输出。\n", | 346 | "- 输入实际 shape 仍然是 `[2, 3]`,多出来的 Tile 宽度是执行组织上的对齐要求,不改变数学输入和输出。\n", |
| 346 | "\n", | 347 | "\n", |
| @@ -377,7 +378,7 @@ | |||
| 377 | "id": "9e25fc52", | 378 | "id": "9e25fc52", |
| 378 | "metadata": {}, | 379 | "metadata": {}, |
| 379 | "source": [ | 380 | "source": [ |
| 380 | - "上面的完整单元用来验证 `row_sum_kernel`。它先在当前设备上构造一个二维输入,再提前分配一维输出 Tensor 并调用 PyPTO kernel。随后用普通 PyTorch 写出等价的按行求和参考结果,打印输入、输出、参考值和最大误差;最后的 `assert_close` 是正式检查。\n" | 381 | + "上面的完整单元用来验证 `sum_kernel`。它先在当前设备上构造一个二维输入,再提前分配一维输出 Tensor 并调用 PyPTO kernel。随后用普通 PyTorch 写出等价的按行求和参考结果,打印输入、输出、参考值和最大误差;最后的 `assert_close` 是正式检查。\n" |
| 381 | ] | 382 | ] |
| 382 | }, | 383 | }, |
| 383 | { | 384 | { |
| @@ -397,17 +398,10 @@ | |||
| 397 | "source": [ | 398 | "source": [ |
| 398 | "## 5.2 最大值和最小值规约:amax 与 amin\n", | 399 | "## 5.2 最大值和最小值规约:amax 与 amin\n", |
| 399 | "\n", | 400 | "\n", |
| 400 | - "求和不是唯一的规约。`pypto.amax` 和 `pypto.amin` 也会沿某个维度把一组元素合成一个元素,只是合成方式从“相加”变成“取最大值”或“取最小值”。\n", | 401 | + "求和不是唯一的规约。`pypto.amax` 和 `pypto.amin` 也会沿某个维度把一组元素合成一个元素,只是合成方式从“相加”变成“取最大值”或“取最小值”。 \n", |
| 401 | - "\n", | 402 | + "对于输入:`[[1, 2, 3],[4, 5, 6]]` \n", |
| 402 | - "对于输入:\n", | 403 | + "沿 `dim=-1` 做 `amax` 可以得到 `[3, 6]`和`[[3], [6]]`两种结果,这取决于是否要改变 shape。 \n", |
| 403 | - "\n", | 404 | + "以下是amax的操作样例,amin操作逻辑一致,不再额外举例。\n" |
| 404 | - "```python\n", | ||
| 405 | - "[[1, 2, 3],\n", | ||
| 406 | - " [4, 5, 6]]\n", | ||
| 407 | - "```\n", | ||
| 408 | - "\n", | ||
| 409 | - "沿 `dim=-1` 做 `amax` 得到 `[3, 6]`,`[[3], [6]]`。这取决于是否要改变 shape。\n", | ||
| 410 | - "amin和amax操作逻辑一致,不再额外举例。\n" | ||
| 411 | ] | 405 | ] |
| 412 | }, | 406 | }, |
| 413 | { | 407 | { |
| @@ -428,7 +422,7 @@ | |||
| 428 | " out_shape.pop(dim)\n", | 422 | " out_shape.pop(dim)\n", |
| 429 | " out_shape = tuple(out_shape)\n", | 423 | " out_shape = tuple(out_shape)\n", |
| 430 | "\n", | 424 | "\n", |
| 431 | - " @pypto.frontend.jit(runtime_options={\"run_mode\": global_run_mode})\n", | 425 | + " @pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", |
| 432 | " def amax_kernel(a: pypto.Tensor([], dtype),\n", | 426 | " def amax_kernel(a: pypto.Tensor([], dtype),\n", |
| 433 | " out: pypto.Tensor([], dtype)):\n", | 427 | " out: pypto.Tensor([], dtype)):\n", |
| 434 | " tile_shapes = [8] * len(a.shape)\n", | 428 | " tile_shapes = [8] * len(a.shape)\n", |
| @@ -466,7 +460,8 @@ | |||
| 466 | " print(f\"Max difference: {max_diff:.6f}\")\n", | 460 | " print(f\"Max difference: {max_diff:.6f}\")\n", |
| 467 | " torch.testing.assert_close(out.cpu().numpy(), expected.cpu().numpy(), rtol=1e-3, atol=1e-3)\n", | 461 | " torch.testing.assert_close(out.cpu().numpy(), expected.cpu().numpy(), rtol=1e-3, atol=1e-3)\n", |
| 468 | "\n", | 462 | "\n", |
| 469 | - " print(\"✓ Basic usage of amax function completed successfully\")" | 463 | + "\n", |
| 464 | + "test_amax_basic()" | ||
| 470 | ] | 465 | ] |
| 471 | }, | 466 | }, |
| 472 | { | 467 | { |
| @@ -734,7 +729,7 @@ | |||
| 734 | "source": [ | 729 | "source": [ |
| 735 | "## 9. 本节小结\n", | 730 | "## 9. 本节小结\n", |
| 736 | "\n", | 731 | "\n", |
| 737 | - "本节学习了规约算子的基本用法和全量 API 覆盖。理解 `dim` 和 `keepdim` 是后续学习 Softmax、RMSNorm、LayerNorm 的关键。规约要始终和 shape 变化一起看:`sum/amax/amin` 会沿某个维度汇总,`maximum/minimum` 等逐元素操作则依赖对应位置或广播。下一节会进一步学习 Tiling 和形状操作,理解 Tensor 在计算前后如何被组织。" | 732 | + "本节学习了规约算子的基本用法。理解 `dim` 和 `keepdim` 是后续学习 Softmax、RMSNorm、LayerNorm 的关键。规约要始终和 shape 变化一起看:`sum/amax/amin` 会沿某个维度汇总,`maximum/minimum` 等逐元素操作则依赖对应位置或广播。下一节会进一步学习 Tiling 和形状操作,理解 Tensor 在计算前后如何被组织。" |
| 738 | ] | 733 | ] |
| 739 | } | 734 | } |
| 740 | ], | 735 | ], |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "d2305ff3", | 5 | "id": "d2305ff3", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.05 Tiling、形状变换、切片与转置\n", | 8 | + "# Tiling、形状变换、切片与转置\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "前面几节已经完成逐元素、矩阵乘法和规约计算。本节把视角从“算什么”扩展到“数据如何组织”。真实算子开发中,除了数学计算,还要频繁处理 Tiling、reshape、切片和转置。\n", | 10 | "前面几节已经完成逐元素、矩阵乘法和规约计算。本节把视角从“算什么”扩展到“数据如何组织”。真实算子开发中,除了数学计算,还要频繁处理 Tiling、reshape、切片和转置。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -36,6 +36,7 @@ | |||
| 36 | "source": [ | 36 | "source": [ |
| 37 | "import os\n", | 37 | "import os\n", |
| 38 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 38 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 39 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 39 | "import time\n", | 40 | "import time\n", |
| 40 | "import torch\n", | 41 | "import torch\n", |
| 41 | "import pypto\n", | 42 | "import pypto\n", |
| @@ -157,7 +158,9 @@ | |||
| 157 | "metadata": {}, | 158 | "metadata": {}, |
| 158 | "source": [ | 159 | "source": [ |
| 159 | "测试矩阵乘法的`cube tile`。`vec tile`类似不再举例。 \n", | 160 | "测试矩阵乘法的`cube tile`。`vec tile`类似不再举例。 \n", |
| 160 | - "验证内容:使用Tiling不影响结果准确性;使用不同的Tiling不影响结果准确性;使用不同的Tiling,执行runtime不同。" | 161 | + "\n", |
| 162 | + "\n", | ||
| 163 | + "验证内容:使用 Tiling 不影响结果准确性。compute_with_cube_tile_shapes_kernel方法使用了set_cube_tile_shapes,但是结果保持不变。" | ||
| 161 | ] | 164 | ] |
| 162 | }, | 165 | }, |
| 163 | { | 166 | { |
| @@ -193,6 +196,14 @@ | |||
| 193 | "test_set_cube_tile_shapes_basic()" | 196 | "test_set_cube_tile_shapes_basic()" |
| 194 | ] | 197 | ] |
| 195 | }, | 198 | }, |
| 199 | + { | ||
| 200 | + "cell_type": "markdown", | ||
| 201 | + "id": "83ea84b6", | ||
| 202 | + "metadata": {}, | ||
| 203 | + "source": [ | ||
| 204 | + "验证内容:使用不同的 Tiling 不影响结果准确性。compute_with_different_tile_shapes方法中,分别使用不同的 Tiling 参数,对相同的输入做矩阵乘法得到out1、out2、out3,结果一致。" | ||
| 205 | + ] | ||
| 206 | + }, | ||
| 196 | { | 207 | { |
| 197 | "cell_type": "code", | 208 | "cell_type": "code", |
| 198 | "execution_count": null, | 209 | "execution_count": null, |
| @@ -224,12 +235,9 @@ | |||
| 224 | " pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128])\n", | 235 | " pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128])\n", |
| 225 | " print(f\"pypto.get_cube_tile_shapes(): {pypto.get_cube_tile_shapes()}\")\n", | 236 | " print(f\"pypto.get_cube_tile_shapes(): {pypto.get_cube_tile_shapes()}\")\n", |
| 226 | " out3[:] = pypto.matmul(x, y, x.dtype)\n", | 237 | " out3[:] = pypto.matmul(x, y, x.dtype)\n", |
| 227 | - "\n", | ||
| 228 | - "\n", | ||
| 229 | " return compute_with_different_tile_shapes\n", | 238 | " return compute_with_different_tile_shapes\n", |
| 230 | "\n", | 239 | "\n", |
| 231 | "\n", | 240 | "\n", |
| 232 | - "\n", | ||
| 233 | "def test_set_different_tile_shapes_result(device_id: int = None):\n", | 241 | "def test_set_different_tile_shapes_result(device_id: int = None):\n", |
| 234 | " \"\"\"Test the impact of different tile shape settings on calculation results\"\"\"\n", | 242 | " \"\"\"Test the impact of different tile shape settings on calculation results\"\"\"\n", |
| 235 | " print(\"=\" * 60)\n", | 243 | " print(\"=\" * 60)\n", |
| @@ -261,6 +269,14 @@ | |||
| 261 | "test_set_different_tile_shapes_result()" | 269 | "test_set_different_tile_shapes_result()" |
| 262 | ] | 270 | ] |
| 263 | }, | 271 | }, |
| 272 | + { | ||
| 273 | + "cell_type": "markdown", | ||
| 274 | + "id": "f97b82bb", | ||
| 275 | + "metadata": {}, | ||
| 276 | + "source": [ | ||
| 277 | + "验证内容:使用不同的Tiling,执行runtime不同。对相同的输入,compute_with_tile_32_kernel方法和compute_with_tile_64_kernel方法使用不同的 Tiling 参数,循环TEST_TIME次(使用者可以自行设定),执行时间不同。" | ||
| 278 | + ] | ||
| 279 | + }, | ||
| 264 | { | 280 | { |
| 265 | "cell_type": "code", | 281 | "cell_type": "code", |
| 266 | "execution_count": null, | 282 | "execution_count": null, |
| @@ -5,209 +5,79 @@ | |||
| 5 | "id": "d13c1afd", | 5 | "id": "d13c1afd", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 03.06 章节实践:实现行 Softmax\n", | 8 | + "# 章节实践\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | - "本节是第三章的综合实践。前面几节分别学习了逐元素计算、矩阵乘法、规约计算、Tiling、形状变换、切片和转置。本节将逐元素、规约和广播组合起来,实现一个按最后一维计算的稳定版 Softmax 算子。\n", | 10 | + "本节是第三章的综合实践。前面几节分别学习了逐元素计算、矩阵乘法、规约计算、Tiling、形状变换、切片和转置。\n", |
| 11 | "\n", | 11 | "\n", |
| 12 | - "本节不引入新的复杂 API,而是把已经学过的内容组织成一个完整算子:`amax/sum` 负责规约,`sub/exp/div` 负责逐元素计算,`keepdim=True` 负责保留广播所需的维度,最后用 PyTorch 的 `torch.softmax` 建立验证闭环。\n" | 12 | + "\n" |
| 13 | ] | 13 | ] |
| 14 | }, | 14 | }, |
| 15 | { | 15 | { |
| 16 | "cell_type": "markdown", | 16 | "cell_type": "markdown", |
| 17 | - "id": "745fd0a7", | 17 | + "id": "e9cd54a8", |
| 18 | "metadata": {}, | 18 | "metadata": {}, |
| 19 | "source": [ | 19 | "source": [ |
| 20 | - "## 1. 实践任务总览\n", | 20 | + "## 1. 实践目标\n", |
| 21 | "\n", | 21 | "\n", |
| 22 | - "Softmax 可以先理解为“把一组分数变成一组比例”。例如一行里有几个分数:\n", | 22 | + "本节将逐元素、规约和广播组合起来,实现一个按最后一维计算的稳定版 Softmax 算子。 \n", |
| 23 | - "\n", | 23 | + "本节不引入新的复杂 API,而是把已经学过的内容组织成一个完整算子,其中涉及规约计算,逐元素计算,最后用 PyTorch 的 `torch.softmax` 进行验证。\n" |
| 24 | - "```python\n", | ||
| 25 | - "[2.0, 1.0, 0.1]\n", | ||
| 26 | - "```\n", | ||
| 27 | - "\n", | ||
| 28 | - "Softmax 会把它们变成类似“概率”的数字:\n", | ||
| 29 | - "\n", | ||
| 30 | - "- 每个数字都大于 0。\n", | ||
| 31 | - "- 加起来约等于 1。\n", | ||
| 32 | - "- 原来分数越大的位置,转换后通常也越大。\n", | ||
| 33 | - "\n", | ||
| 34 | - "因此,行 Softmax 很适合作为第三章的章节实践:它同时用到逐元素计算、规约计算、广播、shape 保持和输出写回。\n" | ||
| 35 | ] | 24 | ] |
| 36 | }, | 25 | }, |
| 37 | { | 26 | { |
| 38 | "cell_type": "markdown", | 27 | "cell_type": "markdown", |
| 39 | - "id": "cfc128c5", | 28 | + "id": "d0dfc350", |
| 40 | "metadata": {}, | 29 | "metadata": {}, |
| 41 | "source": [ | 30 | "source": [ |
| 42 | - "## 2. Softmax 的数学含义\n", | 31 | + "## 2. 章节实践题\n", |
| 32 | + "实现行 Softmax,Host 侧代码已给出,只需编写kernel侧 main_row_softmax_practice 方法的代码。 \n", | ||
| 43 | "\n", | 33 | "\n", |
| 44 | - "Softmax 常用于把一组分数转换为概率分布。对于一行输入 `x`,Softmax 可以写成:\n", | 34 | + "提示: \n", |
| 35 | + " - 行 Softmax 针对一个 矩阵(2维),独立地对每一行应用 Softmax 函数,由于行 Softmax 是对每一行独立算指数,依然存在数值溢出风险(如行内有极大值)。实际代码实现中,行 Softmax 内部会先对当前行减去该行的最大值。 \n", | ||
| 36 | + " - softmax步骤: \n", | ||
| 37 | + " - 步骤 1:输入(向量,数值列表)。 \n", | ||
| 38 | + " - 步骤 2:指数化(对每个值取自然常数 e 的幂)。原因:使所有值为正,并放大差异。 \n", | ||
| 39 | + " - 步骤 3:求和(将所有指数值相加得到分母)。 \n", | ||
| 40 | + " - 步骤 4:归一化(用每个指数值除以总和)。 \n", | ||
| 41 | + " - 步骤 5:输出(0 到 1 之间的概率,总和为 1)。 " | ||
| 42 | + ] | ||
| 43 | + }, | ||
| 44 | + { | ||
| 45 | + "cell_type": "code", | ||
| 46 | + "execution_count": null, | ||
| 47 | + "id": "88f68d33", | ||
| 48 | + "metadata": {}, | ||
| 49 | + "outputs": [], | ||
| 50 | + "source": [ | ||
| 51 | + "# TODO 此处添加kernel函数的实现\n", | ||
| 45 | "\n", | 52 | "\n", |
| 46 | - "```python\n", | ||
| 47 | - "softmax(x_i) = exp(x_i) / sum(exp(x_i))\n", | ||
| 48 | - "```\n", | ||
| 49 | "\n", | 53 | "\n", |
| 50 | - "输出通常具有两个特点:\n", | 54 | + "def main_row_softmax_practice():\n", |
| 55 | + " x = torch.randn((8, 8), dtype=torch.float32, device=device)\n", | ||
| 56 | + " out = torch.empty_like(x)\n", | ||
| 51 | "\n", | 57 | "\n", |
| 52 | - "1. 每个元素大于 0。\n", | 58 | + " row_softmax_practice_kernel(x, out)\n", |
| 53 | - "2. 每一行元素之和接近 1。\n", | ||
| 54 | "\n", | 59 | "\n", |
| 55 | - "在注意力机制中,Softmax 常用于把 `q @ k.T` 得到的分数转换成权重。\n" | 60 | + " ref = torch.softmax(x, dim=-1)\n", |
| 61 | + " max_diff = (out - ref).abs().max().item()\n", | ||
| 62 | + " torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3)\n", | ||
| 63 | + "\n", | ||
| 64 | + " print(\"row_softmax_practice_kernel 验证通过\")\n", | ||
| 65 | + " print(\"输入 shape:\", tuple(x.shape), \"输出 shape:\", tuple(out.shape))\n", | ||
| 66 | + " print(\"输出每行求和:\", out.sum(dim=-1).detach().cpu())\n", | ||
| 67 | + " print(\"最大误差:\", max_diff)\n", | ||
| 68 | + "\n", | ||
| 69 | + "\n", | ||
| 70 | + "main_row_softmax_practice()" | ||
| 56 | ] | 71 | ] |
| 57 | }, | 72 | }, |
| 58 | { | 73 | { |
| 59 | "cell_type": "markdown", | 74 | "cell_type": "markdown", |
| 60 | - "id": "1302f628", | 75 | + "id": "e75458cb", |
| 61 | "metadata": {}, | 76 | "metadata": {}, |
| 62 | "source": [ | 77 | "source": [ |
| 63 | - "## 3. 为什么要减去 row_max\n", | 78 | + "## 3. 查看答案\n", |
| 64 | "\n", | 79 | "\n", |
| 65 | - "直接计算 `exp(x)` 可能遇到数值溢出。例如某些输入值很大时,指数结果会非常大。为了提升数值稳定性,常见做法是先减去每一行最大值:\n", | 80 | + "执行以下代码获取参考答案。\n" |
| 66 | - "\n", | ||
| 67 | - "```python\n", | ||
| 68 | - "row_max = max(x)\n", | ||
| 69 | - "shifted = x - row_max\n", | ||
| 70 | - "```\n", | ||
| 71 | - "\n", | ||
| 72 | - "这个操作不会改变 Softmax 的数学结果,因为同一行所有元素都减去了同一个常数,最终归一化后的比例保持不变。\n", | ||
| 73 | - "\n", | ||
| 74 | - "在 PyPTO 中,我们使用:\n", | ||
| 75 | - "\n", | ||
| 76 | - "```python\n", | ||
| 77 | - "row_max = pypto.amax(x, dim=-1, keepdim=True)\n", | ||
| 78 | - "```\n", | ||
| 79 | - "\n", | ||
| 80 | - "这里 `keepdim=True` 很重要,它让 `row_max` 保持 `[M, 1]` 形状,方便与原始 `[M, N]` 的 `x` 做广播相减。\n" | ||
| 81 | - ] | ||
| 82 | - }, | ||
| 83 | - { | ||
| 84 | - "cell_type": "markdown", | ||
| 85 | - "id": "5b4e2c8e", | ||
| 86 | - "metadata": {}, | ||
| 87 | - "source": [ | ||
| 88 | - "### 为什么减去最大值不会改变 Softmax 结果\n", | ||
| 89 | - "\n", | ||
| 90 | - "可以先用直觉理解:Softmax 关心的是同一行里数字之间的相对大小,而不是它们整体加了多少。\n", | ||
| 91 | - "\n", | ||
| 92 | - "例如:\n", | ||
| 93 | - "\n", | ||
| 94 | - "```python\n", | ||
| 95 | - "[10, 11, 12]\n", | ||
| 96 | - "```\n", | ||
| 97 | - "\n", | ||
| 98 | - "每个数字都减去 12 后变成:\n", | ||
| 99 | - "\n", | ||
| 100 | - "```python\n", | ||
| 101 | - "[-2, -1, 0]\n", | ||
| 102 | - "```\n", | ||
| 103 | - "\n", | ||
| 104 | - "最大的位置仍然是原来的最大位置,数字之间的差距也没有变。这样做能让 `exp` 计算更稳定,因为输入不会那么大。\n" | ||
| 105 | - ] | ||
| 106 | - }, | ||
| 107 | - { | ||
| 108 | - "cell_type": "markdown", | ||
| 109 | - "id": "89630b24", | ||
| 110 | - "metadata": {}, | ||
| 111 | - "source": [ | ||
| 112 | - "## 4. 计算流程拆解\n", | ||
| 113 | - "\n", | ||
| 114 | - "行 Softmax 可以拆成五步:\n", | ||
| 115 | - "\n", | ||
| 116 | - "| 步骤 | PyPTO 表达 | 作用 |\n", | ||
| 117 | - "| --- | --- | --- |\n", | ||
| 118 | - "| 1 | `pypto.amax(x, dim=-1, keepdim=True)` | 得到每一行最大值 |\n", | ||
| 119 | - "| 2 | `x - row_max` | 提升数值稳定性 |\n", | ||
| 120 | - "| 3 | `pypto.exp(shifted)` | 计算指数 |\n", | ||
| 121 | - "| 4 | `pypto.sum(exp, dim=-1, keepdim=True)` | 得到每一行指数和 |\n", | ||
| 122 | - "| 5 | `exp / esum` | 归一化 |\n", | ||
| 123 | - "\n", | ||
| 124 | - "通过这个拆解可以看到,Softmax 并不是单个神秘操作,而是规约和逐元素计算的组合。\n" | ||
| 125 | - ] | ||
| 126 | - }, | ||
| 127 | - { | ||
| 128 | - "cell_type": "markdown", | ||
| 129 | - "id": "938a693b", | ||
| 130 | - "metadata": {}, | ||
| 131 | - "source": [ | ||
| 132 | - "## 4.1 从前面的问题回到 Softmax\n", | ||
| 133 | - "\n", | ||
| 134 | - "本章前面的几个问题都可以在 Softmax 中串起来:\n", | ||
| 135 | - "\n", | ||
| 136 | - "1. `amax(x, dim=-1, keepdim=True)` 是 reduction,用于得到每一行最大值。\n", | ||
| 137 | - "2. `x - row_max` 是 elementwise,同时依赖广播,因为 `row_max.shape` 是 `[batch, 1]`。\n", | ||
| 138 | - "3. `exp(shifted)` 是 elementwise,每个位置独立取指数。\n", | ||
| 139 | - "4. `sum(exp, dim=-1, keepdim=True)` 是 reduction,用于得到每一行指数和。\n", | ||
| 140 | - "5. `exp / row_sum` 是 elementwise,同时依赖广播。\n", | ||
| 141 | - "\n", | ||
| 142 | - "Softmax 的输出满足每一行加起来约等于 1:\n", | ||
| 143 | - "\n", | ||
| 144 | - "```text\n", | ||
| 145 | - "softmax_i = exp_i / sum(exp)\n", | ||
| 146 | - "sum(softmax_i) = sum(exp_i / sum(exp)) = 1\n", | ||
| 147 | - "```\n", | ||
| 148 | - "\n", | ||
| 149 | - "因此,Softmax 是本章概念的组合练习:shape、broadcast、elementwise、reduction、keepdim 都会同时出现。" | ||
| 150 | - ] | ||
| 151 | - }, | ||
| 152 | - { | ||
| 153 | - "cell_type": "markdown", | ||
| 154 | - "id": "b8f3bf28", | ||
| 155 | - "metadata": {}, | ||
| 156 | - "source": [ | ||
| 157 | - "## 4.2 与前面基础算子的关系\n", | ||
| 158 | - "\n", | ||
| 159 | - "行 Softmax 正好把本章前面学过的内容组合起来:\n", | ||
| 160 | - "\n", | ||
| 161 | - "| Softmax 步骤 | 计算类型 | 关键点 |\n", | ||
| 162 | - "| --- | --- | --- |\n", | ||
| 163 | - "| `amax(x, dim=-1, keepdim=True)` | reduction | 按行取最大值,保留 `[8, 1]` 形状 |\n", | ||
| 164 | - "| `x - row_max` | elementwise | 依赖广播,提升数值稳定性 |\n", | ||
| 165 | - "| `exp(shifted)` | elementwise | 每个位置独立取指数 |\n", | ||
| 166 | - "| `sum(exp, dim=-1, keepdim=True)` | reduction | 按行求指数和,保留 `[8, 1]` 形状 |\n", | ||
| 167 | - "| `exp / esum` | elementwise | 依赖广播完成归一化 |\n", | ||
| 168 | - "| `out.move(...)` | 输出写回 | 把最终表达式写入 host 侧传入的输出 Tensor |\n", | ||
| 169 | - "\n", | ||
| 170 | - "因此,章节实践不是新 API 的堆叠,而是把逐元素、规约、广播和输出写回组织成一个完整算子。\n" | ||
| 171 | - ] | ||
| 172 | - }, | ||
| 173 | - { | ||
| 174 | - "cell_type": "markdown", | ||
| 175 | - "id": "5f678f59", | ||
| 176 | - "metadata": {}, | ||
| 177 | - "source": [ | ||
| 178 | - "## 5. 编写 PyPTO Softmax 算子\n", | ||
| 179 | - "\n", | ||
| 180 | - "下面将上述五步写入一个 PyPTO JIT 函数。\n" | ||
| 181 | - ] | ||
| 182 | - }, | ||
| 183 | - { | ||
| 184 | - "cell_type": "markdown", | ||
| 185 | - "id": "c073bb30", | ||
| 186 | - "metadata": {}, | ||
| 187 | - "source": [ | ||
| 188 | - "## 5.1 Softmax 算子规格\n", | ||
| 189 | - "\n", | ||
| 190 | - "| 项目 | 说明 |\n", | ||
| 191 | - "| --- | --- |\n", | ||
| 192 | - "| 输入 `x` | FP32 Tensor,示例 shape 为 `[8, 8]` |\n", | ||
| 193 | - "| 输出 `out` | FP32 Tensor,shape 为 `[8, 8]` |\n", | ||
| 194 | - "| 规约维度 | `dim=-1`,按每一行内部计算 |\n", | ||
| 195 | - "| 中间变量 `row_max` | shape 为 `[8, 1]`,用于数值稳定 |\n", | ||
| 196 | - "| 中间变量 `exp` | shape 为 `[8, 8]` |\n", | ||
| 197 | - "| 中间变量 `esum` | shape 为 `[8, 1]`,用于归一化 |\n", | ||
| 198 | - "| 验证参考 | `torch.softmax(x, dim=-1)` |\n", | ||
| 199 | - "\n", | ||
| 200 | - "这个规格表把 Softmax 拆成了可验证的中间步骤,后面阅读代码时会更清晰。\n" | ||
| 201 | - ] | ||
| 202 | - }, | ||
| 203 | - { | ||
| 204 | - "cell_type": "markdown", | ||
| 205 | - "id": "9ea4badc", | ||
| 206 | - "metadata": {}, | ||
| 207 | - "source": [ | ||
| 208 | - "**预期输出说明**\n", | ||
| 209 | - "\n", | ||
| 210 | - "运行成功后,会看到 `row_softmax_kernel 验证通过`,并打印输出前几项和每行求和结果。Softmax 每行求和应接近 1。\n" | ||
| 211 | ] | 81 | ] |
| 212 | }, | 82 | }, |
| 213 | { | 83 | { |
| @@ -225,7 +95,7 @@ | |||
| 225 | "id": "chapter_practice_final_summary", | 95 | "id": "chapter_practice_final_summary", |
| 226 | "metadata": {}, | 96 | "metadata": {}, |
| 227 | "source": [ | 97 | "source": [ |
| 228 | - "## 7. 本章小结\n", | 98 | + "## 4. 本章小结\n", |
| 229 | "\n", | 99 | "\n", |
| 230 | "通过第三章学习,已经建立 PyPTO 初级计算算子的基本开发方法:\n", | 100 | "通过第三章学习,已经建立 PyPTO 初级计算算子的基本开发方法:\n", |
| 231 | "\n", | 101 | "\n", |
| @@ -11,19 +11,3 @@ def fused_square_practice_kernel( | |||
| 11 | squared = pypto.mul(relu, relu) | 11 | squared = pypto.mul(relu, relu) |
| 12 | out.move(squared) | 12 | out.move(squared) |
| 13 | 13 | ||
| 14 | - | ||
| 15 | -def main_fused_square_practice(device_id: int = None): | ||
| 16 | - device_local = get_device() | ||
| 17 | - x = torch.randn(8, 8, dtype=torch.float16, device=device_local) | ||
| 18 | - scale = torch.full((8, 8), 1.5, dtype=torch.float16, device=device_local) | ||
| 19 | - bias = torch.full((8, 8), -0.1, dtype=torch.float16, device=device_local) | ||
| 20 | - out = torch.empty_like(x) | ||
| 21 | - | ||
| 22 | - fused_square_practice_kernel(x, scale, bias, out) | ||
| 23 | - | ||
| 24 | - ref = torch.maximum(x * scale + bias, torch.zeros_like(x)) | ||
| 25 | - ref = ref * ref | ||
| 26 | - | ||
| 27 | - torch.testing.assert_close(out, ref, rtol=3e-3, atol=3e-3) | ||
| 28 | - max_diff = (out - ref).abs().max().item() | ||
| 29 | - print(f"最大误差: {max_diff:.6f}") | ||
| @@ -12,30 +12,7 @@ def matmul_bias_relu_kernel( | |||
| 12 | bias: pypto.Tensor([], pypto.DT_FP32), | 12 | bias: pypto.Tensor([], pypto.DT_FP32), |
| 13 | out: pypto.Tensor([], pypto.DT_FP32)): | 13 | out: pypto.Tensor([], pypto.DT_FP32)): |
| 14 | extend_params = {"bias_tensor": bias} | 14 | extend_params = {"bias_tensor": bias} |
| 15 | - # Cube Tile Shape(矩阵乘法使用) | ||
| 16 | pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64]) | 15 | pypto.set_cube_tile_shapes([32, 32], [64, 64], [64, 64]) |
| 17 | - # Vector Tile Shape(Maximum/ReLU 使用)—— 修正接口名和参数 | ||
| 18 | pypto.set_vec_tile_shapes(8, 8) | 16 | pypto.set_vec_tile_shapes(8, 8) |
| 19 | - # 第一步:matmul + bias | ||
| 20 | result = pypto.matmul(a, b, pypto.DT_FP32, extend_params=extend_params) | 17 | result = pypto.matmul(a, b, pypto.DT_FP32, extend_params=extend_params) |
| 21 | - # 第二步:ReLU(maximum(result, 0)) | 18 | + out.move(pypto.maximum(result, 0)) |
| 22 | - out.move(pypto.maximum(result, 0)) | ||
| 23 | - | ||
| 24 | - | ||
| 25 | -def test(): | ||
| 26 | - a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32, device=device) | ||
| 27 | - b = torch.tensor([[5, 6], [7, 8]], dtype=torch.float32, device=device) | ||
| 28 | - bias = torch.tensor([[1, 2]], dtype=torch.float32, device=device) | ||
| 29 | - out = torch.empty((2, 2), dtype=torch.float32, device=device) | ||
| 30 | - matmul_bias_relu_kernel(a, b, bias, out) | ||
| 31 | - ref = torch.maximum(torch.matmul(a, b) + bias, torch.tensor(0.0)) | ||
| 32 | - max_diff = (out - ref).abs().max().item() | ||
| 33 | - torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) | ||
| 34 | - print("matmul_bias_relu_kernel 验证通过") | ||
| 35 | - print("device:", device, "run_mode:", RUN_MODE) | ||
| 36 | - print("a shape:", tuple(a.shape), "b shape:", tuple(b.shape), "bias shape:", tuple(bias.shape)) | ||
| 37 | - print("输出:", out.cpu()) | ||
| 38 | - print("参考:", ref.cpu()) | ||
| 39 | - print("最大误差:", max_diff) | ||
| 40 | - | ||
| 41 | -test() | ||
| @@ -1,13 +1,3 @@ | |||
| 1 | -# 03.06 章节实践参考答案:实现稳定版行 Softmax | ||
| 2 | - | ||
| 3 | -# 章节自测答案: | ||
| 4 | -# 1. 逐元素、规约等向量 | ||
| 5 | -# 2. 矩阵乘法等 Cube | ||
| 6 | -# 3. B | ||
| 7 | -# 4. A | ||
| 8 | -# 5. A | ||
| 9 | -# 6. 把计算表达式的最终结果写回调用者传入的输出 Tensor | ||
| 10 | - | ||
| 11 | 1 | ||
| 12 | def row_softmax_practice_kernel( | 2 | def row_softmax_practice_kernel( |
| 13 | x: pypto.Tensor([], pypto.DT_FP32), | 3 | x: pypto.Tensor([], pypto.DT_FP32), |
| @@ -21,20 +11,4 @@ def row_softmax_practice_kernel( | |||
| 21 | out.move(exp / esum) | 11 | out.move(exp / esum) |
| 22 | 12 | ||
| 23 | 13 | ||
| 24 | -def main_row_softmax_practice(): | ||
| 25 | - x = torch.randn((8, 8), dtype=torch.float32, device=device) | ||
| 26 | - out = torch.empty_like(x) | ||
| 27 | 14 | ||
| 28 | - row_softmax_practice_kernel(x, out) | ||
| 29 | - | ||
| 30 | - ref = torch.softmax(x, dim=-1) | ||
| 31 | - max_diff = (out - ref).abs().max().item() | ||
| 32 | - torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3) | ||
| 33 | - | ||
| 34 | - print("row_softmax_practice_kernel 验证通过") | ||
| 35 | - print("输入 shape:", tuple(x.shape), "输出 shape:", tuple(out.shape)) | ||
| 36 | - print("输出每行求和:", out.sum(dim=-1).detach().cpu()) | ||
| 37 | - print("最大误差:", max_diff) | ||
| 38 | - | ||
| 39 | - | ||
| 40 | -main_row_softmax_practice() | ||
| @@ -5,11 +5,11 @@ | |||
| 5 | "id": "intro-title", | 5 | "id": "intro-title", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 4.1 从基础算子走向模型模块\n", | 8 | + "# 从基础算子走向模型模块\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "学完基础算子后,下一步不是记住更多 API,而是学会把已有能力组合起来。真实模型里的一个模块,往往同时包含逐元素计算、归约、矩阵乘、shape 变换、动态维度处理和多 kernel 调用。\n", | 10 | "学完基础算子后,下一步不是记住更多 API,而是学会把已有能力组合起来。真实模型里的一个模块,往往同时包含逐元素计算、归约、矩阵乘、shape 变换、动态维度处理和多 kernel 调用。\n", |
| 11 | "\n", | 11 | "\n", |
| 12 | - "这一节先把学习方法和运行环境准备好。后面会从激活函数、Softmax、LayerNorm、RMSNorm、FFN 一直走到 Attention、Transformer block、Cost Model、ACLGraph 和章节综合实践。每一节都会围绕同一个闭环展开:先明确要计算什么,再拆 shape 和公式,然后写 PyPTO kernel,最后用 PyTorch reference 对照验证。" | 12 | + "这一节先把学习方法和运行环境准备好。后面会从激活函数、Softmax、LayerNorm、RMSNorm、FFN 一直走到 Attention、Transformer block和章节综合实践。每一节都会围绕同一个闭环展开:先明确要计算什么,再拆 shape 和公式,然后写 PyPTO kernel,最后用 PyTorch reference 对照验证。" |
| 13 | ] | 13 | ] |
| 14 | }, | 14 | }, |
| 15 | { | 15 | { |
| @@ -23,19 +23,18 @@ | |||
| 23 | "\n", | 23 | "\n", |
| 24 | "| 章节 | 主要内容 | 你会重点练到什么 |\n", | 24 | "| 章节 | 主要内容 | 你会重点练到什么 |\n", |
| 25 | "| --- | --- | --- |\n", | 25 | "| --- | --- | --- |\n", |
| 26 | + "| 4.1 | 章节介绍、学习路线与运行准备 | 准备好运行环境,并且了解本章内容的学习方法 |\n", | ||
| 26 | "| 4.2 | 自定义激活函数与 Softmax | SiLU、GELU、SwiGLU、GeGLU、稳定 Softmax,以及逐元素计算和归约计算的组合方式。 |\n", | 27 | "| 4.2 | 自定义激活函数与 Softmax | SiLU、GELU、SwiGLU、GeGLU、稳定 Softmax,以及逐元素计算和归约计算的组合方式。 |\n", |
| 27 | "| 4.3 | LayerNorm、RMSNorm 与 FFN | 均值/方差归一化、RMS 归一化、前馈网络里的 matmul + activation + matmul 数据流。 |\n", | 28 | "| 4.3 | LayerNorm、RMSNorm 与 FFN | 均值/方差归一化、RMS 归一化、前馈网络里的 matmul + activation + matmul 数据流。 |\n", |
| 28 | "| 4.4 | 动态 Shape 与控制流 | `pypto.DYNAMIC`、`view`、`valid_shape`、`assemble`、`pypto.loop` 和条件分支。 |\n", | 29 | "| 4.4 | 动态 Shape 与控制流 | `pypto.DYNAMIC`、`view`、`valid_shape`、`assemble`、`pypto.loop` 和条件分支。 |\n", |
| 29 | "| 4.5 | Attention 与 Transformer 组合 | `Q @ K^T -> softmax -> @ V`、Q/K/V 投影、多头拆分/合并、残差连接和多 kernel 组合。 |\n", | 30 | "| 4.5 | Attention 与 Transformer 组合 | `Q @ K^T -> softmax -> @ V`、Q/K/V 投影、多头拆分/合并、残差连接和多 kernel 组合。 |\n", |
| 30 | - "| 4.6 | 系统分析与加速 | Cost Model 输出分析、ACLGraph 图捕获与 replay,理解算子如何进入系统优化链路。 |\n", | 31 | + "| 4.6 | 章节实践 | 综合使用激活函数、归一化、动态 shape、loop、view 和 assemble,完成可验证的融合算子。 |\n", |
4.1应该也需要补充进去吧 ![]() ![]() | |||
| 31 | - "| 4.7 | 章节实践 | 综合使用激活函数、归一化、动态 shape、loop、view 和 assemble,完成可验证的融合算子。 |\n", | ||
| 32 | "\n", | 32 | "\n", |
| 33 | "进入中高级实践后,代码变长的原因通常不是某个 API 变复杂,而是一个模块里同时出现了多种基础能力。阅读时可以先抓住这些复用关系:\n", | 33 | "进入中高级实践后,代码变长的原因通常不是某个 API 变复杂,而是一个模块里同时出现了多种基础能力。阅读时可以先抓住这些复用关系:\n", |
| 34 | "\n", | 34 | "\n", |
| 35 | "- Softmax 既可以单独实现,也会成为 Attention 的核心步骤。\n", | 35 | "- Softmax 既可以单独实现,也会成为 Attention 的核心步骤。\n", |
| 36 | "- dynamic shape 既能单独练习,也会被动态 Attention 和动态 FFN 使用。\n", | 36 | "- dynamic shape 既能单独练习,也会被动态 Attention 和动态 FFN 使用。\n", |
| 37 | "- LayerNorm、GELU、Residual Connection 单独看是小模块,组合起来就是 Transformer block。\n", | 37 | "- LayerNorm、GELU、Residual Connection 单独看是小模块,组合起来就是 Transformer block。\n", |
| 38 | - "- Cost Model 和 ACLGraph 不改变数学结果,但会影响分析、集成和执行效率。\n", | ||
| 39 | "\n", | 38 | "\n", |
| 40 | "也就是说,后面的内容不是互不相关的新知识点,而是在不断复用和放大前面学过的模式:先组合算子,再组合模块,最后理解系统分析和图捕获,并通过章节实践检查自己能否把多类能力组合到同一个 kernel 中。" | 39 | "也就是说,后面的内容不是互不相关的新知识点,而是在不断复用和放大前面学过的模式:先组合算子,再组合模块,最后理解系统分析和图捕获,并通过章节实践检查自己能否把多类能力组合到同一个 kernel 中。" |
| 41 | ] | 40 | ] |
| @@ -47,7 +46,7 @@ | |||
| 47 | "source": [ | 46 | "source": [ |
| 48 | "## 2. 先确认运行环境\n", | 47 | "## 2. 先确认运行环境\n", |
| 49 | "\n", | 48 | "\n", |
| 50 | - "先运行下面这个单元。它不会实现具体算子,只负责把 Notebook 调到一个稳定状态:导入依赖、清理 PyPTO 记录状态、选择 CPU / NPU 设备,并设置 `RUN_MODE`。" | 49 | + "先运行下面这个单元。它不会实现具体算子,只负责把 Notebook 调到一个稳定状态:导入依赖、选择NPU 设备,并设置 `RUN_MODE`。" |
| 51 | ] | 50 | ] |
| 52 | }, | 51 | }, |
| 53 | { | 52 | { |
| @@ -59,6 +58,7 @@ | |||
| 59 | "source": [ | 58 | "source": [ |
| 60 | "import os\n", | 59 | "import os\n", |
| 61 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 60 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 61 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 62 | "import torch\n", | 62 | "import torch\n", |
| 63 | "import pypto\n", | 63 | "import pypto\n", |
| 64 | "import torch_npu\n", | 64 | "import torch_npu\n", |
| @@ -91,8 +91,8 @@ | |||
| 91 | "| `import torch` | Host 侧 | 创建真实输入、输出和 PyTorch reference |\n", | 91 | "| `import torch` | Host 侧 | 创建真实输入、输出和 PyTorch reference |\n", |
| 92 | "| `import pypto` | PyPTO 编译入口 | 定义 JIT kernel、Tensor 描述和算子表达式 |\n", | 92 | "| `import pypto` | PyPTO 编译入口 | 定义 JIT kernel、Tensor 描述和算子表达式 |\n", |
| 93 | "| `import torch_npu` | NPU 运行环境 | 当前环境支持 NPU 时启用真实设备执行 |\n", | 93 | "| `import torch_npu` | NPU 运行环境 | 当前环境支持 NPU 时启用真实设备执行 |\n", |
| 94 | - "| `get_device()` | 设备选择 | 有 NPU 时返回 `npu:设备号`,否则返回 `cpu` |\n", | 94 | + "| `get_device()` | 设备选择 | 返回 `npu:设备号` |\n", |
| 95 | - "| `RUN_MODE` | JIT 配置 | 有 NPU 时用 `pypto.RunMode.NPU`,否则用 `pypto.RunMode.SIM` |\n", | 95 | + "| `RUN_MODE` | JIT 配置 | 标识使用NPU |\n", |
| 96 | "\n", | 96 | "\n", |
| 97 | "其中最容易混淆的是 `torch.Tensor` 和 `pypto.Tensor`:\n", | 97 | "其中最容易混淆的是 `torch.Tensor` 和 `pypto.Tensor`:\n", |
| 98 | "\n", | 98 | "\n", |
| @@ -118,7 +118,7 @@ | |||
| 118 | "5. 看数据组织:是否使用 `view`、`assemble`、`loop`、`valid_shape`。\n", | 118 | "5. 看数据组织:是否使用 `view`、`assemble`、`loop`、`valid_shape`。\n", |
| 119 | "6. 看验证闭环:PyTorch reference 是怎么写的,误差阈值怎么判断。\n", | 119 | "6. 看验证闭环:PyTorch reference 是怎么写的,误差阈值怎么判断。\n", |
| 120 | "\n", | 120 | "\n", |
| 121 | - "这样读,Attention、FFN、动态 shape 和 ACLGraph 都不会显得杂乱。它们只是把前面学过的基础部件放进了更大的结构里。" | 121 | + "这样读,Attention、FFN、动态 shape 都不会显得杂乱。它们只是把前面学过的基础部件放进了更大的结构里。" |
| 122 | ] | 122 | ] |
| 123 | }, | 123 | }, |
| 124 | { | 124 | { |
| @@ -212,14 +212,13 @@ | |||
| 212 | "source": [ | 212 | "source": [ |
| 213 | "## 5. 学完后你应该能够\n", | 213 | "## 5. 学完后你应该能够\n", |
| 214 | "\n", | 214 | "\n", |
| 215 | - "完成 4.2 到 4.7 后,可以回到这里检查自己是否已经做到:\n", | 215 | + "完成学习后,可以回到这里检查自己是否已经做到:\n", |
| 216 | "\n", | 216 | "\n", |
| 217 | "1. 说清楚激活函数、Softmax、归一化、FFN、Attention 分别解决什么问题。\n", | 217 | "1. 说清楚激活函数、Softmax、归一化、FFN、Attention 分别解决什么问题。\n", |
| 218 | "2. 看懂 `pypto.DYNAMIC`、`pypto.loop`、`valid_shape` 和条件分支在 kernel 里的作用。\n", | 218 | "2. 看懂 `pypto.DYNAMIC`、`pypto.loop`、`valid_shape` 和条件分支在 kernel 里的作用。\n", |
| 219 | "3. 理解 `view`、`assemble`、`transpose`、`reshape` 如何组织数据流。\n", | 219 | "3. 理解 `view`、`assemble`、`transpose`、`reshape` 如何组织数据流。\n", |
| 220 | "4. 把多个基础算子组合成可验证的中级和高级网络片段。\n", | 220 | "4. 把多个基础算子组合成可验证的中级和高级网络片段。\n", |
| 221 | "5. 能够用 PyTorch reference 对照验证 PyPTO 实现。\n", | 221 | "5. 能够用 PyTorch reference 对照验证 PyPTO 实现。\n", |
| 222 | - "6. 理解 Cost Model 和 ACLGraph 这类系统能力与数学算子的区别。\n", | ||
| 223 | "7. 能把激活、归一化、动态 shape 和分块写回组合成一个章节综合实践算子。" | 222 | "7. 能把激活、归一化、动态 shape 和分块写回组合成一个章节综合实践算子。" |
| 224 | ] | 223 | ] |
| 225 | }, | 224 | }, |
| @@ -230,7 +229,7 @@ | |||
| 230 | "source": [ | 229 | "source": [ |
| 231 | "## 6. 建议学习顺序\n", | 230 | "## 6. 建议学习顺序\n", |
| 232 | "\n", | 231 | "\n", |
| 233 | - "建议仍然按 4.2 到 4.7 顺序阅读:先会写组合算子,再会写归一化和 FFN 这类模块,然后处理动态 shape 和控制流,接着进入 Attention / Transformer 组合,最后理解系统分析与图捕获,并通过章节实践把这些能力串成一个完整任务。这个顺序的好处是每一节都会复用上一节的能力,复杂度是逐步叠上去的。\n" | 232 | + "先会写组合算子,再会写归一化和 FFN 这类模块,然后处理动态 shape 和控制流,接着进入 Attention / Transformer 组合。这个顺序的好处是每一节都会复用上一节的能力,复杂度是逐步叠上去的。\n" |
| 234 | ] | 233 | ] |
| 235 | }, | 234 | }, |
| 236 | { | 235 | { |
| @@ -246,7 +245,6 @@ | |||
| 246 | "| LayerNorm、RMSNorm、FFN | 如何把 reduction、matmul 和激活函数组织成神经网络子模块 | 为 Transformer block 的前馈分支做准备 |\n", | 245 | "| LayerNorm、RMSNorm、FFN | 如何把 reduction、matmul 和激活函数组织成神经网络子模块 | 为 Transformer block 的前馈分支做准备 |\n", |
| 247 | "| dynamic shape、loop、condition | 如何在输入维度变化时稳定地切块、计算和写回 | 为动态 batch、动态 attention 和系统集成做准备 |\n", | 246 | "| dynamic shape、loop、condition | 如何在输入维度变化时稳定地切块、计算和写回 | 为动态 batch、动态 attention 和系统集成做准备 |\n", |
| 248 | "| Attention 与 Transformer 组合 | 如何组织 Q/K/V、多头拆分、残差连接和多 kernel 调用 | 把前面的小模块组合成完整模型片段 |\n", | 247 | "| Attention 与 Transformer 组合 | 如何组织 Q/K/V、多头拆分、残差连接和多 kernel 调用 | 把前面的小模块组合成完整模型片段 |\n", |
| 249 | - "| Cost Model 与 ACLGraph | 如何观察执行成本,并捕获可重放的执行图 | 从算子正确性走向性能分析和工程集成 |\n", | ||
| 250 | "| 章节实践 | 如何把动态 shape、归一化、激活和输出写回组合成完整算子 | 检查第四章知识是否能独立迁移到综合任务 |" | 248 | "| 章节实践 | 如何把动态 shape、归一化、激活和输出写回组合成完整算子 | 检查第四章知识是否能独立迁移到综合任务 |" |
| 251 | ] | 249 | ] |
| 252 | }, | 250 | }, |
| @@ -265,9 +263,7 @@ | |||
| 265 | "| View / Assemble | 从大 Tensor 取局部块,再写回大 Tensor | 4.4、4.5 |\n", | 263 | "| View / Assemble | 从大 Tensor 取局部块,再写回大 Tensor | 4.4、4.5 |\n", |
| 266 | "| Multi-head | 把 hidden 拆成多个 head 并行计算 | 4.5 |\n", | 264 | "| Multi-head | 把 hidden 拆成多个 head 并行计算 | 4.5 |\n", |
| 267 | "| Residual Connection | 把输入直接加回输出,形成残差路径 | 4.5 |\n", | 265 | "| Residual Connection | 把输入直接加回输出,形成残差路径 | 4.5 |\n", |
| 268 | - "| Cost Model | 模拟和分析执行成本 | 4.6 |\n", | 266 | + "| Chapter Practice | 用一个综合任务检查本章能力 | 4.6 |" |
| 269 | - "| ACLGraph | 捕获执行图并重放,减少重复调度开销 | 4.6 |\n", | ||
| 270 | - "| Chapter Practice | 用一个综合任务检查本章能力 | 4.7 |" | ||
| 271 | ] | 267 | ] |
| 272 | }, | 268 | }, |
| 273 | { | 269 | { |
| @@ -279,21 +275,13 @@ | |||
| 279 | "\n", | 275 | "\n", |
| 280 | "本节练习用于检查中高级实践的整体阅读方法。请结合本节的最小 `add + relu` kernel 和后续章节路线完成以下题目。\n", | 276 | "本节练习用于检查中高级实践的整体阅读方法。请结合本节的最小 `add + relu` kernel 和后续章节路线完成以下题目。\n", |
| 281 | "\n", | 277 | "\n", |
| 282 | - "1. (选择题)为什么学习时要先理解能力之间的复用关系? \n", | 278 | + "1. (填空题)`torch.Tensor` 处在________侧,`pypto.Tensor` 是 kernel 的________。\n", |
| 283 | - " A. 因为中高级能力会在 Softmax、dynamic shape、Attention 等场景中反复组合出现 \n", | 279 | + "2. (选择题)阅读一个中高级 kernel 时,为什么要先看 shape? \n", |
| 284 | - " B. 因为每个主题都完全独立 \n", | ||
| 285 | - " C. 因为可以不看 shape 直接写代码 \n", | ||
| 286 | - " D. 因为 Cost Model 会改变数学公式\n", | ||
| 287 | - "2. (填空题)`torch.Tensor` 处在________侧,`pypto.Tensor` 是 kernel 的________。\n", | ||
| 288 | - "3. (选择题)阅读一个中高级 kernel 时,为什么要先看 shape? \n", | ||
| 289 | " A. shape 决定 matmul、transpose、reshape、reduction 和 broadcast 是否能对齐 \n", | 280 | " A. shape 决定 matmul、transpose、reshape、reduction 和 broadcast 是否能对齐 \n", |
| 290 | " B. shape 可以替代所有验证逻辑 \n", | 281 | " B. shape 可以替代所有验证逻辑 \n", |
| 291 | " C. shape 会自动修复错误 dtype \n", | 282 | " C. shape 会自动修复错误 dtype \n", |
| 292 | " D. shape 与中高级 kernel 无关\n", | 283 | " D. shape 与中高级 kernel 无关\n", |
| 293 | - "4. (填空题)`view -> compute -> assemble` 这个模式通常用于解决________问题。\n", | 284 | + "3. (填空题)`view -> compute -> assemble` 这个模式通常用于解决________问题。\n", |
| 294 | - "5. (选择题)Cost Model 和 ACLGraph 会改变 Softmax 的数学含义吗? \n", | ||
| 295 | - " A. 会 \n", | ||
| 296 | - " B. 不会\n", | ||
| 297 | "\n", | 285 | "\n", |
| 298 | "**执行以下代码获取答案。**\n" | 286 | "**执行以下代码获取答案。**\n" |
| 299 | ] | 287 | ] |
| @@ -317,7 +305,7 @@ | |||
| 317 | "\n", | 305 | "\n", |
| 318 | "这一节先完成三件事:确认运行环境,建立阅读中高级 kernel 的顺序,并用一个最小 `add + relu` kernel 回顾完整验证闭环。\n", | 306 | "这一节先完成三件事:确认运行环境,建立阅读中高级 kernel 的顺序,并用一个最小 `add + relu` kernel 回顾完整验证闭环。\n", |
| 319 | "\n", | 307 | "\n", |
| 320 | - "接下来进入具体代码:先学习如何把基础算子组合成 SiLU、GELU、SwiGLU、GeGLU 和稳定 Softmax,最后在 4.7 中完成章节综合实践。" | 308 | + "接下来进入具体代码:先学习如何把基础算子组合成 SiLU、GELU、SwiGLU、GeGLU 和稳定 Softmax,最后完成章节综合实践。" |
| 321 | ] | 309 | ] |
| 322 | } | 310 | } |
| 323 | ], | 311 | ], |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "title", | 5 | "id": "title", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 4.2 算子组合与 Softmax\n", | 8 | + "# 算子组合与 Softmax\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "这一节开始练习“把公式翻译成 kernel”。你会先实现几类常见激活函数,再实现稳定 Softmax。它们看起来都不长,但能覆盖中高级算子里非常常见的组合方式:elementwise、reduction、广播、动态 batch 和 loop 分块。\n", | 10 | "这一节开始练习“把公式翻译成 kernel”。你会先实现几类常见激活函数,再实现稳定 Softmax。它们看起来都不长,但能覆盖中高级算子里非常常见的组合方式:elementwise、reduction、广播、动态 batch 和 loop 分块。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -72,9 +72,8 @@ | |||
| 72 | "id": "env-title", | 72 | "id": "env-title", |
| 73 | "metadata": {}, | 73 | "metadata": {}, |
| 74 | "source": [ | 74 | "source": [ |
| 75 | - "## 3. 先准备环境与通用工具\n", | 75 | + "## 3. 环境准备\n", |
| 76 | - "\n", | 76 | + "\n" |
| 77 | - "先运行下面这个单元。它会完成四件事:导入依赖、选择设备、清理 PyPTO Notebook 状态、定义 PyTorch reference。后面的 kernel 和测试函数都会复用这些变量。" | ||
| 78 | ] | 77 | ] |
| 79 | }, | 78 | }, |
| 80 | { | 79 | { |
| @@ -86,6 +85,7 @@ | |||
| 86 | "source": [ | 85 | "source": [ |
| 87 | "import os\n", | 86 | "import os\n", |
| 88 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 87 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 88 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 89 | "import torch\n", | 89 | "import torch\n", |
| 90 | "import pypto\n", | 90 | "import pypto\n", |
| 91 | "import numpy as np\n", | 91 | "import numpy as np\n", |
| @@ -144,7 +144,7 @@ | |||
| 144 | "\n", | 144 | "\n", |
| 145 | "| 函数或变量 | 作用 | 阅读时关注点 |\n", | 145 | "| 函数或变量 | 作用 | 阅读时关注点 |\n", |
| 146 | "| --- | --- | --- |\n", | 146 | "| --- | --- | --- |\n", |
| 147 | - "| `get_device()` | 有 `torch_npu` 时选择 NPU,否则使用 CPU/SIM 阅读模式 | `torch.Tensor` 的真实 device 由这里决定。 |\n", | 147 | + "| `get_device()` | 选择 NPU | `torch.Tensor` 的真实 device 由这里决定。 |\n", |
| 148 | "| `RUN_MODE` | 传给 `@pypto.frontend.jit`,决定 kernel 运行模式 | NPU 环境为 `NPU`,无 NPU 时为 `SIM`。 |\n", | 148 | "| `RUN_MODE` | 传给 `@pypto.frontend.jit`,决定 kernel 运行模式 | NPU 环境为 `NPU`,无 NPU 时为 `SIM`。 |\n", |
| 149 | "| `configure_tiling(x)` | 根据输入维度设置 vec tile,服务 elementwise 激活函数 | 只影响执行组织,不改变公式。 |\n", | 149 | "| `configure_tiling(x)` | 根据输入维度设置 vec tile,服务 elementwise 激活函数 | 只影响执行组织,不改变公式。 |\n", |
| 150 | "| `*_golden` | PyTorch reference,用来验证 PyPTO kernel 输出 | reference 的公式和 dtype 要与 kernel 对齐。 |\n", | 150 | "| `*_golden` | PyTorch reference,用来验证 PyPTO kernel 输出 | reference 的公式和 dtype 要与 kernel 对齐。 |\n", |
| @@ -168,7 +168,7 @@ | |||
| 168 | "| SiLU | `x * sigmoid(x)` | `[32, 128]` | `[32, 128]` |\n", | 168 | "| SiLU | `x * sigmoid(x)` | `[32, 128]` | `[32, 128]` |\n", |
| 169 | "| GELU 近似 | `x * sigmoid(1.702 * x)` | `[32, 128]` | `[32, 128]` |\n", | 169 | "| GELU 近似 | `x * sigmoid(1.702 * x)` | `[32, 128]` | `[32, 128]` |\n", |
| 170 | "\n", | 170 | "\n", |
| 171 | - "这两个算子都是 elementwise:每个位置独立计算,不需要跨行、跨列归约。" | 171 | + "这两个算子都是 elementwise:每个位置独立计算,不需要跨行、跨列归约。尤其注意这里的gelu是近似实现。" |
| 172 | ] | 172 | ] |
| 173 | }, | 173 | }, |
| 174 | { | 174 | { |
| @@ -218,11 +218,9 @@ | |||
| 218 | "\n", | 218 | "\n", |
| 219 | "```text\n", | 219 | "```text\n", |
| 220 | "x: [32, 128], BF16\n", | 220 | "x: [32, 128], BF16\n", |
| 221 | - "x_scaled: [32, 128], BF16/内部计算类型\n", | 221 | + "x_scaled: [32, 128], BF16\n", |
| 222 | "out: [32, 128], BF16\n", | 222 | "out: [32, 128], BF16\n", |
| 223 | - "```\n", | 223 | + "```" |
| 224 | - "\n", | ||
| 225 | - "这里特意写成 `x * 1.702`,而不是 `1.702 * x`。原因是部分 PyPTO 版本对“左操作数是 Python float、右操作数是 Tensor”的重载支持不完整,可能触发 `TypeError`;把 Tensor 放在左边更稳妥。这个细节很小,但写组合算子时经常会遇到。" | ||
| 226 | ] | 224 | ] |
| 227 | }, | 225 | }, |
| 228 | { | 226 | { |
| @@ -266,6 +264,7 @@ | |||
| 266 | " print(f\"GELU max difference vs approximation: {max_diff_approx:.6f}\")\n", | 264 | " print(f\"GELU max difference vs approximation: {max_diff_approx:.6f}\")\n", |
| 267 | " assert max_diff_full < 1e-1, \"GELU result mismatch!\"\n", | 265 | " assert max_diff_full < 1e-1, \"GELU result mismatch!\"\n", |
| 268 | "\n", | 266 | "\n", |
| 267 | + "\n", | ||
| 269 | "test_silu()\n", | 268 | "test_silu()\n", |
| 270 | "test_gelu()\n" | 269 | "test_gelu()\n" |
| 271 | ] | 270 | ] |
| @@ -416,24 +415,23 @@ | |||
| 416 | "\n", | 415 | "\n", |
| 417 | "\n", | 416 | "\n", |
| 418 | "def test_geglu() -> None:\n", | 417 | "def test_geglu() -> None:\n", |
| 419 | - " device_local = get_device()\n", | 418 | + " device = get_device()\n", |
| 420 | " shape = (32, 128)\n", | 419 | " shape = (32, 128)\n", |
| 421 | - " gate_torch = torch.randn(shape, dtype=torch.bfloat16, device=device_local)\n", | 420 | + " gate_torch = torch.randn(shape, dtype=torch.bfloat16, device=device)\n", |
| 422 | - " up_torch = torch.randn(shape, dtype=torch.bfloat16, device=device_local)\n", | 421 | + " up_torch = torch.randn(shape, dtype=torch.bfloat16, device=device)\n", |
| 423 | - " out_torch = torch.empty(shape, dtype=torch.bfloat16, device=device_local)\n", | 422 | + " out_torch = torch.empty(shape, dtype=torch.bfloat16, device=device)\n", |
| 424 | - "\n", | 423 | + " # Execute\n", |
| 425 | " geglu_activation_kernel(gate_torch, up_torch, out_torch)\n", | 424 | " geglu_activation_kernel(gate_torch, up_torch, out_torch)\n", |
| 426 | - " expected_full = geglu_golden(gate_torch, up_torch)\n", | ||
| 427 | - " expected_approx = gelu_approx_golden(gate_torch) * up_torch\n", | ||
| 428 | - " max_diff_full = (out_torch - expected_full).abs().max().item()\n", | ||
| 429 | - " max_diff_approx = (out_torch - expected_approx).abs().max().item()\n", | ||
| 430 | "\n", | 425 | "\n", |
| 431 | - " print(f\"GeGLU gate shape: {gate_torch.shape}\")\n", | 426 | + " # Verify\n", |
| 432 | - " print(f\"GeGLU up shape: {up_torch.shape}\")\n", | 427 | + " expected = geglu_golden(gate_torch, up_torch)\n", |
| 433 | - " print(f\"GeGLU output shape: {out_torch.shape}\")\n", | 428 | + " max_diff = (out_torch - expected).abs().max().item()\n", |
| 434 | - " print(f\"GeGLU max difference vs torch.gelu gate: {max_diff_full:.6f}\")\n", | 429 | + "\n", |
| 435 | - " print(f\"GeGLU max difference vs approximation: {max_diff_approx:.6f}\")\n", | 430 | + " print(f\"Gate shape: {gate_torch.shape}\")\n", |
| 436 | - " assert max_diff_full < 1e-1, \"GeGLU result mismatch!\"\n", | 431 | + " print(f\"Up shape: {up_torch.shape}\")\n", |
| 432 | + " print(f\"Output shape: {out_torch.shape}\")\n", | ||
| 433 | + " print(f\"Max difference: {max_diff:.6f}\")\n", | ||
| 434 | + " assert max_diff < 1e-1, \"Result mismatch!\"\n", | ||
| 437 | "\n", | 435 | "\n", |
| 438 | "\n", | 436 | "\n", |
| 439 | "test_swiglu()\n", | 437 | "test_swiglu()\n", |
| @@ -725,20 +723,24 @@ | |||
| 725 | "本节练习用于复盘激活函数组合、稳定 Softmax 和动态 batch 分块写回。题型包含选择题和填空题。\n", | 723 | "本节练习用于复盘激活函数组合、稳定 Softmax 和动态 batch 分块写回。题型包含选择题和填空题。\n", |
| 726 | "\n", | 724 | "\n", |
| 727 | "1. (填空题)SiLU 的公式是________;SwiGLU 会把 SiLU 用在________分支上,再乘以 `up`。\n", | 725 | "1. (填空题)SiLU 的公式是________;SwiGLU 会把 SiLU 用在________分支上,再乘以 `up`。\n", |
| 726 | + "\n", | ||
| 727 | + "\n", | ||
| 728 | "2. (填空题)GeGLU 会把 GELU 或 GELU 近似用在________分支上,再乘以 `up`。\n", | 728 | "2. (填空题)GeGLU 会把 GELU 或 GELU 近似用在________分支上,再乘以 `up`。\n", |
| 729 | + "\n", | ||
| 730 | + "\n", | ||
| 729 | "3. (填空题)门控激活需要 `gate` 和 `up` 两路输入,其中 `gate` 负责________,`up` 负责________。\n", | 731 | "3. (填空题)门控激活需要 `gate` 和 `up` 两路输入,其中 `gate` 负责________,`up` 负责________。\n", |
| 730 | - "4. (选择题)为什么 `x * 1.702` 比 `1.702 * x` 更推荐? \n", | 732 | + "\n", |
| 731 | - " A. Tensor 放左边对部分 PyPTO 版本的运算重载更稳妥 \n", | 733 | + "4. (选择题)Softmax 中 `row_max` 的 shape 为什么是 `[B, S, H, 1]`? \n", |
| 732 | - " B. 二者数学含义完全不同 \n", | ||
| 733 | - " C. 这样可以跳过 sigmoid 计算 \n", | ||
| 734 | - " D. 这样会自动变成 Softmax\n", | ||
| 735 | - "5. (选择题)Softmax 中 `row_max` 的 shape 为什么是 `[B, S, H, 1]`? \n", | ||
| 736 | " A. 因为沿最后一维规约并使用 `keepdim=True` \n", | 734 | " A. 因为沿最后一维规约并使用 `keepdim=True` \n", |
| 737 | " B. 因为 batch 维被删除 \n", | 735 | " B. 因为 batch 维被删除 \n", |
| 738 | " C. 因为最后一维必须固定为 1 个 head \n", | 736 | " C. 因为最后一维必须固定为 1 个 head \n", |
| 739 | " D. 因为输出 dtype 改成了 INT32\n", | 737 | " D. 因为输出 dtype 改成了 INT32\n", |
| 740 | - "6. (填空题)`pypto.loop` 在 Softmax kernel 中主要用于________。\n", | 738 | + "\n", |
| 741 | - "7. (选择题)如果把 `output_tensor[b_offset:b_offset_end, ...]` 写成过大的切片,可能出现什么问题? \n", | 739 | + "\n", |
| 740 | + "5. (填空题)`pypto.loop` 在 Softmax kernel 中主要用于________。\n", | ||
| 741 | + "\n", | ||
| 742 | + "\n", | ||
| 743 | + "6. (选择题)如果把 `output_tensor[b_offset:b_offset_end, ...]` 写成过大的切片,可能出现什么问题? \n", | ||
| 742 | " A. 覆盖不属于当前 tile 的输出区域 \n", | 744 | " A. 覆盖不属于当前 tile 的输出区域 \n", |
| 743 | " B. 自动提升数值精度 \n", | 745 | " B. 自动提升数值精度 \n", |
| 744 | " C. 自动修复动态 shape \n", | 746 | " C. 自动修复动态 shape \n", |
代码比较长的时候最好还是在代码中增加注释,后面几节也是


| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "title", | 5 | "id": "title", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 4.3 归一化与 FFN\n", | 8 | + "# 归一化与 FFN\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "这一节开始把前面学过的 reduction、广播、matmul 和激活函数组合成更接近模型结构的模块:**归一化层** 和 **FFN 模块**。\n", | 10 | "这一节开始把前面学过的 reduction、广播、matmul 和激活函数组合成更接近模型结构的模块:**归一化层** 和 **FFN 模块**。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -17,9 +17,7 @@ | |||
| 17 | "id": "env-title", | 17 | "id": "env-title", |
| 18 | "metadata": {}, | 18 | "metadata": {}, |
| 19 | "source": [ | 19 | "source": [ |
| 20 | - "## 1. 环境准备\n", | 20 | + "## 1. 环境准备\n" |
| 21 | - "\n", | ||
| 22 | - "和前面的章节一样,先清理 Notebook 状态,再准备当前运行模式。" | ||
| 23 | ] | 21 | ] |
| 24 | }, | 22 | }, |
| 25 | { | 23 | { |
| @@ -31,6 +29,7 @@ | |||
| 31 | "source": [ | 29 | "source": [ |
| 32 | "import os\n", | 30 | "import os\n", |
| 33 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 31 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 32 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 34 | "import math\n", | 33 | "import math\n", |
| 35 | "from dataclasses import dataclass\n", | 34 | "from dataclasses import dataclass\n", |
| 36 | "from typing import Literal\n", | 35 | "from typing import Literal\n", |
| @@ -185,9 +184,11 @@ | |||
| 185 | " mean = pypto.sum(x, dim=-1, keepdim=True)\n", | 184 | " mean = pypto.sum(x, dim=-1, keepdim=True)\n", |
| 186 | " mean = mean / hidden_size\n", | 185 | " mean = mean / hidden_size\n", |
| 187 | " centered = x - mean\n", | 186 | " centered = x - mean\n", |
| 187 | + " # 方差:平方和 / hidden_size\n", | ||
| 188 | " squared = centered * centered\n", | 188 | " squared = centered * centered\n", |
| 189 | " var = pypto.sum(squared, dim=-1, keepdim=True)\n", | 189 | " var = pypto.sum(squared, dim=-1, keepdim=True)\n", |
| 190 | " var = var / hidden_size\n", | 190 | " var = var / hidden_size\n", |
| 191 | + " # 标准差\n", | ||
| 191 | " std = pypto.sqrt(var + eps)\n", | 192 | " std = pypto.sqrt(var + eps)\n", |
| 192 | " normalized = centered / std\n", | 193 | " normalized = centered / std\n", |
| 193 | " scaled = normalized * gamma\n", | 194 | " scaled = normalized * gamma\n", |
| @@ -217,6 +218,7 @@ | |||
| 217 | " mean_sq = pypto.sum(squared, dim=-1, keepdim=True)\n", | 218 | " mean_sq = pypto.sum(squared, dim=-1, keepdim=True)\n", |
| 218 | " mean_sq = mean_sq / hidden_size\n", | 219 | " mean_sq = mean_sq / hidden_size\n", |
| 219 | " rms = pypto.sqrt(mean_sq + eps)\n", | 220 | " rms = pypto.sqrt(mean_sq + eps)\n", |
| 221 | + " # 归一化并缩放(RMSNorm 无 beta,无减均值步骤)\n", | ||
| 220 | " normalized = x / rms\n", | 222 | " normalized = x / rms\n", |
| 221 | " return normalized * gamma\n", | 223 | " return normalized * gamma\n", |
| 222 | "\n", | 224 | "\n", |
| @@ -388,7 +390,7 @@ | |||
| 388 | " hidden_size: int\n", | 390 | " hidden_size: int\n", |
| 389 | " intermediate_size: int\n", | 391 | " intermediate_size: int\n", |
| 390 | " activation: Literal[\"gelu\", \"swiglu\", \"relu\"] = \"gelu\"\n", | 392 | " activation: Literal[\"gelu\", \"swiglu\", \"relu\"] = \"gelu\"\n", |
| 391 | - " dtype: pypto.DataType = pypto.DT_FP16\n", | 393 | + " dtype: pypto.DataType = pypto.DT_BF16\n", |
| 392 | " use_dynamic_shape: bool = False\n", | 394 | " use_dynamic_shape: bool = False\n", |
| 393 | " vec_tile_shape: tuple = (64, 128)\n", | 395 | " vec_tile_shape: tuple = (64, 128)\n", |
| 394 | " cube_tile_shape: tuple = (64, 128, 128)\n", | 396 | " cube_tile_shape: tuple = (64, 128, 128)\n", |
| @@ -416,12 +418,14 @@ | |||
| 416 | "\n", | 418 | "\n", |
| 417 | "def gelu_activation_core(x: pypto.tensor) -> pypto.tensor:\n", | 419 | "def gelu_activation_core(x: pypto.tensor) -> pypto.tensor:\n", |
| 418 | " pypto.set_vec_tile_shapes(*x.shape[:2] if len(x.shape) >= 2 else (32, 128))\n", | 420 | " pypto.set_vec_tile_shapes(*x.shape[:2] if len(x.shape) >= 2 else (32, 128))\n", |
| 421 | + " # 先转 FP32 计算以保证精度\n", | ||
| 419 | " x_fp32 = pypto.cast(x, pypto.DT_FP32)\n", | 422 | " x_fp32 = pypto.cast(x, pypto.DT_FP32)\n", |
| 420 | " x_scaled = pypto.mul(x_fp32, GELU_COEFF)\n", | 423 | " x_scaled = pypto.mul(x_fp32, GELU_COEFF)\n", |
| 424 | + " # 手工实现 1 / (1 + exp(-x))\n", | ||
| 421 | " x_neg = pypto.mul(x_scaled, F_NEGA_1)\n", | 425 | " x_neg = pypto.mul(x_scaled, F_NEGA_1)\n", |
| 422 | " exp_neg = pypto.exp(x_neg)\n", | 426 | " exp_neg = pypto.exp(x_neg)\n", |
| 423 | " ones = pypto.full(exp_neg.shape, 1.0, exp_neg.dtype, valid_shape=exp_neg.shape)\n", | 427 | " ones = pypto.full(exp_neg.shape, 1.0, exp_neg.dtype, valid_shape=exp_neg.shape)\n", |
| 424 | - " sigmoid = pypto.div(ones, pypto.add(exp_neg, F_1))\n", | 428 | + " sigmoid = pypto.div(ones, pypto.add(exp_neg, ones))\n", |
| 425 | " activated = pypto.cast(pypto.mul(x_fp32, sigmoid), pypto.DT_BF16)\n", | 429 | " activated = pypto.cast(pypto.mul(x_fp32, sigmoid), pypto.DT_BF16)\n", |
| 426 | " return activated\n", | 430 | " return activated\n", |
| 427 | "\n", | 431 | "\n", |
| @@ -430,11 +434,14 @@ | |||
| 430 | " gate_fp32 = pypto.cast(gate, pypto.DT_FP32)\n", | 434 | " gate_fp32 = pypto.cast(gate, pypto.DT_FP32)\n", |
| 431 | " up_fp32 = pypto.cast(up, pypto.DT_FP32)\n", | 435 | " up_fp32 = pypto.cast(up, pypto.DT_FP32)\n", |
| 432 | " pypto.set_vec_tile_shapes(*gate.shape[:2] if len(gate.shape) >= 2 else (32, 128))\n", | 436 | " pypto.set_vec_tile_shapes(*gate.shape[:2] if len(gate.shape) >= 2 else (32, 128))\n", |
| 437 | + " # 手工实现 sigmoid(gate)\n", | ||
| 433 | " gate_neg = pypto.mul(gate_fp32, F_NEGA_1)\n", | 438 | " gate_neg = pypto.mul(gate_fp32, F_NEGA_1)\n", |
| 434 | " exp_neg = pypto.exp(gate_neg)\n", | 439 | " exp_neg = pypto.exp(gate_neg)\n", |
| 435 | " ones = pypto.full(exp_neg.shape, F_1, exp_neg.dtype, valid_shape=exp_neg.shape)\n", | 440 | " ones = pypto.full(exp_neg.shape, F_1, exp_neg.dtype, valid_shape=exp_neg.shape)\n", |
| 436 | " sigmoid = pypto.div(ones, pypto.add(exp_neg, ones))\n", | 441 | " sigmoid = pypto.div(ones, pypto.add(exp_neg, ones))\n", |
| 442 | + " # Swish = gate * sigmoid(gate)\n", | ||
| 437 | " swish = pypto.mul(gate_fp32, sigmoid)\n", | 443 | " swish = pypto.mul(gate_fp32, sigmoid)\n", |
| 444 | + " # 门控输出 = Swish * up,转回 BF16\n", | ||
| 438 | " return pypto.cast(pypto.mul(swish, up_fp32), pypto.DT_BF16)\n", | 445 | " return pypto.cast(pypto.mul(swish, up_fp32), pypto.DT_BF16)\n", |
| 439 | "\n", | 446 | "\n", |
| 440 | "\n", | 447 | "\n", |
| @@ -454,6 +461,7 @@ | |||
| 454 | " hidden_size, intermediate_size = config.hidden_size, config.intermediate_size\n", | 461 | " hidden_size, intermediate_size = config.hidden_size, config.intermediate_size\n", |
| 455 | " basic_batch = config.basic_batch\n", | 462 | " basic_batch = config.basic_batch\n", |
| 456 | " batch_size = hidden_states.shape[0]\n", | 463 | " batch_size = hidden_states.shape[0]\n", |
| 464 | + " # 向上取整,保证尾块也被覆盖\n", | ||
| 457 | " num_iterations = ceil_div(batch_size, basic_batch)\n", | 465 | " num_iterations = ceil_div(batch_size, basic_batch)\n", |
| 458 | " for idx in pypto.loop(0, num_iterations, 1, name=\"LOOP_FFN_BATCH\", idx_name=\"idx\"):\n", | 466 | " for idx in pypto.loop(0, num_iterations, 1, name=\"LOOP_FFN_BATCH\", idx_name=\"idx\"):\n", |
| 459 | " batch_offset = idx * basic_batch\n", | 467 | " batch_offset = idx * basic_batch\n", |
| @@ -466,11 +474,13 @@ | |||
| 466 | " up = pypto.matmul(hidden_chunk, up_proj_weight, config.dtype)\n", | 474 | " up = pypto.matmul(hidden_chunk, up_proj_weight, config.dtype)\n", |
| 467 | " pypto.set_vec_tile_shapes(*config.vec_tile_shape)\n", | 475 | " pypto.set_vec_tile_shapes(*config.vec_tile_shape)\n", |
| 468 | " activated = gelu_activation_core(up)\n", | 476 | " activated = gelu_activation_core(up)\n", |
| 477 | + " # 重新设置 cube tile(因为上一步激活 core 可能改了 tile 配置)\n", | ||
| 469 | " pypto.set_cube_tile_shapes(\n", | 478 | " pypto.set_cube_tile_shapes(\n", |
| 470 | " [config.cube_tile_shape[0], config.cube_tile_shape[0]],\n", | 479 | " [config.cube_tile_shape[0], config.cube_tile_shape[0]],\n", |
| 471 | " [config.cube_tile_shape[1], config.cube_tile_shape[1]],\n", | 480 | " [config.cube_tile_shape[1], config.cube_tile_shape[1]],\n", |
| 472 | " [config.cube_tile_shape[2], config.cube_tile_shape[2]]\n", | 481 | " [config.cube_tile_shape[2], config.cube_tile_shape[2]]\n", |
| 473 | " )\n", | 482 | " )\n", |
| 483 | + " # set_matrix_size 声明 matmul 的 M/K/N 维度信息\n", | ||
| 474 | " pypto.set_matrix_size([basic_batch, intermediate_size, hidden_size])\n", | 484 | " pypto.set_matrix_size([basic_batch, intermediate_size, hidden_size])\n", |
| 475 | " output_chunk = pypto.matmul(activated, down_proj_weight, config.dtype, b_trans=False)\n", | 485 | " output_chunk = pypto.matmul(activated, down_proj_weight, config.dtype, b_trans=False)\n", |
| 476 | " pypto.assemble(output_chunk, [batch_offset, 0], output)\n", | 486 | " pypto.assemble(output_chunk, [batch_offset, 0], output)\n", |
| @@ -494,6 +504,7 @@ | |||
| 494 | " up = pypto.matmul(hidden_states, up_proj_weight, config.dtype)\n", | 504 | " up = pypto.matmul(hidden_states, up_proj_weight, config.dtype)\n", |
| 495 | " activated = gelu_activation_core(up)\n", | 505 | " activated = gelu_activation_core(up)\n", |
| 496 | " elif config.activation == \"swiglu\":\n", | 506 | " elif config.activation == \"swiglu\":\n", |
| 507 | + " # SwiGLU:同时做 gate 和 up 两路投影\n", | ||
| 497 | " gate = pypto.matmul(hidden_states, gate_proj_weight, config.dtype)\n", | 508 | " gate = pypto.matmul(hidden_states, gate_proj_weight, config.dtype)\n", |
| 498 | " up = pypto.matmul(hidden_states, up_proj_weight, config.dtype)\n", | 509 | " up = pypto.matmul(hidden_states, up_proj_weight, config.dtype)\n", |
| 499 | " activated = swiglu_activation_core(gate, up)\n", | 510 | " activated = swiglu_activation_core(gate, up)\n", |
| @@ -502,6 +513,7 @@ | |||
| 502 | " activated = relu_activation_core(up)\n", | 513 | " activated = relu_activation_core(up)\n", |
| 503 | " else:\n", | 514 | " else:\n", |
| 504 | " raise ValueError(f\"Unsupported activation: {config.activation}\")\n", | 515 | " raise ValueError(f\"Unsupported activation: {config.activation}\")\n", |
| 516 | + " # 下投影:[B, I] @ [I, H] -> [B, H]\n", | ||
| 505 | " result = pypto.matmul(activated, down_proj_weight, config.dtype, b_trans=False)\n", | 517 | " result = pypto.matmul(activated, down_proj_weight, config.dtype, b_trans=False)\n", |
| 506 | " pypto.assemble(result, [0, 0], output)" | 518 | " pypto.assemble(result, [0, 0], output)" |
| 507 | ] | 519 | ] |
| @@ -529,7 +541,7 @@ | |||
| 529 | "| `hidden_size` | 输入和输出 hidden 维度。静态测试是 128,动态测试是 512。 |\n", | 541 | "| `hidden_size` | 输入和输出 hidden 维度。静态测试是 128,动态测试是 512。 |\n", |
| 530 | "| `intermediate_size` | FFN 中间维度,通常大于 hidden 维度。这里使用 1024。 |\n", | 542 | "| `intermediate_size` | FFN 中间维度,通常大于 hidden 维度。这里使用 1024。 |\n", |
| 531 | "| `activation` | 选择 `gelu`、`swiglu` 或 `relu` 分支。 |\n", | 543 | "| `activation` | 选择 `gelu`、`swiglu` 或 `relu` 分支。 |\n", |
| 532 | - "| `dtype` | matmul 和输出使用的数据类型,如 `pypto.DT_BF16` 或 `pypto.DT_FP16`。 |\n", | 544 | + "| `dtype` | matmul 和输出使用的数据类型,如 `pypto.DT_BF16`。 |\n", |
| 533 | "| `use_dynamic_shape` | 是否走动态 batch 版本。 |\n", | 545 | "| `use_dynamic_shape` | 是否走动态 batch 版本。 |\n", |
| 534 | "| `vec_tile_shape` | 激活、逐元素计算使用的向量 tile。 |\n", | 546 | "| `vec_tile_shape` | 激活、逐元素计算使用的向量 tile。 |\n", |
| 535 | "| `cube_tile_shape` | matmul 使用的 cube tile。 |\n", | 547 | "| `cube_tile_shape` | matmul 使用的 cube tile。 |\n", |
| @@ -819,7 +831,7 @@ | |||
| 819 | "| `num_iterations = ceil_div(batch_size, basic_batch)` | 计算需要循环多少次,保证最后一个尾块也能覆盖。 |\n", | 831 | "| `num_iterations = ceil_div(batch_size, basic_batch)` | 计算需要循环多少次,保证最后一个尾块也能覆盖。 |\n", |
| 820 | "| `pypto.loop(...)` | 在 kernel 里显式描述 batch 循环。 |\n", | 832 | "| `pypto.loop(...)` | 在 kernel 里显式描述 batch 循环。 |\n", |
| 821 | "| `pypto.view(hidden_states, ...)` | 从大输入中取出当前 batch 块,形成局部视图。 |\n", | 833 | "| `pypto.view(hidden_states, ...)` | 从大输入中取出当前 batch 块,形成局部视图。 |\n", |
| 822 | - "| `valid_shape=[(batch_size - batch_offset).min(basic_batch), hidden_size]` | 告诉 PyPTO 当前块可能是尾块,不一定满 `basic_batch`;这里要使用 PyPTO 符号值的 `.min(...)`,不能用 Python 内置 `min(...)`。 |\n", | 834 | + "| `valid_shape=[(batch_size - batch_offset).min(basic_batch), hidden_size]` | 告诉 PyPTO 当前块可能是尾块,不一定满 `basic_batch` |\n", |
| 823 | "| `pypto.matmul(hidden_chunk, up_proj_weight, ...)` | 在局部块上做上投影。 |\n", | 835 | "| `pypto.matmul(hidden_chunk, up_proj_weight, ...)` | 在局部块上做上投影。 |\n", |
| 824 | "| `pypto.assemble(output_chunk, [batch_offset, 0], output)` | 把局部结果写回全局输出的对应位置。 |\n", | 836 | "| `pypto.assemble(output_chunk, [batch_offset, 0], output)` | 把局部结果写回全局输出的对应位置。 |\n", |
| 825 | "\n", | 837 | "\n", |
| @@ -837,7 +849,7 @@ | |||
| 837 | "| --- | --- |\n", | 849 | "| --- | --- |\n", |
| 838 | "| `pypto.sum` | LayerNorm / RMSNorm 里的规约计算 |\n", | 850 | "| `pypto.sum` | LayerNorm / RMSNorm 里的规约计算 |\n", |
| 839 | "| `pypto.sqrt` | 归一化中的方差或均方根开根号 |\n", | 851 | "| `pypto.sqrt` | 归一化中的方差或均方根开根号 |\n", |
| 840 | - "| `pypto.cast` | 在 FFN 中做 FP16/FP32 间转换 |\n", | 852 | + "| `pypto.cast` | 在 FFN 中做 BF16/FP32 间转换 |\n", |
| 841 | "| `pypto.matmul` | FFN 的投影层核心 |\n", | 853 | "| `pypto.matmul` | FFN 的投影层核心 |\n", |
| 842 | "| `pypto.maximum` | ReLU 分支的实现 |\n", | 854 | "| `pypto.maximum` | ReLU 分支的实现 |\n", |
| 843 | "| `pypto.view` | 动态 batch 切块 |\n", | 855 | "| `pypto.view` | 动态 batch 切块 |\n", |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "title", | 5 | "id": "title", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 4.4 动态 Shape 与控制流\n", | 8 | + "# 动态 Shape 与控制流\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "这一节开始处理更接近真实工程的执行逻辑:动态 shape、循环和条件分支。\n", | 10 | "这一节开始处理更接近真实工程的执行逻辑:动态 shape、循环和条件分支。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -17,9 +17,7 @@ | |||
| 17 | "id": "env-title", | 17 | "id": "env-title", |
| 18 | "metadata": {}, | 18 | "metadata": {}, |
| 19 | "source": [ | 19 | "source": [ |
| 20 | - "## 1. 环境准备\n", | 20 | + "## 1. 环境准备\n" |
| 21 | - "\n", | ||
| 22 | - "本节练习数量较多,先统一准备环境、参考实现和运行模式。" | ||
| 23 | ] | 21 | ] |
| 24 | }, | 22 | }, |
| 25 | { | 23 | { |
| @@ -31,6 +29,7 @@ | |||
| 31 | "source": [ | 29 | "source": [ |
| 32 | "import os\n", | 30 | "import os\n", |
| 33 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 31 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 32 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 34 | "from dataclasses import dataclass\n", | 33 | "from dataclasses import dataclass\n", |
| 35 | "from typing import Optional\n", | 34 | "from typing import Optional\n", |
| 36 | "import torch\n", | 35 | "import torch\n", |
| @@ -203,6 +202,7 @@ | |||
| 203 | " for idx in pypto.loop(b_loop):\n", | 202 | " for idx in pypto.loop(b_loop):\n", |
| 204 | " b_offset = idx * tile_b\n", | 203 | " b_offset = idx * tile_b\n", |
| 205 | " b_offset_end = (b_offset + tile_b).min(batch_size_dyn)\n", | 204 | " b_offset_end = (b_offset + tile_b).min(batch_size_dyn)\n", |
| 205 | + " # valid_shape 声明当前 tile 的有效数据范围\n", | ||
| 206 | " valid_shape = [b_offset_end - b_offset, 128]\n", | 206 | " valid_shape = [b_offset_end - b_offset, 128]\n", |
| 207 | " x_view = pypto.view(x, [tile_b, 128], [b_offset, 0], valid_shape=valid_shape)\n", | 207 | " x_view = pypto.view(x, [tile_b, 128], [b_offset, 0], valid_shape=valid_shape)\n", |
| 208 | " pypto.set_vec_tile_shapes(1, 128)\n", | 208 | " pypto.set_vec_tile_shapes(1, 128)\n", |
| @@ -262,7 +262,7 @@ | |||
| 262 | " output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32)):\n", | 262 | " output_tensor: pypto.Tensor([pypto.DYNAMIC, ...], pypto.DT_FP32)):\n", |
| 263 | " tile_b = 1\n", | 263 | " tile_b = 1\n", |
| 264 | " bs_dyn, seqlen, head, dim = input_tensor.shape\n", | 264 | " bs_dyn, seqlen, head, dim = input_tensor.shape\n", |
| 265 | - " b_loop = bs_dyn // tile_b\n", | 265 | + " b_loop = (bs_dyn + tile_b - 1) // tile_b\n", |
| 266 | "\n", | 266 | "\n", |
| 267 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | 267 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", |
| 268 | " for idx in pypto.loop(0, b_loop, 1, name=\"LOOP_L0_bIdx\", idx_name=\"idx\"):\n", | 268 | " for idx in pypto.loop(0, b_loop, 1, name=\"LOOP_L0_bIdx\", idx_name=\"idx\"):\n", |
| @@ -339,6 +339,7 @@ | |||
| 339 | "\n", | 339 | "\n", |
| 340 | "def scaled_dot_product_attention_core(q: pypto.Tensor, k: pypto.Tensor, v: pypto.Tensor,\n", | 340 | "def scaled_dot_product_attention_core(q: pypto.Tensor, k: pypto.Tensor, v: pypto.Tensor,\n", |
| 341 | " scale: float, dtype: pypto.DataType) -> pypto.Tensor:\n", | 341 | " scale: float, dtype: pypto.DataType) -> pypto.Tensor:\n", |
| 342 | + " # 交换最后两维\n", | ||
| 342 | " k_t = pypto.transpose(k, 2, 3)\n", | 343 | " k_t = pypto.transpose(k, 2, 3)\n", |
| 343 | " scores = pypto.matmul(q, k_t, out_dtype=dtype)\n", | 344 | " scores = pypto.matmul(q, k_t, out_dtype=dtype)\n", |
| 344 | " scores_scaled = scores * scale\n", | 345 | " scores_scaled = scores * scale\n", |
| @@ -359,6 +360,7 @@ | |||
| 359 | " dim = config.head_dim\n", | 360 | " dim = config.head_dim\n", |
| 360 | " q_len = q.shape[2]\n", | 361 | " q_len = q.shape[2]\n", |
| 361 | " kv_len = k.shape[2]\n", | 362 | " kv_len = k.shape[2]\n", |
| 363 | + " # 缩放因子:1/sqrt(d_k),未配置时自动计算\n", | ||
| 362 | " scale = config.scale if config.scale is not None else (1.0 / (dim ** 0.5))\n", | 364 | " scale = config.scale if config.scale is not None else (1.0 / (dim ** 0.5))\n", |
| 363 | " cube_tiling = 64\n", | 365 | " cube_tiling = 64\n", |
| 364 | " pypto.set_cube_tile_shapes([cube_tiling, cube_tiling], [cube_tiling, cube_tiling], [cube_tiling, cube_tiling])\n", | 366 | " pypto.set_cube_tile_shapes([cube_tiling, cube_tiling], [cube_tiling, cube_tiling], [cube_tiling, cube_tiling])\n", |
| @@ -367,6 +369,7 @@ | |||
| 367 | " for bss_idx in pypto.loop(bs_loop):\n", | 369 | " for bss_idx in pypto.loop(bs_loop):\n", |
| 368 | " bs_offset = bss_idx * tile\n", | 370 | " bs_offset = bss_idx * tile\n", |
| 369 | " bs_offset_end = (bs_offset + tile).min(bs_dyn)\n", | 371 | " bs_offset_end = (bs_offset + tile).min(bs_dyn)\n", |
| 372 | + " # Q/K/V 各沿 batch 维取局部 tile,valid_shape 处理尾块\n", | ||
| 370 | " q_view = pypto.view(q, [tile, head, q_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, q_len, dim])\n", | 373 | " q_view = pypto.view(q, [tile, head, q_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, q_len, dim])\n", |
| 371 | " k_view = pypto.view(k, [tile, head, kv_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, kv_len, dim])\n", | 374 | " k_view = pypto.view(k, [tile, head, kv_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, kv_len, dim])\n", |
| 372 | " v_view = pypto.view(v, [tile, head, kv_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, kv_len, dim])\n", | 375 | " v_view = pypto.view(v, [tile, head, kv_len, dim], [bs_offset, 0, 0, 0], valid_shape=[bs_offset_end - bs_offset, head, kv_len, dim])\n", |
| @@ -432,11 +435,13 @@ | |||
| 432 | " hidden_dyn = x.shape[1]\n", | 435 | " hidden_dyn = x.shape[1]\n", |
| 433 | " b_loop = (batch_dyn + tile_b - 1) // tile_b\n", | 436 | " b_loop = (batch_dyn + tile_b - 1) // tile_b\n", |
| 434 | "\n", | 437 | "\n", |
| 438 | + " # 外层 loop:按 batch 维分块\n", | ||
| 435 | " for b_idx in pypto.loop(b_loop):\n", | 439 | " for b_idx in pypto.loop(b_loop):\n", |
| 436 | " b_offset = b_idx * tile_b\n", | 440 | " b_offset = b_idx * tile_b\n", |
| 437 | " b_offset_end = (b_offset + tile_b).min(batch_dyn)\n", | 441 | " b_offset_end = (b_offset + tile_b).min(batch_dyn)\n", |
| 438 | " valid_b = b_offset_end - b_offset\n", | 442 | " valid_b = b_offset_end - b_offset\n", |
| 439 | " h_loop = (hidden_dyn + tile_h - 1) // tile_h\n", | 443 | " h_loop = (hidden_dyn + tile_h - 1) // tile_h\n", |
| 444 | + " # 内层 loop:按 hidden 维分块\n", | ||
| 440 | " for h_idx in pypto.loop(h_loop):\n", | 445 | " for h_idx in pypto.loop(h_loop):\n", |
| 441 | " h_offset = h_idx * tile_h\n", | 446 | " h_offset = h_idx * tile_h\n", |
| 442 | " h_offset_end = (h_offset + tile_h).min(hidden_dyn)\n", | 447 | " h_offset_end = (h_offset + tile_h).min(hidden_dyn)\n", |
| @@ -573,7 +578,7 @@ | |||
| 573 | " for outside_idx in pypto.loop(5):\n", | 578 | " for outside_idx in pypto.loop(5):\n", |
| 574 | " print(f\"outside_idx: {outside_idx}\")\n", | 579 | " print(f\"outside_idx: {outside_idx}\")\n", |
| 575 | " for inside_idx in pypto.loop(3):\n", | 580 | " for inside_idx in pypto.loop(3):\n", |
| 576 | - " print(f\"inside_idx: {outside_idx}\")\n", | 581 | + " print(f\"inside_idx: {inside_idx}\")\n", |
| 577 | " res = pypto.add(in_t0, in_t0)\n", | 582 | " res = pypto.add(in_t0, in_t0)\n", |
| 578 | " print(f\"res: {res}\")\n", | 583 | " print(f\"res: {res}\")\n", |
| 579 | " if outside_idx < 3:\n", | 584 | " if outside_idx < 3:\n", |
| @@ -673,7 +678,7 @@ | |||
| 673 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | 678 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", |
| 674 | " b = input0.shape[0]\n", | 679 | " b = input0.shape[0]\n", |
| 675 | " tile_b = 1\n", | 680 | " tile_b = 1\n", |
| 676 | - " b_loop = b // tile_b\n", | 681 | + " b_loop = (b + tile_b - 1) // tile_b\n", |
| 677 | " for idx in pypto.loop(b_loop):\n", | 682 | " for idx in pypto.loop(b_loop):\n", |
| 678 | " b_offset = idx * tile_b\n", | 683 | " b_offset = idx * tile_b\n", |
| 679 | " b_offset_end = (idx + 1) * tile_b\n", | 684 | " b_offset_end = (idx + 1) * tile_b\n", |
| @@ -692,7 +697,7 @@ | |||
| 692 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | 697 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", |
| 693 | " b, w, n, c = input0.shape\n", | 698 | " b, w, n, c = input0.shape\n", |
| 694 | " tile_b = 1\n", | 699 | " tile_b = 1\n", |
| 695 | - " b_loop = b // tile_b\n", | 700 | + " b_loop = (b + tile_b - 1) // tile_b\n", |
| 696 | " for idx in pypto.loop(b_loop):\n", | 701 | " for idx in pypto.loop(b_loop):\n", |
| 697 | " b_offset = idx * tile_b\n", | 702 | " b_offset = idx * tile_b\n", |
| 698 | " b_offset_end = min((idx + 1) * tile_b, b)\n", | 703 | " b_offset_end = min((idx + 1) * tile_b, b)\n", |
| @@ -775,7 +780,7 @@ | |||
| 775 | "\n", | 780 | "\n", |
| 776 | "| 练习 | reference 逻辑 |\n", | 781 | "| 练习 | reference 逻辑 |\n", |
| 777 | "| --- | --- |\n", | 782 | "| --- | --- |\n", |
| 778 | - "| `test_nested_loops_with_conditions` | `golden[0] = a[0] + b[0]`?`golden[1] = a[1] - b[1]`? |\n", | 783 | + "| `test_nested_loops_with_conditions` | `golden[0] = a[0] + b[0]` `golden[1] = a[1] - b[1]` |\n", |
| 779 | "| `test_add_scalar_loop_dyn_axis_static_cond` | `flag=False` 时是 `x + y`,`flag=True` 时是 `x + y + val`。 |\n", | 784 | "| `test_add_scalar_loop_dyn_axis_static_cond` | `flag=False` 时是 `x + y`,`flag=True` 时是 `x + y + val`。 |\n", |
| 780 | "| `test_add_scalar_loop_dynamic_axis_dynamic_cond` | 先整体 `x + y`,再只给前两个 batch 加 `val`。 |\n", | 785 | "| `test_add_scalar_loop_dynamic_axis_dynamic_cond` | 先整体 `x + y`,再只给前两个 batch 加 `val`。 |\n", |
| 781 | "| `test_add_scalar_loop_dynamic_axis_dynamic_loop_cond` | 第一个 batch 加 `val`,最后一个 batch 加 `val + 1`。 |\n", | 786 | "| `test_add_scalar_loop_dynamic_axis_dynamic_loop_cond` | 第一个 batch 加 `val`,最后一个 batch 加 `val + 1`。 |\n", |
| @@ -810,7 +815,7 @@ | |||
| 810 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", | 815 | " pypto.set_vec_tile_shapes(1, 4, 1, 64)\n", |
| 811 | " b = input0.shape[0]\n", | 816 | " b = input0.shape[0]\n", |
| 812 | " tile_b = 1\n", | 817 | " tile_b = 1\n", |
| 813 | - " b_loop = b // tile_b\n", | 818 | + " b_loop = (b + tile_b - 1) // tile_b\n", |
| 814 | " for idx in pypto.loop(b_loop):\n", | 819 | " for idx in pypto.loop(b_loop):\n", |
| 815 | " b_offset = idx * tile_b\n", | 820 | " b_offset = idx * tile_b\n", |
| 816 | " b_offset_end = (idx + 1) * tile_b\n", | 821 | " b_offset_end = (idx + 1) * tile_b\n", |
| @@ -865,6 +870,7 @@ | |||
| 865 | " b = input0.shape[0]\n", | 870 | " b = input0.shape[0]\n", |
| 866 | " for idx in pypto.loop(b):\n", | 871 | " for idx in pypto.loop(b):\n", |
| 867 | " t3_sub = input0[idx:idx + 1, ...] + input1[idx:idx + 1, ...]\n", | 872 | " t3_sub = input0[idx:idx + 1, ...] + input1[idx:idx + 1, ...]\n", |
| 873 | + " # is_loop_begin:首块加 val\n", | ||
| 868 | " if pypto.is_loop_begin(idx):\n", | 874 | " if pypto.is_loop_begin(idx):\n", |
| 869 | " output[idx:idx + 1, ...] = t3_sub + val\n", | 875 | " output[idx:idx + 1, ...] = t3_sub + val\n", |
| 870 | " elif pypto.is_loop_end(idx):\n", | 876 | " elif pypto.is_loop_end(idx):\n", |
| @@ -5,7 +5,7 @@ | |||
| 5 | "id": "title", | 5 | "id": "title", |
| 6 | "metadata": {}, | 6 | "metadata": {}, |
| 7 | "source": [ | 7 | "source": [ |
| 8 | - "# 4.5 Attention 与 Transformer 组合\n", | 8 | + "# Attention 与 Transformer 组合\n", |
| 9 | "\n", | 9 | "\n", |
| 10 | "这一节进入更完整的神经网络模块:先把 Attention 的数据流拆清楚,再把多个 JIT kernel 组织成 Transformer block 的雏形。\n", | 10 | "这一节进入更完整的神经网络模块:先把 Attention 的数据流拆清楚,再把多个 JIT kernel 组织成 Transformer block 的雏形。\n", |
| 11 | "\n", | 11 | "\n", |
| @@ -17,9 +17,7 @@ | |||
| 17 | "id": "env-title", | 17 | "id": "env-title", |
| 18 | "metadata": {}, | 18 | "metadata": {}, |
| 19 | "source": [ | 19 | "source": [ |
| 20 | - "## 1. 环境准备\n", | 20 | + "## 1. 环境准备" |
| 21 | - "\n", | ||
| 22 | - "本节会同时使用 matmul、transpose、reshape、softmax、LayerNorm、GELU 和 residual add,先统一准备运行环境。" | ||
| 23 | ] | 21 | ] |
| 24 | }, | 22 | }, |
| 25 | { | 23 | { |
| @@ -31,6 +29,7 @@ | |||
| 31 | "source": [ | 29 | "source": [ |
| 32 | "import os\n", | 30 | "import os\n", |
| 33 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | 31 | "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", |
| 32 | + "os.environ['TORCH_DEVICE_BACKEND_AUTOLOAD'] = '0'\n", | ||
| 34 | "from dataclasses import dataclass\n", | 33 | "from dataclasses import dataclass\n", |
| 35 | "from typing import Optional\n", | 34 | "from typing import Optional\n", |
| 36 | "import torch\n", | 35 | "import torch\n", |
| @@ -166,35 +165,7 @@ | |||
| 166 | "SEQ_LEN = 32\n", | 165 | "SEQ_LEN = 32\n", |
| 167 | "NUM_HEADS = 8\n", | 166 | "NUM_HEADS = 8\n", |
| 168 | "HEAD_DIM = 64\n", | 167 | "HEAD_DIM = 64\n", |
| 169 | - "HIDDEN_SIZE = 512\n", | 168 | + "HIDDEN_SIZE = 512" |
| 170 | - "\n", | ||
| 171 | - "\n", | ||
| 172 | - "@dataclass\n", | ||
| 173 | - "class AttentionConfig:\n", | ||
| 174 | - " num_heads: int = 8\n", | ||
| 175 | - " head_dim: int = 64\n", | ||
| 176 | - " scale: Optional[float] = None\n", | ||
| 177 | - " dtype: pypto.DataType = pypto.DT_BF16\n", | ||
| 178 | - " use_dynamic_shape: bool = False\n", | ||
| 179 | - "\n", | ||
| 180 | - "\n", | ||
| 181 | - "def scaled_dot_product_attention_golden(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,\n", | ||
| 182 | - " scale: float, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n", | ||
| 183 | - " scores = torch.matmul(q, k.transpose(-2, -1))\n", | ||
| 184 | - " scores = scores * scale\n", | ||
| 185 | - " if attn_mask is not None:\n", | ||
| 186 | - " scores = scores + attn_mask\n", | ||
| 187 | - " attn_weights = torch.softmax(scores, dim=-1)\n", | ||
| 188 | - " return torch.matmul(attn_weights, v)\n", | ||
| 189 | - "\n", | ||
| 190 | - "\n", | ||
| 191 | - "def scaled_dot_product_attention_core(q: pypto.Tensor, k: pypto.Tensor, v: pypto.Tensor,\n", | ||
| 192 | - " scale: float, dtype: pypto.DataType) -> pypto.Tensor:\n", | ||
| 193 | - " k_t = pypto.transpose(k, 2, 3)\n", | ||
| 194 | - " scores = pypto.matmul(q, k_t, out_dtype=dtype)\n", | ||
| 195 | - " scores_scaled = scores * scale\n", | ||
| 196 | - " attn_weights = pypto.softmax(scores_scaled, dim=-1)\n", | ||
| 197 | - " return pypto.matmul(attn_weights, v, out_dtype=dtype)\n" | ||
| 198 | ] | 169 | ] |
| 199 | }, | 170 | }, |
| 200 | { | 171 | { |
| @@ -226,6 +197,15 @@ | |||
| 226 | "metadata": {}, | 197 | "metadata": {}, |
| 227 | "outputs": [], | 198 | "outputs": [], |
| 228 | "source": [ | 199 | "source": [ |
| 200 | + "def scaled_dot_product_attention_golden(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,\n", | ||
| 201 | + " scale: float, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:\n", | ||
| 202 | + " scores = torch.matmul(q, k.transpose(-2, -1))\n", | ||
| 203 | + " scores = scores * scale\n", | ||
| 204 | + " if attn_mask is not None:\n", | ||
| 205 | + " scores = scores + attn_mask\n", | ||
| 206 | + " attn_weights = torch.softmax(scores, dim=-1)\n", | ||
| 207 | + " return torch.matmul(attn_weights, v)\n", | ||
| 208 | + "\n", | ||
| 229 | "@pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", | 209 | "@pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", |
| 230 | "def scaled_dot_product_attention_kernel(\n", | 210 | "def scaled_dot_product_attention_kernel(\n", |
| 231 | " q: pypto.Tensor((BATCH_SIZE, NUM_HEADS, SEQ_LEN_Q, HEAD_DIM), pypto.DT_BF16),\n", | 211 | " q: pypto.Tensor((BATCH_SIZE, NUM_HEADS, SEQ_LEN_Q, HEAD_DIM), pypto.DT_BF16),\n", |
| @@ -237,7 +217,9 @@ | |||
| 237 | " pypto.set_vec_tile_shapes(1, 8, 16, HEAD_DIM)\n", | 217 | " pypto.set_vec_tile_shapes(1, 8, 16, HEAD_DIM)\n", |
| 238 | " scores = pypto.matmul(q, pypto.transpose(k, 2, 3), out_dtype=pypto.DT_BF16)\n", | 218 | " scores = pypto.matmul(q, pypto.transpose(k, 2, 3), out_dtype=pypto.DT_BF16)\n", |
| 239 | " scores_scaled = pypto.mul(scores, scale)\n", | 219 | " scores_scaled = pypto.mul(scores, scale)\n", |
| 220 | + " # 归一化\n", | ||
| 240 | " attn_weights = pypto.softmax(scores_scaled, dim=-1)\n", | 221 | " attn_weights = pypto.softmax(scores_scaled, dim=-1)\n", |
| 222 | + " # 用权重V加权\n", | ||
| 241 | " output.move(pypto.matmul(attn_weights, v, out_dtype=pypto.DT_BF16))\n", | 223 | " output.move(pypto.matmul(attn_weights, v, out_dtype=pypto.DT_BF16))\n", |
| 242 | "\n", | 224 | "\n", |
| 243 | "\n", | 225 | "\n", |
| @@ -253,17 +235,10 @@ | |||
| 253 | " print(f\"Scaled attention output shape: {out.shape}\")\n", | 235 | " print(f\"Scaled attention output shape: {out.shape}\")\n", |
| 254 | " max_diff = (out - golden).abs().max().item()\n", | 236 | " max_diff = (out - golden).abs().max().item()\n", |
| 255 | " print(f\"Scaled attention max diff: {max_diff:.6f}\")\n", | 237 | " print(f\"Scaled attention max diff: {max_diff:.6f}\")\n", |
| 256 | - " assert max_diff < 3e-3, \"Scaled attention mismatch!\"\n" | 238 | + " assert max_diff < 3e-3, \"Scaled attention mismatch!\"\n", |
| 257 | - ] | 239 | + "\n", |
| 258 | - }, | 240 | + "\n", |
| 259 | - { | 241 | + "test_scaled_dot_product_attention()" |
| 260 | - "cell_type": "code", | ||
| 261 | - "execution_count": null, | ||
| 262 | - "id": "62fbddf2", | ||
| 263 | - "metadata": {}, | ||
| 264 | - "outputs": [], | ||
| 265 | - "source": [ | ||
| 266 | - "test_scaled_dot_product_attention()\n" | ||
| 267 | ] | 242 | ] |
| 268 | }, | 243 | }, |
| 269 | { | 244 | { |
| @@ -299,78 +274,112 @@ | |||
| 299 | "metadata": {}, | 274 | "metadata": {}, |
| 300 | "outputs": [], | 275 | "outputs": [], |
| 301 | "source": [ | 276 | "source": [ |
| 302 | - "def attention_with_projection_golden(hidden_states: torch.Tensor, q_weight: torch.Tensor,\n", | ||
| 303 | - " k_weight: torch.Tensor, v_weight: torch.Tensor,\n", | ||
| 304 | - " out_weight: torch.Tensor) -> torch.Tensor:\n", | ||
| 305 | - " q = torch.matmul(hidden_states, q_weight)\n", | ||
| 306 | - " k = torch.matmul(hidden_states, k_weight)\n", | ||
| 307 | - " v = torch.matmul(hidden_states, v_weight)\n", | ||
| 308 | - " batch_size, seq_len, _ = q.shape\n", | ||
| 309 | - " q = q.view(batch_size, seq_len, NUM_HEADS, HEAD_DIM).transpose(1, 2)\n", | ||
| 310 | - " k = k.view(batch_size, seq_len, NUM_HEADS, HEAD_DIM).transpose(1, 2)\n", | ||
| 311 | - " v = v.view(batch_size, seq_len, NUM_HEADS, HEAD_DIM).transpose(1, 2)\n", | ||
| 312 | - " scores = torch.matmul(q, k.transpose(-2, -1)) * (1.0 / (HEAD_DIM ** 0.5))\n", | ||
| 313 | - " attn_weights = torch.softmax(scores, dim=-1)\n", | ||
| 314 | - " context = torch.matmul(attn_weights, v)\n", | ||
| 315 | - " context = context.transpose(1, 2).reshape(batch_size, seq_len, NUM_HEADS * HEAD_DIM)\n", | ||
| 316 | - " return torch.matmul(context, out_weight)\n", | ||
| 317 | - "\n", | ||
| 318 | - "\n", | ||
| 319 | "@pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", | 277 | "@pypto.frontend.jit(runtime_options={\"run_mode\": RUN_MODE})\n", |
| 320 | "def attention_with_projection_kernel(\n", | 278 | "def attention_with_projection_kernel(\n", |
| 321 | - " hidden_states: pypto.Tensor((BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE), pypto.DT_BF16),\n", | 279 | + " hidden_states: pypto.Tensor((BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE), pypto.DT_FP32),\n", |
| 322 | - " q_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_BF16),\n", | 280 | + " q_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_FP32),\n", |
| 323 | - " k_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_BF16),\n", | 281 | + " k_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_FP32),\n", |
| 324 | - " v_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_BF16),\n", | 282 | + " v_weight: pypto.Tensor((1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM), pypto.DT_FP32),\n", |
| 325 | - " out_weight: pypto.Tensor((1, NUM_HEADS * HEAD_DIM, HIDDEN_SIZE), pypto.DT_BF16),\n", | 283 | + " out_weight: pypto.Tensor((1, NUM_HEADS * HEAD_DIM, HIDDEN_SIZE), pypto.DT_FP32),\n", |
| 326 | - " output_tensor: pypto.Tensor((BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE), pypto.DT_BF16)):\n", | 284 | + " output_tensor: pypto.Tensor((BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE), pypto.DT_FP32),\n", |
| 285 | + "):\n", | ||
| 327 | " tile_b = 1\n", | 286 | " tile_b = 1\n", |
| 328 | " b_loop = BATCH_SIZE // tile_b\n", | 287 | " b_loop = BATCH_SIZE // tile_b\n", |
| 329 | - " scale = 1.0 / (HEAD_DIM ** 0.5)\n", | 288 | + "\n", |
| 289 | + " scale = 1.0 / (HEAD_DIM**0.5)\n", | ||
| 330 | " pypto.set_cube_tile_shapes([64, 64], [64, 64], [64, 64])\n", | 290 | " pypto.set_cube_tile_shapes([64, 64], [64, 64], [64, 64])\n", |
| 331 | - " pypto.set_vec_tile_shapes(1, 16, 8, HEAD_DIM)\n", | 291 | + " pypto.set_vec_tile_shapes(1, 8, 32, HEAD_DIM)\n", |
| 332 | "\n", | 292 | "\n", |
| 333 | - " q_flat = pypto.matmul(hidden_states, q_weight, out_dtype=pypto.DT_BF16)\n", | 293 | + " # Q/K/V 线性投影\n", |
| 334 | - " k_flat = pypto.matmul(hidden_states, k_weight, out_dtype=pypto.DT_BF16)\n", | 294 | + " q_flat = pypto.matmul(hidden_states, q_weight, out_dtype=pypto.DT_FP32)\n", |
| 335 | - " v_flat = pypto.matmul(hidden_states, v_weight, out_dtype=pypto.DT_BF16)\n", | 295 | + " k_flat = pypto.matmul(hidden_states, k_weight, out_dtype=pypto.DT_FP32)\n", |
| 336 | - " q = pypto.transpose(pypto.reshape(q_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM]), 1, 2)\n", | 296 | + " v_flat = pypto.matmul(hidden_states, v_weight, out_dtype=pypto.DT_FP32)\n", |
| 337 | - " k = pypto.transpose(pypto.reshape(k_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM]), 1, 2)\n", | ||
| 338 | - " v = pypto.transpose(pypto.reshape(v_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM]), 1, 2)\n", | ||
| 339 | "\n", | 297 | "\n", |
| 298 | + " # Q/K/V 多头拆分\n", | ||
| 299 | + " q = pypto.reshape(q_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM])\n", | ||
| 300 | + " k = pypto.reshape(k_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM])\n", | ||
| 301 | + " v = pypto.reshape(v_flat, [BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM])\n", | ||
| 302 | + "\n", | ||
| 303 | + " # 把 head 维提前\n", | ||
| 304 | + " q = pypto.transpose(q, 1, 2)\n", | ||
| 305 | + " k = pypto.transpose(k, 1, 2)\n", | ||
| 306 | + " v = pypto.transpose(v, 1, 2)\n", | ||
| 307 | + "\n", | ||
| 308 | + " # 按 batch tile 循环做 Attention\n", | ||
| 340 | " for idx in pypto.loop(0, b_loop, 1, name=\"LOOP_L0_bIdx\", idx_name=\"idx\"):\n", | 309 | " for idx in pypto.loop(0, b_loop, 1, name=\"LOOP_L0_bIdx\", idx_name=\"idx\"):\n", |
| 341 | " b_offset = idx * tile_b\n", | 310 | " b_offset = idx * tile_b\n", |
| 342 | - " b_offset_end = ((idx + 1) * tile_b).min(BATCH_SIZE)\n", | 311 | + " b_offset_end = min((idx + 1) * tile_b, BATCH_SIZE)\n", |
| 343 | " view_shape = [tile_b, NUM_HEADS, SEQ_LEN, HEAD_DIM]\n", | 312 | " view_shape = [tile_b, NUM_HEADS, SEQ_LEN, HEAD_DIM]\n", |
| 344 | " valid_shape = [b_offset_end - b_offset, NUM_HEADS, SEQ_LEN, HEAD_DIM]\n", | 313 | " valid_shape = [b_offset_end - b_offset, NUM_HEADS, SEQ_LEN, HEAD_DIM]\n", |
| 314 | + " # 取当前 batch tile 的 Q/K/V\n", | ||
| 345 | " q_view = pypto.view(q, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", | 315 | " q_view = pypto.view(q, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", |
| 346 | " k_view = pypto.view(k, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", | 316 | " k_view = pypto.view(k, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", |
| 347 | " v_view = pypto.view(v, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", | 317 | " v_view = pypto.view(v, view_shape, [b_offset, 0, 0, 0], valid_shape=valid_shape)\n", |
| 348 | - " scores = pypto.matmul(q_view, pypto.transpose(k_view, 2, 3), out_dtype=pypto.DT_BF16)\n", | 318 | + "\n", |
| 319 | + " scores = pypto.matmul(q_view, pypto.transpose(k_view, 2, 3), out_dtype=pypto.DT_FP32)\n", | ||
| 349 | " scores_scaled = pypto.mul(scores, scale)\n", | 320 | " scores_scaled = pypto.mul(scores, scale)\n", |
| 350 | " attn_weights = pypto.softmax(scores_scaled, dim=-1)\n", | 321 | " attn_weights = pypto.softmax(scores_scaled, dim=-1)\n", |
| 351 | - " context = pypto.matmul(attn_weights, v_view, out_dtype=pypto.DT_BF16)\n", | 322 | + " context = pypto.matmul(attn_weights, v_view, out_dtype=pypto.DT_FP32)\n", |
| 323 | + "\n", | ||
| 324 | + " # 多头合并 + 输出投影\n", | ||
| 352 | " context = pypto.transpose(context, 1, 2)\n", | 325 | " context = pypto.transpose(context, 1, 2)\n", |
| 353 | " context_flat = pypto.reshape(context, [tile_b, SEQ_LEN, NUM_HEADS * HEAD_DIM])\n", | 326 | " context_flat = pypto.reshape(context, [tile_b, SEQ_LEN, NUM_HEADS * HEAD_DIM])\n", |
| 354 | - " output_view = pypto.matmul(context_flat, out_weight, out_dtype=pypto.DT_BF16)\n", | 327 | + " output_view = pypto.matmul(context_flat, out_weight, out_dtype=pypto.DT_FP32)\n", |
| 355 | " output_tensor[b_offset:b_offset_end, ...] = output_view\n", | 328 | " output_tensor[b_offset:b_offset_end, ...] = output_view\n", |
| 356 | "\n", | 329 | "\n", |
| 357 | "\n", | 330 | "\n", |
| 358 | - "def test_attention_with_projection() -> None:\n", | 331 | + "def attention_with_projection_golden(\n", |
| 359 | - " device_local = get_device()\n", | 332 | + " hidden_states: torch.Tensor,\n", |
| 360 | - " hidden_states = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, dtype=torch.bfloat16, device=device_local)\n", | 333 | + " q_weight: torch.Tensor,\n", |
| 361 | - " q_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device_local)\n", | 334 | + " k_weight: torch.Tensor,\n", |
| 362 | - " k_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device_local)\n", | 335 | + " v_weight: torch.Tensor,\n", |
| 363 | - " v_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device_local)\n", | 336 | + " out_weight: torch.Tensor,\n", |
| 364 | - " out_weight = torch.randn(1, NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, dtype=torch.bfloat16, device=device_local)\n", | 337 | + ") -> torch.Tensor:\n", |
| 365 | - " out = torch.empty(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, dtype=torch.bfloat16, device=device_local)\n", | 338 | + " num_heads = NUM_HEADS\n", |
| 366 | - " attention_with_projection_kernel(hidden_states, q_weight, k_weight, v_weight, out_weight, out)\n", | 339 | + " head_dim = HEAD_DIM\n", |
| 367 | - " golden = attention_with_projection_golden(hidden_states, q_weight, k_weight, v_weight, out_weight)\n", | ||
| 368 | - " print(f\"Attention projection hidden shape: {hidden_states.shape}\")\n", | ||
| 369 | - " print(f\"Attention projection output shape: {out.shape}\")\n", | ||
| 370 | - " max_diff = (out - golden).abs().max().item()\n", | ||
| 371 | - " print(f\"Attention projection max diff: {max_diff:.6f}\")\n", | ||
| 372 | - " torch.allclose(out, golden, rtol=3e-3, atol=3e-3)\n", | ||
| 373 | "\n", | 340 | "\n", |
| 341 | + " q = torch.matmul(hidden_states, q_weight)\n", | ||
| 342 | + " k = torch.matmul(hidden_states, k_weight)\n", | ||
| 343 | + " v = torch.matmul(hidden_states, v_weight)\n", | ||
| 344 | + "\n", | ||
| 345 | + " batch_size, seq_len, _ = q.shape\n", | ||
| 346 | + " q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)\n", | ||
| 347 | + " k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)\n", | ||
| 348 | + " v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)\n", | ||
| 349 | + "\n", | ||
| 350 | + " scale = 1.0 / (head_dim**0.5)\n", | ||
| 351 | + " scores = torch.matmul(q, k.transpose(-2, -1)) * scale\n", | ||
| 352 | + " attn_weights = torch.softmax(scores, dim=-1)\n", | ||
| 353 | + " context = torch.matmul(attn_weights, v)\n", | ||
| 354 | + " context = context.transpose(1, 2).reshape(batch_size, seq_len, num_heads * head_dim)\n", | ||
| 355 | + " output = torch.matmul(context, out_weight)\n", | ||
| 356 | + " return output\n", | ||
| 357 | + "\n", | ||
| 358 | + "\n", | ||
| 359 | + "def test_attention_with_projection(device_id=None, dynamic: bool = False) -> None:\n", | ||
| 360 | + " device = get_device()\n", | ||
| 361 | + "\n", | ||
| 362 | + " # Create tensors\n", | ||
| 363 | + " hidden_states = torch.randn(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, dtype=torch.float32, device=device)\n", | ||
| 364 | + " q_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.float32, device=device)\n", | ||
| 365 | + " k_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.float32, device=device)\n", | ||
| 366 | + " v_weight = torch.randn(1, HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, dtype=torch.float32, device=device)\n", | ||
| 367 | + " out_weight = torch.randn(1, NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, dtype=torch.float32, device=device)\n", | ||
| 368 | + "\n", | ||
| 369 | + " # Execute\n", | ||
| 370 | + " out = torch.empty(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, dtype=torch.float32, device=device)\n", | ||
| 371 | + " attention_with_projection_kernel(hidden_states, q_weight, k_weight, v_weight, out_weight, out)\n", | ||
| 372 | + "\n", | ||
| 373 | + " golden = attention_with_projection_golden(hidden_states, q_weight, k_weight, v_weight, out_weight)\n", | ||
| 374 | + "\n", | ||
| 375 | + " print(f\"Hidden states shape: {hidden_states.shape}\")\n", | ||
| 376 | + " print(f\"Output shape: {out.shape}\")\n", | ||
| 377 | + " \n", | ||
| 378 | + " max_diff = (out - golden).abs().max().item()\n", | ||
| 379 | + " print(f\"Max difference: {max_diff:.6f}\")\n", | ||
| 380 | + " torch.testing.assert_close(out, golden, rtol=3e-3, atol=3e-3)\n", | ||
| 381 | + " print(\"✓ Attention with projections passed for the test case\")\n", | ||
| 382 | + " \n", | ||
| 374 | "\n", | 383 | "\n", |
| 375 | "test_attention_with_projection()" | 384 | "test_attention_with_projection()" |
| 376 | ] | 385 | ] |
| @@ -4,19 +4,8 @@ | |||
| 4 | "cell_type": "markdown", | 4 | "cell_type": "markdown", |
| 5 | "metadata": {}, | 5 | "metadata": {}, |
| 6 | "source": [ | 6 | "source": [ |
| 7 | - "# 4.7 章节实践:动态 Residual + LayerNorm + GELU 近似\n", | 7 | + "# 章节实践\n", |
| 8 | - "\n", | 8 | + "本节是第四章的综合实践,用一个融合算子检查中高级内容的掌握情况。实践目标不是引入新 API,而是把前面几节的能力组合起来:Residual Connection、LayerNorm、GELU 近似、动态 batch、`pypto.loop`、`pypto.view`、`valid_shape`、`pypto.assemble` 和 PyTorch reference 验证。\n" |
| 9 | - "本节是第四章的综合实践,用一个融合算子检查中高级内容的掌握情况。实践目标不是引入新 API,而是把前面几节的能力组合起来:Residual Connection、LayerNorm、GELU 近似、动态 batch、`pypto.loop`、`pypto.view`、`valid_shape`、`pypto.assemble` 和 PyTorch reference 验证。\n", | ||
| 10 | - "\n", | ||
| 11 | - "本节练习的计算目标是:\n", | ||
| 12 | - "\n", | ||
| 13 | - "```text\n", | ||
| 14 | - "residual = x + residual_input\n", | ||
| 15 | - "mean = mean(residual, dim=-1)\n", | ||
| 16 | - "var = mean((residual - mean)^2, dim=-1)\n", | ||
| 17 | - "norm = (residual - mean) / sqrt(var + eps)\n", | ||
| 18 | - "output = gelu_approx(norm * gamma + beta)\n", | ||
| 19 | - "```\n" | ||
| 20 | ] | 9 | ] |
| 21 | }, | 10 | }, |
| 22 | { | 11 | { |
| @@ -37,59 +26,7 @@ | |||
| 37 | "cell_type": "markdown", | 26 | "cell_type": "markdown", |
| 38 | "metadata": {}, | 27 | "metadata": {}, |
| 39 | "source": [ | 28 | "source": [ |
| 40 | - "## 2. 环境准备\n", | 29 | + "## 2. 章节实践题\n", |
| 41 | - "\n", | ||
| 42 | - "先运行下面的环境准备单元。没有可用 NPU 时,本节会保留代码阅读路径,并跳过真实设备验证。\n" | ||
| 43 | - ] | ||
| 44 | - }, | ||
| 45 | - { | ||
| 46 | - "cell_type": "code", | ||
| 47 | - "execution_count": null, | ||
| 48 | - "metadata": {}, | ||
| 49 | - "outputs": [], | ||
| 50 | - "source": [ | ||
| 51 | - "import os\n", | ||
| 52 | - "os.environ['TILE_FWK_DEVICE_ID'] = '0'\n", | ||
| 53 | - "import torch\n", | ||
| 54 | - "import pypto\n", | ||
| 55 | - "import torch_npu\n", | ||
| 56 | - "\n", | ||
| 57 | - "\n", | ||
| 58 | - "def get_device():\n", | ||
| 59 | - " device_id = int(os.environ.get(\"TILE_FWK_DEVICE_ID\", \"0\"))\n", | ||
| 60 | - " return f\"npu:{device_id}\"\n", | ||
| 61 | - "\n", | ||
| 62 | - "\n", | ||
| 63 | - "device = get_device()\n", | ||
| 64 | - "RUN_MODE = pypto.RunMode.NPU if device != \"cpu\" else pypto.RunMode.SIM\n", | ||
| 65 | - "print(\"TILE_FWK_DEVICE_ID:\", os.environ.get(\"TILE_FWK_DEVICE_ID\", \"<未设置,默认 0>\"))\n", | ||
| 66 | - "print(\"device:\", device)\n", | ||
| 67 | - "print(\"run_mode:\", RUN_MODE)\n" | ||
| 68 | - ] | ||
| 69 | - }, | ||
| 70 | - { | ||
| 71 | - "cell_type": "markdown", | ||
| 72 | - "metadata": {}, | ||
| 73 | - "source": [ | ||
| 74 | - "## 3. 先把计算拆开\n", | ||
| 75 | - "\n", | ||
| 76 | - "这个综合算子由四类能力组成:\n", | ||
| 77 | - "\n", | ||
| 78 | - "| 步骤 | PyPTO 能力 | 作用 |\n", | ||
| 79 | - "| --- | --- | --- |\n", | ||
| 80 | - "| Residual | elementwise add | 把主分支和残差分支相加 |\n", | ||
| 81 | - "| LayerNorm | `sum`、`sqrt`、广播 | 沿最后一维做归一化 |\n", | ||
| 82 | - "| GELU 近似 | elementwise activation | 使用 `x * sigmoid(1.702 * x)` 做非线性变换 |\n", | ||
| 83 | - "| Dynamic batch | `loop`、`view`、`valid_shape`、`assemble` | batch 维运行时变化时仍能分块处理 |\n", | ||
| 84 | - "\n", | ||
| 85 | - "实践时不要一次性写完。先确认输入输出 shape,再写一个 batch tile 内的计算,最后把结果 assemble 回完整输出。\n" | ||
| 86 | - ] | ||
| 87 | - }, | ||
| 88 | - { | ||
| 89 | - "cell_type": "markdown", | ||
| 90 | - "metadata": {}, | ||
| 91 | - "source": [ | ||
| 92 | - "## 4. 章节实践题\n", | ||
| 93 | "\n", | 30 | "\n", |
| 94 | "题型包含选择题、填空题和编程题。建议先独立完成,再执行下一单元查看参考答案。\n", | 31 | "题型包含选择题、填空题和编程题。建议先独立完成,再执行下一单元查看参考答案。\n", |
| 95 | "\n", | 32 | "\n", |
| @@ -115,7 +52,7 @@ | |||
| 115 | " C. 把输出 Tensor 转成字符串 \n", | 52 | " C. 把输出 Tensor 转成字符串 \n", |
| 116 | " D. 改变 NPU 设备编号\n", | 53 | " D. 改变 NPU 设备编号\n", |
| 117 | "\n", | 54 | "\n", |
| 118 | - "6. (编程题)补全 `dynamic_residual_norm_gelu_kernel`:输入 `x`、`residual_input`、`gamma`、`beta` 和输出 `out`,其中 `x/residual_input/out` 的 shape 为 `[batch, 128]`,batch 是动态维度。要求: \n", | 55 | + "6. (编程题)实现残差连接、LayerNorm、GELU 激活的融合操作。Host 侧代码已给出,只需编写kernel侧 main_dynamic_residual_norm_gelu 方法的代码。 要求: \n", |
| 119 | " - 使用 `pypto.DYNAMIC` 描述动态 batch。 \n", | 56 | " - 使用 `pypto.DYNAMIC` 描述动态 batch。 \n", |
| 120 | " - 使用 `pypto.loop` 按 batch tile 分块。 \n", | 57 | " - 使用 `pypto.loop` 按 batch tile 分块。 \n", |
| 121 | " - 使用 `pypto.view(..., valid_shape=...)` 取当前 tile。 \n", | 58 | " - 使用 `pypto.view(..., valid_shape=...)` 取当前 tile。 \n", |
| @@ -123,11 +60,38 @@ | |||
| 123 | " - 使用 `pypto.assemble` 写回输出,并用 PyTorch reference 验证。\n" | 60 | " - 使用 `pypto.assemble` 写回输出,并用 PyTorch reference 验证。\n" |
| 124 | ] | 61 | ] |
| 125 | }, | 62 | }, |
| 63 | + { | ||
| 64 | + "cell_type": "code", | ||
| 65 | + "execution_count": null, | ||
| 66 | + "metadata": {}, | ||
| 67 | + "outputs": [], | ||
| 68 | + "source": [ | ||
| 69 | + "# TODO 此处添加kernel函数的实现\n", | ||
| 70 | + "\n", | ||
| 71 | + "\n", | ||
| 72 | + "def main_dynamic_residual_norm_gelu():\n", | ||
| 73 | + " for batch in [8, 13]:\n", | ||
| 74 | + " x = torch.randn((batch, HIDDEN_SIZE), dtype=torch.float32, device=device)\n", | ||
| 75 | + " residual_input = torch.randn((batch, HIDDEN_SIZE), dtype=torch.float32, device=device)\n", | ||
| 76 | + " gamma = torch.ones((HIDDEN_SIZE,), dtype=torch.float32, device=device)\n", | ||
| 77 | + " beta = torch.zeros((HIDDEN_SIZE,), dtype=torch.float32, device=device)\n", | ||
| 78 | + " out = torch.empty_like(x)\n", | ||
| 79 | + " dynamic_residual_norm_gelu_kernel(x, residual_input, gamma, beta, out)\n", | ||
| 80 | + " ref = residual_norm_gelu_golden(x, residual_input, gamma, beta)\n", | ||
| 81 | + " max_diff = (out - ref).abs().max().item()\n", | ||
| 82 | + " torch.testing.assert_close(out, ref, rtol=1e-3, atol=1e-3)\n", | ||
| 83 | + " print(f\"batch={batch} 验证通过\")\n", | ||
| 84 | + " print(\"输出 shape:\", tuple(out.shape))\n", | ||
| 85 | + " print(\"最大误差:\", max_diff)\n", | ||
| 86 | + "\n", | ||
| 87 | + "main_dynamic_residual_norm_gelu()" | ||
| 88 | + ] | ||
| 89 | + }, | ||
| 126 | { | 90 | { |
| 127 | "cell_type": "markdown", | 91 | "cell_type": "markdown", |
| 128 | "metadata": {}, | 92 | "metadata": {}, |
| 129 | "source": [ | 93 | "source": [ |
| 130 | - "## 5. 查看答案\n", | 94 | + "## 3. 查看答案\n", |
| 131 | "\n", | 95 | "\n", |
| 132 | "执行以下代码获取参考答案。\n" | 96 | "执行以下代码获取参考答案。\n" |
| 133 | ] | 97 | ] |
| @@ -138,16 +102,16 @@ | |||
| 138 | "metadata": {}, | 102 | "metadata": {}, |
| 139 | "outputs": [], | 103 | "outputs": [], |
| 140 | "source": [ | 104 | "source": [ |
| 141 | - "!cat ./answer/04.07_answer.py\n" | 105 | + "!cat ./answer/04.06_answer.py\n" |
| 142 | ] | 106 | ] |
| 143 | }, | 107 | }, |
| 144 | { | 108 | { |
| 145 | "cell_type": "markdown", | 109 | "cell_type": "markdown", |
| 146 | "metadata": {}, | 110 | "metadata": {}, |
| 147 | "source": [ | 111 | "source": [ |
| 148 | - "## 6. 本章小结\n", | 112 | + "## 4. 本章小结\n", |
| 149 | "\n", | 113 | "\n", |
| 150 | - "第四章从单个组合算子走向模型模块和系统能力:激活函数与 Softmax 训练公式拆解能力,LayerNorm/RMSNorm/FFN 训练归一化与矩阵乘组合能力,动态 shape 和控制流训练分块写回能力,Attention/Transformer 训练多 Tensor 数据流组织能力,Cost Model/ACLGraph 帮助理解系统分析与加速。章节实践把这些能力重新收束到一个可验证的融合算子中,作为进入更复杂项目实践前的检查点。\n" | 114 | + "第四章从单个组合算子走向模型模块和系统能力:激活函数与 Softmax 训练公式拆解能力,LayerNorm/RMSNorm/FFN 训练归一化与矩阵乘组合能力,动态 shape 和控制流训练分块写回能力,Attention/Transformer 训练多 Tensor 数据流组织能力。章节实践把这些能力重新收束到一个可验证的融合算子中,作为进入更复杂项目实践前的检查点。\n" |
| 151 | ] | 115 | ] |
| 152 | } | 116 | } |
| 153 | ], | 117 | ], |
🟡 Medium Priority
变更行:
01.01_chapter_intro.ipynb第 92 行的章节路线表将 1.3 节(CANN)的关注重点描述为"软件栈定位、核心组成、功能特性和安装方式"。但该 PR 同时重写了01.03_CANN.ipynb,将原有的"## 4. 安装与使用方式"整节(含 CANNLab / 主机安装 / Docker 等安装方式说明)替换为"## 4. 架构介绍"(异构计算架构概念、矩阵乘示例、CANN 架构图),不再包含任何安装相关内容。受影响的行为/契约:章节路线表是读者选择阅读顺序和建立预期的依据。当表格标注某节包含"安装方式"而实际内容中完全不存在时,读者会感到困惑,甚至可能以为漏读了内容。
建议:将表格中 1.3 节"关注重点"列的"安装方式"替换为与实际内容匹配的描述,如"异构计算架构"或"关键能力"。
5.作为算子开发者,为什么通常希望使用PyPTO,而不是直接面对底层硬件指令?\n",