MLIR

TensorFlow uses LLVM's MLIR (Multi-Level IR) for an increasing share of its compiler work: TFLite conversion, TF→XLA lowering, quantization, and StableHLO export. Lives under tensorflow/compiler/mlir/, with smaller pieces in tensorflow/core/ir/ and tensorflow/core/transforms/.

Purpose

• Express TF computations in IR forms that are easier to optimize than GraphDef.
• Provide dialects for the various semantic layers (TF graph, TF executor, TFLite, quantization, StableHLO, MHLO, TFG).
• Run passes that lower between dialects and apply optimizations (canonicalization, fusion, layout, quantization).
• Serve as the basis of:
  • The TFLite converter (replaces TOCO).
  • Newer TF→HLO paths (replacing parts of tf2xla).
  • StableHLO export and import.

Directory layout

tensorflow/compiler/mlir/
├── tensorflow/                 # TF dialect: tf.* operations
├── tf2xla/                     # MLIR-based TF→HLO bridge
├── lite/                       # TFLite (tfl) dialect, converter, quantizer
├── quantization/               # Generic quantization framework
├── stablehlo/                  # StableHLO ↔ MHLO conversions
├── tfrt/                       # Lowering to the TFRT runtime
├── tools/                      # MLIR command-line tools (mlir-translate, etc.)
├── transforms/                 # Cross-dialect passes
├── utils/
└── ...

tensorflow/core/
├── ir/                         # `tfg` dialect (TFGraph) — yet another TF MLIR dialect
└── transforms/                 # Passes for the tfg dialect

TF dialects in this repo

Several dialects coexist, each focused on a layer of abstraction:

Dialect	Purpose	Defined in
tf	One-to-one mirror of TF op semantics; the entry point for tracing.	tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td
tf_executor	Mirrors TF's executor (control deps, frames, switch/merge).	tensorflow/compiler/mlir/tensorflow/ir/tf_executor.td
tf_device	Multi-device annotations (replicate, cluster).	tensorflow/compiler/mlir/tensorflow/ir/tf_device.td
tfl	TFLite ops (mirror of TFLite builtins).	tensorflow/compiler/mlir/lite/ir/tfl_ops.td
tf_quant	Quantization framework ops.	tensorflow/compiler/mlir/quantization/
mhlo/stablehlo	Portable HLO; the canonical export format.	tensorflow/compiler/mlir/stablehlo/
tfg	"TF Graph" dialect — yet another TF representation, targeted at runtime fallbacks.	tensorflow/core/ir/

The dialects are defined using TableGen *.td files; passes are C++ classes registered with the MLIR pass manager.

The TFLite conversion pipeline

graph LR
    SM[SavedModel / Keras / concrete fn]
    TF[tf dialect]
    TFLPrep[Prep passes legalize control flow, constant fold]
    TFL[tfl dialect]
    Quant[Quantization passes optional]
    Flatbuf[FlatBuffer .tflite]

    SM --> TF
    TF --> TFLPrep
    TFLPrep --> TFL
    TFL --> Quant
    Quant --> Flatbuf

The high-level call tf.lite.TFLiteConverter.from_saved_model(...) ends up running this pipeline (tensorflow/compiler/mlir/lite/). Passes include legalising control flow, fusing batch-norm into conv, quantizing activations and weights, and finally writing a flatbuffer that matches tensorflow/lite/schema/schema.fbs.

The TF→HLO MLIR pipeline

tensorflow/compiler/mlir/tf2xla/ contains a parallel TF→HLO path that runs MLIR passes (tf → mhlo/stablehlo) and produces the same HloModule the older tf2xla would. This path is increasingly the default for TPU and CPU XLA paths.

StableHLO

StableHLO is a versioned MLIR dialect for HLO operations, intended to be a portable, stable export format shared across TF, JAX, and PyTorch/XLA. Conversions live in:

• tensorflow/compiler/mlir/stablehlo/ — TF↔StableHLO and MHLO↔StableHLO.
• tensorflow/lite/stablehlo/ — TFLite-side StableHLO support.

Recent commits like "Integrate StableHLO at openxla/stablehlo@…" (visible in git log) update the vendored StableHLO version.

TFRT lowering

tensorflow/compiler/mlir/tfrt/ lowers TF MLIR into the TFRT runtime's IR ops, so a graph can run on the TFRT async runtime. TFRT-internal ops live in MLIR too (the tfrt dialect, defined in TFRT's own repo).

Tools

• tf-opt (tensorflow/compiler/mlir/tools/tf-opt) — generic MLIR opt driver with all TF dialects loaded.
• tf-mlir-translate — converts between MLIR text and TF/HLO/StableHLO formats.
• tf-tfrt-opt — TFRT-flavoured opt.
• tflite_converter-style tools live under tensorflow/compiler/mlir/lite/.

Integration points

• TFLite converter uses MLIR for the actual conversion. See apps/tensorflow-lite.
• jit_compile=True sometimes runs through the MLIR path before producing HloModule.
• TFRT consumes MLIR.
• StableHLO export for sharing models with JAX/PyTorch goes through MLIR.

Entry points for modification

• New TF dialect op — add to tf_ops.td, register a builder, add lowering passes if needed.
• New conversion pass — tensorflow/compiler/mlir/<area>/transforms/<your_pass>.cc. Register with the pass manager.
• TFLite converter changes — tensorflow/compiler/mlir/lite/. New op support usually needs both a TFL op def and a tf→tfl legalisation pattern.
• StableHLO version bump — typically a single commit that updates the pinned StableHLO version (the recurring "Integrate StableHLO at openxla/stablehlo@…" commits are exactly this).

Related

• compilers/xla — the older tf2xla path that MLIR is gradually replacing.
• apps/tensorflow-lite — the converter that lives on top of MLIR.
• systems/tfrt — TFRT runs MLIR-lowered programs.