[xegpu] Generalize fused attention schedule for KernelBench

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",

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

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)

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,
     )

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")