除了应用极广的循环优化,在 AI 编译器底层还存在指令和存储这两种不同优化。
指令优化依赖于硬件提供的特殊加速计算指令。这些指令,如向量化和张量化,能够显著提高计算密度和执行效率。向量化允许我们并行处理数据,而张量化则进一步扩展了这一概念,通过将数据组织成更高维度的结构来实现更大规模的并行计算。这些技术使得算法能够充分利用现代处理器的多核和多线程特性,从而大幅提升性能。
向量化
在之前的循环优化中,已经介绍过了向量化的原理,它是一种数据级并行的优化。其硬件实现如下图所示,将多个连续存储的数据批量加载进向量寄存器中,对整个向量寄存器进行操作,从而同时对多个数据元素进行了计算。
假设两个整数的数组 A 和 B 计算元素和. 并将结果存储到数组 C 当中。在非向量化的代码中,代码形式是这样的:
在向量化的代码中,可能是这样的形式:
要实现加速,需要将其转换为硬件提供的向量化指令,例如:
张量化
在 AI 应用日益广泛的今天,程序运行的数据形式经历了显著的演变。特别是以神经网络为代表的神经网络模型,其内部数据形式为多维矩阵,通常称为张量。例如,在计算机视觉任务中,典型的输入数据具有 N, C, H, W 的维度,其中 N 代表批次大小,C 代表通道数,H 和 W 分别代表图像的高度和宽度。
在神经网络的内部计算过程中,特征图和参与计算的权重(如卷积核)也以类似的 4 维张量形式存在。传统的计算方法,如使用多层循环嵌套逐个计算数据元素,对于神经网络模型而言,效率极其低下。由于神经网络的计算具有高度的数据并行性,同一层内的元素之间几乎没有依赖关系。在理想情况下,如果能在硬件负载允许的范围内一次性将大量数据送入运算单元,将显著提升数据并行性,从而加速计算过程。
针对这一需求,英伟达公司开发了 Tensor Core 技术,专门针对张量计算进行加速。Tensor Core 是一种特殊的硬件单元,设计用于高效执行深度学习中的张量运算,如矩阵乘法和累加操作,这些操作是神经网络模型中的核心组成部分。通过 Tensor Core,GPU 能够在保持精度的同时,大幅度提高张量运算的速度和效率。在 Volta 架构中,一个 SM 由 8 个 FP64 Cuda Cores,16 个 INT32 Cuda Core,16 个 FP32 Cuda Core,和 128 个 Tensor Core 组成,一共有 4 个 SM。
一个典型的 Tensor Core 指令如下所示,它可以在一条指令中计算 C8x8 += A8x32 * B32x8
除了英伟达的 Tensor Core,Intel 也有类似的技术 Vector Neural Network Instructions(VNNI)。利用这些张量指令的一种常见方法是通过调用硬件厂商提供的算子库,如英伟达的 cuBLAS 和 cuDNN,以及英特尔的 oneDNN 等。然而,当神经网络模型中出现新的算子或开发者需要进一步榨取硬件性能时,单纯依赖硬件厂商提供的算子库就显示出其局限性。首先,算子库可能没有包含最新的或自定义的算子,这限制了模型的创新和多样性。其次,算子库的通用优化策略可能无法充分适应特定模型的特性,导致性能提升的空间没有被完全利用。
为了克服这些局限性,研究人员探索了更深层次的优化策略,例如使用自动算子生成工具,如 TVM。在之前的章节我们知道了算子的实现可以分为计算与调度,在调度部分可以使用更加丰富的指令和优化。TVM 可以通过分解一个算子,将其一部分操作映射为硬件实现或者高度优化的手工微内核,从而显著提高性能。我们使用伪代码来描述下如何将大的算子进行拆分,进而直接使用张量化指令。
存储优化关乎于如何高效地管理数据在硬件中的存储和访问。在 AI 芯片硬件中,内存层次结构的设计至关重要。通过优化数据在不同层级内存之间的流动,我们可以减少数据传输的延迟和带宽消耗,从而提升整体的计算效率。例如,GPU 的内存管理策略包括全局内存、共享内存、常量内存等,每种内存类型都有其特定的用途和访问模式。
访存延迟隐藏
在纯串行执行的架构中,计算资源的利用受到显著限制,因为在同一时间点,通常只有一个运算单元处于活跃状态。以程序执行过程为例,其工作流程可概括为以下几个步骤:
这种串行处理模式虽然简单,却无法充分利用现代处理器的并行处理能力,导致计算效率和系统性能受限。
在这种情况下,每个运算部件的时间开销都累加得到最终的时延。为了克服这一局限,现代深度学习系统广泛采用并行计算架构,允许多个运算单元同时工作,显著提高了数据处理速度和整体系统性能。延迟隐藏(Latency Hiding)技术在这一领域得到了广泛的应用。该技术通过将内存操作与计算任务并行化,实现了两者的重叠执行,从而最大化了内存带宽和计算资源的利用效率。通过这种方式,即使在数据加载和写回阶段,也能持续执行计算任务,有效减少了因等待内存操作而产生的空闲时间。
CPU 实现延迟隐藏的过程主要依赖于多线程技术和硬件隐式数据预取机制。在多线程环境中,当一个线程等待内存访问时,CPU 可以切换到另一个线程继续执行计算任务,从而减少 CPU 的空闲时间。此外,现代 CPU 通常具备数据预取单元,能够预测程序接下来可能需要的数据,并提前从内存中加载到缓存中,这样当计算单元需要这些数据时,它们已经准备好了,减少了 CPU 等待内存访问的时间。
GPU 在实现延迟隐藏方面,主要依赖于其高度并行化的架构和先进的调度技术。Wrap Schedule 是一种用于管理多线程执行的技术,它允许 GPU 在等待内存操作完成时,动态地调度其他线程来继续执行计算任务。这种技术通过减少线程间的同步开销,提高了线程的执行效率。GPU 还采用了上下文切换机制,能够在不同的线程上下文之间快速切换,进一步隐藏内存访问的延迟。当一个线程因为内存访问而暂停时,GPU 可以立即切换到另一个准备好的线程继续执行,从而保持了 GPU 核心的持续工作。
NPU 采用了解耦访问/执行(Decoupled Access/Execute,DAE)架构。在 DAE 架构中,内存访问操作和计算操作是分开进行的,允许它们并行执行而不是顺序依赖。NPU 拥有专门的硬件单元来处理数据的加载和存储,这些单元独立于执行计算的核心。当计算核心需要数据时,它会发送请求,然后继续执行其他计算任务,而数据加载操作在后台进行。在这种情况下,一般需要使用双缓冲机制,来缓存不同 LOAD 指令得到的数据。
在这种模式下,执行指令会变为并行方式:
存储分配
从传统编译器的视角来看,内存被划分为几个关键区域,每个区域都有其特定的用途和生命周期。
在 AI 系统中,这种视角下的内存管理显然无法支撑起 AI 应用。AI 系统通常需要处理大量的数据和复杂的算法,这就需要高效的内存分配和回收策略来支持它们的运行。专用硬件如 GPU 和 NPU 具有各自独特的内存管理机制,这些机制针对它们处理任务的特点进行了优化。
GPU 的内存管理机制包括: