Machine learning workloads are often expensive to train, taking weeks to converge. The current generation of frameworks relies on custom back-ends in order to achieve efficiency, making it impractical to train models on less common hardware where no such back-ends exist. Knossos builds on recent work that avoids the need for hand-written libraries, instead compiles machine learning models in much the same way one would compile other kinds of software. In order to make the resulting code efficient, the Knossos complier directly optimises the abstract syntax tree of the program. However in contrast to traditional compilers that employ hand-written optimisation passes, we take a rewriting approach driven by the $A^\star$ search algorithm and a learn value function that evaluates future potential cost reduction of taking various rewriting actions to the program. We show that Knossos can automatically learned optimisations that past compliers had to implement by hand. Furthermore, we demonstrate that Knossos can achieve wall time reduction compared to a hand-tuned compiler on a suite of machine learning programs, including basic linear algebra and convolutional networks. The Knossos compiler has minimal dependencies and can be used on any architecture that supports a \Cpp toolchain. 
 Since cost model the proposed algorithm optimises can be tailored to a particular hardware architecture, the proposed approach can potentially applied to a variety of hardware. While the development of any kind of software can benefit from compliers able to produce fast code, runtime efficiency is particularity important for modern machine learning. In particular, because modern models they can take weeks to train (OpenAI, 2018) , complier optimisations that lead to execution speed-ups are of huge value. In parallel, machine learning is being deployed on a variety of diverse devices ranging from wearables to huge clusters clusters of powerful GPUs. Since each architecture has different performance profile and requires different code optimisations, it is difficult to provide tooling that works fast on all of them. Traditionally, the tension between performance and interoperability is resolved by machine learning frameworks (Paszke et al., 2017; Abadi et al., 2016) . In these frameworks, while code execution is outsourced to hardware-specific back-ends such as XLA (XLA authors, 2016) . While this approach has seen huge initial success, the cost of providing customised back-ends for each target architecture is prohibitive. Moreover, the frameworks also custom front-ends that require the programmer to specify the model being trained as a compute graph. Since the compute graph has semantics separate from the host programming language, this process is often error-prone and time-consuming. In order to address these obstacles, a new generation of tools has recently appeared that transform machine learning code using the same techniques that have been used for compiling traditional software. The need for a separate front-end API for machine learning operations is eliminated by including automatic differentiation as a first-class feature of the complied language (Innes et al., 2019; Frostig et al., 2018) . Instead of custom back-ends, modern machine learning compliers use an intermediate representation and perform extensive code optimisations (Innes et al., 2019; Frostig et al., 2018; van Merrienboer et al., 2018; Wei et al., 2018; Sotoudeh et al., 2019; Rotem et al., 2018) . In addition, program optimisation is being modelled as a machine learning task itself, with the complier learning how to perform rewrites (Chen et al., 2018b; a) . in mind. We formalize program optimisation as a finite-horizon Markov Decision Process (MDP), with the reward signal determined by the cost of executing a program. By solving this MDP, we are able to produce fast code tailor-made for any given task and architecture, without relying on backend-specific hand-written libraries. Knossos works by re-writing programs written in an intermediate representation (IR). Akin to JAX (Frostig et al., 2018) and Zygote (Innes et al., 2019) , all Knossos functions are potentially differentiable, avoiding the syntactic awkwardness that arises from embedding a differentiable program in a host language. The IR can then be transpiled, allowing it to run on any platform that supports a C ++ toolchain. This allows Knossos code to be seamlessly deployed on specialized or embedded hardware without the need of manual tuning, both for training and for deployment of models, enabling a much broader user base than competing approaches. To our knowledge, Knossos is the first compiler that combines RL-based program optimisation, firstclass support for deep learning primitives and the ability to target any architecture supporting the C ++ toolchain. We defer detailed scope comparisons with prior work to Section 4. We empirically demonstrate the benefits of our program optimisation in Section 5, showing that Knossos was able to automatically learn loop fusion, a type of compiler optimisation that previously had to be applied manually. We have introduced Knossos, a new complier targetting machine learning and numerical computation. Thanks to its automatic code optimisation, Knossos produces binaries that achieve better run-times than a traditional, rule-based complier. Knossos can deal with complex code generated by automatic differentiation and automatically discover optimisations that previously required careful complier design. We believe that Knossos will pave the way towards a new generation of future compliers, which will crucially rely on automatically inferring the correct optimisations. It also has a LISP-like surface syntax, which we used to implement our programs. In the future, we plan to provide transpilers, allowing for the compilation of code written in other languages into Knossos. We provide a sample Knossos program in Figure 4 .In order to facilitate Machine Learning workloads, the Knossos IL has native support for automatic differentiation. We use a new unified view of automatic differentiation as generalised transposition (Elliott, 2018) . Rather than having an explicit distinction between forward mode and reverse mode AD, Knossos uses uses a type system together with a set of consistent rewrite rules. Whenever the gradient operator is used as part of a Knossos algorithm, the complier first generates a syntax tree corresponding to the differentiated program and then applies rewrites to optimize the cost of its execution. This means that the resulting AD algorithm is tailor-made and optimized with that exact use case in mind. This is in contrast to systems such as PyTorch, which have hard-coded routines for backward-mode AD. From the perspective of the user, this process is completely transparent in the sense that taking gradients can be applied to any piece of Knossos code. While the details of this process are beyond the scope of this paper, from the perspective of this work, the important feature of AD is that it corresponds to a transformation of the abstract syntax tree. The resulting AST can then be optimised in the same way as any other code.
We combine A* search with reinforcement learning to speed up machine learning code
{"src": "scitldr-AIC_train", "paper_id": "SylyHkHYDB"}