[xegpu] Add modular, parameter-free pipeline by tkarna · Pull Request #293 · llvm/lighthouse

tkarna · 2026-10-01T21:42:51Z

Introduces a new XeGPU pipeline that consists of parameter-free sub-schedules.

The full pipeline is defined as

schedules = [
    xegpu.cleanup_schedule(),
    xegpu.wg_tiling_schedule(),
    xegpu.vectorize_schedule(),
    xegpu.bufferize_schedule(),
    xegpu.outline_gpu_func_schedule(),
    xegpu.vector_to_xegpu_schedule(),
    xegpu.annotate_layouts_schedule(),
    xegpu.xegpu_to_binary_schedule(),
]

Each lowering stage has been defined as a standalone sub-schedule.
The schedules do not take any parameters by default. Instead, the schedules analyze the IR at lowering time and infer required tile sizes etc.
Consequently, the tile selection logic is now defined using the transform dialect; In the schedule, parameters, like tile sizes, are no longer concrete Python int values, but transform dialect parameters. They do not get concrete values until the transform schedule is invoked, i.e. at lowering time.

The long-term goal is that we could replace all the existing xegpu schedules, e.g., mlp_schedule, reduction_schedule, fused_attetnion_schedule, with a single pipeline that can consume any kind of input IR.

This PR is just a first step toward that direction, and still in progress. Currently the pipeline supports basic matmul, reduction and attention payloads (from the KernelBench suite). The matmul support is most mature - it implements the existing GEMM tile size selection mechanism completely. At the moment, the new pipeline is only used in kernel_bench.py script.

Some schedules, like xepgu.wg_tiling(), implement different lowering implementations for gemm, reduction, and attention -like payloads. These implementations are handled by the transform.alternatives op that tries to apply each defined "alternative" block until one succeeds.

As the parameters are not concrete Python values, all calculations that involve parameters must be done inside a dedicated transform op (see e.g., compute_sg_layout.py), or via the transform dialect SMT extension (see annotate_layouts schedule for example).

tkarna · 2026-10-02T12:47:10Z

Currently supported KernelBench kernels. The new pipeline produces identical xegpu-wg level IR as the existing mlp/reduction/fused_attention schedules.

Type	Level	Benchmark	Status	Parameter selection
GEMM	1	1_Square_matrix_multiplication	OK	Automatic
GEMM	1	2_Standard_matrix_multiplication	OK	Automatic
GEMM	1	6_Matmul_with_large_K_dimension	OK	Automatic
GEMM	1	7_Matmul_with_small_K_dimension	OK	Automatic
GEMM	1	9_Tall_skinny_matrix_multiplication	OK	Automatic
GEMM	1	13_Matmul_for_symmetric_matrices	OK	Automatic
GEMM	1	16_Matmul_with_transposed_A	OK	Automatic
GEMM	1	17_Matmul_with_transposed_B	OK	Automatic
GEMM	1	18_Matmul_with_transposed_both	OK	Automatic
GEMM	2	9_Matmul_Subtract_Multiply_ReLU	OK	Automatic
GEMM	2	12_Gemm_Multiply_LeakyReLU	OK	Automatic
GEMM	2	29_Matmul_Mish_Mish	OK	Automatic
GEMM	2	40_Matmul_Scaling_ResidualAdd	OK	Automatic
GEMM	2	59_Matmul_Swish_Scaling	OK	Automatic
GEMM	2	63_Gemm_ReLU_Divide	OK	Automatic
GEMM	2	70_Gemm_Sigmoid_Scaling_ResidualAdd	OK	Automatic
GEMM	2	81_Gemm_Swish_Divide_Clamp_Tanh_Clamp	OK	Automatic
MLP	3	1_MLP	OK	Automatic
MLP	3	2_ShallowWideMLP	OK	Automatic
MLP	3	3_DeepNarrowMLP	OK	Automatic
Reduction	1	23_Softmax	OK	Hard-coded
Reduction	1	24_LogSoftmax	OK	Hard-coded
Attention	1	97_ScaledDotProductAttention	OK	Hard-coded

Interfaces must be registered for every new context, the python class bool _interfaces_attached persists across contexts.

adam-smnk

Could you also add more granular tests per new transform op?
Probably not all really need it but at least the infer ones.

adam-smnk · 2026-10-06T09:10:40Z

+    op_name = op.operation.name
+    wg_tile = None
+    k_tile = None
+    if op_name == "linalg.matmul":


I assume it only supports plain variant without broadcasts, transposes etc.
I'd be good to at least add assert to document these.

I can add some asserts. But in general I think we should adopt test-driven development philosophy. It's not possible to safe-guard such analysis against all possible input IR variations. The test suite shows which cases are supported and for the rest we assume the method is broken. We'll update the test suite as we improve coverage.

Test coverage is often limited, I wouldn't take them as ground truth.

Also, the question is how quickly you realize that it's broken. Will it blow up instantly, 10 passes later, or give incorrect results?
Specialized transform are fine and bugs are unavoidable but at least sth traceable directly in code will make the made assumptions explicit.

Hmm, actually I do not know what the expected behavior is if the input IR has a transpose in the matmul op's indexing maps, or if an operand is produced by a broadcast op. The correct xegpu metadata depends on what the vector/xegpu dialect IR looks like. I can figure that out by cooking an example, which brings me back to TTD - address it when there's a concrete use case.

added assertion and tests

Copilot

Copilot review overview

🟡 Changes recommended

Unresolved parameter-inference and lowering failures prevent reliable execution of the new pipeline.

Review effort: Balanced
Findings: 3 · 6

Open (9)

GPU pipeline drops the requested payload function name · New
Empty GPU pipeline crashes before returning the module · New
Reduction-only payloads fail epilogue extraction · New
Transpose slice shape recovery returns incorrect dimensions · New
Untiled matmul shape analysis ignores indexing maps · New
Parameter inference rejects nested workgroup and K loops · New
Payload lookup misses operations normalized by cleanup · New
Unsupported multi-reduction-to-contract patterns remain enabled · New
Fast path ignores forced subgroup tile sizes · New

What changed in this PR

Introduces modular XeGPU lowering schedules that select parameters during transformation, moving toward a shared pipeline for matmul, reduction, and attention payloads.

Changes:

Adds standalone lowering stages and fallback alternatives.
Introduces transform-based parameter inference and layout calculations.
Integrates the pipeline into KernelBench and adds matmul-analysis tests.

File	Description
test/transform/test_matmul_analysis.py	Tests shape and tile inference.
lighthouse/utils/mlir.py	Stops producer traversal at matmuls.
lighthouse/transform/alternatives.py	Wraps fallback transform regions.
lighthouse/transform/__init__.py	Exports alternatives wrapper.
lighthouse/schedule/xegpu/xegpu_parameter_selector.py	Supports constrained parameter selection.
lighthouse/schedule/xegpu/wg_tiling_schedule.py	Adds tiling alternatives.
lighthouse/schedule/xegpu/vectorize_schedule.py	Adds standalone vectorization.
lighthouse/schedule/xegpu/vector_to_xegpu_schedule.py	Adds vector-to-XeGPU conversion stage.
lighthouse/schedule/xegpu/outline_gpu_func_schedule.py	Infers launch threads and outlines kernels.
lighthouse/schedule/xegpu/matmul_costmodel.py	Supports fixed tile constraints.
lighthouse/schedule/xegpu/lowering_common.py	Generalizes shared lowering helpers.
lighthouse/schedule/xegpu/cleanup_schedule.py	Normalizes input operations.
lighthouse/schedule/xegpu/bufferize_schedule.py	Adds standalone GPU bufferization.
lighthouse/schedule/xegpu/annotate_layouts_schedule.py	Adds layout and prefetch annotation.
lighthouse/schedule/xegpu/__init__.py	Exports modular schedules.
lighthouse/dialects/transform/transform_ext/utils/xegpu_param_selection.py	Connects analysis to parameter selection.
lighthouse/dialects/transform/transform_ext/utils/matmul_analysis.py	Infers matmul shapes and tiles.
lighthouse/dialects/transform/transform_ext/ops/replace_with_fused_attention.py	Accepts transform-valued tile sizes.
lighthouse/dialects/transform/transform_ext/ops/infer_xegpu_reduction_params.py	Adds placeholder reduction parameters.
lighthouse/dialects/transform/transform_ext/ops/infer_xegpu_gemm_params.py	Infers GEMM parameter dictionaries.
lighthouse/dialects/transform/transform_ext/ops/infer_xegpu_attention_params.py	Adds placeholder attention parameters.
lighthouse/dialects/transform/transform_ext/ops/get_param_dict_entry.py	Extracts dictionary parameters.
lighthouse/dialects/transform/transform_ext/ops/extract_handle.py	Supports recoverable extraction failures.
lighthouse/dialects/transform/transform_ext/ops/emit_definite_failure.py	Reports exhausted alternatives.
lighthouse/dialects/transform/transform_ext/ops/compute_sg_layout.py	Computes subgroup layouts.
lighthouse/dialects/transform/transform_ext/ops/compute_num_threads.py	Computes launch thread counts.
lighthouse/dialects/transform/transform_ext/__init__.py	Exports new transform operations.
lighthouse/dialects/transform/smt_ext/ops/constrain_params.py	Attaches interfaces per context.
examples/xegpu/llama3_schedule.py	Updates helper call to keyword syntax.
examples/xegpu/kernel_bench.py	Adds optional modular pipeline execution.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

-            stop_at_stage=stop_at_stage,
-        )
+    if use_gpu_pipeline:
+        pipeline = get_xegpu_pipeline(stop_at_stage)


-            stop_at_stage=stop_at_stage,
-        )
+    if use_gpu_pipeline:
+        pipeline = get_xegpu_pipeline(stop_at_stage)


+        wg_loop = match_and_split(func, ops={"scf.forall"}, nhandles=1)[0]
+        generic_ops = match(wg_loop, ops={"linalg.generic"})
+        elemwise_ops = transform_ext.filter_elementwise(generic_ops)
+        leaf_elemwise = transform_ext.extract_handle(elemwise_ops, -1, silenceable=True)


+    slice_op = _first_producer_named(value, "tensor.extract_slice")
+    source = slice_op.operands[0] if slice_op is not None else value
+    return list(ir.ShapedType(source.type).shape)


+    # TODO use op name as xegpu dialect python bindings are missing
+    op_name = op.operation.name
+    if op_name == "linalg.matmul":
+        return _linalg_matmul_shape_and_transpose(op)


+        if op.parent.name != "scf.forall":
+            # target is not within a scf.forall loop, so it's not workgroup tiled
+            return None, None


+    """Normalize singleton dimensions and fuse elementwise ops in the payload."""
+
+    with schedule_boilerplate() as (schedule, named_seq):
+        op_names = ["linalg.generic", "linalg.matmul"]


+            op_name="builtin.module",
+            deduplicate=True,
+        )
+        lowering_common.vectorize(payload_mod, payload_func=func)


+        if (fixed_wg_tile is not None and wg_tile != fixed_wg_tile) or (
+            fixed_k_tile is not None and params["k_tile"] != fixed_k_tile
+        ):


charithaintc

thanks!
I have some clarification comments.

nit: It would be nice if the new transfrom ext ops show a simple usage of the op in docstring.

charithaintc · 2026-10-06T20:35:14Z

+    """
+    Compute the workgroup thread count from wg and sg tile params.
+
+    Returns `base * prod_i(wg_i // sg_i)` over the tile dimensions, treating an


good to define what base, wg_i and sg_i specifically refer to here. I believe wg_i and sg_i refer to tile sizes?

charithaintc · 2026-10-06T20:43:40Z

+    TransformExtensionDialect.Operation, name="infer_xegpu_attention_params"
+):
+    """
+    Infer XeGPU attention tiling parameters for an attention anchor op.


what is an attention anchor op? Do you mean to say that we match the attention pattern and then use this op to derive tile shapes?

charithaintc · 2026-10-06T20:50:42Z

            n_ctx = ir.RankedTensorType(k.type).shape[-2]
-            tile_size = ir.IntegerAttr(op.tile_size).value
+            tile_size_params = state.get_params(op.tile_size)
+            if len(tile_size_params) != 1 or not isinstance(


how would this look like when we transition to dependent reduction fusion approach? I believe you will still pass the derived attention tiling params to the sequence of transform ops?

charithaintc · 2026-10-06T21:00:53Z

+        lh_transform.cleanup(func)
+
+        # Fuse elementwise ops, also removes unused linalg op results (if any).
+        func = apply_registered_pass(func, "linalg-fuse-elementwise-ops")


I wonder if we can avoid this pass and use your upstreamed change for "unused linalg op results (if any)"

tkarna added 2 commits October 1, 2026 23:14

lowering_common: add gpu-outlining pipeline stage + updates

6531c3b

transform: add transform.alternatives op wrapper

0e1b97f

tkarna changed the title ~~[xegpu] Add pipeline~~ [xegpu] Add modular, parameter-free pipeline Oct 1, 2026

tkarna force-pushed the xegpu-modular-pipeline branch from 8c28598 to 9e6a3dc Compare October 2, 2026 12:46

tkarna force-pushed the xegpu-modular-pipeline branch from 9e6a3dc to 2431d89 Compare October 5, 2026 16:47

tkarna marked this pull request as ready for review October 5, 2026 16:51

tkarna requested a review from adam-smnk October 5, 2026 16:51

tkarna added 22 commits October 5, 2026 23:12

smt_ext: fix constrain_params interface registration

2c4906e

Interfaces must be registered for every new context, the python class bool _interfaces_attached persists across contexts.

transform_ext: add matmul_analysis and xegpu_param_selector helpers

6e365ff

transform_ext: add infer_xegpu_gemm_params transform op

03ae8df

transform_ext: add get_param_dict_entry transform op

79c781a

transform_ext: extract_handle has a silenceable bool arg

b88dc2f

transform_ext: add emit_definite_failure op

512a519

transform_ext: add compute_num_threads op

f1c0601

transform_ext: add compute_sg_layout op

8714c6d

transform_ext: add infer_xegpu_reduction_params

3ea2b12

matmul analysis: add method to infer wg and k tile sizes from IR

e09c9a5

matmul_analysis: add tests

57dcacc

XeGPUParameterSelector: support fixed wg, sg and k tile sizes

9f1e7a4

infer_xegpu_gemm_params: add force wg/sg/k tile arguments

6461582

replace_with_fused_attention: tile_size can be a transform param

0913e2a

lowering_common: vectorize: hoist after cleanup

5d55b73

transform_ext: add infer_xegpu_attention_params op

2b53fda

xegpu: add outline_gpu_func schedule

42b3f59

xegpu: add vectorize schedule

0f4ec1a

xegpu: add bufferize schedule

8e2a325

xegpu: add vector_to_xegpu schedule

5a64e87

xegpu: add annotate_layouts schedule

60c48cf

xegpu: add cleanup schedule

aa04f81

tkarna added 4 commits October 5, 2026 23:12

xegpu: add wg_tiling schedule

37cf38a

fix llama3 schedule

eaa40f3

inspect_payload: stop tracing producers at matmul op

5628489

kernel_bench: use new pipeline via use-gpu-pipeline CLI arg

f634cea

tkarna force-pushed the xegpu-modular-pipeline branch from 2431d89 to f634cea Compare October 5, 2026 20:13

adam-smnk reviewed Oct 6, 2026

View reviewed changes

tkarna added 2 commits October 6, 2026 20:22

address comments

2c64eb0

add _schedule suffix to all schedules

c0327ac

tkarna force-pushed the xegpu-modular-pipeline branch from 812c343 to c0327ac Compare October 6, 2026 17:35

matmul_analysis: add guardrails and tests

0caceac

tkarna requested a balanced review from Copilot October 6, 2026 18:06

Copilot started reviewing on behalf of tkarna October 6, 2026 18:07 View session

Copilot AI reviewed Oct 6, 2026

View reviewed changes

charithaintc reviewed Oct 6, 2026

View reviewed changes

Conversation

tkarna commented Oct 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tkarna commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

adam-smnk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

charithaintc left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tkarna commented Oct 1, 2026 •

edited

Loading

tkarna commented Oct 2, 2026 •

edited

Loading

charithaintc left a comment •

edited

Loading