FForge
← Labsmlir-basics-02MLIR anatomy●○○○

Types, and the tower of dialects

mlir-basics/02_types_and_dialects.mlirneeds: mlir-opt5 steps
0/5 steps
what this lab teaches

The four types you will meet in EVERY ML compiler, and why each exists: tensor<4x8xf32> an immutable VALUE. No address, no aliasing. Great for transformations (fusion, tiling decisions) because rewriting values is safe. Triton-shared, torch-mlir and IREE all start here. memref<4x8xf32> a BUFFER: pointer + sizes + strides + offset. Mutable, can alias. Needed before you can emit loads/stores. Tensor -> memref is called BUFFERIZATION. vector<4xf32> a fixed-size SIMD register value (NEON q-reg = 4 x f32). vector<[4]xf32> a SCALABLE vector: 4 x vscale lanes. SVE/SME. The hardware decides vscale at runtime (1 => 128-bit, 4 => 512-bit). index the machine's pointer-sized integer (i64 here). And the "tower" — the rough order dialects appear in a CPU pipeline: torch / tt (Triton) / stablehlo frontend dialects (what the user wrote) | linalg + tensor structured ops on values (what the math IS) | tiling, fusion, packing <- the performance decisions live here linalg + memref (bufferized) structured ops on buffers | vector + scf SIMD ops + loops (what the core DOES) | arith / math / memref / cf scalar-ish, CFG-ish | llvm 1:1 with LLVM IR -> mlir-translate -> llc Every dialect below exists at least once in IREE and triton-shared's LLVMCPU backend. The ARM work is mostly "make the linalg -> vector step pick good shapes for NEON/SVE/SME".

evidence to leave with
A table of which dialect each op in your file belongs to, and why tensor and memref cannot be mixed
walk it, one step at a time
run this
$ mlir-opt 02_types_and_dialects.mlir --canonicalize | sed -n '/fold_me/,/}/p'
what it printed on the machine that wrote this lab
  func.func @fold_me() -> i32 {
    %c42_i32 = arith.constant 42 : i32
    return %c42_i32 : i32
  }
`canonicalize` is the pass you'll run 50 times per pipeline (count them in
a production backend’s Linalg→LLVM pass). It applies every op's local folds
and simplification patterns until nothing changes.
run this
$ mlir-opt 02_types_and_dialects.mlir --one-shot-bufferize="bufferize-function-boundaries" | sed -n '/tensor_world/,/^  }/p'
what it printed on the machine that wrote this lab
  func.func @tensor_world(%arg0: memref<4xf32, strided<[?], offset: ?>>, %arg1: f32) -> memref<4xf32, strided<[?], offset: ?>> {
    %c1 = arith.constant 1 : index
    memref.store %arg1, %arg0[%c1] : memref<4xf32, strided<[?], offset: ?>>
    return %arg0 : memref<4xf32, strided<[?], offset: ?>>
  }
Read this carefully. `tensor.insert` "returned a new tensor", yet the
bufferizer wrote IN PLACE into the caller's buffer. It was allowed to
because nothing reads the old %t afterwards. Proving that is the whole
job of one-shot-bufferize (it's an alias analysis). When it CAN'T prove
it, it inserts memref.alloc + memref.copy — and an unexpected copy in the
hot loop is one of the most common perf bugs you'll look for.
Also note `strided<[?], offset: ?>`: "I don't know your layout". The
`function-boundary-type-conversion=identity-layout-map` option in
tools/env.sh's LOWER_STD removes that by promising contiguous buffers.
run this
$ mlir-opt 02_types_and_dialects.mlir --expand-strided-metadata | sed -n '/@views/,/^  }/p'
what it printed on the machine that wrote this lab
  %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %arg0 ...
  %reinterpret_cast = memref.reinterpret_cast %base_buffer to offset: [17], sizes: [4, 4], strides: [8, 2] ...
The subview [2,1] with step [1,2] became: offset = 2*8 + 1*1 = 17,
strides = [8*1, 1*2] = [8, 2]. This `reinterpret_cast` is EXACTLY what
triton-shared emits for a Triton pointer `ptr + offs` (lab 07_triton/02).
run this
$ mlir-opt 02_types_and_dialects.mlir --finalize-memref-to-llvm --convert-arith-to-llvm --convert-func-to-llvm --reconcile-unrealized-casts | sed -n '/@memref_world/,/^  }/p'
what it printed on the machine that wrote this lab
  llvm.func @memref_world(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: i64, %arg3: i64, %arg4: i64, %arg5: f32) {
    %0 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
    ...
    %8 = llvm.getelementptr inbounds|nuw %7[%6] : (!llvm.ptr, i64) -> !llvm.ptr, f32
    llvm.store %arg5, %8 : f32, !llvm.ptr
ONE memref<4xf32> argument became FIVE machine arguments:
  (allocated ptr, aligned ptr, offset, size[0], stride[0])

A rank-N memref is 3 + 2N words. Remember this when you call MLIR code from C (lab 02_lowering/04) or read a Triton launcher that packs args (a vendor MLIR backend: the generated Triton launcher).

run this
$ mlir-opt 02_types_and_dialects.mlir --convert-vector-to-llvm --convert-arith-to-llvm --convert-func-to-llvm --reconcile-unrealized-casts | sed -n '/@vectors/,/^  }/p'
what it printed on the machine that wrote this lab
  llvm.func @vectors(%arg0: vector<4xf32>, %arg1: vector<[4]xf32>) -> !llvm.struct<(vector<4xf32>, vector<[4]xf32>)> {
    %0 = llvm.fadd %arg0, %arg0 : vector<4xf32>
    %1 = llvm.fadd %arg1, %arg1 : vector<[4]xf32>
`vector<[4]xf32>` becomes LLVM's `<vscale x 4 x float>`. On AArch64 with
+sve, llc puts it in a z-register. Lab 05_aarch64/02 proves the same
binary is correct at 128, 256, 512 ... 2048-bit vectors.
take it further

1. In @tensor_world, add a use of the OLD %t after the insert (e.g. return both). Re-run STEP 2: where does the copy appear? 2. Compute by hand the strided layout of `memref.subview %m[0, 0] [8, 4] [1, 2]` on memref<8x8xf32>, then check with mlir-opt. 3. How many machine arguments does a memref<?x?x?xf32> become? ============================================================ VALUE semantics: tensors are immutable SSA values. "insert" returns a NEW tensor.

read the source · mlir-basics/02_types_and_dialects.mlir
func.func @tensor_world(%t: tensor<4xf32>, %v: f32) -> tensor<4xf32> { %c1 = arith.constant 1 : index %t2 = tensor.insert %v into %t[%c1] : tensor<4xf32> return %t2 : tensor<4xf32> } // BUFFER semantics: memrefs are pointers + shape + strides. "store" mutates memory. func.func @memref_world(%m: memref<4xf32>, %v: f32) { %c1 = arith.constant 1 : index memref.store %v, %m[%c1] : memref<4xf32> return } // A memref is a VIEW: a subview of every other column has a non-trivial layout. func.func @views(%m: memref<8x8xf32>) -> memref<4x4xf32, strided<[8, 2], offset: 17>> { %s = memref.subview %m[2, 1] [4, 4] [1, 2] : memref<8x8xf32> to memref<4x4xf32, strided<[8, 2], offset: 17>> return %s : memref<4x4xf32, strided<[8, 2], offset: 17>> } // VECTOR types: what ends up in SIMD registers. [4] is fixed (NEON-sized for f32), // [[4]] is SCALABLE: 4 x vscale lanes, the SVE model (lab 05_aarch64). func.func @vectors(%a: vector<4xf32>, %b: vector<[4]xf32>) -> (vector<4xf32>, vector<[4]xf32>) { %a2 = arith.addf %a, %a : vector<4xf32> %b2 = arith.addf %b, %b : vector<[4]xf32> return %a2, %b2 : vector<4xf32>, vector<[4]xf32> } // Constant folding is free once you have values instead of memory: func.func @fold_me() -> i32 { %a = arith.constant 6 : i32 %b = arith.constant 7 : i32 %c = arith.muli %a, %b : i32 return %c : i32 }
◇
Finish the lab: write the evidence note

A lab is done when you can say what it showed on your machine — not when you have read it.