Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper extends the XLA compiler with several memory optimizations. The users can run much larger scale computations without changing any code.
The 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.
strengths_and_weaknesses: Strengths:
- Good direction: integrate memory optimization as compiler passes.
- Impressive results: Run much larger scale computations on a single device without code change.

Weaknesses:
- It is still a prototype with many missing optimizations, as it cannot outperform KeOps

questions: Tiling/Blocking is a widely used technique to save memory.

This 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]?


[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.

[2] Swapadvisor: Pushing deep learning beyond the GPU memory limit via smart swapping, ASPLOS 2020.  
[3] Training deep nets with sublinear memory cost, arXiv 2016.   
[4] Dynamic Tensor Rematerialization, ICLR 2021.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: Strengths:
- 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
- Authors demonstrate that the system can run computations that would have failed with out-of-memory without optimization
- Authors demonstrate that the system can be comparable to manual optimization

Weaknesses:
- The optimization is limited and runs much slower than KeOps system (authors acknowledge and explain this limitation)
- Limited explanation of splitting algorithm
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?
- 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
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: 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).
strengths_and_weaknesses: Strength:
1) 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;
2) The algorithm is integrated with XLA, a mature deep learning compiler, which makes it easier for productionalization;
3) The paper optimizes kNN and SGPR and delivers positive results.

Weakness:
1) 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.
2) 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.
3) 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.


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?

The 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?

limitations: As mentioned before, the limitation includes three components:
1) The novelty of this work, if only limited to various rewriting, is less significant.
2) While focusing on deep learning compiler and frameworks in paper writing, the experiments are not relevant to deep learning workloads.
review_title: null

Meta-review: metareview: This paper implements a set of optimizations on top of the XLA compiler with the explicit purpose of reducing the memory footprint of an algorithm. This is important because while a lot of optimization work in this space has traditionally focused on speed, memory is more commonly the bottleneck when running large computations. The key strengths of the paper are that it presents a genuinely useful artifact and is generally well explained (except for section 4.3 which is not very well explained). The results also show that one can get some significant memory improvements with zero additional effort if your code already compiles with XLA. The main limitations of the paper are limited novelty (none of the transformations are particularly surprising) and limited relevance to NeurIPS---this is really a compilers paper, and it's not even evaluated on any deep learning workloads. These two limitations make this a borderline paper, but all reviewers considered it to be above the bar. 

Minor comment: In section 4.1, in the compiler literature, these match-and-replace optimizations are known as peephole optimizations. The term should be at least mentioned.
confidence: Less certain
Decision: Accept