BladeDISC Documentation

The BladeDISC compilation process follows a multi-stage pass pipeline that transforms high-level TensorFlow/PyTorch IR into optimized machine code for both CPU and GPU. The pipeline is divided into several major phases:

1. TF-to-HLO Passes: Converts TensorFlow dialect to HLO (High-Level Optimizer) dialect. This phase includes standard MLIR-HLO pipelines and BladeDISC-specific passes like DiscLowerTfPass (for custom ops like RandomUniform and TopK) and ReviseArgumentsForStaticRankPass.
2. HLO Graph Optimization & Placement: Optimizes the HLO graph and decides where operations should execute. Key components include the ShapeSimplifier pass (for shape propagation and constraint insertion) and Placement Passes (which mark shape-calculation ops to be executed on the CPU).
3. Bufferize Passes: Transitions the IR from the 'tensor world' to the 'buffer world' by explicitly emitting allocation and deallocation logic. This includes converting mhlo to lmhlo (the bufferized representation) and assigning memory spaces (CPU vs. GPU).
4. LHLO Graph Optimization Passes: Performs optimizations on the bufferized representation, most notably the Fusion Pass (using 'base' or 'stitch' strategies) and the Speculation Pass (generating multiple kernel versions for different runtime conditions like vectorization or implicit broadcasts).
5. Runtime & Library Call Related Passes: Integrates the Runtime Abstraction Layer (RAL). It injects the RAL context and rewrites custom call ops into disc_ral.dispatch ops to ensure a stable ABI and manage stateful resources.
6. CodeGen Passes: Lowers lmhlo.fusion ops into nested loops. This phase is backend-aware (CPU vs. GPU) and uses backbone passes like DiscLhloLegalizeRootsToParallelLoopsPass and InputInlineFusionPass to generate schedules.
7. Loops to GPU / GPU Module to CUBIN: For GPU backends, this phase tiles loops, maps them to GPU blocks/threads, and lowers the GPU dialect to vendor-specific dialects (NVVM for CUDA or ROCm for AMD) before compiling to a binary blob.
8. Host Side Passes: Generates the scheduling logic (kernel launching, data movement, synchronization) and lowers the final IR to the LLVM dialect for binary generation.

The primary interface for the toolkit is the Quantizer class. You instantiate a Quantizer with configuration (such as target backend type or excluded module types) and use it to generate different types of proxy models depending on your stage in the workflow:

• Calibration Proxy: Created via Quantizer.calib(model). Used to calibrate parameters by running forward passes with typical data.
• QAT Proxy: Created via Quantizer.qat(model). Used for Quantization-Aware Training.
• Quantized Model Proxy: Created via Quantizer.quantize(model). Represents the final quantized model ready for inference or export.

from torch_quant import Quantizer

model = MyModel() # torch.nn.Module
quantizer = Quantizer()

# For calibration
calib_model = quantizer.calib(model)

# For QAT
qat_model = quantizer.qat(model)

# For final quantized model
quant_model = quantizer.quantize(model)

To maintain a stable ABI and hide the implementation details of memory structures (like MemRef), the compiler rewrites the entry function's inputs and outputs. Instead of receiving raw buffers, the function receives a context and uses recv_input and send_output API calls to interact with the RAL.

This design allows for partial execution: the compiler can place recv_input calls such that the binary starts executing as soon as specific inputs are ready, or send_output calls to stream results back before all computations are finished.

// Original IR
func @main(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>) -> memref<?x?xf32> {
  %ret = alloc(...)
  use(%arg0, %arg1, %ret, ...)
  return %ret : memref<?x?xf32>
}

// After RAL conversion
func @main(!disc_ral.context %ctx) {
  %arg0 = disc_ral.recv_input(%ctx, 0) // receive the first input
  %arg1 = disc_ral.recv_input(%ctx, 1) // receive the second input
  %ret = alloc(...)
  use(%arg0, %arg1, %ret, ...)
  disc_ral.send_output(%ctx, 0, %ret) // send the first output
}