spacing=nonfrench

#  PartIR: Composing SPMD Partitioning Strategies for Machine Learning

Sami Alabed 33footnotemark: 3 Bart Chrzaszcz11footnotemark: 1 33footnotemark:
3 Juliana Franco11footnotemark: 1 33footnotemark: 3 Dominik
Grewe11footnotemark: 1 33footnotemark: 3 Dougal Maclaurin11footnotemark: 1
33footnotemark: 3
James Molloy11footnotemark: 1 33footnotemark: 3 Tom Natan11footnotemark: 1
33footnotemark: 3 Tamara Norman11footnotemark: 1 33footnotemark: 3 Xiaoyue
Pan11footnotemark: 1 33footnotemark: 3 Adam Paszke11footnotemark: 1
33footnotemark: 3
Norman A. Rink11footnotemark: 1 33footnotemark: 3 Michael
Schaarschmidt11footnotemark: 1 Timur Sitdikov11footnotemark: 1 33footnotemark:
3
Agnieszka Swietlik11footnotemark: 1 33footnotemark: 3 Dimitrios
Vytiniotis11footnotemark: 1 33footnotemark: 3 Joel Wee11footnotemark: 1
33footnotemark: 3 Equal contribution, authors in alphabetical order.
Correspondence to<EMAIL_ADDRESS>done while at Google DeepMind.
33footnotemark: 3 Google DeepMind 22footnotemark: 2 Isomorphic Labs

###### Abstract

Training of modern large neural networks (NN) requires a combination of
parallelization strategies encompassing data, model, or optimizer sharding.
When strategies increase in complexity, it becomes necessary for partitioning
tools to be 1) expressive, allowing the composition of simpler strategies, and
2) predictable to estimate performance analytically. We present PartIR, our
design for a NN partitioning system. PartIR is focused on an incremental
approach to rewriting and is hardware-and-runtime agnostic. We present a
simple but powerful API for composing sharding strategies and a simulator to
validate them. The process is driven by high-level programmer-issued
partitioning tactics, which can be both manual and automatic. Importantly, the
tactics are specified separately from the model code, making them easy to
change. We evaluate PartIR on several different models to demonstrate its
predictability, expressibility, and performance.

##  1 Introduction

The recent growth of NN training requirements has significantly outpaced the
increase in accelerator memory and FLOPS. Google TPU [27, 26] v2 reports 46
TFLOPS / 16 GB of high-bandwidth memory (HBM) per chip, while v4 reports 275
TFLOPS / 32 GB[12] over a period where model parameters - consequently FLOPS -
requirements increased by $10^{4}$ [8]. Due to the need to scale out, today’s
large NNs [5, 44, 11, 20, 58, 60] are trained on many accelerators through a
mixture of parallelism strategies.

User-driven partitioners [32, 67, 53] make it possible to express a wide range
of parallelism strategies in high-level ML frameworks [4, 41, 1], without
requiring ML engineers to explicitly write distributed systems code using low-
level communication primitives – an error-prone and non-portable practice
(e.g. JAX pmap[4]). Despite their significant success in research and
production, such tools couple the ML model code (XLA [15], Triton [59],
OpenXLA [38]) and the partitioning information together by requiring internal
sharding annotations, making it difficult to change the parallelization
strategies when the underlying system configuration changes. They also feature
no support for incremental composition of various strategies and intermediate
feedback to users ahead of actual compilation and profiling – a time-consuming
process in a performance engineer’s workflow.

The complexity of partitioning for end-users has motivated research on
automated partitioning tools [24, 65, 51, 2, 61, 71]. These tools face
different challenges: interpretability, predictability, and scalability;
making them risky to use for large models where sub-optimal partitioning is
costly.

PartIR takes the stance that an incremental partitioner that composes manual
and automatic partitioning tactics, provides cost estimates after every
applied tactic, offers better control over the sharding of ML programs,
reduces the need for model-internal sharding annotations, de-risks automatic
partitioning, and makes partitioning an easier task.

###  1.1 Motivation and key design principles

  * •

Composable sharding strategies: Large NNs use a mix of sharding strategies to
maximize their hardware utilization. This can be as simple as applying batch
(data) parallelism and model parallelism (e.g., Megatron [54] parameter
sharding) or a complicated mix of data, possibly model, and optimizer sharding
composed together [70, 11, 20, 35].

  * •

Verifiable sharding strategies: Each known sharding strategy introduces a
verifiable number of MPI-style communication collectives in the ML program and
expects a certain runtime behavior. For example, with batch parallelism, we
expect one all-reduce collective per parameter and a reduction in per-device
flops approximately equal to the batch parallelism factor. Providing this
information back to the user allows them to confirm if their sharding strategy
implementation is correct, even before compiling their code.

  * •

Decoupling model code and sharding strategies: Baking complex sharding
strategies in the model makes the resulting artifact difficult to test and
maintain while distracting ML engineers with system details; this can be
mitigated by decoupling the sharding strategies from the model code, with the
extra advantage of making the sharding strategy generalizable to various
deployment configurations.

  * •

Native support for composing automated tools: A system that combines both
manual and automated tools benefits from a predetermined, well-known sharding
strategy manually inserted by an expert, and an automated one [51, 2, 71] to
discover unknown sharding strategies. It is particularly useful for
experimental models with unknown optimal sharding or for partitioning a fine-
tuning model built on top of a foundational model.

We designed PartIR as an MLIR-based compiler [31]. It separates the
partitioning stack from the downstream compiler (e.g., XLA) into external
MLIR-based rewriting passes operating on the StableHLO [39] dialect, the MLIR
encoding of XLA HLO [15]. This design allowed us to develop the partitioner
independently of the downstream compiler and to provide an interface to the
high-level libraries (e.g., automatic partitioning) and the JAX-level APIs.

PartIR fully separates the model implementation from its partitioning,
allowing ML engineers to focus on building models that outlive the
partitioning strategies, without concerns of if and how they need to be
partitioned. As an MLIR-based compiler, PartIR could be used for any frontend
(Tensorflow, PyTorch, JAX)111We only present the APIs for JAX since it is our
users’ tool of choice. or any SPMD backend (XLA, OpenXLA), delegating
hardware-specific code generation and execution to the underlying compiler and
runtime.

PartIR is built with incrementality and predictability in mind - the user
expresses the sharding strategies using sequences of PartIR tactics that we
refer to as a schedule. These tactics can never undo sharding decisions
introduced by previous tactics in the same schedule and can invoke manual or
automatic partitioning. These sharding decisions are translated through
PartIR’s stack to device-local code that makes communication collectives
explicit and verifiable. PartIR provides an analytical simulator for debugging
and automated tools that take the hardware description, the inserted
collectives, and the model to provide performance and memory estimates of the
applied sharding.

###  1.2 Contributions

To summarize, our main contributions are:

  * •

The overall design of a partitioner system that allows users to compose SPMD
sharding strategies incrementally via a schedule (3) and decouples the model
from its partitioning to make it possible to step-by-step debug and compose
manual or automated sharding decisions.

  * •

A compiler stack (4) that exposes tiling and replication actions as a compiler
API to be targeted by high-level manual tactics or automated tools. The most
important dialects in our stack are (i) PartIR:Core (5), which introduces
functional tiling and reduction loops on top of an array IR (e.g., StableHLO)
that can be given sequential semantics (for a reference implementation) or
parallel semantics (for SPMD lowering); and (ii) PartIR:HLO (6), which offers
a device-local view of the SPMD computation and high-level collective
communication primitives. In addition, we have formally proved the correctness
of lowering PartIR:Core to SPMD code.

  * •

A pass to propagate partitioning decisions across a module (5.2.2), based on
linear algebra homomorphisms (LAHs). Propagation in PartIR does not rely on
cost-based heuristics. The incrementality/ordering of tactics compensates for
the internal absence of heuristics and keeps the propagation logic simple.
Propagation is implemented as a set of rewrites at the PartIR:Core level that
introduce loops expressing these LAHs. The rewriting logic is generic and
shared across all tensor ops, simply by querying a centralized registry
containing every tensor op’s behavior in terms of LAHs.

  * •

An evaluation of PartIR for predictability and debugability (7.3),
composability across manual or automatic tactics (7.3.1), performance of
partitioned code (7.2), value of incrementality for reducing the number of
manual decisions (7.4), and compile-time overhead (7.5).

PartIR is actively used for scaling a wide range of machine learning models in
research [69, 52, 7, 13, 23, 6]

##  2 Background

We give an overview here of a few fundamental concepts that we will keep
referring to throughout this work.

###  2.1 Device meshes

Figure 1:  Batch (N) and model (M) parallelism. On the left, the communication
along the M axis (e.g., Megatron’s [54] activation reductions). On the right,
the communication along the N axis (e.g., gradient reductions). Each device
parameter shard is color-coded; notably, all devices along the batch axis hold
the same shard.

Distributed execution of large tensor programs employs the concept of a mesh,
exposed in ML frameworks (e.g., in jax.sharding.mesh). For our purposes, a
mesh is an n-dimensional array with named axes that offers a logical view of
the available devices. We may view a system of 16 devices as a 2D mesh
$\\{a{:}2,b{:}8\\}$; or as 3D mesh $\\{a{:}2,b{:}2,c{:}4\\}$; or just as a 1D
mesh $\\{a{:}16\\}$ among many others. The mesh structures used in practice
usually reflect the underlying system’s communication topology to make
reasoning about performance easier and better utilize the fastest networks.
For example, consider a cluster of 4 servers with 8 GPUs each, connected
through a fast interconnect network [36] within a server, but slower Ethernet
connections across servers. One may view this system as a $\\{x{:}4,y{:}8\\}$
system so that communication along devices that span each axis of the mesh
happens exclusively on one of the two networks.

###  2.2 Parallelism strategies

Figure 2: Batch (top) vs. Z3/FSDP parallelism (bottom), assuming two devices.
In Z3, the parameters are all-gathered before their use, and their gradients
are reduced-scattered before being used to update the local parameter shards.

Large-scale NN training is usually performed using a limited set of
parallelism strategies:

##### Batch parallelism (BP)

The input batch is sharded across the devices while the model parameters and
optimizer state are replicated. We expect to see one AllReduce operation per
parameter in the backward pass.

##### Model parallelism (MP)

The model parameters are sharded across the devices – for example in Megatron
sharding for Transformer models [54, 35, 30], 2 AllReduce operations (on
activation tensors) are introduced per Transformer layer (and 2 more for its
backward pass).

##### Optimizer sharding

The optimizer state is sharded to reduce peak-memory. Two popular variants
exist: ZeRO2 (Z2) that additionally partitions the parameter gradients, and
ZeRO3 (Z3) that further partitions the parameters [46, 48, 66]. This paper
uses an SPMD-oriented variant that shards all parameters, also known as fully-
sharded data parallelism (FSDP) [70]. From a communication perspective, Z3
introduces AllGather whenever the parameters are needed (once in forward and
once in backward pass for every parameter), but has the advantage that the
full parameters only live on the device briefly, just before they are used.
Correspondingly, in the backward pass, the gradients for the parameters are
not reduced; they are rather reduce-scattered across the devices – a cheaper
operation. The optimizer is then also sharded and uses the gradient shards and
parameter/optimizer shards to update the parameter shards. Figure 2 gives a
simplified view of a training step with batch and Z3 parallelism.

The aforementioned partitioning strategies can be sequentially combined to
span one or multiple axes of the device mesh. For example, Figure 1
illustrates partitioning over such a 2D mesh with data parallelism over one
axis and model parallelism over the other. For training large models, multiple
strategies are composed, e.g., $BP+MP+Z3$ [44, 35]. Beyond these basic
strategies, there exist many others, such as activation sharding after model
parallelism [64, 30], or multi-query sharding techniques in Transformers [43].
We have implemented and verified (in terms of expected collectives) many of
these combinations in PartIR.

Moreover, the PartIR incremental design helps verify that each strategy was
applied correctly, as the user can inspect the collectives and estimated
performance after each strategy. Some schedules (notably FSDP) rely on
prioritizing some sharding decisions over others (Section 5.2.3) – a natural
property of incremental partitioning in PartIR.

###  2.3 Strategies as program transforms

Let us now inspect how the outlined partitioning strategies work in practice.
For that, we consider a short JAX [4] program composed of two matrix
multiplications222For simplicity, we do not show a full training step with
back-propagation and optimizer code, just a feed-forward function. We even
skip the elementwise non-linear operators, as they are trivial to partition.:

⬇

def f(x, w1, w2):

return (x @ w1) @ w2

y = jax.jit(f)(input, w1, w2) # Trace, compile, and run.

Listing 1: A matmul chain in JAX.

The jax.jit() function extracts a representation of the Python function that
can be converted to XLA HLO [15] and compiled for a specific backend. As an
intermediate step, JAX represents the program using the StableHLO [31] MLIR
dialect, which in this case looks as follows333StableHLO does not feature a
matmul op, and represents it using a dot_general op. We use matmul as
syntactic sugar for terser presentation.:

⬇

func @main(%x: tensor<(256x8xf32>, %w1: tensor<8x16xf32>, %w2:
tensor<(16x8xf32>) {

%x1 = matmul(%x, %w1) : tensor<256x16xf32>

%x2 = matmul(%x1, %w2) : tensor<256x8xf32>

return %x2 : tensor<256x8xf32> }

Listing 2: An unpartitioned matmul chain.

Note that each value is annotated with a type, e.g., %x1 = ...:
tensor<256x16xf32>, denoting an array of shape [256,8] and a float32 element
type.

Assume we are given a 2D mesh $\\{B:4,M:2\\}$ to execute this program – what
kind of parallelization strategy should we use?

##### Batch (data) parallelism

One strategy is to partition the first (256-sized) dimension of input %x
across axis $B$, observing that the program is a pure “map” over that
dimension:

⬇

func @main(%x: tensor<64x8xf32>, %w1: tensor<8x16xf32>, %w2: tensor<16x8xf32>)

attributes {mesh={"B":4, "M":2}} {

%x1 = matmul(%x, %w1) : tensor<64x16xf32>

%x2 = matmul(%x1, %w2) : tensor<64x8xf32>

return %x2 : tensor<64x8xf32> }

Listing 3: Data-parallel matmul chain.

The resulting device-local program above takes a first argument of smaller
shape 64x8 since each device acts in parallel on a slice of %x determined by
the device coordinate along axis $B$. At the same time, the shape of
parameters %w1 and %w2 remains constant across all devices (due to
replication).

##### Adding model parallelism

If, on top of BP, we partition the parameter %w1 on $dim=1$ and %w2 on $dim=0$
along axis $M$, then each device along axis $B$ may perform a smaller
multiplication with a different parameter shard:

⬇

func @main(%x: tensor<64x8xf32>, %w1: tensor<8x8xf32>,

%w2: tensor<8x8xf32>) attributes ... {

%x1 = matmul(%x, %w1) : tensor<64x8xf32>

%x2 = matmul(%x1, %w2) : tensor<64x8xf32>

%x3 = all_reduce <"M"> %x2 : tensor<64x8xf32>

return %x3 : tensor<64x8xf32> }

Listing 4: Data-parallel and sharded matmul chain.

The first matmul is a map over the second dimension of %w1, so no special care
is necessary. We need only to remember that %x1 is partitioned along its
second dimension. For the second matmul, we observe that both of its operands
are partitioned along the contracting dimensions. Hence, we can recover the
original program semantics by inserting a final all_reduce operation across
axis $M$. This sharding of pairs of matrix multiplications is the essence of
the Megatron sharding in Transformers [54].

##### Adding fully sharded parameters

Notice that the parameters are only sharded on axis $M$ (but not $B$) in
Listing LABEL:lst:fsdp-chain. We could modify the function to shard them
further on dimensions 0 and 1, respectively, but we would need to insert two
all_gather operations right before they are needed in the matrix
multiplication:

⬇

func @main(%x: tensor<64x8xf32>, %w1: tensor<2x8xf32>, %w2: tensor<8x2xf32>)
attributes {mesh = {"B":4, "M":2}} {

%w1g = all_gather [{"B"},{}] %w1 : tensor<8x8xf32>

%x1 = matmul(%x, %w1g) : tensor<64x8xf32>

%w2g = all_gather [{}, {"B"}] %w2 : tensor<8x8xf32>

%x2 = matmul(%x1, %w2g) : tensor<64x8xf32>

%x3 = all_reduce <"M"> %x2 : tensor<64x8xf32>

return %x3 : tensor<64x8xf32> }

Listing 5: Data-parallel and fully sharded matmul chain.

These operations gather the shards on the corresponding dimensions. This
sharding of parameters, after batch parallelism, is the essence of (SPMD) Z3
sharding [46]. Note that the parameters are only gathered before use, reducing
the peak memory requirements.

Of course, more complex shardings are possible on top of the aforementioned
chain of strategies. For example, the input and output activation (%x and the
return value) may be additionally sharded on the model axis $M$. Such a
sharding will convert the all_reduce to an reduce_scatter and introduce an
all_gather on the input %x before using it in the first matmul. We omit the
MLIR code for this transformation (this, in effect, is the ES strategy in
Section 7.)

The sequence of examples above highlights that (i) popular sharding strategies
do compose; (ii) by following linear algebraic reasoning, a model can be
partitioned just by sharding its inputs and parameters, and sometimes internal
operations, see Section 8; (iii) it suffices to utilize information about the
operations’ parallel and contracting dimensions. In the rest of the paper, we
show how to express these strategies through compiler actions that enable a
semantics-preserving rewrite of an initial unpartitioned program (as in
LABEL:lst:unpartitioned-chain) into a device-local program (as in
LABEL:lst:data-parallel-chain, LABEL:lst:partitioned-chain and LABEL:lst:fsdp-
chain).

##### A note on scale

Although the illustrating example was small, a typical ML model contains
thousands of tensor operations (e.g., dot-products, convolutions, gather and
scatter ops, and more) and control flow (e.g., conditionals and loops). A full
training step, including back-propagation and optimizer, can reach 10-100k
operations and have thousands of multi-dimensional array arguments for model
parameters and optimizer state. A good API should not burden the user with a
decision per operation or function argument.

##  3 A schedule is all you need

PartIR exposes a schedule-like API, inspired by work in kernel-generating DSLs
[29, 68, 45, 9, 16]. This API allows users to compose partitioning strategies
incrementally. In PartIR, a schedule is a sequence of manual or automatic
tactics. For example, to achieve the final sharding of f
(LABEL:lst:jax_matmul) appearing in LABEL:lst:fsdp-chain, a JAX user would
write:

⬇

#1. Arrange devices in a BxM mesh.

mesh = maps.mesh(device_array, ("B", "M"))

#2. Define sharding strategies as series of tactics.

BP = ManualPartition({"x":0}, axis="B")

MP = ManualPartition({"w1":1}, axis="M")

Z3 = ManualPartition({"w1":0, "w2":1}, axis="B")

schedule = [BP, MP, Z3]

#3. Partition and get metadata & distributed function.

dist_fn, metadata = partir.jit(f, mesh, schedule)

Listing 6: Partitioning strategies as a series of tactics in a schedule.

The first tactic BP partitions the first argument "x" on dimension 0 and
across axis "B"; yielding batch parallelism (LABEL:lst:data-parallel-chain).
The second tactic MP partitions input w1 on dimension $1$. The compiler will
identify that this action is sharding the contracting dimension of a matmul
and will shard w2 on dimension 0 through a process we call propagation invoked
at the end of every tactic (LABEL:lst:partitioned-chain). The final tactic Z3
shards the parameters on the remaining available dimensions and axis $B$
(LABEL:lst:fsdp-chain). The schedule is finally passed to partir.jit together
with the function to partition and the device mesh. partir.jit works like
jax.jit, but injects the PartIR partitioning stack before compilation.

Even though we present JAX API examples, PartIR operates over StableHLO,
making it frontend- (e.g., JAX, PyTorch, TensorFlow) and a backend- (e.g.,
XLA, CUDA, OpenXLA) agnostic, as long as they produce and consume StableHLO.

partir.jit returns a partitioned module exposed as a Python callable ready to
be called with JAX sharded arrays and executed on the devices of the mesh.
Additionally, PartIR produces various metadata that includes a sharding
specification of the function inputs and outputs produced by PartIR, and cost
estimates (e.g., collective communication breakdown by type and simulation
results) recorded after every tactic in the schedule. Dumping the intermediate
modules after every tactic is also very common for debugging.

##### Mixing automatic and manual tactics

PartIR exposes an AutomaticPartition tactic that operates on one or more mesh
axes and can be composed with other manual tactics. For example, users could
use a manual tactic to introduce batch parallelism manually, and rely on an
automatic tactic to discover partitioning along the second axis "M":

⬇

BP = ManualPartition({"x":0}, axis="B")

AutoMP = AutomaticPartition(axis="M")

part_fn, _ = partir.jit(fn, mesh, [BP, AutoMP])

Listing 7: A schedule consisting of manual and automation tactics.

The AutomaticPartition is an interface for any optimization algorithm. We
implemented a Monte Carlo tree search for discovering partitioning strategies
[51, 2], using a cost model that seeks runtime improvement and penalizes
models that exceed device memory limits.

What is essential for composability is that both manual and automatic tactics
issue sequences of (the very same) lower-level PartIR compiler actions that
either (i) shard a value dimension along an axis or (ii) explicitly keep a
value unpartitionable across a mesh axis, or (iii) propagate sharding
information across a module. For example, the schedule in LABEL:lst:schedule-
simple generates a sequence of 7 PartIR actions:

⬇

tile<%x,0,"B">; propagate; //BP tactic

tile<%w1,1,"M">; propagate; //MP tactic

tile<%w1,0,"B">; tile<%w2,1,"B">; propagate; // Z3

In what follows, we focus on the PartIR system architecture, the
implementation of actions as program rewrites, and how PartIR eventually
generates device-local SPMD code.

##  4 System architecture

Figure 3:  PartIR partitioning stack, built on top of MLIR, so it supports
layering of new operators on top of existing ones – hence, we use "+" to
signify the introduction of new operators, and "-" to signify that certain
operators have now become illegal.

Figure 3 gives an architectural overview of our system:

  * •

Programs are generated from tracing JAX functions in the MLIR StableHLO
dialect, as described in Section 2.3.

  * •

Manual or automatic tactics generate sequences of compiler actions that
introduce and propagate functional loops and specialized slicing ops, that
belong in the PartIR:Core dialect.

  * •

PartIR:Core loops and slices are interpreted as sequential loops in the
PartIR:Temporal dialect, whose main use is a reference semantics of
PartIR:Core alongside more niche applications like automatic microbatching
transforms. We do not expand on this dialect further as we focus on SPMD in
this paper.

  * •

PartIR:Core loops and slices are lowered to the PartIR:HLO dialect generating
device-local collective communication ops. The collectives at PartIR:HLO refer
to mesh axes that make their IR encoding independent of the total number of
devices in the mesh (as opposed to collectives in StableHLO and XLA HLO that
directly reference groups of logical device IDs), and makes it easy to reason
about and fuse (Section 6). As a result, we implemented our simulation at that
level. To export, we lower any custom high-level PartIR:HLO ops to StableHLO
computations and hand over the module to XLA for compilation.

As Figure 3 shows, our architecture and separation of dialects allow us to
implement the right rewrites at the right level of abstraction and to also
independently test various components of the partitioner. For example,
PartIR:Core is unaware of SPMD execution, and its rewrite axioms remain very
simple; this is a separate step from SPMD lowering (a pass that we have
additionally proven correct), or the optimization of communication primitives.

##  5 PartIR:Core

PartIR:Core introduces just two operations on top of the legal operations in
StableHLO: (1) a loop op that expresses pure (parallel) loops performing
tiling or reductions, and (2) a slice op that extracts a tensor slice based on
a loop index. We will directly present these constructs and their semantics as
part of describing the PartIR:Core compiler actions.

###  5.1 Value tiling action

A tiling action tile<%value, dim, axis> creates a loop that in each iteration
yields a slice of %value along dimension dim. For example, value tiling %x
along dimension 0 and axis "B" from LABEL:lst:unpartitioned-chain produces:

⬇

func @main(%x: tensor<256x8xf32>, %w1: tensor<8x16xf32>, %w2:
tensor<16x8xf32>) {

%xt = loop "B" [#tile<0>] (%rB: range<4>) {

yield (slice 0 %x[%rB])

} : tensor<256x8xf32>

%x1 = matmul(%xt, %w1)

%x2 = matmul(%x1, %w2)

return %x2 : tensor<256x8xf32> }

The loop operation contains two static attributes: (i) a mesh axis ("B") and
(ii) an action attribute (#tile<0>). It also accepts a single-argument closure
(a region in the MLIR jargon) that represents the loop body: (%rB: range<4>){
... }. The closure takes as an argument a range value (%rB) and performs a
tensor computation returning a value of type tensor<64x8xf32>. The range
argument %rB plays the role of a loop index given a PartIR-specific range
type.

The slice ops consume these loop indices. The meaning of slice 0 %x[%rB] is
that it extracts the %rB-th chunk of the tensor %x along dimension 0. The
tiling here perfectly partitions the tensor dimension into $4$ equally-sized,
contiguous chunks since axis "B" has size $4$. Hence, the result of slice has
shape 64x8. Furthermore, the value tiling action has replaced %x of type
tensor<256x8xf32> with value %xt of the same type – i.e., value tiling is a
semantics- and type-preserving local rewrite.

###  5.2 Propagation action

Value tiling actions create fairly trivial loops making them not particularly
interesting. However, they help bootstrap a powerful propagation pass that
subsequently creates loops around operations consuming or producing these
values and further slices other arguments. This propagation is justified by
program equivalences that directly encode linear algebra homomorphisms.

####  5.2.1 Program equivalences

PartIR propagation is built around program equivalences that involve loop and
slice instructions. Figure 4 presents three admissible program equivalences
for a matrix multiplication. The first two rewrite a matrix multiplication as
a tiling loop with a smaller multiplication inside. The last one introduces a
new form of attribute accompanying the loop, a #sum action attribute. This
signifies that the results of each iteration of the loop should be reduced, as
we are slicing the operands on their contracting dimension.444We assume
addition here, but our implementation supports #sum<@f> custom reductions for
any associative reduction @f. These equivalences are justified as linear
algebra homomorphisms, with the monoidal structure being stacking for the
first two and addition for the last one.

To allow us to implement the rewriting code once and for all, for all
operators, we equip PartIR with a tile-mapping registry (TMR) that enables a
concise encoding of linear algebra homomorphisms as maps between tiling and
reduction attributes. The TMR contains, for every tensor operation with $n$
inputs, a set of specifications of the form

$t_{1}^{\bot},\ldots,t_{n}^{\bot}\hookrightarrow\sigma_{1},\ldots,\sigma_{k}$

where $t^{\bot}$ stands for an optional tiling action, while $\sigma$ stands
for an arbitrary action (including #sum). StableHLO ops and our loops may
return multiple results, hence the generalized form
$\sigma_{1},\ldots,\sigma_{k}$. One such specification asserts that a given
operation can be rewritten as a loop with action(s)
$\sigma_{1},\ldots,\sigma_{k}$ if we slice its operands according to
$t_{1}^{\bot},\ldots,t_{n}^{\bot}$ (a missing action implies no slicing).

We give below the entries for matmul (corresponding to the three equivalences
from Figure 4) and for an element-wise add operation that, unsurprisingly,
asserts that tiling its result requires tiling its operands in the same way:

$\small\begin{array}[]{lcl}\textrm{TMR}({\tt
matmul})&=&\\{({{\small\tt\\#tile}}\langle{\tt
0}\rangle,\bot)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 0}\rangle\\}\\\
&\cup&\\{(\bot,{{\small\tt\\#tile}}\langle{\tt
1}\rangle)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 1}\rangle\\}\\\
&\cup&\\{({{\small\tt\\#tile}}\langle{\tt
1}\rangle,{{\small\tt\\#tile}}\langle{\tt
0}\rangle)\hookrightarrow{\small\tt\\#sum}\\}\\\ \textrm{TMR}({\tt
add})&=&\\{({{\small\tt\\#tile}}\langle{\tt
d}\rangle,{{\small\tt\\#tile}}\langle{\tt
d}\rangle)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt
d}\rangle\\}\end{array}$

It turns out that this abstraction (similar to split annotations [40]) is
sufficient to capture a wide variety of equivalences, and for substantially
more complex ops, e.g. convolutions, scatter and gather, dynamic slices, and
reshapes.

⬇

%t = loop "B" [#tile<0>] (%rB:range<4>) {

%xs = slice 0 %x[%rB]

%z = matmul(%xs, y)

yield %z : tensor<8x8xf32> }

⬇

%t = loop "B" [#tile<1>] (%rB:range<4>) {

%ys = slice 1 %y[%rB]

%z = matmul(%x, %ys)

yield %z : tensor<32x2xf32> }

⬇

%t = loop "B" [#sum] (%rB:range<4>) {

%xs = slice 1 %x[%rB]

%ys = slice 0 %y[%rB]

%z = matmul(%xs, %ys)

yield %z : tensor<32x8xf32> }

Figure 4:  Programs equivalent to %t = matmul(%x, %y). We assume %x :
tensor<32x16xf32> and %y : tensor<16x8xf32>, and a mesh axis "B" of size 4.

####  5.2.2 Propagation pass

Propagation is a pass that greedily propagates known and partially known
information and introduces more loops-based on the linear algebra
homomorphisms encoded in the TMR.

##### Propagation of known information

Forward propagation searches for an entry that matches the actions of loops
that produce operands of an operation, whereas backward 555Not to be confused
with NN optimizer back-propagation. propagation searches for an entry that
matches the way the operation result is sliced downstream. For example, assume
that the value %x1 in LABEL:lst:unpartitioned-chain has been tiled already.

⬇

func @main(%x: tensor<256x8xf32>, %w1: tensor<8x16xf32>, %w2:
tensor<16x8xf32>) -> ... {

%x1 = matmul(%x, %w1)

//value tiling

%x1t = loop "B" [#tile<0>] (%rB: range<4>) {

yield (slice 0 %x1[%rB])

}

%x2 = matmul(%x1t, %w2)

return %x2 : tensor<256x8xf32> }

To propagate tiling forward, we observe that the matmul defining %x2 now takes
an operand whose first dimension is tiled, hence the TMR entry
$({{\small\tt\\#tile}}\langle{\tt
0}\rangle,\bot)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 0}\rangle$ for
matmul matches its operand context. Propagating backward, %x1 is produced by a
matmul and is then sliced along dimension 0, which matches the result of that
same TMR entry. As a result, both operations can be rewritten:

⬇

func @main(%x: ..., %w1: ..., %w2: ...) -> ... {

//result of backward propagation

%x1 = loop "B" [#tile<0>] (%rB: range<4>) {

yield (matmul(slice 0 %x[%rB], %w1))

}

//value tiling

%x1t = loop "B" [#tile<0>] (%rB: range<4>) {

yield (slice 0 %x1[%rB])

}

//result of forward propagation

%x2 = loop "B" [#tile<0>] (%rB: range<4>) {

yield (matmul(slice 0 %x1t[%rB], %w2))

}

return %x2 : tensor<256x8xf32> }

Through propagation, we have arrived at a program where every operation is
within a loop context. These loops may be interpreted sequentially in
PartIR:Temporal for micro-batching or lowered directly to SPMD in PartIR:HLO.

Here is what the fused program looks like:

⬇

func @main(%x: ..., %w1: ..., %w2: ...) -> ... {

%r = loop "B" [#tile<0>] (%rB : range<4>) {

%xs = slice 0 %x[%rB] : tensor<64x8xf32>

%x1s = matmul(%xs, %w1)

%x2s = matmul(%x1s, %w2)

yield %x2s : tensor<64x8xf32>

}

return %r : tensor<256x8xf32> }

Listing 8: Chained matrix multiplication with a tiling loop.

##### Inference from partially known information

Inference is the process of deducing missing operand value tiling based on a
partial match against a TMR entry. Continuing from Listing LABEL:lst:Core-
chain-matmul-tiled-on-batch, consider value-tiling %w2:

⬇

func @main(%x: ..., %w1: ..., %w2: ...) -> ... {

%w2t = loop "M" [#tile<0>] (%rb: range<2>) {

yield (slice 0 %x2[%rb])

}

%x2 = loop "B" [#tile<0>] (%rB: range<4>) {

%xs = slice 0 %x[%rB]

%x1s = matmul(%xs, %w1)

%x2s = matmul(%x1s, %w2t)

yield %x2s

}

return %x2 : tensor<256x8xf32> }

The TMR entry ${{{\small\tt\\#tile}}\langle{\tt
1}\rangle,{{\small\tt\\#tile}}\langle{\tt
0}\rangle)\hookrightarrow{\small\tt\\#sum}}$ is a partial match on the
operands of the second matmul, since the second operand (%w2t) is already
tiled. We can extend it into a full match by value tiling the first operand
(%x1s) and then continuing with forward propagation to eventually (after some
simplification) arrive at:

⬇

func @main(%x: ..., %w1: ..., %w2: ...) -> ... {

%r = loop "B" [#tile<0>] (%rB: range<4>) {

%x2s = loop "M" [#sum] (%rM: range<2>) {

%xs = slice 0 %x[%rB] : tensor<64x8xf32>

%w1s = slice 1 %w1[%rM] : tensor<8x8xf32>

%x1ss = matmul(%xs, %w1s) : tensor<64x8xf32>

%w2s = slice 0 %w2[%rM] : tensor<8x8xf32>

%x2ss = matmul(%x1ss, %w2s) : tensor<64x8xf32>

yield %x2ss : tensor<64x8xf32>

}

yield %x2s : tensor<64x8xf32>

}

return %r : tensor<256x8xf32> }

Notice how in that final program, both %w1 and %w2 end up only used sliced
across axis "M", even though only %w2 was explicitly value-tiled.

Inference is of paramount practical importance in ML programs. For example,
the optimizer state in a training step function can be left to be partitioned
as the parameters are by inference since parameters and optimizer state flow
to element-wise operations that express the parameter update.

####  5.2.3 Conflicts during propagation

In some situations, it is impossible to propagate the tiling. E.g. when a loop
over some axis would need to be inserted in the scope of an existing loop over
the same axis, which we forbid as nested loops along the same axis cannot be
mapped to meshes, or when multiple (partial) TMR matches are found — a
situation that we refer to as a conflict. Consider:

⬇

func @main(%x: ..., %w1: ...) -> ... {

%xt = loop "B" [#tile<0>] (%rB: range<4>) {

yield (slice 0 %x[%rB])

}

%w1t = loop "B" [#tile<1>] (%rB: range<4>) {

yield (slice 1 %w1[%rB])

}

%x1 = matmul(%xt, %w1t) ...

Here, the operands of the matmul defining %x1 are tiled in a way that matches
two TMR entries: $({{\small\tt\\#tile}}\langle{\tt
0}\rangle,\bot)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 0}\rangle$ and
$(\bot,{{\small\tt\\#tile}}\langle{\tt
1}\rangle)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 1}\rangle$.

PartIR will not attempt to resolve conflicts automatically. Instead, the
canonical solution is to perform the rewriting _incrementally_. For example,
performing value tiling on %x and propagating that choice _before_ value
tiling %w1 would yield the following program:

⬇

func @main(%x:..., %w1: ...) -> ... {

%w1t = loop "B" [#tile<1>] (%rB: range<4>) {

yield (slice 1 %w1[%rB])

}

//below is result of propagation after tiling %x

%r = loop "B" [#tile<0>] (%rB : range<4>) {

%xs = slice 0 %x[%rB] : tensor<64x8xf32>

%x1s = matmul(%xs, %w1t)

...

At this point the TMR entry $(\bot,{{\small\tt\\#tile}}\langle{\tt
1}\rangle)\hookrightarrow{{\small\tt\\#tile}}\langle{\tt 1}\rangle$ matches
the definition of %x1s again. Alas, the operation in hand is already nested
inside a loop over axis "B" and no further propagation is possible – creating
a doubly-nested loop over "B" is invalid. This prioritization of BP over
subsequent parameter sharding is exactly what is needed for the ZeRO [46]
sharding strategies.

The prioritization of rewrites, happening naturally at the boundaries of
PartIR manual tactics makes conflicts be fairly rare, and as a result
dramatically reduces the need for many internal sharding decisions (but does
not entirely remove their need, see Section 8.)

####  5.2.4 PartIR:core actions and multiple axes

Our tiling and propagation actions naturally extend to loop nests over
multiple axes, see Appendix B.

##  6 SPMD code generation

⬇

1 //Below, assume mesh: {x1 : 2, x2 : 4, x3 : 8}

2 %result = all_reduce<@red_fn> <"x1", "x2"> %operand : tensor<16x500x40xf32>
-> tensor<16x500x40xf32>

3 %result = all_slice [{"x1"}, {}, {"x2"}] %operand : tensor<16x500x40xf32> ->
tensor<8x500x10xf32>

4 %result = all_slice [{"x1, "x3"}, {}, {"x2"}] %operand :
tensor<16x500x40xf32> -> tensor<1x500x10xf32>

5 %result = all_gather [{"x1"}, {"x3"}, {}] %operand : tensor<8x10x16xf32> ->
tensor<16x80x16xf32>

6 %result = reduce_scatter<@red_fn> [{"x1"}, {}, {"x2"}] %operand :
tensor<16x500x40xf32> -> tensor<8x500x10xf32>

7 %result = all_to_all {0 -> 1} <"x1","x2"> %operand : tensor<16x32xf32> ->s
tensor<128x4xf32>

Listing 9: PartIR:HLO collectives by example. Relevant dimensions and
attributes (e.g. slicing or gathering axes) have been highlighted.

PartIR:HLO expresses per-device SPMD computation. It includes StableHLO tensor
and control flow ops, but introduces specialized collective ops for
communication, which unlike their low-level HLO counterparts, operate on mesh
axes, i.e. communication spans across device groups defined by a different set
of coordinates along these axes. We now introduce the collectives included in
our implementation with examples in LABEL:lst:mhlo-collectives:

Figure 5:  Demonstration of sequences of all_slice and all_gather collectives
on a mesh {x:2, y:2}. Top: all devices hold the same 2D array; bottom: data is
sliced row-ise along axis "y"; right: data is further sliced column-wise along
axis "x". In each case we give the device-local tensor types.

  * •

An all_reduce (line 2) reduces arrays along one or more mesh axes according to
some reduction function (in the example @red_fn), replicating the tensor in
each device.

  * •

An all_slice (lines 3 and 4) includes an array of axes per dimension – the
axes in which each dimension is sliced. In line 2, the first dimension is
sliced along "x1", the second dimension is not sliced, and the third dimension
is sliced along "x2". In the result array, each dimension size is divided by
the size of (the product of) the slicing axes in this dimension – each device
holds a slice of the original array based on the coordinate tuple along the
slicing axes.

  * •

An all_gather (line 5) is dual to all_slice. Its operand is gathered along the
gathering axes in each dimension. Each dimension size of the result is
multiplied by (the product of) the gathering axes in this dimension. Figure 5
demonstrates the semantics of all_slice and all_gather.

Patterns such as all_slice(all_gather(%x)) or all_slice(all_reduce(%x)) are
fused into all_to_all and reduce_scatter(%x), respectively, by a PartIR:HLO
pass.

  * •

A reduce_scatter (line 6) effectively is the fusion of an all_reduce followed
by all_slice. In particular, this example is the fusion of collectives in
lines 2 and 3.

  * •

An all_to_all (line 7) corresponds to a fusion of an all_gather along a
dimension (0) followed by all_slice over the same sequence of axes in another
dimension (1).

###  6.1 From PartIR:Core to PartIR:HLO

Lowering of PartIR:Core to PartIR:HLO is a type-preserving transformation that
transforms (i) slice ops to all_slice collectives, and (ii) inserts all_gather
or all_reduce collectives in the results of tiling loops. We demonstrate this
process with an example:

⬇

func @main(%x: t<256x8>, %w1: t<8x16>, %w2: t<16x8>) {

%r = loop "B" [#tile<0>] (%rB: range<4>) {

%x2s = loop "M" [#sum] (%rM: range<2>) {

%xs = slice 0 %x[%rB] : tensor<64x8xf32>

%w1s = slice 1 %w1[%rM] : tensor<8x8xf32>

%x1ss = matmul(%xs, %w1s) : tensor<64x8xf32>

%w2s = slice 0 %w2[%rM] : tensor<8x8xf32>

%x2ss = matmul(%x1ss, %w2s) : tensor<64x8xf32>

yield %x2ss : tensor<64x8xf32>

}

yield %x2s : tensor<64x8xf32>

}

%s = loop "B" [#tile<0>] (%rB: range<4>) {

%rs = slice 0 %r[%rB] : tensor<64x8xf32>

%xs = slice 0 %s[%rB] : tensor<64x8xf32>

%w = add %rs %xs : tensor<64x8xf32>

}

return %s : tensor<256x8xf32> }

Mechanically following the lowering rules outlined in the beginning of this
section gives:

⬇

func @main(%x: t<256x8>, %w1: t<8x16>, %w2: t<16x8>) {

//First loop nest.

%xs = all_slice [{"B"},{}] %x : tensor<64x8xf32>

%w1s = all_slice [{},{"M"}] %w1 : tensor<8x8xf32>

%x1ss = matmul(%xs, %w1s) : tensor<64x8xf32>

%w2s = all_slice [{"M"},{}] %w2 : tensor<8x8xf32>

%x2ss = matmul(%x1ss, %w2s) : tensor<64x8xf32>

%x2s = all_reduce <"M"> %x2ss : tensor<64x8xf32>

%r = all_gather [{"B"},{}] %x2s : tensor<256x8xf32>

//Second loop nest.

%rs = all_slice [{"B"}, {}] %r : tensor<64x8xf32>

%xs = all_slice [{"B"}, {}] %x : tensor<64x8xf32>

%w = add %rs %xs : tensor<64x8xf32>

%s = all_gather [{"B"}, {}] %w : tensor<256x8xf32>

return %s : tensor<256x8xf32> }

The all_slice(all_gather)(%xs2) in this example is fused away. Additionally,
the function arguments used by all_slice ops and the results produced by
all_gathers are converted to device-local arrays:

⬇

func @main(%xs: tensor<64x8xf32>, %w1s: tensor<8x8xf32>, %w2s:
tensor<8x8xf32>) {

//First loop nest.

%x1ss = matmul(%xs, %w1s) : tensor<64x8xf32>

%x2ss = matmul(%x1ss, %w2s) : tensor<64x8xf32>

%x2s = all_reduce <"M"> %x2ss : tensor<64x8xf32>

//Second loop nest.

%w = add %x2s %xs : tensor<64x8xf32>

return %w : tensor<64x8xf32> }

Simulation is performed on PartIR:HLO programs, after we have applied
collective ops optimizations and additional rewrites to enable compute-
communication overlap [64].

##  7 Evaluation

PartIR is a partitioning stack for users who care about performance,
expressibility, and predictability. Hence, we evaluate the following claims
about PartIR: (1) its partitioning of models achieves comparable performance
to the SOTA tool, across different systems (7.2); (2) it is predictable on a
variety of models (7.3); (3) it effectively composes tactics, manual and
automated (7.3.1); (4) its propagation reduces the number of decisions a user
makes, and resolves conflicts (7.4); and (5) its overhead is small compared to
compilation time (7.5).

###  7.1 Benchmarking setup

We use the following set of models as case studies: T32 a 32 layers, 5B
parameters Transformer model based on Chinchilla [20]; T48 a 48 layers, 32B
parameters variant of T32; I32T is a 32 layers inference Transformer model
[62, 43] with key-value caching; U-Net is a model used in the reverse process
of a diffusion model [19]; GNS is a Graph Network Simulator [49] model for
molecular property prediction [14].

When measuring runtime, we perform a warmup step and then report the
measurements of $10$ restarts, each performing a $100$ training or inference
steps depending on the model; we repeat this process three times and report
the average. All experiments used TPUv3 [12] and Nvidia A100 [37]; we refer to
them as TPU and GPU, respectively. The models and hardware specifications are
reported in Appendix A.

###  7.2 Partitioned models matches SOTA

We validated that the performance of models partitioned with PartIR is
comparable to that of GSPMD [67] in terms of Model FLOPS Utilization (MFU)
[11] and High-Bandwidth Memory (HBM) usage, as reported by the XLA compiler,
by training the T32 (5B) and T48 (48B) models with both PartIR and GSPMD. For
the T32 model, we tested two configurations: 1) partitioned over 32 TPUs and
2) over 16 GPUs, while the T48 partitioned over 128 TPUs. In the PartIR
version, we partitioned both models using a schedule of four tactics
(BP+MP+Z3+EMB) and relied on equivalent expert sharding annotations for the
GSPMD version. The results in Table 1 show that the performance of PartIR is
on par with that of GSPMD, with negligible improvements in MFU ($\approx 1\%$)
in TPU configurations and a small regression ($\approx 2\%$) in the GPU
configuration. The HBM usage are comparable too with a regression of $\approx
1\%$ on GPUs.

|  | MFU (%) | HBM (GB)
---|---|---|---
Mesh | Size | PartIR | GSPMD | PartIR | GSPMD
16x2 TPU | 5B | 58.5 | 58.3 | 14.38 | 14.38
32x4 TPU | 32B | 52.3 | 52.2 | 14.48 | 14.48
8x2 GPU | 5B | 42.2 | 42.9 | 27.02 | 26.73
Table 1:  PartIR and GSPMD comparison in terms of MFU (higher is better) and
HBM usage (lower is better) on GPUs and TPUs.

###  7.3 Composability and predictability

We demonstrate how PartIR users can compose well-known strategies using manual
tactics while getting the expected communication collectives (from the
respective papers), as reported after every tactic is applied. We partitioned
four different models with different schedules and inspected the resulting
number of collectives, shown in Table 2.

Model | Schedule | AG | AR | RS | A2A
---|---|---|---|---|---
T32 | BP | 0 | 290 | 0 | 0
BP+MP | 0 | 418 | 0 | 0
BP+MP+Z2 | 129 | 289 | 129 | 0
BP+MP+Z3 | 259 | 289 | 129 | 0
BP+MP+Z3+EMB | 515 | 354 | 257 | 0
MP | 0 | 128 | 0 | 0
EMB | 256 | 193 | 128 | 0
IT32 | BP | 0 | 0 | 0 | 0
BP+MP | 0 | 98304 | 0 | 0
BP+MP+MQ | 64 | 98304 | 0 | 98240
MP | 0 | 98304 | 0 | 0
UNet | BP | 0 | 503 | 0 | 0
BP+Z2 | 517 | 2 | 501 | 0
BP+Z3 | 799 | 2 | 501 | 0
GNS | ES | 0 | 423 | 0 | 0
Table 2:  Number of collectives introduced by different schedules. AG:
AllGather, AR: AllReduce, RS: ReduceScatter, A2A: All2All.

T32 has 289 parameter tensors (9 for each block + embeddings). With batch
parallelism, we expect one AR for each parameter gradient tensor and one AR
for the loss value. The resulting 290 ARs add up with the 4 ARs per layer that
Megatron [55] requires when composing both strategies together. Both Z2 and Z3
[46] partition parameter gradients and optimizer states along the batch axis
(of embeddings and four-parameter tensors per layer in this experiment),
resulting in 129 of the existing ARs becoming RSs, and the introduction of one
or two AG per parameter tensor in the case of Z2 or Z3, respectively. The
embedding partition strategy (EMB) partitions the embedding tensor along the
d_model dimension, which has the effect of partitioning activations, which in
turn helps leverage collective-compute overlapping [64]. The reasoning is
similar for IT32 and UNet. Note that IT32 is an inference-only benchmark and
does not require any AR for batch parallelism. Moreover, the serving loop in
IT32 greatly amplifies the number of collectives. Multi-query sharding (MQ)
[43] over the batch axis introduces an AR and two A2A per layer, except for
the final loop, which requires an extra AG.

Finally, GNS is made of nodes connected by edges; we partitioned it using edge
sharding (ES) [18], which partitions the GNS’s edges, creating subgraphs
around the edges and distributes them to devices in the network while
replicating the associated GNS’s nodes. Every message passing and propagation
through the GNS introduces a collective to communicate updates from neighbors
in the GNS’s graph. Thus, we expect 2 AR per messaging passing (24) through
each MLP layer (5) of every node and one for the global feature aggregator,
the molecular GNS has an additional 2 AR for the graph encoder and one final
one at the decoder.

####  7.3.1 Partial or full automation in schedules

PartIR users may not wish to partition their models manually, and not every
model architecture has a well-studied set of partitioning strategies. Thus,
users may explore different degrees of automation: fully manual, partially
automatic, or fully automatic. Figure 6 shows the actual runtime results of
combining manual and automatic tactics and using them to partition the models
for $32$ TPU devices.

Figure 6: Training runtime (lower is better) on a 8x4 TPU mesh.

Using automatic partition can alleviate the burden of manually partitioning
T32, where AllAuto results in a partition with comparable performance to a
fully manual schedule. While combining manual with automatic partition gives
us performance improvements for UNet and GNS, it can also come with a
performance penalty: e.g., BP+Auto+Z3 in T32 results in slower runtimes than
the fully manual schedule.

###  7.4 Resolving conflicts with incrementality

Figure 7: Relative slowdown compared to PartIR (higher is worse) for UNet
partitioned on a {8:batch, 2:model} TPU mesh.

We show the importance of incrementality as a solution for compiler-internal
conflicts in Fig. 7. We compared PartIR against PartIR-st (Single Tactic),
which amalgamates all tactics of a schedule into a single tactic (i.e., no
propagation in between tactics); GSPMD, with expert-defined sharding
constraints baked into the model code to resolve conflicts; and GSPMD–, which
does not use internal sharding constraints for conflict resolution. We
evaluated this on UNet with BP and Z2/Z3 – two partitioning strategies that
cause conflicts discussed in Section 5.2.3. Furthermore, we show the addition
of a Megatron-like [54] MP tactic to UNet along the model axis. We performed a
similar experiment with the transformer benchmarks used in 7.2, and (somewhat
surprisingly) obtained comparable results of GSPMD and GSPMD–, potentially
because GSPMD internal heuristics to resolve conflicts have been highly tuned
for this kind of workload.

PartIR achieves faster runtime compared to the baselines; (2) incrementality
is fundamental for our design, PartIR-st generated programs that exceeded the
device memory limit without a natural way to resolve the conflicts generated
by these tactics; (3) in multi-axes settings, even when we do not expect the
tactic to cause conflicts (e.g., BP+MP), the lack of incrementality can cause
a performance regression in the GSPMD case; (4) in comparison, without
sharding constraints, GSPMD– produces programs that fit but are noticeably
slower compared to PartIR; and finally (5) once we did manage to obtain the
optimal sharding constraints for GSPMD (which involved human labor to identify
these places where annotations might be needed) the performance of the two
systems was similar.

###  7.5 PartIR partition time evaluation

Figure 8: PartIR partitioning vs. overall compilation time.

In Figure 8, we show that PartIR partitioning time is a small percentage (max
of 14%) of the overall compilation time (XLA compilation), which is important
for an interactive user workflow. Furthermore, these models usually train for
several days or even months; hence, even a very long partitioning time would
be acceptable for our users.

##  8 Discussion and limitations

##### Padding and spatial partitioning

Value tiling and propagation via loops assume that the number of devices in an
axis must exactly divide a partitioned dimension; otherwise, propagation gets
blocked. In addition, partitioning convolutional NNs on spatial dimensions
requires communicating with neighboring devices [10, 21]. PartIR has limited
support for these features (via slice and loops with padding, striding, and
special lowering) as we have not encountered a high demand. GSPMD [67] offers
a more comprehensive solution.

##### Reshape support

Reshape ops pose a challenge: Consider a reshape from tensor<16xf32> to
tensor<4x4xf32>. First, assume a 1D mesh of 4 devices, {"A":4}. It is
straightforward to express a tiled version of the reshape as a PartIR loop:

⬇

%y = loop "A" [tile<0>] (%r:range<4>) {

%s = slice 0 %x[%r]: tensor<16xf32> -> tensor<4xf32>

%r = reshape %s : tensor<4xf32> -> tensor<1x4xf32>

yield %r } : tensor<4x4xf32>

Now, assume that the mesh axis spans over 8 devices, such that it becomes too
large to return chunks of the tensor<4x4xf32> array, and propagation will get
blocked. Intuitively the solution is to logically reshape mesh {"A":8} to
{"A1":4, "A2":2}, in which case propagation over the "A1" axis will work.
Unfortunately, PartIR does not consider mesh transformations like that in the
middle of propagation. By contrast, GSPMD [67] can, during propagation,
manipulate the mapping of logical IDs to data, as it defines the sharding
directly in terms of logical device IDs, which comes at the cost of lower-
level reasoning about the sharding representation.

##### Program rewriting versus annotations

PartIR departs from the traditional view in HPC of distribution as a data
layout problem. Following that approach, GSPMD [67] separates the propagation
of sharding annotations from code generation, which deals with inserting and
optimizing collectives. By contrast, PartIR rewrites the program
incrementally, which is less efficient (esp. as part of the automatic tactic),
and relies on pattern matching on the IR loop structure for propagation, which
is brittle. On the upside, the PartIR loops and rewriting help with
correctness verification since loops reflect the semantics in the IR and allow
for temporal interpretation and testing without actual partitioning.

##### Explicitly replicating values

PartIR propagation greedily partitions the program (c.f., Section 5.2.2),
which may partition tensors the user wants to replicate. For example, by
partitioning the optimizer state, propagation may also shard the parameters,
which is not desired for Z2, where parameters must be replicated (Section
2.2). For this reason, PartIR exposes an explicit atomic<value, axis> action,
which creates trivial loop rewrites whose sole purpose is to block propagation
and keep a value replicated. For example atomic<%x, "M">) replaces %x with the
loop:

⬇

%xr = loop "M" [any] (%r:range<...>) { yield %x; }

The any attribute works as a “consensus” monoid – all devices are guaranteed
by construction to compute the same value, and the loop range index cannot be
used in the body.

##### Model-internal annotations.

Although sequentialization of decisions lets us handle most conflicts, as seen
in Section 7.4, we are still left with corner cases. For example, consider a
component of matrix diagonalization, where a matrix is multiplied by its
transpose:

⬇

func @main(%x: tensor<(256x256xf32>) -> ... {

%tx = transpose %x {dims=[1,0]} : ...

%y = matmul(%x, %t) ...

If we shard %x on dimension 0 along a given mesh axis "M", then %t (its
transpose) is sharded on dimension 1. That is a propagation “conflict” and
prevents the matrix multiplication sharding. To resolve this conflict, users
must apply an atomic action on the intermediate tensor %tx before partitioning
it, which is allowed as long as users name them, using a primitive called tag.
For example:

⬇

func @main(%x: tensor<(256x256xf32>) {

%tx = transpose %x {dims=[1,0]}

%tag_tx = tag "transposed" %tx

%atomic = loop "M" [noop] (%r:range<...>) { yield %tag_tx; }

%y = matmul(%x, %atomic) ...

As expected, the final partitioned multiplication requires an all_gather for
its second operand:

⬇

func @main(%x: tensor<(16x256xf32>) {

%tx = transpose {dims=[1,0]} %x : tensor<256x16xf32>

%gx = all_gather [{},{"M"}] %tx : tensor<256x256xf32>

%y = matmul(%x, %gx) : tensor<16x256xf32> ...

##  9 Related work

The fast growth of ML models opened the floor to new partitioning tools [53,
47], as frameworks like TensorFlow [1] and PyTorch [41] support only data and
limited model parallelism. JAX [4] provides two main partitioning APIs:
jax.jit, which surfaces GSPMD [67] (which itself builds on GShard [32]); and
shard_map, which disables GSPMD and lets users manually perform communication
across devices. PartIR builds on top of some of JAX’s partitioning
infrastructure (e.g., meshes), but instead of users annotating inputs, outputs
and intermediate values with sharding annotations or communication collectives
in the case of shard_map, our users define their partitioning strategies using
schedules, which resemble ideas found in kernel-generating schedule-based
systems [29, 68, 45, 9, 16]. Although we took inspiration from Halide
schedules [45], PartIR differs in that they are applied to entire programs
that span multiple devices rather than generating small program kernels for a
single device.

Similar to Lift [56] and RISE [17], PartIR is inspired by functional-style IRs
[57] that transforms programs by equality-based rewriting [63] and similar to
the functional array representation in [42]. In contrast to DaCe [3], which
allows the user to apply graph transformer interactively by employing a
representation based on data flow graphs. DistIR [50] exposes a Python API
with explicit device placement and point-to-point communication to program the
distribution directly. PartIR:Core abstracts away distribution and deals with
it during SPMD lowering.

##  10 Conclusion and Future Work

PartIR is an MLIR-based compiler that enables programmers and automatic tools
to partition tensor programs effortlessly. Users interact with PartIR using a
schedule of tactics; each tactic is translated into a series of rewrite
actions to be applied by the compiler. It is implemented as a stack of IRs to
express partitioning at different levels of abstraction. It includes a
powerful propagation system that, because of its incremental design,
partitions programs in a predictable way without relying on cost-based
heuristics to handle conflicts. In the future, we want to support MPMD
partitioning across multiple device meshes to support pipelining [22, 33, 34]
or heterogeneous device clusters, which will require new sets of sharding and
optimization tactics (e.g., scheduling.)

## References

  * [1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng.  TensorFlow: A system for Large-Scale machine learning.  In 12th USENIX Symposium on Operating Systems Design and Implementation (OSDI 16), pages 265–283, Savannah, GA, November 2016. USENIX Association.
  * [2] Sami Alabed, Dominik Grewe, Juliana Franco, Bart Chrzaszcz, Tom Natan, Tamara Norman, Norman A. Rink, Dimitrios Vytiniotis, and Michael Schaarschmidt.  Automatic discovery of composite spmd partitioning strategies in partir, 2022.
  * [3] Tal Ben-Nun, Johannes de Fine Licht, Alexandros Nikolaos Ziogas, Timo Schneider, and Torsten Hoefler.  Stateful dataflow multigraphs: A data-centric model for performance portability on heterogeneous architectures.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’19, 2019.
  * [4] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018\.
  * [5] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  CoRR, abs/2005.14165, 2020.
  * [6] Emanuele Bugliarello, Aida Nematzadeh, and Lisa Anne Hendricks.  Weakly-supervised learning of visual relations in multimodal pretraining.  In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Singapore, December 2023. Association for Computational Linguistics.
  * [7] Emanuele Bugliarello, Laurent Sartran, Aishwarya Agrawal, Lisa Anne Hendricks, and Aida Nematzadeh.  Measuring progress in fine-grained vision-and-language understanding.  In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1559–1582, Toronto, Canada, July 2023. Association for Computational Linguistics.
  * [8] Veronika Samborska Charlie Giattino, Edouard Mathieu and Max Roser.  Data Page: Computation used to train notable artificial intelligence systems.  https://ourworldindata.org/grapher/artificial-intelligence-training-computation, 2023\.  Retrieved from [online resource].
  * [9] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Haichen Shen, Eddie Q. Yan, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy.  TVM: end-to-end optimization stack for deep learning.  CoRR, abs/1802.04799, 2018.
  * [10] Youlong Cheng and HyoukJoong Lee.  Train ml models on large images and 3d volumes with spatial partitioning on cloud tpus | google cloud blog.  https://cloud.google.com/blog/products/ai-machine-learning/ train-ml-models-on-large-images-and-3d-volumes -with-spatial-partitioning-on-cloud-tpus, Sep 2019.  [Accessed 06-12-2023].
  * [11] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel.  Palm: Scaling language modeling with pathways, 2022.
  * [12] Google Developers.  Cloud TPU System Architecture.  https://cloud.google.com/tpu/docs/system-architecture-tpu-vm, 2023\.  [Last updated 2023-11-06 UTC.].
  * [13] Carl Doersch, Yi Yang, Mel Vecerik, Dilara Gokay, Ankush Gupta, Yusuf Aytar, Joao Carreira, and Andrew Zisserman.  Tapir: Tracking any point with per-frame initialization and temporal refinement.  ICCV, 2023.
  * [14] Jonathan Godwin, Michael Schaarschmidt, Alexander L Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia.  Simple GNN regularisation for 3d molecular property prediction and beyond.  In International Conference on Learning Representations, 2022.
  * [15] Google XLA team.  XLA: Optimizing compiler for machine learning, 2017.
  * [16] Bastian Hagedorn, Archibald Samuel Elliott, Henrik Barthels, Rastislav Bodik, and Vinod Grover.  Fireiron: A data-movement-aware scheduling language for gpus.  In Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques, PACT ’20, page 71–82, New York, NY, USA, 2020. Association for Computing Machinery.
  * [17] Bastian Hagedorn, Johannes Lenfers, Thomas Kundefinedhler, Xueying Qin, Sergei Gorlatch, and Michel Steuwer.  Achieving high-performance the functional way: A functional pearl on expressing high-performance optimizations as rewrite strategies.  Proc. ACM Program. Lang., 4(ICFP), aug 2020.
  * [18] Yilin He, Chaojie Wang, Hao Zhang, Bo Chen, and Mingyuan Zhou.  Edge partition modulated graph convolutional networks, 2022.
  * [19] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
  * [20] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre.  Training compute-optimal large language models, 2022.
  * [21] Le Hou, Youlong Cheng, Noam Shazeer, Niki Parmar, Yeqing Li, Panagiotis Korfiatis, Travis M Drucker, Daniel J Blezek, and Xiaodan Song.  High resolution medical image analysis with spatial partitioning.  arXiv preprint arXiv:1909.03108, 2019.
  * [22] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, and Zhifeng Chen.  Gpipe: Efficient training of giant neural networks using pipeline parallelism.  In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 103–112, 2019.
  * [23] Daniel Jarrett, Miruna Pislar, Michiel A Bakker, Michael Henry Tessler, Raphael Koster, Jan Balaguer, Romuald Elie, Christopher Summerfield, and Andrea Tacchetti.  Language agents as digital representatives in collective decision-making.  In NeurIPS 2023 Foundation Models for Decision Making Workshop, 2023\.
  * [24] Zhihao Jia, Matei Zaharia, and Alex Aiken.  Beyond data and model parallelism for deep neural networks.  In Ameet Talwalkar, Virginia Smith, and Matei Zaharia, editors, Proceedings of Machine Learning and Systems 2019, MLSys 2019, Stanford, CA, USA, March 31 - April 2, 2019. mlsys.org, 2019.
  * [25] Norm Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, et al.  Tpu v4: An optically reconfigurable supercomputer for machine learning with hardware support for embeddings.  In Proceedings of the 50th Annual International Symposium on Computer Architecture, pages 1–14, 2023.
  * [26] Norman P. Jouppi, Doe Hyun Yoon, Matthew Ashcraft, Mark Gottscho, Thomas B. Jablin, George Kurian, James Laudon, Sheng Li, Peter Ma, Xiaoyu Ma, Thomas Norrie, Nishant Patil, Sushma Prasad, Cliff Young, Zongwei Zhou, and David Patterson.  Ten lessons from three generations shaped google’s tpuv4i : Industrial product.  In 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), pages 1–14, 2021.
  * [27] Norman P. Jouppi, Cliff Young, Nishant Patil, David A. Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, Richard C. Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon.  In-datacenter performance analysis of a tensor processing unit.  CoRR, abs/1704.04760, 2017.
  * [28] Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
  * [29] Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe.  The tensor algebra compiler.  Proc. ACM Program. Lang., 1(OOPSLA), oct 2017.
  * [30] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro.  Reducing activation recomputation in large transformer models.  Proceedings of Machine Learning and Systems, 5, 2023.
  * [31] Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko.  Mlir: Scaling compiler infrastructure for domain specific computation.  In 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pages 2–14, 2021.
  * [32] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  CoRR, abs/2006.16668, 2020.
  * [33] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia.  Pipedream: generalized pipeline parallelism for DNN training.  In Tim Brecht and Carey Williamson, editors, Proceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP 2019, Huntsville, ON, Canada, October 27-30, 2019, pages 1–15. ACM, 2019.
  * [34] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia.  Memory-efficient pipeline-parallel dnn training.  In International Conference on Machine Learning, pages 7937–7947. PMLR, 2021.
  * [35] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia.  Efficient large-scale language model training on GPU clusters.  CoRR, abs/2104.04473, 2021.
  * [36] NVIDIA.  Nvidia nvlink and nvswitch.  https://www.nvidia.com/en-gb/data-center/nvlink/, 2021.  Accessed: 2021-10-07.
  * [37] Nvidia.  NVIDIA A100 Tensor Core GPU.  https://www.nvidia.com/en-gb/data-center/a100/, 2023.  [Last updated 2023-11-06 UTC.].
  * [38] OpenXLA.  OpenXLA: A machine learning compiler for GPUs, CPUs, and ML accelerators .  https://github.com/openxla/xla, 2023.  [Last updated 2023-11-06 UTC.].
  * [39] OpenXLA.  StableHLO: Backward compatible ML compute opset inspired by HLO/MHLO.  https://github.com/openxla/stablehlo, 2023.  [Last updated 2023-11-06 UTC.].
  * [40] Shoumik Palkar and Matei Zaharia.  Optimizing data-intensive computations in existing libraries with split annotations.  In Proceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP ’19, page 291–305, New York, NY, USA, 2019. Association for Computing Machinery.
  * [41] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  In H. Wallach, H. Larochelle, A. Beygelzimer, F. d‘Alché Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
  * [42] Adam Paszke, Daniel D. Johnson, David Duvenaud, Dimitrios Vytiniotis, Alexey Radul, Matthew J. Johnson, Jonathan Ragan-Kelley, and Dougal Maclaurin.  Getting to the point: Index sets and parallelism-preserving autodiff for pointful array programming.  Proc. ACM Program. Lang., 5(ICFP), aug 2021.
  * [43] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean.  Efficiently scaling transformer inference.  Proceedings of Machine Learning and Systems, 5, 2023.
  * [44] Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, H. Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John Mellor, Irina Higgins, Antonia Creswell, Nat McAleese, Amy Wu, Erich Elsen, Siddhant M. Jayakumar, Elena Buchatskaya, David Budden, Esme Sutherland, Karen Simonyan, Michela Paganini, Laurent Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean-Baptiste Lespiau, Maria Tsimpoukelli, Nikolai Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Toby Pohlen, Zhitao Gong, Daniel Toyama, Cyprien de Masson d’Autume, Yujia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew Johnson, Blake A. Hechtman, Laura Weidinger, Iason Gabriel, William S. Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem Ayoub, Jeff Stanway, Lorrayne Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving.  Scaling language models: Methods, analysis & insights from training gopher.  CoRR, abs/2112.11446, 2021.
  * [45] Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe.  Halide: A language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines.  In Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI ’13, page 519–530, New York, NY, USA, 2013. Association for Computing Machinery.
  * [46] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.  Zero: Memory optimization towards training A trillion parameter models.  CoRR, abs/1910.02054, 2019.
  * [47] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He.  Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters.  In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD ’20, page 3505–3506, New York, NY, USA, 2020. Association for Computing Machinery.
  * [48] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He.  Zero-offload: Democratizing billion-scale model training.  CoRR, abs/2101.06840, 2021.
  * [49] Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia.  Learning to simulate complex physics with graph networks.  In International Conference on Machine Learning, pages 8459–8468. PMLR, 2020.
  * [50] Keshav Santhanam, Siddharth Krishna, Ryota Tomioka, Andrew Fitzgibbon, and Tim Harris.  Distir: An intermediate representation for optimizing distributed neural networks.  In Proceedings of the 1st Workshop on Machine Learning and Systems, EuroMLSys ’21, page 15–23, New York, NY, USA, 2021. Association for Computing Machinery.
  * [51] Michael Schaarschmidt, Dominik Grewe, Dimitrios Vytiniotis, Adam Paszke, Georg Stefan Schmid, Tamara Norman, James Molloy, Jonathan Godwin, Norman Alexander Rink, Vinod Nair, et al.  Automap: Towards ergonomic automated parallelism for ml models.  arXiv preprint arXiv:2112.02958, 2021.
  * [52] Michael Schaarschmidt, Morgane Riviere, Alex M. Ganose, James S. Spencer, Alexander L. Gaunt, James Kirkpatrick, Simon Axelrod, Peter W. Battaglia, and Jonathan Godwin.  Learned force fields are ready for ground state catalyst discovery, 2022\.
  * [53] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, Ryan Sepassi, and Blake Hechtman.  Mesh-TensorFlow: Deep learning for supercomputers.  In Neural Information Processing Systems, 2018.
  * [54] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism.  CoRR, abs/1909.08053, 2019.
  * [55] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism, 2020.
  * [56] Michel Steuwer, Christian Fensch, Sam Lindley, and Christophe Dubach.  Generating performance portable code using rewrite rules: From high-level functional expressions to high-performance opencl code.  In Proceedings of the 20th ACM SIGPLAN International Conference on Functional Programming, ICFP 2015, page 205–217, New York, NY, USA, 2015\. Association for Computing Machinery.
  * [57] Michel Steuwer, Toomas Remmelg, and Christophe Dubach.  Lift: A functional data-parallel ir for high-performance gpu code generation.  In Proceedings of the 2017 International Symposium on Code Generation and Optimization, CGO ’17, page 74–85. IEEE Press, 2017.
  * [58] Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, Andrew Poulton, Viktor Kerkez, and Robert Stojnic.  Galactica: A large language model for science.  arXiv preprint arXiv:2211.09085, 2022.
  * [59] Philippe Tillet, Hsiang-Tsung Kung, and David Cox.  Triton: an intermediate language and compiler for tiled neural network computations.  In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019.
  * [60] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [61] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, Xi Luo, Dheevatsa Mudigere, Jongsoo Park, Misha Smelyanskiy, and Alex Aiken.  Unity: Accelerating DNN training through joint optimization of algebraic transformations and parallelization.  In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 267–284, Carlsbad, CA, July 2022. USENIX Association.
  * [62] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  CoRR, abs/1706.03762, 2017.
  * [63] Eelco Visser, Zine-el-Abidine Benaissa, and Andrew Tolmach.  Building program optimizers with rewriting strategies.  In Proceedings of the Third ACM SIGPLAN International Conference on Functional Programming, ICFP ’98, page 13–26, New York, NY, USA, 1998. Association for Computing Machinery.
  * [64] Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ilbeyi, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, Sameer Kumar, Tongfei Guo, Yuanzhong Xu, and Zongwei Zhou.  Overlap communication with dependent computation via decomposition in large deep learning models.  In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. ACM, 2023.
  * [65] Siyu Wang, Yi Rong, Shiqing Fan, Zhen Zheng, Lansong Diao, Guoping Long, Jun Yang, Xiaoyong Liu, and Wei Lin.  Auto-map: A DQN framework for exploring distributed execution plans for DNN workloads.  CoRR, abs/2007.04069, 2020.
  * [66] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Hongjun Choi, Blake Hechtman, and Shibo Wang.  Automatic cross-replica sharding of weight update in data-parallel training, 2020.
  * [67] Yuanzhong Xu, HyoukJoong Lee, Dehao Chen, Blake A. Hechtman, Yanping Huang, Rahul Joshi, Maxim Krikun, Dmitry Lepikhin, Andy Ly, Marcello Maggioni, Ruoming Pang, Noam Shazeer, Shibo Wang, Tao Wang, Yonghui Wu, and Zhifeng Chen.  GSPMD: general and scalable parallelization for ML computation graphs.  CoRR, abs/2105.04663, 2021.
  * [68] Rohan Yadav, Alex Aiken, and Fredrik Kjolstad.  Distal: The distributed tensor algebra compiler.  In Proceedings of the 43rd ACM SIGPLAN International Conference on Programming Language Design and Implementation, PLDI 2022, page 286–300, New York, NY, USA, 2022. Association for Computing Machinery.
  * [69] Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Teh, Alvaro Sanchez-Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin.  Pre-training via denoising for molecular property prediction, 2022.
  * [70] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al.  Pytorch fsdp: experiences on scaling fully sharded data parallel.  arXiv preprint arXiv:2304.11277, 2023.
  * [71] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Joseph E. Gonzalez, and Ion Stoica.  Alpa: Automating inter- and intra-operator parallelism for distributed deep learning.  CoRR, abs/2201.12023, 2022.

##  Appendix A Experiments extended

###  A.1 MFU

We evaluated the MFU to compare SOTA performance in Section 7.2, here we
document what is MFU. MFU is defined as a ratio of actual FLOPs per second and
theoretical FLOPs per second:

$100\times\frac{\text{model FLOPs}/\text{step time
(s)}}{\\#\text{devices}\times\text{peak FLOPs per second}}$

An MFU of 100% indicates that the model is utilizing the full computational
capacity of the hardware.

###  A.2 Hardware details

Our experiments were run on Nvidia A100 [37] as the GPU baseline and TPUv3
[12] for the TPU baseline. Here we outline their specs as outlined in their
respective websites:

##### Nvidia A100

We use the A100 40GB memory HBM2 version; it performs 156 TFLOPS on float32
and 312 TFLOPS for bfloat16. The A100s are connected using Nvidia’s NVLInk,
which can do 600GB/s data transfer.

##### TPUv3

For the TPU experiments, we use the TPUv3 chips; each chip comes with two
tensor cores. Each tensor core has 16GiB HBM2 memory capacity, each capable of
61.5 TFLOPs on float32 and 123 TFLOPs on float16. The chips are connected over
four links, each capable of doing 70GB/s data transfer [25].

###  A.3 Benchmark details

The details of the benchmarks used are as follows:

U-Net

A model variant used in the reverse process of a diffusion model [19]. It uses
9 residual blocks for the down-sampling convolutions and 12 for up-sampling,
and between them, there are two residual blocks and one attention layer. The
attention layer has 16 heads and a resolution of $16\times 8\times 4$. Pairs
of convolutions have a hidden channel size 4 times larger than the
input/output channel size. This allows for efficient partitioning along the
channel dimensions.

GNS

A Graph Network Simulator [49] model used in molecular property prediction
[14], configured with 5 MLP layers of hidden size 1024. The network is
configured with 24 message-passing steps and a latent size of 512. We set the
number of nodes to 2048 and vary the number of edges between 8192 and 65536 in
powers of 2.

T32)

A 32 layers Transformer model based on Chinchilla [20], with additional
normalization layer. We configure with input batch size 48, 32 heads and
$d_{model}=4096$, with the standard 32k vocabulary size for the embedding.

T48)

A scaled-up variant of T32, scaled to 48 layers, and configured with input
batch size 64, 64 heads and $d_{model}=8192$, with the standard 32k vocabulary
size for the embedding.

All models were trained using the Adam optimizer [28]. We turned off XLA
rematerialization in all runs.

###  A.4 Composing automatic and manual tactic

Table 3 provides the the full metrics of using AutomaticPartition to find
sharding strategies.

|  | Mem | Est.Runtime(ms) | AG | AR | RS | A2A
---|---|---|---|---|---|---|---
Model | Strategy |  |  |  |  |  |
GNS | ES | 10379.47 | 294.13 | 0 | 423 | 0 | 0
ES+AutoMP | 8424.38 | 146.43 | 220 | 752 | 97 | 0
ES+AutoBP | 8141.38 | 101.47 | 72 | 679 | 72 | 0
AllAuto | 2508.92 | 118.12 | 476 | 854 | 272 | 0
IT32 | BP | 18302.16 | 1139.31 | 0 | 0 | 0 | 0
BP+MP | 5607.73 | 1447.83 | 0 | 98304 | 0 | 0
BP+MP+MQ | 5439.73 | 1498.92 | 64 | 98304 | 0 | 98240
MP | 5151.44 | 4327.35 | 0 | 98304 | 0 | 0
T32 | BP | 100343.69 | 4803.34 | 0 | 290 | 0 | 0
BP+AutoMP+Z3 | 40472.80 | 4902.41 | 547 | 161 | 353 | 0
BP+MP | 59826.45 | 4856.25 | 0 | 418 | 0 | 0
BP+MP+Z2 | 50124.45 | 4856.25 | 129 | 289 | 129 | 0
BP+MP+Z3 | 45068.63 | 4960.32 | 259 | 289 | 129 | 0
BP+MP+Z3+EMB | 47541.60 | 4946.35 | 515 | 354 | 257 | 0
MP | 177148.23 | 10837.42 | 0 | 128 | 0 | 0
EMB | 176974.51 | 10934.86 | 256 | 193 | 128 | 0
UNet | BP | 2406.68 | 25.80 | 0 | 503 | 0 | 0
BP+AutoMP | 1693.65 | 20.51 | 162 | 482 | 68 | 0
BP+Z2 | 933.36 | 25.80 | 517 | 2 | 501 | 0
BP+Z3 | 309.48 | 37.73 | 799 | 2 | 501 | 0
AllAuto | 1126.94 | 15.74 | 415 | 565 | 88 | 2
Table 3:  Collective and simulator cost estimation of various tactics

###  A.5 Simulation results

Here, we show how close our simulator can approximate real performance on the
TPUv3 32 device described in Section 7.3.1. While the details of the simulator
we leave to another paper, it is indeed a very simple simulator due to most of
the heavy lifting done by our PartIR:HLO dialect. PartIR:HLO generates device-
local communication ops that reference the mesh and device, and PartIR keeps a
registry of popular compilation devices (and it is easy to extend, requiring
only high-level device specs). Combining both the PartIR:HLO code, which has
tensor shapes, and communication ops on each device, with the target specs,
our simulator iterates over each SPMD context, tracks the live memory, and
counts flops usage for the communication ops also tracks the byte transfers.
Using the target device specs, it would estimate its statistics. This simple
analytical cost model worked well for our AutomaticPartition tactic and user
debugging. While the absolute values are not guaranteed to be correct, the
relative improvements should still be sound. For example, applying a BP
tactic, you would expect the memory and flops usage to be reduced by a factor
of the number of devices in the mesh if the BP was applied correctly. A search
algorithm that will seek to improve the relative partitioning of the model
will also benefit from this simple analytical function.

Hardware measurement follows the methodology described in Section 7.1. All
experiments were run on a 32 TPUv3 as described in the evaluation section
Section 7.3.1.

####  A.5.1 Runtime estimation

First, we look at the difference between the estimated and actual runtime for
the various strategies described in the paper. Figure 9 shows the difference
between the estimated and actual measured seconds on hardware. First, we
notice that the error is within acceptable range (sub milliseconds) for
specific configurations, especially for GNS and UNET. The simulator
underestimated T32’s most complicated strategy by 3 seconds at the worst case;
these often can be attributed to the layout pass of XLA and XLA optimization.
IT32, on the other hand, the simulator overestimated its runtime (meaning we
thought it would be slower than it is); this is due to the key-value caching
optimization we implemented in the model; without caching, the numbers are
much closer to the T32 error range. Overall, both automatic tools and human
users use the estimated runtime as a proxy for relative improvement of
applying the sharding strategy rather than absolute value estimates.

Absolute value estimates require a learned cost model that learns what the
backend compiler optimizer will do, different network topology profiles, and
various specialized ops performances (e.g., caching). Improving the simulator
is beyond the scope of our work, but PartIR allows the end to plug in any
simulator, and we are currently looking into a learned cost model for this
exact reason.

Figure 9: PartIR’s simulator runtime estimation compared to the actual
measured runtime, in seconds, closer to the zero better.

####  A.5.2 Memory estimation

Related to runtime estimation is memory estimation; this is more important for
automatic tools to learn if any proposed solution does not fit in the given
device specification. We implement a live range analysis of a tensor usage in
a given SPMD context at the PartIR:HLO level, where we follow a tensor as long
as it is being used; we also implement a simple fusion heuristic that will
predict what the backend compiler will do to the tensor. The results are shown
in Figure 10. Here, we can see the results are much closer to 0 (i.e., no
error); often, we prefer to over-estimate the memory usage to discourage
solutions close to the boundaries. XLA and other backend compilers have
specialized optimization when a model is on the boundaries - such as remating
parts of the computation or reducing the fusion of ops. The simulator forces
the automatic partitioners to avoid these regions and focus on models that fit
comfortably, not to kick off aggressive memory optimization from the backend
compiler.

Figure 10: PartIR’s simulator estimates of memory compared to the actual
measured memory, in GB, closer to the zero better.

####  A.5.3 Partitioning time

We showed the partitioning time of manual partitioning tactics in Section 7.5;
here, we show the time it takes to run AutomaticPartition. AutomaticPartition
tactic depends on the algorithm implementation used for the search; we
implemented an algorithm similar to the one described in [51, 2]. As PartIR is
agnostic to the optimization algorithm, we are looking at different algorithms
to implement, such as ILP solvers or RL-based solutions. Different algorithms
have different search profiles, and naturally, by increasing the number of
axes (and, as a result, the number of available decisions to make), PartIR
automatic partitioning time does increase, as can be seen in Figure 11. Even
in the worst-case scenario, we see a search time of 1250 seconds, which might
be a lot compared to the manual tactic (hence we allow composing manual and
automatic tactics and benefit from both). However, we want to argue that: 1)
the search time reduces the cognitive workload on the ML practitioner and
saves them time from writing their sharding strategies and testing them; 2)
these models train for weeks on end, a search time of 20 minutes is within the
amortized acceptable range, especially when the shardings found can outperform
known strategies; 3) we are actively developing the search mechanism and
exploring different techniques, we expect this to only get faster with time.

Figure 11: PartIR’s automatic partitioning tactic search time compared to
manual partitioning.

###  A.6 Schedules

Here, we provide the tactics we used for all the experiments; we note that the
schedule is simply a list of these tactics as they appear in the order. For
example, to replicate the schedule BP+MP, define the tactic as bellow, then
pass to partir.jit(fun, schedule=[bp, mp]).

##### Batch parallelism and zero

Regardless of the module, both tactics can easily be expressed as:

  * •

BP: shard the 0th dimension of the inputs, assuming the NN’s update function
takes X where its 0th dimension is the batch dimension:

⬇

BP = ManualPartition(inputs={"X": 0}, axis="batch")

  * •

Z3: shard the 0th dimension of the gradients, and optimizer state, but
replicate the parameters:

⬇

Z2 = ManualPartition(inputs={’params’: partir.REPLICATED,’opt_state’:
partir.FIRST_DIVISIBLE_DIM}, axis="batch")

  * •

Z3: shard the 0th dimension of the parameters, gradients, and optimizer state,

⬇

Z3 = ManualPartition(inputs={’params’: partir.FIRST_DIVISIBLE_DIM,
’opt_state’: partir.FIRST_DIVISIBLE_DIM}, axis="batch")

Our tactics take a callback that applies a sharding based on the parameter
name in the NN, this is useful for Megatron and model parallelism tactics.

##### T32/T48/IT32 tactics

There is no difference in tactics applied to any of the transformer models,
showing the usefulness of separating the model code from the sharding
strategies.

MP: here we apply the Megatron [54] sharding strategy.

⬇

def _model_sharding(param_name):

if not ’w’: #only shard the weights

return UNKNOWN #Let the infer-propagate decides.

if multi_head_attention_regex.contains(param_name):

return 0 #shard the linear weights on 0th dimension

if ’dense_up’ in param_name:

#dense layer on its cols

return 2

if ’qkv_einsum’ in param_name:

#the queries on first dimension

return 1

MP = ManualPartition(inputs={’params’: apply(_model_sharding)}, axis="model")

##### UNet tactics

  * •

MP: here we try to mimic the Megatron sharding strategy.

⬇

def _model_sharding(param_name):

if not ’w’: #only shard the weights

return UNKNOWN #Let the infer-propagate decides.

if multi_head_attention_regex.contains(param_name):

return 0 #shard the linear weights on 0th dimension

if ’conv_residual_block’ in param_name:

#shard the convolutions on their weights not stride.

return 3 if is_2d_conv(param_name) else 2.

MP = ManualPartition(inputs={’params’: apply(_model_sharding)}, axis="model")

##### GNS tactics

  * •

ES: Simply performs search on the batch axis.

⬇

#Shard both the sender and receivers of GNS that is using jraph library.

#predictions is a custom name inside the GNS implementation.

ES = ManualPartition(inputs={’edges’: {"predictions": 0,
"predictions_targets": 0}}, axis="batch")

##### Automatic tactics

These tactics are the simplest.

  * •

AutoBP: Simply performs search on the batch axis.

⬇

AutoBP = AutomaticPartition(axes=["batch"])

  * •

AutoMP: Simply performs search on the model axis.

⬇

AutoMP = AutomaticPartition(axes=["model"])

  * •

AutoAll: Perform search on all axes:

⬇

AutoAll = AutomaticPartition(axes=["batch", "model"])

##  Appendix B PartIR:Core extended

Here we discuss extended details around PartIR:Core omitted from Section 5.
This section mostly focuses on the multi-axis case, how does the propagation
and loop construct handle nesting for multiple axes (a 2D or higher dimensions
mesh).

###  B.1 Multi-axis propagation and deep-tiling

####  B.1.1 Multi-axis analysis for propagation

Propagation becomes involved in the presence of multiple axes. Checking
whether some operand is overly tiled or a result is overly sliced requires a
deeper look into the definitions of operands or uses of results. For example,
consider the following:

⬇

%x = loop "a" [#sum] (%ra: range<4>) {

%t = loop "b" [#tile<0>] (%rb: range<2>) { ... }

yield %t

}

%z = matmul(%x, %y)

To realize that the left matmul operand is tiled, we have to look internally
under the reduction loop.

The situation is even more complex when both tiling loops and slices are
involved:

⬇

%x = loop "a" [#tile<0>] (%ra: range<4>) {

%t = loop "b" [#tile<0>] (%rb: range<2>) { ... }

yield %t

}

%w = loop "a" [#tile<0>] (%ra: range<4>) {

%sx = slice 0 %x[%ra]

%z = matmul(%sx, %y)

...

}

In this situation the left matmul operand is not even the result of a loop op,
it is just a slice! However, if we keep looking backwards, we discover that
the argument %x of the slice instruction is produced by a tiling loop over "a"
that can cancel out the slice, and internally there is yet another tiling
loop, i.e. over "b". Consequently, we could rewrite the matmul as a loop over
"b".

What the examples highlight is that when multiple axes are involved, it is
necessary to traverse nests of loop operations as well as chains of slice
operations to avoid missing relevant matches. A dual situation applies to
slice operations in backward propagation, where we need to apply a forward
analysis starting at the uses of the result of an operation. To deduce whether
there is a match on the TMR entry one would have to look inside the loop over
axis "a" to determine that dimension 0 is indeed tiled across axis "b". A dual
situation arises with overly sliced results, where one may need to look inside
chains of slice operations. For this reason PartIR employs a simple static
analysis to determine the per-dimension excess tiling (or slicing) for the
operands (or results) of an operation.

####  B.1.2 Deep tiling

A big design goal of PartIR is to support incremental partitioning (sometimes
called “recursive partitioning” [67]) that never undoes previous actions.
Consider a value %x that is already sliced (e.g. due to previous value tiling
or propagation) along axis "a":

⬇

%xt = loop "a" [#tile<1>] (%ra: range<4>) {

%xs = slice 1 %x[%ra]) ;

...

}

Imagine a user or a tool needing to further tile the value %x across axis "b"
and dimension $1$. It would be wrong to perform a flat value tiling – instead
we need to gather up any existing sliced uses of %x and apply what we call a
“deep” tiling action. In code:

⬇

//WRONG: would undo previous actions

//%xtt = loop "b" [#tile<1>] (%rb: range<2>) { yield (%slice 1 %x[%rb]) }

//CORRECT: deep tiling over both previous and new axes!

%xtt = loop "a" [#tile<1>] (%ra: range<4>) {

%t = loop "b" [#tile<1>] (%rb: range<2>) {

%xsa = slice 1 %x[%ra]

%xsb = slice 1 %xsa[%rb]

yield %xsb

}

yield %t

}

%xt = loop "a" [#tile<1>] (%ra: range<4>) {

%tmp = slice 1 %xtt[%ra];

...

}

Using the multi-axis analysis from Section B.1.1, PartIR is able to deduce the
full tiling context that needs to be inserted during value tiling, including
previous slicing uses and the new tiling requested. This applies to both user-
initiated value tilings as well as inference-initiated value tilings.