[xegpu] Generalize fused attention schedule for KernelBench

[tkarna]

• Generalizes op matching in fused_attention_schedule so that it also applies to KernelBench level1 97 ScaledAttention kernel.
• Fixes duplicated handle issue in trace_producers trasform ext op.
• Adds a new filter_contraction_ops transform ext op to match batch_matmul with optionally fused elementwise ops.

[tkarna]

Ping @charithaintc

[Copilot]

Pull request overview

This PR updates the xegpu fused-attention transform schedule and transform_ext utilities to more flexibly match attention patterns (including KernelBench ScaledAttention), while fixing producer-tracing handle duplication and adding a new contraction-op filtering transform.

Changes:

• Generalize fused_attention_schedule op matching by converting linalg.mul/linalg.batch_matmul to linalg.generic and using a new contraction filter to identify matmul-like ops.
• Fix trace_producers de-duplication logic so the producer chain doesn’t return repeated handles.
• Add filter_contraction_ops transform_ext op to identify contraction ops (named or generic) via indexing-map structure, and export it via transform_ext/__init__.py.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
lighthouse/schedule/xegpu/fused_attention_schedule.py	Generalizes matching and rewires Q/K/V/scale handle discovery using contraction filtering + producer tracing.
lighthouse/dialects/transform/transform_ext/ops/trace_producers.py	Adjusts visited tracking to prevent duplicated producer handles.
lighthouse/dialects/transform/transform_ext/ops/filter_contraction_ops.py	Introduces new transform op to filter contraction (matmul-like) ops, including structural detection for linalg.generic.
lighthouse/dialects/transform/transform_ext/init.py	Exports the new filter_contraction_ops helper from transform_ext.

---

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