HackerPapers

arXiv:1810.09868cs.PL

Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Keno Fischer, Elliot Saba

Abstract

Google's Cloud TPUs are a promising new hardware architecture for machine learning workloads. They have powered many of Google's milestone machine learning achievements in recent years. Google has now made TPUs available for general use on their cloud platform and as of very recently has opened them up further to allow use by non-TensorFlow frontends. We describe a method and implementation for offloading suitable sections of Julia programs to TPUs via this new API and the Google XLA compiler. Our method is able to completely fuse the forward pass of a VGG19 model expressed as a Julia program into a single TPU executable to be offloaded to the device. Our method composes well with existing compiler-based automatic differentiation techniques on Julia code, and we are thus able to also automatically obtain the VGG19 backwards pass and similarly offload it to the TPU. Targeting TPUs using our compiler, we are able to evaluate the VGG19 forward pass on a batch of 100 images in 0.23s which compares favorably to the 52.4s required for the original model on the CPU. Our implementation is less than 1000 lines of Julia, with no TPU specific changes made to the core Julia compiler or any other Julia packages.

arXiv comments
Submitted to SysML 2019
Journal reference
DOI
Subjects
cs.PLProgramming Languages cs.LGMachine Learning stat.MLMachine Learning (Statistics)

On Hacker News

Linked by 3 identifiable commenters in 4 comments across 4 threads, from Jan 2019 to Jul 2019. Found in 4 HN submissions from Oct 2018 to Oct 2018. Open submission discussions.

Comments linking this paper, per year.

Discussed in comments4 threads

  1. A Differentiable Programming System to Bridge ML and Scientific Computing 1 comment

    • JAX is a sophisticated and well implemented high level tracer, combined with a backend compiler that makes use of XLA. High level tracing has some fundamental draw backs in key features we want (efficient scalar AD for example), so the approach being advocated here is doing AD…

      KenoFischerView on HNParent context

  2. Julia on Google TPU: Shakespeare RNN 1 comment

    • There's very little TensorFlow here. We're using it as a glorified grpc client basically to ship XLA to the TPU. The actual mechanism is described in https://arxiv.org/abs/1810.09868. The model definition itself uses the layers from Flux, but there's a couple assumptions that…

      KenoFischerView on HNParent context

  3. Differentiating SSA-Form Programs in Julia (2018) [pdf] 1 comment

    • As I recall, after the publication of [1] on Julia+TPUs there was some talk about doing a comparison of some ML models on the CPU, GPU, and TPU. Is this something you're working on? [1] https://arxiv.org/pdf/1810.09868.pdf

      chrispeelView on HNParent context

  4. High Performance Numeric Programming with Swift: Explorations and Reflections 1 comment

    • >Julia's secret sauce is LLVM That most certainly is not true. See this for example: https://arxiv.org/pdf/1810.09868.pdf

      byt143View on HNParent context

Submitted to Hacker News

Connected through shared HN threads, not inferred similarity or a recommendation.