[Paper Review] DISC: A Dynamic Shape Compiler for Machine Learning Workloads

Kai Zhu, Wenyi Zhao+9 more|arXiv (Cornell University)|Mar 9, 2021

Parallel Computing and Optimization Techniques9 references4 citations

TL;DR

DISC is a dynamic shape compiler for machine learning workloads that extends MLIR with a fully dynamic shape IR (DHLO), generates efficient runtime control flow at compile time to avoid interpretation overhead, and enables host-device co-optimization through shape propagation and constraint collection. It achieves up to 3.3× speedup over TensorFlow/PyTorch and 1.8× over Nimble on dynamic shape workloads like Transformers.

ABSTRACT

Many recent machine learning models show dynamic shape characteristics. However, existing AI compiler optimization systems suffer a lot from problems brought by dynamic shape models, including compilation overhead, memory usage, optimization pipeline and deployment complexity. This paper provides a compiler system to natively support optimization for dynamic shape workloads, named DISC. DISC enriches a set of IR to form a fully dynamic shape representation. It generates the runtime flow at compile time to support processing dynamic shape based logic, which avoids the interpretation overhead at runtime and enlarges the opportunity of host-device co-optimization. It addresses the kernel fusion problem of dynamic shapes with shape propagation and constraints collecting methods. This is the first work to demonstrate how to build an end-to-end dynamic shape compiler based on MLIR infrastructure. Experiments show that DISC achieves up to 3.3x speedup than TensorFlow/PyTorch, and 1.8x than Nimble.

Motivation & Objective

• Address the performance bottleneck in dynamic shape machine learning workloads, where existing compilers suffer from high compilation overhead, memory bloat, and limited fusion opportunities.
• Overcome the limitations of static shape optimization in XLA and similar systems, which cannot handle dynamic shapes efficiently and often disable fusion for such workloads.
• Provide a solution that natively supports dynamic shape computation without relying on runtime interpretation or padding-based workarounds.
• Enable efficient kernel fusion and host-device co-optimization in dynamic shape scenarios by generating control flow code at compile time.
• Demonstrate end-to-end dynamic shape compilation using MLIR infrastructure, with support for multiple frameworks and fallback to static optimization when possible.

Proposed method

• Extend MLIR's HLO dialect to create DHLO (Dynamic HLO), a fully dynamic shape representation that natively supports dynamic shape computations without requiring new IR from scratch.
• Generate the entire runtime control flow (including shape inference, buffer management, and host-side logic) at compile time, avoiding VM interpretation overhead and enabling host-device co-optimization.
• Use shape propagation between producers and consumers to identify fusion candidates for adjacent operations with compatible dynamic shapes.
• Collect shape constraints during lowering to DHLO, enabling larger-scale fusion by reasoning about shape compatibility without full shape knowledge at compile time.
• Support mixed static and dynamic optimization by automatically falling back to static compilation when a subgraph has known static shapes.
• Integrate with TensorFlow and PyTorch via the DHLO hub, enabling end-to-end compilation of dynamic shape models.

Experimental results

Research questions

• RQ1 How can a machine learning compiler efficiently represent and optimize dynamic shape computations without relying on runtime interpretation?
• RQ2 What techniques can be used to enable kernel fusion in dynamic shape workloads when full shape information is unavailable at compile time?
• RQ3 How can host-device control flow be co-optimized to reduce runtime overhead in dynamic shape execution?
• RQ4 To what extent can a dynamic shape compiler achieve performance comparable to static compilers while maintaining flexibility?
• RQ5 Can a compiler system based on MLIR infrastructure natively support end-to-end dynamic shape compilation with minimal runtime cost?

Key findings

• DISC achieves an average 2.27× speedup over TensorFlow and PyTorch on dynamic shape workloads, with a peak speedup of 3.3× on certain models.
• For memory-intensive operations in Transformers, DISC achieves a 2.61× speedup over Nimble, primarily due to more efficient fusion enabled by shape hints and constraints.
• DISC reduces CPU time for runtime control flow to 36.6% of Nimble’s, despite a modest reduction in kernel count, due to efficient code generation and co-optimization.
• DISC achieves 1.8× speedup over Nimble on average, demonstrating the advantage of compile-time flow generation over VM-based interpretation.
• When disabled from falling back to static optimization, DISC achieves 85% of the performance of static compilation on average, with a range from 74.5% to 91.4%, indicating a performance gap due to limited fusion opportunities without full shape knowledge.
• The number of kernels is reduced from 13,924 in Nimble to 10,734 in DISC on the Transformer model, showing improved fusion efficiency despite no explicit kernel count reduction as the primary performance driver.