diff --git a/lighthouse/ingress/mlir_gen/gpu_attention_payload.py b/lighthouse/ingress/mlir_gen/gpu_attention_payload.py index 4ca878ff..0dce3c56 100644 --- a/lighthouse/ingress/mlir_gen/gpu_attention_payload.py +++ b/lighthouse/ingress/mlir_gen/gpu_attention_payload.py @@ -3,10 +3,22 @@ import math from mlir import ir -from mlir.dialects import arith, bufferization, linalg, memref, tensor +from mlir.dialects import ( + arith, + bufferization, + linalg, + math as mlir_math, + memref, + tensor, +) from lighthouse.utils.mlir import func_cif -from lighthouse.ingress.mlir_gen.utils import emit_buf_to_tensor +from lighthouse.ingress.mlir_gen.utils import ( + affine_map, + emit_buf_to_tensor, + parallel, + reduction, +) def generate_gpu_attention_payload( @@ -105,24 +117,82 @@ def payload(output, Q_arg, K_arg, V_arg): scaled_qkt_init = tensor.empty(qkt_shape_3d, dtype) 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 numerator, emitted in decomposed form but WITHOUT + # the final division. This is the flash-attention reordering: keep + # the unnormalized probabilities P = exp(x - rowmax) and the per-row + # sum, and defer the division until after the P*V matmul (the row-sum + # normalization commutes past P*V since it is a per-row scalar). + # + # Softmax reduces over the last dimension (j, dim=2). The iteration + # space is 3D (batch, i, j); reductions reduce over j. + reduce_shape_3d = (batch_dim, n_ctx) + d0 = ir.AffineDimExpr.get(0) + d1 = ir.AffineDimExpr.get(1) + d2 = ir.AffineDimExpr.get(2) + par_map_3d = affine_map(3, [d0, d1, d2]) + red_map_3d = affine_map(3, [d0, d1]) + + # Step 4a: Row max reduction: max[b, i] = max_j scaled_qkt[b, i, j]. + neg_inf = arith.constant(dtype, float("-inf")) + max_init = tensor.empty(reduce_shape_3d, dtype) + max_acc = linalg.fill(neg_inf, outs=[max_init]) + + @linalg.generic( + [scaled_qkt], + [max_acc], + [par_map_3d, red_map_3d], + [parallel, parallel, reduction], ) + def row_max(x, acc): + return arith.MaximumFOp(x, acc) - # Step 5: Multiply attention weights by V using batch_matmul - # attention_weights: (batch_dim, n_ctx, n_ctx) @ V: (batch_dim, n_ctx, n_head) + # Step 4b: Numerator P = exp(x - rowmax). + p_init = tensor.empty(qkt_shape_3d, dtype) + + @linalg.generic( + [scaled_qkt, row_max], + [p_init], + [par_map_3d, red_map_3d, par_map_3d], + [parallel, parallel, parallel], + ) + def numerator(x, m, _out): + centered = arith.SubFOp(x, m).result + return mlir_math.exp(centered) + + # Step 4c: Row sum reduction over the numerator: + # sum[b, i] = sum_j P[b, i, j]. + sum_init = tensor.empty(reduce_shape_3d, dtype) + sum_acc = linalg.fill(zero, outs=[sum_init]) + + @linalg.generic( + [numerator], + [sum_acc], + [par_map_3d, red_map_3d], + [parallel, parallel, reduction], + ) + def row_sum(p, acc): + return arith.AddFOp(p, acc) + + # Step 5: Multiply the UNNORMALIZED numerator P by V using batch_matmul. + # P: (batch_dim, n_ctx, n_ctx) @ V: (batch_dim, n_ctx, n_head) # Result: (batch_dim, n_ctx, n_head) - output_3d_init = tensor.empty(collapsed_shape_3d, dtype) - output_3d_init_filled = linalg.fill(zero, outs=[output_3d_init]) + pv_init = tensor.empty(collapsed_shape_3d, dtype) + pv_init_filled = linalg.fill(zero, outs=[pv_init]) + + pv = linalg.batch_matmul(numerator, V_3d, outs=[pv_init_filled]) + + # Step 6: Apply the deferred division by the row sum: + # out[b, i, k] = pv[b, i, k] / sum[b, i]. + result_3d_init = tensor.empty(collapsed_shape_3d, dtype) - result_3d = linalg.batch_matmul( - attention_weights, V_3d, outs=[output_3d_init_filled] + @linalg.generic( + [pv, row_sum], + [result_3d_init], + [par_map_3d, red_map_3d, par_map_3d], + [parallel, parallel, parallel], ) + def result_3d(pv_val, s, _out): + return arith.DivFOp(pv_val, s) # Materialize 3D result back to 3D output memref bufferization.materialize_in_destination( diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index 65a2c91e..fd3f3f58 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -5,11 +5,8 @@ from mlir.dialects.transform import structured, loop, xegpu from mlir.dialects.transform import bufferization as transform_bufferization from mlir.dialects.bufferization import LayoutMapOption -from mlir.dialects.transform.vector import ( - apply_patterns_vector_cast_away_vector_leading_one_dim, - apply_patterns_vector_drop_unit_dims_with_shape_cast, -) +import lighthouse.transform as lh_transform from lighthouse.pipeline.helper import ( canonicalize, match, @@ -18,14 +15,13 @@ apply_registered_pass, ) from lighthouse.schedule import schedule_boilerplate -from lighthouse.dialects.transform.transform_ext import ( - replace_with_fused_attention, -) +from .lowering_common import get_named_func, vectorize def fused_attention_schedule( stop_at_stage: str | None = None, parameters: dict | None = None, + payload_func_name: str = "payload", ) -> ir.Module: """ Generate transform schedule for attention kernel. @@ -48,6 +44,7 @@ def fused_attention_schedule( - wg_rows: Number of Q*K^T*V rows computed by each work group - sg_rows: Number of Q*K^T*V rows computed by each subgroup - subgroup_size: Size of subgroup + payload_func_name: Name of the payload function to transform Returns: MLIR module containing the transform schedule @@ -70,6 +67,7 @@ def fused_attention_schedule( payload_mod, parameters=parameters, stop_at_stage=stop_at_stage or "", + payload_func_name=payload_func_name, ) except PipelineInterrupt: pass @@ -79,202 +77,341 @@ def fused_attention_schedule( return schedule -def bundle_xegpu_fused_attention_schedule( - mod: ir.Value[transform.AnyOpType], +def tile_and_fuse_parallel_dim( + func: ir.Value[transform.AnyOpType], + anytype: transform.AnyOpType, parameters: dict, - stop_at_stage: str = "", -) -> ir.Value[transform.AnyOpType]: - """Schedule for lowering attention payload to xegpu wg level.""" - - if stop_at_stage == "initial": - raise PipelineInterrupt() - - anytype = transform.AnyOpType.get() - # Match all matmul operations - there should be 2: - # 1. Q @ K^T - # 2. attention_weights @ V - matmul_ops = match_and_split(mod, ops={"linalg.batch_matmul"}, nhandles=2) - - # Get the last matmul (attention_weights @ V) - last_matmul = matmul_ops[1] - func = transform.get_parent_op( - anytype, - last_matmul, - op_name="func.func", - deduplicate=True, +) -> None: + """Outer tiling: partition the whole computation along the parallel batch + dimension (batch = Z * H) into an scf.forall over batch tiles. + + The final division op produces the output; its iteration space is + (batch, n_ctx, n_head) and is fully parallel, so tiling it along dim 0 with + tile_using_forall creates an scf.forall over batch tiles. Every producer + (the two matmuls, the two reductions, the scale multiply, the transpose and + the initialization fills) is then fused into that scf.forall so each + workgroup computes attention for its own slice of batch elements. + """ + # Match the structured ops that remain after generalization + fusion. Order + # of the 5 linalg.generic ops in the payload: + # 0: QK^T matmul 1: row max 2: row sum 3: P*V matmul 4: div + qkt, max_reduction, sum_reduction, pv_matmul, div_op = match_and_split( + func, ops={"linalg.generic"}, nhandles=5 ) - - # Tile the last matmul in both batch and M dimensions. - wg_rows = parameters["wg_rows"] - - tiled_matmul, forall_loop = structured.structured_tile_using_forall( + # The scale multiply (linalg.mul) and the transpose still sit between the + # QK^T matmul and the softmax reductions; grab them so we can pull the whole + # QK^T producer chain into the scf.forall too. + scale_mul = match_and_split(func, ops={"linalg.mul"}, nhandles=1)[0] + transpose_op = match_and_split(func, ops={"linalg.transpose"}, nhandles=1)[0] + + wg_tile_size = parameters.get("wg_rows", 128) + _, forall_loop = structured.structured_tile_using_forall( anytype, anytype, - last_matmul, + div_op, num_threads=[], tile_sizes=[], - static_tile_sizes=(1, wg_rows, 0, 0), - ) - # Fuse the zero initialization of the output of the last matmul (tensor.empty) into the forall loop. - tiled_matmul_init = transform.get_producer_of_operand( - anytype, forall_loop, operand_number=0 - ) - _, forall_loop = structured.structured_fuse_into_containing_op( - anytype, - anytype, - producer_op=tiled_matmul_init, - containing_op=forall_loop, + static_tile_sizes=(1, wg_tile_size, 0), ) - transform.apply_cse(func) - canonicalize(func) - - # Decompose softmax into generic ops - softmax_ops = match_and_split(func, ops={"linalg.softmax"}, nhandles=1) - softmax_op = softmax_ops[0] - structured.structured_decompose_interface(anytype, softmax_op) - transform.apply_cse(func) - canonicalize(func) - - # Fuse all linalg.generic ops from softmax decomposition (4 ops: max, sub+exp, sum, div) - # Match and fuse in reverse order (from consumer to producer) - generic_ops = match_and_split(func, ops={"linalg.generic"}, nhandles=4) - for generic_op in reversed(generic_ops): + # Fuse consumers-to-producers. Both the sum reduction and the P*V matmul feed + # the final division; whichever is fused last is cloned closest to the + # division. Fuse the sum reduction before the P*V matmul so the in-loop order + # stays (row max, row sum, P*V). The inner reduction fusion needs the sum to + # precede the P*V matmul: fusing the sum into the tiled max loop requires + # every *other* user of the max result (here the P*V matmul) to post-dominate + # the sum, which only holds if the sum comes first. + for producer in [ + pv_matmul, + sum_reduction, + max_reduction, + scale_mul, + qkt, + transpose_op, + ]: _, forall_loop = structured.structured_fuse_into_containing_op( - anytype, - anytype, - producer_op=generic_op, - containing_op=forall_loop, + anytype, anytype, producer_op=producer, containing_op=forall_loop ) - transform.apply_cse(func) - canonicalize(func) - # Max and add reductions use linalg.fill to intialize the reduction output. Fuse these fill ops as well. + # Also fuse the linalg.fill ops that initialize the matmul/reduction outputs + # into the scf.forall. Besides keeping each workgroup self-contained, this + # rewrites the reduction inits from `tensor.extract_slice(fill)` (outside the + # loop) into `fill(extract_slice(empty))` (inside the loop), so the + # dependant-reduction fusion's zero-init legality check can see through them. + # There are 5 fills: QK^T out, scale tensor, max init, sum init, P*V out. fill_ops = match_and_split(func, ops={"linalg.fill"}, nhandles=5) - # Max fill is the third fill op and add fill is the fourth fill op (based on the pattern of decomposition) - max_fill_op = fill_ops[2] - add_fill_op = fill_ops[3] - for fill_op in [max_fill_op, add_fill_op]: + for fill_op in fill_ops: _, forall_loop = structured.structured_fuse_into_containing_op( - anytype, - anytype, - producer_op=fill_op, - containing_op=forall_loop, + anytype, anytype, producer_op=fill_op, containing_op=forall_loop ) transform.apply_cse(func) canonicalize(func) - # Fuse the remaining operations into the scf.forall loop. - linalg_mul_op = match_and_split(func, ops={"linalg.mul"}, nhandles=1)[0] - first_matmul = transform.get_producer_of_operand( - anytype, linalg_mul_op, operand_number=0 + +def tile_and_fuse_reduction_dim( + func: ir.Value[transform.AnyOpType], + anytype: transform.AnyOpType, + parameters: dict, +) -> None: + """Build the fused inner reduction loop (online softmax) inside the batch + scf.forall. + + Within each batch tile, tile the row-max reduction along the K/V sequence + (reduction) dimension and fuse the dependant sum reduction and the P*V + matmul into that loop, then pull the QK^T producer chain in as well. The + result is a single scf.for reduction loop implementing online softmax. + """ + # Re-match the ops now living inside the scf.forall. Fusing producers with + # multiple uses can duplicate them, so threading handles through the outer + # fusion is unreliable; matching by op type inside the loop is robust. After + # outer tiling the 5 linalg.generic ops appear in this order: + # 0: QK^T matmul 1: row max 2: row sum 3: P*V matmul 4: div + qkt, max_reduction, sum_reduction, pv_matmul, _div = match_and_split( + func, ops={"linalg.generic"}, nhandles=5 ) - scale_fill_op = transform.get_producer_of_operand( - anytype, linalg_mul_op, operand_number=1 + scale_mul = match_and_split(func, ops={"linalg.mul"}, nhandles=1)[0] + transpose_op = match_and_split(func, ops={"linalg.transpose"}, nhandles=1)[0] + + # Grab the QK^T-out and scale fills so step 6 can sink them into the scf.for + # alongside the QK^T producer chain. In program order after outer tiling the + # 5 fills are: 0 QK^T out, 1 scale tensor, 2 row-max init, 3 P*V out, + # 4 row-sum init. Only fills 0 and 1 belong to the transient QK^T block; the + # other three are loop-carried softmax accumulators and stay in the forall. + # These two handles are not touched by the reduction tiling in steps 3-5 + # (they feed qkt/scale_mul, which are only fused in step 6), so they survive + # just like the scale_mul/qkt handles above. + qkt_out_fill, scale_fill, _, _, _ = match_and_split( + func, ops={"linalg.fill"}, nhandles=5 ) - transpose_op = transform.get_producer_of_operand( - anytype, first_matmul, operand_number=1 + + reduction_step_size = parameters.get("inner_loop_tile_size", 64) + + # Step 3: Tile the producer reduction (row max) along its reduction dim (the + # K/V sequence length, the last iteration dim of the 3D iteration space) and + # annotate the resulting loop so the fusion op recognises it as a tiled + # reduction. + _, max_loop = structured.TileUsingForOp( + max_reduction, sizes=[0, 0, reduction_step_size] + ).results + transform.annotate(max_loop, "__reduction_loop__") + + # Step 4: Fuse the first dependant consumer into the tiled max loop. Note the + # handles: `sum_reduction`/`pv_matmul` above are named for the payload order + # in the comment, but positions 2 and 3 are actually the P*V matmul and the + # row sum respectively, so this call fuses the P*V and returns it as the + # in-loop reduction handle. + fused_loop, fused_pv = structured.structured_tile_and_fuse_dependant_reduction_ops( + anytype, anytype, max_loop, sum_reduction ) - matmul_fill_op = transform.get_producer_of_operand( - anytype, first_matmul, operand_number=2 + transform.annotate(fused_loop, "__reduction_loop__") + + # Step 5: Fuse the remaining dependant consumer (the row sum) into the loop. + # Both fusions have to complete before any unfusion: the unfusion interposes a + # new value between a fused reduction and the loop's yield, and the fusion + # pattern requires every input of the tiled producer to also be an input of + # the consumer, which that new value breaks. + fused_loop, _ = structured.structured_tile_and_fuse_dependant_reduction_ops( + anytype, anytype, fused_loop, pv_matmul ) - for op in [ - linalg_mul_op, - scale_fill_op, - first_matmul, - matmul_fill_op, - transpose_op, - ]: - _, forall_loop = structured.structured_fuse_into_containing_op( + + # Step 5b: Hoist the softmax epilogue out of the fused P*V reduction, leaving + # it with just the contraction `mulf` and its `addf` combiner. Only the + # `subf`/`exp` prologue moves out -- the contraction multiply spans more + # iteration dims than any of its operands, so it stays fused and no + # rank-increasing intermediate is materialized. The row sum is left fused on + # purpose. + # + # The `fused_pv` handle from step 4 was invalidated by the step 5 fusion + # (which consumes the loop the P*V now lives in), so re-match inside the final + # loop. It holds 7 generics at this point: #0 the row max, #1/#2 and #4/#5 the + # all-parallel online-correction ops, #3 the P*V contraction (the only 4-loop + # reduction, carrying the softmax chain) and #6 the row sum. + loop_generics = match_and_split(fused_loop, ops={"linalg.generic"}, nhandles=7) + structured.structured_unfuse_elementwise_from(anytype, anytype, loop_generics[3]) + + transform.apply_cse(func) + canonicalize(func) + + # Step 6: Fuse the QK^T producer chain into the scf.for reduction loop. The + # loop slices the scaled QK^T tensor, so fuse from the closest producer + # outward: the scale multiply, then the QK^T matmul, then the K transpose. + # The QK^T-out and scale fills are pulled in right after the op that consumes + # them (qkt reads qkt_out_fill as its init, scale_mul reads scale_fill as an + # input) so they become per-iteration 1x128x64 inits instead of full-tile + # 1x128x4096 buffers. After vectorization each collapses to its splat + # constant, eliminating the two 1x128x4096 scratch allocs entirely. + for producer in [scale_mul, scale_fill, qkt, qkt_out_fill, transpose_op]: + _, fused_loop = structured.structured_fuse_into_containing_op( anytype, anytype, - producer_op=op, - containing_op=forall_loop, + producer_op=producer, + containing_op=fused_loop, ) transform.apply_cse(func) canonicalize(func) + +def bundle_xegpu_fused_attention_schedule( + mod: ir.Value[transform.AnyOpType], + parameters: dict, + stop_at_stage: str = "", + payload_func_name: str = "payload", +) -> ir.Value[transform.AnyOpType]: + """Schedule for lowering attention payload to xegpu wg level.""" + + if stop_at_stage == "initial": + raise PipelineInterrupt() + + anytype = transform.AnyOpType.get() + + # The payload emits attention with softmax already in decomposed form and the + # final division deferred until after the P*V matmul (flash-attention style): + # 1. QK^T (batch_matmul) + # 2. scale (linalg.mul) + # 3. row max (generic reduction) + # 4. P = exp(x-max) (generic elementwise, the softmax numerator) + # 5. row sum (generic reduction over P) + # 6. P*V (batch_matmul) + # 7. out = PV / sum (generic elementwise, the deferred division) + # + # The computation is tiled at two levels: an outer scf.forall over the + # parallel batch dimension (tile_and_fuse_parallel_dim), and an inner + # scf.for over the softmax/P*V reduction dimension implementing online + # softmax via the dependant-reduction fusion (tile_and_fuse_reduction_dim). + + # Match both matmuls (QK^T and P*V) and scope the rest of the schedule to + # the payload function that contains them. + matmul_ops = match(mod, ops={"linalg.batch_matmul"}) + func = transform.get_parent_op( + anytype, + matmul_ops, + op_name="func.func", + deduplicate=True, + ) + # transform.print_(target=func, name="initial payload function") + + # Step 1: Generalize both matmuls (QK^T and P*V) into linalg.generic ops so + # the reduction-fusion machinery can operate on them uniformly. + structured.structured_generalize(anytype, matmul_ops) + transform.apply_cse(func) + canonicalize(func) + + # Step 2: Fuse elementwise ops into their producers/consumers. This folds the + # softmax numerator P = exp(x - max) into both the sum reduction and the P*V + # matmul, leaving the two matmuls, the two reductions and the final division + # as the remaining structured ops. + func = apply_registered_pass(func, "linalg-fuse-elementwise-ops") + transform.apply_cse(func) + canonicalize(func) + + # Outer tiling: partition the whole computation along the parallel batch + # dimension so each workgroup computes attention for its own batch slice. + tile_and_fuse_parallel_dim(func, anytype, parameters) + # transform.print_(target=func, name="after outer tiling and fusion") + if stop_at_stage == "outer-tiled": raise PipelineInterrupt() - # Vectorize - func = structured.VectorizeChildrenAndApplyPatternsOp( + # Inner tiling: within each batch tile, build a single fused reduction loop + # over the softmax / P*V reduction dimension. Fusing the linalg.fill ops into + # the scf.forall above rewrote the reduction inits into in-loop + # `linalg.fill`s (rather than `tensor.extract_slice` of an outside fill), so + # the dependant-reduction fusion's zero-init legality check now succeeds. + tile_and_fuse_reduction_dim(func, anytype, parameters) + # Drop the batch unit extent (tile size 1) from the linalg ops in the loop + # nest before vectorizing. Tiling the batch dim by 1 left every op with a + # leading 1x... shape; folding it here (via rank-reducing slices) means + # vectorization emits 128x64 / 128 vectors directly instead of 1x128x64 + # vectors wrapped in shape_casts, and bufferization allocates unit-dim-free + # memrefs (e.g. memref<128xf16> instead of memref<1x128xf16>). + # The fold_unit_extent_dims patterns only rewrite linalg.generic ops, so + # first generalize the remaining named/category ops in the loop nest (the K + # transpose, the scale linalg.mul and the div/mul linalg.elementwise ops) + # into generics. Without this the unit dim survives on exactly those ops. + func = apply_registered_pass( func, - fold_type_extensions_into_contract=True, - ).result + "linalg-morph-ops", + options={"category-to-generic": True}, + ) transform.apply_cse(func) canonicalize(func) - # Try to remove any unit dimensions that may have been introduced due to tiling (e.g. batch dim of 1) + # transform.print_(target=func, name="after generalizing") + with ir.InsertionPoint(transform.apply_patterns(func).patterns): - apply_patterns_vector_cast_away_vector_leading_one_dim() - apply_patterns_vector_drop_unit_dims_with_shape_cast() + structured.apply_patterns_linalg_fold_unit_extent_dims_via_slices() + structured.apply_patterns_linalg_fold_unit_extent_dims_via_reshapes() + transform.apply_cse(func) + canonicalize(func) + # transform.print_(target=func, name="after inner tiling and fusion") + + if stop_at_stage == "inner-tiled": + raise PipelineInterrupt() + # Vectorize the fused loop nest: rewrite the remaining linalg ops (the tiled + # matmuls, reductions and elementwise ops inside the scf.for) into vector + # ops. The shared helper also runs loop hoisting (LICM plus + # hoist_loop_invariant_subsets) on the scf.for loops, which pulls the + # loop-invariant vector transfers -- notably the Q tile read, invariant + # across the K/V reduction loop -- out of the reduction loop. + func = vectorize(mod, payload_func_name) + + # Try to remove any unit dimensions that may have been introduced due to tiling (e.g. batch dim of 1) + lh_transform.simplify_vector_ops(func) + + # Hoist again after the unit-dim cleanup. The row-max accumulator is carried + # as a tensor<1x128xf16> through the reduction loop and read/written by a + # transfer_read/transfer_write pair per iteration; those transfers only + # become a hoistable matching subset pair once the leading unit dim is gone. + # Hoisting them turns the accumulator into a plain vector<128xf16> iter_arg, + # which also lets the `tensor.empty` backing it (and its per-workgroup slice) + # fold away entirely. + reduction_loops = match(func, ops={"scf.for"}) + lh_transform.loop_hoisting(reduction_loops) + lh_transform.cleanup(func) + + # transform.print_(target=func, name="after vectorization") if stop_at_stage == "vectorized": raise PipelineInterrupt() - # Bufferize + # Bufferize: convert tensors to memrefs across function boundaries with an + # identity layout map, then fold memref.subviews into the vector transfer + # ops. mod = apply_registered_pass(mod, "eliminate-empty-tensors") identity_layout = LayoutMapOption.IdentityLayoutMap mod = transform_bufferization.OneShotBufferizeOp( mod, - allow_return_allocs_from_loops=True, + allow_return_allocs_from_loops=False, bufferize_function_boundaries=True, function_boundary_type_conversion=identity_layout, ).result - # fold memref.subviews into vector.transfer_read/write ops mod = apply_registered_pass(mod, "fold-memref-alias-ops") transform.apply_cse(mod) canonicalize(mod) - if stop_at_stage == "bufferized": - raise PipelineInterrupt() - - # Extract q, k, v memrefs from the bufferized IR - # Match vector.contract ops to find the q, k, v loads - for_all = match(mod, ops={"scf.forall"}) - func = transform.get_parent_op(anytype, for_all, op_name="func.func") - contract_ops = match_and_split(func, ops={"vector.contract"}, nhandles=2) - - # First vector.contract is Q @ K^T - # Its first operand is the q load (vector.transfer_read) - # Its second operand is the k load (vector.transfer_read) - first_contract = contract_ops[0] - q_load = transform.get_producer_of_operand( - anytype, first_contract, operand_number=0 - ) - k_load = transform.get_producer_of_operand( - anytype, first_contract, operand_number=1 - ) - - # Second vector.contract is attention_weights @ V - # Its second operand is the v load (vector.transfer_read) - second_contract = contract_ops[1] - v_load = transform.get_producer_of_operand( - anytype, second_contract, operand_number=1 - ) - - # Match arith.mulf to get the scale parameter - # The scale is the second operand of arith.mulf (the constant) - mulf_op = match_and_split(func, ops={"arith.mulf"}, nhandles=1)[0] - scale = transform.get_producer_of_operand(anytype, mulf_op, operand_number=1) - - # Apply the fused attention optimization. This replaces the second vector.contract - # (attention_weights @ V) with a tiled loop that implements online softmax for - # efficient memory usage - tile_size = parameters.get( - "inner_loop_tile_size", 64 - ) # Tile size for reduction dimension (K/V sequence length) - replace_with_fused_attention( - q_load=q_load, - k_load=k_load, - v_load=v_load, - scale=scale, - output=second_contract, - tile_size=tile_size, - ) + # Promote small memref.allocs (the per-workgroup scratch buffers) to the + # stack in the payload function. Not needed while the loop hoisting above + # keeps all accumulators in registers -- the bufferized payload has no + # memref.allocs left at all. + # func = match(mod, ops={"func.func"}) + # func = apply_registered_pass( + # func, + # "promote-buffers-to-stack", + # options={ + # "max-alloc-size-in-bytes": "16384", + # "max-rank-of-allocated-memref": "3", + # }, + # ) + + # Bufferization replaced the payload module, so re-match the payload function + # instead of reusing the pre-bufferization handle. Hoist what became loop + # invariant once tensors turned into memrefs: the per-iteration subview and + # index computations feeding the vector transfers. + func = get_named_func(mod, payload_func_name) + func = apply_registered_pass(func, "loop-invariant-code-motion") transform.apply_cse(func) canonicalize(func) - if stop_at_stage == "inner-tiled": + # transform.print_(target=func, name="after bufferization") + if stop_at_stage == "bufferized": raise PipelineInterrupt() # Convert forall to parallel @@ -306,6 +443,7 @@ def bundle_xegpu_fused_attention_schedule( mod = apply_registered_pass(mod, "gpu-kernel-outlining") transform.apply_cse(mod) + # transform.print_(target=mod, name="after gpu kernel outlining") if stop_at_stage == "gpu-outlining": raise PipelineInterrupt() @@ -320,6 +458,8 @@ def bundle_xegpu_fused_attention_schedule( gpu_mod_ops = match_and_split(mod, ops={"gpu.module"}) for gpu_mod in gpu_mod_ops: gpu_func = match(gpu_mod, ops={"gpu.func"}) + # allocas = match(gpu_func, ops={"memref.alloca"}) + # transform_ext.update_address_space(allocas, address_space=3) gpu_func = apply_registered_pass(gpu_func, "convert-vector-to-xegpu") transform.apply_cse(gpu_func) gpu_func = apply_registered_pass(gpu_func, "loop-invariant-code-motion") @@ -334,7 +474,7 @@ def bundle_xegpu_fused_attention_schedule( q_inst_data = [8, 16] k_sg_layout = [num_subgroups, 1] - k_sg_data = [16, n_head] + k_sg_data = [64, n_head] k_inst_data = [16, 16] v_sg_layout = k_sg_layout @@ -342,7 +482,7 @@ def bundle_xegpu_fused_attention_schedule( v_inst_data = k_inst_data kt_sg_layout = [1, num_subgroups] - kt_sg_data = [n_head, 16] + kt_sg_data = [n_head, 64] kt_inst_data = [16, 16] kt_order = [0, 1] @@ -350,13 +490,13 @@ def bundle_xegpu_fused_attention_schedule( out_sg_data = q_sg_data out_inst_data = q_inst_data - layout_128x16_sg_layout = [num_subgroups, 1] - layout_128x16_sg_data = [16, 16] - layout_128x16_inst_data = [8, 16] + layout_128x64_sg_layout = [num_subgroups, 1] + layout_128x64_sg_data = [16, 64] + layout_128x64_inst_data = [8, 16] - qk_sg_layout = layout_128x16_sg_layout - qk_sg_data = layout_128x16_sg_data - qk_inst_data = layout_128x16_inst_data + qk_sg_layout = layout_128x64_sg_layout + qk_sg_data = layout_128x64_sg_data + qk_inst_data = layout_128x64_inst_data # Set layout attributes for xegpu.store_nd ops. store_nd_op = match_and_split(gpu_func, ops={"xegpu.store_nd"}, nhandles=1)[0] @@ -368,88 +508,86 @@ def bundle_xegpu_fused_attention_schedule( ) # Set layout for xegpu.load_nd ops (9 total: 1 Q, 4 K, 4 V) - load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=9) + load_nd_ops = match_and_split(gpu_func, ops={"xegpu.load_nd"}, nhandles=3) # First load_nd: Q layout xegpu.set_anchor_layout( load_nd_ops[0], sg_layout=q_sg_layout, sg_data=q_sg_data, inst_data=q_inst_data ) - # Next 4 load_nd ops: K layout - for load_op in load_nd_ops[:4]: - xegpu.set_anchor_layout( - load_op, - sg_layout=k_sg_layout, - sg_data=k_sg_data, - inst_data=k_inst_data, - ) + # Next load is K load + xegpu.set_anchor_layout( + load_nd_ops[1], + sg_layout=k_sg_layout, + sg_data=k_sg_data, + inst_data=k_inst_data, + ) - # Last 4 load_nd ops: V layout - for load_op in load_nd_ops[4:]: - xegpu.set_anchor_layout( - load_op, - sg_layout=v_sg_layout, - sg_data=v_sg_data, - inst_data=v_inst_data, - ) + # Last load is V load + xegpu.set_anchor_layout( + load_nd_ops[2], + sg_layout=v_sg_layout, + sg_data=v_sg_data, + inst_data=v_inst_data, + ) - # Set layout for xegpu.dpas ops (8 total: 4 for Q@K, 4 for P@V) - dpas_ops = match_and_split(gpu_func, ops={"xegpu.dpas"}, nhandles=8) - - # Layouts for first 4 dpas ops (Q@K^T): - for qk_dpas_op in dpas_ops[:4]: - # Index 0: Q layout - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=q_sg_layout, - sg_data=q_sg_data, - inst_data=q_inst_data, - index=0, - ) - # Index 1: K^T layout - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=kt_sg_layout, - sg_data=kt_sg_data, - inst_data=kt_inst_data, - order=kt_order, - index=1, - ) - # Index 2: QK output layout (128x16) - xegpu.set_anchor_layout( - qk_dpas_op, - sg_layout=layout_128x16_sg_layout, - sg_data=layout_128x16_sg_data, - inst_data=layout_128x16_inst_data, - index=2, - ) + # Set layout for xegpu.dpas ops (2 total: 1 for Q@K, 1 for P@V) + dpas_ops = match_and_split(gpu_func, ops={"xegpu.dpas"}, nhandles=2) + qk_dpas_op = dpas_ops[0] + pv_dpas_op = dpas_ops[1] - # Layouts for second 4 dpas ops (P@V): - for pv_dpas_op in dpas_ops[4:]: - # Index 0: QK (attention weights) layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=qk_sg_layout, - sg_data=qk_sg_data, - inst_data=qk_inst_data, - index=0, - ) - # Index 1: V layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=v_sg_layout, - sg_data=v_sg_data, - inst_data=v_inst_data, - index=1, - ) - # Index 2: Output layout - xegpu.set_anchor_layout( - pv_dpas_op, - sg_layout=out_sg_layout, - sg_data=out_sg_data, - inst_data=out_inst_data, - index=2, - ) + # Layouts for first dpas (Q@K^T): + # Index 0: Q layout + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=q_sg_layout, + sg_data=q_sg_data, + inst_data=q_inst_data, + index=0, + ) + # Index 1: K^T layout + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=kt_sg_layout, + sg_data=kt_sg_data, + inst_data=kt_inst_data, + order=kt_order, + index=1, + ) + # Index 2: QK output layout (128x16) + xegpu.set_anchor_layout( + qk_dpas_op, + sg_layout=layout_128x64_sg_layout, + sg_data=layout_128x64_sg_data, + inst_data=layout_128x64_inst_data, + index=2, + ) + + # Layouts for second dpas op (P@V): + # Index 0: QK (attention weights) layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=qk_sg_layout, + sg_data=qk_sg_data, + inst_data=qk_inst_data, + index=0, + ) + # Index 1: V layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=v_sg_layout, + sg_data=v_sg_data, + inst_data=v_inst_data, + index=1, + ) + # Index 2: Output layout + xegpu.set_anchor_layout( + pv_dpas_op, + sg_layout=out_sg_layout, + sg_data=out_sg_data, + inst_data=out_inst_data, + index=2, + ) if stop_at_stage == "xegpu-wg": raise PipelineInterrupt() diff --git a/lighthouse/schedule/xegpu/softmax_schedule.py b/lighthouse/schedule/xegpu/softmax_schedule.py index b573393e..bc739cd3 100644 --- a/lighthouse/schedule/xegpu/softmax_schedule.py +++ b/lighthouse/schedule/xegpu/softmax_schedule.py @@ -110,68 +110,39 @@ def bundle_xegpu_softmax_schedule( ) structured.structured_decompose_interface(anytype, softmax_ops) - linalg_ops = match_and_split( - func, ops={"linalg.generic", "linalg.fill"}, nhandles=6 - ) - max_reduction = linalg_ops[1] - max_center_and_exp_op = linalg_ops[2] - sum_reduction = linalg_ops[4] - div_op = linalg_ops[5] + # Fuse the elementwise center-and-exp op into its consumers so that only + # the three reductions/division generics remain: max reduction, sum + # reduction and the final division. + func = apply_registered_pass(func, "linalg-fuse-elementwise-ops") reduction_step_size = parameters["reduction_step_size"] - # Tile the division op first. - _, div_loop = structured.TileUsingForOp( - div_op, sizes=[0, reduction_step_size] - ).results - - # Fuse max_center_and_exp_op into the div loop - _, fused_loop = structured.structured_fuse_into_containing_op( - anytype, - anytype, - producer_op=max_center_and_exp_op, - containing_op=div_loop, + # Match the three remaining linalg.generic ops. + max_reduction, sum_reduction, div_op = match_and_split( + func, ops={"linalg.generic"}, nhandles=3 ) - # Tile the sum reduction. - _, _, _, sum_loop = structured.structured_tile_reduction_using_for( - [anytype], - anytype, - anytype, - anytype, - target=sum_reduction, - tile_sizes=[0, reduction_step_size], - ) + # Tile the producer reduction (max) along its reduction dim and annotate the + # resulting loop so the fusion op can recognise it as a tiled reduction. + _, max_loop = structured.TileUsingForOp( + max_reduction, sizes=[0, reduction_step_size] + ).results + transform.annotate(max_loop, "__reduction_loop__") - func = transform.get_parent_op( - anytype, - fused_loop, - op_name="func.func", - deduplicate=True, + # Fuse the consumer reduction (sum) into the tiled max reduction loop. This + # returns the fused loop and the sum reduction as it now exists inside it + # (the fusion clones it into the loop and erases the original). + _, fused_sum = structured.structured_tile_and_fuse_dependant_reduction_ops( + anytype, anytype, max_loop, sum_reduction ) - # Re-match and split linalg generic ops, there are 5 at this point - linalg_ops = match_and_split(func, ops={"linalg.generic"}, nhandles=5) - max_center_and_exp_op = linalg_ops[1] + # The fused sum still carries the center-and-exp elementwise chain inlined in + # its reduction body. Split it back out into a separate all-parallel op so the + # reduction keeps only its `addf` combiner. + structured.structured_unfuse_elementwise_from(anytype, anytype, fused_sum) - # Fuse max_center_and_exp_op into the sum reduction loop - _, fused_sum_loop = structured.structured_fuse_into_containing_op( - anytype, - anytype, - producer_op=max_center_and_exp_op, - containing_op=sum_loop, - ) - - # Tile the max reduction. - max_reduction = linalg_ops[0] - structured.structured_tile_reduction_using_for( - [anytype], - anytype, - anytype, - anytype, - target=max_reduction, - tile_sizes=[0, reduction_step_size], - ) + # Tile the division op. + structured.TileUsingForOp(div_op, sizes=[0, reduction_step_size]) # Cleanup after tiling and fusion transform.apply_cse(func) @@ -196,7 +167,7 @@ def bundle_xegpu_softmax_schedule( identity_layout = LayoutMapOption.IdentityLayoutMap mod = transform_bufferization.OneShotBufferizeOp( mod, - allow_return_allocs_from_loops=True, + allow_return_allocs_from_loops=False, bufferize_function_boundaries=True, function_boundary_type_conversion=identity_layout, ).result @@ -274,13 +245,10 @@ def bundle_xegpu_softmax_schedule( # Set layout attributes for xegpu.store_nd and xegpu.store_matrix ops. store_nd_ops = match_and_split(gpu_func, ops={"xegpu.store_nd"}, nhandles=1) - store_matrix_ops = match_and_split(gpu_func, ops={"xegpu.store_matrix"}, nhandles=4) sg_layout = [parameters["sg_rows"], 1] sg_data = [parameters["sg_rows"], parameters["reduction_step_size"]] for store_op in store_nd_ops: xegpu.set_anchor_layout(store_op, sg_layout=sg_layout, sg_data=sg_data) - for store_op in store_matrix_ops: - xegpu.set_anchor_layout(store_op, sg_layout=sg_layout, sg_data=sg_data) if stop_at_stage == "xegpu-wg": raise PipelineInterrupt()