[ML Compiler] BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach

[ljwoo94]

Info

• ACM article From Alibaba (2023)

Abstraction

• Dynamic tensor shape
  • Unlike Image processing, language models such as LLM need to handle dynamic tensor shape since the input length/shape is not static.
  • the tensor could have varied values at runtime.
• Problem of current ML Compilers
  • Can only handle static shape
  • Can handle dynamic shape but, expose performance problems for operator fusion and code generation
• This paper asserts compilers need to design a cross-level symbolic shape to tackle absence of shape values.
  • Operations rely on adjacent shape rather than exact value.
  • Need compile-time and runtime combined code generation approach

https://github.com/alibaba/BladeDISC

[ljwoo94]

Introduction

• Due to Increase of data diversity, importance of dynamic shape model optimization is growing.
• Common practice (접근) to deal with dynamic shape
  • Pad all dimensions to the max value within each batch
  • Requires dynamic recompilation of graph whenever new shape comes.
  • Overhead & Memory Consumption
• Examples
  • XLA-Optimization (in TensorFlow)
  • x20 more execution time than native TF due to the recompilation time.
  • JIT compilation (Just-In-Time - JAVA/C#)
  • long warm-up process due to the frequent compilations. (more service time for each new service machine)
  • Large compilation cache
  • AOT compilation (Ahead-of-Time - C/C++)
  • Not practical t enumerate and compile all possible shapes
  • Padding in max range results redundant computations.
  • BERT on TensorRT : x2.41 performance degradation
• Contributions of BladeDISC
  • abstracts the hidden shape constraint information of dynamic shape.
  • Found hidden shape information in computation logic and abstracts it for graph level optimization by building global shape constraint
  • Advanced fusion decision approach using shape constraint information
  • Locality checking(locality on registers) requires equality relationship between tensor shape of operations, not a exact shape value.
  • combination of compile time and runtime optimization for code generation
  • Compiles for multi-versions for each operation to enable runtime kernel selection for pipeline schedule.

[ljwoo94]

Background

Current Works of AI Compilers

• Due to growing computing power, memory-intensive computation and bandwidth become a rising factor of end-to-end performance.
• Therefore, static shape compilers applied stitch fusion to reduce off-chip memory traffic.
  • divide to-be-fused operators into several groups -> generate code independently -> collect groups into same kernels and save the outputs into different levels of memory according to locality with other groups.
  • 즉, Locality에 따라서, Group에 독립적일 데이터는 kernel에 가까운 register에, 다른 Group에도 사용되는 data는 공유 될 수 있는 shared 혹은 global memory에 저장하는 의미인듯. GPU의 관점으로 보인다.
  • locality checking is done by shape value propagation between groups. ( # of data items)
• JIT compilation overhead & redundant computation due to padding are key obstacles of static shape compilers when they handle dynamic shape.
  • Insufficient Fusion : cause non-computation overhead
  • Inefficient Codegen : cannot use HW power fully.
• TensorRT : Ranged-shape optimization
  • cannot use when input shape range is not able to obtain before execution

Challenges of Dynamic Shape Optimization

• Challenges due to absence of shape value
• Graph Transformation (eg. Operator Fusion)
  • SOTA (Stitch Fusion) requires data shape value to check locality.
  • GEMM Batching also requires guaranteed same shape.
• Dynamic-shaped Code Generation
  • If the tensor shape is static, it is possible to schedule codes to utilize HW resources. But, arbitrary tensor shape requires extra process to identify shape equality. (ex. implicit Broadcast in GPGPU)

[ljwoo94]

Global Symbolic Shape Information

• Compile is static action while shape values are only known at runtime.
• Fortunately, shape constraint information, that allows graph-level optimization, is hidden in the tensor computation graph.

Shape information Abstraction

• Shape Relationship

  • relationship between different tensors that is important for operator fusion. (eg. tensor shape equality)
  • 2 most important dim relationships : dim equality, dim collapse equality
  • dim equality
  • def: 2 dims are equal
  • methods:
    • Input-output infer
    • dim equality between input & output tensors.
    • Sibling Constraint
    • dim equality between several input tensors.
    • Shape value extraction
    • dim equality according to specific shape value.
    • ops like reshape map the output shape. This phenomenon allows us to consider two different dims are same when they mapped to the same value.
  • dim collapse equality
  • def: computation of group of dims is equal to another single dim.
  • methods:
    • trace integer arithmetic operators (mul, div)
    • trace shape transformation operators (reshape, transpose)

• Shape Property

  • property of a single tensor. (for optimizations like vectorization)
  • dim property
  • BladeDISC analyze constant division and value range of division. But, these properties can be easily added. (in code)
  • Heuristic ways(algebraic computing properties) of analysis.
    • For example, the result of calculation <?x?> dot <?x4> is <?x4>. If we can assume min value of ? is 512, then we can consider the size of tensor after reshape <?x4> to <y> is multiple of 4 and min value is 2048.

• Shape = a set of dimensions. Thus, shape information can be represented by fine-grained dim-information, which means if all the dims are equal for two tensors, the shapes are equal.

• These shape analysis information still has problems in 1)Reusability and 2)Stability in graph representation. Thus, BaldeDISC used [MLIR](https://en.wikipedia.org/wiki/MLIR(software)) to attach information on data types._

• MLIR = Multi-Level Intermediate Representation

  • making it very cheap to define and introduce new abstraction levels, and provide “in the box” infrastructure to solve common compiler engineering problems
  • What is MLIR?

[ljwoo94]

Advanced Fusion Decision

• Pattern based fusion (TVM) is difficult to achieve large-grained fusion for complex memory-intensive subgraph.
• Memory Intensive(AStitch) is not able to fuse compute-intensive operators.
• BladeDISC Method:
  • pipelined fusion decision : form basic fusions and merge them to build larger fusions.
  • use shape constraint information for locality checking of memory intensive stitch fusion.

Operator Fusion Pipeline

• small fusions at each pass and merge in the next pass.
• existing base fusion components can be reused.
• Element-wise fusion -> compute-intensive fusion -> Stitch Fusion
  • memory-intensive 와 compute intensive를 묶으면, 잘나오는 경우가 있음 : (Resources를 다르게 사용하기에 Good performance.)
  • 아마도 다른 resource unit을 사용하기 때문에, Pipeline이 된 것으로 보임.
  • 그러나, 이를 강제하는 것은 Memory intensive Op에게 memory access efficiency 가 낮아 문제 발생할 수 있음

Stitch-Fusion Decision