[Mlir-commits] [mlir] [mlir][xegpu] Add xegpu-canonicalize pass to un-flatten gather/scatter (PR #218292)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Aug 23 14:41:54 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
XeGPU layouts (sg_layout, inst_data, lane_layout, ...) are expressed in terms of the N-D shape of the accessed data. Frontends, however, often emit vector.gather / vector.scatter with 1-D operands obtained by shape_cast-ing N-D indices and masks. That forces layout propagation to reason backwards through the surrounding shape_casts, which is either impossible or yields layouts that cannot be distributed. A flattened 128x64 gather asserts in xegpu::LayoutAttr::expandDim, reached via inferShapeCastSourceLayout during xegpu-propagate-layout{layout-kind=subgroup}.
This PR adds an xegpu-canonicalize pass to hold XeGPU-specific canonicalizations that are not profitable in general, and so are applied explicitly rather than from the dialect canonicalization patterns. Its only pattern so far restores the N-D form of a flattened gather/scatter.
---
Full diff: https://github.com/llvm/llvm-project/pull/218292.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td (+34)
- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h (+2)
- (modified) mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp (+1)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt (+1)
- (added) mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp (+197)
- (added) mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir (+101)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index 36f0a131b345b..acc3a2fd7bf85 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -11,6 +11,40 @@
include "mlir/Pass/PassBase.td"
+def XeGPUCanonicalize : Pass<"xegpu-canonicalize"> {
+ let summary = "XeGPU specific canonicalization of vector/xegpu code";
+ let description = [{
+ Canonicalizations that bring the IR into the form the rest of the XeGPU
+ lowering expects. They are not profitable in general, so they are applied
+ explicitly here instead of from the dialect canonicalization patterns.
+
+ Currently the pass restores the N-D form of flattened `vector.gather` /
+ `vector.scatter`. XeGPU layouts (`sg_layout`, `inst_data`, `lane_layout`,
+ ...) are expressed in terms of the N-D shape of the data, so a flattened
+ gather forces layout propagation to reason through the surrounding
+ `vector.shape_cast` ops, which is either impossible or produces layouts
+ that cannot be distributed.
+
+ ```mlir
+ // Before:
+ %cst = arith.constant dense<0.0> : vector<8192xbf16>
+ %flat_idx = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
+ %flat_mask = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
+ %flat_res = vector.gather %src[%c0] [%flat_idx], %flat_mask, %cst
+ : memref<?xbf16>, vector<8192xindex>, vector<8192xi1>, vector<8192xbf16>
+ into vector<8192xbf16>
+ %res = vector.shape_cast %flat_res : vector<8192xbf16> to vector<128x64xbf16>
+
+ // After:
+ %cst = arith.constant dense<0.0> : vector<128x64xbf16>
+ %res = vector.gather %src[%c0] [%idx], %mask, %cst
+ : memref<?xbf16>, vector<128x64xindex>, vector<128x64xi1>,
+ vector<128x64xbf16> into vector<128x64xbf16>
+ ```
+ }];
+ let dependentDialects = ["arith::ArithDialect", "vector::VectorDialect"];
+}
+
def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
let summary = "Propagate and assign XeGPU layout information";
let description = [{
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index 388bd6145df21..a32c8befdd334 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -59,6 +59,8 @@ struct UnrollOptions {
}
};
+/// Appends the XeGPU specific canonicalization patterns into `patterns`.
+void populateXeGPUCanonicalizePatterns(RewritePatternSet &patterns);
/// Appends patterns for optimizing block load operations into `patterns`.
void populateXeGPUPeepHoleOptimizerPatterns(RewritePatternSet &patterns);
/// Appends patterns for array length optimization into `patterns`.
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 4dc9e2acfe235..3705ad75707d2 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -66,6 +66,7 @@ void buildGPUPassPipeline(OpPassManager &pm,
laneLayoutOptions.indexBitWidth = options.use64bitIndex ? 64 : 32;
laneLayoutOptions.layoutKind = "lane";
pm.addNestedPass<ModuleOp>(createCSEPass());
+ pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUCanonicalize());
if (options.enableVectorToXeGPU)
pm.addNestedPass<gpu::GPUModuleOp>(createConvertVectorToXeGPU());
if (options.xegpuOpLevel == "workgroup") {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
index 2ed81ae05ab34..8b53c7699436f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIRXeGPUTransforms
XeGPUArrayLengthOptimization.cpp
XeGPUBlocking.cpp
+ XeGPUCanonicalize.cpp
XeGPUContiguityAnalysis.cpp
XeGPUSgToLaneDistribute.cpp
XeGPUUnroll.cpp
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
new file mode 100644
index 0000000000000..eb062a3a4f7aa
--- /dev/null
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
@@ -0,0 +1,197 @@
+//===- XeGPUCanonicalize.cpp - XeGPU specific canonicalization --*- C++ -*-===//
+//
+// 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/XeGPU/Transforms/Passes.h"
+#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+namespace xegpu {
+#define GEN_PASS_DEF_XEGPUCANONICALIZE
+#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
+} // namespace xegpu
+} // namespace mlir
+
+#define DEBUG_TYPE "xegpu-canonicalize"
+
+using namespace mlir;
+
+namespace {
+
+/// Returns the `vector.shape_cast` that flattened a value of type `ndType`
+/// into `flat`, if that is how `flat` was produced.
+static vector::ShapeCastOp getFlattenCast(Value flat, VectorType ndType) {
+ auto shapeCast = flat.getDefiningOp<vector::ShapeCastOp>();
+ if (shapeCast && shapeCast.getSourceVectorType() == ndType)
+ return shapeCast;
+ return nullptr;
+}
+
+static DenseElementsAttr getDenseConstant(Value flat) {
+ DenseElementsAttr elements;
+ if (matchPattern(flat, m_Constant(&elements)))
+ return elements;
+ return nullptr;
+}
+
+static vector::BroadcastOp getSplatBroadcast(Value flat) {
+ auto broadcast = flat.getDefiningOp<vector::BroadcastOp>();
+ if (broadcast && !isa<VectorType>(broadcast.getSourceType()))
+ return broadcast;
+ return nullptr;
+}
+
+/// Returns true if `unflatten` can reshape `flat` to `ndType`.
+static bool canUnflatten(Value flat, VectorType ndType) {
+ return getFlattenCast(flat, ndType) || getDenseConstant(flat) ||
+ getSplatBroadcast(flat);
+}
+
+/// Reshape `flat` to `ndType` without introducing a 1-D to N-D
+/// `vector.shape_cast`. Only handles the forms a frontend actually emits when
+/// flattening: the flattening cast itself, a constant, and a splat.
+/// `canUnflatten` must hold.
+static Value unflatten(PatternRewriter &rewriter, Value flat,
+ VectorType ndType) {
+ assert(cast<VectorType>(flat.getType()).getRank() == 1 &&
+ "expected a 1-D vector");
+
+ if (auto shapeCast = getFlattenCast(flat, ndType))
+ return shapeCast.getSource();
+
+ if (DenseElementsAttr elements = getDenseConstant(flat))
+ return arith::ConstantOp::create(rewriter, flat.getLoc(), ndType,
+ elements.reshape(ndType));
+
+ auto broadcast = getSplatBroadcast(flat);
+ assert(broadcast && "expected canUnflatten to hold");
+ return vector::BroadcastOp::create(rewriter, flat.getLoc(), ndType,
+ broadcast.getSource());
+}
+
+/// Restore the N-D form of a flattened `vector.gather` / `vector.scatter`.
+///
+/// XeGPU layouts are expressed in terms of the N-D shape of the accessed data,
+/// so a flattened gather/scatter forces layout propagation to reason through
+/// the surrounding `vector.shape_cast` ops - which is either impossible or
+/// yields layouts that cannot be distributed.
+///
+/// ```mlir
+/// // Before:
+/// %cst = arith.constant dense<0.0> : vector<8192xbf16>
+/// %flat_idx = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
+/// %flat_mask = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
+/// %flat_res = vector.gather %src[%c0] [%flat_idx], %flat_mask, %cst
+/// : memref<?xbf16>, vector<8192xindex>, vector<8192xi1>, vector<8192xbf16>
+/// into vector<8192xbf16>
+/// %res = vector.shape_cast %flat_res : vector<8192xbf16> to vector<128x64xbf16>
+///
+/// // After:
+/// %cst = arith.constant dense<0.0> : vector<128x64xbf16>
+/// %res = vector.gather %src[%c0] [%idx], %mask, %cst
+/// : memref<?xbf16>, vector<128x64xindex>, vector<128x64xi1>,
+/// vector<128x64xbf16> into vector<128x64xbf16>
+/// ```
+template <typename OpTy>
+struct UnflattenGatherScatter : public OpRewritePattern<OpTy> {
+ using OpRewritePattern<OpTy>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(OpTy op,
+ PatternRewriter &rewriter) const override {
+ constexpr bool isGather = std::is_same_v<OpTy, vector::GatherOp>;
+
+ if (op.getIndexVectorType().getRank() != 1)
+ return rewriter.notifyMatchFailure(op, "index vector is not 1-D");
+
+ // The N-D shape comes from the index operand's producer: this only undoes
+ // a flattening that already happened, it never invents a shape.
+ auto indexCast =
+ op.getIndices().template getDefiningOp<vector::ShapeCastOp>();
+ if (!indexCast || indexCast.getSourceVectorType().getRank() < 2)
+ return rewriter.notifyMatchFailure(
+ op, "index vector is not a shape_cast of an N-D vector");
+ VectorType ndIndexType = indexCast.getSourceVectorType();
+ VectorType ndMaskType =
+ ndIndexType.cloneWith(std::nullopt, rewriter.getI1Type());
+ VectorType ndType = ndIndexType.cloneWith(
+ std::nullopt, op.getVectorType().getElementType());
+
+ // Check everything before creating any IR: a partially applied rewrite
+ // would leave dead ops behind.
+ if (!canUnflatten(op.getMask(), ndMaskType))
+ return rewriter.notifyMatchFailure(op, "cannot un-flatten the mask");
+
+ if constexpr (isGather) {
+ if (!canUnflatten(op.getPassThru(), ndType))
+ return rewriter.notifyMatchFailure(op,
+ "cannot un-flatten the pass-thru");
+
+ // Every use must cast back to N-D, else the rewrite would just move the
+ // shape_casts to the result.
+ SmallVector<vector::ShapeCastOp> resultCasts;
+ for (Operation *user : op->getUsers()) {
+ auto resultCast = dyn_cast<vector::ShapeCastOp>(user);
+ if (!resultCast || resultCast.getResultVectorType() != ndType)
+ return rewriter.notifyMatchFailure(
+ op, "result is not exclusively shape_cast back to N-D");
+ resultCasts.push_back(resultCast);
+ }
+ if (resultCasts.empty())
+ return rewriter.notifyMatchFailure(op, "result is unused");
+
+ Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
+ Value passThru = unflatten(rewriter, op.getPassThru(), ndType);
+ auto ndGather = vector::GatherOp::create(
+ rewriter, op.getLoc(), ndType, op.getBase(), op.getOffsets(),
+ indexCast.getSource(), mask, passThru, op.getAlignmentAttr());
+ for (vector::ShapeCastOp resultCast : resultCasts)
+ rewriter.replaceOp(resultCast, ndGather.getResult());
+ rewriter.eraseOp(op);
+ } else {
+ if (!canUnflatten(op.getValueToStore(), ndType))
+ return rewriter.notifyMatchFailure(
+ op, "cannot un-flatten the stored value");
+
+ Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
+ Value valueToStore = unflatten(rewriter, op.getValueToStore(), ndType);
+ // Only operand types change, and this keeps the optional tensor result
+ // untouched.
+ rewriter.modifyOpInPlace(op, [&] {
+ op.getIndicesMutable().assign(indexCast.getSource());
+ op.getMaskMutable().assign(mask);
+ op.getValueToStoreMutable().assign(valueToStore);
+ });
+ }
+ return success();
+ }
+};
+
+struct XeGPUCanonicalizePass final
+ : public xegpu::impl::XeGPUCanonicalizeBase<XeGPUCanonicalizePass> {
+ void runOnOperation() override {
+ RewritePatternSet patterns(&getContext());
+ xegpu::populateXeGPUCanonicalizePatterns(patterns);
+ if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+ return signalPassFailure();
+ }
+};
+
+} // namespace
+
+void xegpu::populateXeGPUCanonicalizePatterns(RewritePatternSet &patterns) {
+ patterns.add<UnflattenGatherScatter<vector::GatherOp>,
+ UnflattenGatherScatter<vector::ScatterOp>>(
+ patterns.getContext());
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir b/mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir
new file mode 100644
index 0000000000000..650f334b2046b
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir
@@ -0,0 +1,101 @@
+// RUN: mlir-opt %s -split-input-file -xegpu-canonicalize | FileCheck %s
+
+// CHECK-LABEL: @gather_2d_from_flat
+// CHECK-SAME: %[[SRC:.+]]: memref<?xbf16, strided<[1], offset: ?>>,
+// CHECK-SAME: %[[IDX:.+]]: vector<128x64xindex>, %[[MASK:.+]]: vector<128x64xi1>
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[PASS_THRU:.+]] = arith.constant dense<0.000000e+00> : vector<128x64xbf16>
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[RES:.+]] = vector.gather %[[SRC]][%[[C0]]] [%[[IDX]]], %[[MASK]], %[[PASS_THRU]]
+// CHECK-SAME: into vector<128x64xbf16>
+// CHECK-NOT: vector.shape_cast
+// CHECK: return %[[RES]] : vector<128x64xbf16>
+func.func @gather_2d_from_flat(%src: memref<?xbf16, strided<[1], offset: ?>>,
+ %idx: vector<128x64xindex>, %mask: vector<128x64xi1>) -> vector<128x64xbf16> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<8192xbf16>
+ %flat_idx = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
+ %flat_mask = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
+ %flat_res = vector.gather %src[%c0] [%flat_idx], %flat_mask, %cst
+ : memref<?xbf16, strided<[1], offset: ?>>, vector<8192xindex>, vector<8192xi1>,
+ vector<8192xbf16> into vector<8192xbf16>
+ %res = vector.shape_cast %flat_res : vector<8192xbf16> to vector<128x64xbf16>
+ return %res : vector<128x64xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @scatter_2d_from_flat
+// CHECK-SAME: %[[SRC:.+]]: memref<?xbf16, strided<[1], offset: ?>>,
+// CHECK-SAME: %[[IDX:.+]]: vector<128x128xindex>, %[[MASK:.+]]: vector<128x128xi1>,
+// CHECK-SAME: %[[VAL:.+]]: vector<128x128xbf16>
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-NOT: vector.shape_cast
+// CHECK: vector.scatter %[[SRC]][%[[C0]]] [%[[IDX]]], %[[MASK]], %[[VAL]]
+func.func @scatter_2d_from_flat(%src: memref<?xbf16, strided<[1], offset: ?>>,
+ %idx: vector<128x128xindex>, %mask: vector<128x128xi1>,
+ %val: vector<128x128xbf16>) {
+ %c0 = arith.constant 0 : index
+ %flat_idx = vector.shape_cast %idx : vector<128x128xindex> to vector<16384xindex>
+ %flat_mask = vector.shape_cast %mask : vector<128x128xi1> to vector<16384xi1>
+ %flat_val = vector.shape_cast %val : vector<128x128xbf16> to vector<16384xbf16>
+ vector.scatter %src[%c0] [%flat_idx], %flat_mask, %flat_val
+ : memref<?xbf16, strided<[1], offset: ?>>, vector<16384xindex>, vector<16384xi1>,
+ vector<16384xbf16>
+ return
+}
+
+// -----
+
+// A splat mask / pass-thru is rebuilt at the N-D shape.
+// CHECK-LABEL: @gather_2d_splat_operands
+// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>, %[[P:.+]]: i1
+// CHECK-DAG: %[[MASK:.+]] = vector.broadcast %[[P]] : i1 to vector<8x16xi1>
+// CHECK-DAG: %[[PASS_THRU:.+]] = arith.constant dense<1.000000e+00> : vector<8x16xf32>
+// CHECK: vector.gather {{.*}}, %[[MASK]], %[[PASS_THRU]] {{.*}} into vector<8x16xf32>
+func.func @gather_2d_splat_operands(%src: memref<?xf32>, %idx: vector<8x16xindex>,
+ %p: i1) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.000000e+00> : vector<128xf32>
+ %flat_idx = vector.shape_cast %idx : vector<8x16xindex> to vector<128xindex>
+ %flat_mask = vector.broadcast %p : i1 to vector<128xi1>
+ %flat_res = vector.gather %src[%c0] [%flat_idx], %flat_mask, %cst
+ : memref<?xf32>, vector<128xindex>, vector<128xi1>, vector<128xf32>
+ into vector<128xf32>
+ %res = vector.shape_cast %flat_res : vector<128xf32> to vector<8x16xf32>
+ return %res : vector<8x16xf32>
+}
+
+// -----
+
+// The mask cannot be un-flattened, so the gather is left alone: rewriting it
+// would only move the shape_cast from the index to the mask operand.
+// CHECK-LABEL: @gather_opaque_mask_untouched
+// CHECK: vector.gather {{.*}} into vector<128xf32>
+func.func @gather_opaque_mask_untouched(%src: memref<?xf32>, %idx: vector<8x16xindex>,
+ %mask: vector<128xi1>, %pass_thru: vector<128xf32>) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %flat_idx = vector.shape_cast %idx : vector<8x16xindex> to vector<128xindex>
+ %flat_res = vector.gather %src[%c0] [%flat_idx], %mask, %pass_thru
+ : memref<?xf32>, vector<128xindex>, vector<128xi1>, vector<128xf32>
+ into vector<128xf32>
+ %res = vector.shape_cast %flat_res : vector<128xf32> to vector<8x16xf32>
+ return %res : vector<8x16xf32>
+}
+
+// -----
+
+// The flat result is consumed as-is, so the gather is left alone: rewriting it
+// would only move the shape_casts to the result.
+// CHECK-LABEL: @gather_flat_use_untouched
+// CHECK: vector.gather {{.*}} into vector<128xf32>
+func.func @gather_flat_use_untouched(%src: memref<?xf32>, %idx: vector<8x16xindex>,
+ %mask: vector<8x16xi1>, %pass_thru: vector<128xf32>) -> vector<128xf32> {
+ %c0 = arith.constant 0 : index
+ %flat_idx = vector.shape_cast %idx : vector<8x16xindex> to vector<128xindex>
+ %flat_mask = vector.shape_cast %mask : vector<8x16xi1> to vector<128xi1>
+ %res = vector.gather %src[%c0] [%flat_idx], %flat_mask, %pass_thru
+ : memref<?xf32>, vector<128xindex>, vector<128xi1>, vector<128xf32>
+ into vector<128xf32>
+ return %res : vector<128xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/218292
More information about the Mlir-commits
mailing list