在大规模深度学习与科学计算场景中，稀疏张量运算构成了计算负载的重要组成部分。传统的手写内核虽然能够实现极高的性能，但缺乏可移植性且开发成本高昂。近年来，以 TACO（Tensor Algebra Compiler）和 MLIR 稀疏张量方言为代表的稀疏计算编译器取得了显著进展，它们能够将高层爱因斯坦表示法自动降级为高度优化的机器码。本文聚焦于稀疏编译器如何实现自动向量化，并深入探讨硬件亲和性与内存布局优化的关键技术路径。

稀疏编译器的核心抽象

现代稀疏编译器采用 sparsity-agnostic 的前端设计，用户只需使用爱因斯坦求和约定书写张量表达式，例如 C[i][j] += A[i][k] * B[k][j]，编译器即可自动推导稀疏迭代空间并生成相应的数据结构和循环嵌套。这种抽象将算法描述与存储格式解耦，使得同一套源代码可以适配多种稀疏格式，包括 CSR（压缩稀疏行）、CSC（压缩稀疏列）、COO（坐标）以及 CSF（压缩稀疏纤维）等。MLIR 的稀疏张量方言通过 #sparse_tensor.encoding 属性描述每维的稀疏级别类型（dense、compressed、singleton）、维度顺序以及索引位宽，为代码生成器提供了丰富的格式建模能力。

迭代空间理论是稀疏编译的数学基础。以 TACO 为代表的系统将计算表示为非零坐标上的集合表达式，通过逐维协同迭代的方式生成嵌套循环。每一层循环对应一个稀疏维度，循环体仅对非零元素执行计算，从而避免了遍历整个稠密张量的无效开销。然而，这种基于指针追踪的迭代模式给自动向量化带来了根本性挑战 —— 传统的 SIMD 向量化器依赖规则化、单元步幅的内存访问模式，而稀疏迭代往往呈现不规则的间接访问。

自动向量化的关键技术

将稀疏计算转化为可向量化形式的核心思路是引入临时稠密工作区。编译器在稀疏输出的某一维上分配一个稠密缓冲区，将该维度上的更新先收集到工作区中，形成规则的、连续的内存访问模式。完成该维度的所有非零更新后，再将工作区的结果 scatter 回稀疏格式。以稀疏矩阵向量乘法 SpMV 为例，当输出向量在维度 j 上稀疏时，编译器可以创建长度为 j 的工作区向量 W，在内层循环中以向量化方式完成稠密累加，最后将非零结果写回稀疏输出。这种技术在 MLIR 的稀疏编译流程中被称为 workspace-based optimization，已被证明能够显著提升 SIMD 利用率。

循环规范化是另一项重要技术。稀疏迭代的原始形态通常是 while 循环配合指针操作，难以被传统向量化通道识别。编译器通过将稀疏操作降级为具备规则归纳变量的 for 循环，为自动向量化器提供可识别的模式。降级过程中，编译器生成掩码加载（masked load）和掩码存储操作，以处理不规则尾部或非满向量情况。MLIR 的向量方言已经支持可扩展向量（scalable vector），例如 ARM SVE 架构的向量长度可由硬件运行时确定，使得生成的代码能够适配不同的 SIMD 宽度。

阻塞（blocking）和重排序（reordering）策略进一步提升了向量化的局部性。HiCOO（分层压缩稀疏矩阵）和 CSF 等阻塞格式将非零元素组织成固定大小的块，使得内层循环能够处理连续的内存段。维度顺序的重排同样关键 —— 将访问频率最高的维度置于最内层，可以让最热点的内循环具备单元步幅访问，从而充分利用向量加载单元的带宽。BFS-MCS 和 Lexi-Order 等重排序算法通过分析非零模式，减小填充并提升热点子块的空间局部性，直接转化为更高的 SIMD 利用率和多核加速比。

硬件亲和性与内存布局优化

稀疏代码的性能瓶颈往往不在计算本身，而在于内存层次结构的有效利用。数据布局的选择直接影响缓存命中率、TLB 行为以及内存带宽利用率。编译器在选择稀疏编码时需要考虑维度顺序：最内层迭代维度应当对应连续存储的稀疏格式，以获得规则的内存访问模式；外层维度则适合使用压缩存储以节省空间。例如，对 CSR 格式的行优先遍历自然地使内层循环对列索引形成稠密访问，这是 SpMV 在 CPU 上能够高效向量化的根本原因。

平铺（tiling）与缓冲化（bufferization）是现代稀疏编译器的标准优化步骤。通过对维度进行分块，编译器生成恰好容纳于 L2 缓存的子问题，从而最大化数据复用。结合显式的局部缓冲区（workspace、reuse buffer、write buffer），系统能够将离散的稀疏访问转化为批量的顺序写回，显著减少离芯片内存流量。MLIR 的缓冲化通道负责将张量操作降级为 memref 操作，并确保生成的缓冲区具备硬件友好的步幅和对齐属性，为后续的 LLVM 向量化通道提供良好的 IR 形态。

预取策略在稀疏计算中尤为敏感。传统硬件预取器基于固定的访问 stride 进行预测，但对不规则的稀疏模式往往失效甚至适得其反，消耗宝贵的 MSHR（Miss Status Holding Register）资源。研究者发现，在某些不规则稀疏场景下禁用不准确的硬件预取器，反而为软件管理的定制预取留出更多带宽和缓冲资源。这提示编译器后端在生成代码时需要考虑目标硬件的预取行为，甚至为特定稀疏模式插入显式预取指令。

自动调参与配置搜索

由于稀疏编码、维度顺序、平铺粒度、向量化宽度等可调参数组合构成了庞大的搜索空间，静态手工配置难以获得最优性能。MLIR 社区探索了基于机器学习的自动配置器（autoconfigurator），将每种稀疏内核的配置编码为特征向量 —— 包括每维的格式、顺序、并行化 / 向量化层级等，然后使用自动化搜索为特定硬件平台和输入稀疏模式选择最佳配置。这种方法在稀疏 - 稠密矩阵乘法等核心算子上取得了显著提升，验证了配置空间探索的工程价值。

综上所述，稀疏张量编译器的自动向量化与硬件亲和性优化是一个多层次、系统性的工程问题。编译器需要在高层抽象保持格式无关性，同时在后端代码生成阶段精心构造稠密工作区、规范化循环结构、应用阻塞与平铺策略，并通过自动调参机制适配目标硬件。掌握这些技术路径，是实现高性能稀疏计算可移植化的关键。

资料来源