张量形状与数据类型

构建计算图的结构，为机器学习 (machine learning)模型提供了骨架。为了充实这个骨架，编译器需要关于沿边流动的数据的准确信息。深度学习 (deep learning)编译器要求这些数据具有严格的类型和形状。与变量可以动态改变类型的标准 Python 代码不同，中间表示 (IR) 将张量形状和数据类型 (dtypes) 视为严格的约定。这些约定使得编译器能够在代码实际运行之前很久，就计算内存需求并选取正确的硬件指令。

数据类型的意义

张量的数据类型决定了数值在内存中的存储方式以及硬件处理单元如何对其操作。数据科学家可能使用泛型“浮点数”的观念，但编译器必须精确。它需要了解一个值是32位浮点数（FP32）、16位半精度浮点数（FP16），还是8位整数（INT8）。

这种精确性对以下两个主要原因而言不可或缺：

1. 内存带宽： 移动数据通常比计算数据耗费更多。一个 FP16 张量占用的内存带宽是 FP32 张量的一半。编译器使用数据类型信息来优化内存访问模式。
2. 指令选取： GPU 或 TPU 等硬件加速器通常具有针对特定类型的专用单元。例如，NVIDIA Tensor Cores 主要对 FP16 或 BF16 输入进行操作。如果 IR 指定 FP32，编译器可能会插入类型转换操作或回退到较慢的通用 CUDA 核心。

在 IR 中，表示加法操作的节点并非泛型。它明确是 add.f32 或 add.i8。如果一个框架生成一个图，其中 FP32 张量与 INT32 张量相加，编译器的第一阶段，即类型推断，必须检测到这种不匹配。它将报错或插入一个明确的 cast 节点来对齐 (alignment)类型，从而确保生成有效的机器码。

张量形状与秩

除了数据类型之外，张量的形状是其最具定义性的属性。形状定义了维度（秩）和每个维度的范围。

• 秩： 维度的数量。标量秩为0，向量 (vector)秩为1，一批图像通常秩为4（批量、通道、高度、宽度）。
• 维度： 每个轴的具体大小。

编译器依靠形状信息进行静态内存分配。如果编译器得知一个张量的形状是 $(32, 128)$ 且使用 float32（4字节），它能预先算出该缓冲区恰好需要 $32 \times 128 \times 4 = 16,384$ 字节。这使得生成的二进制文件能够在堆或栈上高效分配内存，而无需在执行期间进行 malloc 等动态内存管理调用的额外开销。

形状推断

编译器前端的主要职责之一是形状推断。这是将形状信息从输入节点通过整个图传播，以确定每个中间张量和输出张量形状的过程。

思考一个卷积操作。输出形状并非任意；它是输入形状、权重 (weight)形状、填充、步长和膨胀的确定性函数。编译器在其分析阶段实现这些公式。

计算卷积中输出空间维度的标准公式是：

$O = \left\lfloor \frac{I - K + 2P}{S} \right\rfloor + 1$

其中：

• $O$ 是输出大小
• $I$ 是输入大小
• $K$ 是核（滤波器）大小
• $P$ 是填充
• $S$ 是步长

通过将此逻辑应用于有向无环图（DAG）中的每个节点，编译器构建了张量大小的完整映射。

形状推断传播的示意图。编译器根据输入维度和操作符属性计算输出几何形状。

如果推断出的形状不兼容，例如尝试执行维度为 $(M, K)$ 和 $(J, N)$ 的矩阵乘法，其中 $K \neq J$，编译器会将此识别为结构错误。此验证可保护运行时免受无效内存访问导致的崩溃。

处理广播

NumPy 和 PyTorch 等框架允许隐式广播，其中一个较小的张量在逐元素操作期间自动扩展以匹配较大张量的形状。例如，将标量偏置 (bias)添加到向量 (vector)。

尽管对用户来说很方便，但隐式行为对低级代码生成存在问题。编译器通常会执行一个阶段以使这些广播明确。它可能会在 IR 中插入特定的 broadcast 或 expand 节点。

明确的广播确保后端代码生成器明确知道如何处理内存步长。编译器不会物理复制数据来扩展张量（这会浪费内存），而是使用零步长内存访问模式，在遍历较大张量的维度时重复读取相同的值。

布局敏感性

IR 中张量的形状是数学上定义的，但其在内存中的物理布局可以不同。用于图像的 4D 张量通常表示为 NCHW（批量、通道、高度、宽度）或 NHWC（批量、高度、宽度、通道）。

• NCHW： 通常受使用 cuDNN 的 GPU 实现偏好。
• NHWC： 通常受 CPU 后端和 TPU 处理单元偏好，以实现更好的向量 (vector)化。

IR 中“形状”的定义意味着维度的逻辑顺序。然而，优化编译器通常会透明地重写这些布局。我们将在第三章的“内存布局转换”部分讨论这种转换的机制。目前，只需明白 IR 中的形状元数据是数据逻辑组织的事实来源，无论字节在 RAM 中如何物理排列。