Fusion raises

A typical compiler rewrites the program until it can run. This one does not. The tape stays the spec. A faster form is an offer a run may take. It never edits the graph.

That sounds like a refusal to go fast. It is why another compiler still sees a convolution.

The decision#

Write convolution as ordinary ops: pad, unfold, permute, reshape, matmul. A fused kernel has to earn its place — wrong float math, or real cost. Accuracy lives in the instruction set. Speed lives in a catalog of matches. See what earns an instruction.

The matcher looks at the graph, not at who wrote it. A hand-rolled chain and a library call that recorded the same ops look the same. Matches are gathered once. Each consumer picks the ones it can use. A region nobody picks just runs the original ops. A value you asked to observe, inside a match, blocks the match — unless the match names it as a result and writes it back, which is how observing batch-norm's statistics survives the fusion.

What it opened#

module @topos {
  func.func @main(%arg0: tensor<2x1x3x3xf32>, %arg1: tensor<2xf32>, %arg2: tensor<1x1x6x6xf32>) -> (tensor<1x2x2x2xf32>) {
    %v7 = stablehlo.transpose %arg0, dims = [1, 2, 3, 0] : (tensor<2x1x3x3xf32>) -> tensor<1x3x3x2xf32>
    %v8 = stablehlo.reshape %v7 : (tensor<1x3x3x2xf32>) -> tensor<9x2xf32>
    %v9_kernel = stablehlo.reshape %v8 : (tensor<9x2xf32>) -> tensor<1x3x3x2xf32>
    %v9_windows = stablehlo.convolution(%arg2, %v9_kernel) dim_numbers = [b, f, 0, 1]x[i, 0, 1, o]->[b, 0, 1, f], window = {stride = [1, 1], pad = [[0, 0], [0, 0]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x1x6x6xf32>, tensor<1x3x3x2xf32>) -> tensor<1x4x4x2xf32>
    %v9 = stablehlo.reshape %v9_windows : (tensor<1x4x4x2xf32>) -> tensor<16x2xf32>
    %v10 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<2xf32>) -> tensor<16x2xf32>
    %v11 = stablehlo.add %v9, %v10 : tensor<16x2xf32>
    %v12 = stablehlo.reshape %v11 : (tensor<16x2xf32>) -> tensor<1x4x4x2xf32>
    %v13 = stablehlo.transpose %v12, dims = [0, 3, 1, 2] : (tensor<1x4x4x2xf32>) -> tensor<1x2x4x4xf32>
    %v14 = stablehlo.constant dense<0.0> : tensor<1x2x4x4xf32>
    %v15 = stablehlo.maximum %v13, %v14 : tensor<1x2x4x4xf32>
    %v27_seed = stablehlo.constant dense<0xFF800000> : tensor<f32>
    %v27 = "stablehlo.reduce_window"(%v15, %v27_seed) ({
    ^bb0(%v27_left: tensor<f32>, %v27_right: tensor<f32>):
      %v27_larger = stablehlo.maximum %v27_left, %v27_right : tensor<f32>
      stablehlo.return %v27_larger : tensor<f32>
    }) {window_dimensions = array<i64: 1, 1, 2, 2>, window_strides = array<i64: 1, 1, 2, 2>} : (tensor<1x2x4x4xf32>, tensor<f32>) -> tensor<1x2x2x2xf32>
    return %v27 : tensor<1x2x2x2xf32>
  }
}

The same match has two jobs. At home it calls a fused kernel. When we print the plan for another compiler, it becomes that compiler's named op.

The im2col-and-multiply chain becomes stablehlo.convolution. A max-pool window becomes reduce_window. The batch-norm formulas become batch_norm_training and batch_norm_inference.

It looks backwards: the library knew "convolution," recorded the pieces, and the matcher finds convolution again. That is how a small instruction set still hands an industrial backend the op it is good at. The other path is a new instruction for every fused form, each with its own backward rule.

The tape is still the spec. A raise, like a fusion, is an offer. Observing an unnamed piece of the chain refuses both; a named result is kept and written back instead.

Spelled today#

Plan::patterns is the matches this plan took. PatternKind names the ones we know: window-times-multiply, window-reduce, the two batch-norms. Home fusion calls a kernel. Plan::emit_stablehlo writes the named op. Same match. describe prints the choice.

This section may rot; the rest must not.

Not this#

• Rewriting the tape into fused opcodes. Speed is not a new instruction.
• A raise that hides a value you asked to observe.
• A Convolution instruction because a layer has that name. Convolution is a formula.

See Vision for fusion as an offer, What earns an instruction for why the fused form is not an opcode, and Several exports, one spec for the read-list the matcher runs on.