Memory safe computations with XLA compiler
Authors: Artem Artemev
Yuze An
Tilman Roeder
Mark van der Wilk
Venue: NeurIPS 2022 Accept (2022)
Keywords: xla
compiler
gaussian processes
sparse gaussian processes
k-nearest neighbour

Abstract: Software packages like TensorFlow and PyTorch are designed to support linear algebra operations, and their speed and usability determine their success. However, by prioritising speed, they often neglect memory requirements. As a consequence, the implementations of memory-intensive algorithms that are convenient in terms of software design can often not be run for large problems due to memory overflows. Memory-efficient solutions require complex programming approaches with significant logic outside the computational framework. This impairs the adoption and use of such algorithms. To address this, we developed an XLA compiler extension that adjusts the computational data-flow representation of an algorithm according to a user-specified memory limit. We show that k-nearest neighbour, sparse Gaussian process regression methods and Transformers can be run on a single device at a much larger scale, where standard implementations would have failed. Our approach leads to better use of hardware resources. We believe that further focus on removing memory constraints at a compiler level will widen the range of machine learning methods that can be developed in the future.

Memory Safe Computations with XLA Compiler
Artem Artemev Imperial College London
Yuze An Imperial College London yuze an21 @imperial.laaaa.. ac uk
Secondmind
a. ar Artemev200@iperiaal. ac.
Tilman Roeder
Mark van der Wilk Imperial College London
Imperial College London
tilman. rllmam.rooderr100mmpprr ac.
m. vdvilk@imperia. ac. uk
Abstract
Software packages like TensorFlow and PyTorch are designed to support linear algebra operations, and their speed and usability determine their success. However, by prioritising speed, they often neglect memory requirements. As consequence, the implementations of memory-inlensive algorithms that are convenient in terms software design can often not be run for large problems due to memory over- flows. Memory-efficient solutions require complex programming approaches with significant logic outside the computational framework. This impairs the adoption and use of such algorithms. To address this, developed XLA compiler extension that adjusts the computational data-flow representation an algorithm according to user-specified memory limit. We show that k-nearest neighbour, sparse Gaussian process regression methods and Transformers can be run on a single device at much larger scale, where standard implementations would have failed. Our approach leads to better use of hardware resources. We believe that further focus on removing memory constraints compiler level will widen the range of machine learning methods that can be developed in the future.
Introduction
Progress in science is inextricably linked with advances scientific computing, in terms of both software and hardware. This is particularly noticeable in machine learning through the huge impact of numerical software packages supporting automatic differentiation (Baydin et al., 2018), Packages such TensorFlow (Abadi et al., 2016), PyTorch (Paszke et al., 2019), or JAX (Bradbury al., 2018) greatly accelerated 1) the implementation of gradient-based optimisation procedures by elimi- nating error-prone manual differentiation, and 2) the execution of code by leveraging modern and heterogeneous hardware (e.g. GPU, TPU or IPU). A large portion of this impact is attributable to the accessible and user friendly form that these features were delivered in. This contributed the growth in the machine learning community, in terms of methodological researchers, as well as the wider scientific audience and practitioners.
The aforementioned software frameworks work by chaining together efficient implementations of mathematical operations (known kernels). By providing implementations that are tailored various types of hardware, a speed-opimised implementation can be obtained. While speed is certainly important pursue, many algorithms face a different challenge: hardware memory constraints. Often, these have a larger impact, as memory constraint violations can lead to the execution terminating before an answer is obtained. This make-or-break property particularly noticeable on GPUs, where
'The code is available at https://githut: http:///////gtttt....aaaaaa////a com/awaar/tessorloo..
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

allocating more memory than is physically available leads to an immediate termination of execution, and larger amounts of physical memory comes significant cost.
Now that numerical computation frameworks are widely used, they strongly influence what machine learning algorithms are adopted. This happens through hard limitations, well usability con- siderations through what is easily implementable. Currently, the emphasis on optimising runtime causes many algorithms be severely memory limited, or too cumbersome to implement. This is particularly noticeable in methods that rely on linear algebra computations, c.g. kernel methods, nearest neighbour methods for geometric deep learning (Bronstein et al., 2017), or transformers.
In this work, we aim to remove these limitations, by developing a tool that optimises code to be more memory efficient, with particular focus on linear algebra operations. This optimisation is transparent to the user, and therefore allows many algorithms to be run at scales that were previously impossible, while leaving implementations simple as before. This allows wider range of algorithms to take advantage of "the bitter "esson""GGnneee
methods that leverage computation are ultimately the most effective" (Sutton, 2019)-while making them more accessible the wider community, as computational frameworks have sought to do all along.
Our method is implemented as extension to the XLA compiler (Leary and Wang, 2017), which we chose due to its wide use and support for optimising computations specified in TensorFlow and JAX. We demonstrate the benefits of our method by scaling algorithms where simple implementations do not scale due to memory bottlenecks, such as senarest-neighhor regression (Titsias, 2009), and Transformer (Vaswani et al., 2017). With our extensions, these methods scale to far larger problems, without changing a single line of their implementation in Python. Our Gaussian process experiment shows that simply scaling up 13 year old method can outperform much more recent methods, indicating that older methods may be undervalued in recent literature.
sparse Gaussian process
2 Motivation: Memory-Constrained Machine Learning
Since memory overflows cause the execution of code to be immediately halted without producing any result, memory constraints form the key obstacle for scaling many machine learning algorithms. Researchers and ML practitioners often reach this limit, as memory scarce resource that comes at considerable cost, particularly in GPUs. This is particularly noticeable in algorithms where minibatching is undesirable, or that rely on pairwise distances, like k-Nearest Neighbours (kNN) or Gaussian processes (Rasmussen and Williams, 2006). Even in modern deep learning, memory constraints cause problems by limiting batch sizes, layer widths, or sizes of attention mechanisms (Vaswani et al., 7). In all of these examples, linear algebra operations cause the bottleneck. For KNN, kernel methods, and transformers the root of the problem is a pairwise matrix needs to be computed between inputs, giving quadratic memory cost.
Often, more memory efficient implementations can be programmed at the cost of increased software complexity. This ranges from minor annoyances, for example accumulating minibatch gradients in an outer loop for large-batch training, to complex engineering efforts that have been published as scientific contributions in their own right, for example in scaling Gaussian processes to > 105 datapoints (Gal et al., 2014; Wang et al., 201 9a; Meanti et al., 2020).
Our goal is to provide tool that finds memory-efficient ways to execute algorithms, without the need for increasing software complexity. This will allow scientists and practitioners to access the benefits of scale in existing methods more easily, and without incurring the cost of expensive large-memory hardware. For the main demonstration of our approach, we will automatically obtain a memory- efficient implementation of sparse Gaussian process regression (Titsias, 2009), which was previously implemented with considerable difficulty (Gal et al., 2014). The increase in scale makes the method competitive in comparisons where it was previously dismissed as not scalable enough (Wang et al., 2019a), showing the value of reducing the barriers to scaling.
3 Related Work
A popular approach to address memory issues is distributing computation across multiple resources like a group of GPUs or computer cluster with network protocol connectivity between machines (Buyya, 1999; Dean and Ghemawat, 2008), More specifically, sharding allows large tensors to be
2

---

split up and distributed across multiple devices, which increases the total amount of memory available for an algorithm, but comes the cost of requiring more hardware resources. Most computational frameworks" (Abadi et al., 2016; Shazeer et al., 2018; Bradbury et al., 2018; Paszke et al., 2019) support sharding, although some manual specification of how to distribute computation required. This complicates an implementation and requires the user to have wide engineering skill set. Automatic sharding tools such as Tofu (Wang et al., 2019b) simplify implementation somewhat, although the specification of computations through custom interface is still required. While sharding approach does allow scaling of certain implementations, it remains wasteful for algorithms that can be implemented in a more enemory-efficieee way, but where it is simply cumbersome to do so.
Compilers have been introduced to allow humans to express programs in an elegant way, while generating programs that actually run well on specific hardware Aho et al. (2006). Our goal of obtaining memory-efficien implementations, while keeping code convenient for humans, is therefore suited to be addressed by adding memory optimisations to compiler. Compilers are already being used to optimise computational graphs, notably in JAX, TensorFlow and PyTorch by XLA (Leary and Wang, 2017), TVM (Chen et al., 2018), Glow (Rotem et al., 2018) for PyTorch only. TVM performs similar optimisations to XLA, but unlike XLA, it is not seamlessly integrated into popular frameworks and requires additional user effort.
The optimisations in XLA mainly focus on increasing code speed, for example through common expreexeession elimination (CSE), dead code elimination (DCE), operations fusion, and other more specific modifications. The main advantage of XLA is that optimises computations in a way that is completely transparent to the user who specifies the computational graph. Although XIA and TVM implement low-level memory optimisations, they do not adapt code handling large tensors to satisfy memory constraints. For the matrix and linear algebra tasks that we consider, KeOps (Feydy et al., 2020; Charlier et al., 2021) currently provides the most efficient memory management. To achieve any benefits, user must specify a series of computations using KeOps classes, which form layer above the PyTorch framework. KeOps works similarly to a compiler, by first building a symbolic representation of the computation, which allows the computation to be broken into memory-efficient sections, that are then run with custom CUDA kernels.
In terms of prior work, KcOps is closest in aim and achievement to ours. We aim to address three of its limitations. Firstly, KeOps requires users to reimplement their algorithms using KeOps classes While the programming interface is elegant, needing to mix KeOps and other computational frameworks does add complexity. Secondly, for KeOps to be able to optimise an operation, it has to be reimplemented within KeOps, which significantly duplicates effort. Finally, because of the former drawback, KeOps does not inherit the support for a wide range of hardware from e.g. JAX/TensorFlow.
4 Memory Efficient Matrix and Linear Algebra Operations in XLA
Compilers are a promising way for improving runtime properties of code, without requiring user intervention and while leaving code elegant. The specific matrix and linear algebra optimisations that we consider have not yet been implemented in any of the frameworks discussed above. They could be implemented in any of TVM. KeOps, or XLA. We choose to extend XLA over TVM, due to XLA's better integration with common computational frameworks. In addition, we choose to extend XLA over KeOps. because it 1) does not require algorithms to be rewritten in separate framework, 2) can optimise computational graphs in their entirety, rather than just what is implemented in the separate framework, and 3) can take advantage of the full capabilities that already exist in JAX/Tensorllow.
We introduce several optimisation strategies (known as optimisation passes the XLA codebase) into the XLA optimisation pipeline. We aim to constrain the program's 's memory footprint with minimal sacrifices in the execution speed. The optimisation passes examine the entire computational data-flow graph (High Level Optimiser Internal Representation, or HLO IR), search for weak spots, and try to eliminate them. Abstractions at a similar level to HLO IR have been shown to be convenient for optimising linear algebra operations (Barthels et al., 2021). We add mattchannd-rellaee operations, e.g. to introduce a more efficient distance computation, reshuffling operations for expressions that are invariant to evaluation order, and splitting with large tensors to reduce memory usage.
Publishhed under permissive open-source licenses, like Apache BSD.
3

---

Listing 1 Chain multiplication example C = ABv for A, B R"*" and v R".
Qjax.ji1
def matrix wtwp_maati__.eto_mmlll (A, B, v):
C A
return C
4.1 Match and replace
The match and replace optimisation pass, also known as peephole optimisation (Aho et al., 2006), searches for expressions in data-flow graph for which we know in advance that an equivalent and more efficient version exists. For example, we search for expressions that compute Euclidean distance in naive form between vectors of length n and m with dimension d. The naive Euclidean distance computation uses broadcasting over the dimension d and creates a temporary tensor with entries (Ind ymd)) of size m X d. This can be replaced with En and + Wond 2nndymdd where the largest tensor has size x m.
Replacing sub-parts of the graph is a standard procedure in compilers like XLA, although many linear algebra tricks have been missing. While the Euclidean distance is the only match-and and-replace optimisation we implement, other operations can easily be added, for example, efficiently adding diagonals to matrices without allocating a dense square tensor where only the diagonal is non-zero.
4.2 Reordering
computational data-flow graph is an ordered sequence of operations, with the order of operations influencing the memory usage. In some cases, reordering sub-parts of the data-flow graph can lead to reductions in the memory footprint. The classical example reordering is the optimisation of matrix chain multiplications. For example, consider the matrix expression C = ABv for matrices A, B E RT**, and v E IR". In the listing 1, the order of operations determines that the matrix multiplications are performed from left to right, i.e. C = (AB)v, which gives the most inefficient execution order with runtime complexity O(n") and memory complexity On²²). Changing the order to C = A(Bv) improves time complexity to O( and practical memory complexity because the intermediate multiplication result of Bv is a vector not a matrix as in the case of AB multiplication.
The optimisation of matrix chain multiplication is possible due to the associativity of matrix multipli- cation, such that the result of the matrix multiplication chain does not depend on where parentheses are placed. There are many efficient and sophisticated algorithms for addressing this task (Chin, 1978; Czumaj, 1996; Barthels et al., 2018; 8; Schwartz and Weiss, 2019). We implement a simple matrix chain multiplication optimisation algorithm using dynamic programming in bottom-up fashion (Barthels et al., 2018, $2). Our version of matrix chain multiplication algorithm supports expressions with matrix multiplications, sum reductions and transposes.
4.3
Data-flow graph splitting
Often, part of a computational data-flow graph can be divided into multiple independent copies, such that each copy of the data-flow graph or its part act on slice of the input tensor, and the results are combined afterwards in some fashion. This splitting approach is also known as a MapReduce technique (Dean and Ghemawat, 2008), where a computation is divided into smaller and less expensive parts (map) and then combined into the final result (reduce). The splitting technique is common for distributing the computational load. The focus existing solutions is on exploiting hardware parallelism utilising multiple devices. Instead, we use the same techniques for reducing total memory consumption, which is possible because the memory for individual map operations can be freed before the whole result is computed.
An optimisation pass starts by running a depth-firs first search from the final result of the computation. The operations dot or reduce_* * are special, as they often indicate that a computation involving a large tensor can give a smaller result. Once a dot or reduce_ operation marked as fully traversed, we recursively search the traversed paths for operands that are impractically large tensors, until we reach operands that are deemed small enough. Along the way, we keep track of which operations are applied, and along which axes they are trivially parallelisable. The result is collection of sub-graphs,

---

Algorithm High-level description of the depth-first search visitor-handler that splits the data-flow graph up the reduction dot operation. Symbol denotes directed computational path in the data-flow graph. Steps 9-12 are done recursively traversing back visited operations in the data-flow graph.
1: procedure HANDLEDoT(dot: #loInstruction)
2:
if butput_sizee(e size 2 tensor size cearrrssiieeteeeeeeeeeee threshold, i.e. dot is splittable then
3:
Exit and continue traversing succeeding operations in the data-flow graph and search for size-reducing operation for the output tensor of dot.
4:
if dot. rhs is not splittable and dot. -1s is not splittable then
5:
Exit and continue traversing the data-flow graph.
6:
if' bot.rhs * dot dos and both operands are splittable then Exit and continue traversing the data graph.
7:
8:
Let operand_ to_ to_sslit dot. rhs or aperand_to_spp = dot 1hs depending on previous splittability conditions.
9:
Let split_dims {di..... d., ), split_producers
{op............................. op s.t.
Vop split_producers exists a path op ~ sperand_to_spli1 toapddd___ppppi
11:
Vop split_producerrs, input_sized
I tensor_size sittor_siiitttrehhlleeeeeo threshold
Vop E split_producers. Ed E split_diiie dims which is splittable on the path op operand_to_split
13:
if split_diimg Ø or
Ø then Exit and continue traversing the data-Ilow graph.
14:
15:
Let best_split__de limmii_dimmmi = d E split t_mis, and ops S splii_pprouucct s.t.
16:
mindesplit_ _ime d split_size Vop E ops, the path op ~y operand_ to_ split is splittable on best_split_di
Vensor_split_siee
17:
18:
bestt_ Let split_sliiie size best__plitt_ silpiittdmmmmmmmiiieeee size peperand_to_spall bensor_split_szee Create while loop Hlolnstruction, s.t.
19:
20:
The loop iterates splits of size split split_size size best_aplit_di of paths ops
operand_to tt__ppitttttt
The loop applies dot reduction operation on the slice of operand_ serrand_to_spiit
22:
The slice result of the dot reduction operation is put into the replica of the original dot. result
23:
Replace ops ~
instructions with created while loop.
that start at operations that produce large tensors, and end at operations that reduce them again, together with candidate axes that they can be split across. According to some heuristics which ensure appropriate sizes for intermediate results, we then turn this entire computation into a while loop, where each iteration computes a manageable part of the final result (fig. 1).
Checking if the axis is splittable is necessary as not all operations act independently on each dimension. For example, element-wist operations be split on any axis, whereas the triangular solve operation can be split on "batching" dimensions only. Next, the data-flow graph splitting procedure selects the dimension of the largest size which contributes the most to memory.
As we discussed earlier, the decision about where to split the graph depends on the tensor size. We offer two XLA options the user for out-of-memory mitigation: rensor size threshold and tensor split size upper bound. Tensor size threshold is a criterion designed for detecting which operations should be marked candidates for splitting. Tensor split size upper bound serves as threshold on the largest allowed chunk size for splitting. These options are set equal by default. The command-lime snippet at listing 2 shows how variable, and the snippet is indifferent to the machine learning framework used by the script. Minimal user effort is required for using our XLA compiler extension. The user is involved only in defining what the suitable threshold and splitting sizes are.
would use these options by passing them via environment
Listing 2 Example of how user set options for the extended XLA using the environment variable.
SLA_FLAGST -x0s= tensor, httpassrrssiiieeeeshesosooeeeeee threshold=1GB -ew.laaeessorspliittsseeeeee python train.py
5

---

I [a]
G
A A[...,
i:dr ...]
R O[...]
[[a]
S
I' [i:j]
G'
A [..., 1: ...]
R'
or [...]
Loop
Figure 1: The scheme demonstrates transformation of the data-flow graph the top to the data-flon flo--flow graph at the bottom. The graph on the top consists of G and R blocks which are generator and reducer operations respectively, I is the tensor input of G, A is the tensor output of G and is the tensor output of R. The bracket notation A ([..., a, ...] means that the tensor A has a dimension of size a and A have other dimensions. The i is a slicing operation. operations a computational path of the data-flow graph between a tensor A and an operation R. The eXLA splitting optimisation procedure converts the top graph into the loop of independent iterations performing the same chain of operations on a small slice i: i:j.
R denotes
arbitrary amount of
One strong benefit of our compiler- baasee solution, is that the computational graph represents the whole pipeline of computations, including forward and backward propagation. Our splitting procedure will be applicd automatically, regardless how many derivatives need to be computed. In addition, procedure encompasses two splitting schemes that the machine learning literature distinguishes: model-based and data-based splitting schemes of the data-flow graph. Model-based splitting schemes involve partitioning the model over its parameters, whereas data-based splitting schemes batch over inputs. The proposed splitting approach supports both schemes. However, our implementation of model-based splitting supports only splitting memory intensive expressions involving model parameters. It is not able to split model parameters themselves into smaller chunks.
4.4 XLA limitations
While we still believe that XLA is the right framework for our extensions, several limitations came light during implementation. One limitation that is shared with all current frameworks, is that they only have a weak linear algebra type system, where matrices are represented as arrays without additional properties. Solutions that support stronger type systems (Bezanson al., 2017: Barthels et al., 2021) l) may be able to implement wider variety of match-and-replace optimisations.
Another limitation comes from the default memory allocation manager not being aware of memory limits. Its current behaviour is to execute nodes in the computational graph, and therefore allocate any required memory, as soon the required inputs have been computed. This means that even if tensors are split to manageable sizes, memory overflows can still occur several are executed simultaneously. To prevent this from happening, had use limits that were smaller than our total GPU memory.
Experiments
This section shows how existing software packages take advantage our extension to XLA (cXLA). We demonstrate our optimisations non-parametrie k-nearest neighbours (kNN) and sparse Gaussian process regression (SGPR) models, and Transformers.
5.1 Mltrix-Vector Multiplication
We start by demonstrating the improved efficiency that eXLA offers to large-scale matrix-vector kernel matrix, and y,u E IR". Such
multiplications of the form = Kv, where K is an computations are common in Conjugate- Conjugate-Gradients Gradirrtdddiaaaeee based Gaussian process approximations (Gibbs and Mackay, 1997; Wang et al., 2019a; Artemev et al., 2021), where K,, k(2,, Vj) and k is some kernel function. We choose the common Squared Exponential.
We implement this equation using GPfoww (Matthews et al., 2017), TensorFlow- package that provides a convenient software interface for Gaussian processes and kernel functions. Without eXLA, the entire K would be stored in memory, leading to 222 memory cost. This makes running

---

Elapsed time, seconds
80
20GB
Memory, bytes
2.0
15GB
1.5
60
10G8
5GB
40
1G8
100MB
0.5
20
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Size,
Figure 2: GPU memory consumption and elapsed time of n squared exponential kernel matrix- vector multiplication.
Dataset
Distance
n
KeOps
eJAX
cTF
JAX
TF
Random
I
lc4
100
983 263
277 364
284 777
281 695
280 826
Random
2
le4
3 662 188
292 804
294971
288 098
294 776
Random
L
1e6
100
24367
2433
2530
Random
T,2
1e6
3
123 765
2512
2605
Ø
Ø
MNIST
L2
6e4
784
41 084
32 290
33 455
25 544
26 138
MNIST
L2
6e4
784
40697
2356
2985
2498
2988
Fashion
I
6e4
784
40 309
32 382
33 428
25 558
26 128
Fashion
L
6e4
784
40 40982
2357
2984
2498
2989
Glove-50
Cosine
1.18e6
50
3 464 257
2103
1929
Glove-100 100
Cosine
1.18e6
100
631 420
2053
1871
Ø
Ø
Glove-200
Cosine
1.18e6
200
398 293
1967
1724
Table 1: Query processing rates (queries per second) for kNN. and are the number of data points and the data dimension respectively. Runs which failed due to memory overflow are denoted by ø. Runs with eXLA are denoted eJAX and eTF respectively.
large datasets infeasible, where e.g. = 10 would lead to memory requirement of 8TB, which impractical even for the largest of modern GPUs with 40GB of memory. memory efficient split/implicit implementation is necessary to scale to large datasets, as was impressively done by Wang al. (2019a), but is cumbersome.
We ran our implementation with eXLA enabled, which allows user to control the memory of an algorithm. We evaluated the expression in double precision on Tesla V100 GPU with 32 GB of memory, and applied range of memory limits. In fig. 2 we report the peak memory consumption and execution time of evaluating the kernel matrix-vector product for different sizes, with different memory limits applied. We see that the memory constraints are not violated, and that dataset sizes used that far beyond the 32 GB memory capacity.
5.2 K-Nearest Neighbours
K-nearest neighbours is fundamental machine learning algorithm, with a similar large memory cost. A kNN query selects & closest data points in the dataset to each query point. Brute-force implementations compute pairwise distances between m query points and n data points, resulting in the distance matrix of size implemented using column- wise sort operation on the distance matrix. Our benchmarks show that eXLA scales the brute-force approach and does not fail for large problems, i.e. large and
This is followed by topk operation, which is often naively
compare TensorFlow and JAX implementations with and without eXLA optimisations, and KeOps implementation. We use randomly generated data, common benchmarks like MNIST and Fashion-MNIST, and Glove-50, Glove-10 100 and Glove-200 from the ANN-benchmark toolkit Aumüller et al. (2020). We use 1e4 query points in all benchmarks.

---

Dataset
Model
RMSE
NLPD
Time (hours)
GPUs
SGPR-1000
0.0485 - 2
-..00 + 4e-3
1.32 - 1e-2
SGPR-2000
0.0463 + 3e-4
- - 7e-3
4.18 4e-2
houseelectric
SGPR-3000
0.0438 2e-4
8.97 7e-2
SGPR-4000
0.0429 + 4c-5
-1.222 1e-3
15.59 = 7c-2
Iterative GP*
0.054
-0.207 1e-3
1.55 + 2c-2
8
Iterative GP**
0.050
79.96
8
SGPR-1000
0.2820 + 2e-3
-0.165 - 5
0.28 = 3e-3
SGPR-5000
0.1890 = 3e-4
-0.234 + 1e-3
5.09 = 3e-2
3droad
SGPR-8000
0.1750 2e-4
-0.303 1e-3
11.39 + 1e-1
SGPR- 10000
0.1700 + 20-4
-0.321
+ 4e-4
17.81 + 4e-2
Iterative GP*
0.110 17c-3 1.239 + 25c-3
1.00 + 20-3
Iterative GP**
0.106
Ø
7.06
8
Table 2: The performance of SGPR with inducing points included to the trainable parameter set on houseelectric and 3droad datasets. Iterative GP* and Iterative GP** trained with lengthscale per dimension and shared lengthscale across dimensions respectively. Iterative GP values are from Wang al. (2019a), with unreported metries denoted as Ø.
Our results are listed table (see the appendix for a full table that reproduces Feydy et al. (2020, table 3)). In all benchmarks, though this may not be optimal for performance. We observe that eXLA prevents memory overflows in JAX and TensorFlow. In addition, performance is comparable higher. We acknowledge that KeOps performs significantly better than any JAX or TensorFlow implementation. This is explained by 1) JAX/TF not having efficient implementations for certain functions (e.g. topk runs full sorting algorithm), and 2) KeOps having implemented additional optimisations, which could also be added XLA. However, we note that also achieved our goal of improving the memory and time performance of JAXTensooFloww implementation without changing the code.
the tensor size threshold for eXLA 100MB for simplicity,
5.3 Sparse Gaussian Process Regression
Gaussian processes (Rasmussen and Williams, 2006) are considered the gold standard method for performing regression with uncertainty estimates. A straightforward implementation requires taking a matrix decomposition of an n X n kernel matrix (like those considered in section 5.1), which leads to (nn²) time cost, and an O(n²) memory cost. Scaling Gaussian process is challenging, which often attributed to the time cost. In reality however, large datasets cause memory overflows far before long runtimes become obstacle.
0.45
TF-eXLA (trainable ips)
0.6
TF-eXLA (trainable (ps)
0.5
TF (trainable ips)
TF (trainable ips)
RIASE
0.35
0.2
0.30
0.1
0.20
0.7
0.3
2000
4000
6000
6000
10000
0
2000
4000
6000
8000
10000
Number of inducing points
Number of inducing points
Figure 3: Root mean squared error (RMSE) and negative log predictive density (NLPD) performance test metrics of SGPR for 3droad dataset as the number inducing points increased. The red shaded region emphasizes the capacity of the SGPR model which user can run using standard GPflow and TensorFlow release packages.

---

Approximate methods have been introduced to deal with both the time and space issues. While there are many variants L Candela and Rasmussen, 2005), we consider the sparse variational approximation (Titsias, 2009) for which a naive implementation has O(nm² + m²) time cost, and O(nm + m²) memory cost. Here, m denotes the number of inducing variables, which controls the quality of the approximation. Under certain conditions, the method provides reliable hyperparameter selection (Bauer et al., 2016), and very accurate posterior approximations (Burt et al., 2019, 2020) while using < In practice, these standard implementations may still have their performance limited by how large m can become before a memory overflow occurs, A more memory-efficient implementation with memory cost of O(m²) does exist (Gal et al., 2014), but is cumbersome to implement that it is not widely used or compared against.
Fortunately, the splitting optimisation we implemented in eXLA can discover the same procedure that was engineered by Gal et al. (2014). Moreover, since eXLA operates on the entire computation graph, it optimises gradients as well as the optimisation objective function with no additional effort, while Gal et al. (2014) needed to implement gradients manually. We demonstrate the utility of eXLA by scaling the GPflow (Matthews et al., 2017, 2.3.1 release version) implementation of Sparse Gaussian process regression (SGPR, Titsias, 2009), without any modifications of the code.
With our eXLA optimisations, SGPR was able to scale to much larger datasets, with more inducing points. We conduct experiments on Tesla V100 GPU with 32 GB of memory, and run on two of the largest UCI datasets that are commonly considered in Gaussian process research: 3droad and bouseelectric with total number of data points 434,874 and 2,049,280 respectively. We set the tensor size threshold and the tensor split size in e eXIA to 1GB. We primarily compare to Wang et al. (2019a), who use Conjugate Gradients approximation (Gibbs and Mackay, 1997) to achieve the most impressive scaling of a Gaussian process approximation to date, using an impressively engineered implementation that manually splits and distributes parts of the computation.
In fig. 3 we compare GPflow's SGPR implementation with and without eXLA as we increase the number of inducing points. We see that until about 800 inducing points, the normal and eXLA runs result in the same predictive metrics, as desired, After 800 inducing points, runs without XLA fail with an "out of memory" error, while with eXLA we scaled 10 inducing points. Simply scaling the method in this way leads to significant performance improvements.
We compare predictive accuracies directly with the scalable Conjugate Gradients implementation of Wang et al. (2019a). In that paper, SGPR was discussed as method that would not scale, probably due to the difficulty of implementing it in memory memory-effcieent way in Gal et al. (2014). Table 2 shows that simply scaling SGPR using eXLA makes it competitive with the Conjugate Gradients implementation of Wang et al. (2019a), without needing additional hardware.
5.4 Transformers
Neural network architectures based on the attention mechanism, such as Transformers (Vaswani et al., 2017), have become a common tool for solving natural language processing tasks. Transformers have also been applied successfully in computer vision classification tasks (Dosovitskiy al., 2020) and image generation (Ramesh et al., 2022). The attention mechanism Ati(Q, K, V) = softmax
[K" V. with Q, K E [Rxxaxd,, and V E Rbxexd. is an expression that lies in the
core of Transformer models. The memory consumption by the intermediate result QKT grows quadratically with dimension 8, i.e. O(bs2). Rabe and Staats (2021) investigated ways of reducing the memory overhead by sequentially computing parts of the attention operation with large 8. Their technique is very similar the eXI.A's splitting optimisation pass from section 4.3. The advantage of eXLA is the automatic application of such schemes, depending on sizes of the input dimensions b, and d as well as user's eXLA setting for the tensor size limits.
We test eXLA on a sequence sequence-to-sequenee to-sequence language translation Transformer. We use the configuration of the translation Transformer from the TensorFlow website?. We increase the size of the attention mechanism by increasing the size of the randomly generated input sequence until the code fails with memory overflow; we then show that running the same implementation with eXLA avoids this failure. In the experiment with the language translation Transformer this corresponds to increasing the input sequence. We run experiments single Nvidia V100 GPU with 32GB memory. Figure 4
hhttp:: http://wwww.tssrroro..// hww.://wwww..rggeee////otttta///

---

1e10
3.0
Elapsed time, seconds
10
XLA
Memory, bytes
2.5
TF
2.0
10a
eXLA
1.5
1.0
10-1
0.5
0.0
1000
2000
3000
4000
5000
6000
7000
1000
2000
3000
4000
5000
6000
7000
Sequence length
Sequence length
Figure 4: Memory and elapsed time metrics for the language Transformer. The grey vertical line marks the maximum sequence size where implementations in TensorFlow and XLA fail with OOM, but eXLA continues work.
shows that out of the box implementations of translation Transformer using TensorFlow with its default compiler (TF) and TensorFlow with the XLA backend fail to run when the sequence length larger than 2000. The Transformer model compiled with eXLA optimisations managed to run with sequences up 7000 with tensor limit set 10GB and the tensor split size set 1GB. The drop in the fig. 4 after 2000 is explained by eXLA's splitting optimiser modifying the graph with tensor splits equal 1GB. After eXLA optimisation applied, there are no tensors the graph larger than 1GB which can be allocated at the same time together on the GPU. In these situations a practitioner might use larger split size the hardware allows it, which turn will boost the computational performance.
6 Discussion
We showed that our XLA compiler optimisations (eXLA) prevented memory overflows in algo- rithms with large tensor linear algebra operations. Our compiler extension automatically adjusts computational data-flow graphs to control memory utilisation. As demonstrated in the experiments section, machine learning models compiled with eXLA could run greater scale, whereas their out-of-the-box implementations failed with Out of Memory errors. Crucially, we used existing software packages without modifying any code.
In addition showing that our compiler extensions work as intended, our experiments also provide directly useful empirical results for Gaussian processes. We managed to an "old" method (SGPR, Titsias, 2009), with unchanged code, to obtain empirical results that outperformed a state-of-the art method (Wang et al., 2019a). This shows that existing methods may simply need to be scaled in order to provide high-quality solutions, which are even guaranteed by theory in certain cases (Burt et al., 2020). We also provide proof-of-concepp experiments that show that eXLA can help with running Transformers (Vaswani et al., 2017) which are notoriously memory hungry, on cheaper hardware without increasing software complexity.
The exciting possiblity of eXL.A is that it opens up the possibility to probe behaviour of machine learning models in regimes that were previously infeasible, and on cheap hardware. For example, one could train very wide neural networks, to empirically compare to behaviour predicted by NTK theory (Lee et al., 2018; Matthews al., 2018; Jacot et al., 2018; Novak et al., 2020).
The current implementation of eXLA is still only demonstration of what compiler optimisations could achieve, and many more optimisations can be added. We believe that increasing the capability of compilers like XLA will greatly increase the efficiency of researchers and practitioners. We hope that community driven compiler projects will contribute to the community in a similar way to how existing numerical frameworks already do.
Acknowledgements
Thanks to David R. Burt and Sebastian W. Ober for the feedback on the draft of this paper, Fergus Simpson and Alan Saul for their advice on Transformer models. We also would like to thank Lev Walkin and Renat Idrisov for discussions about compilers in the beginning of this project.
10

---

References
Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., et al. (2016). Tensorflow: A system for large-scale machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16). 1,3
Aho, A. V., Lam, M. S., Sethi, R., and Ullman, J. D. (2006). Compilers: Principles, Techniques, and Tools (2nd Edition). Addison- A Longman Publishing Co., Inc. 3,4
Artemev, A., Burt, D. R., and Van Der Wilk, M. (2021). Tighter bounds the log marginal likelihood of Gaussian process regression using conjugate gradients. In Proceedings of the 38th International Conference Machine Learning (ICML), volume 139. 6
Aumüller, M., Bernhardsson, E., and Faithfull, A. (2020). ANN-Benchmarks: benchmarking tool for approximate nearest neighbor algorithms. Information Systems, 87.
Barthels, H., Copik, M., and Bientinesi, P. (2018), The generalized matrix chain algorithm. In Proceedings of the 2018 International Symposiun Code Generation and Optimization. 4
Barthels, H., Psarras, C., and Bientinesi, P. (2021). Linnea: Automatic generation of efficient linear algebra programs. ACM Transactions Mathematical Software (TOMS), 47. 3,
Bauer, M., van der Wilk, M., and Rasmussen, C. (2016). Understanding probabilistic sparse gaus- sian process approximations. In Advances in Neural Information Processing Systems, volume 29.
Baydin, A. G., Pearlmutter, B. A., Radul, A. A., and Siskind, J. M. (2018). Automatic differentiation in machine learning: a survey. Journal of Marchine Learning Research, 18.
Bezanson, J., Edelman, A., Karpinski, S., and Shah, V. B. (2017). Julia: A fresh approach to numerical computing. SIAM review, 59.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke,
A., VanderPlas, J., Wanderman-Miine S., and Zhang, Q. (2018). JAX: composable transformations of Nuthon+NumPy programs. 3,
Bronstein, M. M., Bruna, LeCun, Y., Szlam, A., and Vandergheynst, P. (2017). Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34. 2
Burt, D., Rasmussen, C. E., and van der Wilk, M. (2019), Rates of Convergence for Sparse Variational Gaussian Process Regression. In Proceedings of the 36th International Conference on Machine Learning (ICML), volume 97, 9
Burt, D. R., Rasmussen, C. E., and van der Wilk, M. (2020). Convergence of sparse variational inference in Gaussian processes regression. Journal of Machine Learning Research, 21. 9, 10
Buyya, R. (1999). High performance cluster computing. New Jersey: F f'rentice. 2
Charlier, B., Feydy, J., Glaunès, A., Collin, F.-D., and Durif, G. (2021). Kernel operations the GPU, with autodiff, without memory overflows. Journal of Machine Learning Research, 22.
Chen, T., Moreau, T., Jiang, Z., Zheng, L., Yan, E., Shen, H., Cowan, M.. Wang, L., Hu, Y., Ceze, L., et al. (2018). TVM: An automated end-to-end optimizing compiler for deep learning. 13th USENIX Symposium Operating Systems Design and Implementation (OSDI 18). 3
Chin, Y. (1978). An O(n) algorithm for determining a near-optimal computation order of matrix chain products. Communications of the ACM, 21.
Czumaj. A. (1996). Very fast approximation of the matrix chain product problem. Journal of Algorithms, 21.
Dean, and Ghemawat, (2008). Mapreduce: simplified data processing on large clusters. Communications of the ACM, 51. 2,

---

Dosovitskiy, A., Beyer, L... Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., et al. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 19299. 9
Feydy, J., Glaunès, J., Charlier, B., and Bronstein, (2020). Fast geometric learning with symbolic matrices. Advances in Neural Information Processing Systems, 33. 3, 8
Gal, Y. Van Der Wilk, M., and Rasmussen, C. E. (2014). Distributed variational inference in sparse Gaussian process regression and latent variable models. In Advances in Neural Information Processing Systems, volume 27. 2,
Gibbs, M. and Mackay, D. (1997). Efficient implementation of Gaussian processes. Technical report, Cavendish Laboratory, University of Cambridge. 6,
Jacot, A., Gabriel, F., and Hongler, C. (2018). Neural Tangent Kernel: Convergence and General- ization Neural Networks. In Advances in Newral Information Processing Systems, volume 31. 10
Leary, C. and Wang, T. (2017). XLA: Tensorflow, compiled. TensorFlow Dev Summit. 2,
Lee, J., Sohl-dickstein, J., Pennington, J., Novak, R., Schoenholz, S., and Bahri, Y. (2018). Deep neural networks as Gaussian processes. In International Conference on Learning Representations. 10
Matthews, A. G. d. G., Hron, J., Rowland, M., Turner, R. E., and Ghahramani, Z. (2018). Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations. 10
Matthews, A. G. d. G., van der Wilk, M., Nickson, T., Fujii, K., Boukouvalas, A., León- Villagrl,, P., Ghahramani, Z., and Hensman, J. (2017). GPflow: A Gaussian process library using TensorFlow. Journal of Machine Learning Research, 18. 6, 9, 14
Meanti, G., Carratino, L., Rosasco, L., and Rudi, A. (2020). Kernel methods through the roof: Handling billions of points efficiently. Advances Neural Information Processing Systems, volume 33. 2
Novak, R., Xiao, L., Hron, J., Lee, J., Alemi, A. A., Sohl-Dickstein, J., and Schoenholz, S. (2020).
Neural tangents: Fast and easy infinite neural networks in python. In International Conference Learning Representations.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. (2019). Pytorch: An imperative style, high-performmnce deep learning library. In Advances Neural Information Processing Systems, volume 32. 1,3
Quiñonero- Canieel,, J. and Rasmussen, C. E. (2005). A unifying view of sparse approximate gaussian process regression. Journal of Machine Learning Research, 9
Rabe, M. N. and Staats, C. (2021). Self-attention does not need o(n²) memory. arXiv preprint 2rXiv:2/12.0050) 9
Ramesh, A., Dhariwal, P., Nichol, A., Chu, C., and Chen, M. (2022). Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.0611 06125.
Rasmussen, C. E. and Williams, C. K. (2006). Gaussian processes for machine learning. Gaussian Processes for Machine Learning. 2,
Rotem, N., Fix, J., Abdulrasool, S.. Catron, G., Deng, S., Dzhabarov, R., Gibson, N., Hegeman, J., Lele, M., Levenstein, R., et (2018). Glow: Graph lowering compiler techniques for neural networks. arXiv preprint /00.00000000000000
Schwartz, O. and Weiss, E. (2019). Revisiting "Computation of Matrix Chain Products". SIAM Journal on Computing, 48. 4
12

---

Shazeer, N.. Cheng, Y., Parmar, N., Tran. D., Vaswani, A., Koanantakool, P., Hawkins, P., Lee, H., Hong. M., Young, C., et al. (2018). Mesh-tensorflow: Deep learning for supereomputers. In Advances in Neural Information Processing Systems, volume 31. 3
Sutton, R. (2019). The bitter lesson. Incomplete Ideas (blog), 13.
Titsias, M. (2009). Variational learning of inducing variables in sparse Gaussian processes. In Proceedings of the 12th International Conference on Artificial Intelligence and Statistics, volume 2, 10
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez. A. N., Kaiser, L. u., and Polosukhin, I. (2017). Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. 2, 9, 10
Wang, K., Pleiss, G., Gardner, J., Tyree, S., Weinberger, Q., and Wilson, A. G. (2019a). Exact Gaussian processes on million data points. In Advances in Neural Information Processing
Systems, volume 32. 2, 6, 8, 9, 10
Wang, M., Huang. C.-c., and Li, (2019b). Supporting very large models using automatic dataflow graph partitioning. In Proceedings of the 14th EuroSys Conference 2019. 3
13


Reviews:
[{"review_id": "aVqFsMt59Ix", "reviewer": "Reviewer_1TM1", "summary": "To save memory usage in machine learning workloads, this paper paper presents a graph-level rewriting algorithm that is integrated with the XLA compiler. The algorithm is tested against two non-deep learning workloads, k-nearest neighbor (kNN) and sparse Gaussian process regression (SGPR).", "questions": "The first confusion point to the reviewer is that the work seems to be focusing very much on deep learning compilation, as mentioned on line 20-21, line 48-49, line 98-100, Section 4.4, all of which are deep learning frameworks and compilers. However, the experiments are only conducted on two much simpler non-deep learning algorithms. Is this possible to demonstrate any experiments with deep learning workloads?\n\nThe second question is about literature review. The paper seems to propose that KeOps as major prior work, while ignoring all significant works in memory reduction techniques in deep learning. Is there any consideration not to do so?\n", "limitations": "As mentioned before, the limitation includes three components:\n1) The novelty of this work, if only limited to various rewriting, is less significant.\n2) While focusing on deep learning compiler and frameworks in paper writing, the experiments are not relevant to deep learning workloads.\n", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1) The primary strength of this paper is that it proposes to solve the memory constraint problem, which is an existing painpoint of deep learning systems;\n2) The algorithm is integrated with XLA, a mature deep learning compiler, which makes it easier for productionalization;\n3) The paper optimizes kNN and SGPR and delivers positive results.\n\nWeakness:\n1) There exists quite a lot of work reducing memory footprint in deep learning [1, 2, 3], all of which are quite non-trivial and made significant contributions to the field of deep learning compiler. This paper does not seem to provide comparison with those existing literatures.\n2) The algorithms present in this paper, including match-and-replace, reordering and dataflow graph splitting, seems to be very familiar but less novel. For example, changing computation order (Sec 4.2) (AB)v = A(Bv) seems to be a classic trick to save memory.\n3) Even though XLA is a deep learning compiler, and also dataflow rewriting (Sec 4.3) is a classic mechanism in deep learning compilation, the experiments are conducted only on much simpler workloads like kNN, which only contains a few matrix operators and does not have a strong motivation to conduct generic graph-level rewriting. Therefore, it's less clear to the reviewer that the proposed algorithm could be a perfect fit for deep learning compilation.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "Yxjr0MfcUct", "reviewer": "Reviewer_Mdhs", "summary": "This paper extends the XLA compiler with several memory optimizations. The users can run much larger scale computations without changing any code.\nThe evaluation shows that this exertion enables k-nearest neighbor and sparse Gaussian process regression methods to run at a much larger scale on a single device.", "questions": "Tiling/Blocking is a widely used technique to save memory.\n\nThis paper [1] applies this idea to self-attention, which is a very important and memory-intensive operator for long sequence transformers. Can the technique in eXLA also automatically discover the tiling pattern in [1]?\n\n\n[1] Self-attention Does Not Need O(n^2) Memory, arXiv 2021", "limitations": "There are other memory optimization techniques such as swapping[2] and gradient-checkpoint/rematerialization[3][4]. These techniques naturally fit into compiler optimizations. You can also try these methods in XLA.\n\n[2] Swapadvisor: Pushing deep learning beyond the GPU memory limit via smart swapping, ASPLOS 2020.  \n[3] Training deep nets with sublinear memory cost, arXiv 2016.   \n[4] Dynamic Tensor Rematerialization, ICLR 2021.   ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Good direction: integrate memory optimization as compiler passes.\n- Impressive results: Run much larger scale computations on a single device without code change.\n\nWeaknesses:\n- It is still a prototype with many missing optimizations, as it cannot outperform KeOps\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "NhiLINDhd_", "reviewer": "Reviewer_esE8", "summary": "Authors propose and implement XLA compiler extension that replaces algorithms using a lot of memory with more memory efficient algorithms. This allows to run computations that use much larger data sizes on a single limited-memory device. Authors demonstrate their tool on kNN and sparse Gaussian process regression methods.", "questions": "- It would be good if authors ran their system on a set of models that may or may not contain optimizable code and reported whether the system did not yield any negative results. What I am asking for is a testing run: does the system crash, fail, break the underlying code, slow down the underlying code when run on inputs that it may not expect?\n- It would be good if authors explained or walked through Algorithm 1. While it is clear what algorithm is supposed to do, the details as written are hard to follow. Some detailed questions: What exactly does line 5 of the algorithm mean? What does \"define best_split_dim\" mean in line 13? Line 14 says to iterate over best_split_dim, but best_split_dim was not assigned anything. Line 14 bullets are not clear either.", "limitations": "I think that authors adequately addressed the limitations", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Authors implement XLA compiler extension that automatically allows restructuring underlying computation graph to use less memory and therefore execute larger computations on a single device\n- Authors demonstrate that the system can run computations that would have failed with out-of-memory without optimization\n- Authors demonstrate that the system can be comparable to manual optimization\n\nWeaknesses:\n- The optimization is limited and runs much slower than KeOps system (authors acknowledge and explain this limitation)\n- Limited explanation of splitting algorithm", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept