[Mlir-commits] [mlir] [MLIR][XeGPU] Legalize elementwise vector width before the XeVM conversions (PR #217131)
Sang Ik Lee
llvmlistbot at llvm.org
Tue Aug 18 14:48:16 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/217131
>From 7f8777d605b8a7e77a17af1f6dd85293332d7d42 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 12 Aug 2026 22:18:05 +0000
Subject: [PATCH] [MLIR][XeGPU] Legalize elementwise vector width before the
XeVM conversions
`arith-expand-ops` builds its expansions at whatever width the linearized per-lane
tiles happen to have. For micro-scaling kernels that is 32 components, and nothing
in the pipeline caps vector width, so the backend sees ops it cannot legalize:
%0 = arith.divf %a, %b : vector<32xbf16>
// LLVM ERROR: unable to legalize instruction:
// %:vfid(<32 x s32>) = G_FPEXT %:vfid(<32 x s16>)
Add `xegpu-legalize-vector-width`, which unrolls elementwise ops down to the target
width with the existing `vector::populateVectorUnrollPatterns`, scheduled right
after `arith-expand-ops` since the ops do not exist before it.
Only `ElementwiseMappable` single-result ops are unrolled. Data movement is left
alone, because those vectors are packed payloads whose component count is not a
compute width: a `vector<64xf4E2M1FN>` DPAS operand is 256 bits and becomes
`vector<8xi32>` when lowered, so splitting it would be a pessimization. That
distinction falls out of the trait, so no op list is needed. Ops with a sub-byte
element type are also exempt, since gluing their halves back together materializes
sub-byte vector data movement the backend rejects.
The limit is in components rather than bits on purpose: `vector<32xbf16>` is
exactly 512 bits yet illegal as a compute operand, while the 256-bit
`vector<64xf4E2M1FN>` payload must be preserved. It lives in uArchBase.h as
`kDefaultMaxVectorComponents` so it can become a uArch query later, with a
`max-vector-components` option to override it.
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 35 +++++
.../mlir/Dialect/XeGPU/uArch/uArchBase.h | 23 ++++
.../GPU/Pipelines/GPUToXeVMPipeline.cpp | 9 ++
.../Dialect/XeGPU/Transforms/CMakeLists.txt | 1 +
.../Transforms/XeGPULegalizeVectorWidth.cpp | 126 ++++++++++++++++++
.../Dialect/XeGPU/legalize-vector-width.mlir | 124 +++++++++++++++++
6 files changed, 318 insertions(+)
create mode 100644 mlir/lib/Dialect/XeGPU/Transforms/XeGPULegalizeVectorWidth.cpp
create mode 100644 mlir/test/Dialect/XeGPU/legalize-vector-width.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 36f0a131b345b..bd0668fdc42a1 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -87,6 +87,41 @@ def XeGPUVectorLinearize : Pass<"xegpu-vector-linearize"> {
"scf::SCFDialect", "ub::UBDialect", "vector::VectorDialect"];
}
+def XeGPULegalizeVectorWidth : Pass<"xegpu-legalize-vector-width"> {
+ let summary = "Split elementwise vector ops that exceed the target vector width";
+ let description = [{
+ Unrolls elementwise operations whose vectors have more components than the
+ target can handle into a sequence of narrower operations, glued together
+ with `vector.extract_strided_slice` / `vector.insert_strided_slice`.
+
+ Only operations with `ElementwiseMappable` traits and a single result are
+ considered. Data movement operations (`vector.shuffle`, length-changing
+ `vector.bitcast`, `xegpu.load_nd`, `xegpu.store_nd`, `xegpu.dpas_mx`, ...)
+ are deliberately left alone: their vectors are *packed payloads* whose
+ component count is not a compute width. For example a
+ `vector<64xf4E2M1FN>` DPAS operand is only 256 bits and becomes a
+ `vector<8xi32>` once lowered, so splitting it would be a pessimization.
+ This distinction is exactly what the `ElementwiseMappable` traits express,
+ so no explicit op list is needed.
+
+ Note that a bit-width based threshold would not work here: a
+ `vector<32xbf16>` is exactly 512 bits yet is illegal as a compute operand,
+ while the 256-bit `vector<64xf4E2M1FN>` payload above must be preserved.
+ The limit is therefore expressed in vector components.
+
+ This runs after `arith-expand-ops`, which is what materializes the wide
+ arithmetic (for example the `f8E8M0FNU` scale expansion), and before the
+ XeVM conversions.
+ }];
+ let options = [
+ Option<"maxVectorComponents", "max-vector-components", "unsigned",
+ /*default=*/"0",
+ "Maximum number of components in a vector operand of an elementwise "
+ "op. A value of 0 selects the target default.">
+ ];
+ let dependentDialects = ["arith::ArithDialect", "vector::VectorDialect"];
+}
+
def XeGPUPeepHoleOptimizer : Pass<"xegpu-optimize-peephole"> {
let summary = "Optimize XeGPU block load operations";
let description = [{
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index e1879bb3ffe1c..d791a63b56fd1 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -34,6 +34,29 @@ namespace mlir {
namespace xegpu {
namespace uArch {
+/// Maximum number of components allowed in a vector operand of an elementwise
+/// (compute) operation.
+///
+/// The XeVM targets supported today are backed by SPIR-V, whose `OpTypeVector`
+/// is limited to 16 components, and the backend has no legalization rules for
+/// wider compute vectors. This is a limitation of the currently supported
+/// chips rather than of the XeVM dialect itself, which is why it is declared
+/// here alongside the other microarchitectural queries instead of being
+/// hardcoded in a pass.
+///
+/// Note this is a *component count*, deliberately not a bit width: a
+/// `vector<32xbf16>` is exactly 512 bits yet is illegal as a compute operand,
+/// while a 256-bit `vector<64xf4E2M1FN>` DPAS payload is fine because it is
+/// packed into a `vector<8xi32>` when lowered.
+///
+/// This is also unrelated to `getSubgroupSize()`, which happens to be 16 as
+/// well, and to `getMaxLaneAccessSizeBytes()`, which is 16 *bytes* of block IO.
+///
+/// TODO: Promote this to a virtual `uArch` query once the interface grows a
+/// vector-width entry, and have callers resolve it from the target attribute
+/// rather than using this default.
+constexpr unsigned kDefaultMaxVectorComponents = 16;
+
// An enum class to represent the scope of an instruction
enum class InstructionScope { Lane, Subgroup, Workgroup, Cluster };
enum class InstructionKind {
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 4dc9e2acfe235..8e9b35ee08788 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -114,6 +114,15 @@ void buildGPUPassPipeline(OpPassManager &pm,
pm.addNestedPass<gpu::GPUModuleOp>(
arith::createArithExpandOpsPass(arithExpandOptions));
}
+ // The expansions above build their arithmetic at whatever width the
+ // linearized tiles happen to have, which for micro-scaling kernels is wider
+ // than the target supports (e.g. `arith.divf` on `vector<32xbf16>`, or the
+ // `f8E8M0FNU` expansion working on `vector<32xi32>`). Split those back down
+ // to the target vector width. This has to run after `arith-expand-ops`,
+ // since the ops in question do not exist before it.
+ pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPULegalizeVectorWidth());
+ pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
+ pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToXeVM());
ConvertXeGPUToXeVMPassOptions xegpuToXeVMOptions;
xegpuToXeVMOptions.use64bitIndex = options.use64bitIndex;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 2ed81ae05ab34..68c547e7aef49 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRXeGPUTransforms
XeGPUArrayLengthOptimization.cpp
XeGPUBlocking.cpp
XeGPUContiguityAnalysis.cpp
+ XeGPULegalizeVectorWidth.cpp
XeGPUSgToLaneDistribute.cpp
XeGPUUnroll.cpp
XeGPUWgToSgDistribute.cpp
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULegalizeVectorWidth.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULegalizeVectorWidth.cpp
new file mode 100644
index 0000000000000..237ec4cd7c1c3
--- /dev/null
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULegalizeVectorWidth.cpp
@@ -0,0 +1,126 @@
+//===- XeGPULegalizeVectorWidth.cpp - Split wide elementwise vector ops ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
+#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace xegpu {
+#define GEN_PASS_DEF_XEGPULEGALIZEVECTORWIDTH
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
+} // namespace xegpu
+} // namespace mlir
+
+#define DEBUG_TYPE "xegpu-legalize-vector-width"
+
+using namespace mlir;
+
+namespace {
+
+/// Returns the largest divisor of `numComponents` that does not exceed
+/// `maxComponents`, so that unrolling produces only whole tiles and never a
+/// ragged tail.
+static int64_t getUnrollFactor(int64_t numComponents, int64_t maxComponents) {
+ for (int64_t candidate = maxComponents; candidate > 1; --candidate)
+ if (numComponents % candidate == 0)
+ return candidate;
+ return 1;
+}
+
+/// Returns true for vectors of sub-byte, non-boolean elements (`f4E2M1FN`,
+/// `i4`, ...).
+///
+/// These are packed payload types: they are bit-packed into byte or word
+/// vectors when lowered, and SPIR-V has no corresponding scalar type at all.
+/// An elementwise op producing or consuming one is a quantization boundary
+/// whose own lowering already handles width (see `TruncfToXeVMPattern`, which
+/// splits into `xevm.truncf` instruction groups and concatenates the results as
+/// `i8` vectors). Unrolling such an op here would instead glue the pieces back
+/// together with `insert_strided_slice` on the sub-byte type, materializing
+/// sub-byte vector data movement that cannot be translated.
+///
+/// `i1` is excluded because vector masks are ordinary compute values.
+static bool isSubBytePayload(Type type) {
+ auto vecType = dyn_cast<VectorType>(type);
+ if (!vecType)
+ return false;
+ Type elemType = vecType.getElementType();
+ if (!elemType.isIntOrFloat())
+ return false;
+ unsigned width = elemType.getIntOrFloatBitWidth();
+ return width < 8 && width != 1;
+}
+
+/// Native shape function driving `vector::populateVectorUnrollPatterns`.
+///
+/// Only elementwise, single-result operations are legalized. Everything else
+/// -- `vector.shuffle`, length-changing `vector.bitcast`, `xegpu.load_nd`,
+/// `xegpu.dpas_mx`, ... -- carries packed payloads whose component count is not
+/// a compute width and must be preserved.
+static std::optional<SmallVector<int64_t>>
+getNativeVectorShape(Operation *op, int64_t maxComponents) {
+ if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+ return std::nullopt;
+
+ if (llvm::any_of(op->getOperandTypes(), isSubBytePayload) ||
+ llvm::any_of(op->getResultTypes(), isSubBytePayload))
+ return std::nullopt;
+
+ auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0]);
+ if (!vecType || vecType.getRank() == 0 || vecType.isScalable())
+ return std::nullopt;
+
+ // Elementwise ops have matching shapes across operands and result, so the
+ // result alone determines legality.
+ int64_t trailing = vecType.getShape().back();
+ if (trailing <= maxComponents)
+ return std::nullopt;
+
+ SmallVector<int64_t> nativeShape(vecType.getRank(), 1);
+ nativeShape.back() = getUnrollFactor(trailing, maxComponents);
+ if (nativeShape.back() == trailing)
+ return std::nullopt;
+
+ return nativeShape;
+}
+
+struct XeGPULegalizeVectorWidthPass final
+ : public xegpu::impl::XeGPULegalizeVectorWidthBase<
+ XeGPULegalizeVectorWidthPass> {
+ using XeGPULegalizeVectorWidthBase::XeGPULegalizeVectorWidthBase;
+
+ void runOnOperation() override {
+ // A pass option of 0 means "use the target default". Once `uArch` grows a
+ // vector-width query this is the single place that needs to consult it.
+ int64_t maxComponents = maxVectorComponents
+ ? static_cast<int64_t>(maxVectorComponents)
+ : xegpu::uArch::kDefaultMaxVectorComponents;
+ if (maxComponents < 1) {
+ getOperation()->emitError(
+ "max-vector-components must be greater than zero");
+ return signalPassFailure();
+ }
+
+ RewritePatternSet patterns(&getContext());
+ vector::UnrollVectorOptions options;
+ options.setNativeShapeFn(
+ [maxComponents](Operation *op) -> std::optional<SmallVector<int64_t>> {
+ return getNativeVectorShape(op, maxComponents);
+ });
+ vector::populateVectorUnrollPatterns(patterns, options);
+
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ return signalPassFailure();
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/XeGPU/legalize-vector-width.mlir b/mlir/test/Dialect/XeGPU/legalize-vector-width.mlir
new file mode 100644
index 0000000000000..307350fc1f365
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/legalize-vector-width.mlir
@@ -0,0 +1,124 @@
+// RUN: mlir-opt %s --xegpu-legalize-vector-width --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --xegpu-legalize-vector-width=max-vector-components=8 --split-input-file | FileCheck %s --check-prefix=CHECK8
+
+// Elementwise ops wider than the target width are split; the wide value only
+// survives at the non-elementwise boundary.
+
+// CHECK-LABEL: func.func @split_elementwise
+// CHECK-COUNT-2: arith.divf {{.*}} : vector<16xbf16>
+// CHECK-NOT: arith.divf {{.*}} : vector<32xbf16>
+
+// CHECK8-LABEL: func.func @split_elementwise
+// CHECK8-COUNT-4: arith.divf {{.*}} : vector<8xbf16>
+// CHECK8-NOT: arith.divf {{.*}} : vector<32xbf16>
+func.func @split_elementwise(%a: vector<32xbf16>, %b: vector<32xbf16>) -> vector<32xbf16> {
+ %0 = arith.divf %a, %b : vector<32xbf16>
+ return %0 : vector<32xbf16>
+}
+
+// -----
+
+// Already-narrow ops are left untouched.
+
+// CHECK-LABEL: func.func @already_legal
+// CHECK-NEXT: arith.divf {{.*}} : vector<16xbf16>
+// CHECK-NEXT: return
+func.func @already_legal(%a: vector<16xbf16>, %b: vector<16xbf16>) -> vector<16xbf16> {
+ %0 = arith.divf %a, %b : vector<16xbf16>
+ return %0 : vector<16xbf16>
+}
+
+// -----
+
+// Ops changing element width are still elementwise and must be split, tracking
+// the result shape.
+
+// CHECK-LABEL: func.func @split_cast
+// CHECK-COUNT-2: arith.extui {{.*}} : vector<16xi8> to vector<16xi32>
+// CHECK-NOT: arith.extui {{.*}} : vector<32xi8> to vector<32xi32>
+func.func @split_cast(%a: vector<32xi8>) -> vector<32xi32> {
+ %0 = arith.extui %a : vector<32xi8> to vector<32xi32>
+ return %0 : vector<32xi32>
+}
+
+// -----
+
+// `vector.shuffle` is data movement, not elementwise: a wide shuffle is a
+// packed payload and must be preserved.
+
+// CHECK-LABEL: func.func @shuffle_untouched
+// CHECK-NEXT: vector.shuffle
+// CHECK-NEXT: return
+func.func @shuffle_untouched(%a: vector<32xf4E2M1FN>, %b: vector<32xf4E2M1FN>) -> vector<64xf4E2M1FN> {
+ %0 = vector.shuffle %a, %b [0, 32, 1, 33, 2, 34, 3, 35, 4, 36, 5, 37, 6, 38, 7, 39,
+ 8, 40, 9, 41, 10, 42, 11, 43, 12, 44, 13, 45, 14, 46, 15, 47,
+ 16, 48, 17, 49, 18, 50, 19, 51, 20, 52, 21, 53, 22, 54, 23, 55,
+ 24, 56, 25, 57, 26, 58, 27, 59, 28, 60, 29, 61, 30, 62, 31, 63]
+ : vector<32xf4E2M1FN>, vector<32xf4E2M1FN>
+ return %0 : vector<64xf4E2M1FN>
+}
+
+// -----
+
+// A length-changing `vector.bitcast` is not elementwise either.
+
+// CHECK-LABEL: func.func @bitcast_untouched
+// CHECK-NEXT: vector.bitcast {{.*}} : vector<32xi8> to vector<64xf4E2M1FN>
+// CHECK-NEXT: return
+func.func @bitcast_untouched(%a: vector<32xi8>) -> vector<64xf4E2M1FN> {
+ %0 = vector.bitcast %a : vector<32xi8> to vector<64xf4E2M1FN>
+ return %0 : vector<64xf4E2M1FN>
+}
+
+// -----
+
+// Sub-byte results are packed payloads, not compute values. Splitting the
+// quantizing `arith.truncf` would glue the halves back together with
+// `insert_strided_slice` on `f4E2M1FN`, which lowers to untranslatable
+// sub-byte vector movement. Its own XeVM lowering handles width instead.
+
+// CHECK-LABEL: func.func @subbyte_result_untouched
+// CHECK-NEXT: arith.truncf {{.*}} : vector<32xbf16> to vector<32xf4E2M1FN>
+// CHECK-NEXT: return
+func.func @subbyte_result_untouched(%a: vector<32xbf16>) -> vector<32xf4E2M1FN> {
+ %0 = arith.truncf %a : vector<32xbf16> to vector<32xf4E2M1FN>
+ return %0 : vector<32xf4E2M1FN>
+}
+
+// -----
+
+// Same in the dequantizing direction, where the sub-byte type is the operand.
+
+// CHECK-LABEL: func.func @subbyte_operand_untouched
+// CHECK-NEXT: arith.extf {{.*}} : vector<32xf4E2M1FN> to vector<32xbf16>
+// CHECK-NEXT: return
+func.func @subbyte_operand_untouched(%a: vector<32xf4E2M1FN>) -> vector<32xbf16> {
+ %0 = arith.extf %a : vector<32xf4E2M1FN> to vector<32xbf16>
+ return %0 : vector<32xbf16>
+}
+
+// -----
+
+// `i1` masks are ordinary compute values and must still be split.
+
+// CHECK-LABEL: func.func @mask_is_not_subbyte
+// CHECK-COUNT-2: arith.cmpi eq, {{.*}} : vector<16xi8>
+// CHECK-NOT: arith.cmpi eq, {{.*}} : vector<32xi8>
+func.func @mask_is_not_subbyte(%a: vector<32xi8>, %b: vector<32xi8>) -> vector<32xi1> {
+ %0 = arith.cmpi eq, %a, %b : vector<32xi8>
+ return %0 : vector<32xi1>
+}
+
+// -----
+
+// Widths that are not a multiple of the limit fall back to the largest divisor
+// so that no ragged tail is produced. For 24 components with a limit of 16 the
+// largest usable divisor is 12.
+
+// CHECK-LABEL: func.func @non_multiple_width
+// CHECK-COUNT-2: arith.negf {{.*}} : vector<12xf32>
+// CHECK-NOT: arith.negf {{.*}} : vector<24xf32>
+func.func @non_multiple_width(%a: vector<24xf32>) -> vector<24xf32> {
+ %0 = arith.negf %a : vector<24xf32>
+ return %0 : vector<24xf32>
+}
More information about the Mlir-commits
mailing list