【AI 编译·I】Mosaic: Exploiting Instruction-Level Parallelism on Deep Learning Accelerators with iTex Tessellation

本文选自 ASPLOS’25，计算所国重实验室郭崎组的论文。作者单位包括腾讯，寒武纪。元信息见最后。

概述

这篇论文提出了 Mosaic，一种面向深度学习加速器（DLA）的 自底向上、基于镶嵌（tessellation）的编译器框架。论文的动机在于当前主流深度学习编译器大多采用自顶向下的平铺（tiling）方法，只能生成同构指令映射，难以充分挖掘指令级并行性（ILP），尤其在 CUDA Core 与 Tensor Core 协同场景中表现受限。
Mosaic 的核心创新包括：

提出 iTex 抽象，形式化描述单条指令的计算语义；
构建派生与变换 iTex，从而实现任务空间与指令空间之间的灵活映射；
提出启发式镶嵌算法，结合 independent 与 reduction tessellation 实现异构指令映射；
在代码生成与调度阶段，通过 zig-zag 执行模式 和 指令交错（interleaving） 提高 ILP 暴露度。
方法评估涵盖多种 GEMM/BMM 算子、不同数据精度与多个真实大语言模型（如 Llama, Phi），并在 NVIDIA、AMD GPU 和 Intel CPU 上实现端到端评估，结果显示 Mosaic 可达或超越 Tensor Core 理论峰值吞吐。
局限性方面，Mosaic 当前调度能力仍受限于厂商 PTX 编译器，且在硬件资源有限的场景下仍可能受到寄存器压力和发射冲突的影响。此外，该框架主要面向 GEMM 类算子，对控制流密集任务的适应性仍待进一步扩展。

论文翻译与注解

摘要： 深度学习在诸多应用领域取得了显著成功，但其背后伴随着极高的计算复杂度。为满足日益增长的计算需求，通用硬件平台（如 CPU 与 GPU）提供了丰富的计算资源，包括标量、向量与张量单元，这些资源可并行执行深度学习任务。然而，现有自顶向下、基于划分的深度学习编译器通常采用同质化的方式将张量计算任务映射至硬件算术指令，难以同时调度不同类型的计算单元，从而限制了性能的进一步提升。为解决这一问题，本文提出 Mosaic 编译器，一种自底向上的基于镶嵌（tessellation）策略的深度学习编译框架。该方法直接将给定的张量计算任务以多样化指令进行镶嵌，从而构建异构的指令到任务映射关系，以充分挖掘不同计算单元间的指令级并行性（ILP）。该框架的核心在于提出的 iTex 抽象，它通过形式化的仿射函数建模指令操作与语义之间的映射关系。在 iTex 抽象的基础上，本文进一步提出了一种启发式方法，用于高效生成多样化的镶嵌方案。此外，本文还引入了 iTex 调度技术，以协调不同指令的执行，从而降低结构性冲突风险，并最大化可利用的指令级并行性。实验评估结果表明，Mosaic 在多个硬件平台上相较于高度优化的厂商库实现平均加速比为 1.08× 至 1.28×；在针对从大语言模型（LLM）中提取的实际算子进行测试时，相较于已有最佳基线系统实现平均加速比为 1.34×。更重要的是，Mosaic 可实现最高达 GPU 张量核心理论峰值吞吐率的 106%，充分展示了其在指令级并行性挖掘方面的高效性。

①计算单元：☞标量、向量、张量单元，参考超标量处理器设计。
②镶嵌 (tessellation)：是 tiling 的泛化版本，不必均匀，可以是异构单元，可以是不同类型的硬件指令对应的任务块。在本文中☞将一个张量计算任务空间划分为多个 iTex 实例（即单条硬件指令能完成的子任务），以实现空间上的完全覆盖。
> 即把整个计算任务像拼图一样拆成一块块，每块交给一条能胜任的硬件指令去执行。如何拼得好，决定了效率有多高。
③Tiling（平铺） 是一种在**编译优化、并行计算和张量程序编译器**中非常常见的技术，用于提高局部性（locality）和并行性（parallelism），尤其在矩阵运算、卷积、张量乘法等大规模数据计算中被广泛使用。Tiling 是将大规模计算任务划分为多个较小、规则的“块”，每个块可以被单独执行，从而提升性能。

引言

近年来，深度学习取得了显著进展，尤其是大语言模型（LLMs）的兴起，在多个应用领域中展现出卓越的性能。这一算法演进过程的主要特征是其计算需求持续增长，同时呈现出显著的异构性。具体而言，这类算法融合了多种计算特性，包括标量计算（如控制流）、向量计算（如 Softmax 与 ReLU 等操作），以及张量计算（如卷积与矩阵乘法等高计算密度任务）。
为支持标量、向量及张量等多种类型的计算，现代深度学习硬件广泛采用在单一核心中集成多种计算单元的设计。两种典型的深度学习加速器（DLA）架构包括：
（1）自 Volta 架构起的 NVIDIA GPU，在其流式多处理器子分区（SMSP）中集成了多个被称为 CUDA 核心的向量单元以及名为 Tensor Core 的张量单元；
（2）第四代 Intel Xeon Scalable 处理器，在单个处理器核心内部集成了高级向量扩展（AVX） 单元与高级矩阵扩展（AMX）单元。此类设计使得在不同计算单元上并行执行计算成为可能，从而显著提升整体计算性能。
关于各类深度学习加速器在计算单元、对应指令集及理论峰值吞吐量等方面的详细比较，可参考表 1（略）。然而，当前主流深度学习编译器（包括 TVM、MLIR 与 Triton）受限于自顶向下、基于划分的编译策略，尚未能有效调度并利用不同类型的计算单元。
更具体地说，这些编译器通常将计算表示为嵌套循环结构，并通过一系列循环变换将原始计算任务划分为适配单一类型硬件指令的均匀计算块（tiles），由此形成一种同质的指令与任务之间的映射关系，如图 1(a) 所示。该同质映射模式限制了内核仅使用某一类型的硬件指令完成给定的张量计算任务，导致部分计算单元处于高度负载状态，而其他计算单元则处于空闲，从而未能充分发挥系统整体的计算能力。

看到这里知道本文讲的是，同一计算单元当前只能完成一种类型的指令计算。倘若这个内核“喜欢”某种指令，一直占着其他计算单元，则其他计算单元空缺，发挥不了计算能力。
①“自顶向下（Top-down）” 是一种思维方式和系统设计策略，在计算机科学、编译器设计、算法规划等领域中广泛使用。自顶向下就是从全局开始，逐层细化，最后落实到具体实现。在本文☞论文批评的是当前深度学习编译器普遍采用的 top-down + tiling-based 方法：这些方法从任务的算子级别（如 GEMM）出发，将任务空间用规则 tile 块划分，然后将这些 tile 分配给统一的一类硬件指令，每个 tile 都跑一样的指令，结果就是同质映射，难以挖掘 ILP。

本文提出 Mosaic，一种新型自底向上、基于镶嵌（tessellation）机制的编译器，可通过多种硬件指令对给定张量计算任务的整个计算空间进行镶嵌。Mosaic 以不同硬件指令的计算语义为构建单元，对任务空间进行异构化的划分，如图 1(b) 所示。该方法实现了从硬件指令到计算任务的异构映射，支持在不同计算单元上并行调度，从而高效挖掘指令级并行性（ILP）。
实现该目标需解决以下关键挑战：
1. 面向镶嵌编译过程的抽象建模
在镶嵌问题中，两个核心组成部分不可或缺：镶嵌的基本单元（即 texel）以及如何利用这些 texel 高效地覆盖整个计算空间。前者需要建立一个能够精确表达硬件原生命令所具备的计算语义的抽象模型。该模型的构建十分复杂，因为不同指令的语义差异较大。例如，GPU CUDA 核中的 FFMA 指令可用于执行向量乘法，也可适用于不同维度的矩阵乘法。对于后者，需要高效方法从大量可能的组合中选择适合的 texel 排列方式以完成整个计算任务的镶嵌。
2. 面向 ILP 挖掘的编译技术
即便获得了合理的镶嵌方案，仅靠简单的指令序列生成仍难以实现跨计算单元的并行执行，主要受限于底层硬件的结构性冲突。要充分发挥各计算单元的潜力，必须交错调度不同指令并重叠其执行过程。此外，为减轻寄存器分配压力、避免寄存器组冲突带来的延迟，最大化寄存器复用也至关重要。实现上述目标需要对指令执行过程进行细致编排，从而充分暴露潜在的 ILP。
为应对这些挑战，Mosaic 提出了三个关键模块：
（1）iTex 抽象：即“指令 texel”抽象，建模硬件指令操作与其计算语义之间的关系。具体而言，Mosaic 采用形式化的仿射函数推导出硬件指令可能支持的所有计算语义，为镶嵌编译过程提供基础构建块。
（2）基于 iTex 的镶嵌编译流程：Mosaic 将编译过程建模为一个精确覆盖（Exact-cover）问题，旨在通过组合不同的 iTex 实现整个张量计算空间的完整镶嵌。为解决该计算复杂问题，Mosaic 引入了一种启发式方法生成高质量的镶嵌方案。
（3）iTex 调度机制：包括“Zig-Zag 执行”与“指令交错”两种调度策略，用于协调各类指令的执行顺序，从而减少结构性冲突并最大限度地挖掘 ILP。
为了验证 Mosaic 的有效性，我们在多个基准测试中进行了广泛实验，涵盖高优化的 GEMM 运算、真实世界的大语言模型（LLM）算子及端到端网络。针对 NVIDIA 平台，我们选择了 RTX 4090、L20 和 A100，分别代表客户端、服务器与数据中心级 GPU。实验结果显示，Mosaic 在这些平台上相较于高度优化的厂商库（如 cuBLAS 与 CUTLASS）平均加速比为 1.08× 至 1.28×；在 LLM 实际算子测试中，相较于 cuBLAS 和 CUTLASS 提升 1.43×，相比当前领先的深度学习编译器（如 Triton 和 TensorIR）提升 1.25×。尤为显著的是，Mosaic 的性能最高可达 GPU 张量核心理论峰值吞吐率的 106%，充分体现其对 ILP 的高效利用。此外，我们亦在 AMD MI100 GPU 与 Intel Xeon 8481C CPU 上开展实验，进一步验证了 Mosaic 的通用性。

意思就是通过加一层语义，来让编译流程中有机会更好的映射到硬件单元。
①Zig-Zag：Zig-Zag 执行是一种特殊的指令调度顺序，在空间上呈现“之”字形排列，目的是为了提升操作数重用率，降低寄存器银行冲突。这样安排能让连续的指令尽量重用已有的寄存器数据，而不是每次都从新位置加载，节省访问、减少冲突。
②指令交错是一种把不同类型指令混合排布的策略，避免相同类型的指令连续出现导致资源冲突或执行停顿（stall）。例如FFMA → Load → FFMA → Store → FFMA 这种顺序执行指令。

动机

在本节中，我们以微基准测试为指导，首先剖析典型深度学习加速器（即 NVIDIA Ada RTX 4090 GPU）的体系结构，以揭示其潜在的指令级并行性（ILP）机会。随后，我们分析限制可利用 ILP 的内在结构性冲突。最后，我们给出一个手写示例，展示在面临这些机遇与挑战时的实际应对方法，从而为我们的设计提供动机。我们通过构建以 SASS 汇编代码编写的微基准测试来剖析 GPU 的体系结构和周期级行为。SASS 代码的反汇编与重汇编通过开源工具 CuAssembler 得以实现。

2.1 剖析体系结构以揭示潜在 ILP

图 2 展示了 NVIDIA 4090 GPU 中 SMSP 的体系结构。SMSP 是 GPU 的基本构建单元，负责执行一个 warp 中 32 个线程的机器指令。每个 SMSP 包含以下组件：

（a）一个调度器（图顶部），每个周期顺序发射可执行指令；
（b）一个寄存器文件，结构上划分为两个寄存器 bank，每个 bank配备两个 512 位读端口和一个 512 位写端口；
（c）一个操作数收集器，兼作寄存器文件的只读缓存；
（d）三种主要计算单元：
- Vector0，为支持 FP32 或 INT32 操作而设计的 16 通道向量单元；
- Vector1，为支持 FP32 或 2 倍精度 FP16 操作的 16 通道向量单元；
- 以及一个张量计算单元。
  Vector0 和 Vector1 通常被称为 CUDA 核心，而张量计算单元则被称为张量核心。
  
  在 NVIDIA GPU 中开发指令级并行性的主要机遇在于 CUDA 核心与张量核心的并行执行能力。具体而言，这包括由 CUDA 核心执行的 FFMA（融合浮点乘加） 向量操作以及由张量核心执行的 HMMA（矩阵乘加） 操作。FFMA 指令在一个 warp 的 32 个线程中分派执行，因此每个包含 16 通道的 CUDA 核心将在 2 个周期内完成一次 warp 的 FFMA 指令。由于每个 SMSP 拥有两条向量路径，调度器可以在交替的周期中将 FFMA 指令分派至每个 CUDA 核心。同时，HMMA 指令通常由张量核心在 16 个周期内完成。在理想情况下，一旦发射一条 HMMA 指令，SMSP 可以同时调度多达 15 条 FFMA 指令与其并行执行，从而达到最佳的指令级并行性，在 FP32/TF32 精度下，在 NVIDIA 4090 GPU 上实现理论上的 1.93 倍加速比。图 3(a) 展示了 CUDA 核心与张量核心之间理想的时空并行性示意图。

2.2 限制可利用 ILP 的结构性冲突

尽管微基准测试揭示了 NVIDIA GPU 中指令级并行性的潜力，但实际开发过程中常受到内在结构性冲突的制约。这些冲突源于执行指令所需的硬件资源不足或资源之间存在竞争，从而限制了可开发的并行性。尽管 GPU 拥有多个功能独立的计算单元，仍需重点关注两类结构性冲突：寄存器bank冲突与指令发射冲突。

2.2.1 寄存器bank冲突

当一条指令被发射执行时，首先需要将数据从寄存器读取到相应的计算单元。如果有多条指令同时尝试访问寄存器，超过了寄存器的读取能力，就会发生寄存器bank冲突（Register Bank Conflicts），从而阻碍指令的执行。在 GPU 上，尤其是在开发指令级并行性（ILP）时，寄存器bank冲突是一个关键问题。

①具体例子：FFMA R6, R1, R2, R3 // 读取 R1, R2, R3，写 R6
假设寄存器 R1, R2, R3 都在同一个 Bank，每个 Bank 只能同时读取 2 个寄存器，但是同一个指令要读 3 个寄存器，剩下一个必须等下一个周期 ⇒ 就产生了 Bank 冲突

我们通过微基准测试分析了 NVIDIA Ada GPU 上寄存器文件的设计。结果显示，每个寄存器为 1024 位，而一个 SMSP（Streaming Multiprocessor Sub Partition）配备了两个寄存器bank，每个bank拥有两个读取端口和一个写入端口，每个端口每个周期可以处理 512 位的数据。因此，要通过一个读取端口访问一个寄存器，需要两个时钟周期。

意思是每个端口 %% 一次最多只能读或写 %% 512 位，所以一个 1024 位的寄存器 → 必须拆分成两个 512 位访问。

我们进一步研究了 FFMA（融合乘加）和 HMMA（半精度矩阵乘加）指令对寄存器的需求，从而展示了寄存器bank冲突如何限制可利用的指令级并行性（ILP）。
FFMA 指令涉及三次读操作（即访问三个寄存器）和一次写操作。图 4(a) 展示了连续发射 FFMA 指令时所引发的寄存器bank冲突。最初，第一条 FFMA 指令（图中以蓝色表示）在前两个周期内通过两个寄存器bank的三个读端口访问寄存器 R1、R2 和 R3。随后，当第二条 FFMA 指令（粉色）被发射时，仅有一个读端口可用，因此只能完成其中一个寄存器的读取请求，剩余两个请求需延迟至第三个周期才能完成。由此，第二条 FFMA 指令的寄存器读取延迟至第四周期结束，而第三条 FFMA 指令（黄色）的操作数读取则被顺延至第六周期。

Maxwell 架构引入了可编程的寄存器重用缓存机制，能够有效缓解 FFMA 指令中的寄存器bank冲突，如图 4(b) 所示。通过使用 .reuse 指令，可以将寄存器 R3 标记为可重用。这样，在第一条 FFMA 指令访问 R3 后，其值将被缓存在只读寄存器缓存中，使得后续两条 FFMA 指令可以直接从缓存中获取 R3 的值，从而缓解限制可利用指令级并行性的寄存器bank冲突。最终，这三条 FFMA 指令的寄存器读取可在第四周期内完成。
HMMA 指令同样涉及三次读取和一次写入操作，也可以利用寄存器重用缓存机制。特别地，HMMA 指令的操作数是片段（fragment），这是一种在 NVIDIA GPU 编程接口中内建的数据结构，对程序员而言是透明的。微基准测试表明，一条典型的 HMMA 指令包含一个为期 8 个周期的读取阶段和一个为期 16 个周期的执行阶段。在读取阶段，HMMA 指令仅使用每个寄存器bank的一个读端口，释放了另外两个读端口。这种资源可用性为穿插执行 FFMA 指令提供了可能，从而提升了寄存器的使用效率，并进一步增强了指令级并行性。

2.2.2 发射冲突

发射冲突是指多条指令在同时发射过程中因共享硬件资源而引发的冲突。尽管 NVIDIA Ada GPU 采用了简单的顺序单发射调度器，限制了 HMMA、FFMA 及其他指令在同一周期内的并行发射，但当指令竞争相同的执行单元或涉及对内存与寄存器的访问时，仍可能发生发射冲突。

如果多个指令（即使不同时发射）：要用相同的 寄存器资源、要抢占同一个 执行单元、或者频繁访问 内存/缓存，就会导致发射冲突。

为了最大化 HMMA 与 FFMA 指令之间的指令级并行性（ILP）并避免发射冲突，需要采取以下三项关键策略：
（1）尽可能减少非计算类指令的数量，以确保 GPU 始终专注于执行计算任务；
（2）精心安排指令顺序，充分利用寄存器重用缓存，以降低由前述寄存器bank冲突引起的延迟；
（3）按照特定模式交错执行 HMMA 与 FFMA 指令，使不同指令的执行阶段得以重叠。

2.2.3 动机例子

为此，我们手工设计了一个基于 warp 级别的内核（如图 3(d) 所示），用于规避上述结构性瓶颈，并挖掘 FFMA 与 HMMA 指令之间潜在的 ILP 机会。具体而言，该内核首先利用寄存器重用缓存来消除不必要的寄存器bank冲突。随后，通过精确调度指令执行，解决发射冲突问题，从而最大限度地发挥向量计算单元与张量计算单元的资源利用率。该内核的时空执行图展示于图 3© 中。相较于图 3(b) 中现有深度学习编译器仅使用 Tensor Core 的情况，该手工内核在相同周期内额外完成了 75% 的计算任务。该示例充分展示了在 NVIDIA Ada GPU 上挖掘指令级并行性所带来的潜在性能优势，并进一步激发了我们提出相关设计的动机。

Mosaic 概述

本文提出了一种全新的自底向上、基于镶嵌（tessellation）策略的深度学习编译器 Mosaic。图 5 展示了 Mosaic 的整体架构。Mosaic 以张量计算任务（例如指定 M、N、K 维度的 GEMM 计算）为输入，直接将该计算任务划分为不同类型的指令，通过异构的指令到任务映射机制以挖掘在深度学习加速器（DLA）上的指令级并行性（ILP）。
Mosaic 编译器主要由以下三个核心组件构成：

iTex 抽象模型：该模块通过形式化的仿射函数建模硬件指令的操作与语义之间的关系，同时作为 Mosaic 的基础中间表示。
任务镶嵌模块：该模块将编译过程视为一个空间镶嵌问题，通过无缝拼接多种 iTex 实现对输入张量计算任务的全面镶嵌。此模块同时承担 Mosaic 的前端功能。
代码生成与调度模块：该模块首先为 iTex 生成面向具体硬件的指令（即算术运算指令）及其相关的内存访问指令；随后，借助多种 iTex 调度策略，组织指令的执行顺序，以减少潜在的结构性瓶颈，并最大化可利用的指令级并行性。该模块构成了 Mosaic 的后端部分。

任务镶嵌

4.1 iTex 抽象

预备知识
为提升表述清晰度，本文首先引入若干关键概念及其记号说明：

空间（Space）：例如， $\text{GEMM}((2_p, 2_q)_m, 4_n, 2_k)$ 表示一个 4 × 4 × 2 的 GEMM 任务空间，其中 m 维是一个嵌套的 2 × 2 空间。空间中的元素按照 colexicographic 顺序排列。
映射（Map）：映射是指将一个空间中的元素以准仿射（quasi-affine）关系映射到另一个空间。例如，给定一个融合乘加（FMA）计算空间 FMA(32 $_q$ )，映射 FMA( $32_q$ ) → FMA(4 $_m$ , 8 $_n$ ) 以及 FMA[i] → FMA[i mod 4, i / 4]，均表示将一维计算空间转换为一个 4 × 8 的计算空间，其中每 4 个来自 q 维的元素被重组为 m 维的一行。
任务（Task）：一个任务由三部分组成：以循环结构表示的任务空间、表示被访问数据的操作数空间（operand spaces），以及任务空间与操作数空间之间的映射关系。在简化表达时，可省略操作数空间与其映射。
接下来介绍 iTex 抽象模型。该模型由三部分组成：用于表示指令操作数空间的集合、描述计算语义的任务、以及从指令操作数空间到任务操作数空间的映射。简而言之，iTex 回答了“单条硬件指令具备何种计算语义”的问题，这是传统文本式指令表示方法所无法表达的能力。一条指令可能对应多个有效的 iTex，每个 iTex 表示该指令所能完成的某个任务。Mosaic 编译器采用三类 iTex 表示方式：规范 iTex、派生 iTex 和 变换 iTex，如图 6 所示，具体将在后文详细讨论。

规范 iTex（Canonical iTex）
规范 iTex 描述的是指令本身所固有的语义，由编译器预定义。例如，NVIDIA GPU 中的 FFMA 向量指令，其对应的规范 iTex 可表示为：
操作数空间：rs1(32 $_l$ )、rs2(32 $_l$ )、rs3(32 $_l$ )、rd(32 $_l$ ) ← 就是一个指令的操作数
任务：FMA(32 $_l$ ) → A(32 $_l$ )、…
映射关系：rs1(32 $_l$ ) → A(32 $_l$ )、…
该 iTex 表明，FFMA 指令完成一个 32 通道宽度的向量融合乘加操作，每个基础任务由一个向量通道（lane）完成。此外，该表示也揭示了指令的向量操作数与任务操作数之间为恒等映射关系。
派生 iTex（Derived iTex）
尽管可以直接使用规范 iTex 对给定的计算任务进行镶嵌，但实际应用中，任务语义与某条指令的规范 iTex 往往并不一致，导致无法直接完成镶嵌。因此，Mosaic 编译器需要回答一个关键问题：为什么某条指令能够完成一个不同于其规范语义的任务？ 例如，如何利用向量化的 FFMA 指令完成一个 GEMM 任务。
为解决这一问题，Mosaic 自动生成派生 iTex（Derived iTex），即通过一系列派生规则将某条指令的规范任务空间映射到目标任务空间，从而匹配给定任务的语义。这些派生规则记为函数 $\varphi$ ，每条规则表示为从一个任务空间到另一个任务空间的变换关系。部分典型规则如下：
从 FMA 到 OuterProduct：
$\varphi(\text{FMA}(m), \gamma) = \text{FMA}(m) \rightarrow \text{Outer}(\gamma, \frac{m}{\gamma})$
从 OuterProduct 到 GEMM：

$\varphi(\text{Outer}(m, n)) = \text{Outer}(m, n) \rightarrow \text{GEMM}(m, n, 1)$

从 GEMM 到 GEMV：
$φ(GEMM(m,n,k))=GEMM(m,n,k)→GEMV(1,n,k)\varphi(\text{GEMM}(m, n, k)) = \text{GEMM}(m, n, k) \rightarrow \text{GEMV}(1, n, k)$
恒等映射（Identity Rule）：
$φ(x)=x→x\varphi(x) = x \rightarrow x$

通过这些派生规则，Mosaic 能够系统性地从规范指令语义出发，推导出指令可以支持的更广泛的任务空间，从而在任务与指令之间建立灵活且高效的语义对应关系，实现异构任务镶嵌与更高的指令级并行性。
更具体地说，我们展示了如何为 FFMA 指令导出一个执行 $4 \times 8$ 外积（Outer Product）任务的 iTex，如图 6(b) 所示。当设定参数 $\gamma = 4$ 时，其派生映射为：

$\varphi(\text{FMA}(32l), 4) = \text{FMA}(32l) \rightarrow \text{Outer}(4l_0, 8l_1)$

关于指令操作数空间到任务操作数空间之间的映射，以操作数 rs2 为例，通过级联组合（chain-composing）规范 iTex 的操作数映射、FMA 任务的操作数映射、派生映射及派生任务的操作数映射，最终可构造出如下派生操作数映射：

$\text{Outer}(4l_0, 8l_1) \rightarrow B_{\text{Outer}}(8l_1) \circ \text{FMA}(32l) \rightarrow \text{Outer}(4l_0, 8l_1) \circ B_{\text{FMA}}(32l) \rightarrow \text{FMA}(32l) \circ rs2(32l) \rightarrow B_{\text{FMA}}(32l) = rs2(4l_0, 8l_1) \rightarrow B_{\text{Outer}}(8l_1)$

该派生映射明确描述了使用向量化 FFMA 指令完成 Outer Product 任务所需的内在收缩（contraction）语义要求。具体而言，此映射表明：B 操作数中的每一个值需要在 rs2 操作数中被复制 4 次，以满足语义正确性。
变换 iTex（Transformed iTex）
派生 iTex 可用于任务镶嵌，因为它们的任务空间与目标计算任务相匹配。然而，仅使用派生 iTex 进行镶嵌往往效率不高。例如，在图 6(b) 所示的场景中，rs2 的 $l_1$ 维与操作数 B 的 $l_1$ 维具有恒等映射关系。因此，加载 B 至 rs2 时只能使用每通道 32 位的向量加载操作，这会导致内存复制指令数量增加，从而降低加载效率，并减少可用于发射计算指令的槽位，进而影响 ILP 的提升。
为缓解此问题，Mosaic 进一步从派生 iTex 构造变换 iTex（Transformed iTex）。变换 iTex 是通过在派生任务空间基础上再次施加映射生成的，其技术路径与派生 iTex 的生成方式一致。如图 6© 所示的变换映射为：

$\text{Outer}(4, 8) \rightarrow \text{Outer}(4, (2, 8))$

相较于原始派生 iTex，该变换 iTex 将 $l_1$ 维引入了跨度为 2 的扩展。这一变换允许多个连续 iTex 实例共享每通道 64 位的加载指令，从而提升内存复制效率，并为算术指令保留更多的发射资源。

如果无法理解可以详细看下图 6

4.2 镶嵌建模（Tessellation Formulation）

在获得多种 iTex 之后，接下来的目标是使用这些 iTex 对张量计算任务进行空间覆盖，要求覆盖无重叠亦无空缺，即构建一个有效的 iTex 镶嵌（tessellation）。Mosaic 设计了两类镶嵌策略：均匀镶嵌（uniform tessellation）和非均匀镶嵌（non-uniform tessellation，又称非周期镶嵌）。
均匀镶嵌指使用某一固定形状的 iTex $T(m, n)$ 对任务空间 $T(M, N)$ 进行规则划分，其结果为：

$T\left((m, \frac{M}{m}), (n, \frac{N}{n})\right)$

而非均匀镶嵌则涵盖所有其他形式的有效镶嵌，其唯一约束条件是：任意两个 iTex 实例的任务空间不能相交，且它们的并集应正好覆盖目标计算任务空间。这种高度灵活的表示形式带来了极大的搜索空间，增加了设计复杂度。
为应对这一挑战，Mosaic 将问题形式化为精确覆盖问题（Exact-Cover Problem），以显著收敛搜索空间。该方法定义的搜索流程如下：

首先，为每个 iTex 构造其所有可能的均匀镶嵌；
然后，从这些均匀镶嵌中选择一组 iTex，用于构建一个完整的非均匀镶嵌，覆盖整个任务空间。
该策略能有效剔除所有不对齐的镶嵌方案，显著减少可行解的规模，从而提升搜索效率。
在 Mosaic 中，我们引入了两种关键的镶嵌策略：独立镶嵌（independent-tessellation） 与归约镶嵌（reduction-tessellation），用于依据设定的比例参数（Ratio）在两个 iTex 间进行均匀镶嵌组合，从而生成非均匀镶嵌方案。为增强 Mosaic 的适应性与鲁棒性，系统亦支持 Ratio 取值为无穷大或零的极端情况。
具体而言，独立镶嵌策略沿输出张量的维度进行任务划分，例如 GEMM 运算中的 $M$ 与 $N$ 维；而归约镶嵌策略则沿归约轴（reduction axis）进行划分。两种策略各有优劣：独立镶嵌可能降低 $M$ 和 $N$ 维度上的数据重用率；归约镶嵌则可能带来额外的归约开销，并显著增加累加器的寄存器使用压力，从而加重寄存器资源负担。因此，在实际的镶嵌计划中，我们结合两种策略，以实现性能与资源的平衡权衡。

4.3 镶嵌计划生成（Tessellation Plan Generation）

虽然上述精确覆盖问题的所有有效解可视为潜在的镶嵌计划（Tessellation Plans），但在实际任务规模下，这一解空间仍然巨大。这主要源于精确覆盖问题本身具有指数级时间复杂度（ $\mathcal{O}(2^N)$ ），而实际任务空间常为如 $T(4096^3)$ 级别。即便是当前最先进的 Knuth 的 Algorithm X 也难以高效求解此类问题。
然而，在深度学习编译场景中，许多理论可行解在实际中并不高效。原因有二：

深度学习加速器（DLA）广泛采用单程序多数据（SPMD）模型，并具有分层多核架构。这些特性天然引入了一些重复模式作为基本单元（如 NVIDIA GPU 的 SMSP 或 Intel CPU 的核），不仅降低了编程复杂度，还提高了指令缓存的利用率。
在这些基本单元内，若使用过于灵活的非均匀镶嵌策略（例如将空间划分为三种以上的 iTex 类型），则由于计算与内存访问指令之间对齐困难，可能会引入额外开销，显著降低整体性能。
为解决上述问题，我们设计了一种启发式算法（见算法 1），用于压缩解空间并生成高效可行的镶嵌计划（Tessellation Plans, 简称 TPs）。该算法接收四个输入：
$\text{iTexes}^V$ ：向量类指令的派生及变换 iTex 集合；
$\text{iTexes}^T$ ：张量类指令的派生及变换 iTex 集合；
$\text{Ratio}$ ：控制两类 iTex 混合比例的参数；
${\text{Spaces}}$ ：任务空间的层级结构列表，其中最底层表示具体的张量计算任务空间。
算法整体分为三个步骤

iTex 选择与平铺（Selection and Tiling）
非均匀镶嵌生成（Non-uniform Tessellation）
层级镶嵌调度（Hierarchy Tessellation）
最终输出完整的镶嵌计划集 $\text{TPs}$ 。
在算法中，独立镶嵌与归约镶嵌作为两种基本策略被融合使用，用于在空间维度与归约维度之间取得性能与资源的最优折中。这种灵活的镶嵌计划设计，使 Mosaic 能够在复杂架构与不规则任务之间实现高效的指令映射与执行调度。

代码生成与调度

5.1 代码生成（Code Generation）

Mosaic 的代码生成过程采用与镶嵌计划生成类似的层级结构。在生成非均匀镶嵌计划的底层（例如 warp 级别），Mosaic 首先将所得 iTex 实例直接翻译为对应的硬件算术指令。在更高层级（即采用均匀镶嵌计划的层级），编译器则引入循环变量或并行变量（如 blockIdx），用于索引重复维度的任务空间。
在不同层级之间，Mosaic 融合了多种常见优化策略，包括使用局部层级的存储空间（如共享内存）以及内存流水线（memory pipelining）技术。最终，系统会根据每个 iTex 实例的空间坐标与其操作数映射关系，推导并合并内存访问模式，并注入相应的内存访问指令。
归约优化（Reduction Optimization）
在进行归约镶嵌（reduction-tessellation）时，Mosaic 会为不同类型的 iTex 分别分配独立的累加器。这是因为各 iTex 可能具有不同的操作数映射关系，导致最终归约操作无法简单合并，因此必须显式处理归约开销。
对于 GPU 平台，Mosaic 提供三种归约策略，以应对不同类型 iTex 的累加器归一需求：

直接归约（Direct Reduction）：当多个 iTex 拥有相同的操作数映射关系时，可直接对其寄存器累加器进行归约。
局部 shuffle 归约（Local Shuffle Reduction）：当映射关系不一致时，使用 warp shuffle 指令对某一 iTex 的累加器进行重映射以匹配另一个 iTex。
共享内存归约（Shared Memory Reduction）：通过共享内存将一个 iTex 的累加器转换为另一 iTex 可用的格式。
这三种策略的适用范围与开销依次递增，分别在效率与通用性之间形成折中。

5.2 iTex 调度（iTex Scheduling）

尽管代码生成阶段已根据镶嵌计划为每个 iTex 实例生成了相应指令，但这些 iTex 实例在空间与时间维度上仍处于无序状态。若使用简单的线性顺序执行策略，极易引发结构性瓶颈，例如寄存器bank冲突与指令发射冲突，从而限制可挖掘的指令级并行性（ILP）。
为此，Mosaic 提出iTex 调度机制，旨在精细编排 iTex 的执行顺序，最大程度地减少潜在结构冲突，并释放更多可用的并行性。需要指出的是，当前 Mosaic 在 NVIDIA GPU 上仅支持中间表示 PTX 代码级别的调度。
Zig-zag 执行策略（Zig-zag Execution）
Mosaic 为算术类指令引入了一种空间“Zig-zag”式执行模式。对于 GEMM 任务而言，该执行模式可确保连续两条相同类型的指令理论上至少共享一个可重用操作数，从而在降低寄存器bank冲突概率的同时，提升 ILP 的暴露程度。

指令交错（Instruction Interleaving）

指令交错是一种经典的调度技术，旨在重新排列指令执行顺序，避免由指令发射冲突带来的执行停顿，进一步提升可利用的 ILP。
目前，Mosaic 采用一种简化的启发式策略：

延长内存依赖距离并对内存访问指令进行分组；
对算术指令进行排序，使其尽可能贴近微基准测试中推导出的最优执行模式。
通过上述调度机制，Mosaic 在结构性瓶颈与执行效率之间取得了更优平衡，充分释放了硬件资源所支持的潜在并行性。

评估方法

基本信息

~4K 行 Python：实现 iTex 主逻辑
~15K 行 CUDA/C++/HIP：在 NVIDIA 的 CuTe/CATLASS 上扩展

测试

NVIDIA GPU（Ada RTX 4090，Ada L20，Ampere A100）
AMD GPU（CDNA1 MI100）
Intel CPU（Intel Xeon Platinum 8481C）

基准

GEMM 与 BMM（批量矩阵乘法）
选择原因：

最核心的计算算子
此类算子可同时调用平台上的张量计算单元与向量计算单元，方便验证 Mosaic 有效性
目前这两个是在深度学习加速器（DLA）上优化程度最高的算子，所以 Mosaic 效果好的话更具有说服力
附加基准：

Phi-3B、Mistral-7B、Llama-3-8B、Llama-2-13B

评测阶段

LLM 微调的预填充阶段（pre-fill），采用 TF32-FP32 与 BF16-FP32 精度；
推理阶段的预填充任务，采用 FP16-FP16 与 NT8-INT32 精度。

预解码阶段普遍使用键值缓存（KV-cache）技术，计算范式从原始的 GEMM 转变为 GEMV，进而无法利用张量核心，因此我们将该阶段从评估中剔除。

本文略过实验效果

讨论

通过对 Mosaic 编译性能的系统性分析，我们识别出两个关键的体系结构改进方向，可进一步提升其运行效率：

进一步缓解寄存器bank冲突。在现代 GPU 架构中，静态物理寄存器分配机制可能导致 Tensor Core 与 CUDA Core 之间不可避免的寄存器冲突。若能引入动态寄存器重命名机制，有望缓解该资源竞争问题，从而提升 Mosaic 的指令调度效率。此外，我们的实验观察表明，NVIDIA GPU 上的向量与张量计算单元可能共享同一个寄存器重用缓存（register reuse cache），这在混合调度场景下会进一步加剧寄存器压力。若能为每类计算单元提供独立的寄存器重用缓存，可在硬件代价极小的前提下显著降低该冲突。
减少指令发射冲突。目前，Mosaic 依赖 iTex 调度策略来规避指令发射冲突，但该策略仍受限于厂商提供的底层编译器接口。若硬件支持类似 AMD RDNA3 架构的双发射调度器（dual-issue scheduler），或具备乱序调度能力的 GPU 指令调度器（out-of-order scheduler），则有望进一步减少发射瓶颈，释放更多潜在并行性。

元信息概括

领域：AI 编译→调度优化
创新点：提出了基于指令语义重建模的 iTex 抽象机制，并结合自底向上的 镶嵌式调度（tessellation scheduling），打破传统同构指令映射限制，构建异构指令到任务的映射关系，从而实现在向量计算单元与张量计算单元间的高效协同，显著提高指令级并行性（ILP）与硬件利用率。
动机：当前异构计算平台（如 GPU、DLA 等 XPU）在执行深度学习算子（尤其是 GEMM/BMM）时，常因编译器对指令粒度调度能力不足，导致 CUDA Core 与 Tensor Core 无法协同工作，计算单元资源闲置严重，限制了性能上限。
方法简述：

提出 iTex 抽象，形式化建模指令与任务的语义映射关系；
基于 派生规则与变换规则 构造多种指令任务空间（iTexes）；
引入 Exact-Cover 编译问题建模 与启发式求解算法，生成任务空间的高效镶嵌；
设计 Zig-zag 调度模式 和 指令交错技术，最大化指令发射效率并缓解寄存器银行冲突；
实现一个具备完整编译流程的 Mosaic 编译器框架。

工具+平台：NVIDIA RTX 4090、L20、A100，AMD MI100，Intel Xeon 8481C+CuTe / CUTLASS
局限性：

当前 Mosaic 只能在 PTX 层级进行调度，仍受限于厂商编译器（如 NVCC）的接口能力；
尚未支持 decode 阶段的优化，该阶段多使用 GEMV 运算，不易调动 Tensor Core；
硬件结构限制，如共享寄存器缓存与单发射调度器，仍制约进一步并行性提升；
未来可结合硬件支持的乱序调度器或寄存器重命名机制，进一步释放 Mosaic 的潜力。

Mosaic： Exploiting Instruction-Level Parallelism on Deep Learning Accelerators with iTex Tessellation