diff --git a/lighthouse/dialects/transform/transform_ext/__init__.py b/lighthouse/dialects/transform/transform_ext/__init__.py index 45aa49d0..fdfc3713 100644 --- a/lighthouse/dialects/transform/transform_ext/__init__.py +++ b/lighthouse/dialects/transform/transform_ext/__init__.py @@ -19,6 +19,7 @@ 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.filter_contraction_ops import filter_contraction_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 @@ -32,6 +33,7 @@ "convert_func_results_to_args", "extract_handle", "filter_by_name", + "filter_contraction_ops", "filter_elementwise", "filter_num_loops", "filter_reduction_ops", diff --git a/lighthouse/dialects/transform/transform_ext/ops/filter_contraction_ops.py b/lighthouse/dialects/transform/transform_ext/ops/filter_contraction_ops.py new file mode 100644 index 00000000..c8342839 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/ops/filter_contraction_ops.py @@ -0,0 +1,76 @@ +from mlir import ir +from mlir.dialects import transform, linalg + +from lighthouse.dialects.transform.transform_ext.utils.make_filter_handles_op import ( + make_filter_handles_op, +) +from lighthouse.utils.mlir import opview, indexing_maps, linalg_inputs, dim_position + + +def _map_dims(m: ir.AffineMap) -> set[int]: + """Set of iteration-dim positions referenced by an affine map's results.""" + dims = set() + for r in m.results: + pos = dim_position(r) + if pos is not None: + dims.add(pos) + return dims + + +def _is_structural_contraction(ov: ir.OpView) -> bool: + """Detect a matmul-like op from its indexing maps alone (body-agnostic). + + A contraction has a reduction dimension (one that is absent from the + output) that is shared by at least two inputs. This ignores the body, so it + also matches contractions with extra elementwise operands (e.g. a scaled or + dequantized input) that the strict multiply-accumulate matcher rejects. + """ + maps = indexing_maps(ov) + inputs = linalg_inputs(ov) + if not maps or not inputs: + return False + num_inputs = len(inputs) + input_maps = maps[:num_inputs] + output_maps = maps[num_inputs:] + + out_dims: set[int] = set() + for m in output_maps: + out_dims |= _map_dims(m) + + # Reduction dims are the iteration dims not present in any output. + reduction_dims = set(range(maps[0].n_dims)) - out_dims + for k in reduction_dims: + if sum(1 for m in input_maps if k in _map_dims(m)) >= 2: + return True + return False + + +def is_contraction_op(op: ir.Operation | ir.OpView) -> bool: + """Check whether the op is a linalg contraction (matmul-like) op. + + Recognizes both named contractions (e.g. linalg.batch_matmul) and their + generic form, independent of rank. A contraction contracts a reduction + dimension shared by two inputs, which distinguishes it from the surrounding + elementwise or plain (single-input) reduction ops. + """ + ov = opview(op) + if "linalg" not in ov.operation.name: + return False + return linalg.isa_contraction_op(ov) or _is_structural_contraction(ov) + + +FilterContractionOpsOp = make_filter_handles_op( + "filter_contraction_ops", is_contraction_op +) + + +def filter_contraction_ops(target: ir.Value[transform.AnyOpType]) -> ir.Value: + """ + snake_case wrapper to create a FilterContractionOpsOp. + + Args: + target: Handle to target op(s). + Returns: + Handle to the contraction-op subset of `target`. + """ + return FilterContractionOpsOp(target=target).ops diff --git a/lighthouse/dialects/transform/transform_ext/ops/trace_producers.py b/lighthouse/dialects/transform/transform_ext/ops/trace_producers.py index a87b659d..2aec0bd1 100644 --- a/lighthouse/dialects/transform/transform_ext/ops/trace_producers.py +++ b/lighthouse/dialects/transform/transform_ext/ops/trace_producers.py @@ -41,13 +41,13 @@ def apply( # Walk the SSA producer graph via operand -> owner edges. # Use BFS to guarantee closest-first ordering by graph distance. producers: list[ir.Operation] = [] - visited_ids: set[int] = set() + visited: set[ir.Operation] = set() worklist = deque() for operand in leaf.operands: owner_op = defining_op(operand) - if owner_op is not None and id(owner_op) not in visited_ids: - visited_ids.add(id(owner_op)) + if owner_op is not None and owner_op not in visited: + visited.add(owner_op) worklist.append(owner_op) while worklist: @@ -56,8 +56,8 @@ def apply( for operand in producer.operands: owner_op = defining_op(operand) - if owner_op is not None and id(owner_op) not in visited_ids: - visited_ids.add(id(owner_op)) + if owner_op is not None and owner_op not in visited: + visited.add(owner_op) worklist.append(owner_op) results.set_ops(op.ops, producers) diff --git a/lighthouse/schedule/xegpu/fused_attention_schedule.py b/lighthouse/schedule/xegpu/fused_attention_schedule.py index 5595f077..950a1a9a 100644 --- a/lighthouse/schedule/xegpu/fused_attention_schedule.py +++ b/lighthouse/schedule/xegpu/fused_attention_schedule.py @@ -151,6 +151,13 @@ def bundle_xegpu_fused_attention_schedule( # 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) + # Convert linalg.mul and linalg.batch_matmul to linalg.generic + structured.structured_generalize( + anytype, + structured.structured_match( + anytype, func, ops=["linalg.mul", "linalg.batch_matmul"] + ), + ) # Normalize possible singleton dimensions so tile+fuse logic works. with ir.InsertionPoint(transform.apply_patterns(func).patterns): @@ -197,20 +204,51 @@ def bundle_xegpu_fused_attention_schedule( # 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) + + linalg_ops = structured.structured_match( + anytype, func, ops=["linalg.generic", "linalg.batch_matmul"] + ) + contraction_ops = transform_ext.filter_contraction_ops(linalg_ops) + + # Match max reduction op. Assumes there's only one arith.max* op. + arith_max_op = match_and_split( + func, ops=["arith.maximumf", "arith.maxnumf"], nhandles=1 + )[0] + max_reduction = transform.get_parent_op( + anytype, arith_max_op, op_name="linalg.generic" + ) + + def get_producers_by_name(target, op_names): + producers = transform_ext.trace_producers(target) + return transform_ext.filter_by_name(producers, op_names=op_names) + + # Trace the scalar from max reduction producer chain. + max_producer_generics = get_producers_by_name( + max_reduction, op_names="linalg.generic" + ) + # Assume the scaling happens in the first ancestor. + max_producer = transform_ext.extract_handle(max_producer_generics, 0) + max_scale_mul_op = match_and_split(max_producer, ops={"arith.mulf"}, nhandles=1)[0] + scale_producers = transform_ext.trace_producers(max_scale_mul_op) + scale_const_op = transform_ext.extract_handle( + transform_ext.filter_by_name(scale_producers, op_names="arith.constant"), 0 + ) + + matmul_ops = transform.split_handle(2 * [anytype], contraction_ops) 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 + # Find the tensor.extract_slice producers for the Q@K^T matmul. + qk_extract_slice_producers = get_producers_by_name( + qk_matmul, op_names="tensor.extract_slice" ) - k = transform.get_producer_of_operand(anytype, k_transpose, operand_number=0) - v = transform.get_producer_of_operand(anytype, pv_matmul, operand_number=1) + q = transform_ext.extract_handle(qk_extract_slice_producers, 0) + k = transform_ext.extract_handle(qk_extract_slice_producers, 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) + # Find handle to v as the first tensor.extract_slice producer of PV matmul + pv_extract_slice_producers = get_producers_by_name( + pv_matmul, op_names="tensor.extract_slice" + ) + v = transform_ext.extract_handle(pv_extract_slice_producers, 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. @@ -221,7 +259,7 @@ def bundle_xegpu_fused_attention_schedule( q=q, k=k, v=v, - scale=scale, + scale=scale_const_op, output=pv_matmul, tile_size=reduction_tile, ) diff --git a/test/transform/test_filter_contraction_ops.py b/test/transform/test_filter_contraction_ops.py new file mode 100644 index 00000000..e6e1324e --- /dev/null +++ b/test/transform/test_filter_contraction_ops.py @@ -0,0 +1,102 @@ +# RUN: %PYTHON %s | FileCheck %s + +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 + + +def apply_filter(payload: str, name: str): + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + module = ir.Module.parse(payload) + with schedule_boilerplate() as (sched, named_seq): + candidates = lh_transform.match_op( + named_seq.bodyTarget, structured.MatchInterfaceEnum.LinalgOp + ) + filtered = transform_ext.filter_contraction_ops(candidates) + transform.print_(target=filtered, name=name) + transform.yield_() + sched.body.operations[0].apply(module.operation) + + +# Named matmul and batch_matmul mixed with elementwise ops. +NAMED = """ +#id = affine_map<(d0, d1) -> (d0, d1)> +module { + func.func @main( + %a: tensor<8x8xf32>, + %m0: tensor<4x8xf32>, + %m1: tensor<8x4xf32>, + %x: tensor<2x4x8xf32>, + %y: tensor<2x8x4xf32>) + -> (tensor<8x8xf32>, tensor<4x4xf32>, tensor<2x4x4xf32>) { + %e0 = tensor.empty() : tensor<8x8xf32> + %add = linalg.add ins(%a, %a : tensor<8x8xf32>, tensor<8x8xf32>) + outs(%e0 : tensor<8x8xf32>) -> tensor<8x8xf32> + + %e1 = tensor.empty() : tensor<8x8xf32> + %gen = linalg.generic {indexing_maps = [#id, #id], iterator_types = ["parallel", "parallel"]} + ins(%a : tensor<8x8xf32>) + outs(%e1 : tensor<8x8xf32>) { + ^bb0(%i: f32, %o: f32): + linalg.yield %i : f32 + } -> tensor<8x8xf32> + + %e2 = tensor.empty() : tensor<4x4xf32> + %mm = linalg.matmul ins(%m0, %m1 : tensor<4x8xf32>, tensor<8x4xf32>) + outs(%e2 : tensor<4x4xf32>) -> tensor<4x4xf32> + + %e3 = tensor.empty() : tensor<2x4x4xf32> + %bm = linalg.batch_matmul ins(%x, %y : tensor<2x4x8xf32>, tensor<2x8x4xf32>) + outs(%e3 : tensor<2x4x4xf32>) -> tensor<2x4x4xf32> + return %gen, %mm, %bm : tensor<8x8xf32>, tensor<4x4xf32>, tensor<2x4x4xf32> + } +} +""" + +# CHECK-LABEL: IR printer: NAMED +# CHECK: linalg.matmul +# CHECK: linalg.batch_matmul +# CHECK-NOT: linalg.add +# CHECK-NOT: linalg.generic +apply_filter(NAMED, name="NAMED") + + +# A linalg.generic with elementwise (extf) ops fused into the matmul body, +# preceded by a linalg.fill that must not be matched. +FUSED = """ +#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d4)> +#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4, d3)> +#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)> +module { + func.func @main( + %a: tensor<1x1x128x1024xbf16>, + %b: tensor<1x1x1024x512xbf16>) -> tensor<1x1x128x512xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<1x1x128x512xf32> + %fill = linalg.fill ins(%cst : f32) outs(%e : tensor<1x1x128x512xf32>) -> tensor<1x1x128x512xf32> + %mm = linalg.generic {indexing_maps = [#map, #map1, #map2], + iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"]} + ins(%a, %b : tensor<1x1x128x1024xbf16>, tensor<1x1x1024x512xbf16>) + outs(%fill : tensor<1x1x128x512xf32>) { + ^bb0(%in: bf16, %in_18: bf16, %out: f32): + %ea = arith.extf %in : bf16 to f32 + %eb = arith.extf %in_18 : bf16 to f32 + %mul = arith.mulf %ea, %eb : f32 + %acc = arith.addf %out, %mul : f32 + linalg.yield %acc : f32 + } -> tensor<1x1x128x512xf32> + return %mm : tensor<1x1x128x512xf32> + } +} +""" + +# CHECK-LABEL: IR printer: FUSED +# CHECK: linalg.generic +# CHECK-NOT: linalg.fill +apply_filter(FUSED, name="FUSED")