How BladeDISC++ Cuts Memory Peaks for Dynamic‑Shape Deep Learning Models

Background and Challenges

As deep‑learning models become increasingly dynamic—varying image sizes, batch sizes, sequence lengths, and even data‑dependent shapes—traditional static‑shape compilers struggle to optimize memory usage. Existing compilers such as TVM and OpenXLA focus on static‑shape optimizations and lack effective memory‑reduction techniques for dynamic scenarios.

Uncertain tensor shapes : Compilers cannot know tensor sizes at compile time, making efficient code generation difficult.

Dynamic memory allocation : Without predetermined shape information, memory blocks cannot be pre‑allocated, leading to fragmentation.

Complex optimization algorithms : Scheduling and fusion decisions depend on shape information, which is unavailable in dynamic contexts.

BladeDISC++ addresses these issues by aiming to keep memory peaks below a threshold, enabling larger training batches with limited GPU memory.

BladeDISC++ Innovation

BladeDISC++ is a dynamic‑shape AI compiler built on MLIR. It serves as the backend for TorchAcc, which captures PyTorch graphs via GraphCapture and converts them to StableHLO. BladeDISC++ then applies three major optimization categories:

Memory optimization : Analyzes buffer lifetimes to perform automatic offloading and recomputation.

Computation optimization : Includes operator fusion, graph simplification, and custom operator support.

Communication optimization : Utilizes multiple CUDA streams for asynchronous communication.

It supports major hardware platforms including Nvidia, AMD, and Intel CPUs.

BladeDISC++ Overview

In dynamic‑shape scenarios, tensor shapes are unknown while the graph topology remains fixed. BladeDISC++ performs joint compile‑time and run‑time optimization through the following core stages:

Operation Fusion : Merges memory‑intensive or GEMM operators to reduce kernel launch overhead and improve shared‑memory utilization.

Operation Scheduling : Reorders operators based on symbolic shape analysis to lower the total size of live tensors and thus the memory peak.

Automatic Rematerialization (Auto Remat) : Uses offloading and recomputation to release tensors when memory exceeds a threshold, regenerating them as needed.

Symbolic Shape Graph

The MLIR IR shows a dynamic‑shape graph where unknown dimensions are denoted by ?. BladeDISC++ introduces SymbolicDimOp to bind symbolic dimensions to tensors, e.g., %arg0: tensor<?, [$S0]>. It then constructs SymbolicExpr to represent element counts, such as expr1 = 11008 * @S1; and expr2 = 1024 * @S0;. By simplifying these expressions using operator semantics (e.g., Reshape preserves element count), BladeDISC++ can infer relationships like @S0 = 12 * @S1, allowing size comparisons.

Operation Scheduling

By adjusting execution order, the compiler reduces memory peaks. For example, scheduling A→C→E→B→D→F yields a total live tensor size of 5, whereas A→C→B→D→E→F results in size 6. In dynamic‑shape contexts, BladeDISC++ uses symbolic expressions to compare memory requirements, selecting the schedule with the lower estimated peak.

Just‑In‑Time Auto Rematerialization (JIT Remat)

Static‑shape Auto Remat predicts memory peaks and inserts EvictOp and RegenerateOp to offload or recompute tensors. BladeDISC++ extends this to dynamic shapes by performing a compile‑time search for potentially evictable tensors and inserting placeholder ops. At runtime, actual tensor shapes are combined with the symbolic graph to estimate memory usage and decide whether to offload or recompute, minimizing end‑to‑end performance loss.

Performance Evaluation

Experiments on a trimmed Llama2‑7B model (4 hidden layers) compare three settings: (1) Dynamic‑shape BladeDISC without memory optimization, (2) Static‑shape BladeDISC with memory optimization, and (3) BladeDISC++ with memory optimization in dynamic‑shape mode. Results show that at batch size 14, BladeDISC++ matches the memory peak of static‑shape while maintaining similar throughput. At batch size 18, the unoptimized dynamic version runs out of memory, the static‑shape version runs but with reduced throughput due to padding, whereas BladeDISC++ keeps memory usage comparable to static‑shape and improves throughput by about 11%.

Conclusion

Dynamic‑shape workloads face tensor shape uncertainty, which BladeDISC++ tackles by constructing a Symbolic Shape Graph and applying symbolic‑based operation scheduling and joint compile‑time/run‑time JIT auto‑rematerialization. Experiments on Llama2 demonstrate that BladeDISC++ effectively reduces memory consumption to levels comparable with static‑shape compilers while achieving higher throughput, representing a pioneering approach for dynamic‑shape AI compilation.