diff --git a/examples/xegpu/fused_attention.py b/examples/xegpu/fused_attention.py index fab111e0..70ad7941 100644 --- a/examples/xegpu/fused_attention.py +++ b/examples/xegpu/fused_attention.py @@ -259,6 +259,13 @@ def parse_cli(): default=64, help="Tile size for the inner reduction dimension (K/V sequence length)", ) + parser.add_argument( + "--reference-flash", + action="store_true", + help="Emit the flash loop from the hand-written generator " + "(transform_ext.replace_with_fused_attention) instead of deriving it from " + "the payload chain. Kept as a reference point for comparing the two.", + ) parser.add_argument( "--q-load-tile", type=int, @@ -349,6 +356,7 @@ def parse_cli(): "sg_rows": args.sg_rows, "subgroup_size": args.subgroup_size, "reduction_tile": args.reduction_tile, + "reference_flash": args.reference_flash, "q_load_tile": args.q_load_tile, "v_load_tile": args.v_load_tile, "prefetch_tile": args.prefetch_tile, diff --git a/examples/xegpu/nanoGPT_payload.py b/examples/xegpu/nanoGPT_payload.py index 5761bc2f..0615ea89 100644 --- a/examples/xegpu/nanoGPT_payload.py +++ b/examples/xegpu/nanoGPT_payload.py @@ -21,12 +21,13 @@ """ from mlir import ir -from mlir.dialects import linalg, bufferization, tensor, arith, gpu, memref +from mlir.dialects import linalg, bufferization, tensor, arith, gpu, math, memref from lighthouse.ingress.mlir_gen.utils import ( emit_buf_to_tensor, affine_map, parallel, + reduction, ) from lighthouse.ingress.mlir_gen.gpu_utils import emit_gpu_util_funcs from lighthouse.ingress.mlir_gen.gpu_layer_norm_payload import emit_layer_norm_generics @@ -229,11 +230,29 @@ def heads_view(self, buf2d, n_ctx, n_head, d_head): def attention_4d( self, Qh, Kh, Vh, n_head, n_ctx, d_head, out_view, out_view_memref ): - # Emits the SAME linalg op sequence as generate_gpu_attention_payload - # (batch_matmul QK^T -> scale-mul -> softmax -> batch_matmul @V), so the - # fused-attention schedule's matchers/rewrite apply verbatim. After the - # per-region fused tiling, all these ops fuse into one scf.forall -> one - # GPU kernel (the flash/online-softmax kernel). Counts as one 'fa'. + # batch_matmul QK^T -> scale-mul -> softmax -> batch_matmul @V, with the + # softmax spelled out in the flash-attention form: + # + # m = max_j s l = sum_j p + # p = exp(s - m) o = p @ V out = o / l + # + # i.e. the normalizing divide comes *after* the contraction. That is + # algebraically identical to `softmax(s) @ V` -- dividing by the per-row `l` + # commutes with a contraction that reduces the other axis -- but it leaves + # the `max -> exp -> {sum, @V}` dependency chain explicit, which is what + # `transform_ext.fuse_dependant_reduction_ops` consumes to derive the online + # one-pass loop (see `_fuse_attention_in_region` in nanoGPT_schedule.py). + # A `linalg.softmax` would not do: its decomposition normalizes *before* the + # contraction, leaving @V reading the normalized P and breaking the chain. + # + # Same chain shape as `generate_gpu_attention_payload`, but not the same ops: + # this one stays f16 throughout and keeps @V a named `linalg.batch_matmul`, + # while that one accumulates in f32 and writes @V as a `linalg.generic` so the + # narrowing of P can sit in its body. Hence the extra `generalize` on the + # `_fuse_attention_in_region` path -- the fusion needs @V as a generic. + # + # After the per-region fused tiling, all these ops fuse into one scf.forall + # -> one GPU kernel (the flash/online-softmax kernel). Counts as one 'fa'. # Inputs Qh/Kh/Vh are (n_head,n_ctx,d_head) f16 strided views (heads_view); the @V result # is materialized into `out_view`, a (n_head,n_ctx,d_head) strided view of a (n_ctx,n_embd) buffer, # so the merge back to 2D is also a free view (no from_heads kernel). @@ -252,15 +271,62 @@ def attention_4d( scaled = linalg.mul( qkt, scale_t, outs=[tensor.empty((n_head, n_ctx, n_ctx), f16)] ) - aw = linalg.softmax( - result=[ir.RankedTensorType.get((n_head, n_ctx, n_ctx), f16)], - input=scaled, - output=tensor.empty((n_head, n_ctx, n_ctx), f16), - dimension=2, + + # (head, row, col) -> (head, row, col) and -> (head, row): the per-row + # statistics are broadcast over the reduced axis. + dims = [ir.AffineDimExpr.get(i) for i in range(3)] + ew_map = affine_map(3, dims) + row_map = affine_map(3, dims[:2]) + row_shape = (n_head, n_ctx) + + # m = max_j s + neg_inf = arith.constant(f16, float("-inf")) + max_acc = linalg.fill(neg_inf, outs=[tensor.empty(row_shape, f16)]) + + @linalg.generic( + [scaled], [max_acc], [ew_map, row_map], [parallel, parallel, reduction] ) - # @V: (n_head,n_ctx,n_ctx) @ (n_head,n_ctx,d_head) -> (n_head,n_ctx,d_head) f16, materialized into the (n_ctx,n_embd) view. - out_filled = linalg.fill(zero, outs=[out_view]) - out = linalg.batch_matmul(aw, Vh, outs=[out_filled]) + def row_max(s, acc): + return arith.MaximumFOp(s, acc) + + # p = exp(s - m), read by both the row sum and the @V contraction. + @linalg.generic( + [scaled, row_max], + [tensor.empty((n_head, n_ctx, n_ctx), f16)], + [ew_map, row_map, ew_map], + [parallel, parallel, parallel], + ) + def probs(s, m, out): + return math.ExpOp(arith.SubFOp(s, m).result) + + # l = sum_j p + sum_acc = linalg.fill(zero, outs=[tensor.empty(row_shape, f16)]) + + @linalg.generic( + [probs], [sum_acc], [ew_map, row_map], [parallel, parallel, reduction] + ) + def row_sum(p, acc): + return arith.AddFOp(p, acc) + + # @V: (n_head,n_ctx,n_ctx) @ (n_head,n_ctx,d_head) -> (n_head,n_ctx,d_head) + # f16, still unnormalized. + unnorm_init = linalg.fill( + zero, outs=[tensor.empty((n_head, n_ctx, d_head), f16)] + ) + unnormalized = linalg.batch_matmul(probs, Vh, outs=[unnorm_init]) + + # out = o / l, the deferred normalization, materialized into the + # (n_ctx,n_embd) view. `l` broadcasts over d_head, the contraction's free + # axis. + @linalg.generic( + [unnormalized, row_sum], + [out_view], + [ew_map, row_map, ew_map], + [parallel, parallel, parallel], + ) + def out(o, denom, dst): + return arith.DivFOp(o, denom) + bufferization.materialize_in_destination( None, out, out_view_memref, restrict=True, writable=True ) diff --git a/examples/xegpu/nanoGPT_schedule.py b/examples/xegpu/nanoGPT_schedule.py index b52341d8..b5096d1a 100644 --- a/examples/xegpu/nanoGPT_schedule.py +++ b/examples/xegpu/nanoGPT_schedule.py @@ -21,8 +21,8 @@ DPAS tile sizes come from `mm_params`) - layernorm-> `_tile_one_layernorm` (tile rows, fuse the 2 reductions + 2 zero-fills into the loop) - - fused attn-> `_tile_one_fused_attention_region` (tile @V batch_matmul into - a forall, fuse QK^T/scale/softmax/@V in; flash rewrite later) + - fused attn-> `_tile_one_fused_attention_region` (tile the chain's leaf + into a forall and fuse the rest in; flash loop folded after) - elementwise -> a single `structured_tile_using_forall` over rows (b) Shared tail (same for every kernel): vectorize -> bufferize (tensors -> memrefs) -> convert the forall grids to `gpu.launch` -> outline each into @@ -36,7 +36,7 @@ from mlir import ir from mlir.dialects import transform -from mlir.dialects.transform import structured, xegpu +from mlir.dialects.transform import structured, tensor, xegpu from mlir.dialects.transform import bufferization as transform_bufferization from mlir.dialects.transform.vector import ( apply_patterns_vector_cast_away_vector_leading_one_dim, @@ -144,14 +144,14 @@ def _tile_one_layernorm( canonicalize(ln_forall) -def _tile_one_fused_attention_region(anytype, pv_bmm, softmax_op, fa_params): - """Tile + fuse one attention region (QK^T -> scale -> softmax -> @V) into a - SINGLE scf.forall, so it vectorizes/bufferizes into one kernel body that - `replace_with_fused_attention` later rewrites into the flash loop. +def _tile_one_fused_attention_region(anytype, divide_op, fa_params): + """Tile + fuse one attention region (QK^T -> scale -> softmax -> @V -> divide) + into a SINGLE scf.forall, so it vectorizes/bufferizes into one kernel body that + `_fuse_attention_in_region` then folds into the flash loop. - Operates on PRE-SPLIT, per-region - handles (pv_bmm, softmax_op) so it is region-local and works at any - multiplicity. All further producers are pulled in via get_producer_of_operand + Operates on a PRE-SPLIT, per-region handle (the region's deferred normalizing + divide, which is its leaf op) so it is region-local and works at any + multiplicity. All producers are pulled in via get_producer_of_operand (SSA-walk = inherently scoped to this region).""" prod = transform.get_producer_of_operand @@ -161,39 +161,31 @@ def fuse(p, c): )[1] wg_rows = fa_params["wg_rows"] - # 1. Tile the @V batch_matmul in (batch=1, M=wg_rows) -> forall grid. - tiled_pv, forall = structured.structured_tile_using_forall( + # 1. Tile the region's leaf -- the `o / l` divide -- in (batch=1, M=wg_rows). + # It is all-parallel over (head, row, d_head), so three tile sizes. + tiled_div, forall = structured.structured_tile_using_forall( anytype, anytype, - pv_bmm, + divide_op, num_threads=[], tile_sizes=[], - static_tile_sizes=(1, wg_rows, 0, 0), + static_tile_sizes=(1, wg_rows, 0), ) func = transform.get_parent_op( anytype, forall, op_name="func.func", deduplicate=True ) - # 2. Fuse the @V output init fill (producer of forall operand 0). - forall = fuse(prod(anytype, forall, operand_number=0), forall) - transform.apply_cse(func) - canonicalize(func) - # 3. Decompose this region's softmax. linalg.softmax -> 4 generics + 2 fills: - # max = reduce_max(scaled) [+ -inf fill] - # num = exp(scaled - max) - # den = reduce_sum(num) [+ 0 fill] - # div = num / den (feeds @V) - structured.structured_decompose_interface(anytype, softmax_op) transform.apply_cse(func) canonicalize(func) # Grab the whole producer chain UP FRONT via SSA walk (region-local; no count # matching). Fusing op X invalidates only X's handle, so collect all, then fuse - # each once in consumer->producer topological order. - # tiled_pv operand 0 is the aw extract_slice inside the forall; hop through it - # to the func-scope softmax `div` that it slices. - aw_slice = prod(anytype, tiled_pv, operand_number=0) - div = prod(anytype, aw_slice, operand_number=0) # num / den (softmax out) - num = prod(anytype, div, operand_number=0) # exp generic - den = prod(anytype, div, operand_number=1) # sum-reduce generic + # each once in consumer->producer topological order. The divide's two operands + # are extract_slices inside the forall, so hop through each to its producer. + o_slice = prod(anytype, tiled_div, operand_number=0) + pv = prod(anytype, o_slice, operand_number=0) # @V batch_matmul (unnormalized) + l_slice = prod(anytype, tiled_div, operand_number=1) + den = prod(anytype, l_slice, operand_number=0) # row-sum generic + num = prod(anytype, pv, operand_number=0) # exp generic, feeds @V and the sum + pv_fill = prod(anytype, pv, operand_number=2) # 0 fill (@V acc) den_fill = prod(anytype, den, operand_number=1) # 0 fill (sum acc) mx = prod(anytype, num, operand_number=1) # max-reduce generic mx_fill = prod(anytype, mx, operand_number=1) # -inf fill (max acc) @@ -202,13 +194,15 @@ def fuse(p, c): qkt = prod(anytype, scaled, operand_number=0) # QK^T batch_matmul kt = prod(anytype, qkt, operand_number=1) # K^T transpose qkt_fill = prod(anytype, qkt, operand_number=2) # 0 fill (qkt acc) + # `num` is fused after both its consumers (@V and the row sum) are inside. for p in ( - div, + pv, den, num, mx, scaled, qkt, + pv_fill, den_fill, mx_fill, scale_fill, @@ -222,31 +216,135 @@ def fuse(p, c): def _fuse_attention_in_region(anytype, forall, fa_params): - """Rewrite one attention region's tensor-level batch_matmul pair (QK^T, @V) - into the flash loop via the transform op. Scoped to `forall` so counts are - exact at any multiplicity. Runs right after the region was tiled, i.e. still - on tensors, so the shared vectorize tail lowers the emitted loop.""" + """Fold one attention region's chain into the flash loop. + + `fuse_dependant_reduction_ops` does the work: it moves the elementwise term and + one consumer reduction into an already-tiled producer reduction loop and inserts + the online correction that rescales that reduction's running accumulator + whenever the running max changes. Inside `forall` the chain reads: + + %s = linalg.mul(batch_matmul(q, k^T), fill(scale)) the scaled scores + %m = max_j %s the producer reduction + %p = exp(%s - %m) the elementwise term + %o = batch_matmul(%p, v) consumer reduction + %l = sum_j %p consumer reduction + %out = %o / %l the deferred divide + + Every match is scoped to `forall` so counts are exact at any multiplicity. Runs + right after the region was tiled, i.e. still on tensors, so the shared vectorize + tail lowers the emitted loop. + + NB: non-causal only -- there is no `causal` parameter yet. + """ prod = transform.get_producer_of_operand - bmms = match_and_split(forall, ops={"linalg.batch_matmul"}, nhandles=2) - qk_bmm, pv_bmm = bmms[0], bmms[1] - q = prod(anytype, qk_bmm, operand_number=0) - # K reaches the QK^T matmul through the linalg.transpose that forms K^T. - k = prod(anytype, prod(anytype, qk_bmm, operand_number=1), operand_number=0) - v = prod(anytype, pv_bmm, operand_number=1) - # The scale is the fill value of the linalg.mul rhs operand. - mul_op = match_and_split(forall, ops={"linalg.mul"}, nhandles=1)[0] - scale = prod(anytype, prod(anytype, mul_op, operand_number=1), operand_number=0) - # NB: the merged fused-attention op is non-causal only -- there is - # no `causal` parameter yet, so the model runs as non-causal attention. - transform_ext.replace_with_fused_attention( - q=q, - k=k, - v=v, - scale=scale, - output=pv_bmm, - tile_size=fa_params["inner_loop_tile_size"], + tile_size = fa_params["inner_loop_tile_size"] + + max_op, p_op, sum_op, _ = match_and_split( + forall, ops={"linalg.generic"}, nhandles=4 + ) + _, pv_bmm = match_and_split(forall, ops={"linalg.batch_matmul"}, nhandles=2) + # The fusion op wants both the elementwise term and the consumer reduction as + # linalg.generic ops, so generalize the contraction. + pv_op = structured.structured_generalize(anytype, pv_bmm) + + # Tile the row max along the key/value axis. This is the producer reduction loop + # the rest of the chain gets folded into; the marker attribute is what the + # fusion op recognizes it by. + _, reduction_loop = structured.structured_tile_using_for( + anytype, + [anytype], + max_op, + dynamic_sizes=[], + interchange=[], + static_sizes=[0, 0, tile_size], + scalable_sizes=[False, False, False], + ) + transform.annotate(reduction_loop, transform_ext.REDUCTION_LOOP_ATTR_NAME) + + # First chain: max -> p -> row sum. `p` also feeds the contraction, so the op + # fuses a clone of it and leaves the original in place for the second chain. + # Fusing replaces the loop, but the replacement inherits the marker attribute, + # so it is ready to serve as the producer reduction of the second chain. + reduction_loop = transform_ext.fuse_dependant_reduction_ops( + p_op, sum_op, reduction_loop ) + # Second chain: max -> p -> @V, into that same loop. The first fusion consumed + # the handle to `p`; the original is still the contraction's operand. + p_op = prod(anytype, pv_op, operand_number=0) + reduction_loop = transform_ext.fuse_dependant_reduction_ops( + p_op, pv_op, reduction_loop + ) + transform.apply_cse(forall) + + # Sink the score computation into the reduction loop as well, so only one + # [wg_rows, tile_size] score tile -- rather than the full [wg_rows, n_ctx] + # matrix -- is ever live. + for producer_name in ["linalg.mul", "linalg.batch_matmul", "linalg.transpose"]: + producer_op = match_and_split(forall, ops={producer_name}, nhandles=1)[0] + _, reduction_loop = structured.structured_fuse_into_containing_op( + anytype, + anytype, + producer_op=producer_op, + containing_op=reduction_loop, + ) + + # The Q @ K^T zero accumulator and the scale tensor are still filled at full + # [wg_rows, n_ctx] extent outside the loop, even though the ops now inside it + # only ever read a [wg_rows, tile_size] slice. Sink those two fills as well so + # neither tensor is materialized whole. Reach them through the in-loop consumer + # they initialize -- one hop for the slice the fusion left behind, one more for + # the fill itself. The remaining fills initialize the loop's running + # accumulators and must stay outside. + scale_mul_op = match_and_split(reduction_loop, ops={"linalg.mul"}, nhandles=1)[0] + qk_matmul = match_and_split( + reduction_loop, ops={"linalg.batch_matmul"}, nhandles=1 + )[0] + for consumer_op, operand_number in [(scale_mul_op, 1), (qk_matmul, 2)]: + fill_slice = prod(anytype, consumer_op, operand_number=operand_number) + fill_op = prod(anytype, fill_slice, operand_number=0) + _, reduction_loop = structured.structured_fuse_into_containing_op( + anytype, + anytype, + producer_op=fill_op, + containing_op=reduction_loop, + ) + + transform.apply_cse(forall) + canonicalize(forall) + + # Strip the head dim, which the region tiling cut down to 1. Everything + # downstream -- the XeGPU layouts and the blocking/distribution passes -- is + # built for rank-2 tiles, and the vector-level `cast_away_leading_one_dim` + # patterns cannot finish the job: they have no pattern for multi_reduction, + # broadcast or transpose, so the softmax row reductions and the correction's + # broadcasts would keep a unit dim and drag shape_casts (and rank-3 XeGPU + # layouts) along with them. `fold_unit_extent_dims` only rewrites + # linalg.generic, hence the generalize; and it leaves two-step slice chains + # behind, which plain canonicalization does not compose, hence the tensor + # patterns. + named_ops = match( + forall, + ops={ + "linalg.batch_matmul", + "linalg.mul", + "linalg.transpose", + "linalg.fill", + "linalg.elementwise", + }, + ) + structured.structured_generalize(anytype, named_ops) + with ir.InsertionPoint(transform.apply_patterns(forall).patterns): + structured.apply_patterns_linalg_fold_unit_extent_dims_via_slices() + transform.apply_patterns_canonicalization() + transform.apply_cse(forall) + with ir.InsertionPoint(transform.apply_patterns(forall).patterns): + tensor.apply_patterns_tensor_merge_consecutive_insert_extract_slice() + tensor.apply_patterns_tensor_drop_redundant_insert_slice_rank_expansion() + tensor.apply_patterns_tensor_fold_tensor_subset_ops() + transform.apply_patterns_canonicalization() + transform.apply_cse(forall) + def xegpu_fa_annotation(gf, fa_params): """Attach XeGPU layouts to one fused-attention gpu.func.""" @@ -422,16 +520,17 @@ def _bundle( # Generic build order: each layernorm contributes [mean, var, normalize] (3), # in block build order; each elementwise contributes 1. We reconstruct the # per-op handle slices from `kinds`. - # 'fa' softmax generics do NOT exist yet (fa is tiled last, softmax still - # un-decomposed), so they are not in this pool. The fa core's linalg.transpose - # /linalg.mul/batch_matmul are not linalg.generic, so also excluded. (The head + # Each 'fa' region contributes 4 bare generics -- row max, exp term, row sum + # and the deferred normalizing divide (see Builder.attention_4d) -- so they ARE + # in this pool, in that build order. The fa core's linalg.transpose /linalg.mul + # /batch_matmul are not linalg.generic, so those stay excluded. (The head # reshape is a pure memref VIEW -- no generic, no kernel; see Builder.heads_view.) - ngen_total = 3 * n_ln + n_ew + ngen_total = 3 * n_ln + n_ew + 4 * n_fa gen_handles = transform.split_handle( (anytype,) * ngen_total, match(mod, ops={"linalg.generic"}) ) # Walk kinds to assign generic handles to ops. - ln_slices, ew_handles = [], [] + ln_slices, ew_handles, fa_slices = [], [], [] gi = 0 for k in kinds: if k == "ln": @@ -442,7 +541,11 @@ def _bundle( elif k == "ew": ew_handles.append(gen_handles[gi]) gi += 1 - # mm / sm / fa contribute no bare linalg.generic here + elif k == "fa": + # (row_max, probs, row_sum, divide) + fa_slices.append(tuple(gen_handles[gi : gi + 4])) + gi += 4 + # mm contributes no bare linalg.generic here # 1) Tile layernorms FIRST, using preserved (mean,var,normalize) handles. # Doing this BEFORE EW/matmul tiling keeps the bare linalg.fill pool exactly @@ -477,22 +580,18 @@ def _bundle( for mm in mms: _tile_one_matmul(mm, mm_params) - # 5) Fused-attention regions. Done last so the generic pre-split above ran while - # each fa softmax was still one linalg.softmax (its decomposition generics - # don't exist yet, so they can't inflate ngen_total). Pre-split the 2*n_fa - # batch_matmuls (build order [QK^T, @V] per region) + n_fa softmaxes by count, - # then tile+fuse each region into one forall (decompose happens in-region). - if n_fa: - fa_bmms = match_and_split(mod, ops={"linalg.batch_matmul"}, nhandles=2 * n_fa) - fa_softmaxes = match_and_split(mod, ops={"linalg.softmax"}, nhandles=n_fa) - for r in range(n_fa): - _, fa_forall = _tile_one_fused_attention_region( - anytype, fa_bmms[2 * r + 1], fa_softmaxes[r], fa_params - ) - # Rewrite the region into the flash online-softmax loop while it is - # still on tensors, so the shared vectorize tail lowers it like any - # other tiled region. - _fuse_attention_in_region(anytype, fa_forall, fa_params) + # 5) Fused-attention regions. Done last so every other kernel is already tiled + # when the reduction fusion runs. Each region is driven from its pre-split + # leaf handle -- the deferred normalizing divide -- and tiled into one forall; + # the flash loop is then folded out of the chain inside it. + for divide_op in fa_slices: + _, fa_forall = _tile_one_fused_attention_region( + anytype, divide_op[3], fa_params + ) + # Fold the region into the flash online-softmax loop while it is still on + # tensors, so the shared vectorize tail lowers it like any other tiled + # region. + _fuse_attention_in_region(anytype, fa_forall, fa_params) func = match(mod, ops={"func.func"}) lh_transform.cleanup(func) @@ -512,6 +611,12 @@ def _bundle( with lh_transform.foreach(match(mod, ops={"scf.for"})) as reduction_loop: lh_transform.loop_hoisting(reduction_loop) transform.yield_() + # Each reduction fusion left the elementwise term's full-extent result as a + # loop accumulator that nothing reads (every tile but the last is scaled by a + # stale running max). Only now, with vectorization having replaced the in-loop + # destination slice with transfer ops, does liveness see them as unused; left in + # place they bufferize into real stores of a stale tensor. + func = apply_registered_pass(func, "remove-dead-values") lh_transform.cleanup(func) # Drop any leading unit dims left over from the (1, wg_rows, 0, 0) tiling of # the attention regions so the QK^T/@V vector.contracts stay 2D. diff --git a/lighthouse/dialects/transform/transform_ext/__init__.py b/lighthouse/dialects/transform/transform_ext/__init__.py index 45aa49d0..6cea25b6 100644 --- a/lighthouse/dialects/transform/transform_ext/__init__.py +++ b/lighthouse/dialects/transform/transform_ext/__init__.py @@ -7,34 +7,40 @@ from .ops.get_tile_sizes import get_tile_sizes from .ops.param_cmp_eq import param_cmp_eq from .ops.replace import replace +from .ops.replace_with_fused_attention import replace_with_fused_attention from .ops.convert_func_results_to_args import convert_func_results_to_args +from .ops.enable_fastmath_optimizations import enable_fastmath_optimizations from .ops.extract_handle import extract_handle from .ops.get_tileable_consumers import get_tileable_consumers from .ops.get_tiling_sizes import get_tiling_sizes from .ops.trace_producers import trace_producers from .ops.reverse_handles import reverse_handles from .ops.update_address_space import update_address_space -from .ops.replace_with_fused_attention import replace_with_fused_attention from .ops.filter_num_loops import filter_num_loops from .ops.filter_elementwise import filter_elementwise from .ops.filter_by_name import filter_by_name from .ops.filter_reduction_ops import filter_reduction_ops +from .ops.fuse_dependant_reduction_ops import fuse_dependant_reduction_ops from .ops.get_leading_unit_tile_sizes import get_leading_unit_tile_sizes from .ops.move_offsets_to_subview import move_offsets_to_subview from .ops.clear_tile_and_fuse_annotations import clear_tile_and_fuse_annotations from .ops.get_fusion_roots import get_fusion_roots from .ops.propagate_tile_sizes import propagate_tile_sizes +from .utils.dependant_reduction_legality import REDUCTION_LOOP_ATTR_NAME __all__ = [ + "REDUCTION_LOOP_ATTR_NAME", "TransformExtensionDialect", "assign_tile_sizes", "clear_tile_and_fuse_annotations", "convert_func_results_to_args", + "enable_fastmath_optimizations", "extract_handle", "filter_by_name", "filter_elementwise", "filter_num_loops", "filter_reduction_ops", + "fuse_dependant_reduction_ops", "get_fusion_roots", "get_leading_unit_tile_sizes", "get_named_attribute", diff --git a/lighthouse/dialects/transform/transform_ext/ops/enable_fastmath_optimizations.py b/lighthouse/dialects/transform/transform_ext/ops/enable_fastmath_optimizations.py new file mode 100644 index 00000000..7ea994b7 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/ops/enable_fastmath_optimizations.py @@ -0,0 +1,192 @@ +from mlir import ir +from mlir.dialects import arith, ext, math, transform +from mlir.dialects.transform import DiagnosedSilenceableFailure + +from lighthouse.dialects.transform.transform_ext import TransformExtensionDialect + +#: The `fastmath` attribute's default. It is present in the attribute dictionary +#: but elided when printed, and carries no intent, so it counts as absent. +NO_FASTMATH = "#arith.fastmath" + +#: Ops that get the `fastmath` attribute. Deliberately narrow -- only what the +#: current usecases need. +#: +#: - `math.exp`, for the `exp(a)/exp(b)` fold below. +#: - `arith.subf`, because the correction's terms come out as `0.0 - m` (the +#: fusion substitutes the additive neutral for `E`'s data operand). LLVM only +#: narrows that to a negation under `nsz` -- `0.0 - x` and `-x` disagree at +#: `x = 0.0` -- so without the flags the subtraction survives to the ISA. +FASTMATH_OP_NAMES = frozenset({"math.exp", "arith.subf"}) + + +def _collect(root: ir.Operation, names) -> list[ir.Operation]: + """Every op under `root` whose name is in `names`, in pre-order.""" + ops: list[ir.Operation] = [] + + def visit(op: ir.Operation) -> ir.WalkResult: + if op.name in names: + ops.append(op) + return ir.WalkResult.ADVANCE + + root.walk(visit, ir.WalkOrder.PRE_ORDER) + return ops + + +def _annotate_fastmath(root: ir.Operation) -> None: + """Set `fastmath` on every `FASTMATH_OP_NAMES` op under `root`. + + Ops that already carry a non-default `fastmath` attribute are left alone, so an + explicitly-chosen weaker flag set is never widened. + """ + with root.context: + fast = ir.Attribute.parse("#arith.fastmath") + for op in _collect(root, FASTMATH_OP_NAMES): + existing = op.attributes.get("fastmath") + if existing is not None and str(existing) != NO_FASTMATH: + continue + op.attributes["fastmath"] = fast + + +def _fastmath_allows_transform(op: ir.Operation) -> bool: + """True if `op`'s `fastmath` flags permit `exp(a)/exp(b)` -> `exp(a-b)`. + + The rewrite is not IEEE-preserving and relies on three flags: + + - `reassoc`, which licenses regrouping the expression at all. + - `nnan` and `ninf`, because the two forms disagree once an exponential + overflows or underflows. For `a = b = 1000`, `exp(a)/exp(b)` is `inf/inf`, + i.e. NaN, while `exp(a-b)` is 1. + + Only the full `fastmath` set is accepted. That is stricter than the + three flags above require -- `` would also be sound -- but + it is the only set the payloads here produce, so the check stays a plain + comparison. + """ + if "fastmath" not in op.attributes: + return False + return str(op.attributes["fastmath"]) == "#arith.fastmath" + + +def _exp_operand(value: ir.Value) -> ir.Value | None: + """The argument of `value`'s defining `math.exp`, if it is a fast-math one.""" + producer = value.owner + if not isinstance(producer, ir.Operation): + producer = getattr(producer, "operation", None) + if producer is None or producer.name != "math.exp": + return None + if not _fastmath_allows_transform(producer): + return None + return producer.operands[0] + + +def _fold_exp_div(div_op: ir.Operation, rewriter: transform.TransformRewriter) -> bool: + """Rewrite one `exp(a) / exp(b)` into `exp(a - b)`. True if it applied. + + Gated on the two `math.exp` producers alone, not on the `arith.divf`: the flags + the rewrite needs describe the exponentials overflowing, and the divide is not in + `FASTMATH_OP_NAMES`. + """ + numerator = _exp_operand(div_op.operands[0]) + denominator = _exp_operand(div_op.operands[1]) + if numerator is None or denominator is None: + return False + + # exp(a)/exp(b) and exp(a-b) only agree when a and b -- hence a-b -- share the + # result type; a mixed-precision chain is left alone. + if numerator.type != denominator.type: + return False + if numerator.type != div_op.results[0].type: + return False + + # Both replacements are flagged directly: step 1 has already walked the payload, + # so ops created here never pass through `_annotate_fastmath`. The subtract wants + # the flags for the same reason every other one does -- `a - b` is + # `(0.0 - m_new) - (0.0 - m_old)`, which only collapses to `m_old - m_new` once + # reassociation is licensed. + fast = arith.FastMathFlags.fast + with ir.InsertionPoint(div_op), div_op.location: + difference = arith.SubFOp(numerator, denominator, fastmath=fast) + replacement = math.ExpOp(difference.result, fastmath=fast) + # The two original exps are left in place; they are dead only if nothing else + # reads them, which DCE/CSE afterwards decides. + rewriter.replace_op(div_op, [replacement.result]) + return True + + +class EnableFastmathOptimizationsOp( + TransformExtensionDialect.Operation, name="enable_fastmath_optimizations" +): + """ + Enable fast-math and apply the rewrites it licenses, under `target`. + + Two steps, in this order: + + 1. Set `fastmath` on every op in `FASTMATH_OP_NAMES` -- a deliberately + narrow set, see there for what is in it and why. Ops that already carry a + non-default flag set keep it, so an explicitly-chosen weaker one is never + widened. + 2. Rewrite `exp(a) / exp(b)` into `exp(a - b)`, trading a transcendental and + a divide for a subtract. Only applied where both `math.exp` producers carry + `fastmath`; see `_fastmath_allows_transform` for which flags the + rewrite relies on. + + The order matters and is why the two are one op: the fold keys off the + annotation. Ops without the flag are skipped rather than rejected, so this is + safe to run over a whole function. + + Ideally must be run after vectorization. + + Args: + target: Handle to root ops to work within (e.g. func.func). + Returns: + Handle to the same target roots. + """ + + target: ext.Operand[transform.AnyOpType] + optimized_ops: ext.Result[transform.AnyOpType[()]] = ext.infer_result() + + @classmethod + def attach_interface_impls(cls, context=None): + cls.TransformOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + cls.MemoryEffectsOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + + class TransformOpInterfaceModel(transform.TransformOpInterface): + @staticmethod + def apply( + op: "EnableFastmathOptimizationsOp", + rewriter: transform.TransformRewriter, + results: transform.TransformResults, + state: transform.TransformState, + ) -> DiagnosedSilenceableFailure: + targets = state.get_payload_ops(op.target) + + for target in targets: + _annotate_fastmath(target) + for div_op in _collect(target, {"arith.divf"}): + _fold_exp_div(div_op, rewriter) + + results.set_ops(op.optimized_ops, targets) + return DiagnosedSilenceableFailure.Success + + @staticmethod + def allow_repeated_handle_operands( + _op: "EnableFastmathOptimizationsOp", + ) -> bool: + return False + + class MemoryEffectsOpInterfaceModel(ir.MemoryEffectsOpInterface): + @staticmethod + def get_effects(op: ir.Operation): + return ( + transform.only_reads_handle(op.op_operands) + + transform.produces_handle(op.results) + + transform.modifies_payload() + ) + + +def enable_fastmath_optimizations( + target: ir.Value[transform.AnyOpType], +) -> ir.Value[transform.AnyOpType]: + """snake_case wrapper to create EnableFastmathOptimizationsOp.""" + op = EnableFastmathOptimizationsOp(target=target) + return op.optimized_ops diff --git a/lighthouse/dialects/transform/transform_ext/ops/fuse_dependant_reduction_ops.py b/lighthouse/dialects/transform/transform_ext/ops/fuse_dependant_reduction_ops.py new file mode 100644 index 00000000..9c63efff --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/ops/fuse_dependant_reduction_ops.py @@ -0,0 +1,191 @@ +import sys + +from mlir import ir +from mlir.dialects import ext, linalg, scf, transform +from mlir.dialects.transform import DiagnosedSilenceableFailure + +from lighthouse.dialects.transform.transform_ext import TransformExtensionDialect +from lighthouse.dialects.transform.transform_ext.utils import ( + dependant_reduction_legality as legality, +) +from lighthouse.dialects.transform.transform_ext.utils.dependant_reduction_fusion import ( + fuse_dependant_reduction_ops as apply_fusion, +) + + +def _single(payload_ops, what: str): + """The one payload op behind a handle, or None if the handle is not singular.""" + ops = list(payload_ops) + if len(ops) != 1: + print( + f"fuse_dependant_reduction_ops: requires exactly one {what}, got " + f"{len(ops)}", + file=sys.stderr, + ) + return None + return ops[0].opview if isinstance(ops[0], ir.Operation) else ops[0] + + +class FuseDependantReductionOpsOp( + TransformExtensionDialect.Operation, name="fuse_dependant_reduction_ops" +): + """ + Fuses a dependency chain ``R1 -> E -> R2`` into a single online (one-pass) + reduction loop, where ``R1`` is the producer reduction (`tiled_reduction_loop`), + ``E`` is an elementwise op (`elementwise_op`) consuming ``R1``'s result as a + broadcast input, and ``R2`` is the consumer reduction (`reduction_op`) reducing + ``E``'s output over the same axis. + + ``E`` is kept as a *separate* elementwise op rather than folded into ``R2``'s + body; this keeps the per-element term explicit and easier to vectorize later. + Because ``E`` is exactly the per-element term that ``R2`` reduces, the online + correction factor is derived directly from ``E``. + + The producer must already be tiled along the shared reduction dimension into + an ``scf.for`` annotated with the ``__reduction_loop__`` unit attribute; the + tile size is read from that loop's step. This op does not tile the producer + itself -- use ``transform.structured.tile_using_for`` followed by + ``transform.annotate`` for that. + TODO: Remove the annotation requirement for the producer reduction loop. + + Whichever copy is fused computes each tile against the *running* ``R1`` + accumulator, which the online correction only accounts for on ``R2``'s + accumulator, so its own full-extent result is stale in every tile but the last + and must not be read off the loop. That result surfaces as a dead loop result, + which the schedule drops with ``remove-dead-values`` after vectorization. + + So if ``E`` has users besides ``R2``, this op fuses a *clone* and leaves the + original where it is: there it still reads ``R1``'s final result off the loop and + recomputes the term correctly for those users. With ``R2`` as its only user, + ``E`` itself is fused and nothing is left behind. + + One application fuses one consumer reduction. ``E`` may have several, though -- + any consumer reducing it over the same axis is a candidate ``R2`` in its own + right -- so applying the op once per consumer folds an attention chain (one + ``exp`` term feeding both a row sum and a `P @ V` contraction) into a single loop, + each application adding its own online correction. The returned loop keeps the + ``__reduction_loop__`` attribute, so it can be passed straight back in as + `tiled_reduction_loop` with no re-annotation. + + For more details about the fusion algorithm refer: + https://discourse.llvm.org/t/linalg-rfc-an-approach-for-tiling-and-fusing-dependent-reductions-in-linalg/91698 + + Return modes: + Each handle must point to exactly one payload op, otherwise this produces + a silenceable failure -- as it does when the loop is not an ``scf.for``, + when either op is not a ``linalg.generic``, or when the triple does not + satisfy the fusion legality conditions. The rejected condition is reported + on stderr. + + Args: + elementwise_op: Handle to the elementwise term ``E``. + reduction_op: Handle to the consumer reduction ``R2``. + tiled_reduction_loop: Handle to ``R1``'s already-tiled ``scf.for``. + Returns: + Handle to the fused loop. + """ + + elementwise_op: ext.Operand[transform.AnyOpType] + reduction_op: ext.Operand[transform.AnyOpType] + tiled_reduction_loop: ext.Operand[transform.AnyOpType] + fused_loop: ext.Result[transform.AnyOpType[()]] = ext.infer_result() + + @classmethod + def attach_interface_impls(cls, context=None): + cls.TransformOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + cls.MemoryEffectsOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + + class TransformOpInterfaceModel(transform.TransformOpInterface): + @staticmethod + def apply( + op: "FuseDependantReductionOpsOp", + rewriter: transform.TransformRewriter, + results: transform.TransformResults, + state: transform.TransformState, + ) -> DiagnosedSilenceableFailure: + r1_loop = _single( + state.get_payload_ops(op.tiled_reduction_loop), + "tiled reduction loop", + ) + e = _single(state.get_payload_ops(op.elementwise_op), "elementwise op") + r2 = _single(state.get_payload_ops(op.reduction_op), "reduction op") + if r1_loop is None or e is None or r2 is None: + return DiagnosedSilenceableFailure.SilenceableFailure + + if not isinstance(r1_loop, scf.ForOp): + print( + "fuse_dependant_reduction_ops: expected the tiled reduction " + "loop to be an scf.for op", + file=sys.stderr, + ) + return DiagnosedSilenceableFailure.SilenceableFailure + for name, candidate in (("elementwise", e), ("reduction", r2)): + if not isinstance(candidate, linalg.GenericOp): + print( + f"fuse_dependant_reduction_ops: expected the {name} op to " + f"be a linalg.generic op", + file=sys.stderr, + ) + return DiagnosedSilenceableFailure.SilenceableFailure + + if not legality.collect_inner_reduction_generics(r1_loop): + print( + "fuse_dependant_reduction_ops: the reduction loop body does " + "not contain any reduction linalg.generic", + file=sys.stderr, + ) + return DiagnosedSilenceableFailure.SilenceableFailure + + result_to_inner = legality.map_loop_results_to_inner_reductions(r1_loop) + try: + e_tiled_dim, tile_size = legality.check_legal_fusion_triple( + r1_loop, result_to_inner, e, r2 + ) + fused = apply_fusion(rewriter, r1_loop, e, r2, e_tiled_dim, tile_size) + except legality.FusionRejected as rejected: + print( + f"fuse_dependant_reduction_ops: could not fuse the elementwise " + f"op and the consumer reduction into the producer reduction " + f"loop -- {rejected}", + file=sys.stderr, + ) + return DiagnosedSilenceableFailure.SilenceableFailure + + results.set_ops(op.fused_loop, [fused]) + return DiagnosedSilenceableFailure.Success + + @staticmethod + def allow_repeated_handle_operands(_op: "FuseDependantReductionOpsOp") -> bool: + return False + + class MemoryEffectsOpInterfaceModel(ir.MemoryEffectsOpInterface): + @staticmethod + def get_effects(op: "FuseDependantReductionOpsOp"): + return ( + transform.consumes_handle(op.op_operands) + + transform.produces_handle(op.results) + + transform.modifies_payload() + ) + + +def fuse_dependant_reduction_ops( + elementwise_op: ir.Value[transform.AnyOpType], + reduction_op: ir.Value[transform.AnyOpType], + tiled_reduction_loop: ir.Value[transform.AnyOpType], +) -> ir.Value[transform.AnyOpType]: + """ + snake_case wrapper to create a FuseDependantReductionOpsOp. + + Args: + elementwise_op: Handle to the elementwise term ``E``. + reduction_op: Handle to the consumer reduction ``R2``. + tiled_reduction_loop: Handle to ``R1``'s already-tiled ``scf.for``. + Returns: + Handle to the fused loop. + """ + op = FuseDependantReductionOpsOp( + elementwise_op=elementwise_op, + reduction_op=reduction_op, + tiled_reduction_loop=tiled_reduction_loop, + ) + return op.fused_loop diff --git a/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py b/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py index 614e51b0..4dbae123 100644 --- a/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py +++ b/lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py @@ -285,11 +285,12 @@ def apply( # f32 for numerical accuracy; only the matmul operands keep their # narrower element types. k_element_type = ir.RankedTensorType(k.type).element_type - # P is the lhs operand of the contraction being replaced, so its - # element type is the precision the rest of the graph expects. - p_element_type = ir.RankedTensorType( - output_op.operands[0].type - ).element_type + # P is the lhs of the `P @ V` contraction, so it carries V's element + # type: a `linalg.batch_matmul`'s two operands must agree. Derived from + # `v` rather than from `output_op`'s lhs so that `output` may be either + # the contraction itself or, when the payload defers the normalizing + # divide past it, that divide. + p_element_type = ir.RankedTensorType(v.type).element_type out_element_type = ir.RankedTensorType( output_op.results[0].type ).element_type diff --git a/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_fusion.py b/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_fusion.py new file mode 100644 index 00000000..386ff7bc --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_fusion.py @@ -0,0 +1,617 @@ +"""The dependent-reduction (flash-attention style) fusion rewrite. + +Fuses an ``R1 -> E -> R2`` chain into ``R1``'s already-tiled reduction loop, +turning a two-pass reduction into an online one-pass loop. See +`dependant_reduction_legality` for the chain's shape and the conditions checked +before this runs. + +The rewrite *rebuilds* the loop rather than mutating it: MLIR cannot grow an +``scf.for``'s ``iter_args`` in place, so a new loop is created carrying the +original accumulators plus one for ``E``'s result and one for ``R2``'s, the +original body is cloned into it, and the fused ops are appended. + +For softmax ``m = max_j x``, ``p = exp(x - m)``, ``s = sum_j p`` the result is:: + + %loop:3 = scf.for %iv = 0 to 512 step 32 + iter_args(%mArg = %mInit, %eArg = %eInit, %sArg = %sInit) { + %xt = tensor.extract_slice %X[0, %iv] [64, 32] [1, 1] + %mOld = tensor.extract_slice %mArg[0] [64] [1] + %mNew = linalg.generic ins(%xt) outs(%mOld) { arith.maximumf } // R1 + %et = tensor.extract_slice %eArg[0, %iv] [64, 32] [1, 1] + %p = linalg.generic ins(%xt, %mNew) outs(%et) { subf; exp } // fused E + %sOld = tensor.extract_slice %sArg[0] [64] [1] + %tNew = linalg.generic ins(%mNew) outs(...) { subf 0, m; exp } // correction + %tOld = linalg.generic ins(%mOld) outs(...) { subf 0, m; exp } + %f = linalg.elementwise kind=div ins(%tNew, %tOld) + %sc = linalg.elementwise kind=mul ins(%sOld, %f) + %sNew = linalg.generic ins(%p) outs(%sc) { arith.addf } // fused R2 + scf.yield %mNew, insert_slice(%p into %eArg), insert_slice(%sNew into %sArg) + } + +``%eArg`` carries ``E``'s full-extent result only so the fused op has a destination; +every tile but the last is stale (computed against the *running* accumulator), so +nothing may read it off the loop. It is dead on arrival, but not provably so until +vectorization has replaced the in-loop destination slice with transfer ops -- so the +schedule runs ``remove-dead-values`` after vectorizing to unwind it. + +The op fused is ``E`` itself, unless ``E`` has a consumer besides ``R2``: then a +clone is fused and the original stays put for that consumer, still reading ``R1``'s +final result. See ``needs_elementwise_clone``. +""" + +from mlir import ir +from mlir.dialects import arith, linalg, scf, tensor + +from lighthouse.utils.mlir import opview +from lighthouse.dialects.transform.transform_ext.utils import ir_rewrite as irr +from lighthouse.dialects.transform.transform_ext.utils import linalg_structured as ls +from lighthouse.dialects.transform.transform_ext.utils.dependant_reduction_legality import ( + FusionRejected, + collect_r1_as_elementwise_inputs, + find_r2_elementwise_operand, + map_loop_results_to_inner_reductions, + needs_elementwise_clone, +) + +__all__ = ["fuse_dependant_reduction_ops"] + + +def _tile_bounds( + shaped: ir.Value, + imap: ir.AffineMap, + tiled_dim: int, + iv: ir.Value, + tile_size: int, +) -> tuple[list[int], list[ir.Value], list[int]]: + """Offsets/sizes selecting the current reduction tile of `shaped`. + + Every tensor position whose map result references `tiled_dim` is cut to the + tile (offset = the loop IV, size = `tile_size`); the rest span their full + extent. Returns the offsets split into the static list (with + ``kDynamic`` where the IV goes) and the matching dynamic values, plus the + sizes -- the form the slice builders take. + """ + shape = list(ir.ShapedType(shaped.type).shape) + offsets: list = [0] * len(shape) + sizes: list = list(shape) + for pos, expr in enumerate(imap.results): + if isinstance(expr, ir.AffineDimExpr) and expr.position == tiled_dim: + offsets[pos] = iv + sizes[pos] = tile_size + + static_offsets: list[int] = [] + dynamic_offsets: list[ir.Value] = [] + for offset in offsets: + if isinstance(offset, int): + static_offsets.append(offset) + else: + static_offsets.append(ir.ShapedType.get_dynamic_size()) + dynamic_offsets.append(offset) + return static_offsets, dynamic_offsets, sizes + + +def _tile_slice( + source: ir.Value, + imap: ir.AffineMap, + tiled_dim: int, + iv: ir.Value, + tile_size: int, +) -> ir.Value: + """Extract the current reduction tile of `source`, per its indexing map. + + This is what re-slices a fused op's full-extent operands down to the tile the + enclosing loop is on. An operand not spanning `tiled_dim` (a broadcast running + accumulator) comes back as a full-extent slice. + """ + static_offsets, dynamic_offsets, sizes = _tile_bounds( + source, imap, tiled_dim, iv, tile_size + ) + result_type = ir.RankedTensorType.get( + sizes, ir.ShapedType(source.type).element_type + ) + return tensor.extract_slice( + result_type, + source, + dynamic_offsets, + [], + [], + static_offsets=static_offsets, + static_sizes=sizes, + static_strides=[1] * len(sizes), + ) + + +def _insert_tile_slice( + tile: ir.Value, + dest: ir.Value, + imap: ir.AffineMap, + tiled_dim: int, + iv: ir.Value, + tile_size: int, +) -> ir.Value: + """Insert `tile` back into `dest` at the current reduction tile's offsets. + + The mirror of `_tile_slice`: it yields the loop-carried destination for a + fused op whose result spans `tiled_dim`. For a result that does not span it + (a reduction accumulator) this degenerates to a full-extent insert. + """ + static_offsets, dynamic_offsets, sizes = _tile_bounds( + dest, imap, tiled_dim, iv, tile_size + ) + return tensor.insert_slice( + tile, + dest, + dynamic_offsets, + [], + [], + static_offsets=static_offsets, + static_sizes=sizes, + static_strides=[1] * len(sizes), + ) + + +def _clone_generic_with_operands( + src: ir.OpView, + inputs: list[ir.Value], + outputs: list[ir.Value], + result_types: list[ir.Type], +) -> ir.OpView: + """Rebuild `src` over new operands, copying its maps, iterators and body. + + Retiling changes only operand *extents*, never the indexing maps or iterator + types, so both are carried over verbatim. The body is cloned op by op: a + linalg body is straight-line scalar arithmetic with no nested regions. + """ + op = linalg.GenericOp( + result_tensors=result_types, + inputs=inputs, + outputs=outputs, + indexing_maps=src.indexing_maps, + iterator_types=src.iterator_types, + ) + src_body = src.regions[0].blocks[0] + body = op.regions[0].blocks.append(*[a.type for a in src_body.arguments]) + with ir.InsertionPoint(body): + irr.clone_block_body(src_body, list(body.arguments), skip_terminator=False) + return op + + +def _emit_elementwise( + kind: linalg.ElementwiseKind, + scalar_op, + inputs: list[ir.Value], + dest: ir.Value, +) -> ir.Value: + """Emit a ``linalg.elementwise`` of `kind` over `inputs` into `dest`. + + The Python builder fills in the default identity indexing maps but leaves the + region empty (the C++ region builder is not bound), so the scalar body is emitted + here via `scalar_op` -- which therefore has to agree with `kind`. + """ + op = linalg.ElementwiseOp( + result_tensors=[dest.type], inputs=inputs, outputs=[dest], kind=kind + ) + element_type = ir.ShapedType(dest.type).element_type + arg_types = [element_type] * (len(inputs) + 1) + body = op.regions[0].blocks.append(*arg_types) + with ir.InsertionPoint(body): + linalg.yield_([scalar_op(body.arguments[0], body.arguments[1])]) + return op.results[0] + + +def _cast_tensor(value: ir.Value, element_type: ir.Type) -> ir.Value: + """Elementwise ``extf``/``truncf`` of a whole tensor to `element_type`. + + The caller has established that the conversion exists (`irr.cast_float` accepts + the pair), so the emitted body cannot fail to build. + """ + shaped = ir.RankedTensorType(value.type) + identity = ir.AffineMapAttr.get(ir.AffineMap.get_identity(shaped.rank)) + dest = tensor.empty(irr.mixed_sizes(value), element_type) + op = linalg.GenericOp( + result_tensors=[dest.type], + inputs=[value], + outputs=[dest], + indexing_maps=ir.ArrayAttr.get([identity, identity]), + iterator_types=ir.ArrayAttr.get( + [ + ir.Attribute.parse("#linalg.iterator_type") + for _ in range(shaped.rank) + ] + ), + ) + body = op.regions[0].blocks.append(shaped.element_type, element_type) + with ir.InsertionPoint(body): + linalg.yield_([irr.cast_float(body.arguments[0], element_type)]) + return op.results[0] + + +def _emit_correction_term( + e: ir.OpView, + bindings: list[tuple[ls.Operand, ir.Value]], + compute_type: ir.Type, +) -> ir.Value | None: + """Emit `e`'s body at the current insertion point, isolated on its accumulators. + + `bindings` binds each accumulator-reading ``E`` input operand to the scalar to + substitute for it -- the current-tile value for the "new" factor, the previous + running value for the "old" one. + + Every op is rebuilt at `compute_type` rather than at the type it had in ``E``, + so the term is evaluated in the precision `_correction_factor` picked; the bound + values are already of that type, and anything the body captures from an enclosing + scope is converted to it. + + Because ``E`` is a separate op its body can be cloned directly; no backward + slice is needed, ``E``'s yielded value *is* the term. Data block arguments are + replaced by a stand-in constant chosen per consuming op + (`operand_eliminating_constant`): ``1.0`` for ``mulf``/``divf``, ``0.0`` for + ``addf``/``subf``. What makes that sound is not the constant itself but + separability, which legality has already established: with + ``E(x, m) = f(x) * g(m)``, substituting any ``c`` for ``x`` leaves + ``E(c, m_new) / E(c, m_old) = g(m_new) / g(m_old)``. For softmax + ``E = exp(x - m)`` and ``c = 0`` that is ``exp(-m_new) / exp(-m_old)``, the same + ratio as ``exp(x - m_new) / exp(x - m_old)`` for any ``x``. + + TODO: Caveat: the cancellation needs ``f(c)`` to be finite and non-zero, which is not + checked. A body where the substitution zeroes the data factor -- ``mulf`` fed by + a ``subf`` of two data arguments, say ``(x1 - x2) * exp(-m)`` -- is accepted by + the separability analysis but yields ``0 / 0`` here. + + Returns the cloned term value, or None if a data argument feeds an op with no + stand-in constant, or a captured value cannot be converted to `compute_type`. + """ + body = e.regions[0].blocks[0] + ops = list(body.operations) + term = ops[-1].operands[0] + + # Bind the accumulator block arguments; data arguments stay unmapped and are + # neutralized per consuming op below. + value_map: dict = {} + for operand, value in bindings: + value_map[ls.matching_block_argument(e, operand)] = value + + for op in ops[:-1]: + ov = opview(op) + temporary: list = [] + for operand in ov.operands: + if operand in value_map: + continue + if isinstance(operand, ir.BlockArgument) and operand.owner == body: + neutral = irr.operand_eliminating_constant(ov, compute_type) + if neutral is None: + return None + value_map[operand] = arith.constant(compute_type, neutral) + temporary.append(operand) + else: + # Captured from an enclosing scope: kept, but at `compute_type`. Not + # dropped afterwards -- unlike a neutral element it is reusable. + converted = irr.cast_float(operand, compute_type) + if converted is None: + return None + value_map[operand] = converted + if irr.clone_op_with_map(ov, value_map, result_type=compute_type) is None: + return None + # Drop the per-op substitutions so the next consumer of the same data + # argument gets its own neutral element. + for operand in temporary: + del value_map[operand] + + return value_map.get(term, term) + + +def _correction_factor( + e: ir.OpView, + r2: ir.OpView, + accumulators: list[tuple[ls.Operand, ir.Value, ir.Value]], + r2_accumulator: ir.Value, +) -> ir.Value | None: + """Build the online rescale factor for the current tile: ``term(new)/term(old)``. + + `accumulators` holds ``(fused E operand, new value, old value)`` per running + accumulator the fused ``E`` reads. The two ``term`` evaluations and the + division are ``linalg.generic``s over ``R2``'s *parallel* iteration space, so + the factor takes ``R2``'s accumulator shape -- for a contraction that space is + wider than ``E``'s (a GEMM's extra N dim) and the accumulator must broadcast + over it. + + Each accumulator is read through its own ``E`` indexing map, translated from + ``E``'s dims into ``R2``'s and with ``R2``'s reduction dim projected out. + Legality guarantees those maps do not reference the reduction axis, so the + projection is lossless. + + ``E``'s body and ``R2``'s accumulator need not share an element type (narrow + probabilities feeding a wide contraction accumulator is the usual attention + shape), so the term is evaluated in whichever of the two is wider and cast at the + boundaries: the accumulator inputs are converted on entry to the body, and the + finished factor is converted back to ``R2``'s accumulator type for the rescale + multiply. Evaluating in the wider type keeps a ratio of two exponentials off f16's + exponent range even when the payload's term is f16. + """ + # Align E's output map with R2's map for E's result position-by-position to + # get the E-dim -> R2-dim correspondence. + r2_e_operand = find_r2_elementwise_operand(r2, e) + e_out_map = ls.indexing_map_for(e, ls.dps_init_operands(e)[0]) + r2_e_map = ls.indexing_map_for(r2, r2_e_operand) + if len(e_out_map.results) != len(r2_e_map.results): + return None + e_dim_to_r2: dict[int, int] = {} + for e_expr, r2_expr in zip(e_out_map.results, r2_e_map.results): + if not isinstance(e_expr, ir.AffineDimExpr) or not isinstance( + r2_expr, ir.AffineDimExpr + ): + return None + e_dim_to_r2[e_expr.position] = r2_expr.position + + r2_red_dim = ls.reduction_dims(r2)[0] + accumulator_type = ir.RankedTensorType(r2_accumulator.type) + rank = accumulator_type.rank + + indexing_maps = [] + for operand, _, _ in accumulators: + in_r2 = irr.remap_dims( + ls.indexing_map_for(e, operand), e_dim_to_r2, ls.num_loops(r2) + ) + if in_r2 is None: + return None + projected = irr.project_dims(in_r2, {r2_red_dim}) + if projected is None or projected.n_dims != rank: + return None + indexing_maps.append(projected) + # The factor takes R2's accumulator shape, with an identity map over its rank. + indexing_maps.append(ir.AffineMap.get_identity(rank)) + + maps_attr = ir.ArrayAttr.get([ir.AffineMapAttr.get(m) for m in indexing_maps]) + iterator_types = ir.ArrayAttr.get( + [ir.Attribute.parse("#linalg.iterator_type") for _ in range(rank)] + ) + element_type = accumulator_type.element_type + # Legality has already checked the two types have a common widening. + compute_type = irr.wider_float_type( + ir.ShapedType(e.results[0].type).element_type, element_type + ) + if compute_type is None: + return None + init = tensor.empty(irr.mixed_sizes(r2_accumulator), compute_type) + + def build_term(pick) -> ir.Value | None: + inputs = [pick(acc) for acc in accumulators] + op = linalg.GenericOp( + result_tensors=[init.type], + inputs=inputs, + outputs=[init], + indexing_maps=maps_attr, + iterator_types=iterator_types, + ) + # A block argument type follows its operand, so the accumulators enter the + # body at their own element type and are converted inside it. + arg_types = [ir.ShapedType(v.type).element_type for v in inputs] + body = op.regions[0].blocks.append(*arg_types, compute_type) + with ir.InsertionPoint(body): + bindings = [] + for (operand, _, _), arg in zip(accumulators, list(body.arguments)[:-1]): + converted = irr.cast_float(arg, compute_type) + if converted is None: + return None + bindings.append((operand, converted)) + term = _emit_correction_term(e, bindings, compute_type) + # `E` may yield a bare accumulator or a captured value; either way the + # yielded type has to be the output's. + term = None if term is None else irr.cast_float(term, compute_type) + if term is None: + return None + linalg.yield_([term]) + return op.results[0] + + term_new = build_term(lambda acc: acc[1]) + term_old = build_term(lambda acc: acc[2]) + if term_new is None or term_old is None: + return None + + factor = _emit_elementwise( + linalg.ElementwiseKind.div, + lambda a, b: arith.divf(a, b), + [term_new, term_old], + init, + ) + if compute_type == element_type: + return factor + # Narrow the ratio back to R2's accumulator type, which the rescale multiply and + # the fused R2 both work in. + return _cast_tensor(factor, element_type) + + +def fuse_dependant_reduction_ops( + rewriter, + r1_loop: ir.OpView, + e: ir.OpView, + r2: ir.OpView, + e_tiled_dim: int, + tile_size: int, +) -> ir.OpView: + """Fuse the ``R1 -> E -> R2`` chain into `r1_loop`, returning the fused loop. + + `e_tiled_dim` is the ``E`` loop dim carrying ``R2``'s reduction axis and + `tile_size` the loop's step, both as returned by + ``check_legal_fusion_triple``. Every payload mutation goes through `rewriter` + so the transform state keeps tracking the ops it replaces. + """ + r1_loop, e, r2 = opview(r1_loop), opview(e), opview(r2) + r2_red_dim = ls.reduction_dims(r2)[0] + num_old_results = len(list(r1_loop.results)) + + # Fuse a *clone* of E when another consumer needs the term too, leaving the + # original in place for them (see `needs_elementwise_clone`). The clone goes + # immediately before E, which keeps it above every other user of an R1 loop + # result, so the new loop built at its position still dominates them all. + fused_source = e + if needs_elementwise_clone(e, r2): + r2_e_operand = find_r2_elementwise_operand(r2, e) + with ir.InsertionPoint(e), e.location: + clone = opview(e.operation.clone()) + r2_e_operand.set(clone.results[0]) + fused_source = clone + + # The new loop replaces `fused_source`'s position: below E's and R2's inits + # (so they dominate it) yet at or above every user of an R1 loop result, which + # legality requires to post-dominate E. + anchor = fused_source + e_init_operand = ls.dps_init_operands(fused_source)[0] + r2_init_operand = ls.dps_init_operands(r2)[0] + + # Hoist the operand definitions the new loop will need above it. Tiling + # routinely leaves them below: E's and R2's destinations (a `tensor.empty` / + # `linalg.fill`) and, when an enclosing `scf.forall` tiled a parallel dim, the + # `tensor.extract_slice` feeding a data input -- e.g. the per-batch slice of + # `V` in an attention chain. Operands that depend on the loop are skipped; + # moving them would try to move the loop above itself. + to_hoist = [ + operand.value + for consumer in (fused_source, r2) + for operand in ls.operands_of(consumer) + if not irr.depends_on_op(operand.value, r1_loop) + ] + if not irr.move_value_definitions(to_hoist, anchor): + raise FusionRejected( + "could not move the fused ops' operand definitions above the reduction loop" + ) + + e_dest = e_init_operand.value + r2_dest = r2_init_operand.value + result_to_inner = map_loop_results_to_inner_reductions(r1_loop) + accumulator_operands, accumulator_result_indices = collect_r1_as_elementwise_inputs( + r1_loop, fused_source + ) + + # --- build the replacement loop ------------------------------------------ + with ir.InsertionPoint(anchor), r1_loop.location: + new_loop = scf.ForOp( + r1_loop.lowerBound, + r1_loop.upperBound, + r1_loop.step, + list(r1_loop.initArgs) + [e_dest, r2_dest], + ) + # Carry the original loop's attributes over, `__reduction_loop__` included, so + # the fused loop is still recognisable as a tiled reduction loop for a + # subsequent application (which is how a second consumer reduction is fused). + for name, attr in irr.op_attributes(r1_loop).items(): + new_loop.operation.attributes[name] = attr + + iv = new_loop.induction_variable + e_arg = new_loop.inner_iter_args[num_old_results] + r2_arg = new_loop.inner_iter_args[num_old_results + 1] + + value_map: dict = {r1_loop.induction_variable: iv} + value_map.update(zip(r1_loop.inner_iter_args, new_loop.inner_iter_args)) + + old_body_ops = list(r1_loop.body.operations) + with ir.InsertionPoint(new_loop.body), r1_loop.location: + # The original body, cloned; its inner reductions now accumulate into the + # new loop's `iter_arg`s. + for op in old_body_ops[:-1]: + irr.clone_op_deep_with_map(op, value_map) + + # --- the fused E, re-sliced to the current tile --- + # Each accumulator the fused E reads maps to its current-tile value (the + # cloned inner reduction's result) and its previous running value (that + # reduction's own DPS init). Using the init rather than the raw `iter_arg` + # keeps the old value readable *after* the inner reduction has run: the + # init is an immutable SSA value, so bufferization copies the accumulator + # instead of letting the reduction write in place. + accumulator_values = {} + for operand, result_idx in zip( + accumulator_operands, accumulator_result_indices + ): + inner = result_to_inner[result_idx] + new_value = value_map[inner.results[0]] + old_value = value_map.get( + ls.dps_init_operands(inner)[0].value, + ls.dps_init_operands(inner)[0].value, + ) + accumulator_values[operand] = (new_value, old_value) + + e_inputs = [] + for operand in ls.dps_input_operands(fused_source): + if operand in accumulator_values: + # A running accumulator: broadcast over the reduction axis, so it + # is read whole rather than sliced. + e_inputs.append(accumulator_values[operand][0]) + else: + e_inputs.append( + _tile_slice( + operand.value, + ls.indexing_map_for(fused_source, operand), + e_tiled_dim, + iv, + tile_size, + ) + ) + e_out_map = ls.indexing_map_for(fused_source, e_init_operand) + e_dest_tile = _tile_slice(e_arg, e_out_map, e_tiled_dim, iv, tile_size) + fused_e = _clone_generic_with_operands( + fused_source, e_inputs, [e_dest_tile], [e_dest_tile.type] + ) + + # --- the online correction on R2's running accumulator --- + r2_out_map = ls.indexing_map_for(r2, r2_init_operand) + r2_acc_tile = _tile_slice(r2_arg, r2_out_map, r2_red_dim, iv, tile_size) + accumulators = [ + (operand, new_value, old_value) + for operand, (new_value, old_value) in accumulator_values.items() + ] + if not accumulators: + raise FusionRejected( + "fused E does not consume a running accumulator of the loop" + ) + factor = _correction_factor(fused_source, r2, accumulators, r2_acc_tile) + if factor is None: + raise FusionRejected("could not build the correction factor from E") + scaled = _emit_elementwise( + linalg.ElementwiseKind.mul, + lambda a, b: arith.mulf(a, b), + [r2_acc_tile, factor], + r2_acc_tile, + ) + + # --- the fused R2, accumulating this tile into the rescaled sum --- + e_result = fused_source.results[0] + r2_inputs = [] + for operand in ls.dps_input_operands(r2): + if operand.value == e_result: + r2_inputs.append(fused_e.results[0]) + else: + r2_inputs.append( + _tile_slice( + operand.value, + ls.indexing_map_for(r2, operand), + r2_red_dim, + iv, + tile_size, + ) + ) + fused_r2 = _clone_generic_with_operands( + r2, r2_inputs, [scaled], [r2.results[0].type] + ) + + # --- the new yield --- + yielded = [value_map[o] for o in old_body_ops[-1].operands] + yielded.append( + _insert_tile_slice( + fused_e.results[0], e_arg, e_out_map, e_tiled_dim, iv, tile_size + ) + ) + yielded.append( + _insert_tile_slice( + fused_r2.results[0], r2_arg, r2_out_map, r2_red_dim, iv, tile_size + ) + ) + scf.YieldOp(yielded) + + # --- retire the originals ------------------------------------------------- + # R2 first, then E: erasing R2 drops the only use of the fused E's full-extent + # result, leaving the loop result that replaces it dead (as it must be -- it + # holds a stale term in every tile but the last). + rewriter.replace_op(r2, [new_loop.results[num_old_results + 1]]) + rewriter.replace_op(fused_source, [new_loop.results[num_old_results]]) + rewriter.replace_op(r1_loop, list(new_loop.results)[:num_old_results]) + return new_loop diff --git a/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_legality.py b/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_legality.py new file mode 100644 index 00000000..a84a6d31 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/dependant_reduction_legality.py @@ -0,0 +1,674 @@ +"""Legality analysis for dependent-reduction (flash-attention style) fusion. + +Decides whether an ``R1 -> E -> R2`` chain can be fused into ``R1``'s +already-tiled reduction loop, where + + * ``R1`` is a reduction already tiled along its reduction axis into an + ``scf.for`` marked with ``__reduction_loop__``, + * ``E`` is an all-parallel elementwise op consuming ``R1``'s result as a + broadcast input plus the data inputs it shares with ``R1``, + * ``R2`` is a reduction over ``E``'s result along the same axis. + +The canonical case is two-pass softmax (``max`` then ``sum(exp)``) becoming the +online one-pass form; an attention chain is the same shape with a `P @ V` +contraction as a second ``R2``. + +Rejections raise `FusionRejected` carrying the reason, which the transform op +surfaces as a silenceable failure. +""" + +from mlir import ir +from mlir.dialects import arith, linalg, math, tensor + +from lighthouse.utils.mlir import defining_op, op_users, opview +from lighthouse.dialects.transform.transform_ext.utils import ir_rewrite as irr +from lighthouse.dialects.transform.transform_ext.utils import linalg_structured as ls + +__all__ = [ + "REDUCTION_LOOP_ATTR_NAME", + "FusionRejected", + "check_legal_fusion_triple", + "collect_inner_reduction_generics", + "collect_r1_as_elementwise_inputs", + "find_r2_elementwise_operand", + "map_loop_results_to_inner_reductions", + "needs_elementwise_clone", +] + +#: Unit attribute marking an ``scf.for`` as a tiled reduction loop. A plain +#: ``scf.for`` carries no iterator-type metadata, so this marker is the only way +#: to tell that the loop iterates a reduction axis. The producer of the IR +#: (e.g. ``transform.structured.tile_using_for`` followed by +#: ``transform.annotate``) is responsible for tagging it. +REDUCTION_LOOP_ATTR_NAME = "__reduction_loop__" + +#: Element types the online correction is defined for. +_SUPPORTED_FLOAT_TYPES = (ir.F16Type, ir.BF16Type, ir.F32Type, ir.F64Type) + + +class FusionRejected(Exception): + """The chain does not satisfy the fusion legality conditions.""" + + +# --- chain structure --------------------------------------------------------- + + +def collect_r1_as_elementwise_inputs( + r1_loop: ir.OpView, e: ir.OpView +) -> tuple[list[ls.Operand], list[int]]: + """The `e` inputs consuming a result of the ``R1`` loop, with those result indices. + + Pre-fusion, the elementwise term consumes the loop's ``iter_arg``-carried + results; the legality check needs both the operands (to verify their indexing + maps) and the result index (to recover the inner reduction producing each). + """ + operands: list[ls.Operand] = [] + result_indices: list[int] = [] + loop_results = list(r1_loop.results) + for operand in ls.dps_input_operands(e): + for i, result in enumerate(loop_results): + if operand.value == result: + operands.append(operand) + result_indices.append(i) + break + return operands, result_indices + + +def collect_inner_reduction_generics(loop: ir.OpView) -> list[ir.OpView]: + """The reduction ``linalg.generic``s in the loop body, in program order. + + A loop may carry several running accumulators (e.g. a fused softmax + ``(max, sum)`` loop), each produced by its own inner reduction. + """ + result = [] + for op in loop.body.operations: + ov = opview(op) + if isinstance(ov, linalg.GenericOp) and ls.num_reduction_loops(ov) != 0: + result.append(ov) + return result + + +def map_loop_results_to_inner_reductions(loop: ir.OpView) -> list[ir.OpView | None]: + """Map each loop result to the inner reduction generic producing its tile. + + Indexed by loop result number; an entry is None when the result is not + yielded from a reduction generic's tile (e.g. a passthrough result). + + Each running accumulator is yielded as ``tensor.insert_slice %red into %arg``, + so each ``scf.yield`` operand is inspected: an ``insert_slice`` whose source + is defined by a reduction generic identifies that generic. + """ + result: list[ir.OpView | None] = [None] * len(list(loop.results)) + terminator = loop.body.operations[len(loop.body.operations) - 1] + for idx, yielded in enumerate(terminator.operands): + insert_op = defining_op(yielded) + if insert_op is None or not isinstance(insert_op.opview, tensor.InsertSliceOp): + continue + source_op = defining_op(insert_op.opview.source) + if source_op is None: + continue + ov = source_op.opview + if isinstance(ov, linalg.GenericOp) and ls.num_reduction_loops(ov) != 0: + result[idx] = ov + return result + + +def find_r2_elementwise_operand(r2: ir.OpView, e: ir.OpView) -> ls.Operand: + """The single `r2` input reading `e`'s result. + + ``R2`` may have several inputs (a GEMM-like contraction does), but exactly one + of them must be ``E``'s result. + """ + found = None + e_result = e.results[0] + for operand in ls.dps_input_operands(r2): + if operand.value != e_result: + continue + if found is not None: + raise FusionRejected("R2 consumes E's result more than once") + found = operand + if found is None: + raise FusionRejected("no R2 input is E's result") + return found + + +def needs_elementwise_clone(e: ir.OpView, r2: ir.OpView) -> bool: + """Whether the fusion must work on a *clone* of `e` rather than `e` itself. + + True if ``E`` has any user besides ``R2``. Fusing ``E`` computes each tile + against the *running* ``R1`` accumulator instead of its final result, so + ``E``'s own output is stale in every tile but the last and no user outside the + loop may read it. The original is left in place for those users, where it still + reads ``R1``'s final result and recomputes the term correctly. + """ + r2_op = r2.operation + return any(user != r2_op for user in op_users(e.results[0])) + + +# --- separability analysis ------------------------------------------------------------ + +#: `v` is independent of every block argument (a loop-invariant scalar). +_CONST = 1 << 0 +#: `v = h(x)`: no dependence on the accumulators. +_IND = 1 << 1 +#: `v = alpha(m) + h(x)`: the accumulators enter additively. +_ADD = 1 << 2 +#: `v = g(m) * h(x)`: the accumulators enter multiplicatively. +_MUL = 1 << 3 + + +def _close_facts(facts: int) -> int: + """Close a fact set under the implications between the facts. + + A value the accumulators never reach trivially fits both separable shapes, + with a trivial accumulator part (``h(x) = 0 + h(x)`` and ``h(x) = 1 * h(x)``). + """ + if facts & _CONST: + facts |= _IND + if facts & _IND: + facts |= _ADD | _MUL + return facts + + +def check_elementwise_separability( + e: ir.OpView, accumulator_args: list[ir.BlockArgument] +) -> None: + """Verify `e` is *multiplicatively separable* in the accumulators it consumes. + + That is, that its body computes ``E(x, m) = f(x) * g(m)``, where ``m`` are the + values read through `accumulator_args` and ``x`` is everything else. This is + exactly the condition under which the online correction is well defined: the + factor evaluates ``E`` twice, at the new and the old accumulator, with a + stand-in constant for the data operands, and separability is what makes ``f`` + cancel -- + + E(v, m_new) / E(v, m_old) = g(m_new) / g(m_old) for every `v` + + so a single scalar per parallel slice repairs ``R2``'s accumulator. Without it + the ratio depends on the stand-in and no single scalar is correct: for + ``E = (x - m)^2`` with a tile holding ``x = 1`` and ``m`` moving from ``1`` to + ``3``, the true sum is ``4`` while any rescale of the stale ``0`` yields ``0``. + + Implemented as an abstract interpretation over the fact sets above, in one + forward pass -- ``E``'s body is straight-line. Accumulator block arguments are + seeded ``_ADD | _MUL`` (a bare ``m`` is both ``m + 0`` and ``m * 1``), other + block arguments ``_IND``, values from an enclosing scope ``_CONST``. The analysis + is one-sided, as an abstract interpretation must be: ``_MUL`` on the yielded value + proves separability, its absence only fails to. Accepting additionally requires + the yield *not* to hold ``_IND``, which would mean it never reads an accumulator + at all -- a constant-``1`` correction, i.e. not a dependent reduction. + + Tracking both shapes rather than just ``_MUL`` is what ``exp`` needs: + ``exp(alpha(m) + h(x)) = exp(alpha(m)) * exp(h(x))`` converts an additive + dependence into a multiplicative one, which is how softmax's ``exp(x - m)`` + reaches ``_MUL`` at all. Variance's ``(x - m)^2`` does not: ``mulf`` needs both + operands ``_MUL`` and ``x - m`` carries only ``_ADD``. + """ + body = e.regions[0].blocks[0] + facts: dict = {} + + # Seed the block arguments. All tracked accumulators are seeded at once, so + # the pass proves *joint* separability `E(x, m1, m2) = f(x) * g(m1, m2)` -- + # the correction substitutes new/old for all of them simultaneously. + accumulators = set(accumulator_args) + for barg in body.arguments: + facts[barg] = _close_facts((_ADD | _MUL) if barg in accumulators else _IND) + + def facts_of(value: ir.Value) -> int: + # Values from an enclosing scope are invariant over E's iteration space. + return facts.get(value, _close_facts(_CONST)) + + def has(value: ir.Value, fact: int) -> bool: + return bool(facts_of(value) & fact) + + ops = list(body.operations) + for op in ops[:-1]: + ov = opview(op) + # Only single-result scalar arithmetic is modelled; anything else (a + # comparison, a select, a call) is opaque, so its results are recorded with + # no facts, which rejects the chain unless they are dead. Recorded rather + # than skipped: an absent value is read as loop-invariant by `facts_of`, + # which for an unmodelled result would be a wrong (and unsound) answer. + if len(ov.results) != 1: + for opaque in ov.results: + facts[opaque] = 0 + continue + result = ov.results[0] + f = 0 + + def binary(preserved: int) -> int: + bits = 0 + lhs, rhs = ov.operands[0], ov.operands[1] + if has(lhs, _CONST) and has(rhs, _CONST): + bits |= _CONST + if has(lhs, _IND) and has(rhs, _IND): + bits |= _IND + # `(a1 + a2)(m) + (h1 + h2)(x)` for the additive family, and + # `(g1 g2)(m) * (h1 h2)(x)` for the multiplicative one. + if has(lhs, preserved) and has(rhs, preserved): + bits |= preserved + return bits + + def unary(from_fact: int, to_fact: int) -> int: + # `_CONST`/`_IND` always survive a unary op. + arg = ov.operands[0] + bits = facts_of(arg) & (_CONST | _IND) + if has(arg, from_fact): + bits |= to_fact + return bits + + if isinstance(ov, arith.ConstantOp): + # Invariant over E's iteration space, exactly like a value from an + # enclosing scope. Without this every value downstream of an in-body + # constant would carry no facts, rejecting a naturally written scaled + # term like `exp(x * cst - m)`. + f = _CONST + elif isinstance(ov, (arith.AddFOp, arith.SubFOp)): + f = binary(_ADD) + elif isinstance(ov, (arith.MulFOp, arith.DivFOp)): + f = binary(_MUL) + elif isinstance(ov, arith.NegFOp): + # `-(alpha + h) = (-alpha) + (-h)` and `-(g * h) = g * (-h)`. + f = facts_of(ov.operands[0]) + elif isinstance(ov, (math.ExpOp, math.Exp2Op)): + f = unary(_ADD, _MUL) + elif isinstance(ov, (math.LogOp, math.Log2Op)): + # `log(g * h) = log(g) + log(h)`, the inverse bridge. + f = unary(_MUL, _ADD) + elif isinstance(ov, (math.AbsFOp, math.SqrtOp, math.RsqrtOp)): + # `|g * h| = |g| * |h|`, and likewise for the (r)sqrt of a product. + f = unary(_MUL, _MUL) + elif isinstance(ov, math.PowFOp): + base, exponent = ov.operands[0], ov.operands[1] + f = facts_of(base) & facts_of(exponent) & (_CONST | _IND) + # `(g * h)^p = g^p * h^p` needs a *fixed* `p`: were the exponent to + # vary with the data, `g(m)^p(x)` would still depend on `x`. + if has(base, _MUL) and has(exponent, _CONST): + f |= _MUL + # `c^(alpha + h) = c^alpha * c^h`. + if has(base, _CONST) and has(exponent, _ADD): + f |= _MUL + + facts[result] = _close_facts(f) + + terminator = ops[-1] + if len(terminator.operands) != 1: + raise FusionRejected("E does not yield exactly one value") + term = terminator.operands[0] + if not has(term, _MUL): + raise FusionRejected( + "E is not multiplicatively separable in the accumulators it consumes, " + "so no per-slice scalar can correct R2's running accumulator" + ) + # A yield that never reads an accumulator would make the correction the + # constant 1; the chain is then not a dependent reduction at all. + if has(term, _IND): + raise FusionRejected("E does not depend on any consumed accumulator") + + +# --- per-link checks --------------------------------------------------------- + + +def find_elementwise_dim_for_r2_reduction_dim( + e: ir.OpView, r2: ir.OpView, r2_e_operand: ls.Operand, r2_red_dim: int +) -> int: + """The `e` loop dim carrying `r2`'s reduction axis. + + ``E``'s output map and ``R2``'s input map both describe the same tensor + (``E``'s result), so they align position-by-position: at tensor dim ``i``, + ``R2``'s map yields an ``R2`` loop dim and ``E``'s output map yields an ``E`` + loop dim. The ``E`` dim sitting where ``R2`` reads its reduction dim is the + axis along which ``E`` must be re-sliced when fused into ``R1``'s tiled loop. + """ + r2_map = ls.indexing_map_for(r2, r2_e_operand) + e_out_map = ls.indexing_map_for(e, ls.dps_init_operands(e)[0]) + if len(r2_map.results) != len(e_out_map.results): + raise FusionRejected( + f"R2's map for E's result and E's output map have different rank " + f"(R2: {r2_map}, E out: {e_out_map})" + ) + for r2_expr, e_expr in zip(r2_map.results, e_out_map.results): + if not isinstance(r2_expr, ir.AffineDimExpr) or not isinstance( + e_expr, ir.AffineDimExpr + ): + raise FusionRejected( + f"a non-dim affine expr in R2's map for E's result or in E's " + f"output map (R2: {r2_map}, E out: {e_out_map})" + ) + if r2_expr.position == r2_red_dim: + return e_expr.position + raise FusionRejected( + f"R2's reduction dim d{r2_red_dim} does not appear in its map for E's " + f"result: {r2_map}" + ) + + +def check_inner_reduction_against_elementwise( + r1: ir.OpView, e: ir.OpView, e_tiled_dim: int, inner_results: set +) -> None: + """Check one inner reduction `r1` of the loop can be fused against `e`. + + ``r1`` must reduce over exactly one innermost loop, and every ``r1`` input + (resolved *through* its tile ``extract_slice`` to the source tensor) must also + appear as an ``e`` input. + + Fusion rewrites ``r1``'s indexing maps into ``e``'s iteration space, so the two + ops' loop dims have to correspond. That correspondence ``phi`` is not given but + *derived*, by unifying the maps of the inputs they share: at each position of a + shared input's map, the ``r1`` dim indexing that tensor dim pairs with the ``e`` + dim indexing it, and the pairing must be consistent across all shared inputs. + For softmax's ``max`` and ``exp(x - m)``, both read ``x`` as + ``(d0, d1) -> (d0, d1)``, so ``phi = {0: 0, 1: 1}``. + + ``phi`` must additionally be + + * total over ``r1``'s loop dims, so every ``r1`` map -- its init map included + -- can be translated, and + * reduction-aligned: ``phi[r1_red_dim] == e_tiled_dim``, i.e. ``r1`` reduces + the same axis that ``R2`` reduces through ``e``. + + Note the comparison is against ``E``, not ``R2``: with the elementwise term + left unfused, ``E`` is the op sharing ``R1``'s data inputs (e.g. ``x`` in + ``exp(x - m)``), while ``R2`` reduces ``E``'s result and need not share any + input with ``R1``. + """ + r1_red_dims = ls.reduction_dims(r1) + if len(r1_red_dims) != 1: + raise FusionRejected( + f"inner R1 does not have exactly one reduction iterator " + f"({len(r1_red_dims)})" + ) + if r1_red_dims[0] != ls.num_loops(r1) - 1: + raise FusionRejected("reduction iterator is not the innermost loop in inner R1") + + # Every input of R1 must also appear as an input of E. Unifying the two ops' + # maps for each shared input derives phi: R1 loop dim -> E loop dim. + phi: dict[int, int] = {} + + def try_add_mapping(r1_dim: int, e_dim: int) -> bool: + if r1_dim not in phi: + phi[r1_dim] = e_dim + return True + return phi[r1_dim] == e_dim + + e_inputs = ls.dps_input_operands(e) + for in1 in ls.dps_input_operands(r1): + # The inner R1 reads a `tensor.extract_slice` of the real input tensor; + # resolve through the slice to compare against E's (untiled) inputs. + in1_source = irr.resolve_slice_source(in1.value) + # Inputs that are results of a sibling inner reduction (e.g. the sum + # reduction reading the max result) are running accumulators broadcast + # over the reduction axis; they need not appear in E, so skip them. + if in1_source in inner_results: + continue + in_e = None + for candidate in e_inputs: + if irr.resolve_slice_source(candidate.value) == in1_source: + in_e = candidate + break + if in_e is None: + raise FusionRejected(f"R1 input is not also an input of E: {in1_source}") + m1 = ls.indexing_map_for(r1, in1) + m_e = ls.indexing_map_for(e, in_e) + if len(m1.results) != len(m_e.results): + raise FusionRejected( + f"shared input has maps of different rank in R1 vs E " + f"(R1: {m1}, E: {m_e})" + ) + for e1, e2 in zip(m1.results, m_e.results): + if not isinstance(e1, ir.AffineDimExpr) or not isinstance( + e2, ir.AffineDimExpr + ): + raise FusionRejected( + f"shared input map has a non-dim affine expr (R1: {m1}, E: {m_e})" + ) + if not try_add_mapping(e1.position, e2.position): + raise FusionRejected( + f"inconsistent dim mapping between R1 and E derived from " + f"shared inputs (R1.d{e1.position} -> " + f"{{E.d{phi[e1.position]}, E.d{e2.position}}})" + ) + + # phi must be total over R1's loop dims so R1's init map (and any other R1 + # map) can be translated into E's iteration space during fusion. + if len(phi) != ls.num_loops(r1): + raise FusionRejected( + f"derived dim mapping does not cover all of R1's loop dims " + f"(covered {len(phi)} of {ls.num_loops(r1)})" + ) + if phi.get(r1_red_dims[0]) != e_tiled_dim: + raise FusionRejected( + "R1's reduction dim is not aligned with the E dim carrying R2's " + "reduction axis under the derived dim mapping" + ) + + +# --- the triple -------------------------------------------------------------- + + +def check_legal_fusion_triple( + r1_loop: ir.OpView, + result_to_inner: list[ir.OpView | None], + e: ir.OpView, + r2: ir.OpView, +) -> tuple[int, int]: + """Verify the ``R1 -> E -> R2`` chain can be fused into `r1_loop`. + + `r1_loop` is the already-tiled producer reduction: an ``scf.for`` carrying + running accumulators as ``iter_arg``s whose results `e` consumes. + `result_to_inner` maps each loop result index to the inner per-tile reduction + producing it (None for non-reduction results). The loop's ``step`` is the tile + size and ``ub - lb`` the full reduction extent. + + Returns ``(e_tiled_dim, tile_size)``: the `e` loop dim carrying ``R2``'s + reduction axis, which the retiling step needs, and the tile size. + + Raises `FusionRejected` with the reason on any violated condition. + TODO: These conditions are too strict, reconsider and relax later. + """ + # A structured op on tensors has one init per result (`num_dps_inits` is derived + # from the result count), so one condition covers both. + if ls.num_dps_inits(r2) != 1: + raise FusionRejected( + f"R2 does not have exactly one result/init ({ls.num_dps_inits(r2)})" + ) + + # (E1) E must be all-parallel with exactly one result/init: it is the + # elementwise term feeding R2, not a reduction of its own. + if ls.num_dps_inits(e) != 1: + raise FusionRejected( + f"E does not have exactly one result/init ({ls.num_dps_inits(e)})" + ) + if ls.num_reduction_loops(e) != 0: + raise FusionRejected( + f"E is not all-parallel (it has {ls.num_reduction_loops(e)} " + f"reduction loops)" + ) + + # (E2) The three ops must share a block: the rewrite builds the replacement loop + # at E's position and reasons about program order there, and both dominance + # helpers answer only within one block. + if e.operation.block != r2.operation.block or ( + e.operation.block != r1_loop.operation.block + ): + raise FusionRejected("R1 loop, E and R2 are not all in the same block") + + # (E3) R2 may have MULTIPLE inputs (e.g. a GEMM-like contraction), but exactly + # one must be E's result. + r2_e_operand = find_r2_elementwise_operand(r2, e) + + # (E4) R2 must have exactly one reduction loop, and it must be its innermost. + r2_red_dims = ls.reduction_dims(r2) + if len(r2_red_dims) != 1: + raise FusionRejected( + f"R2 does not have exactly one reduction iterator ({len(r2_red_dims)})" + ) + if r2_red_dims[0] != ls.num_loops(r2) - 1: + raise FusionRejected("reduction iterator is not the innermost loop in R2") + + # (E5) Every R2 input must be reduced along R2's reduction axis, i.e. its + # indexing map must reference `r2_red_dim`. An input that does *not* is + # broadcast across the reduction (a per-parallel-slice value); such an operand + # would have to be re-derived per tile rather than simply re-sliced, which the + # rewrite does not model. Requiring all inputs to carry the axis is what makes + # "re-slice every R2 input to the current tile" a complete description. + for operand in ls.dps_input_operands(r2): + imap = ls.indexing_map_for(r2, operand) + carries_red_dim = False + for expr in imap.results: + if not isinstance(expr, ir.AffineDimExpr): + raise FusionRejected( + f"R2 input indexing map has a non-dim affine expr: {imap}" + ) + if expr.position == r2_red_dims[0]: + carries_red_dim = True + if not carries_red_dim: + raise FusionRejected( + f"R2 input is not reduced along R2's reduction axis (map does " + f"not reference dim {r2_red_dims[0]}): {imap}" + ) + + # (E6) Locate the E loop dim carrying R2's reduction axis. This is the axis + # along which E is re-sliced when cloned into R1's tiled loop, and the axis + # R1's reduction dim must align with. + e_tiled_dim = find_elementwise_dim_for_r2_reduction_dim( + e, r2, r2_e_operand, r2_red_dims[0] + ) + + # The producer loop must have static, constant bounds: the tile size is its + # `step` and the full reduction extent is `ub - lb`. That extent must equal + # R2's (static) reduction extent, and the tile size must evenly divide it. + lb = irr.constant_int_value(r1_loop.lowerBound) + ub = irr.constant_int_value(r1_loop.upperBound) + step = irr.constant_int_value(r1_loop.step) + if lb is None or ub is None or step is None or step <= 0: + raise FusionRejected( + "R1 reduction loop does not have constant, positive bounds/step" + ) + full_extent = ub - lb + tile_size = step + r2_red_range = ls.static_loop_ranges(r2)[r2_red_dims[0]] + if ir.ShapedType.is_dynamic_size(r2_red_range): + raise FusionRejected( + "R2 reduction range is dynamic; fusion requires a static reduction extent" + ) + if r2_red_range != full_extent: + raise FusionRejected( + f"R2's reduction extent ({r2_red_range}) differs from the R1 loop " + f"extent ({full_extent})" + ) + if full_extent % tile_size != 0: + raise FusionRejected( + f"tile size {tile_size} does not evenly divide the reduction extent " + f"{full_extent}" + ) + + # (E7) E's extent along the axis carrying R2's reduction must match the R1 + # loop extent too, so re-slicing E to the tile is well defined. + e_red_range = ls.static_loop_ranges(e)[e_tiled_dim] + if ir.ShapedType.is_dynamic_size(e_red_range) or e_red_range != full_extent: + raise FusionRejected( + f"E's extent along the axis carrying R2's reduction ({e_red_range}) " + f"differs from the R1 loop extent ({full_extent})" + ) + + # The E operands consuming a loop result. There must be at least one: this is + # the R1 -> E data dependence that makes the chain fusable. + r1_as_e_operands, r1_result_indices = collect_r1_as_elementwise_inputs(r1_loop, e) + if not r1_as_e_operands: + raise FusionRejected("E does not consume any result of the R1 loop") + + # (E8) E must be multiplicatively separable in the accumulators it reads, or + # the online correction it is asked to derive does not exist. + accumulator_args = [ + ls.matching_block_argument(e, operand) for operand in r1_as_e_operands + ] + check_elementwise_separability(e, accumulator_args) + + # The set of all inner reduction results in the loop, used to skip + # sibling-reduction inputs in the per-inner-reduction phi check. + inner_results = {inner.results[0] for inner in result_to_inner if inner is not None} + + # Each R1 result consumed by E must be broadcast across the E axis carrying + # R2's reduction -- i.e. its indexing map must not reference `e_tiled_dim`. + # Conservatively also require a pure dim-expr projection. This is what makes + # the running accumulator a per-parallel-slice scalar the correction can + # rescale. + for operand in r1_as_e_operands: + imap = ls.indexing_map_for(e, operand) + for expr in imap.results: + if not isinstance(expr, ir.AffineDimExpr): + raise FusionRejected( + f"R1-as-E-input indexing map has a non-dim affine expr: {imap}" + ) + if expr.position == e_tiled_dim: + raise FusionRejected( + f"R1's result is not broadcast across the E axis carrying " + f"R2's reduction (map references dim {expr.position}): {imap}" + ) + + # For each consumed loop result, recover the inner reduction producing it and + # verify the reduction-dim alignment and shared-input phi mapping against E. + for operand, result_idx in zip(r1_as_e_operands, r1_result_indices): + inner = result_to_inner[result_idx] + if inner is None: + raise FusionRejected( + f"consumed loop result {result_idx} is not produced by an inner " + f"reduction generic" + ) + check_inner_reduction_against_elementwise(inner, e, e_tiled_dim, inner_results) + + # (E9) R2's region must be a single-combiner sum reduction. + _, combiners = irr.match_reduction(ls.region_output_args(r2), 0) + if not combiners: + raise FusionRejected("R2's region does not match a reduction pattern") + if len(combiners) != 1: + raise FusionRejected( + f"R2's reduction has {len(combiners)} combiners, expected exactly 1" + ) + if not isinstance(combiners[0].opview, arith.AddFOp): + raise FusionRejected(f"R2's combiner is not arith.addf: {combiners[0].name}") + + # (E10) R2's init must be the additive identity (zero), produced directly or + # through a linalg.fill of zero into an empty tensor. + r2_init = ls.dps_init_operands(r2)[0].value + if not irr.is_defined_as_zero(r2_init): + raise FusionRejected("R2's init is not the additive identity (zero)") + + # (E11) Restrict to supported floating-point element types. + element_type = ir.ShapedType(r2.results[0].type).element_type + if not isinstance(element_type, _SUPPORTED_FLOAT_TYPES): + raise FusionRejected( + f"R2's element type {element_type} is not a supported " + f"floating-point type (f16/bf16/f32/f64)" + ) + + # (E12) E's element type and R2's accumulator must have a common widening: the + # correction term is E's body re-evaluated in the wider of the two, with casts on + # the way in and out. `f16` and `bf16` have no single-step conversion between + # them, so such a pair is rejected here rather than mid-rewrite. + e_element_type = ir.ShapedType(e.results[0].type).element_type + if irr.wider_float_type(e_element_type, element_type) is None: + raise FusionRejected( + f"E's element type {e_element_type} and R2's accumulator type " + f"{element_type} have no common widening to evaluate the correction " + f"term in" + ) + + # (E13) Any other user of any R1 loop result (besides E) must post-dominate E + # so re-routing the values through the fused op is safe. Note the anchor is E, + # not R2: E directly consumes the R1 loop results and is cloned into the loop + # first. + for r1_result in r1_loop.results: + for user in op_users(r1_result): + if user == e.operation: + continue + if not irr.post_dominates(user, e): + raise FusionRejected( + f"user of an R1 result does not post-dominate E: {user.name}" + ) + + # Note there is no condition on the users of E's result: any user besides R2 + # makes the fusion work on a clone and leave the original E -- and therefore + # those users -- exactly where they are. See `needs_elementwise_clone`. + return e_tiled_dim, tile_size diff --git a/lighthouse/dialects/transform/transform_ext/utils/ir_rewrite.py b/lighthouse/dialects/transform/transform_ext/utils/ir_rewrite.py new file mode 100644 index 00000000..4622d068 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/ir_rewrite.py @@ -0,0 +1,548 @@ +"""Generic IR helpers the bindings do not provide. + +Replacements for the C++ utilities the reduction fusion relies on: ``IRMapping`` +plus a mapping-aware ``clone``, ``DominanceInfo``/``PostDominanceInfo``, +``getBackwardSlice``, ``AffineMap`` dim rewriting, and the zero/neutral-element +predicates from ``FoldAddIntoDest``. +""" + +from mlir import ir +from mlir.dialects import arith, linalg, tensor + +from lighthouse.utils.mlir import defining_op, opview + +__all__ = [ + "backward_slice", + "cast_float", + "clone_block_body", + "clone_op_deep_with_map", + "clone_op_with_map", + "constant_int_value", + "depends_on_op", + "float_width", + "is_defined_as_zero", + "match_reduction", + "mixed_sizes", + "move_value_definitions", + "op_attributes", + "operand_eliminating_constant", + "post_dominates", + "project_dims", + "properly_dominates", + "remap_dims", + "resolve_slice_source", + "wider_float_type", +] + + +# --- cloning ----------------------------------------------------------------- + + +def op_attributes(op: ir.Operation | ir.OpView) -> dict[str, ir.Attribute]: + """The op's discardable + inherent attributes as a name -> attr dict.""" + attrs = opview(op).operation.attributes + return {attrs[i].name: attrs[i].attr for i in range(len(attrs))} + + +def clone_op_with_map( + op: ir.Operation | ir.OpView, + value_map: dict, + *, + result_type: ir.Type | None = None, +) -> ir.Operation | None: + """Clone `op` at the current insertion point, remapping operands via `value_map`. + + Stands in for ``OpBuilder::clone(op, IRMapping)``, which the bindings do not + expose (``Operation.clone`` takes only an insertion point). Regions are *not* + copied, so this is limited to the region-free scalar ops that make up a + linalg body -- which is all the fusion clones one op at a time. Results are + recorded into `value_map`, so cloning a block in order threads the + substitution through. + + `result_type` retypes the clone's results, which is how a body is re-emitted in + a different precision; an ``arith.constant``'s value attribute is rebuilt to + match. It only makes sense for the single-type scalar float ops a linalg body is + made of -- the caller is responsible for having remapped the operands to that + same type. + """ + ov = opview(op) + if any(len(r.blocks) for r in ov.operation.regions): + return None + attributes = op_attributes(ov) + result_types = [r.type for r in ov.results] + if result_type is not None: + result_types = [result_type] * len(result_types) + value = attributes.get("value") + if value is not None and ir.FloatAttr.isinstance(value): + attributes["value"] = ir.FloatAttr.get( + result_type, ir.FloatAttr(value).value + ) + cloned = ir.Operation.create( + ov.operation.name, + results=result_types, + operands=[value_map.get(o, o) for o in ov.operands], + attributes=attributes, + ) + value_map.update(zip(ov.results, cloned.results)) + return cloned + + +def clone_op_deep_with_map( + op: ir.Operation | ir.OpView, value_map: dict +) -> ir.Operation: + """Deep-clone `op` (regions included) at the current insertion point. + + Operands are remapped through `value_map`, both on the op itself and on any + value its regions capture from an enclosing scope. Results are recorded into + `value_map`, so cloning a block's ops in order threads the substitution + through -- which is how the reduction loop's body is copied into its + replacement. + + Unlike `clone_op_with_map` this handles ops carrying regions (the inner + reduction ``linalg.generic``s), at the cost of going through + ``Operation.clone`` and patching operands afterwards. + """ + ov = opview(op) + cloned = ov.operation.clone() + for i, operand in enumerate(ov.operands): + if operand in value_map: + cloned.operands[i] = value_map[operand] + + def remap_captured(inner: ir.Operation) -> ir.WalkResult: + for i, operand in enumerate(inner.operands): + if operand in value_map: + inner.operands[i] = value_map[operand] + return ir.WalkResult.ADVANCE + + cloned.walk(remap_captured) + value_map.update(zip(ov.results, cloned.results)) + return cloned + + +def move_value_definitions( + values: list[ir.Value], before_op: ir.Operation | ir.OpView +) -> bool: + """Move the definitions of `values` to just before `before_op`. + + Stands in for ``mlir::moveValueDefinitions``. Collects the backward slice of + each value, keeps the ops that do not already dominate `before_op`, and moves + them in program order so their relative order -- and therefore their + def-before-use -- is preserved. Returns False if any op to move lives in a + different block, which this cannot safely relocate. + """ + anchor = opview(before_op).operation + block = anchor.block + to_move: dict = {} + stack = [defining_op(v) for v in values] + while stack: + cur = stack.pop() + if cur is None or cur.__hash__() in to_move: + continue + if properly_dominates(cur, anchor): + continue + if cur.block != block: + return False + to_move[cur.__hash__()] = cur + for operand in cur.operands: + producer = defining_op(operand) + if producer is not None: + stack.append(producer) + + if not to_move: + return True + # Program order within the block, so relative order survives the move. + ordered = [op for op in block.operations if op.operation.__hash__() in to_move] + for op in ordered: + op.operation.move_before(anchor) + return True + + +def clone_block_body( + src_block: ir.Block, + arg_values: list[ir.Value], + *, + skip_terminator: bool = True, + value_map: dict | None = None, +) -> dict: + """Clone `src_block`'s ops at the current insertion point. + + `arg_values` binds the source block arguments positionally; entries may be + None to leave an argument unbound (the caller then pre-seeds `value_map`, or + the clone of a consuming op substitutes for it). Returns the value map, so + the caller can look up the clone of any source value. + """ + vmap = {} if value_map is None else value_map + for arg, val in zip(src_block.arguments, arg_values): + if val is not None: + vmap[arg] = val + ops = list(src_block.operations) + if skip_terminator: + ops = ops[:-1] + for op in ops: + clone_op_with_map(op, vmap) + return vmap + + +# --- ordering / dominance ---------------------------------------------------- + + +def _ancestor_in_block(op: ir.Operation | ir.OpView, block: ir.Block): + """The ancestor of `op` that sits directly in `block`, or None.""" + cur = opview(op).operation + while cur is not None: + parent_block = cur.block + if parent_block is None: + return None + if parent_block == block: + return cur + owner = parent_block.owner + cur = owner.operation if owner is not None else None + return None + + +def properly_dominates( + a: ir.Operation | ir.OpView, b: ir.Operation | ir.OpView +) -> bool: + """Whether `a` properly dominates `b`, for ops related through one block. + + The fusion legality pins the chain to a single block, so dominance reduces to + program order there: `a` dominates `b` when the ancestor of `a` in that block + precedes the ancestor of `b`. An op enclosing `b` dominates it. Returns False + when no common block is found, which is the conservative answer for every + caller here. + """ + a_op, b_op = opview(a).operation, opview(b).operation + if a_op == b_op: + return False + block = a_op.block + if block is None: + return False + b_anchor = _ancestor_in_block(b_op, block) + if b_anchor is None: + return False + if b_anchor == a_op: + # `a` encloses `b`. + return True + return a_op.is_before_in_block(b_anchor) + + +def post_dominates(a: ir.Operation | ir.OpView, b: ir.Operation | ir.OpView) -> bool: + """Whether `a` post-dominates `b`, for ops related through one block. + + Within a single block with no control flow, post-dominance is the reverse of + program order: `a` post-dominates `b` when `a` comes at or after `b`. Returns + False when no common block is found (the conservative answer). + """ + a_op, b_op = opview(a).operation, opview(b).operation + if a_op == b_op: + return True + block = a_op.block + if block is None: + return False + b_anchor = _ancestor_in_block(b_op, block) + if b_anchor is None: + return False + if b_anchor == a_op: + return True + return b_anchor.is_before_in_block(a_op) + + +# --- slices ------------------------------------------------------------------ + + +def backward_slice(value: ir.Value) -> set: + """The ops transitively producing `value`, as a set of operation hashes. + + A plain DFS over operands, standing in for ``getBackwardSlice``. Regions are + traversed only through their ops' operands, which suffices for the + straight-line tensor IR the fusion inspects. + """ + slice_ops: set = set() + def_op = defining_op(value) + if def_op is None: + return slice_ops + stack = [def_op] + while stack: + cur = stack.pop() + key = cur.__hash__() + if key in slice_ops: + continue + slice_ops.add(key) + for operand in cur.operands: + producer = defining_op(operand) + if producer is not None: + stack.append(producer) + return slice_ops + + +def depends_on_op(value: ir.Value, op: ir.Operation | ir.OpView) -> bool: + """Whether `value` transitively depends on `op`. + + Tells the operands that may be hoisted above the reduction loop from the ones + computed *by* it. + """ + target = opview(op).operation + def_op = defining_op(value) + if def_op is not None and def_op == target: + return True + return target.__hash__() in backward_slice(value) + + +def resolve_slice_source(value: ir.Value) -> ir.Value: + """Resolve `value` through any chain of ``tensor.extract_slice`` to its source. + + Inside a tiled reduction loop the inner reduction reads slices of the real + input tensors, so comparing its inputs against an untiled op's inputs means + looking through those tile slices. + """ + while True: + def_op = defining_op(value) + if def_op is None or not isinstance(def_op.opview, tensor.ExtractSliceOp): + return value + value = def_op.opview.source + + +def match_reduction( + iter_carried_args: list[ir.BlockArgument], red_pos: int +) -> tuple[ir.Value | None, list]: + """Match a generic reduction, returning ``(reduced_value, combiner_ops)``. + + A port of ``mlir::matchReduction``. Relies on the same invariants: the first + combiner is a binary op taking the iteration-carried value and the reduced + value; the def-use chain from it is single-use, side-effect free and + immediately nested in the reduction region; and it ends at the terminator. + Returns ``(None, [])`` when no reduction is matched. + + Matching is limited to a single combiner op, as upstream does. + """ + combiners: list = [] + carried = iter_carried_args[red_pos] + uses = list(carried.uses) + if len(uses) != 1: + return None, [] + + combiner = uses[0].owner.operation + if len(combiner.operands) != 2: + return None, [] + reduced = ( + combiner.operands[1] + if combiner.operands[0] == carried + else combiner.operands[0] + ) + + # The reduced value must not itself depend on a carried value, or the chain + # is not a plain accumulate. + region_block = carried.owner + carried_set = set(iter_carried_args) + if reduced in carried_set: + return None, [] + slice_ops = backward_slice(reduced) + if any( + operand in carried_set + for op in region_block.operations + if op.operation.__hash__() in slice_ops + for operand in op.operands + ): + return None, [] + + # Walk the def-use chain to the terminator, gathering combiners in order. + while not combiner.has_trait(ir.IsTerminatorTrait): + if len(combiner.results) != 1: + return None, [] + combiner_uses = list(combiner.results[0].uses) + if len(combiner_uses) != 1: + return None, [] + if combiner.block != region_block: + return None, [] + combiners.append(combiner) + combiner = combiner_uses[0].owner.operation + + if len(combiners) != 1: + return None, [] + return reduced, combiners + + +# --- constants --------------------------------------------------------------- + + +def _constant_value(value: ir.Value): + """The numeric value of an ``arith.constant`` (scalar or splat), else None.""" + def_op = defining_op(value) + if def_op is None or not isinstance(def_op.opview, arith.ConstantOp): + return None + attr = def_op.opview.value + if isinstance(attr, (ir.FloatAttr, ir.IntegerAttr)): + return attr.value + if isinstance(attr, ir.DenseElementsAttr) and attr.is_splat: + splat = attr.get_splat_value() + return splat.value if hasattr(splat, "value") else None + return None + + +def constant_int_value(value: ir.Value) -> int | None: + """The integer value of an ``arith.constant``, or None if not constant. + + Stands in for ``getConstantIntValue``. + """ + constant = _constant_value(value) + return constant if isinstance(constant, int) else None + + +def is_defined_as_zero(value: ir.Value) -> bool: + """Whether `value` is statically known to be zero. + + Either a constant zero scalar/splat, or chained through a ``linalg.fill`` / + ``linalg.copy`` of a zero value. Mirrors the helper in ``FoldAddIntoDest``. + """ + if value is None: + return False + constant = _constant_value(value) + if constant is not None and constant == 0: + return True + def_op = defining_op(value) + if def_op is None: + return False + ov = def_op.opview + if isinstance(ov, (linalg.FillOp, linalg.CopyOp)): + inputs = list(ov.inputs) + return len(inputs) == 1 and is_defined_as_zero(inputs[0]) + return False + + +#: Supported float element types with their bit widths. `f16` and `bf16` share a +#: width but not a format, so neither widens into the other. +_FLOAT_WIDTHS = ( + (ir.F16Type, 16), + (ir.BF16Type, 16), + (ir.F32Type, 32), + (ir.F64Type, 64), +) + + +def float_width(element_type: ir.Type) -> int | None: + """Bit width of a supported float type, else None.""" + for cls, width in _FLOAT_WIDTHS: + if isinstance(element_type, cls): + return width + return None + + +def wider_float_type(a: ir.Type, b: ir.Type) -> ir.Type | None: + """The wider of two float types, or None if there is no common widening. + + Picks the precision a mixed-precision body is evaluated in. Equal-width types + of different format (``f16`` vs ``bf16``) have no single-step conversion between + them, so they are refused rather than guessed at. + """ + width_a, width_b = float_width(a), float_width(b) + if width_a is None or width_b is None: + return None + if a == b: + return a + if width_a == width_b: + return None + return a if width_a > width_b else b + + +def cast_float(value: ir.Value, element_type: ir.Type) -> ir.Value | None: + """`value` converted to `element_type` via ``extf``/``truncf``, or unchanged. + + Returns None when the two types have no such conversion, which is exactly when + `wider_float_type` refuses them. + """ + if value.type == element_type: + return value + have, want = float_width(value.type), float_width(element_type) + if have is None or want is None or have == want: + return None + if want > have: + return arith.extf(element_type, value) + return arith.truncf(element_type, value) + + +def operand_eliminating_constant( + op: ir.Operation | ir.OpView, element_type: ir.Type +) -> ir.Attribute | None: + """The constant to substitute for one operand of `op` to drop its magnitude. + + ``0`` for the additive family (``addf``/``subf``/``addi``/``subi``) and ``1`` for + the multiplicative one (``mulf``/``divf``/``muli``) -- i.e. the neutral element of + the family, which `arith::getNeutralElement` also gives for the commutative ops. + Unlike that helper this answers for the non-commutative ones too, where the + constant is *not* an identity in the left operand's position: ``0 - x`` is ``-x`` + and ``1 / x`` is the reciprocal, not ``x``. + + So the substitution generally changes the value, and the caller has to have its + own reason why that is harmless -- for the fusion's correction term it is that the + substituted part cancels in the new/old ratio, see `_emit_correction_term`. + Returns None for op kinds with no such constant. + """ + ov = opview(op) + if isinstance(ov, (arith.AddFOp, arith.SubFOp)): + return ir.FloatAttr.get(element_type, 0.0) + if isinstance(ov, (arith.MulFOp, arith.DivFOp)): + return ir.FloatAttr.get(element_type, 1.0) + if isinstance(ov, (arith.AddIOp, arith.SubIOp)): + return ir.IntegerAttr.get(element_type, 0) + if isinstance(ov, arith.MulIOp): + return ir.IntegerAttr.get(element_type, 1) + return None + + +# --- affine maps ------------------------------------------------------------- + + +def remap_dims(imap: ir.AffineMap, dim_map: dict[int, int], num_dims: int): + """Rewrite a pure dim-projection map through `dim_map`, or None if not pure. + + Stands in for ``AffineMap::replaceDimsAndSymbols``, which the bindings do not + expose. The fusion legality restricts every map it rewrites to a pure + projection of plain dim exprs, so rebuilding from dim positions is exact. + Returns None if a result is not a plain dim expr or its position is unmapped. + """ + results = [] + for expr in imap.results: + if not isinstance(expr, ir.AffineDimExpr): + return None + if expr.position not in dim_map: + return None + results.append(ir.AffineDimExpr.get(dim_map[expr.position])) + return ir.AffineMap.get(num_dims, 0, results) + + +def project_dims(imap: ir.AffineMap, projected: set[int]): + """Drop `projected` dims from the map's domain, renumbering the rest. + + Stands in for ``projectDims(map, dims, /*compress=*/true)``. The caller + guarantees the map does not reference the projected dims (fusion legality + checks exactly that), so the projection is lossless. Returns None if the map + is not a pure dim projection or does reference a projected dim. + """ + num_dims = imap.n_dims + renumber: dict[int, int] = {} + next_pos = 0 + for pos in range(num_dims): + if pos in projected: + continue + renumber[pos] = next_pos + next_pos += 1 + return remap_dims(imap, renumber, next_pos) + + +# --- shapes ------------------------------------------------------------------ + + +def mixed_sizes(value: ir.Value) -> list: + """Sizes of a shaped `value`: ints for static dims, ``tensor.dim`` otherwise. + + Stands in for ``tensor::getMixedSizes``. + """ + shaped = ir.ShapedType(value.type) + sizes = [] + for pos, extent in enumerate(shaped.shape): + if ir.ShapedType.is_dynamic_size(extent): + index = arith.constant(ir.IndexType.get(), pos) + sizes.append(tensor.dim(value, index)) + else: + sizes.append(extent) + return sizes diff --git a/lighthouse/dialects/transform/transform_ext/utils/linalg_structured.py b/lighthouse/dialects/transform/transform_ext/utils/linalg_structured.py new file mode 100644 index 00000000..32deb4cf --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/linalg_structured.py @@ -0,0 +1,164 @@ +"""Structured-op (``LinalgOp`` interface) accessors missing from the bindings. + +The bindings expose only the raw ``indexing_maps`` / ``iterator_types`` +attributes, so the per-operand queries the reduction fusion relies on -- indexing +map, body block argument, loop ranges -- are rebuilt here. + +Operand order for a structured linalg op is inputs first, then outputs (DPS +inits), with one output per result. That single fact is what lets every accessor +below be derived from an operand index. + +`ir.OpOperand` carries only ``owner``/``operand_number`` and cannot read or write +the operand's value, so `Operand` below stands in for it. +""" + +from mlir import ir +from mlir.dialects import linalg + +from lighthouse.utils.mlir import indexing_maps, opview + +__all__ = [ + "Operand", + "dps_init_operands", + "dps_input_operands", + "indexing_map_for", + "iterator_types", + "matching_block_argument", + "num_dps_inits", + "num_loops", + "num_reduction_loops", + "operands_of", + "reduction_dims", + "region_output_args", + "static_loop_ranges", +] + + +class Operand: + """A readable/writable reference to one operand of an op. + + Holds the owning op and the operand index, and reads or writes the operand + value through them. Two references are equal when they name the same operand + of the same op, so they work as dict keys. + """ + + __slots__ = ("owner", "index") + + def __init__(self, owner: ir.Operation | ir.OpView, index: int): + self.owner = opview(owner) + self.index = index + + @property + def value(self) -> ir.Value: + """The value currently bound to this operand.""" + return self.owner.operands[self.index] + + def set(self, value: ir.Value) -> None: + """Rebind this operand to `value`.""" + self.owner.operands[self.index] = value + + def _key(self): + return (self.owner.operation.__hash__(), self.index) + + def __eq__(self, other) -> bool: + return isinstance(other, Operand) and self._key() == other._key() + + def __hash__(self) -> int: + return hash(self._key()) + + def __repr__(self) -> str: + return f"Operand(#{self.index} of {self.owner.operation.name})" + + +def operands_of(op: ir.Operation | ir.OpView) -> list[Operand]: + """All operands of `op` as `Operand` references, in operand order.""" + ov = opview(op) + return [Operand(ov, i) for i in range(len(ov.operands))] + + +def iterator_types(op: ir.Operation | ir.OpView) -> list[str]: + """Iterator types of a structured linalg op as ``"parallel"``/``"reduction"``. + + The attribute holds ``#linalg.iterator_type<...>`` attrs, which are compared + against the built enum attr rather than parsed. + """ + ov = opview(op) + build = ir.AttrBuilder.get("linalg.IteratorTypeEnum") + parallel = build(linalg.IteratorType.parallel, context=ov.context) + return ["parallel" if it == parallel else "reduction" for it in ov.iterator_types] + + +def num_loops(op: ir.Operation | ir.OpView) -> int: + """Number of iteration dims (loops) of a structured linalg op.""" + return len(iterator_types(op)) + + +def reduction_dims(op: ir.Operation | ir.OpView) -> list[int]: + """Positions of the reduction iterators, in order.""" + return [i for i, it in enumerate(iterator_types(op)) if it == "reduction"] + + +def num_reduction_loops(op: ir.Operation | ir.OpView) -> int: + """Number of reduction iterators.""" + return len(reduction_dims(op)) + + +def num_dps_inits(op: ir.Operation | ir.OpView) -> int: + """Number of DPS init (``outs``) operands, i.e. one per result.""" + return len(list(opview(op).results)) + + +def dps_input_operands(op: ir.Operation | ir.OpView) -> list[Operand]: + """The ``ins`` operands as `Operand` references.""" + ov = opview(op) + return operands_of(ov)[: len(ov.operands) - num_dps_inits(ov)] + + +def dps_init_operands(op: ir.Operation | ir.OpView) -> list[Operand]: + """The ``outs`` operands as `Operand` references.""" + ov = opview(op) + return operands_of(ov)[len(ov.operands) - num_dps_inits(ov) :] + + +def indexing_map_for(op: ir.Operation | ir.OpView, operand: Operand) -> ir.AffineMap: + """The indexing map matching `operand`, i.e. ``getMatchingIndexingMap``.""" + return indexing_maps(op)[operand.index] + + +def matching_block_argument( + op: ir.Operation | ir.OpView, operand: Operand +) -> ir.BlockArgument: + """The body block argument matching `operand`. + + Body arguments come one per operand in operand order, so the operand index + is the argument index. + """ + return opview(op).regions[0].blocks[0].arguments[operand.index] + + +def region_output_args(op: ir.Operation | ir.OpView) -> list[ir.BlockArgument]: + """The body block arguments matching the DPS init operands.""" + ov = opview(op) + args = list(ov.regions[0].blocks[0].arguments) + return args[len(args) - num_dps_inits(ov) :] + + +def static_loop_ranges(op: ir.Operation | ir.OpView) -> list[int]: + """Static extent of every loop dim, ``ShapedType::kDynamic`` where unknown. + + Recovered by matching each operand's indexing map against its shaped type: a + plain dim expr at map result `i` pins that loop dim to the operand's extent + along tensor dim `i`. Composite exprs carry no single extent and are skipped. + """ + ov = opview(op) + ranges = [ir.ShapedType.get_dynamic_size()] * num_loops(ov) + for operand, imap in zip(operands_of(ov), indexing_maps(ov)): + try: + shape = ir.ShapedType(operand.value.type).shape + except (ValueError, TypeError): + # A scalar operand pins no loop extent. + continue + for pos, expr in enumerate(imap.results): + if isinstance(expr, ir.AffineDimExpr): + ranges[expr.position] = shape[pos] + return ranges diff --git a/lighthouse/ingress/mlir_gen/gpu_attention_payload.py b/lighthouse/ingress/mlir_gen/gpu_attention_payload.py index b6faa006..ef486aaf 100644 --- a/lighthouse/ingress/mlir_gen/gpu_attention_payload.py +++ b/lighthouse/ingress/mlir_gen/gpu_attention_payload.py @@ -4,11 +4,46 @@ from mlir import ir from mlir.dialects import arith, bufferization, linalg, memref, tensor +from mlir.dialects import math as math_dialect from lighthouse.utils.mlir import func_cif from lighthouse.ingress.mlir_gen.utils import emit_buf_to_tensor +def _element_type(value: ir.Value) -> ir.Type: + """`value`'s element type if it is shaped, else its own type.""" + try: + return ir.ShapedType(value.type).element_type + except (ValueError, TypeError): + return value.type + + +def _generic(ins, out, indexing_maps, iterator_types, body): + """Emit a single-result ``linalg.generic`` over statically shaped tensors. + + `body` receives one scalar block argument per input followed by the output's, + and returns the value to yield. Block argument types are taken per operand, so + the operands and the result may differ in element type (a mixed-precision + body). + """ + maps = ir.ArrayAttr.get([ir.AffineMapAttr.get(m) for m in indexing_maps]) + iterators = ir.ArrayAttr.get( + [ir.Attribute.parse(f"#linalg.iterator_type<{it}>") for it in iterator_types] + ) + op = linalg.GenericOp( + result_tensors=[out.type], + inputs=ins, + outputs=[out], + indexing_maps=maps, + iterator_types=iterators, + ) + arg_types = [_element_type(v) for v in (*ins, out)] + block = op.regions[0].blocks.append(*arg_types) + with ir.InsertionPoint(block): + linalg.yield_([body(*block.arguments)]) + return op.results[0] + + def generate_gpu_attention_payload( func_name: str, batch_size: int, @@ -23,6 +58,11 @@ def generate_gpu_attention_payload( Computes attention: output = softmax(Q @ K^T / sqrt(d_head)) @ V + The softmax is emitted in the flash-attention form -- ``max``, ``exp``, row + ``sum`` and the ``P@V`` contraction as separate ops, with the normalizing + divide *after* the contraction -- rather than as a `linalg.softmax`. See + step 4 for why. + Args: func_name: Name of the payload function batch_size: Batch size @@ -84,44 +124,137 @@ def payload(output, Q_arg, K_arg, V_arg): # Step 2: Compute Q @ K^T using batch_matmul # Q: (batch_dim, n_ctx, d_head) @ K^T: (batch_dim, d_head, n_ctx) # Result: (batch_dim, n_ctx, n_ctx) + # Mixed precision: both contractions keep narrow operands (f16 feeds the + # DPAS units) but accumulate in f32, and the softmax runs in f32. The + # extensions the mixed-type contraction puts in its body are folded back + # into the `vector.contract` at vectorization by + # `fold_type_extensions_into_contract`. + compute_type = ir.F32Type.get() + narrow = dtype != compute_type + qkt_shape_3d = (batch_dim, n_ctx, n_ctx) - qkt_init = tensor.empty(qkt_shape_3d, dtype) + qkt_init = tensor.empty(qkt_shape_3d, compute_type) # Initialize with zeros for matmul accumulation - zero = arith.constant(dtype, 0.0) + zero = arith.constant(compute_type, 0.0) qkt_init_filled = linalg.fill(zero, outs=[qkt_init]) - # Batch matmul: Q @ K^T + # Batch matmul: Q @ K^T, f16 operands accumulating in f32. qkt = linalg.batch_matmul(Q_3d, K_transposed, outs=[qkt_init_filled]) # Step 3: Scale by 1/sqrt(d_head) scale_factor = 1.0 / math.sqrt(d_head) - scale_const = arith.constant(dtype, scale_factor) + scale_const = arith.constant(compute_type, scale_factor) # Create a tensor filled with the scale factor - scale_tensor_init = tensor.empty(qkt_shape_3d, dtype) + scale_tensor_init = tensor.empty(qkt_shape_3d, compute_type) scale_tensor = linalg.fill(scale_const, outs=[scale_tensor_init]) # Elementwise multiply qkt with scale tensor - scaled_qkt_init = tensor.empty(qkt_shape_3d, dtype) + scaled_qkt_init = tensor.empty(qkt_shape_3d, compute_type) scaled_qkt = linalg.mul(qkt, scale_tensor, outs=[scaled_qkt_init]) - # Step 4: Apply softmax along the last dimension (dim=2 in 3D) - softmax_init = tensor.empty(qkt_shape_3d, dtype) - attention_weights = linalg.softmax( - result=[ir.RankedTensorType.get(qkt_shape_3d, dtype)], - input=scaled_qkt, - output=softmax_init, - dimension=2, + # Step 4: softmax over the last dimension, written in the + # flash-attention form -- with the normalizing divide deferred past the + # P@V contraction: + # + # m = max_k s l = sum_k P + # P = exp(s - m) O = P @ V out = O / l + # + # Algebraically identical to `softmax(s) @ V`: dividing by the per-row + # `l` commutes with a contraction that reduces the *other* axis. Written + # this way rather than as a `linalg.softmax` -- whose decomposition + # normalizes *before* the contraction -- it leaves the + # `max -> exp -> {sum, P@V}` chain explicit, which is what + # `transform_ext.fuse_dependant_reduction_ops` folds into one loop. + d0, d1, d2 = (ir.AffineDimExpr.get(i) for i in range(3)) + # (batch, row, col) -> (batch, row, col) and -> (batch, row): the + # per-row statistics are broadcast over the reduced axis. + elementwise_map = ir.AffineMap.get(3, 0, [d0, d1, d2]) + row_map = ir.AffineMap.get(3, 0, [d0, d1]) + row_shape_3d = (batch_dim, n_ctx) + + # m = max_k s + neg_inf = arith.constant(compute_type, float("-inf")) + m_init = linalg.fill( + neg_inf, outs=[tensor.empty(row_shape_3d, compute_type)] + ) + row_max = _generic( + [scaled_qkt], + m_init, + [elementwise_map, row_map], + ["parallel", "parallel", "reduction"], + lambda s, acc: arith.maximumf(s, acc), + ) + + # P = exp(s - m), kept in f32 and read directly by both consumers: a + # cast in between would hide the chain from the fusion. + probs = _generic( + [scaled_qkt, row_max], + tensor.empty(qkt_shape_3d, compute_type), + [elementwise_map, row_map, elementwise_map], + ["parallel", "parallel", "parallel"], + lambda s, m, out: math_dialect.exp(arith.subf(s, m)), + ) + + # l = sum_k P, all in f32. + l_init = linalg.fill(zero, outs=[tensor.empty(row_shape_3d, compute_type)]) + row_sum = _generic( + [probs], + l_init, + [elementwise_map, row_map], + ["parallel", "parallel", "reduction"], + lambda p, acc: arith.addf(p, acc), ) - # Step 5: Multiply attention weights by V using batch_matmul - # attention_weights: (batch_dim, n_ctx, n_ctx) @ V: (batch_dim, n_ctx, d_head) - # Result: (batch_dim, n_ctx, d_head) - output_3d_init = tensor.empty(collapsed_shape_3d, dtype) + # Step 5: O = P @ V, still unnormalized. A `linalg.generic` rather than a + # `linalg.batch_matmul` so the narrowing of P can live *inside* the body: + # trunc P, widen it and V back to f32, multiply, accumulate. The two + # widenings are what `fold_type_extensions_into_contract` matches -- it + # folds both into the `vector.contract`, leaving a narrow x narrow -> f32 + # contraction (the DPAS shape) with just the `truncf` outside. + # + # A named matmul cannot express this: it casts every operand to the + # accumulator type, so an f32 P leaves nothing to fold on the lhs and V + # gets widened instead, losing the f16 DPAS. Keeping the cast in the body + # also keeps P readable straight from the elementwise term, so the fusion + # still sees the chain. + b, m, n, k = (ir.AffineDimExpr.get(i) for i in range(4)) + pv_lhs_map = ir.AffineMap.get(4, 0, [b, m, k]) + pv_rhs_map = ir.AffineMap.get(4, 0, [b, k, n]) + pv_out_map = ir.AffineMap.get(4, 0, [b, m, n]) + + def contract(p, v, acc): + lhs, rhs = p, v + if narrow: + lhs = arith.extf(compute_type, arith.truncf(dtype, p)) + rhs = arith.extf(compute_type, v) + return arith.addf(acc, arith.mulf(lhs, rhs)) + + output_3d_init = tensor.empty(collapsed_shape_3d, compute_type) output_3d_init_filled = linalg.fill(zero, outs=[output_3d_init]) + unnormalized = _generic( + [probs, V_3d], + output_3d_init_filled, + [pv_lhs_map, pv_rhs_map, pv_out_map], + ["parallel", "parallel", "parallel", "reduction"], + contract, + ) + + # Step 6: out = O / l, the deferred normalization, narrowed back to the + # payload's element type. `l` is broadcast over d_head, the contraction's + # free axis. + def normalize(o, denom, out): + normalized = arith.divf(o, denom) + if narrow: + return arith.truncf(dtype, normalized) + return normalized - result_3d = linalg.batch_matmul( - attention_weights, V_3d, outs=[output_3d_init_filled] + result_3d = _generic( + [unnormalized, row_sum], + tensor.empty(collapsed_shape_3d, dtype), + [elementwise_map, row_map, elementwise_map], + ["parallel", "parallel", "parallel"], + normalize, ) # Materialize 3D result back to 3D output memref diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index 5595f077..de36fa2b 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -72,12 +72,13 @@ def fused_attention_schedule( of form [1, ..., wg_rows], depending on the number of leading parallel dimensions, and the `sg_rows` tiling is applied over the n_ctx dimension. - In step 2., the inner attention block is tiled and fused over the - reduction dimension (n_ctx) of the final P@V operation, controlled by the - `reduction_tile` parameter. The Q@K^T and softmax operations are fused into - the P@V loop, implementing online softmax. This happens at tensor level, so - the tiling and fusion decisions stay at the level the rest of the schedule - works on and the emitted loop is lowered by the regular vectorization step. + In step 2., the inner attention block is folded into a single loop over the + key/value axis, controlled by the `reduction_tile` parameter: the row max is + tiled along that axis and the rest of the chain -- the exp term, the row sum, + the P@V contraction and the score computation -- is fused into it, giving the + online (flash) softmax. This happens at tensor level, so the tiling and fusion + decisions stay at the level the rest of the schedule works on and the emitted + loop is lowered by the regular vectorization step. Prefetching of K and V tiles is controlled by the `prefetch_tile` and `nb_prefetch` parameters. @@ -131,6 +132,198 @@ def fused_attention_schedule( return schedule +def _derive_flash_attention(anytype, func, layer_params): + """Derive the flash loop from the payload chain with the reduction fusion. + + `fuse_dependant_reduction_ops` moves the elementwise term and one consumer + reduction into an already-tiled producer reduction loop and inserts the online + correction that rescales that reduction's running accumulator whenever the + running max changes. Applied once per consumer reduction -- the row sum and the + `P@V` contraction -- it folds the whole chain into a single loop. + + Inside the WG forall the chain reads, in program order: + + %s = linalg.mul(batch_matmul(q, k^T), fill(scale)) the scaled scores + %m = max_k %s the producer reduction + %p = exp(%s - %m) the elementwise term + %o = contract(%p, v) consumer reduction + %l = sum_k %p consumer reduction + %out = %o / %l the deferred divide + """ + # Tile size for the reduction dimension (the K/V sequence length). + reduction_tile = layer_params["reduction_tile"] + + # Program order at this point (the WG tiling reorders the two consumer + # reductions relative to the payload's build order): row max, exp term, the P@V + # contraction, row sum, deferred divide. P@V is already a `linalg.generic` -- + # the payload writes it that way so the narrowing of P can sit in its body (see + # `generate_gpu_attention_payload`) -- so unlike Q@K^T it needs no `generalize` + # to be a fusable consumer. + max_op, p_op, pv_op, sum_op, _divide = match_and_split( + func, ops={"linalg.generic"}, nhandles=5 + ) + + # Tile the row max along the key/value axis. This is the producer reduction + # loop the rest of the chain gets folded into; the marker attribute is what + # the fusion op recognizes it by. + _, reduction_loop = structured.structured_tile_using_for( + anytype, + [anytype], + max_op, + dynamic_sizes=[], + interchange=[], + static_sizes=[0, 0, reduction_tile], + scalable_sizes=[False, False, False], + ) + transform.annotate(reduction_loop, transform_ext.REDUCTION_LOOP_ATTR_NAME) + + # First chain: max -> p -> row sum. `p` also feeds the contraction, so the op + # fuses a clone of it and leaves the original in place for the second chain. + # Fusing replaces the loop, but the replacement inherits the marker attribute, + # so it is ready to serve as the producer reduction of the second chain. + reduction_loop = transform_ext.fuse_dependant_reduction_ops( + p_op, sum_op, reduction_loop + ) + + # Second chain: max -> p -> P@V, into that same loop. The first fusion + # consumed the handle to `p`; the original is still the contraction's operand. + p_op = transform.get_producer_of_operand(anytype, pv_op, operand_number=0) + reduction_loop = transform_ext.fuse_dependant_reduction_ops( + p_op, pv_op, reduction_loop + ) + transform.apply_cse(func) + + # Sink the score computation into the reduction loop as well, so only one + # [wg_rows, tile_size] score tile -- rather than the full [wg_rows, n_ctx] + # matrix -- is ever live. + for producer_name in ["linalg.mul", "linalg.batch_matmul", "linalg.transpose"]: + producer_op = match_and_split(func, ops={producer_name}, nhandles=1)[0] + _, reduction_loop = structured.structured_fuse_into_containing_op( + anytype, + anytype, + producer_op=producer_op, + containing_op=reduction_loop, + ) + + # The Q @ K^T zero accumulator and the scale tensor are still filled at full + # [wg_rows, n_ctx] extent outside the loop, even though the ops now inside it + # only ever read a [wg_rows, tile_size] slice. Sink those two fills as well so + # neither tensor is materialized whole. Reach them through the in-loop consumer + # they initialize -- one hop for the slice the fusion left behind, one more for + # the fill itself. The three remaining fills initialize the loop's running + # accumulators and must stay outside. + scale_mul_op = match_and_split(reduction_loop, ops={"linalg.mul"}, nhandles=1)[0] + qk_matmul = match_and_split( + reduction_loop, ops={"linalg.batch_matmul"}, nhandles=1 + )[0] + for consumer_op, operand_number in [(scale_mul_op, 1), (qk_matmul, 2)]: + fill_slice = transform.get_producer_of_operand( + anytype, consumer_op, operand_number=operand_number + ) + fill_op = transform.get_producer_of_operand( + anytype, fill_slice, operand_number=0 + ) + _, reduction_loop = structured.structured_fuse_into_containing_op( + anytype, + anytype, + producer_op=fill_op, + containing_op=reduction_loop, + ) + + transform.apply_cse(func) + canonicalize(func) + + # Strip the batch dim, which the work-group tiling cut down to 1. Everything + # downstream -- the XeGPU layouts below and the blocking/distribution passes in + # `xegpu_to_binary` -- is built for rank-2 tiles, and the vector-level + # `cast_away_leading_one_dim` patterns cannot finish the job: they have no + # pattern for multi_reduction, broadcast or transpose, so the softmax row + # reductions and the correction's broadcasts would keep a unit dim and drag + # shape_casts (and rank-3 XeGPU layouts) along with them. + # + # Done here, after the reduction fusion rather than before it, for two reasons: + # `fuse_dependant_reduction_ops` gets to run on the shape it already handles, + # and every op above is still matched by name -- generalizing first would turn + # the whole chain into linalg.generic. + # + # `fold_unit_extent_dims` only rewrites linalg.generic, hence the generalize; + # and it leaves two-step slice chains behind (16x4096x64 -> 1x128x64 -> + # 128x64), which plain canonicalization does not compose, hence the tensor + # patterns. The scf.for's own iter_args keep their unit dim -- no pattern + # retypes a loop signature -- but that no longer matters: with every op inside + # rank 2, vectorization reads through them rank-reducing and the accumulators + # come out as rank-1/2 vectors. + named_ops = match( + func, + ops={ + "linalg.batch_matmul", + "linalg.mul", + "linalg.transpose", + "linalg.fill", + "linalg.elementwise", + }, + ) + structured.structured_generalize(anytype, named_ops) + with ir.InsertionPoint(transform.apply_patterns(func).patterns): + structured.apply_patterns_linalg_fold_unit_extent_dims_via_slices() + transform.apply_patterns_canonicalization() + transform.apply_cse(func) + with ir.InsertionPoint(transform.apply_patterns(func).patterns): + tensor.apply_patterns_tensor_merge_consecutive_insert_extract_slice() + tensor.apply_patterns_tensor_drop_redundant_insert_slice_rank_expansion() + tensor.apply_patterns_tensor_fold_tensor_subset_ops() + transform.apply_patterns_canonicalization() + transform.apply_cse(func) + + +def _replace_with_reference_flash_attention(anytype, func, layer_params): + """Emit the flash loop from the hand-written generator instead of deriving it. + + `replace_with_fused_attention` builds the whole online-softmax loop from + scratch given Q, K, V and the scale, replacing the chain's leaf. It is kept as + a reference point: the derived path (the default, see + `fuse_dependant_reduction_ops`) should converge on the same loop, and diffing + the two at `--dump-kernel=reduction-tiled` is how that is tracked. + + Both paths consume the same payload. The generator's own output *is* the + normalized result, so its `output` is the payload's deferred divide -- the + chain's leaf -- rather than the `P@V` contraction; everything upstream of it + (max, exp, row sum, `P@V`) is left dead for DCE. + """ + prod = transform.get_producer_of_operand + # The Q, K, V tensors and the scale constant are found by walking the SSA chain + # of the two batch matmuls inside the WG forall: + # + # Q@K^T: linalg.batch_matmul(q_slice, linalg.transpose(k_slice)) + # scale: linalg.mul(qkt, linalg.fill(scale_constant)) + # P@V: linalg.generic(probs, v_slice) -- a generic, not a named matmul, + # so that P's narrowing can live in + # its body; V is its rhs either way + qk_matmul = match_and_split(func, ops={"linalg.batch_matmul"}, nhandles=1)[0] + q = prod(anytype, qk_matmul, operand_number=0) + k_transpose = prod(anytype, qk_matmul, operand_number=1) + k = prod(anytype, k_transpose, operand_number=0) + mul_op = match_and_split(func, ops={"linalg.mul"}, nhandles=1)[0] + scale = prod(anytype, prod(anytype, mul_op, operand_number=1), operand_number=0) + # Generics in program order: row max, exp term, P@V, row sum, deferred divide. + # The divide is the chain's leaf; V is P@V's rhs. + _max, _exp, pv_op, _sum, divide_op = match_and_split( + func, ops={"linalg.generic"}, nhandles=5 + ) + v = prod(anytype, pv_op, operand_number=1) + + transform_ext.replace_with_fused_attention( + q=q, + k=k, + v=v, + scale=scale, + output=divide_op, + tile_size=layer_params["reduction_tile"], + ) + transform.apply_cse(func) + lh_transform.cleanup(func) + + def bundle_xegpu_fused_attention_schedule( mod: ir.Value[transform.AnyOpType], params: ScheduleParameters, @@ -148,9 +341,11 @@ def bundle_xegpu_fused_attention_schedule( # Match payload function func = get_payload_func(mod, op_name=["linalg.generic", "linalg.batch_matmul"]) - # Match linalg.softmax operation if any and decompose it into generic ops - softmax_ops = structured.structured_match(anytype, func, ops=["linalg.softmax"]) - structured.structured_decompose_interface(anytype, softmax_ops) + # The payload spells the softmax out as `max -> exp -> {sum, P@V} -> divide` + # (see `generate_gpu_attention_payload`), so there is no `linalg.softmax` to + # decompose: its decomposition normalizes *before* the contraction, which + # would leave the P@V reading the normalized P and so break the dependency + # chain the reduction fusion needs. # Normalize possible singleton dimensions so tile+fuse logic works. with ir.InsertionPoint(transform.apply_patterns(func).patterns): @@ -190,43 +385,23 @@ def bundle_xegpu_fused_attention_schedule( if stop_at_stage == "tiled": raise PipelineInterrupt() - # Apply reduction tiling and fusion, still at tensor level. The Q, K, V - # tensors and the scale constant are found by walking the SSA chain of the - # two batch matmuls inside the WG forall: + # Build the fused (flash) attention inner loop -- a single loop over the + # key/value axis -- while still at linalg level. Two paths produce it, and both + # consume the same payload, so their output can be diffed at + # `--dump-kernel=reduction-tiled`: # - # Q@K^T: linalg.batch_matmul(q_slice, linalg.transpose(k_slice)) - # scale: linalg.mul(qkt, linalg.fill(scale_constant)) - # P@V: linalg.batch_matmul(softmax_out, v_slice) - matmul_ops = match_and_split(func, ops={"linalg.batch_matmul"}, nhandles=2) - qk_matmul, pv_matmul = matmul_ops[0], matmul_ops[1] - - q = transform.get_producer_of_operand(anytype, qk_matmul, operand_number=0) - k_transpose = transform.get_producer_of_operand( - anytype, qk_matmul, operand_number=1 - ) - k = transform.get_producer_of_operand(anytype, k_transpose, operand_number=0) - v = transform.get_producer_of_operand(anytype, pv_matmul, operand_number=1) - - # The scale is the fill value of the linalg.mul rhs operand. - mul_op = match_and_split(func, ops={"linalg.mul"}, nhandles=1)[0] - scale_fill = transform.get_producer_of_operand(anytype, mul_op, operand_number=1) - scale = transform.get_producer_of_operand(anytype, scale_fill, operand_number=0) - - # Replace the P@V batch matmul with a loop over the K/V sequence length that - # implements online softmax, fusing Q@K^T and the softmax into it. - reduction_tile = layer_params[ - "reduction_tile" - ] # Tile size for reduction dimension (K/V sequence length) - transform_ext.replace_with_fused_attention( - q=q, - k=k, - v=v, - scale=scale, - output=pv_matmul, - tile_size=reduction_tile, - ) - transform.apply_cse(func) - lh_transform.cleanup(func) + # * the default derives it from the payload's chain with + # `fuse_dependant_reduction_ops` (below); + # * `reference_flash` emits it from the hand-written generator instead, as a + # reference point for how close the derived version gets. + # Tile size for the reduction dimension (the K/V sequence length); also drives + # the K/V prefetch and the XeGPU layouts further down. + reduction_tile = layer_params["reduction_tile"] + + if layer_params.get("reference_flash", False): + _replace_with_reference_flash_attention(anytype, func, layer_params) + else: + _derive_flash_attention(anytype, func, layer_params) if stop_at_stage == "reduction-tiled": raise PipelineInterrupt() @@ -240,6 +415,18 @@ def bundle_xegpu_fused_attention_schedule( # accumulators are carried as vector iter_args, i.e. in registers. reduction_loop = match(func, ops={"scf.for"}) lh_transform.loop_hoisting(reduction_loop) + + # Turn on fast math and take the rewrite it licenses: the online rescale factor + # arrives as `exp(-m_new) / exp(-m_old)` and becomes the single + # `exp(m_old - m_new)` a hand-written flash-attention kernel uses, dropping a + # transcendental and a divide per iteration. Run after vectorization -- before + # it the two exponentials and the divide live in three separate linalg.generics, + # so there is no single op to match. + transform_ext.enable_fastmath_optimizations(func) + transform.apply_cse(func) + canonicalize(func) + + func = apply_registered_pass(func, "remove-dead-values") lh_transform.cleanup(func) if stop_at_stage == "vectorized": diff --git a/test/transform/test_fuse_dependant_reduction_ops.py b/test/transform/test_fuse_dependant_reduction_ops.py new file mode 100644 index 00000000..fbcd63cf --- /dev/null +++ b/test/transform/test_fuse_dependant_reduction_ops.py @@ -0,0 +1,463 @@ +# RUN: %PYTHON %s | FileCheck %s + +"""Tests for `transform_ext.fuse_dependant_reduction_ops`. + +The op fuses a dependency chain ``R1 -> E -> R2`` into ``R1``'s already-tiled +reduction loop, turning a two-pass reduction into an online (one-pass) one. Three +scenarios are covered: + + 1. **softmax** -- ``max`` then ``sum(exp)``, with the normalizing divide reading + ``E``'s full extent. Because ``E`` has a consumer besides ``R2``, the op fuses + a *clone* and leaves the original outside for the divide. + 2. **flash attention** -- one ``exp`` term feeding both a row sum and a `P @ V` + contraction. Applying the op once per consumer reduction folds both into a + single loop, leaving only the normalization outside. + 3. **mixed-precision softmax** -- the same chain with an f16 term feeding an f32 + sum accumulator, checking the correction factor is evaluated in the wider of + the two element types. +""" + +from mlir import ir +from mlir.dialects import transform +from mlir.dialects.transform import structured + +import lighthouse.dialects as lh_dialects +from lighthouse import transform as lh_transform +from lighthouse.dialects.transform import transform_ext +from lighthouse.schedule.builders import schedule_boilerplate + + +# --------------------------------------------------------------------------- +# Payloads +# --------------------------------------------------------------------------- + +#: Softmax over the trailing (reduction) axis, untiled: +#: m = max_j x, p = exp(x - m), s = sum_j p, out = p / s +SOFTMAX = """ +#rowcol = affine_map<(d0, d1) -> (d0, d1)> +#row = affine_map<(d0, d1) -> (d0)> + +func.func @softmax(%x: tensor<64x512xf32>) -> tensor<64x512xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %ninf = arith.constant 0xFF800000 : f32 + %row_init = tensor.empty() : tensor<64xf32> + %full_init = tensor.empty() : tensor<64x512xf32> + + // R1: m = max_j x + %m_init = linalg.fill ins(%ninf : f32) outs(%row_init : tensor<64xf32>) -> tensor<64xf32> + %m = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%x : tensor<64x512xf32>) outs(%m_init : tensor<64xf32>) { + ^bb0(%in: f32, %out: f32): + %mx = arith.maximumf %in, %out : f32 + linalg.yield %mx : f32 + } -> tensor<64xf32> + + // E: p = exp(x - m) + %p = linalg.generic {indexing_maps = [#rowcol, #row, #rowcol], + iterator_types = ["parallel", "parallel"]} + ins(%x, %m : tensor<64x512xf32>, tensor<64xf32>) + outs(%full_init : tensor<64x512xf32>) { + ^bb0(%in: f32, %mv: f32, %out: f32): + %d = arith.subf %in, %mv : f32 + %e = math.exp %d : f32 + linalg.yield %e : f32 + } -> tensor<64x512xf32> + + // R2: s = sum_j p + %s_init = linalg.fill ins(%zero : f32) outs(%row_init : tensor<64xf32>) -> tensor<64xf32> + %s = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%p : tensor<64x512xf32>) outs(%s_init : tensor<64xf32>) { + ^bb0(%in: f32, %out: f32): + %a = arith.addf %in, %out : f32 + linalg.yield %a : f32 + } -> tensor<64xf32> + + // The normalizing divide, downstream of the chain. This is the extra consumer + // of `E` that forces the fusion to work on a clone. + %out = linalg.generic {indexing_maps = [#rowcol, #row, #rowcol], + iterator_types = ["parallel", "parallel"]} + ins(%p, %s : tensor<64x512xf32>, tensor<64xf32>) + outs(%full_init : tensor<64x512xf32>) { + ^bb0(%in: f32, %sv: f32, %o: f32): + %d = arith.divf %in, %sv : f32 + linalg.yield %d : f32 + } -> tensor<64x512xf32> + return %out : tensor<64x512xf32> +} +""" + +#: Softmax whose term is f16 while the sum accumulates in f32, the usual attention +#: mix: `m = max_j x`, `p = exp(x - m)` in f16, `s = sum_j p` in f32. +MIXED_SOFTMAX = """ +#rowcol = affine_map<(d0, d1) -> (d0, d1)> +#row = affine_map<(d0, d1) -> (d0)> + +func.func @mixed_softmax(%x: tensor<64x512xf16>) -> tensor<64xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %ninf = arith.constant 0xFC00 : f16 + %row_init_f16 = tensor.empty() : tensor<64xf16> + %row_init_f32 = tensor.empty() : tensor<64xf32> + %full_init = tensor.empty() : tensor<64x512xf16> + + // R1: m = max_j x, in f16. + %m_init = linalg.fill ins(%ninf : f16) outs(%row_init_f16 : tensor<64xf16>) -> tensor<64xf16> + %m = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%x : tensor<64x512xf16>) outs(%m_init : tensor<64xf16>) { + ^bb0(%in: f16, %out: f16): + %mx = arith.maximumf %in, %out : f16 + linalg.yield %mx : f16 + } -> tensor<64xf16> + + // E: p = exp(x - m), also f16. + %p = linalg.generic {indexing_maps = [#rowcol, #row, #rowcol], + iterator_types = ["parallel", "parallel"]} + ins(%x, %m : tensor<64x512xf16>, tensor<64xf16>) + outs(%full_init : tensor<64x512xf16>) { + ^bb0(%in: f16, %mv: f16, %out: f16): + %d = arith.subf %in, %mv : f16 + %e = math.exp %d : f16 + linalg.yield %e : f16 + } -> tensor<64x512xf16> + + // R2: s = sum_j p, widened into an f32 accumulator. + %s_init = linalg.fill ins(%zero : f32) outs(%row_init_f32 : tensor<64xf32>) -> tensor<64xf32> + %s = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%p : tensor<64x512xf16>) outs(%s_init : tensor<64xf32>) { + ^bb0(%in: f16, %out: f32): + %w = arith.extf %in : f16 to f32 + %a = arith.addf %w, %out : f32 + linalg.yield %a : f32 + } -> tensor<64xf32> + return %s : tensor<64xf32> +} +""" + +#: `softmax(x) @ v`, with the softmax written as an unfused two-pass reduction. +ATTENTION = """ +#rowcol = affine_map<(d0, d1) -> (d0, d1)> +#row = affine_map<(d0, d1) -> (d0)> +#ik = affine_map<(d0, d1, d2) -> (d0, d2)> +#kj = affine_map<(d0, d1, d2) -> (d2, d1)> +#ij = affine_map<(d0, d1, d2) -> (d0, d1)> + +func.func @attention(%x: tensor<64x512xf32>, %v: tensor<512x128xf32>) + -> tensor<64x128xf32> { + %zero = arith.constant 0.000000e+00 : f32 + %ninf = arith.constant 0xFF800000 : f32 + %row_init = tensor.empty() : tensor<64xf32> + %p_init = tensor.empty() : tensor<64x512xf32> + + // R1: m = max_k x + %m_init = linalg.fill ins(%ninf : f32) outs(%row_init : tensor<64xf32>) -> tensor<64xf32> + %m = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%x : tensor<64x512xf32>) outs(%m_init : tensor<64xf32>) { + ^bb0(%in: f32, %out: f32): + %mx = arith.maximumf %in, %out : f32 + linalg.yield %mx : f32 + } -> tensor<64xf32> + + // E: p = exp(x - m), read by BOTH reductions below. + %p = linalg.generic {indexing_maps = [#rowcol, #row, #rowcol], + iterator_types = ["parallel", "parallel"]} + ins(%x, %m : tensor<64x512xf32>, tensor<64xf32>) + outs(%p_init : tensor<64x512xf32>) { + ^bb0(%in: f32, %mv: f32, %out: f32): + %d = arith.subf %in, %mv : f32 + %e = math.exp %d : f32 + linalg.yield %e : f32 + } -> tensor<64x512xf32> + + // R2a: l = sum_k p + %l_init = linalg.fill ins(%zero : f32) outs(%row_init : tensor<64xf32>) -> tensor<64xf32> + %l = linalg.generic {indexing_maps = [#rowcol, #row], + iterator_types = ["parallel", "reduction"]} + ins(%p : tensor<64x512xf32>) outs(%l_init : tensor<64xf32>) { + ^bb0(%in: f32, %out: f32): + %a = arith.addf %in, %out : f32 + linalg.yield %a : f32 + } -> tensor<64xf32> + + // R2b: o = p @ v, a contraction over the same axis. + %o_init = tensor.empty() : tensor<64x128xf32> + %o_fill = linalg.fill ins(%zero : f32) outs(%o_init : tensor<64x128xf32>) -> tensor<64x128xf32> + %o = linalg.generic {indexing_maps = [#ik, #kj, #ij], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%p, %v : tensor<64x512xf32>, tensor<512x128xf32>) + outs(%o_fill : tensor<64x128xf32>) { + ^bb0(%pv: f32, %vv: f32, %out: f32): + %mul = arith.mulf %pv, %vv : f32 + %add = arith.addf %out, %mul : f32 + linalg.yield %add : f32 + } -> tensor<64x128xf32> + + // The deferred normalization, downstream of both reductions. + %out_init = tensor.empty() : tensor<64x128xf32> + %out = linalg.generic {indexing_maps = [#rowcol, #row, #rowcol], + iterator_types = ["parallel", "parallel"]} + ins(%o, %l : tensor<64x128xf32>, tensor<64xf32>) + outs(%out_init : tensor<64x128xf32>) { + ^bb0(%in: f32, %lv: f32, %o2: f32): + %d = arith.divf %in, %lv : f32 + linalg.yield %d : f32 + } -> tensor<64x128xf32> + return %out : tensor<64x128xf32> +} +""" + + +# --------------------------------------------------------------------------- +# Schedules +# --------------------------------------------------------------------------- + + +def online_softmax_schedule(tile_size: int = 32) -> ir.Module: + """Tile `R1` along its reduction axis, then fuse the `E -> R2` chain into it. + + The fusion op does not tile `R1` itself -- it reads the tile size off the + loop's step -- so the schedule tiles first and annotates the resulting loop. + """ + with schedule_boilerplate() as (sched, seq): + generics = lh_transform.match_op(seq.bodyTarget, "linalg.generic") + anyop = transform.AnyOpType.get() + # In program order: the max (R1), the exp term (E), the sum (R2) and the + # normalizing divide. + r1, e, r2, _divide = transform.split_handle([anyop] * 4, generics) + + # Tile R1 along the reduction dim only, giving the `scf.for` the fusion + # needs, and mark it as a reduction loop. + _tiled_r1, r1_loop = structured.TileUsingForOp(r1, sizes=[0, tile_size]).results + transform.annotate(r1_loop, transform_ext.REDUCTION_LOOP_ATTR_NAME) + + fused = transform_ext.fuse_dependant_reduction_ops(e, r2, r1_loop) + transform.annotate(fused, "online_softmax_loop") + transform.yield_([]) + return sched + + +def mixed_softmax_schedule(tile_size: int = 32) -> ir.Module: + """Same as `online_softmax_schedule`, for a payload with no trailing divide.""" + with schedule_boilerplate() as (sched, seq): + generics = lh_transform.match_op(seq.bodyTarget, "linalg.generic") + anyop = transform.AnyOpType.get() + # In program order: the max (R1), the exp term (E) and the sum (R2). + r1, e, r2 = transform.split_handle([anyop] * 3, generics) + + _tiled_r1, r1_loop = structured.TileUsingForOp(r1, sizes=[0, tile_size]).results + transform.annotate(r1_loop, transform_ext.REDUCTION_LOOP_ATTR_NAME) + + fused = transform_ext.fuse_dependant_reduction_ops(e, r2, r1_loop) + transform.annotate(fused, "mixed_softmax_loop") + transform.yield_([]) + return sched + + +def flash_attention_schedule(tile_size: int = 32) -> ir.Module: + """Tile `R1`, then fuse both consumer reductions into its loop in turn.""" + with schedule_boilerplate() as (sched, seq): + generics = lh_transform.match_op(seq.bodyTarget, "linalg.generic") + anyop = transform.AnyOpType.get() + # In program order: the max (R1), the exp term (E), the row sum (R2a), the + # contraction (R2b) and the normalizing divide. + r1, e, r2a, r2b, _divide = transform.split_handle([anyop] * 5, generics) + + _tiled_r1, r1_loop = structured.TileUsingForOp(r1, sizes=[0, tile_size]).results + transform.annotate(r1_loop, transform_ext.REDUCTION_LOOP_ATTR_NAME) + + # First chain: the row sum. This fuses a *clone* of E, since E still feeds + # the contraction, and leaves the original E in place for it. + loop = transform_ext.fuse_dependant_reduction_ops(e, r2a, r1_loop) + # The first fusion consumed the handle to E; the original E is still the + # contraction's operand, so re-derive it from there. + e_again = transform.get_producer_of_operand(anyop, r2b, 0) + # Second chain: the contraction, into the same (now replaced) loop. + loop = transform_ext.fuse_dependant_reduction_ops(e_again, r2b, loop) + transform.annotate(loop, "flash_attention_loop") + transform.yield_([]) + return sched + + +def apply_schedule(payload_str: str, build_schedule, tile_size: int) -> ir.Module: + """Parse `payload_str` and apply `build_schedule` at `tile_size`.""" + payload = ir.Module.parse(payload_str) + # Bound to a local: the schedule module must outlive `apply`. + schedule = build_schedule(tile_size) + schedule.body.operations[0].apply(payload.operation) + assert payload.operation.verify() + return payload + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_softmax_structure() -> None: + """Fuse softmax and print the resulting online loop.""" + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + print(apply_schedule(SOFTMAX, online_softmax_schedule, 32)) + + +# The fused loop carries three accumulators: the running max, the (stale, +# write-only) full-extent E result and the running sum. +# CHECK-LABEL: func.func @softmax +# CHECK-SAME: %[[X:[a-zA-Z0-9_]+]]: tensor<64x512xf32> +# CHECK: %[[LOOP:.+]]:3 = scf.for %[[IV:[a-zA-Z0-9_]+]] = +# CHECK-SAME: iter_args(%[[MARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}, %[[EARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}, %[[SARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}) +# CHECK-SAME: -> (tensor<64xf32>, tensor<64x512xf32>, tensor<64xf32>) + +# R1 over this tile. %[[MOLD]] is its DPS init, i.e. the previous running max. +# CHECK: %[[MOLD:.+]] = tensor.extract_slice %[[MARG]] +# CHECK: %[[MNEW:.+]] = linalg.generic +# CHECK-SAME: outs(%[[MOLD]] : +# CHECK: arith.maximumf + +# The fused clone of E, cut from the full 512 extent to the 32-wide tile and +# reading the *new* max. +# CHECK: %[[XT:.+]] = tensor.extract_slice %[[X]][0, %[[IV]]] [64, 32] [1, 1] +# CHECK: %[[ET:.+]] = tensor.extract_slice %[[EARG]][0, %[[IV]]] [64, 32] [1, 1] +# CHECK: %[[P:.+]] = linalg.generic +# CHECK-SAME: ins(%[[XT]], %[[MNEW]] : tensor<64x32xf32>, tensor<64xf32>) +# CHECK-SAME: outs(%[[ET]] : tensor<64x32xf32>) +# CHECK: arith.subf +# CHECK: math.exp + +# The online correction: E isolated on the max, evaluated at the new and the old +# max, their ratio rescaling the running sum. +# CHECK: %[[SOLD:.+]] = tensor.extract_slice %[[SARG]] +# CHECK: %[[TNEW:.+]] = linalg.generic {{.*}}ins(%[[MNEW]] : +# CHECK: arith.subf %{{.+}}, %in +# CHECK: math.exp +# CHECK: %[[TOLD:.+]] = linalg.generic {{.*}}ins(%[[MOLD]] : +# CHECK: arith.subf %{{.+}}, %in +# CHECK: math.exp +# CHECK: %[[F:.+]] = linalg.elementwise kind=#linalg.elementwise_kind
ins(%[[TNEW]], %[[TOLD]] +# CHECK: %[[SCALED:.+]] = linalg.elementwise kind=#linalg.elementwise_kind ins(%[[SOLD]], %[[F]] + +# The fused R2 accumulates this tile's terms into the rescaled running sum. +# CHECK: %[[SNEW:.+]] = linalg.generic +# CHECK-SAME: ins(%[[P]] : tensor<64x32xf32>) +# CHECK-SAME: outs(%[[SCALED]] : tensor<64xf32>) +# CHECK: arith.addf +# CHECK: tensor.insert_slice %[[P]] into %[[EARG]][0, %[[IV]]] [64, 32] [1, 1] +# CHECK: tensor.insert_slice %[[SNEW]] into %[[SARG]] +# CHECK: } {__reduction_loop__, online_softmax_loop} + +# The original E survives outside the loop, reading the *final* max off it, so the +# normalizing divide sees a correctly recomputed numerator. +# CHECK: %[[PFINAL:.+]] = linalg.generic +# CHECK-SAME: ins(%[[X]], %[[LOOP]]#0 : tensor<64x512xf32>, tensor<64xf32>) +# CHECK: arith.subf +# CHECK: math.exp +# CHECK: linalg.generic +# CHECK-SAME: ins(%[[PFINAL]], %[[LOOP]]#2 : tensor<64x512xf32>, tensor<64xf32>) +# CHECK: arith.divf + + +def test_attention_structure() -> None: + """Fuse both consumer reductions of an attention chain into one loop.""" + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + print(apply_schedule(ATTENTION, flash_attention_schedule, 32)) + + +# Both consumer reductions end up in one loop, which now carries five accumulators, +# in this order: the running max, the first chain's full-extent (stale, write-only) E +# result, the running row sum, the second chain's E result, and the running +# contraction accumulator. +# CHECK-LABEL: func.func @attention +# CHECK-SAME: %[[X:[a-zA-Z0-9_]+]]: tensor<64x512xf32> +# CHECK-SAME: %[[V:[a-zA-Z0-9_]+]]: tensor<512x128xf32> +# CHECK: %[[LOOP:.+]]:5 = scf.for %[[IV:[a-zA-Z0-9_]+]] = +# CHECK-SAME: iter_args(%[[MARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}, +# CHECK-SAME: -> (tensor<64xf32>, tensor<64x512xf32>, tensor<64xf32>, tensor<64x512xf32>, tensor<64x128xf32>) + +# R1 over this tile, then the first chain: E's clone, its correction and the row +# sum accumulating into the rescaled running sum. +# CHECK: %[[MOLD:.+]] = tensor.extract_slice %[[MARG]] +# CHECK: %[[MNEW:.+]] = linalg.generic +# CHECK-SAME: outs(%[[MOLD]] : +# CHECK: arith.maximumf +# CHECK: %[[P1:.+]] = linalg.generic +# CHECK-SAME: ins(%{{.+}}, %[[MNEW]] : tensor<64x32xf32>, tensor<64xf32>) +# CHECK: math.exp +# CHECK: linalg.elementwise kind=#linalg.elementwise_kind
+# CHECK: %[[LSCALED:.+]] = linalg.elementwise kind=#linalg.elementwise_kind +# CHECK: linalg.generic +# CHECK-SAME: ins(%[[P1]] : tensor<64x32xf32>) +# CHECK-SAME: outs(%[[LSCALED]] : tensor<64xf32>) +# CHECK: arith.addf + +# The second chain. No clone this time: the row sum is gone, so the contraction is +# E's only remaining user and E itself is fused. Its correction runs over the +# contraction's *wider* 64x128 accumulator -- the per-row factor broadcast over N. +# CHECK: %[[P2:.+]] = linalg.generic +# CHECK-SAME: ins(%{{.+}}, %[[MNEW]] : tensor<64x32xf32>, tensor<64xf32>) +# CHECK: math.exp +# CHECK: linalg.generic {{.*}}ins(%[[MNEW]] : tensor<64xf32>) outs(%{{.+}} : tensor<64x128xf32>) +# CHECK: linalg.elementwise kind=#linalg.elementwise_kind
+# CHECK: %[[OSCALED:.+]] = linalg.elementwise kind=#linalg.elementwise_kind + +# The contraction reads V sliced along the shared reduction axis. +# CHECK: %[[VT:.+]] = tensor.extract_slice %[[V]][%[[IV]], 0] [32, 128] [1, 1] +# CHECK: linalg.generic +# CHECK-SAME: ins(%[[P2]], %[[VT]] : tensor<64x32xf32>, tensor<32x128xf32>) +# CHECK-SAME: outs(%[[OSCALED]] : tensor<64x128xf32>) +# CHECK: } {__reduction_loop__, flash_attention_loop} + +# Only the normalization is left outside, reading the two running results. +# CHECK: linalg.generic +# CHECK-SAME: ins(%[[LOOP]]#4, %[[LOOP]]#2 : tensor<64x128xf32>, tensor<64xf32>) +# CHECK: arith.divf + + +def test_mixed_precision_softmax() -> None: + """Fuse a softmax whose f16 term feeds an f32 sum accumulator.""" + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + print(apply_schedule(MIXED_SOFTMAX, mixed_softmax_schedule, 32)) + + +# The running max and the term stay f16; only the sum accumulator is f32. +# CHECK-LABEL: func.func @mixed_softmax +# CHECK-SAME: %[[X:[a-zA-Z0-9_]+]]: tensor<64x512xf16> +# CHECK: %[[LOOP:.+]]:3 = scf.for %[[IV:[a-zA-Z0-9_]+]] = +# CHECK-SAME: iter_args(%[[MARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}, %{{[a-zA-Z0-9_]+}} = %{{[a-zA-Z0-9_]+}}, %[[SARG:[a-zA-Z0-9_]+]] = %{{[a-zA-Z0-9_]+}}) +# CHECK-SAME: -> (tensor<64xf16>, tensor<64x512xf16>, tensor<64xf32>) +# CHECK: %[[MOLD:.+]] = tensor.extract_slice %[[MARG]] +# CHECK: %[[MNEW:.+]] = linalg.generic +# CHECK: arith.maximumf %{{.+}} : f16 +# CHECK: %[[P:.+]] = linalg.generic +# CHECK: math.exp %{{.+}} : f16 + +# The correction is evaluated in f32 -- the wider of E's f16 and R2's f32 +# accumulator -- so the f16 running max enters the body and is widened there. +# CHECK: linalg.generic {{.*}}ins(%[[MNEW]] : tensor<64xf16>) outs(%{{.+}} : tensor<64xf32>) +# CHECK: %[[WNEW:.+]] = arith.extf %in : f16 to f32 +# CHECK: arith.subf %{{.+}}, %[[WNEW]] : f32 +# CHECK: math.exp %{{.+}} : f32 +# CHECK: linalg.generic {{.*}}ins(%[[MOLD]] : tensor<64xf16>) outs(%{{.+}} : tensor<64xf32>) +# CHECK: %[[WOLD:.+]] = arith.extf %in : f16 to f32 +# CHECK: arith.subf %{{.+}}, %[[WOLD]] : f32 +# CHECK: math.exp %{{.+}} : f32 + +# The factor already has R2's element type, so it rescales the running sum directly. +# CHECK: linalg.elementwise kind=#linalg.elementwise_kind
ins(%{{.+}} : tensor<64xf32>, tensor<64xf32>) +# CHECK: %[[SCALED:.+]] = linalg.elementwise kind=#linalg.elementwise_kind +# CHECK-SAME: tensor<64xf32>, tensor<64xf32> +# CHECK: linalg.generic +# CHECK-SAME: ins(%[[P]] : tensor<64x32xf16>) +# CHECK-SAME: outs(%[[SCALED]] : tensor<64xf32>) +# CHECK: arith.addf +# CHECK: } {__reduction_loop__, mixed_softmax_loop} + + +if __name__ == "__main__": + test_softmax_structure() + test_attention_structure() + test_mixed_precision_softmax()