[Mlir-commits] [mlir] [mlir][xegpu] Add xegpu-canonicalize pass to un-flatten gather/scatter (PR #218292)
Jianhui Li
llvmlistbot at llvm.org
Mon Aug 31 17:59:20 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/218292
>From 2edf8ae2fc2e97b773808ed08ac80e6bf92c8141 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 23 Aug 2026 21:14:54 +0000
Subject: [PATCH 1/4] [mlir][xegpu] Add xegpu-canonicalize pass to un-flatten
gather/scatter
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 asserted in xegpu::LayoutAttr::expandDim, reached via
inferShapeCastSourceLayout during xegpu-propagate-layout.
Add an xegpu-canonicalize pass holding 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 when the 1-D index
operand is a shape_cast of an N-D vector and every other vector operand
can be un-flattened to that same shape:
%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>
becomes
%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>
The N-D shape always comes from the index operand's producer, so the
pattern only ever undoes a flattening that already happened. It bails out
unless every non-index vector operand is reshapable and, for gather,
unless every use of the flat result casts back to the N-D shape -
otherwise the rewrite would merely move the shape_casts to the result.
All checks run before any IR is created, so a rejected match leaves no
dead ops behind.
The pass is added to the GPU-to-XeVM pipeline ahead of
convert-vector-to-xegpu and xegpu-wg-to-sg-distribute, so that both see
the N-D shapes.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 34 +++
.../Dialect/XeGPU/Transforms/Transforms.h | 2 +
.../GPU/Pipelines/GPUToXeVMPipeline.cpp | 1 +
.../Dialect/XeGPU/Transforms/CMakeLists.txt | 1 +
.../XeGPU/Transforms/XeGPUCanonicalize.cpp | 197 ++++++++++++++++++
.../Dialect/XeGPU/xegpu-canonicalize.mlir | 101 +++++++++
6 files changed, 336 insertions(+)
create mode 100644 mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
create mode 100644 mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir
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>
+}
>From 0cdb2ea2868e6e302235fb28f57bc7a03cd22726 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 23 Aug 2026 23:23:06 +0000
Subject: [PATCH 2/4] fix fomrat
---
.../lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp | 10 ++++++----
1 file changed, 6 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
index eb062a3a4f7aa..92b3d677c0946 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
@@ -91,12 +91,14 @@ static Value unflatten(PatternRewriter &rewriter, Value flat,
/// ```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
+/// %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>
+/// %res = vector.shape_cast %flat_res : vector<8192xbf16> to
+/// vector<128x64xbf16>
///
/// // After:
/// %cst = arith.constant dense<0.0> : vector<128x64xbf16>
>From e6ed37f41b95021b0dc7bc606e137c54854a3b73 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 1 Sep 2026 00:26:23 +0000
Subject: [PATCH 3/4] [mlir][VectorToXeGPU] Un-flatten gather/scatter inside
the conversion
Addresses review feedback: this is not a canonicalization, so it should not be
a pass named `xegpu-canonicalize`. Nothing about `vector<128x64>` is more
canonical than `vector<8192>` - restoring the N-D shape is only wanted because
the XeGPU lowering needs it, and on any other target it would be noise.
Move it into `ConvertVectorToXeGPUPass` as a step ahead of the conversion
patterns, next to `promoteAllocasToSLM`, which is there for the same reason:
a target-specific preparation the conversion depends on.
This also fixes a real gap in the previous shape. The transform only ran as a
separate pass injected into the GPU-to-XeVM pipeline, so `mlir-opt
-convert-vector-to-xegpu` on its own still produced an `xegpu.load` over a
flattened index vector. It now happens wherever the conversion runs.
No functional change to the transform itself. The pass, its Passes.td entry,
its `populateXeGPUCanonicalizePatterns` declaration, its CMake entry and its
pipeline injection are all removed, leaving the XeGPU dialect untouched by this
change. The tests move to
`test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir` and now check the
resulting `xegpu.load` / `xegpu.store` shapes rather than an intermediate
`vector.gather`.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/Transforms/Passes.td | 34 ---
.../Dialect/XeGPU/Transforms/Transforms.h | 2 -
.../VectorToXeGPU/VectorToXeGPU.cpp | 162 ++++++++++++++
.../GPU/Pipelines/GPUToXeVMPipeline.cpp | 1 -
.../Dialect/XeGPU/Transforms/CMakeLists.txt | 1 -
.../XeGPU/Transforms/XeGPUCanonicalize.cpp | 199 ------------------
.../unflatten-gather-scatter.mlir} | 114 ++++++----
7 files changed, 237 insertions(+), 276 deletions(-)
delete mode 100644 mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
rename mlir/test/{Dialect/XeGPU/xegpu-canonicalize.mlir => Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir} (50%)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
index acc3a2fd7bf85..36f0a131b345b 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td
@@ -11,40 +11,6 @@
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 a32c8befdd334..388bd6145df21 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -59,8 +59,6 @@ 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/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9863206f14fe1..4529cb0dc7818 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -21,6 +21,7 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -974,6 +975,162 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
}
};
+// 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.
+//
+// Before:
+// %cst = arith.constant dense<0.0> : vector<8192xbf16>
+// %fi = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
+// %fm = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
+// %fr = vector.gather %src[%c0] [%fi], %fm, %cst : memref<?xbf16>,
+// vector<8192xindex>, vector<8192xi1>, vector<8192xbf16>
+// into vector<8192xbf16>
+// %res = vector.shape_cast %fr : 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();
+ }
+};
+
+// Un-flatten every gather/scatter that was flattened to 1-D operands, so that
+// the conversion patterns and the XeGPU layouts downstream of them see the N-D
+// shape of the accessed data.
+static LogicalResult unflattenGatherScatter(Operation *root) {
+ RewritePatternSet patterns(root->getContext());
+ patterns.add<UnflattenGatherScatter<vector::GatherOp>,
+ UnflattenGatherScatter<vector::ScatterOp>>(root->getContext());
+ return applyPatternsGreedily(root, std::move(patterns));
+}
+
// Returns `memrefTy` with its memory space replaced by `newMemSpace`.
static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
return MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
@@ -1054,6 +1211,11 @@ struct ConvertVectorToXeGPUPass
// load_matrix/store_matrix lowerings have well-typed memref operands.
promoteAllocasToSLM(getOperation());
+ // Undo any flattening of gather/scatter operands, so that the conversion
+ // below sees the N-D shape the XeGPU layouts are expressed in.
+ if (failed(unflattenGatherScatter(getOperation())))
+ return signalPassFailure();
+
RewritePatternSet patterns(&getContext());
populateVectorToXeGPUConversionPatterns(patterns);
populatePrepareVectorToMMAPatterns(patterns);
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 3705ad75707d2..4dc9e2acfe235 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -66,7 +66,6 @@ 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 8b53c7699436f..2ed81ae05ab34 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt
@@ -1,7 +1,6 @@
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
deleted file mode 100644
index 92b3d677c0946..0000000000000
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUCanonicalize.cpp
+++ /dev/null
@@ -1,199 +0,0 @@
-//===- 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/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
similarity index 50%
rename from mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir
rename to mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
index 650f334b2046b..06bff7cae192d 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-canonicalize.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
@@ -1,16 +1,12 @@
-// RUN: mlir-opt %s -split-input-file -xegpu-canonicalize | FileCheck %s
+// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | 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: ?>>,
+// Frontends often emit a gather/scatter with 1-D operands, obtained by
+// shape_cast-ing N-D indices and masks. The conversion undoes that first, so
+// the resulting xegpu.load / xegpu.store keeps the N-D shape of the accessed
+// data, which is the shape the XeGPU layouts downstream are expressed in.
+
+gpu.module @xevm_module {
+gpu.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>
@@ -20,19 +16,26 @@ func.func @gather_2d_from_flat(%src: memref<?xbf16, strided<[1], offset: ?>>,
: 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>
+ gpu.return %res : vector<128x64xbf16>
}
-// -----
-
-// CHECK-LABEL: @scatter_2d_from_flat
+// CHECK-LABEL: @gather_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-SAME: %[[IDX:.+]]: vector<128x64xindex>, %[[MASK:.+]]: vector<128x64xi1>
+// The flat pass-thru constant is reshaped rather than shape_cast.
+// CHECK: %[[PASS_THRU:.+]] = arith.constant dense<0.000000e+00> : vector<128x64xbf16>
+// CHECK-NOT: vector.shape_cast
+// CHECK: %[[VEC:.+]] = xegpu.load %{{.+}}[%{{.+}}], %[[MASK]]
+// CHECK-SAME: : i64, vector<128x64xindex>, vector<128x64xi1> -> vector<128x64xbf16>
+// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]]
// 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: ?>>,
+// CHECK: gpu.return %[[RES]] : vector<128x64xbf16>
+}
+
+// -----
+
+gpu.module @xevm_module {
+gpu.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
@@ -42,18 +45,23 @@ func.func @scatter_2d_from_flat(%src: memref<?xbf16, strided<[1], offset: ?>>,
vector.scatter %src[%c0] [%flat_idx], %flat_mask, %flat_val
: memref<?xbf16, strided<[1], offset: ?>>, vector<16384xindex>, vector<16384xi1>,
vector<16384xbf16>
- return
+ gpu.return
+}
+
+// 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-NOT: vector.shape_cast
+// CHECK: xegpu.store %[[VAL]], %{{.+}}[%{{.+}}], %[[MASK]]
+// CHECK-SAME: : vector<128x128xbf16>, i64, vector<128x128xindex>, vector<128x128xi1>
}
// -----
// 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>,
+gpu.module @xevm_module {
+gpu.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>
@@ -63,16 +71,24 @@ func.func @gather_2d_splat_operands(%src: memref<?xf32>, %idx: vector<8x16xindex
: 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>
+ gpu.return %res : vector<8x16xf32>
+}
+
+// CHECK-LABEL: @gather_2d_splat_operands(
+// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>, %[[P:.+]]: i1
+// CHECK-DAG: %[[PASS_THRU:.+]] = arith.constant dense<1.000000e+00> : vector<8x16xf32>
+// CHECK-DAG: %[[MASK:.+]] = vector.broadcast %[[P]] : i1 to vector<8x16xi1>
+// CHECK: %[[VEC:.+]] = xegpu.load %{{.+}}[%[[IDX]]], %[[MASK]]
+// CHECK-SAME: : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
+// CHECK: arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]]
}
// -----
-// The mask cannot be un-flattened, so the gather is left alone: rewriting it
+// The mask cannot be un-flattened, so the gather is left flat: 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>,
+gpu.module @xevm_module {
+gpu.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>
@@ -80,16 +96,25 @@ func.func @gather_opaque_mask_untouched(%src: memref<?xf32>, %idx: vector<8x16xi
: 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>
+ gpu.return %res : vector<8x16xf32>
+}
+
+// CHECK-LABEL: @gather_opaque_mask_untouched(
+// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>,
+// CHECK-SAME: %[[MASK:.+]]: vector<128xi1>, %[[PASS_THRU:.+]]: vector<128xf32>
+// CHECK: %[[FLAT_IDX:.+]] = vector.shape_cast %[[IDX]] : vector<8x16xindex> to vector<128xindex>
+// CHECK: %[[VEC:.+]] = xegpu.load %{{.+}}[%[[FLAT_IDX]]], %[[MASK]]
+// CHECK-SAME: : i64, vector<128xindex>, vector<128xi1> -> vector<128xf32>
+// CHECK: %[[SEL:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]]
+// CHECK: vector.shape_cast %[[SEL]] : vector<128xf32> to vector<8x16xf32>
}
// -----
-// The flat result is consumed as-is, so the gather is left alone: rewriting it
+// The flat result is consumed as-is, so the gather is left flat: 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>,
+gpu.module @xevm_module {
+gpu.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>
@@ -97,5 +122,16 @@ func.func @gather_flat_use_untouched(%src: memref<?xf32>, %idx: vector<8x16xinde
%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>
+ gpu.return %res : vector<128xf32>
+}
+
+// CHECK-LABEL: @gather_flat_use_untouched(
+// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>,
+// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>, %[[PASS_THRU:.+]]: vector<128xf32>
+// CHECK: %[[FLAT_IDX:.+]] = vector.shape_cast %[[IDX]] : vector<8x16xindex> to vector<128xindex>
+// CHECK: %[[FLAT_MASK:.+]] = vector.shape_cast %[[MASK]] : vector<8x16xi1> to vector<128xi1>
+// CHECK: %[[VEC:.+]] = xegpu.load %{{.+}}[%[[FLAT_IDX]]], %[[FLAT_MASK]]
+// CHECK-SAME: : i64, vector<128xindex>, vector<128xi1> -> vector<128xf32>
+// CHECK: %[[RES:.+]] = arith.select %[[FLAT_MASK]], %[[VEC]], %[[PASS_THRU]]
+// CHECK: gpu.return %[[RES]] : vector<128xf32>
}
>From 1449228428e9cd8d14308109fb638c380512533f Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 1 Sep 2026 00:53:08 +0000
Subject: [PATCH 4/4] [mlir][VectorToXeGPU] Simplify un-flattening: cast the
result instead of checking uses
The gather rewrite inspected every use of the flat result and bailed unless all
of them cast back to N-D. Instead, always build the N-D gather and cast its
result back to the flat type. `ShapeCastOp::fold` collapses
`shape_cast(shape_cast(x))` and then drops the resulting no-op, so where a use
did cast back to N-D the pair cancels - within this pass, since the greedy
driver applies folders.
That removes the use scan, the `resultCasts` vector and two bail-outs, and makes
the pattern purely local: it only looks at its own operands.
It also widens the rewrite. A gather whose flat result is genuinely consumed as
1-D used to be skipped; it is now rewritten, trading the two operand casts
(index and mask) for one result cast. So the access becomes N-D and the IR gets
one op lighter, rather than being left alone.
`@gather_flat_use_untouched` never tested the removed check - its pass-thru is an
opaque block argument, so it bailed on the operand. Rename it to
`@gather_opaque_pass_thru_untouched` to say what it actually covers, and add
`@gather_flat_result_use` for the newly rewritten case.
Also soften an overstated comment: layout propagation reasoning through the
surrounding shape_casts is not impossible, it just adds complexity and tends to
yield layouts that lower to unoptimized code.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 25 ++++---------
.../unflatten-gather-scatter.mlir | 37 +++++++++++++++++--
2 files changed, 40 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 4529cb0dc7818..753085a427ef5 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -1030,8 +1030,8 @@ static Value unflatten(PatternRewriter &rewriter, Value flat,
//
// 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.
+// surrounding `vector.shape_cast` ops. That adds complexity and tends to yield
+// layouts that lower to unoptimized code.
//
// Before:
// %cst = arith.constant dense<0.0> : vector<8192xbf16>
@@ -1081,27 +1081,16 @@ struct UnflattenGatherScatter : public OpRewritePattern<OpTy> {
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);
+ // Keep the uses on the flat type instead of inspecting them. Where a use
+ // already casts back to N-D, `ShapeCastOp::fold` collapses the two casts
+ // and then drops the resulting no-op, so nothing is left behind.
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getVectorType(),
+ ndGather);
} else {
if (!canUnflatten(op.getValueToStore(), ndType))
return rewriter.notifyMatchFailure(
diff --git a/mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir b/mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
index 06bff7cae192d..14f2bc64ba268 100644
--- a/mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/unflatten-gather-scatter.mlir
@@ -111,10 +111,10 @@ gpu.func @gather_opaque_mask_untouched(%src: memref<?xf32>, %idx: vector<8x16xin
// -----
-// The flat result is consumed as-is, so the gather is left flat: rewriting it
-// would only move the shape_casts to the result.
+// The pass-thru cannot be un-flattened, so the gather is left flat: rewriting
+// it would only move the shape_cast from the index to the pass-thru operand.
gpu.module @xevm_module {
-gpu.func @gather_flat_use_untouched(%src: memref<?xf32>, %idx: vector<8x16xindex>,
+gpu.func @gather_opaque_pass_thru_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>
@@ -125,7 +125,7 @@ gpu.func @gather_flat_use_untouched(%src: memref<?xf32>, %idx: vector<8x16xindex
gpu.return %res : vector<128xf32>
}
-// CHECK-LABEL: @gather_flat_use_untouched(
+// CHECK-LABEL: @gather_opaque_pass_thru_untouched(
// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>,
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>, %[[PASS_THRU:.+]]: vector<128xf32>
// CHECK: %[[FLAT_IDX:.+]] = vector.shape_cast %[[IDX]] : vector<8x16xindex> to vector<128xindex>
@@ -135,3 +135,32 @@ gpu.func @gather_flat_use_untouched(%src: memref<?xf32>, %idx: vector<8x16xindex
// CHECK: %[[RES:.+]] = arith.select %[[FLAT_MASK]], %[[VEC]], %[[PASS_THRU]]
// CHECK: gpu.return %[[RES]] : vector<128xf32>
}
+
+// -----
+
+// How the result is used does not matter: the rewrite always casts it back to
+// the flat type. Here the use is flat, so that cast stays - but it replaces the
+// two operand casts, so the access still ends up N-D and one op lighter.
+gpu.module @xevm_module {
+gpu.func @gather_flat_result_use(%src: memref<?xf32>, %idx: vector<8x16xindex>,
+ %mask: vector<8x16xi1>) -> vector<128xf32> {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<128xf32>
+ %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, %cst
+ : memref<?xf32>, vector<128xindex>, vector<128xi1>, vector<128xf32>
+ into vector<128xf32>
+ gpu.return %res : vector<128xf32>
+}
+
+// CHECK-LABEL: @gather_flat_result_use(
+// CHECK-SAME: %[[SRC:.+]]: memref<?xf32>, %[[IDX:.+]]: vector<8x16xindex>,
+// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
+// CHECK: %[[PASS_THRU:.+]] = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+// CHECK: %[[VEC:.+]] = xegpu.load %{{.+}}[%[[IDX]]], %[[MASK]]
+// CHECK-SAME: : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
+// CHECK: %[[SEL:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]]
+// CHECK: %[[RES:.+]] = vector.shape_cast %[[SEL]] : vector<8x16xf32> to vector<128xf32>
+// CHECK: gpu.return %[[RES]] : vector<128xf32>
+}
More information about the Mlir-commits
mailing list