MLIR Dialects

A dialect encapsulates a set of operations, types, and attributes for a specific domain. This allows high-level, intuitive representation of computations.

• Examples: The tensor dialect for n-dimensional array operations, the linalg dialect for linear algebra on buffers, or a custom acme_npu dialect for a proprietary neural processor.
• Purpose: Enables compiler passes to reason about and optimize domain-specific constructs before lowering to lower-level hardware instructions.

Dialects exist at different abstraction levels, enabling a stepwise transformation from high-level, hardware-agnostic representations down to low-level, hardware-specific instructions.

• High-Level (e.g., tosa): Portable, standardized tensor operations.
• Mid-Level (e.g., linalg, vector): Loop and vector-oriented transformations.
• Low-Level (e.g., llvm, spirv): Close to machine code (LLVM IR) or GPU targets (SPIR-V).
• Flow: A model graph in a framework dialect (like torch) is progressively lowered through these dialect levels via legalization patterns.

Dialects are not siloed. The MLIR framework provides mechanisms for legalizing or converting operations from one dialect to another, which is the core of the compilation process.

• Dialect Conversion Framework: The primary infrastructure for defining rewrite patterns that map an operation in a source dialect to one or more operations in a target dialect.
• Type Conversion: Handles the transformation of custom types (e.g., a tensor<f32> to a memref<f32>).
• Unregistered Operations: The compiler will fail if an operation cannot be converted to a dialect supported by the final target backend.

A primary use case is creating hardware-specific dialects to represent the unique capabilities and constraints of an accelerator. This is central to hardware-aware compression and optimization.

• Target-Specific Operations: Define ops like acme.npu.conv2d_int4 that explicitly use 4-bit integer arithmetic.
• Hardware Intrinsic Mapping: Directly model specialized hardware units (e.g., tensor cores, systolic arrays).
• Constraint Encoding: Embed hardware constraints (e.g., memory alignment, supported data layouts) directly into the dialect's type system and operation semantics.

A concrete flow for deploying an AI model illustrates dialect roles:

1. Start: Model imported into MLIR as tosa (Tensor Operator Set Architecture) dialect – a portable, high-level specification.
2. Optimize: tosa ops are converted to linalg dialect for loop fusion, tiling, and data layout transformations.
3. Vectorize: linalg ops are lowered to the vector dialect for explicit SIMD instruction mapping.
4. Target: vector ops are finally legalized to a custom acme_npu dialect, which represents the exact instructions for the target Neural Processing Unit.
5. Emit: The acme_npu dialect is translated to the accelerator's machine code.

Dialects work in concert with other MLIR components:

• Passes: Transformation algorithms that operate on MLIR code, often by matching and rewriting patterns within a dialect or across dialects.
• Regions & Blocks: Dialects use MLIR's structural constructs to represent control flow (e.g., scf dialect for loops, cf for branches).
• Attributes: Dialect-specific metadata attached to operations (e.g., padding = [1,1] on a convolution op).
• ODS (Operation Definition Specification): A declarative language using TableGen to define dialects, automating the generation of C++ classes, parsers, and verifiers.

Vendors and research groups create custom MLIR dialects to model the unique capabilities and constraints of their specific AI accelerators (NPUs, DPUs). These dialects act as the Hardware Abstraction Layer (HAL) within the MLIR stack.

• Represent vendor-specific operations (e.g., custom activation functions, specialized memory operations).
• Capture hardware intrinsic properties like supported data types, memory hierarchies, and systolic array dimensions.
• Enable hardware-aware compression by allowing optimization passes that understand the cost model of the target silicon, such as favoring operations that map to efficient hardware-specific kernels.

MLIR is a compiler infrastructure framework for creating reusable and extensible compiler toolchains. Unlike traditional compilers with a fixed set of intermediate representations (IRs), MLIR allows the definition of custom, domain-specific dialects. Its core innovation is a unified infrastructure that supports multiple levels of abstraction—from high-level, graph-based operations down to low-level, hardware-specific instructions—within a single framework. This enables progressive lowering and sophisticated transformations across different representation levels.

Dialect lowering is the fundamental transformation process in MLIR where operations from a higher-level, more abstract dialect are progressively converted into operations of a lower-level dialect, eventually targeting a hardware-specific or executable format. This is a multi-step process:

• Legalization: Converting illegal operations for a target into legal ones.
• Conversion: Using rewrite patterns to map one dialect's ops to another's.
• Lowering culminates in generating LLVM IR, which can then be compiled to machine code. This process is where hardware-aware optimizations, like operator fusion for a specific NPU, are applied.

TOSA is a standardized, portable MLIR dialect that defines a complete set of tensor operations for machine learning inference. Developed by the ML industry, it serves as a stable, hardware-agnostic intermediate representation.

• Purpose: Provides a common target for framework compilers (like from TensorFlow or PyTorch) and a common source for hardware backend compilers.
• Key Feature: Its operations have well-defined numerical behavior, ensuring predictable accuracy across different platforms.
• Role in MLIR: TOSA is a crucial mid-level dialect in a lowering pipeline, sitting between framework-specific dialects (like torch) and vendor-specific hardware dialects.

The Linalg dialect is an MLIR dialect designed for structured, loop-based computations on dense tensors, such as matrix multiplications, convolutions, and generic affine loops. It is a key abstraction layer for high-performance code generation.

• Structured Ops: Operations like linalg.matmul and linalg.conv_2d capture computation semantics without specifying low-level loop ordering or tiling.
• Transformations: The dialect is designed for powerful polyhedral-style transformations (tiling, fusion, promotion) that can be applied automatically or guided by a compiler.
• Bridge Role: It often serves as an intermediate step when lowering from graph-based dialects (e.g., TOSA) down to low-level parallel loops and, eventually, hardware-specific instructions.

The Affine dialect is an MLIR dialect for representing nested loops with static control flow and affine loop bounds and memory accesses. It is essential for automatic optimization and parallelization of compute-intensive kernels.

• Affine Maps: Uses mathematical affine expressions to describe loop bounds and array indices, enabling advanced dependence analysis.
• Optimizations: Facilitates transformations like loop fusion, tiling, unrolling, and vectorization in a provably correct manner.
• Position in Pipeline: Typically appears after the Linalg dialect, where structured operations are lowered to explicit affine loops before further lowering to LLVM or GPU dialects.

These are MLIR dialects that model the execution model and intrinsic operations of specific hardware accelerators. They are the final MLIR abstraction before generating machine code.

• GPU Dialect: Models GPU execution hierarchy (thread, warp, block, grid) and includes ops for memory allocation, synchronization, and kernel launching.
• Vendor NPU Dialects: Custom dialects (e.g., for Qualcomm, Google, or Apple NPUs) that expose hardware-specific operations like custom tensor instructions, memory layouts, and synchronization primitives.
• Function: They allow compiler optimizations (like memory coalescing for GPUs) to be expressed at the IR level and provide a clean target for the final lowering to machine code via the LLVM backend or a vendor-specific code generator.