[Paper Review] Generalizing Shape Analysis with Gradual Types

Migeed, Zeina, Reed, James+2 more|arXiv (Cornell University)|Apr 17, 2019

Parallel Computing and Optimization Techniques42 references16 citations

TL;DR

Relay is a high-level, statically typed, functional intermediate representation (IR) for deep learning compilers that unifies and generalizes existing IRs to support expressive models, composable optimizations, and portable hardware targeting. It achieves competitive performance across CPUs, GPUs, and accelerators by reframing domain-specific optimizations as standard compiler passes and enabling extensible, platform-agnostic code generation.

ABSTRACT

Frameworks for writing, compiling, and optimizing deep learning (DL) models have recently enabled progress in areas like computer vision and natural language processing. Extending these frameworks to accommodate the rapidly diversifying landscape of DL models and hardware platforms presents challenging tradeoffs between expressivity, composability, and portability. We present Relay, a new compiler framework for DL. Relay's functional, statically typed intermediate representation (IR) unifies and generalizes existing DL IRs to express state-of-the-art models. The introduction of Relay's expressive IR requires careful design of domain-specific optimizations, addressed via Relay's extension mechanisms. Using these extension mechanisms, Relay supports a unified compiler that can target a variety of hardware platforms. Our evaluation demonstrates Relay's competitive performance for a broad class of models and devices (CPUs, GPUs, and emerging accelerators). Relay's design demonstrates how a unified IR can provide expressivity, composability, and portability without compromising performance.

Motivation & Objective

• Address the growing challenge of supporting diverse deep learning models and hardware platforms in a unified, extensible compiler framework.
• Overcome limitations in existing DL IRs that sacrifice expressivity, composability, or portability when extending for new models, optimizations, or hardware.
• Enable seamless integration of domain-specific optimizations (e.g., quantization, fusion) and hardware-specific code generation without compromising performance.
• Provide a principled, extensible foundation for future deep learning compiler research by unifying model representation, optimization, and code generation.

Proposed method

• Design Relay as a functional, statically typed IR that supports first-class functions, control flow, and complex data structures like trees and graphs.
• Reframe common DL optimizations—such as quantization, shape inference, and operator fusion—as standard compiler passes using established compiler techniques.
• Introduce extensible mechanisms for defining new operators and optimizations, enabling modular and composable extension of the compiler stack.
• Implement a platform-agnostic operator representation that allows the same IR to target diverse backends, including CPUs, GPUs, and custom accelerators.
• Use a unified IR to compile entire models with complex control flow (e.g., RNNs, TreeLSTMs) into lean, optimized binaries without relying on host-language scripting.
• Leverage decades of traditional compiler research to compose optimizations in a way that is both correct and efficient, avoiding combinatorial explosion in implementation.

Experimental results

Research questions

• RQ1 Can a functional, statically typed IR unify and generalize existing deep learning IRs to express state-of-the-art models with complex control flow and data structures?
• RQ2 Can domain-specific optimizations like quantization and operator fusion be composed effectively when reframed as standard compiler passes?
• RQ3 Can a single, extensible IR achieve high performance across diverse hardware platforms (CPUs, GPUs, FPGAs, accelerators) without sacrificing expressivity or composability?
• RQ4 To what extent can Relay outperform existing frameworks in optimizing and deploying models with complex control flow and heterogeneous hardware targets?

Key findings

• Relay achieves competitive performance on vision and NLP workloads, matching or exceeding state-of-the-art frameworks: 2.3× speedup over MxNet on GRU and 2× speedup on TreeLSTM.
• Relay outperforms PyTorch by 1.4× on CharRNN and 2× on TreeLSTM due to its ability to compile complex control flow into a single optimized binary.
• Quantized inference on ARM platforms (Raspberry Pi 3 and Firefly RK3399) shows minimal accuracy loss (e.g., 69.4% vs. 70.7% for ResNet-18) with significant performance gains under aggressive quantization (8/16).
• On FPGA-based DNN accelerators, Relay enables effective evaluation of different hardware designs, showing batching improves throughput on memory-bound NLP workloads like TreeLSTM but not on compute-bound ResNets.
• Operator fusion and layout transformation yield significant speedups on CPU and GPU: up to 2.3× on MobileNet and ResNet due to fusion of elementwise operations.
• Relay successfully targets multiple hardware backends—including TPU, Inferentia, and custom FPGAs—demonstrating strong portability without sacrificing optimization power.