Venue: NeurIPS 2022 Conference

Reviews:

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: The paper outlines an approach to streamlining the process of tensor program optimization for different hardware substrates, by integrating a prior over optimization approaches. The latency of the compiled program (or a prediction thereof) is added to the prior probability of the optimization transforms as a potential function for random walk metropolis hastings, to explore alternative optimization approaches. Adding hardware specific optimizations is feasible and unlocks performance improvements on modern DL models, relative to TVM and PyTorch.
strengths_and_weaknesses: Originality
The approach seems to be a novel one, and contributes to an important field of compiler optimizations for ML models. These systems facilitate performant deployment of models from training devices (GPU, high end CPU) to inference endpoints (mobile CPU, GPU). Previous work in this area could include Tensor Comprehensions, TVM, Halide, and beyond.

Quality
The "user experience" of specifying kernels in the probabilistic DSL looks a little bit awkward. An upside, as a python-embedded traced DSL, we can interleave host language control flow to study conditions such as spatial-vs-reduction. A downside is that the resulting kernel code may need to be specialized to particular tensor shapes and ranks. It also appears (but I could have misunderstood?) that transformation modules are specific to tensor programs, which could limit the applicability to performance-critical code segments. The decision to use MCMC to sample this space is an interesting one, because I would expect the outcomes to be highly nonlinear relative to the decisions being sampled. A non-smooth posterior is a nightmare for most MCMC samplers, and I'd wonder whether there's any value to MCMC relative to some kind of prior-regularized evolutionary approach (or is that basically how we are viewing the MH work here?). Do you need to scale the prior relative to the runtime "energy"?

Clarity
This is a "systems" type of paper, where observing code and before/after different optimizations is probably an effective mode of presentation, but is necessarily limited by the page limit. The "kinds of things" that can be done with SpaceCraft are reasonably well presented, but I did not feel that the "way those things are done" was very clearly communicated. Also unclear was the time spent exploring for optimizations relative to a deterministic system like TVM.

Significance
~2x speedup relative to PyTorch on BERT-Large seems significant, but without very-convenient tie-in to major ML systems may not be very usable by the community. The claim on 278-80 about ease of adding a novel TC-specific module are enticing, but probably deserve to be fleshed out further in the body of the paper rather than supplementary.
questions: How easily is the new system integrated with existing ML frameworks? Can we easily write code in PyTorch or JAX that calls into this system? Does optimization need to be re-run with every process startup, or is there a cache of previous work maintained?

How important is the prior, and how is it specified? None of the examples suggest the user can provide the prior probability of selecting any given categorical decision or relative probabilities of different transformation modules being applicable. So, is a uniform prior over categories assumed?
Do you need to scale the prior relative to the runtime "energy"?

It seems like some discussion of Bayesopt / RL based approaches would be in order. I am surprised if noone has tried RL / contextual bandits in this space, but even if that's the case it's worth calling out as future work.

It seems like the system supports only categorical/discrete decisions. Are there cases in compiler optimization where a continuous random variable is useful? If we had such a case would it be easily supported by SpaceCraft?

Do the programs produced yield interesting numerical issues? Sometimes, reordering arithmetic operations such as sum reductions can result in catastrophic cancellations.
limitations: Main limitation discussed is invalid output programs (OOM, etc), and validation being a viable approach to rectify this.
It'd be interesting to learn how long optimization takes, possibly see curves of "best runtime by # steps".
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: SpaceCraft is developed as an improvement over previous TVM IR such that it can unlock two key features: (1) not having to define the template like initial AutoTVM and its variants, and (2) achieving same programmability as Auto Scheduling such as Ansor. Overall, the paper devise a new abstraction with probabilistic language abstractions and search algorithms that improve upon previous black-box optimizations.

First of all, it is important to emphasize that the paper makes a novel and timely contribution to the code-generation for tensor programs. There has been two major directions in code-generation. One being the black-box optimization building on TVM or TVM-like templates. Another being the auto-scheduling techniques that aim to minimize human intervention in devising the templates. The abstraction introduced in the paper initiates a third thrust of probabilistic programs. It seems also important to note that it builds on TVM.

The paper enables a stochastic search space construction which the contributions are straight-forward. Especially, Figure 3 clearly demonstrates an example. However, it would assist readers better if there could be some table that enumerates the additional methods used in the probabilistic program such as Sample-Tile and so on. Also, it would be good to add an example that contrasts different methods for tensor program generation.

One issue that I have with the paper is that unlike AutoTVM for instance, it does not outline the optimization details. For example, it is not intuitive as how to incorporate pre-trained cost-models from datasets such as TenSet. It seems that lack of these information in this respect leaves readers with lots of questions. How do we incorporate pre-trained cost-models? Is the “Learning-driven search” part exactly equivalent to the AutoTVM and its variants?, In that case doesn’t that make transfer learning more difficult as the search spaces become more heterogeneous? While these are just some hypothetical questions that I enumerated, adding more details to the Section 4 would prevent the confusion.

Regarding evaluation, it seems that the largest benefits come from (1) multi-level tiling and (2) use tensor cores. However, reasons behind this nor insights were shared with the readers. This seems to make it difficult for readers to gain more insights as to how to optimize the search.

Lastly, the paper is focused on CPUs and GPUs. Considering the explosion of number of accelerators in the recent decade, it would be interesting to see how this can be incorporated into compilation for these diverse set of accelerators.

In sum, despite the above issues, the paper makes in important contribution that would enable interesting works on search and optimization. To this end, I would like to recommend the paper for a Weak Accept. However, with more concreteness to the paper addressing the issues above, I would be willing to reevaluate.
strengths_and_weaknesses: + timely contribution of probabilistic programming to tensor program compilation
+ abstraction proposed in Section 3 can instigate interesting research in optimization that would benefit the readers.

- unclear on details in Section 4 about learning
- unclear on how it relates to prior works in more detail (not just high level statements in 3.3)
questions: * Can you fortify the Section 4?
* Would this new abstraction require changes over prior works on black-box optimization given the probabilistic programs would render much diverse set of search spaces?
* How can transfer learning in AutoTVM be enabled? Also, datasets such as TenSet?
* How would the Figure 8 look like with some fused operations such as Conv + Pool + Relu? Could you share some results if possible with some of the common patterns in DNNs?
limitations: NA
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: Program optimization relies on rewriting a given program into an equivalent, but more efficient version of itself. These equivalent programs are searched within hand-engineered search spaces which are typically hard to customize. This paper introduces SpaceCraft, a domain-specific probabilistic-programming-language-abstraction which enables easy customization and construction of tensor program search spaces. SpaceCraft achieves this goal by providing modularity for composing program transformations. Further, the authors combine SpaceCraft with MAP estimation for automatic tensor programs optimization. SpaceCraft can cover the optimization search space used in many prior work and optimization with SpaceCraft delivers similar performance to prior work. On growing the search-space with a hardware specific module, the optimization yields a significant improvement over prior work.
strengths_and_weaknesses: # Strengths:

1. The program tackles an very useful problem - construction of the program search space. Typically,  construction of the "right" search space requires many trial and errors. SpaceCraft can significantly help reduce the complexity of this process.
2. Clarity: The figures effectively convey the intended message. 
3. The use of probabilistic programs for tensor program optimization is novel.
4. SpaceCraft covers the prior methods under a single framework.

# Weaknesses:

1. **Presentation**:  The presentation of the paper requires further refinement. Apart from plenty of minor grammatical mistakes (specified below), the paper also consist of many confusing sentences: 
	1. The paper fails to specify what exactly are the hurdles with search space construction/customization in prior work. The paper simply states previous methods "... requires surgical modification...".  What are the surgical modifications? Why are they costly/infeasible?  
	2. In Line 44-52, paper states that " These two observations lead to a simple yet powerful abstraction ..." but do not to specify what exactly this abstraction is.  
	3. In Line 124-126, the paper states that SpaceCraft captures long-term structural arithmatic dependency between random variables and tensor program. The paper would benefit from specifically state why SpaceCraft achieves this. 
	4. The presentation in Section 3.3 is very confusing. In this section, the reader needs to be aware of the authors intentions to understand some sentences. For example,  in Lines 163-165, what is implied by "it" is not clear. 

2.  **Experiments**: There are multiple drawbacks with the experiments section in the paper:
	1. A central benefit of SpaceCraft is that it simplifies the construction and modification of the optimization search spaces. This benefit however, and not been evaluated at all. This benefit should be evaluated by comparing the human efforts required for introducing new modules in prior frameworks versus the effort required for the same in SpaceCraft. 
	2. The paper states (Line 278) that it takes 2 days to introduce a new module, but what is the reader supposed to compare this number with? What is the time it would take for introducing similar modules in prior work? 
	3. The paper does not compare their optimization results to newer optimization techniques such as AutoTVM [1] and CHAMELEON [2] making it unclear how well this framework fares w.r.t. the current state-of-the-art. 
	4. Comparisons in Figure 10 (b) are unfair as the search space of SpaceCraft and TVM are different. Are there fundamental limitations to TVM which make expanding the search space with operations from "Use-Tensor-Core" module infeasible? 
	5. All the plots in Figure 8, 9, and 10  present results with normalized performance without clarifying what is the normalized performance measures. Further, the paper would benefit if more interpret-able statistics such as wall clock time or TFLOPS are used.
	6. The paper does not contain any measurement of the time it takes to perform the proposed automatic optimization. This is important as a high optimization cost would make the usage of SpaceCraft infeasible for many situations.   

## Minor Weaknesses


1. In the conclusions, the paper attempt to say they will open-source the framework. However, in the checklist, they mention the code and data is proprietary. In that case, how will they eventually open source the framework?
2.   In the checklist, the authors state that limitation of the work are described in section 7. However, there are no sentences that discuss any limitations of the work in the entire paper.  


## Errors: 

Line 18: has become

Line 70 : Then, the framework finds

Line 83: optimization programs

Line 84: Therefore/Hence, we need

Line 87: build a generic

Line 288: We will open-source our framework 

Line 481: in the Appendix. 


## References:


[1]  Learning to Optimize Tensor Programs, Tianqi Chen et al., NeurIPS 2018

[2] Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation, Byung Hoon Ahn et al., ICLR 2020
questions: Can the authors can address the weaknesses raised in the previous section? 

Specifically:
 
1. Can they address the issues with the text of the paper?
2. Can the address the drawbacks of their experimental design? 
3. Can the authors explain why certain important experiments have not been conducted?
limitations: No. The authors have not discussed limitations of this approach. The authors can discuss trade-offs of using MAP estimation vs. techniques such as reinforcement learning or simulated annealing used in prior work. The authors can also highlight the cost of learning a new DSL for crafting search spaces in SpaceCraft.
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: This paper introduces SpaceCraft, which improves the efficiency of designing tensor programs’ schedulers via establishing the search space in a composable and modular way. The key of SpaceCraft is to use the probabilistic programming method to describe the search space. Designers can improve their efficiency by combining SpaceCraft’s reusable program transformation modules with prior domain knowledge. The paper also introduces a learning-driven searching framework that can search efficiently under the guidance of a learned cost model. The experiments demonstrate that SpaceCraft achieves comparable optimization ability against TVM. 
strengths_and_weaknesses: This paper proposes a new programming model, which abstracts the search space as a probabilistic language and allows incorporation of domain knowledge, for tensor program optimization. The proposed method is technically sound. It seems that prior wok such as AutoTVM and Ansor can be easily implemented by using SpaceCraft, and SpaceCraft works well for domain-specific accelerators such as Tensor Core. The related work, including template-guided auto-tuning and auto-scheduling, is adequately cited.

This work is evaluated in terms of both expressiveness and optimization performance. Although it achieves comparable performance against TVM, it would be quite helpful to compare the performance with state-of-the-art approaches including FlexTensor and Ansor. Moreover, as SpaceCraft can provide a better search space, I suggest to directly compare the quality of produced search spaces against previous work, for example, ratio of high-performance candidate programs or performance difference of candidate programs, etc.

The advantage of the proposed learning-driven search over existing AutoTVM is unclear, which employs the cost model as well. I would like to see more ablation study on the effectiveness of the proposed learning framework.
questions: Please address the weaknesses listed above. Besides, in Figure 10(b), does TVM use Tensor Core as well? 
limitations: Supporting tensor computation unit is critical while challenging for existing exploration-based approached. Both FlexTensor and Ansor cannot support Tensor Core in a convenient manner. The readers would expect to see more explanation on the Use-Tensor-Core module.
review_title: null

Meta-review: metareview: With many of the changes proposed by the authors in response to reviewer feedback, I think this will be a very good paper, arguably warranting even higher scores than the updated reviewer scores would indicate, which are already unanimously in favor of acceptance post-rebuttal. The performance improvements obtained here seem very significant, and I think that this paper represents an enormous engineering effort, which I think has potentially enormous value to researchers and practitioners training large end-to-end deep learning models. Despite the paper's significant focus on engineering and systems concerns, I think the potential value to machine learning here is quite clear, and therefore clearly a fit for NeurIPS.
confidence: Certain
Decision: Accept