3. 数据布局与表示 (Data Layout & Representation)

当计算逻辑 (第1、2章) 确定后，性能的瓶颈往往转移到数据的物理组织形式上。

本章关注张量数据在内存地址空间中的排布与访问模式。核心目标是使数据的物理布局与硬件的访问特性 (如 SIMD 通道、Cache Line) 相匹配，以最大化空间局部性 (Spatial Locality)。编译器需要从全图范围选择最优的逻辑布局 (如 NCHW vs NHWC)，并在微观层面执行数据打包 (Packing)、对齐 (Padding)和缓冲区化 (Bufferization)，以消除因格式不匹配导致的昂贵的数据重排开销。

数据布局优化是深度学习编译器的核心问题之一。TVM[1]、Glow[20]、XLA等主流编译器都投入大量研究来优化张量的内存布局。ALT[2]等最新研究表明，联合优化数据布局和循环变换可带来显著性能提升。

3.1 全局布局优化 (Global Layout Optimization)

3.1.1 Layout Propagation & Assignment (布局传播与指派)

背景

不同算子对内存布局有不同偏好 (例如：Conv2d 在 GPU 上偏好 NHWC，在 NPU 上可能偏好 NC1HWC0)。如果在图中频繁插入 Transpose 或 Reshape，数据搬运开销将抵消计算收益。

布局传播旨在为整个子图选择统一的最佳布局。TVM[1]、Glow[20]、VTensor[3]等编译器都实现了布局传播优化来减少格式转换开销。

技术原理

布局优化本质上是一个全局约束满足问题 (Global Constraint Satisfaction Problem)或最小代价路径搜索问题[1,2,4]。其技术实现通常包含三个步骤：

1. 约束定义与硬件亲和性 (Hardware Affinity)： 编译器首先标记每个算子对布局的"偏好"[1,4,20,21]：

   • GPU Tensor Core：偏好 NHWC (Channel-last)。原因在于 Tensor Core 计算矩阵乘时，需要 Inner Dimension (通常是 Channel) 在内存中连续，以实现合并访问 (Memory Coalescing) 并利用向量化加载指令 (如 LDG.128)[16,17]。
   • Ascend Cube Core：偏好 Fractal (NC1HWC0/5D)。这是为了适配矩阵单元内部的 16x16 分块读取逻辑[21,22]。
   • CPU Vector Unit：可能偏好 NCHWc (Blocked Channel)，配合 SIMD 宽度[25]。

2. 双向传播 (Bi-directional Propagation)： 从对布局敏感的"锚点算子"(如 Conv) 开始 (Source-Sink Propagate)，编译器在计算图上执行前向和后向遍历，传播布局约束。

   • Forward：输入数据的布局决定后续算子的布局 (如 Input 是 NHWC，则 Conv 也选 NHWC)。
   • Backward：最终输出的要求反推前序算子 (如 Output 要求 NCHW，则最后一个 Conv 最好输出 NCHW)。
   • Layout Transform Elimination：利用数学性质消除冗余转换 (例如：Transpose (Transpose (x)) == x)。
   • 冲突解决：当一个算子的输入来自 NHWC，但自身强制要求 NCHW 时，传播受阻。

3. 代价最小化求解 (Cost Minimization)： 当发生冲突时，编译器需要在图的边 (Edge) 上插入 Layout Transform (Transpose/Permute) 算子。 优化目标是最小化总代价：

   编译器通常使用动态规划或贪心算法 (如 Union-Find) 来决定在何处“切一刀”插入转置，使得转置次数最少且算子效率最高。

MLIR 实现

在 MLIR 中，这通常发生在 tosa 到 linalg 的降级过程中，或者通过专门的 LayoutOptimization Pass 实现。

// 原始图：未指定布局的 Conv2D -> Relu -> MaxPool
// 默认可能是 NCHW，但目标硬件(GPU)偏好 NHWC

func.func @layout_opt(%input: tensor<1x3x224x224xf32>) {
  
  // 1. [Layout Assignment] 
  // 编译器分析发现 Conv2D 在 NHWC 下能使用 TensorCore，收益极高
  // 因此决定插入 Layout Transform，并将后续算子全部染成 NHWC
  
  // 插入 NCHW -> NHWC
  %input_perm = tosa.transpose %input {perms = [0, 2, 3, 1]} 
              : (tensor<1x3x224x224xf32>) -> tensor<1x224x224x3xf32>

  // 2. [Propagation]
  // Conv2D 选用 NHWC 变体
  %conv = linalg.conv_2d_nhwc_hwcf ins(%input_perm, %weight) ...
  
  // ReLU 自动适配 NHWC (Element-wise 对布局不敏感，随大流)
  %relu = linalg.generic ... ins(%conv) ... 
  
  // MaxPool 选用 NHWC 变体
  %pool = linalg.pooling_nhwc_max ... ins(%relu) ...

  // 3. [Finalize] 如果需要 NCHW 输出，则在最后转回；否则直接输出 NHWC
  return %pool
}

补充说明 (Ascend NPU)

在 Ascend 场景下，布局传播尤为关键。如果传播失败，导致在计算图中频繁插入 TransData(格式转换算子)，会严重阻塞流水线。

因此 Ascend 的编译器会极力推导 DefaultFormat 为 5D (Fractal)，并让尽可能多的算子 (包括 ReLU, Add 等) 直接在 5D 数据上运行，实现 Layout-agnostic Fusion。

---

3.2 数据打包与微布局 (Data Packing & Micro-layout)

3.2.1 Tensor Packing (Block Layout / Fractal Format)

背景

标准的 Row-major (行优先) 或 Column-major (列优先) 布局在处理 2D/3D 卷积或矩阵乘法时，往往无法提供最佳的空间局部性。

Tensor Packing (数据打包) 指将高维张量的逻辑维度物理重排为嵌套的块状结构 (Blocked Structure)。这种变换将逻辑上相邻的二维子矩阵 (Tile) 存储在物理连续的内存地址中[6,7,8]。

K车材等人[8]的经典研究表明，分块数据布局可显著改善Cache和TLB性能。现代张量编译器[1,2,6,7]将这一技术进一步发扬光大。

技术原理

1. 最大化缓存行利用率 (Cache Line Utilization)：

   问题：在 Row-major 中，访问矩阵的一列 (Column stride) 会导致巨大的内存跳跃，极易引发 Cache Miss。

   原理：Packing 将一个 的 Tile 里的数据连续存放。当 GPU/NPU 读取 Tile 的第一个元素时，整个 Cache Line (通常 64 Bytes、128 Bytes 或 256 Bytes) 加载进来的数据正好是该 Tile 的后续元素。这使得二维的空间局部性转化为了物理内存的一维连续性。

2. 向量化加载友好 (Vectorization Friendliness)：

   SIMD/SIMT 单元通常一次加载 128/256/512 位。Block Layout 保证了数据的Inner Dimension (最内层维度) 长度固定且对齐 (例如总是 16 或 32)，这允许编译器直接生成对齐的向量加载指令 (Aligned Vector Load)，无需处理边缘情况。

硬件视角：Ascend NPU 的分形格式 (Fractal Format)

在 Ascend NPU (达芬奇架构) 上，Tensor Packing 不仅仅是优化，更是驱动 Cube Unit (矩阵加速单元) 的前置条件。

1. 分形格式 (Fractal Layout - ZnZ/NC1HWC0)： Ascend 的 Cube Unit 硬件上只能处理 的微矩阵 (Fractal Block)。

   • NC1HWC0: 针对 Activation (Feature Map)，逻辑上的 [N, C, H, W] 被物理重排为 5D 格式 [N, C1, H, W, C0]。其中 C0 固定为 16 (FP16)，C1 = C / 16。
   • ZnZ (Z-in-Z): 针对 Weight (权重矩阵)。逻辑上的 2D 矩阵被重排为小“Z”字型嵌套大“Z”字型的结构，以确保 Cube 在计算矩阵乘时，内存读取是连续的。

2. 强制对齐与补零 (Alignment & Padding)： 如果逻辑维度不能被 16 整除，编译器必须在 Packing 阶段隐式补零 (Padding)。例如，Channel=3 的 RGB 图片在 Ascend 上物理占用实际上是 Channel=16 的空间 (C0=16)，虽然浪费了 13/16 的空间，但换取了 Cube 单元的极致吞吐。

MLIR 实现：tensor.pack / tensor.unpack

MLIR 引入了 tensor.pack 和 tensor.unpack 算子来显式表达这种布局变换。在 Ascend 编译流中，这通常对应于将 linalg 降级为 NPU 专用 Dialect 时的步骤。

// 场景：为了适配 Ascend Cube Unit (16x16 FP16)，将 Row-major 转换为 Block Layout
// 逻辑形状：[1024, 1024]
// 物理形状：[64, 64, 16, 16] (Outer_M, Outer_N, Inner_m, Inner_n)

func.func @ascend_packing(%input: tensor<1024x1024xf16>) -> tensor<64x64x16x16xf16> {
  
  // 定义 Packing 策略
  // inner_dims_pos = [0, 1] 表示对两个维度都进行切分
  // inner_tiles = [16, 16] 对应 Cube Unit 的硬件限制
  %packed = tensor.pack %input
      inner_dims_pos = [0, 1]
      inner_tiles = [16, 16]
      : tensor<1024x1024xf16> -> tensor<64x64x16x16xf16>

  return %packed
}

// 融合说明：
// 在实际编译中，编译器会尝试将这个 tensor.pack 操作向上融合 (Fuse Up)
// 到产生 %input 的算子(如 Previous Conv/MatMul)的 Epilogue 中，
// 使得前一个算子直接写出 16x16 的分形格式，避免单独的内存搬运。

3.2.2 Swizzled Layout (布局重排)

背景

在现代加速器 (比如 GPU 和 NPU) 中，片上高速缓存 (Shared Memory / L1 / UB) 通常被划分为多个存储体 (Banks) (例如 32 个 Bank)。当一个 Warp/Wavefront 中的多个线程同时访问不同地址，但这些地址映射到同一个 Bank 时，就会发生 Bank Conflict (存储体冲突)。硬件必须将这些访问串行化 (Serialization)，导致有效带宽成倍下降[11,12,13]。

技术原理

Swizzling (地址重排/混洗) 是一种通过数学变换打乱数据在 Shared Memory 中物理存放顺序的技术，目的是让逻辑上连续的访问在物理上分散到不同的 Bank 中[11,12,13]。这是CUDA和ROCm编程中的标准优化技术。

1. XOR Swizzling (异或混洗)： 这是目前最主流的方案 (NVIDIA Ampere+ 硬件原生支持)。

   • 冲突场景：假设 stride=32，Bank 映射通常是 Address % 32。如果线程 访问 A[0][0]， 访问 A[1][0] (地址 32)，它们都会命中 Bank 0。
   • 解法：引入 XOR 映射。
   • 通过让列索引与行索引的高位进行异或，原本每一列都映射到 Bank 0 的局面被打乱，变为对角线式分布，从而消除冲突。

2. Vectorized Access Optimization： Swizzled Layout 还是为了适配硬件特殊的加载指令 (如 NVIDIA 的 ldmatrix)。这些指令要求数据以特定的“交错”格式存放，以便一个 Warp 能够用一条指令加载一个 或 的矩阵块到 Tensor Core 寄存器。

3. Ascend NPU (Implicit Swizzle)：

   在 Ascend 架构中，Swizzling 通常由 MTE (Memory Transfer Engine) 在数据搬运过程中自动完成。当数据从线性内存 (如 UB) 搬运到分形内存 (如 L0 Cube Buffer) 时，硬件电路会自动进行地址重排，对软件透明。

MLIR 示例

MLIR 通过 nvgpu 方言支持显式的 GPU Swizzling，而对于 Ascend，通常在 Bufferization 后的 Copy 阶段隐式处理。

// 场景：将数据从 Global 拷贝到 Shared，启用 Swizzle 以优化 Tensor Core 读取
func.func @gpu_swizzle_copy(%global: memref<?xf16>, %shared: memref<?xf16, #gpu.address_space<workgroup>>) {
  // swizzle 属性指示编译器生成 XOR 地址计算逻辑
  %token = nvgpu.device_async_copy %global[%i] to %shared[%j] 
           dst_elements = 8 
           {bypass_l1, src_elements = 8} 
           : memref<?xf16> to memref<?xf16, #gpu.address_space<workgroup>>
  
  // 后续使用 ldmatrix 读取时，硬件会自动解码 Swizzled 地址
}

3.3 填充与对齐 (Padding & Alignment)

3.3.1 Dimension Padding (维度填充)

背景

真实世界的模型 Tensor 维度往往是不规则的 (如 Channel=3 或 Vocab=30522)。然而，高性能硬件单元通常有严格的对齐要求 (Alignment Requirements)。如果不满足这些要求，硬件效率会大幅下降甚至无法运行。

Intel的向量化指南[18]明确指出，数据对齐对于SIMD指令的性能至关重要。MLIR社区[19]也广泛讨论了Padding在向量化中的作用。

技术原理

Padding (填充) 的核心价值在于以少量的空间浪费换取控制流简化和硬件兼容性。

1. 消除 Loop Tail (通用GPU)： 如果矩阵尺寸不能被 Tiling Size 整除，编译器需要生成主循环和低效的余数循环 (Loop Tail)。通过 Padding 将尺寸补齐 (如补齐到 128 的倍数)，编译器可以生成单一的完美循环，消除分支跳转，利于流水线满载。

2. 强制 C0 对齐 (Ascend NPU 硬约束)： 达芬奇架构的 Cube Unit 物理上是一个 的阵列。它无法处理非 16 倍数的维度 (如 Channel=3)。

   • On-the-fly Padding (随路填充)：Ascend 编译器利用 MTE 引擎支持“搬运时填充”。从 HBM 读取 3 个数，MTE 在写入片上缓存时自动补 13 个 0。这实现了零带宽开销 (HBM 带宽只消耗了 3 个数的量，虽然片上占了 16 个数的空间)。

3. 避免 Stride Conflict (GPU)： 如果矩阵的 Stride 恰好是 Shared Memory Bank 数量的整数倍，会导致列访问冲突。Padding Stride (增加一个哑元) 可以错开 Bank 映射。

MLIR 实现

tensor.pad 操作用于显式填充，通常在 Tiling 之前或之后进行，以确保每个 Tile 都是完整的 (Full Tile)。

// 场景：处理输入 Channel=3 的卷积，适配 Ascend C0=16 要求
func.func @ascend_padding_c0(%img: tensor<1x3x224x224xf16>) -> tensor<1x16x224x224xf16> {
  
  // 显式 Padding：逻辑上将 C=3 扩展为 C=16
  // 这告诉编译器：计算循环的边界是 16，不是 3
  // low/high 指定了在维度前后填充 0 的数量
  %padded = tensor.pad %img low[0,0,0,0] high[0,13,0,0] {
    ^bb0(...):
      %c0 = arith.constant 0.0 : f16
      tensor.yield %c0 : f16
  } : tensor<1x3x... > to tensor<1x16... >

  // 后续 Packing 会将这个 Padded Tensor 重排为 NC1HWC0
  // 后端优化：识别到 Padding 为全 0，生成 MTE 指令参数 src_stride=3, dst_stride=16
  return %padded
}

---

3.4 缓冲区化与原地更新 (Bufferization & In-place)

3.4.1 One-Shot Bufferization

背景

在编译器的中高层 IR (如 MLIR Linalg/Tensor 级别)，数据通常被表示为不可变张量 (Immutable Tensors)，遵循 SSA (静态单赋值) 原则。这便于数学分析，但与底层硬件的冯·诺依曼架构 (可变内存、指针读写) 不符。

Bufferization 的任务是将 Tensor 转换为 Memref (内存缓冲区)[26,27]：

• Naive Bufferization：为每个算子的输出都 malloc 新内存。结果：内存爆炸，拷贝频繁。
• One-Shot Bufferization：对整个函数进行全局分析，寻找 原地更新 (In-place Update) 的机会，即复用输入 Buffer 来存储输出结果[27,28,29]。

MLIR的One-Shot Bufferization[27]是现代编译器缓冲区优化的代表性工作，TeMCO[28]等研究进一步推进了跨操作的内存优化。

技术原理

1. Destination-Passing Style (DPS, 目标传递风格)： 这是现代 Bufferization 的核心范式。算子不再“返回”一个新的 Tensor，而是接收一个“输出 Buffer”作为参数，并将结果写进去。

   • Tensor IR: %result = op (%input)
   • Buffer IR: op (%input, %output_buffer)

2. RaW Conflict Detection (读写冲突检测)： 编译器必须证明复用是安全的。

   • 冲突场景：如果 Tensor A 被 Op1 读取，同时被 Op2 写入 (复用)。如果 Op2 先于 Op1 执行，Op1 就会读到错误的数据。
   • 分析算法：构建 Interference Graph (干涉图)。如果输入 Tensor 在写入点之后不再被读取 (Last Use)，则标记为可复用；否则，插入显式的 memref.copy 进行保护。

3. 硬件视角：

   • GPU：Bufferization 通常映射为动态的 malloc/free 或显式的 Shared Memory 分配。
   • Ascend NPU：由于 NPU 运行时 (Runtime) 通常不支持高效的动态内存分配，Bufferization 是后续 Static Memory Planning 的前置步骤。编译器倾向于极度激进的 In-place 复用，以减少 Workspace 的总需求量，确保模型能塞进有限的 Device Memory。

MLIR 实现

MLIR 通过 bufferization dialect 和 one-shot-bufferize pass 实现这一过程。

// 转换前：Tensor Level (Value Semantics)
// 原始图：SSA 形式，%t1 是不可变的
func.func @tensor_calc(%A: tensor<1024xf32>, %B: tensor<1024xf32>) -> tensor<1024xf32> {
  
  // 1. Init Tensor (逻辑上的空张量)
  %init = tensor.empty() : tensor<1024xf32>
  
  // 2. Generic Op (DPS 风格: outs 指定了潜在的复用目标)
  %t1 = linalg.generic { ... } 
      ins(%A, %B : ...) 
      outs(%init : ...) { ... }
  
  return %t1
}

// 转换后：Memref Level (In-place Bufferization)
// 优化后：指针形式，In-place 更新
func.func @buffer_calc(%A: memref<1024xf32>, %B: memref<1024xf32>, %Out: memref<1024xf32>) {
  
  // 编译器分析发现 %init 只是占位符，且 %Out Buffer 在此处可写
  // 因此，直接将计算结果写入传入的 %Out 指针，完全消除了中间 malloc
  
  linalg.generic { ... }
      ins(%A, %B : memref<1024xf32>, memref<1024xf32>)
      outs(%Out : memref<1024xf32>) { ... }
  
  // 函数无返回值 (void)，结果副作用在 %Out 中
  return
}

性能影响案例

• Case: Y = ReLU (X)
• Copying: Alloc (Y); Load (X); Compute; Store (Y); Free (Y)。带宽消耗 。
• In-place: Load (X); Compute; Store (X). 带宽消耗 (读写同一地址，Cache 命中率极高)，且省去了 Alloc/Free 的系统调用开销 (约 5-10us)。

---

参考文献 (References)

数据布局优化 (Data Layout Optimization)

[1] Tianqi Chen, et al. "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning." USENIX OSDI, 2018. [Link]

[2] (CGO 2023). "ALT: Breaking the Wall between Data Layout and Loop Transformations." ACM CGO, 2023. [Link]

[3] (Journal 2023). "VTensor: Using Virtual Tensors to Build a Layout-Oblivious DNN Compiler." 2023. [Link]

[4] Y. Li, Z. Liu. "The Deep Learning Compiler: A Comprehensive Survey." arXiv, 2020. [Link]

[5] (Electronics 2025). "An Improved Strategy for Data Layout in Convolutional Neural Networks." Electronics, 2025. [Link]

数据打包与分块布局 (Tensor Packing & Block Layout)

[6] M. Moroz. "Writing an Optimizing Tensor Compiler from Scratch." 2024. [TensorFrost]

[7] (PLDI 2024). "A Tensor Compiler with Automatic Data Packing for Simple FHE." ACM PLDI, 2024. [Link]

[8] K. 车, et al. "Tiling, Block Data Layout, and Memory Hierarchy Performance." IEEE TC, 2003. [Link]

[9] (ACM 2024). "Optimizing Tensor Train Decomposition in DNNs for RISC-V." ACM, 2024. [Link]

[10] (PyTorch Blog 2024). "Accelerated PyTorch Inference with torch.compile on AWS Graviton." PyTorch Blog, 2024. [Link]

Swizzling与Bank Conflict (Swizzling & Bank Conflict)

[11] L. Mao. "CUDA Shared Memory Swizzling." Lei Mao's Log Book. [Tutorial]

[12] (ROCm Blog 2025). "Avoiding LDS Bank Conflicts on AMD GPUs Using CK-Tile." ROCm Blog, 2025. [Link]

[13] (Blog 2025). "Flash Attention from Scratch Part 4: Bank Conflicts & Swizzling." 2025. [Link]

[14] S. Markidis, et al. "NVIDIA Tensor Core Programmability, Performance & Precision." arXiv, 2018. [Link]

[15] J. Yan, et al. "Demystifying Tensor Cores to Optimize Half-Precision Matrix Multiplication." IEEE IPDPS, 2020. [Link]

填充与对齐 (Padding & Alignment)

[16] (PACMPL 2024). "A Verified Compiler for a Functional Tensor Language." PACMPL, 2024. [Link]

[17] (Intel 2021). "Data Alignment to Assist Vectorization." Intel Developer Guide, 2021. [Link]

[18] (MLIR Discourse 2021). "Padding for Vectorization - MLIR Discourse." LLVM/MLIR, 2021. [Discussion]

[19] (arXiv 2025). "Generate Vectorized Codes for General Tensor Permutation." arXiv, June 2025. [Latest]

[20] (IACR 2025). "A Tensor Compiler for Autovectorizing Homomorphic Encryption." IACR, 2025. [Latest]

缓冲区化与原地更新 (Bufferization & In-place)

[21] (MLIR Docs). "Bufferization - MLIR Documentation." MLIR. [Official]

[22] (MLIR Open Meeting 2022). "One-Shot Function Bufferization of Tensor Programs." MLIR Open Meeting, 2022. [Link]

[23] (ACM 2024). "TeMCO: Tensor Memory Compiler Optimization across Multiple Operations." ACM, Aug 2024. [Link]

[24] (MLSys 2023). "Safe Optimized Static Memory Allocation for Parallel Deep Learning." MLSys, 2023. [Link]

[25] (APXML Course). "Static Memory Planning for ML Graphs." Compiler Course. [Link]

硬件特定布局优化 (Hardware-specific Layout Optimization)

[26] N. Rotem, et al. "Glow: Graph Lowering Compiler Techniques for Neural Networks." arXiv, 2018. [Link]

[27] (CGO 2021). "AKG: Automatic Kernel Generation for Neural Processing Units." ACM CGO, 2021. [Link]

[28] (USENIX ATC 2025). "Accelerating Model Training on Ascend Chips." USENIX ATC, 2025. [Latest]

[29] (arXiv 2024). "Ascend HiFloat8 Format for Deep Learning." arXiv, Sept 2024. [Link]

[30] (ACM 2009). "A SIMD Optimization Framework for Retargetable Compilers." ACM, 2009. [Link]

综述与系统论文 (Surveys & Systems)

[31] (ICOMP 2023/2024). "Compiler Technologies in Deep Learning Co-Design: A Survey." ICOMP, 2023/2024. [Link]

[32] (GitHub). "Awesome Tensor Compilers." GitHub Repository. [Collection]

[33] H. Anderson, et al. "Learning to Optimize Halide with Tree Search and Random Programs." arXiv, 2019. [Link]

[34] J. Zheng, et al. "FlexTensor: An Automatic Schedule Exploration Framework." 2020. [Link]

[35] (arXiv 2025). "Linear Layouts: Robust Code Generation of Efficient Convolution Kernels." arXiv, May 2025. [Latest]