[Mlir-commits] [mlir] [mlir][xegpu] Add xegpu-canonicalize pass to un-flatten gather/scatter (PR #218292)
Jianhui Li
llvmlistbot at llvm.org
Sun Aug 23 14:41:16 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/218292
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.
>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] [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>
+}
More information about the Mlir-commits
mailing list