the path · 0/15
start the path

the kernel path · XLA · lesson 07 of 7

Levels as a first-class idea

One elementwise add, written three ways in one framework: whole tensors, then loops over buffers, then register-width vectors. What is computed never changes. What has been decided changes at every step.

the goal Read one computation in three MLIR dialects, say what each rung wrote down that the rung above left open, and define legalization precisely enough to predict which ops a pass expands and which it leaves alone.

mastery work · this chapter0/3
  2. go →
manual items are your word; auto items complete from your streaks, labs, and can-you ticks · stored in your browser only
§ 01

The top of the ladder

The vocabulary this lesson runs on was settled one chapter back. The StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → chapter at /l/stablehlo pins what MLIR is and what a dialect is; take both as given and put them to work on a program small enough that nothing hides.

Two arrays of shape (8, 1024), added. At the top of the ladder that is one operation over one value, and the type carries the whole shape.

Read it once more for what it does not say. No loop, no order, no memory, no tile, no vector width. A backend is free to choose all of that, and every rung below is a record of it choosing.

jax 0.4.38 on CPU: jax.jit(lambda a, b: a + b).lower(x, x).as_text(), x of shape (8, 1024) float32
module @jit__lambda_ attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<8x1024xf32>, %arg1: tensor<8x1024xf32>) -> (tensor<8x1024xf32> {jax.result_info = ""}) {
    %0 = stablehlo.add %arg0, %arg1 : tensor<8x1024xf32>
    return %0 : tensor<8x1024xf32>
  }
}
§ 02

The middle, where the loops get written down

XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s CPU tiled pipeline has a pass whose entire job is one rung of this ladder, and its test file shows both ends of the step at once. The input is linalg.elementwise kind=add over three memref<8x1024xf32> operands. Values have become buffers, which is what the ins and outs lists are about: an op that used to produce a result now writes into one you hand it.

The expected output is written as FileCheck patterns, so %[[IV0]] is a capture name rather than an SSA value. Read past that syntax and the assertion is precise. Two scf.for loops with bounds 8 and 1024, the inner one stepping by 8, a vector.transfer_read per operand pulling vector<8xf32> out of the memref, one arith.addf, one vector.transfer_write.

Count what arrived at this rung. Loop order, loop bounds, a tile of eight elements along the minor dimension, the register width, and the exact points where memory gets touched. The add itself is untouched. Every new fact is a decision about how, and the pass wrote it into the IR rather than keeping it in a data structure off to the side.

openxla/xla at commit a6c8e17, xla/backends/cpu/codegen/tiled/transforms/tests/linalg_elementwise_to_vector_pass.mlir: input op and expected output, verbatim
func.func @elementwise_add_to_vector(
    %arg0 : memref<8x1024xf32>,
    %arg1 : memref<8x1024xf32>,
    %arg2 : memref<8x1024xf32>) {
  // CHECK-DAG: %[[MASK:.*]] = ub.poison : f32
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
  // CHECK-DAG: %[[C1024:.*]] = arith.constant 1024 : index
  // CHECK: scf.for %[[IV0:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
  // CHECK:   scf.for %[[IV1:.*]] = %[[C0]] to %[[C1024]] step %[[C8]] {
  // CHECK:     %[[LHS:.*]] = vector.transfer_read %arg0[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x1024xf32>, vector<8xf32>
  // CHECK:     %[[RHS:.*]] = vector.transfer_read %arg1[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x1024xf32>, vector<8xf32>
  // CHECK:     %[[OUT:.*]] = arith.addf %[[LHS]], %[[RHS]] : vector<8xf32>
  // CHECK:     vector.transfer_write %[[OUT]], %arg2[%[[IV0]], %[[IV1]]]
  // CHECK-SAME:  {in_bounds = [true]} : vector<8xf32>, memref<8x1024xf32>
  // CHECK:   }
  // CHECK: }
  linalg.elementwise kind=#linalg.elementwise_kind<add>
    ins(%arg0, %arg1 : memref<8x1024xf32>, memref<8x1024xf32>)
    outs(%arg2 : memref<8x1024xf32>)
  return
}
all four cases in the test, including the ragged 8x100 tail and a two-op linalg.generic fused into one loop body · 97 lines 
func.func @elementwise_add_to_vector(
    %arg0 : memref<8x1024xf32>,
    %arg1 : memref<8x1024xf32>,
    %arg2 : memref<8x1024xf32>) {
  // CHECK-DAG: %[[MASK:.*]] = ub.poison : f32
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
  // CHECK-DAG: %[[C1024:.*]] = arith.constant 1024 : index
  // CHECK: scf.for %[[IV0:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
  // CHECK:   scf.for %[[IV1:.*]] = %[[C0]] to %[[C1024]] step %[[C8]] {
  // CHECK:     %[[LHS:.*]] = vector.transfer_read %arg0[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x1024xf32>, vector<8xf32>
  // CHECK:     %[[RHS:.*]] = vector.transfer_read %arg1[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x1024xf32>, vector<8xf32>
  // CHECK:     %[[OUT:.*]] = arith.addf %[[LHS]], %[[RHS]] : vector<8xf32>
  // CHECK:     vector.transfer_write %[[OUT]], %arg2[%[[IV0]], %[[IV1]]]
  // CHECK-SAME:  {in_bounds = [true]} : vector<8xf32>, memref<8x1024xf32>
  // CHECK:   }
  // CHECK: }
  linalg.elementwise kind=#linalg.elementwise_kind<add>
    ins(%arg0, %arg1 : memref<8x1024xf32>, memref<8x1024xf32>)
    outs(%arg2 : memref<8x1024xf32>)
  return
}

//------

func.func @elementwise_add_to_vector_non_multiple_of_8(
    %arg0 : memref<8x100xf32>,
    %arg1 : memref<8x100xf32>,
    %arg2 : memref<8x100xf32>) {
  // CHECK-DAG: %[[MASK:.*]] = ub.poison : f32
  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
  // CHECK-DAG: %[[C96:.*]] = arith.constant 96 : index
  // CHECK: scf.for %[[IV0:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
  // CHECK:   scf.for %[[IV1:.*]] = %[[C0]] to %[[C96]] step %[[C8]] {
  // CHECK:     %[[LHS:.*]] = vector.transfer_read %arg0[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x100xf32>, vector<8xf32>
  // CHECK:     %[[RHS:.*]] = vector.transfer_read %arg1[%[[IV0]], %[[IV1]]],
  // CHECK-SAME:  %[[MASK]] {in_bounds = [true]} : memref<8x100xf32>, vector<8xf32>
  // CHECK:     %[[OUT:.*]] = arith.addf %[[LHS]], %[[RHS]] : vector<8xf32>
  // CHECK:     vector.transfer_write %[[OUT]], %arg2[%[[IV0]], %[[IV1]]]
  // CHECK-SAME:  {in_bounds = [true]} : vector<8xf32>, memref<8x100xf32>
  // CHECK:   }
  // CHECK: %[[UNROLL_LHS:.*]] = vector.transfer_read %arg0[%[[IV0]], %[[C96]]], %[[MASK]]
  // CHECK-SAME: {in_bounds = [true]} : memref<8x100xf32>, vector<4xf32>
  // CHECK: %[[UNROLL_RHS:.*]] = vector.transfer_read %arg1[%[[IV0]], %[[C96]]], %[[MASK]]
  // CHECK-SAME: {in_bounds = [true]} : memref<8x100xf32>, vector<4xf32>
  // CHECK: %[[UNROLL_OUT:.*]] = arith.addf %[[UNROLL_LHS]], %[[UNROLL_RHS]] : vector<4xf32>
  // CHECK: vector.transfer_write %[[UNROLL_OUT]], %arg2[%[[IV0]], %[[C96]]]
  // CHECK-SAME: {in_bounds = [true]} : vector<4xf32>, memref<8x100xf32>
  // CHECK: }
  linalg.elementwise kind=#linalg.elementwise_kind<add>
    ins(%arg0, %arg1 : memref<8x100xf32>, memref<8x100xf32>)
    outs(%arg2 : memref<8x100xf32>)
  return
}

//------

#map = affine_map<(d0, d1) -> (d0, d1)>
func.func @fused(%arg0: memref<8x1024xf32>,
                 %arg1: memref<8x1024xf32>,
                 %arg2: memref<8x1024xf32>) {
  // CHECK: scf.for
  // CHECK: scf.for
  // CHECK-NEXT: vector.transfer_read
  // CHECK-NEXT: vector.transfer_read
  // CHECK-NEXT: arith.mulf
  // CHECK-NEXT: arith.addf
  // CHECK-NEXT: vector.transfer_write
  linalg.generic
    {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
    ins(%arg0, %arg1 : memref<8x1024xf32>, memref<8x1024xf32>)
    outs(%arg2 : memref<8x1024xf32>) {
  ^bb0(%lhs: f32, %rhs: f32, %out: f32):
    %mul = arith.mulf %lhs, %rhs : f32
    %res = arith.addf %mul, %rhs : f32
    linalg.yield %res : f32
  }
  return
}

// -----

func.func @elementwise_add_to_vector_small_minor(
    %arg0 : memref<8x3xf32>,
    %arg1 : memref<8x3xf32>,
    %arg2 : memref<8x3xf32>) {
  linalg.elementwise kind=#linalg.elementwise_kind<add>
    ins(%arg0, %arg1 : memref<8x3xf32>, memref<8x3xf32>)
    outs(%arg2 : memref<8x3xf32>)
  return
}
§ 03

Several dialects, one module

Six op prefixes appear in that one file: func, linalg, scf, vector, arith, and ub. Nothing was translated into another language or written to another file. Ops from one dialect were replaced in place by ops from others, in the same module, under the same verifier, and a module halfway through a pipeline legitimately holds several dialects at once.

That is what it means for levels to be first-class. Each of those dialects is defined in the same framework as every other, so a pass can rewrite between two of them the way an ordinary pass rewrites within one. The alternative, which the previous lesson walked through, is a single fixed level that every frontend must reach in one jump.

The consequence for tensor programs is direct. A fact survives as long as some dialect on the way down can still express it, so shapes, whole-value semantics, and loop structure can each be dropped at the rung that no longer needs them, instead of all at once at the door.

dialectwhat it addswhere it shows up here
stablehlowhole-tensor ops with the shape in the typethe lowered jax dump above, and the type system lesson at /l/stablehlo
linalgstructured ops over buffers or tensors: what to compute, not yet in what orderthe input side of the pass test
scfstructured control flow, for and if and while, as regions rather than labelsthe output side of the same test
vectorregister-width values, and the transfers that move them in and out of memoryvector<8xf32> and the two transfer ops
ariththe arithmetic itself, on scalars and on vectors alikearith.addf here, and the tanh expansion below
tputhe ops that name one chip, which is Mosaicthe module a pallas_call prints, taught at /l/mosaic
the dialects this site quotes, and what each one contributes
§ 04

Legalization, and what a target refuses to take

A pass that lowers between dialects needs a definition of acceptable output. MLIR calls that the conversion target, and every op is marked against it as legal, illegal, or dynamic, that last one meaning only some instances qualify. The framework documentation gives its own example, arith.addi legal only on 32-bit integers, and rewriting an illegal op into legal ones is what legalization means.

A test in the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → tree makes the idea concrete with one operation and three types. At f64, math.tanh is left exactly as written, because the target accepts it. At f32 it becomes 46 operations, 16 of them constants of a rational approximation. At f16 the pass widens with arith.extf, runs the f32 expansion, and narrows back with arith.truncf.

So legality belongs to an op together with its types and a target, never to the op alone. When a dump answers one line of your source with a wall of arithmetic, a legalization pattern like this one is the usual reason, and the identity it applied is written into the pass, not derived on the spot.

openxla/xla at commit a6c8e17, xla/mlir_hlo/tests/legalize-trigonometric-to-approximation.mlir: the RUN line names the pass, then f64 untouched and f32 expanded
// RUN: mlir-hlo-opt --mhlo-legalize-trigonometric-to-approximation --split-input-file %s | FileCheck %s

// CHECK-LABEL: @tanh_f64
func.func @tanh_f64(%arg0 : f64) -> f64 {
  // CHECK: tanh
  %res = math.tanh %arg0 : f64
  func.return %res : f64
}

// -----

// CHECK-LABEL: @tanh_f32
// CHECK-SAME: (%[[ARG:.*]]: f32) -> f32
func.func @tanh_f32(%arg0 : f32) -> f32 {
  // CHECK-DAG: %[[C:.*]] = arith.constant -2.76076837E-16 : f32
  // CHECK-DAG: %[[C0:.*]] = arith.constant 2.00018794E-13 : f32
  // CHECK-DAG: %[[C1:.*]] = arith.constant -8.60467184E-11 : f32
  // CHECK-DAG: %[[C2:.*]] = arith.constant 5.12229725E-8 : f32
  // CHECK-DAG: %[[C3:.*]] = arith.constant 1.48572235E-5 : f32
all three tanh cases, f64 through f16, with the whole f32 approximation · 77 lines 
// RUN: mlir-hlo-opt --mhlo-legalize-trigonometric-to-approximation --split-input-file %s | FileCheck %s

// CHECK-LABEL: @tanh_f64
func.func @tanh_f64(%arg0 : f64) -> f64 {
  // CHECK: tanh
  %res = math.tanh %arg0 : f64
  func.return %res : f64
}

// -----

// CHECK-LABEL: @tanh_f32
// CHECK-SAME: (%[[ARG:.*]]: f32) -> f32
func.func @tanh_f32(%arg0 : f32) -> f32 {
  // CHECK-DAG: %[[C:.*]] = arith.constant -2.76076837E-16 : f32
  // CHECK-DAG: %[[C0:.*]] = arith.constant 2.00018794E-13 : f32
  // CHECK-DAG: %[[C1:.*]] = arith.constant -8.60467184E-11 : f32
  // CHECK-DAG: %[[C2:.*]] = arith.constant 5.12229725E-8 : f32
  // CHECK-DAG: %[[C3:.*]] = arith.constant 1.48572235E-5 : f32
  // CHECK-DAG: %[[C4:.*]] = arith.constant 6.37261954E-4 : f32
  // CHECK-DAG: %[[C5:.*]] = arith.constant 0.00489352457 : f32
  // CHECK-DAG: %[[C6:.*]] = arith.constant 1.19825836E-6 : f32
  // CHECK-DAG: %[[C7:.*]] = arith.constant 1.18534706E-4 : f32
  // CHECK-DAG: %[[C8:.*]] = arith.constant 0.00226843474 : f32
  // CHECK-DAG: %[[C9:.*]] = arith.constant 0.00489352504 : f32
  // CHECK-DAG: %[[C10:.*]] = arith.constant 4.000000e-04 : f32
  // CHECK-DAG: %[[C11:.*]] = arith.constant 7.90531111 : f32
  // CHECK-DAG: %[[C12:.*]] = arith.constant -7.90531111 : f32
  // CHECK-DAG: %[[C13:.*]] = arith.constant 1.000000e+00 : f32
  // CHECK-DAG: %[[C14:.*]] = arith.constant -1.000000e+00 : f32
  // CHECK-DAG: %[[TMP0:.*]] = arith.mulf %[[ARG]], %[[ARG]] : f32
  // CHECK-DAG: %[[TMP1:.*]] = arith.mulf %[[TMP0]], %[[C]] : f32
  // CHECK-DAG: %[[TMP2:.*]] = arith.addf %[[TMP1]], %[[C0]] : f32
  // CHECK-DAG: %[[TMP3:.*]] = arith.mulf %[[TMP0]], %[[TMP2]] : f32
  // CHECK-DAG: %[[TMP4:.*]] = arith.addf %[[TMP3]], %[[C1]] : f32
  // CHECK-DAG: %[[TMP5:.*]] = arith.mulf %[[TMP0]], %[[TMP4]] : f32
  // CHECK-DAG: %[[TMP6:.*]] = arith.addf %[[TMP5]], %[[C2]] : f32
  // CHECK-DAG: %[[TMP7:.*]] = arith.mulf %[[TMP0]], %[[TMP6]] : f32
  // CHECK-DAG: %[[TMP8:.*]] = arith.addf %[[TMP7]], %[[C3]] : f32
  // CHECK-DAG: %[[TMP9:.*]] = arith.mulf %[[TMP0]], %[[TMP8]] : f32
  // CHECK-DAG: %[[TMP10:.*]] = arith.addf %[[TMP9]], %[[C4]] : f32
  // CHECK-DAG: %[[TMP11:.*]] = arith.mulf %[[TMP0]], %[[TMP10]] : f32
  // CHECK-DAG: %[[TMP12:.*]] = arith.addf %[[TMP11]], %[[C5]] : f32
  // CHECK-DAG: %[[TMP13:.*]] = arith.mulf %[[ARG]], %[[TMP12]] : f32
  // CHECK-DAG: %[[TMP14:.*]] = arith.mulf %[[TMP0]], %[[C6]] : f32
  // CHECK-DAG: %[[TMP15:.*]] = arith.addf %[[TMP14]], %[[C7]] : f32
  // CHECK-DAG: %[[TMP16:.*]] = arith.mulf %[[TMP0]], %[[TMP15]] : f32
  // CHECK-DAG: %[[TMP17:.*]] = arith.addf %[[TMP16]], %[[C8]] : f32
  // CHECK-DAG: %[[TMP18:.*]] = arith.mulf %[[TMP0]], %[[TMP17]] : f32
  // CHECK-DAG: %[[TMP19:.*]] = arith.addf %[[TMP18]], %[[C9]] : f32
  // CHECK-DAG: %[[TMP20:.*]] = arith.divf %[[TMP13]], %[[TMP19]] : f32
  // CHECK-DAG: %[[TMP21:.*]] = math.absf %[[ARG]] : f32
  // CHECK-DAG: %[[TMP22:.*]] = arith.cmpf olt, %[[TMP21]], %[[C10]] : f32
  // CHECK-DAG: %[[TMP23:.*]] = arith.select %[[TMP22]], %[[ARG]], %[[TMP20]] : f32
  // CHECK-DAG: %[[TMP24:.*]] = arith.cmpf ugt, %[[ARG]], %[[C11]] : f32
  // CHECK-DAG: %[[TMP25:.*]] = arith.cmpf ult, %[[ARG]], %[[C12]] : f32
  // CHECK-DAG: %[[IS_NAN:.*]] = arith.cmpf une, %[[ARG]], %[[ARG]] : f32
  // CHECK-DAG: %[[TMP26:.*]] = arith.select %[[TMP24]], %[[C13]], %[[TMP23]] : f32
  // CHECK-DAG: %[[TMP27:.*]] = arith.select %[[TMP25]], %[[C14]], %[[TMP26]] : f32
  // CHECK-DAG: %[[RESULT:.*]] = arith.select %[[IS_NAN]], %[[ARG]], %[[TMP27]] : f32
  // CHECK: return %[[RESULT]] : f32
  %res = math.tanh %arg0 : f32
  func.return %res : f32
}

// -----

func.func @tanh_f16(%arg0 : f16) -> f16 {
  // CHECK-LABEL: func @tanh_f16
  // CHECK-SAME: (%[[ARG:.*]]: f16) -> f16
  // CHECK: %{{.*}} = arith.extf %[[ARG]] : f16 to f32
  // CHECK: %[[RES:.*]] = arith.truncf %{{.*}} : f32 to f16
  // CHECK: return %[[RES]] : f16
  %res = math.tanh %arg0 : f16
  func.return %res : f16
}
§ 05

What the ladder never does

Walk the rungs back up and each one added a decision. Buffers, then loop order, then a tile, then a vector width, then instructions. Not one of them changed what gets computed, and that is the contract every lowering pass signs.

Which is exactly why the ceiling this unit states holds all the way down. Two-pass softmax stays two-pass in every dialect on the descent, because no legalization pattern says that a reduce and a divide over the same values are a streaming recurrence. The chapter states the reason, and the fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → taxonomy lesson shows the boundary in a real dump; the ladder is not a second chance at it.

Lowering rewrites how. Nothing in the pipeline is allowed to rewrite what.

The bottom rung is the one-level world of the previous lesson, where the dialect distinction disappears along with everything it was carrying. What MLIR keeps explicit as ops and types, LLVM keeps as address arithmetic and metadata, and after that there is only the machine.

before you move on

Check yourself

01 Which facts does the vector form of that add state that the linalg form left open?

The loop order and both bounds, the tile of eight along the minor dimension, the register width as vector<8xf32>, and the exact instructions where memory is read and written. What is computed is identical in the two forms; only the how got written down.

02 A pass leaves math.tanh alone at f64 and expands it into 46 operations at f32. Which legality action is that, and what does the f16 case add to the picture?

Dynamic legality: the op is legal for some types and illegal for others, so the pass expands only the illegal instances. The f16 case shows the second move available to a legalization pattern, which is converting a type it cannot expand into one it can, with arith.extf in and arith.truncf out around the f32 path.

03 Progressive lowering added buffers, loops, tiles, and vector widths to one add. Why does the same machinery never turn two-pass softmax into a one-pass kernel?

Because every rung is semantics-preserving by contract. Lowering commits to how a computation runs and is forbidden from changing what it computes, while the streaming rewrite is an algebraic identity about exponentials that no pass in the pipeline holds. The chapter at /l/xla states the same limit from the fusion side.

assigned

Readings